From 96822ae6338c9f6e50e087e0171fe9b73c069a9b Mon Sep 17 00:00:00 2001 From: Laurent Courty Date: Thu, 24 Sep 2026 14:13:51 -0600 Subject: [PATCH 1/2] Rename discharge buffers --- src/itzi_core/array_definitions.py | 16 +++++++-------- src/itzi_core/report.py | 4 ++-- src/itzi_core/surfaceflow.py | 30 ++++++++++++++-------------- tests/ea8b/helpers.py | 4 ++-- tests/ea8b/test_hotstart.py | 2 +- tests/test_hotstart_integration.py | 6 +++--- tests/test_hotstart_state_loading.py | 10 +++++----- tests/test_hotstart_timed_inputs.py | 2 +- tests/test_steep_plane.py | 8 ++++---- 9 files changed, 41 insertions(+), 41 deletions(-) diff --git a/src/itzi_core/array_definitions.py b/src/itzi_core/array_definitions.py index 0cb346d..bb864d1 100644 --- a/src/itzi_core/array_definitions.py +++ b/src/itzi_core/array_definitions.py @@ -258,41 +258,41 @@ class ArrayDefinition: var_loc="edge", ), ArrayDefinition( - key="old_discharge_east", + key="discharge_east", csdms_name="land_surface_water_flow__x_component_of_z_integral_of_velocity", cf_name="", category=(ArrayCategory.INTERNAL,), - description="Water unit-width discharge at eastern cell edge (previous timestep).", + description="Committed water unit-width discharge at eastern cell edge.", unit="m2 s-1", cf_unit="", var_loc="edge", ), ArrayDefinition( - key="old_discharge_south", + key="discharge_south", csdms_name="land_surface_water_flow__y_component_of_z_integral_of_velocity", cf_name="", category=(ArrayCategory.INTERNAL,), - description="Water unit-width discharge at southern cell edge (previous timestep).", + description="Committed water unit-width discharge at southern cell edge.", unit="m2 s-1", cf_unit="", var_loc="edge", ), ArrayDefinition( - key="new_discharge_east", + key="discharge_east_work", csdms_name="land_surface_water_flow__x_component_of_z_integral_of_velocity", cf_name="", category=(ArrayCategory.INTERNAL,), - description="Water unit-width discharge at eastern cell edge.", + description="Solver workspace for water unit-width discharge at eastern cell edge.", unit="m2 s-1", cf_unit="", var_loc="edge", ), ArrayDefinition( - key="new_discharge_south", + key="discharge_south_work", csdms_name="land_surface_water_flow__y_component_of_z_integral_of_velocity", cf_name="", category=(ArrayCategory.INTERNAL,), - description="Water unit-width discharge at southern cell edge.", + description="Solver workspace for water unit-width discharge at southern cell edge.", unit="m2 s-1", cf_unit="", var_loc="edge", diff --git a/src/itzi_core/report.py b/src/itzi_core/report.py index 2fa2d2b..54ac212 100644 --- a/src/itzi_core/report.py +++ b/src/itzi_core/report.py @@ -165,11 +165,11 @@ def get_output_arrays(self, data: SimulationData) -> dict[str, np.ndarray]: ) elif arr_key == "flow_rate_x": output_arrays[arr_key] = rastermetrics.calculate_flux( - raw["new_discharge_east"], cell_dy + raw["discharge_east_work"], cell_dy ) elif arr_key == "flow_rate_y": output_arrays[arr_key] = rastermetrics.calculate_flux( - raw["new_discharge_south"], cell_dx + raw["discharge_south_work"], cell_dx ) elif arr_key == "created_volume": output_arrays[arr_key] = accum_arrays["error_depth_accum"] * cell_area diff --git a/src/itzi_core/surfaceflow.py b/src/itzi_core/surfaceflow.py index e312e2a..8e89a40 100644 --- a/src/itzi_core/surfaceflow.py +++ b/src/itzi_core/surfaceflow.py @@ -87,7 +87,7 @@ def step( self.arr_err = np.isnan(self.dom.get_array("water_depth")) if np.any(self.arr_err): raise NullError - self.swap_flow_arrays() + self._commit_discharge_buffers() return self def solve_dt(self): @@ -133,8 +133,8 @@ def update_h( """Calculate new water depth, average velocity and Froude number""" solve_h( arr_ext=self.dom.get_padded("ext"), - arr_qe=self.dom.get_padded("new_discharge_east"), - arr_qs=self.dom.get_padded("new_discharge_south"), + arr_qe=self.dom.get_padded("discharge_east_work"), + arr_qs=self.dom.get_padded("discharge_south_work"), arr_bct=self.dom.get_padded("boundary_type"), arr_bcv=self.dom.get_padded("boundary_value"), arr_h=self.dom.get_padded("water_depth"), @@ -159,20 +159,20 @@ def update_h( def solve_q(self): """Solve flow inside the domain using C/Cython function""" - arr_qe_new = self.dom.get_padded("new_discharge_east") - arr_qs_new = self.dom.get_padded("new_discharge_south") + arr_qe_work = self.dom.get_padded("discharge_east_work") + arr_qs_work = self.dom.get_padded("discharge_south_work") arr_boundaries_accum = self.dom.get_padded("boundaries_accum") solve_q( arr_z=self.dom.get_padded("ground_elevation"), arr_n=self.dom.get_padded("friction"), arr_h=self.dom.get_padded("water_depth"), - arr_qe=self.dom.get_padded("old_discharge_east"), - arr_qs=self.dom.get_padded("old_discharge_south"), + arr_qe=self.dom.get_padded("discharge_east"), + arr_qs=self.dom.get_padded("discharge_south"), arr_hfe=self.dom.get_padded("flow_depth_east"), arr_hfs=self.dom.get_padded("flow_depth_south"), arr_bctype=self.dom.get_padded("boundary_type"), - arr_qe_new=arr_qe_new, - arr_qs_new=arr_qs_new, + arr_qe_new=arr_qe_work, + arr_qs_new=arr_qs_work, dt=self._dt, dx=self.dx, dy=self.dy, @@ -183,8 +183,8 @@ def solve_q(self): max_slope=self.max_slope, ) accumulate_boundary_fluxes( - arr_qe_new=arr_qe_new, - arr_qs_new=arr_qs_new, + arr_qe_new=arr_qe_work, + arr_qs_new=arr_qs_work, arr_bcaccum=arr_boundaries_accum, dt=self._dt, dx=self.dx, @@ -192,8 +192,8 @@ def solve_q(self): ) return self - def swap_flow_arrays(self): - """Swap flow arrays from calculated to input""" - self.dom.swap_arrays("old_discharge_east", "new_discharge_east") - self.dom.swap_arrays("old_discharge_south", "new_discharge_south") + def _commit_discharge_buffers(self): + """Promote solver work buffers to the committed discharge state.""" + self.dom.swap_arrays("discharge_east", "discharge_east_work") + self.dom.swap_arrays("discharge_south", "discharge_south_work") return self diff --git a/tests/ea8b/helpers.py b/tests/ea8b/helpers.py index c9b92a1..2e90887 100644 --- a/tests/ea8b/helpers.py +++ b/tests/ea8b/helpers.py @@ -32,8 +32,8 @@ # scheduler semantics keep the resumed run within XPSTORM acceptance while a # few southern-discharge cells can drift slightly above the historical restart tolerances. "water_depth": 7.0e-3, - "old_discharge_east": 2.9e-3, - "old_discharge_south": 1.6e-3, + "discharge_east": 2.9e-3, + "discharge_south": 1.6e-3, } diff --git a/tests/ea8b/test_hotstart.py b/tests/ea8b/test_hotstart.py index 45119d4..7b3908f 100644 --- a/tests/ea8b/test_hotstart.py +++ b/tests/ea8b/test_hotstart.py @@ -121,7 +121,7 @@ def test_ea8b_hotstart_roundtrip( final_state = np.load(final_state_path, allow_pickle=False) - for key in ["water_depth", "old_discharge_east", "old_discharge_south"]: + for key in ["water_depth", "discharge_east", "discharge_south"]: arr_resumed = simulation.raster_domain.get_array(key) arr_uninterrupted = final_state[f"raster_{key}"] np.testing.assert_allclose( diff --git a/tests/test_hotstart_integration.py b/tests/test_hotstart_integration.py index 9c228d3..afdd560 100644 --- a/tests/test_hotstart_integration.py +++ b/tests/test_hotstart_integration.py @@ -174,8 +174,8 @@ def assert_final_state_matches(simulation: Simulation, reference: Simulation) -> "water_depth", "max_water_depth", "max_flow_speed", - "old_discharge_east", - "old_discharge_south", + "discharge_east", + "discharge_south", ]: arr_resumed = simulation.raster_domain.get_array(key) arr_reference = reference.raster_domain.get_array(key) @@ -184,7 +184,7 @@ def assert_final_state_matches(simulation: Simulation, reference: Simulation) -> def assert_state_differs(simulation: Simulation, reference: Simulation) -> None: mismatch_found = False - for key in ["water_depth", "old_discharge_east", "old_discharge_south"]: + for key in ["water_depth", "discharge_east", "discharge_south"]: arr_resumed = simulation.raster_domain.get_array(key) arr_reference = reference.raster_domain.get_array(key) if not np.allclose(arr_resumed, arr_reference): diff --git a/tests/test_hotstart_state_loading.py b/tests/test_hotstart_state_loading.py index 3f31d7e..2d2e7c0 100644 --- a/tests/test_hotstart_state_loading.py +++ b/tests/test_hotstart_state_loading.py @@ -78,11 +78,11 @@ def populated_raster_domain(self, raster_domain, domain_5by5) -> RasterDomain: raster_domain.update_array("ground_elevation", domain_5by5.arr_dem_flat.copy()) raster_domain.update_array("friction", domain_5by5.arr_n.copy()) # Set internal arrays (required for save_state) - # Discharges at the eastern and southern cell edges from the previous step. - old_discharge_east = np.full(raster_domain.shape, 1.1, dtype=np.float32) - old_discharge_south = np.full(raster_domain.shape, 1.2, dtype=np.float32) - raster_domain.update_array("old_discharge_east", old_discharge_east) - raster_domain.update_array("old_discharge_south", old_discharge_south) + # Committed discharges at the eastern and southern cell edges. + discharge_east = np.full(raster_domain.shape, 1.1, dtype=np.float32) + discharge_south = np.full(raster_domain.shape, 1.2, dtype=np.float32) + raster_domain.update_array("discharge_east", discharge_east) + raster_domain.update_array("discharge_south", discharge_south) return raster_domain def test_load_state_round_trip(self, populated_raster_domain: RasterDomain) -> None: diff --git a/tests/test_hotstart_timed_inputs.py b/tests/test_hotstart_timed_inputs.py index 5ca352b..2f4a618 100644 --- a/tests/test_hotstart_timed_inputs.py +++ b/tests/test_hotstart_timed_inputs.py @@ -97,7 +97,7 @@ def _run_to_end(simulation: "Simulation", *, skip_initialize: bool = False) -> N def _assert_final_state_matches(resumed: "Simulation", reference: "Simulation") -> None: - for key in ["water_depth", "old_discharge_east", "old_discharge_south"]: + for key in ["water_depth", "discharge_east", "discharge_south"]: np.testing.assert_allclose( resumed.raster_domain.get_array(key), reference.raster_domain.get_array(key), diff --git a/tests/test_steep_plane.py b/tests/test_steep_plane.py index 263f735..7b4a0fa 100644 --- a/tests/test_steep_plane.py +++ b/tests/test_steep_plane.py @@ -107,8 +107,8 @@ def test_rain_on_steep_plane_uses_capped_downhill_flow(tmp_path): (high_cap, tmp_path / "high_cap.csv"), ]: water_depth = simulation.get_array("water_depth") - eastward_flow = simulation.get_array("old_discharge_east") - southward_flow = simulation.get_array("old_discharge_south") + eastward_flow = simulation.get_array("discharge_east") + southward_flow = simulation.get_array("discharge_south") assert simulation.sim_time == simulation.end_time assert np.all(np.isfinite(water_depth)) @@ -120,7 +120,7 @@ def test_rain_on_steep_plane_uses_capped_downhill_flow(tmp_path): # Both runs solve the same first wet face, so the GMS flux differs only by sqrt(max_slope). center = (4, 4) - low_flux = low_cap.get_array("old_discharge_east")[center] - high_flux = high_cap.get_array("old_discharge_east")[center] + low_flux = low_cap.get_array("discharge_east")[center] + high_flux = high_cap.get_array("discharge_east")[center] assert low_flux > 0 assert high_flux / low_flux == pytest.approx(np.sqrt(0.8 / 0.2), rel=1e-5) From d0c60b759b9aaa1a1102e4fedac5fe4d9a92ed07 Mon Sep 17 00:00:00 2001 From: Laurent Courty Date: Thu, 24 Sep 2026 14:15:07 -0600 Subject: [PATCH 2/2] Fix reported discharge arrays --- src/itzi_core/report.py | 4 ++-- tests/test_report.py | 42 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 2 deletions(-) diff --git a/src/itzi_core/report.py b/src/itzi_core/report.py index 54ac212..352bc95 100644 --- a/src/itzi_core/report.py +++ b/src/itzi_core/report.py @@ -165,11 +165,11 @@ def get_output_arrays(self, data: SimulationData) -> dict[str, np.ndarray]: ) elif arr_key == "flow_rate_x": output_arrays[arr_key] = rastermetrics.calculate_flux( - raw["discharge_east_work"], cell_dy + raw["discharge_east"], cell_dy ) elif arr_key == "flow_rate_y": output_arrays[arr_key] = rastermetrics.calculate_flux( - raw["discharge_south_work"], cell_dx + raw["discharge_south"], cell_dx ) elif arr_key == "created_volume": output_arrays[arr_key] = accum_arrays["error_depth_accum"] * cell_area diff --git a/tests/test_report.py b/tests/test_report.py index 465a85d..787e41a 100644 --- a/tests/test_report.py +++ b/tests/test_report.py @@ -114,6 +114,48 @@ 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) + report = Report( + start_time=start_time, + temporal_type=TemporalType.ABSOLUTE, + raster_output_provider=raster_provider, + vector_output_provider=MemoryVectorOutputProvider(), + mass_balance_output_provider=None, + out_map_names=out_map_names, + dt=timedelta(seconds=1), + ) + data = SimulationData( + sim_time=start_time, + time_step=1.0, + time_steps_counter=1, + continuity_data=CONTINUITY_DATA, + raw_arrays={ + "discharge_east": np.array([[2.0]], dtype=np.float32), + "discharge_south": np.array([[3.0]], dtype=np.float32), + "discharge_east_work": np.array([[11.0]], dtype=np.float32), + "discharge_south_work": np.array([[13.0]], dtype=np.float32), + }, + accumulation_arrays={}, + cell_dx=4.0, + cell_dy=5.0, + drainage_network_attributes=None, + ) + + report.step(data) + + np.testing.assert_array_equal( + raster_provider.output_maps_dict["flow_rate_x"][0][1], + np.array([[10.0]], dtype=np.float32), + ) + np.testing.assert_array_equal( + raster_provider.output_maps_dict["flow_rate_y"][0][1], + np.array([[12.0]], dtype=np.float32), + ) + + def test_drainage_topology_is_written_before_attributes() -> None: start_time = datetime(2000, 1, 1, tzinfo=UTC) raster_provider = MemoryRasterOutputProvider({})