diff --git a/src/itzi_core/providers/memory_output.py b/src/itzi_core/providers/memory_output.py index c29db03..243f85b 100644 --- a/src/itzi_core/providers/memory_output.py +++ b/src/itzi_core/providers/memory_output.py @@ -25,20 +25,17 @@ class MemoryRasterOutputProvider(RasterOutputProvider): """Save rasters in memory as numpy arrays.""" - def __init__(self, out_map_names: Mapping[str, str]) -> None: - """Initialize output provider with simulation configuration.""" - # user-selected map names. - self.out_map_names = out_map_names - self.output_maps_dict: dict[str, list[tuple[datetime | timedelta, np.ndarray]]] = { - key: [] for key in self.out_map_names - } + def __init__(self) -> None: + self.output_maps_dict: dict[str, list[tuple[datetime | timedelta, np.ndarray]]] = {} def write_arrays( self, array_dict: Mapping[str, np.ndarray], sim_time: datetime | timedelta ) -> None: for arr_key, arr in array_dict.items(): if isinstance(arr, np.ndarray): - self.output_maps_dict[arr_key].append((deepcopy(sim_time), arr.copy())) + self.output_maps_dict.setdefault(arr_key, []).append( + (deepcopy(sim_time), arr.copy()) + ) class MemoryVectorOutputProvider(VectorOutputProvider): diff --git a/tests/ea8b/conftest.py b/tests/ea8b/conftest.py index 022f7ee..8fee5be 100644 --- a/tests/ea8b/conftest.py +++ b/tests/ea8b/conftest.py @@ -173,7 +173,7 @@ def ea8b_simulation(ea8b_data, test_data_path, ea8b_temp_path): simulation_end_time=sim_config.end_time, ) ) - raster_output_provider = MemoryRasterOutputProvider(sim_config.output_map_names) + raster_output_provider = MemoryRasterOutputProvider() vector_output_provider = MemoryVectorOutputProvider() simulation = ( diff --git a/tests/ea8b/helpers.py b/tests/ea8b/helpers.py index 2e90887..71c01e8 100644 --- a/tests/ea8b/helpers.py +++ b/tests/ea8b/helpers.py @@ -96,7 +96,7 @@ def build_resumed_simulation( simulation_end_time=sim_config.end_time, ) ) - raster_output_provider = MemoryRasterOutputProvider(sim_config.output_map_names) + raster_output_provider = MemoryRasterOutputProvider() vector_output_provider = MemoryVectorOutputProvider() simulation = ( diff --git a/tests/test_5by5.py b/tests/test_5by5.py index ef7a99f..28163e5 100644 --- a/tests/test_5by5.py +++ b/tests/test_5by5.py @@ -58,7 +58,7 @@ def _build_diagnostic_simulation( surface_flow_parameters=SurfaceFlowParameters(hmin=0.0001, dtmax=dtmax, cfl=0.2), infiltration_model=InfiltrationModelType.NULL, ) - raster_output = MemoryRasterOutputProvider(sim_config.output_map_names) + raster_output = MemoryRasterOutputProvider() simulation = ( SimulationBuilder(sim_config, domain_5by5.arr_mask, np.float32) .with_domain_data(domain_5by5.domain_data) @@ -161,7 +161,7 @@ def _run_center_pulse_simulation( infiltration_model=InfiltrationModelType.NULL, ) - raster_output = MemoryRasterOutputProvider(sim_config.output_map_names) + raster_output = MemoryRasterOutputProvider() simulation = ( SimulationBuilder(sim_config, domain_5by5.arr_mask, np.float32) .with_domain_data(domain_data) @@ -322,7 +322,7 @@ def sim_5by5(domain_5by5, helpers) -> Simulation: ) # Create output provider - raster_output = MemoryRasterOutputProvider(sim_config.output_map_names) + raster_output = MemoryRasterOutputProvider() # Build simulation simulation = ( diff --git a/tests/test_analytic.py b/tests/test_analytic.py index 0e88394..173a989 100644 --- a/tests/test_analytic.py +++ b/tests/test_analytic.py @@ -155,7 +155,7 @@ def mcdo_norain_sim(test_data_path, test_data_temp_path): ), surface_flow_parameters=SurfaceFlowParameters(dtmax=2, cfl=0.5), ) - raster_output = MemoryRasterOutputProvider(config.output_map_names) + raster_output = MemoryRasterOutputProvider() simulation = ( SimulationBuilder(config, array_mask, np.float32) .with_domain_data(domain_data) @@ -285,7 +285,7 @@ def mcdo_rain_sim(test_data_path, test_data_temp_path): surface_flow_parameters=SurfaceFlowParameters(dtmax=2, cfl=0.5), dtinf=1, ) - raster_output = MemoryRasterOutputProvider(config.output_map_names) + raster_output = MemoryRasterOutputProvider() simulation = ( SimulationBuilder(config, array_mask, np.float32) .with_domain_data(domain_data) diff --git a/tests/test_boundaries.py b/tests/test_boundaries.py index 439f8a4..3944905 100644 --- a/tests/test_boundaries.py +++ b/tests/test_boundaries.py @@ -57,7 +57,7 @@ def sim_5by5_open_boundaries(domain_5by5, helpers): ) # Create output provider - raster_output = MemoryRasterOutputProvider(sim_config.output_map_names) + raster_output = MemoryRasterOutputProvider() # Build simulation simulation = ( diff --git a/tests/test_ea8a.py b/tests/test_ea8a.py index 2a51d80..849de85 100644 --- a/tests/test_ea8a.py +++ b/tests/test_ea8a.py @@ -314,7 +314,7 @@ def ea_test8a_sim(ea_test8a_xarray_data, test_data_path, test_data_temp_path): ) # Create output providers - raster_output = MemoryRasterOutputProvider(sim_config.output_map_names) + raster_output = MemoryRasterOutputProvider() vector_output = MemoryVectorOutputProvider() # Create simulation domain data from input provider diff --git a/tests/test_hotstart_integration.py b/tests/test_hotstart_integration.py index afdd560..94aae6a 100644 --- a/tests/test_hotstart_integration.py +++ b/tests/test_hotstart_integration.py @@ -94,9 +94,7 @@ def build_simulation( domain_5by5: Domain fixture hotstart_bytes: Optional hotstart archive bytes """ - raster_output = raster_output_provider or MemoryRasterOutputProvider( - sim_config.output_map_names - ) + raster_output = raster_output_provider or MemoryRasterOutputProvider() builder = ( SimulationBuilder(sim_config, domain_5by5.arr_mask, np.float32) @@ -384,7 +382,7 @@ def test_resume_allows_output_provider_change( sim_config = baseline_hotstart_run["sim_config"] hotstart_bytes = baseline_hotstart_run["checkpoints"]["split_30"]["hotstart_bytes"] - resumed_output = MemoryRasterOutputProvider(sim_config.output_map_names) + resumed_output = MemoryRasterOutputProvider() sim_b = build_simulation( sim_config, domain_5by5, @@ -414,7 +412,7 @@ def test_resume_allows_output_map_name_change( ["water_depth", "max_water_depth", "flow_rate_x", "flow_rate_y", "created_volume"], ) sim_b_config = sim_a_config.model_copy(update={"output_map_names": resumed_output_map_names}) - resumed_output = MemoryRasterOutputProvider(resumed_output_map_names) + resumed_output = MemoryRasterOutputProvider() sim_b = build_simulation( sim_b_config, domain_5by5, @@ -425,7 +423,6 @@ def test_resume_allows_output_map_name_change( run_simulation_to_end(sim_b, skip_initialize=True) assert sim_b.report.out_map_names == resumed_output_map_names - assert resumed_output.out_map_names == resumed_output_map_names assert resumed_output.output_maps_dict["water_depth"] assert resumed_output.output_maps_dict["max_water_depth"] np.testing.assert_allclose( @@ -466,7 +463,7 @@ def test_resume_applies_new_record_step_cadence( resumed_record_step = timedelta(seconds=10) resumed_config = original_config.model_copy(update={"record_step": resumed_record_step}) - resumed_output = MemoryRasterOutputProvider(resumed_config.output_map_names) + resumed_output = MemoryRasterOutputProvider() sim_b = build_simulation( resumed_config, domain_5by5, diff --git a/tests/test_hotstart_state_loading.py b/tests/test_hotstart_state_loading.py index 2d2e7c0..eb1d258 100644 --- a/tests/test_hotstart_state_loading.py +++ b/tests/test_hotstart_state_loading.py @@ -196,7 +196,7 @@ class TestSimulationBuilderHotstart: @staticmethod def _create_hotstart_bytes(domain_5by5, sim_config: SimulationConfig) -> io.BytesIO: - raster_output = MemoryRasterOutputProvider(sim_config.output_map_names) + raster_output = MemoryRasterOutputProvider() simulation = ( SimulationBuilder(sim_config, domain_5by5.arr_mask, np.float32) .with_domain_data(domain_5by5.domain_data) @@ -217,7 +217,7 @@ def _build_with_hotstart( hotstart_bytes: io.BytesIO, domain_data: DomainData | None = None, ): - raster_output = MemoryRasterOutputProvider(sim_config.output_map_names) + raster_output = MemoryRasterOutputProvider() builder = ( SimulationBuilder(sim_config, domain_5by5.arr_mask, np.float32) .with_domain_data(domain_data or domain_5by5.domain_data) @@ -274,7 +274,7 @@ def test_with_hotstart_from_bytes( valid_hotstart_bytes: io.BytesIO, ) -> None: """Builder should store HotstartLoader from BytesIO.""" - raster_output = MemoryRasterOutputProvider(sim_config.output_map_names) + raster_output = MemoryRasterOutputProvider() builder = ( SimulationBuilder(sim_config, domain_5by5.arr_mask, np.float32) .with_domain_data(domain_5by5.domain_data) @@ -298,7 +298,7 @@ def test_with_hotstart_from_file( valid_hotstart_bytes.seek(0) hotstart_file.write_bytes(valid_hotstart_bytes.read()) - raster_output = MemoryRasterOutputProvider(sim_config.output_map_names) + raster_output = MemoryRasterOutputProvider() builder = ( SimulationBuilder(sim_config, domain_5by5.arr_mask, np.float32) .with_domain_data(domain_5by5.domain_data) @@ -607,7 +607,7 @@ def test_build_rejects_drainage_mismatch_hotstart_has_drainage( swmm_inp=None, # No drainage ) - raster_output = MemoryRasterOutputProvider(config_no_drainage.output_map_names) + raster_output = MemoryRasterOutputProvider() with pytest.raises( HotstartError, match="Hotstart contains drainage state but current configuration" @@ -672,7 +672,7 @@ def test_build_rejects_drainage_mismatch_config_has_drainage( swmm_inp="fake.inp", # Has drainage ) - raster_output = MemoryRasterOutputProvider(config_with_drainage.output_map_names) + raster_output = MemoryRasterOutputProvider() with pytest.raises( HotstartError, match="Hotstart has no drainage state but current configuration" diff --git a/tests/test_hotstart_timed_inputs.py b/tests/test_hotstart_timed_inputs.py index 2f4a618..7ef6a02 100644 --- a/tests/test_hotstart_timed_inputs.py +++ b/tests/test_hotstart_timed_inputs.py @@ -237,7 +237,7 @@ def _build_provider_simulation( builder = ( SimulationBuilder(sim_config, domain_5by5.arr_mask, np.float32) .with_input_provider(input_provider) - .with_raster_output_provider(MemoryRasterOutputProvider(sim_config.output_map_names)) + .with_raster_output_provider(MemoryRasterOutputProvider()) .with_vector_output_provider(MemoryVectorOutputProvider()) ) if hotstart_bytes is not None: @@ -282,7 +282,7 @@ def _build_map_aware_provider_simulation( builder = ( SimulationBuilder(sim_config, domain_5by5.arr_mask, np.float32) .with_input_provider(provider) - .with_raster_output_provider(MemoryRasterOutputProvider(sim_config.output_map_names)) + .with_raster_output_provider(MemoryRasterOutputProvider()) .with_vector_output_provider(MemoryVectorOutputProvider()) ) if hotstart_bytes is not None: @@ -544,7 +544,7 @@ def test_resume_rejects_changed_input_map_without_provider(domain_5by5) -> None: builder = ( SimulationBuilder(changed_config, domain_5by5.arr_mask, np.float32) .with_domain_data(domain_5by5.domain_data) - .with_raster_output_provider(MemoryRasterOutputProvider(changed_config.output_map_names)) + .with_raster_output_provider(MemoryRasterOutputProvider()) .with_vector_output_provider(MemoryVectorOutputProvider()) .with_hotstart(hotstart_bytes) ) @@ -616,7 +616,7 @@ def test_resume_rejects_pending_timed_input_without_provider(domain_5by5) -> Non builder = ( SimulationBuilder(sim_config, domain_5by5.arr_mask, np.float32) .with_domain_data(domain_5by5.domain_data) - .with_raster_output_provider(MemoryRasterOutputProvider(sim_config.output_map_names)) + .with_raster_output_provider(MemoryRasterOutputProvider()) .with_vector_output_provider(MemoryVectorOutputProvider()) .with_hotstart(hotstart_bytes) ) diff --git a/tests/test_max_values.py b/tests/test_max_values.py index 36140c6..2759c00 100644 --- a/tests/test_max_values.py +++ b/tests/test_max_values.py @@ -61,7 +61,7 @@ def sim_5by5_max_values(domain_5by5, helpers) -> Simulation: ) # Create output provider - raster_output = MemoryRasterOutputProvider(sim_config.output_map_names) + raster_output = MemoryRasterOutputProvider() # Build simulation simulation = ( diff --git a/tests/test_rainfall_precision.py b/tests/test_rainfall_precision.py index ef9605f..8856a37 100644 --- a/tests/test_rainfall_precision.py +++ b/tests/test_rainfall_precision.py @@ -126,7 +126,7 @@ def _run_rain_only_simulation( simulation = ( SimulationBuilder(config, np.zeros(shape, dtype=np.bool_), dtype) .with_domain_data(domain_data) - .with_raster_output_provider(MemoryRasterOutputProvider(config.output_map_names)) + .with_raster_output_provider(MemoryRasterOutputProvider()) .with_vector_output_provider(MemoryVectorOutputProvider()) .with_mass_balance_output_provider(mass_balance_output) .build() @@ -175,7 +175,7 @@ def _run_combined_removal_simulation( infiltration_model=infiltration_model, ) mass_balance_output = CaptureMassBalanceOutputProvider() - raster_output = MemoryRasterOutputProvider(output_map_names) + raster_output = MemoryRasterOutputProvider() dtype = np.float32 simulation = ( SimulationBuilder(config, np.zeros(shape, dtype=np.bool_), dtype) diff --git a/tests/test_report.py b/tests/test_report.py index 787e41a..279eb07 100644 --- a/tests/test_report.py +++ b/tests/test_report.py @@ -38,7 +38,7 @@ def test_get_output_arrays_returns_a_fresh_selection() -> None: start_time = datetime(2000, 1, 1, tzinfo=UTC) out_map_names = {"water_depth": "depth", "max_water_depth": "depth_max"} - raster_provider = MemoryRasterOutputProvider(out_map_names) + raster_provider = MemoryRasterOutputProvider() report = Report( start_time=start_time, temporal_type=TemporalType.ABSOLUTE, @@ -74,7 +74,7 @@ def test_get_output_arrays_returns_a_fresh_selection() -> None: def test_maxima_are_selected_independently_of_base_arrays() -> None: start_time = datetime(2000, 1, 1, tzinfo=UTC) out_map_names = {"max_water_depth": "depth_max", "max_flow_speed": "speed_max"} - raster_provider = MemoryRasterOutputProvider(out_map_names) + raster_provider = MemoryRasterOutputProvider() report = Report( start_time=start_time, temporal_type=TemporalType.ABSOLUTE, @@ -117,7 +117,7 @@ def test_maxima_are_selected_independently_of_base_arrays() -> None: def test_flow_rate_outputs_use_committed_discharges() -> None: start_time = datetime(2000, 1, 1, tzinfo=UTC) out_map_names = {"flow_rate_x": "flow_x", "flow_rate_y": "flow_y"} - raster_provider = MemoryRasterOutputProvider(out_map_names) + raster_provider = MemoryRasterOutputProvider() report = Report( start_time=start_time, temporal_type=TemporalType.ABSOLUTE, @@ -158,7 +158,7 @@ def test_flow_rate_outputs_use_committed_discharges() -> None: def test_drainage_topology_is_written_before_attributes() -> None: start_time = datetime(2000, 1, 1, tzinfo=UTC) - raster_provider = MemoryRasterOutputProvider({}) + raster_provider = MemoryRasterOutputProvider() vector_provider = MemoryVectorOutputProvider() report = Report( start_time=start_time, diff --git a/tests/test_stats.py b/tests/test_stats.py index 1417f7d..af4312f 100644 --- a/tests/test_stats.py +++ b/tests/test_stats.py @@ -104,7 +104,7 @@ def sim_5by5_stats(domain_5by5, helpers, tmp_path_factory): ) # Create output provider - raster_output = MemoryRasterOutputProvider(sim_config.output_map_names) + raster_output = MemoryRasterOutputProvider() # Build simulation simulation = ( @@ -297,7 +297,7 @@ def _run_timed_stats_simulation( simulation = ( SimulationBuilder(sim_config, domain_5by5.arr_mask, np.float32) .with_input_provider(input_provider) - .with_raster_output_provider(MemoryRasterOutputProvider(sim_config.output_map_names)) + .with_raster_output_provider(MemoryRasterOutputProvider()) .with_vector_output_provider(MemoryVectorOutputProvider()) .with_mass_balance_output_provider(CSVMassBalanceOutputProvider(file_name=str(stats_file))) .build() diff --git a/tests/test_steep_plane.py b/tests/test_steep_plane.py index a265f7c..42c0fbb 100644 --- a/tests/test_steep_plane.py +++ b/tests/test_steep_plane.py @@ -70,7 +70,7 @@ def _run_steep_plane(max_slope: float, stats_file: Path): infiltration_model=InfiltrationModelType.NULL, ) array_mask = np.zeros((rows, cols), dtype=np.bool_) - raster_output = MemoryRasterOutputProvider(config.output_map_names) + raster_output = MemoryRasterOutputProvider() simulation = ( SimulationBuilder(config, array_mask, np.float32) .with_domain_data(domain_data) @@ -142,7 +142,7 @@ def _build_regime_switch_simulation( cols=cols, crs_wkt="", ) - raster_output = MemoryRasterOutputProvider(config.output_map_names) + raster_output = MemoryRasterOutputProvider() builder = ( SimulationBuilder(config, np.zeros((rows, cols), dtype=np.bool_), np.float32) .with_domain_data(domain_data) diff --git a/tests/test_wse.py b/tests/test_wse.py index 9580e55..496e27d 100644 --- a/tests/test_wse.py +++ b/tests/test_wse.py @@ -61,7 +61,7 @@ def sim_5by5_wse(domain_5by5, helpers) -> Simulation: ) # Create output provider - raster_output = MemoryRasterOutputProvider(sim_config.output_map_names) + raster_output = MemoryRasterOutputProvider() # Build simulation simulation = ( @@ -185,7 +185,7 @@ def test_timed_memory_input_updates_water_depth_from_wse(domain_5by5) -> None: }, } ) - raster_output = MemoryRasterOutputProvider(sim_config.output_map_names) + raster_output = MemoryRasterOutputProvider() simulation = ( SimulationBuilder(sim_config, domain_5by5.arr_mask, np.float32) .with_input_provider(input_provider)