Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 35 additions & 31 deletions src/itzi_core/simulation_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -87,11 +87,11 @@ def __init__(
self.dtype = dtype

# Optional components (set via builder methods)
self.raster_input_provider: RasterInputProvider | None = None
self.domain_data: DomainData | None = None
self.raster_output_provider: RasterOutputProvider | None = None
self.vector_output_provider: VectorOutputProvider | None = None
self.mass_balance_output_provider: MassBalanceOutputProvider | None = None
self._raster_input_provider: RasterInputProvider | None = None
self._domain_data: DomainData | None = None
self._raster_output_provider: RasterOutputProvider | None = None
self._vector_output_provider: VectorOutputProvider | None = None
self._mass_balance_output_provider: MassBalanceOutputProvider | None = None

# Hotstart data (set via with_hotstart)
self.hotstart_loader: HotstartLoader | None = None
Expand All @@ -115,28 +115,32 @@ def with_hotstart(
return self

def with_input_provider(self, provider: RasterInputProvider) -> Self:
self.raster_input_provider = provider
self.domain_data = provider.get_domain_data()
if self._domain_data is not None:
raise ValueError("Cannot add an input provider when DomainData is already set.")
self._raster_input_provider = provider
self._domain_data = provider.get_domain_data()
return self

def with_domain_data(self, domain_data: DomainData) -> Self:
"""Set domain data directly for simulations without input provider."""
self.domain_data = domain_data
if self._domain_data is not None:
raise ValueError("DomainData is already set.")
self._domain_data = domain_data
return self

def with_raster_output_provider(self, provider: RasterOutputProvider) -> Self:
self.raster_output_provider = provider
self._raster_output_provider = provider
return self

def with_vector_output_provider(self, provider: VectorOutputProvider) -> Self:
self.vector_output_provider = provider
self._vector_output_provider = provider
return self

def with_mass_balance_output_provider(
self,
provider: MassBalanceOutputProvider,
) -> Self:
self.mass_balance_output_provider = provider
self._mass_balance_output_provider = provider
return self

def _validate_hotstart_congruence(self, hotstart_loader: HotstartLoader) -> None:
Expand Down Expand Up @@ -204,7 +208,7 @@ def _validate_resume_config_congruence(
"Hotstart input map changes are not supported for evolved stage inputs: "
f"{', '.join(sorted(changed_stage_keys))}"
)
if changed_input_keys and self.raster_input_provider is None:
if changed_input_keys and self._raster_input_provider is None:
raise HotstartError(
"Hotstart changed input map names require an input provider for resume"
)
Expand Down Expand Up @@ -235,44 +239,44 @@ def _validate_surface_flow_parameter_congruence(

def _validate_domain_congruence(self, hotstart_domain: DomainData) -> None:
"""Validate that domain metadata matches between hotstart and builder."""
assert self.domain_data is not None # Already validated in build()
assert self._domain_data is not None # Already validated in build()

# Check spatial bounds
if not np.isclose(self.domain_data.north, hotstart_domain.north):
if not np.isclose(self._domain_data.north, hotstart_domain.north):
raise HotstartError(
f"Domain north mismatch: builder={self.domain_data.north}, "
f"Domain north mismatch: builder={self._domain_data.north}, "
f"hotstart={hotstart_domain.north}"
)
if not np.isclose(self.domain_data.south, hotstart_domain.south):
if not np.isclose(self._domain_data.south, hotstart_domain.south):
raise HotstartError(
f"Domain south mismatch: builder={self.domain_data.south}, "
f"Domain south mismatch: builder={self._domain_data.south}, "
f"hotstart={hotstart_domain.south}"
)
if not np.isclose(self.domain_data.east, hotstart_domain.east):
if not np.isclose(self._domain_data.east, hotstart_domain.east):
raise HotstartError(
f"Domain east mismatch: builder={self.domain_data.east}, "
f"Domain east mismatch: builder={self._domain_data.east}, "
f"hotstart={hotstart_domain.east}"
)
if not np.isclose(self.domain_data.west, hotstart_domain.west):
if not np.isclose(self._domain_data.west, hotstart_domain.west):
raise HotstartError(
f"Domain west mismatch: builder={self.domain_data.west}, "
f"Domain west mismatch: builder={self._domain_data.west}, "
f"hotstart={hotstart_domain.west}"
)

# Check dimensions
if self.domain_data.rows != hotstart_domain.rows:
if self._domain_data.rows != hotstart_domain.rows:
raise HotstartError(
f"Domain rows mismatch: builder={self.domain_data.rows}, "
f"Domain rows mismatch: builder={self._domain_data.rows}, "
f"hotstart={hotstart_domain.rows}"
)
if self.domain_data.cols != hotstart_domain.cols:
if self._domain_data.cols != hotstart_domain.cols:
raise HotstartError(
f"Domain cols mismatch: builder={self.domain_data.cols}, "
f"Domain cols mismatch: builder={self._domain_data.cols}, "
f"hotstart={hotstart_domain.cols}"
)

# Check CRS
if self.domain_data.crs_wkt != hotstart_domain.crs_wkt:
if self._domain_data.crs_wkt != hotstart_domain.crs_wkt:
raise HotstartError(
"Domain CRS mismatch: builder and hotstart have different coordinate reference systems."
)
Expand Down Expand Up @@ -312,10 +316,10 @@ def _validate_drainage_congruence(

def build(self) -> Simulation:
"""Build a simulation and explicitly load or prime provider-backed inputs."""
domain_data = self.domain_data
raster_output_provider = self.raster_output_provider
vector_output_provider = self.vector_output_provider
input_provider = self.raster_input_provider
domain_data = self._domain_data
raster_output_provider = self._raster_output_provider
vector_output_provider = self._vector_output_provider
input_provider = self._raster_input_provider
hotstart_loader = self.hotstart_loader

if domain_data is None:
Expand Down Expand Up @@ -363,7 +367,7 @@ def build(self) -> Simulation:
temporal_type=self.sim_config.temporal_type,
raster_output_provider=raster_output_provider,
vector_output_provider=vector_output_provider,
mass_balance_output_provider=self.mass_balance_output_provider,
mass_balance_output_provider=self._mass_balance_output_provider,
out_map_names=self.sim_config.output_map_names,
dt=self.sim_config.record_step,
)
Expand Down
32 changes: 32 additions & 0 deletions tests/test_simulation_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@

from itzi_core.const import TemporalType
from itzi_core.data_containers import SimulationConfig, SurfaceFlowParameters
from itzi_core.domain_data import DomainData
from itzi_core.providers.memory_input import MemoryRasterInputProvider
from itzi_core.simulation_builder import SimulationBuilder


Expand All @@ -45,6 +47,36 @@ def _builder(domain_5by5, config: SimulationConfig) -> SimulationBuilder:
)


def _input_provider(
domain_data: DomainData, config: SimulationConfig
) -> MemoryRasterInputProvider:
return MemoryRasterInputProvider(
{
"domain_data": domain_data,
"simulation_start_time": config.start_time,
"simulation_end_time": config.end_time,
}
)


def test_input_provider_cannot_overwrite_domain_data(domain_5by5, helpers) -> None:
config = _simulation_config(helpers)
builder = _builder(domain_5by5, config)
input_provider = _input_provider(domain_5by5.domain_data, config)

with pytest.raises(ValueError, match="DomainData is already set"):
builder.with_input_provider(input_provider)


def test_domain_data_cannot_overwrite_input_provider_domain_data(domain_5by5, helpers) -> None:
config = _simulation_config(helpers)
input_provider = _input_provider(domain_5by5.domain_data, config)
builder = SimulationBuilder(config, domain_5by5.arr_mask).with_input_provider(input_provider)

with pytest.raises(ValueError, match="DomainData is already set"):
builder.with_domain_data(domain_5by5.domain_data)


def test_configured_raster_outputs_require_a_provider(domain_5by5, helpers) -> None:
config = _simulation_config(helpers).model_copy(
update={"output_map_names": {"water_depth": "depth"}}
Expand Down
Loading