diff --git a/benchmarks/benchmark_surface_flow.py b/benchmarks/benchmark_surface_flow.py index 22580f6..b74ffd6 100644 --- a/benchmarks/benchmark_surface_flow.py +++ b/benchmarks/benchmark_surface_flow.py @@ -12,8 +12,8 @@ GNU Lesser General Public License for more details. """ -import math import datetime +import math import numpy as np import pytest @@ -47,6 +47,11 @@ def gen_eggbox( num_cells_ids = ["1M", "10M"] cell_size_params = [1, 5, 10, 20] cell_size_ids = ["1m", "5m", "10m", "20m"] +DIAGNOSTIC_STEPS: dict[str, tuple[tuple[bool, bool], ...]] = { + "all": ((True, True), (True, True)), + "record_every_2": ((False, False), (True, True)), + "none": ((False, False), (False, False)), +} def setup_eggbox_simulation(num_cells=10_000, cell_size=5): @@ -95,11 +100,14 @@ def setup_eggbox_simulation(num_cells=10_000, cell_size=5): return surface_flow -def benchmark_surface_flow_n_steps(eggbox_simulation, n_steps=10): - for _ in range(n_steps): +def benchmark_surface_flow_n_steps(eggbox_simulation, diagnostic_steps): + for compute_vdir, compute_froude in diagnostic_steps: eggbox_simulation.solve_dt() - eggbox_simulation.step() - return n_steps + eggbox_simulation.step( + compute_vdir=compute_vdir, + compute_froude=compute_froude, + ) + return len(diagnostic_steps) def benchmark_surface_flow_n_seconds(eggbox_simulation, n_seconds=30): @@ -117,11 +125,14 @@ def benchmark_surface_flow_n_seconds(eggbox_simulation, n_seconds=30): @pytest.mark.parametrize( "cell_size", [5], ids=["5m"] ) # Set as parameter to get it in the output json -def test_benchmark_surface_flow_n_steps(benchmark, num_cells, cell_size, n_steps=10): +@pytest.mark.parametrize("diagnostic_mode", DIAGNOSTIC_STEPS) +def test_benchmark_surface_flow_n_steps(benchmark, num_cells, cell_size, diagnostic_mode): """Run the benchmark for a given number of cells and cell size""" eggbox_sim = setup_eggbox_simulation(num_cells=num_cells, cell_size=cell_size) - benchmark(benchmark_surface_flow_n_steps, eggbox_sim, n_steps) - benchmark.extra_info["lattice_updates"] = n_steps * num_cells + diagnostic_steps = DIAGNOSTIC_STEPS[diagnostic_mode] + benchmark(benchmark_surface_flow_n_steps, eggbox_sim, diagnostic_steps) + benchmark.extra_info["lattice_updates"] = len(diagnostic_steps) * num_cells + benchmark.extra_info["diagnostic_mode"] = diagnostic_mode @pytest.mark.parametrize("num_cells", num_cells_params, ids=num_cells_ids) diff --git a/benchmarks/benchmark_update_h.py b/benchmarks/benchmark_update_h.py index 0c8b59b..fbd0520 100644 --- a/benchmarks/benchmark_update_h.py +++ b/benchmarks/benchmark_update_h.py @@ -16,12 +16,12 @@ import numpy as np import pytest - from itzi_core.compute.partial_inertia_h import ( - set_solve_h_tile_size, get_solve_h_tile_size, + set_solve_h_tile_size, solve_h, ) + from itzi_core.data_containers import SurfaceFlowParameters NUM_CELLS_TO_SHAPE: dict[int, tuple[int, int]] = { @@ -39,6 +39,12 @@ (128, 128), (256, 64), ] +DIAGNOSTIC_MODES: dict[str, tuple[bool, bool]] = { + "all": (True, True), + "none": (False, False), + "vdir": (True, False), + "froude": (False, True), +} def full_padded_array(shape: tuple[int, int], fill_value: np.float32) -> np.ndarray: @@ -108,16 +114,30 @@ def setup_update_h_args(num_cells: int) -> tuple: UPDATE_H_TILE_SIZES, ids=[f"tile_{tile_rows}x{tile_cols}" for tile_rows, tile_cols in UPDATE_H_TILE_SIZES], ) -def test_benchmark_update_h(benchmark, num_cells: int, tile_rows: int, tile_cols: int) -> None: +@pytest.mark.parametrize("diagnostic_mode", DIAGNOSTIC_MODES) +def test_benchmark_update_h( + benchmark, + num_cells: int, + tile_rows: int, + tile_cols: int, + diagnostic_mode: str, +) -> None: solve_h_args = setup_update_h_args(num_cells) + compute_vdir, compute_froude = DIAGNOSTIC_MODES[diagnostic_mode] previous_tile_rows, previous_tile_cols = get_solve_h_tile_size() set_solve_h_tile_size(tile_rows, tile_cols) try: - benchmark(solve_h, *solve_h_args) + benchmark( + solve_h, + *solve_h_args, + compute_vdir=compute_vdir, + compute_froude=compute_froude, + ) finally: set_solve_h_tile_size(previous_tile_rows, previous_tile_cols) benchmark.extra_info["lattice_updates"] = num_cells benchmark.extra_info["tile_rows"] = tile_rows benchmark.extra_info["tile_cols"] = tile_cols + benchmark.extra_info["diagnostic_mode"] = diagnostic_mode diff --git a/src/itzi_core/compute/partial_inertia_h.pyx b/src/itzi_core/compute/partial_inertia_h.pyx index ada601b..1f3a25e 100644 --- a/src/itzi_core/compute/partial_inertia_h.pyx +++ b/src/itzi_core/compute/partial_inertia_h.pyx @@ -64,6 +64,8 @@ cdef inline void solve_h_tile( DTYPE_t dy, DTYPE_t dt, DTYPE_t g, + bint compute_vdir, + bint compute_froude, int r_start, int r_end, int c_start, @@ -121,17 +123,19 @@ cdef inline void solve_h_tile( vx = .5 * (ve + vw) vy = .5 * (vs + vn) - # velocity magnitude and direction + # velocity magnitude and maximum v = c_sqrt(vx*vx + vy*vy) # sqrt faster than hypot arr_v[r, c] = v arr_vmax[r, c] = max(v, arr_vmax[r, c]) - vdir = c_atan(-vy, vx) * 180. / PI - # Branchless. Add 360 only to negative numbers - vdir = vdir + 360. * (vdir < 0) - arr_vdir[r, c] = vdir + if compute_vdir: + vdir = c_atan(-vy, vx) * 180. / PI + # Branchless. Add 360 only to negative numbers + vdir = vdir + 360. * (vdir < 0) + arr_vdir[r, c] = vdir # Froude number - use epsilon to avoid division by zero - arr_fr[r, c] = v / c_sqrt(g * fmax(h_new, eps)) * (h_new > 0.) + if compute_froude: + arr_fr[r, c] = v / c_sqrt(g * fmax(h_new, eps)) * (h_new > 0.) @cython.wraparound(False) # Disable negative index check @@ -158,7 +162,10 @@ def solve_h( DTYPE_t dx, DTYPE_t dy, DTYPE_t dt, - DTYPE_t g + DTYPE_t g, + *, + bint compute_vdir=True, + bint compute_froude=True, ): """Update the water depth and max depth Adjust water depth according to in-domain 'boundary' condition @@ -215,6 +222,8 @@ def solve_h( dy=dy, dt=dt, g=g, + compute_vdir=compute_vdir, + compute_froude=compute_froude, r_start=r_start, r_end=r_end, c_start=c_start, diff --git a/src/itzi_core/simulation.py b/src/itzi_core/simulation.py index 3a910a6..2fe445a 100644 --- a/src/itzi_core/simulation.py +++ b/src/itzi_core/simulation.py @@ -190,6 +190,13 @@ def update(self) -> Self: except DtError as e: raise DtError(f"{step_start}: Time-step computation error detected in simulation: {e}") step_end = self.schedule.select_step_end(step_start + self.surface_flow.dt) + is_final_ts = step_end == self.end_time + is_record_due = step_end == self.schedule.deadline("record") + should_write_report = is_record_due or is_final_ts + is_vdir_requested = self.report.out_map_names.get("vdir") is not None + is_froude_requested = self.report.out_map_names.get("froude") is not None + compute_vdir = should_write_report and is_vdir_requested + compute_froude = should_write_report and is_froude_requested # surface flow # # update arrays of infiltration, rainfall etc. @@ -202,9 +209,12 @@ def update(self) -> Self: # surface_flow.step() raise NullError in case of NaN/NULL cell # if this happen, stop simulation try: - self.surface_flow.step() + self.surface_flow.step( + compute_vdir=compute_vdir, + compute_froude=compute_froude, + ) except NullError: - raise NullError(f"{step_start}: Null value detected in simulation, terminating") + raise NullError(f"{step_start}: Null value detected in simulation") # Align timed inputs to the interval end before closing and reporting it # under that time label. Due submodels will consume that label on the next update cycle. @@ -222,9 +232,6 @@ def update(self) -> Self: steps_since_start = self.time_steps_counters["since_start"] + 1 steps_since_report = self.time_steps_counters["since_last_report"] + 1 is_first_ts = step_start == self.start_time - is_final_ts = step_end == self.end_time - is_record_due = step_end == self.schedule.deadline("record") - should_write_report = is_record_due or is_final_ts is_ts_over_threshold = steps_since_report % 200 == 0 is_error_comp_due = is_first_ts or is_ts_over_threshold or should_write_report if is_error_comp_due: @@ -361,7 +368,11 @@ def set_array( return self def get_array(self, arr_id: str) -> np.ndarray: - """Here form BMI interface.""" + """Return an array through the BMI interface. + + Between reports, ``vdir`` and ``froude`` contain their values from the + most recent report step when those outputs are enabled. + """ return self.raster_domain.get_array(arr_id) def get_continuity_data(self) -> ContinuityData: diff --git a/src/itzi_core/surfaceflow.py b/src/itzi_core/surfaceflow.py index 8ede559..8663060 100644 --- a/src/itzi_core/surfaceflow.py +++ b/src/itzi_core/surfaceflow.py @@ -13,15 +13,16 @@ """ from __future__ import annotations + import math from datetime import timedelta from typing import TYPE_CHECKING import numpy as np -from itzi_core.compute.partial_inertia_q import solve_q, accumulate_boundary_fluxes from itzi_core.compute.partial_inertia_h import solve_h -from itzi_core.itzi_error import NullError, DtError +from itzi_core.compute.partial_inertia_q import accumulate_boundary_fluxes, solve_q +from itzi_core.itzi_error import DtError, NullError if TYPE_CHECKING: from itzi_core.data_containers import SurfaceFlowParameters @@ -52,19 +53,26 @@ def __init__( self.dx = domain.dx self.dy = domain.dy self.cell_surf = self.dx * self.dy - - self._dt = None # 1e-6 second self._dt_fudge = timedelta.resolution.total_seconds() + self._dt: float = self._dt_fudge def update_flow_dir(self): """Deprecated.""" return self - def step(self): + def step( + self, + *, + compute_vdir: bool = True, + compute_froude: bool = True, + ): """Run a full simulation time-step""" self.solve_q() - self.update_h() + self.update_h( + compute_vdir=compute_vdir, + compute_froude=compute_froude, + ) # in case of NaN/NULL cells, raise a NullError self.arr_err = np.isnan(self.dom.get_array("water_depth")) if np.any(self.arr_err): @@ -112,7 +120,12 @@ def dt(self, newdt: timedelta): else: self._dt = newdt_s - def update_h(self): + def update_h( + self, + *, + compute_vdir: bool = True, + compute_froude: bool = True, + ): """Calculate new water depth, average velocity and Froude number""" solve_h( arr_ext=self.dom.get_padded("ext"), @@ -134,6 +147,8 @@ def update_h(self): dy=self.dy, dt=self._dt, g=self.g, + compute_vdir=compute_vdir, + compute_froude=compute_froude, ) assert not np.any(self.dom.get_array("water_depth") < 0) return self diff --git a/tests/test_5by5.py b/tests/test_5by5.py index fe8a8c0..358b834 100644 --- a/tests/test_5by5.py +++ b/tests/test_5by5.py @@ -33,6 +33,85 @@ from itzi_core.simulation import Simulation +def _build_diagnostic_simulation( + domain_5by5, + helpers, + output_keys: list[str], + *, + end_seconds: float, + record_seconds: float, + dtmax: float, + initial_depth: np.ndarray | None = None, +) -> Simulation: + start_time = datetime(2000, 1, 1) + sim_config = SimulationConfig( + start_time=start_time, + end_time=start_time + timedelta(seconds=end_seconds), + record_step=timedelta(seconds=record_seconds), + temporal_type=TemporalType.RELATIVE, + input_map_names=helpers.make_input_map_names( + dem="z", + friction="n", + water_depth="start_h", + ), + output_map_names=helpers.make_output_map_names("diagnostics", output_keys), + surface_flow_parameters=SurfaceFlowParameters(hmin=0.0001, dtmax=dtmax, cfl=0.2), + infiltration_model=InfiltrationModelType.NULL, + ) + raster_output = MemoryRasterOutputProvider(sim_config.output_map_names) + simulation = ( + SimulationBuilder(sim_config, domain_5by5.arr_mask, np.float32) + .with_domain_data(domain_5by5.domain_data) + .with_raster_output_provider(raster_output) + .with_vector_output_provider(MemoryVectorOutputProvider()) + .build() + ) + simulation.set_array("dem", domain_5by5.arr_dem_flat.copy()) + simulation.set_array("friction", domain_5by5.arr_n.copy()) + simulation.set_array( + "water_depth", + np.zeros_like(domain_5by5.arr_start_h) if initial_depth is None else initial_depth.copy(), + ) + return simulation + + +def _run_diagnostic_regression(domain_5by5, helpers, *, force_all: bool): + simulation = _build_diagnostic_simulation( + domain_5by5, + helpers, + ["water_depth", "v", "vmax", "vdir", "froude"], + end_seconds=1.0, + record_seconds=0.4, + dtmax=0.3, + initial_depth=domain_5by5.arr_start_h, + ) + if force_all: + scheduled_step = simulation.surface_flow.step + + def always_compute_step(*, compute_vdir: bool, compute_froude: bool): + return scheduled_step(compute_vdir=True, compute_froude=True) + + simulation.surface_flow.step = always_compute_step + + simulation.initialize() + while simulation.sim_time < simulation.end_time: + simulation.update() + + output_maps = simulation.report.raster_provider.output_maps_dict + result = { + "outputs": { + key: [(time, array.copy()) for time, array in output_maps[key]] + for key in ("vdir", "froude") + }, + "water_depth": simulation.get_array("water_depth").copy(), + "v": simulation.get_array("v").copy(), + "vmax": simulation.get_array("vmax").copy(), + "steps": simulation.time_steps_counters["since_start"], + } + simulation.finalize() + return result + + def _run_center_pulse_simulation( domain_5by5, helpers, @@ -95,6 +174,100 @@ def _run_center_pulse_simulation( return final_depth +@pytest.mark.parametrize( + ("diagnostic_keys", "report_flags"), + [ + ([], (False, False)), + (["vdir"], (True, False)), + (["froude"], (False, True)), + (["vdir", "froude"], (True, True)), + ], + ids=["neither", "vdir", "froude", "both"], +) +def test_scheduler_computes_only_requested_report_diagnostics( + domain_5by5, + helpers, + diagnostic_keys: list[str], + report_flags: tuple[bool, bool], +): + """Non-report steps retain diagnostics, including before an off-cadence final report.""" + simulation = _build_diagnostic_simulation( + domain_5by5, + helpers, + ["water_depth", *diagnostic_keys], + end_seconds=10.0, + record_seconds=4.0, + dtmax=3.0, + ) + scheduled_step = simulation.surface_flow.step + calls = [] + + def tracked_step(*, compute_vdir: bool, compute_froude: bool): + step_end = simulation.sim_time + simulation.dt + vdir_before = simulation.get_array("vdir").copy() + froude_before = simulation.get_array("froude").copy() + result = scheduled_step( + compute_vdir=compute_vdir, + compute_froude=compute_froude, + ) + if not compute_vdir: + np.testing.assert_array_equal(simulation.get_array("vdir"), vdir_before) + if not compute_froude: + np.testing.assert_array_equal(simulation.get_array("froude"), froude_before) + calls.append( + ( + step_end - simulation.start_time, + (compute_vdir, compute_froude), + ) + ) + return result + + simulation.surface_flow.step = tracked_step + simulation.initialize() + simulation.get_array("vdir").fill(-123.0) + simulation.get_array("froude").fill(-456.0) + while simulation.sim_time < simulation.end_time: + simulation.update() + simulation.finalize() + + assert calls == [ + (timedelta(seconds=3), (False, False)), + (timedelta(seconds=4), report_flags), + (timedelta(seconds=7), (False, False)), + (timedelta(seconds=8), report_flags), + (timedelta(seconds=10), report_flags), + ] + + output_maps = simulation.report.raster_provider.output_maps_dict + expected_report_times = [ + timedelta(seconds=0), + timedelta(seconds=4), + timedelta(seconds=8), + timedelta(seconds=10), + ] + assert [time for time, _ in output_maps["water_depth"]] == expected_report_times + for key in ("vdir", "froude"): + expected_times = expected_report_times if key in diagnostic_keys else [] + assert [time for time, _ in output_maps[key]] == expected_times + + +def test_lazy_diagnostic_reports_match_always_compute_reference(domain_5by5, helpers): + optimized = _run_diagnostic_regression(domain_5by5, helpers, force_all=False) + reference = _run_diagnostic_regression(domain_5by5, helpers, force_all=True) + + assert optimized["steps"] == reference["steps"] + for key in ("water_depth", "v", "vmax"): + np.testing.assert_allclose(optimized[key], reference[key], rtol=1e-6, atol=1e-7) + for key in ("vdir", "froude"): + optimized_outputs = optimized["outputs"][key] + reference_outputs = reference["outputs"][key] + assert [time for time, _ in optimized_outputs] == [time for time, _ in reference_outputs] + for (_, optimized_array), (_, reference_array) in zip( + optimized_outputs, reference_outputs, strict=True + ): + np.testing.assert_allclose(optimized_array, reference_array, rtol=1e-6, atol=1e-7) + + @pytest.fixture(scope="module") def sim_5by5(domain_5by5, helpers) -> Simulation: """Run a 5x5 simulation for 60s with 30s record step. diff --git a/tests/test_flow.py b/tests/test_flow.py index 9c5838c..51ffcf3 100644 --- a/tests/test_flow.py +++ b/tests/test_flow.py @@ -211,6 +211,71 @@ def test_vectorizable_velocity_calculation(): assert v_optimized == pytest.approx(expected_v) +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize( + ("compute_vdir", "compute_froude"), + [(True, True), (False, False), (True, False), (False, True)], +) +def test_solve_h_optional_diagnostics(dtype, compute_vdir, compute_froude): + """Direction and Froude writes are independent of live velocity updates.""" + shape = (5, 5) + arr_ext = np.zeros(shape, dtype=dtype) + arr_qe = np.full(shape, 0.5, dtype=dtype) + arr_qs = np.full(shape, 0.3, dtype=dtype) + arr_bct = np.zeros(shape, dtype=np.uint8) + arr_bcv = np.zeros(shape, dtype=dtype) + arr_h = np.full(shape, 0.1, dtype=dtype) + arr_hmax = arr_h.copy() + arr_hfix = np.zeros(shape, dtype=dtype) + arr_herr = np.zeros(shape, dtype=dtype) + arr_hfe = np.full(shape, 0.05, dtype=dtype) + arr_hfs = np.full(shape, 0.05, dtype=dtype) + arr_v = np.zeros(shape, dtype=dtype) + arr_vdir = np.full(shape, -123.0, dtype=dtype) + arr_vmax = np.ones(shape, dtype=dtype) + arr_fr = np.full(shape, -456.0, dtype=dtype) + + solve_h( + arr_ext=arr_ext, + arr_qe=arr_qe, + arr_qs=arr_qs, + arr_bct=arr_bct, + arr_bcv=arr_bcv, + arr_h=arr_h, + arr_hmax=arr_hmax, + arr_hfix=arr_hfix, + arr_herr=arr_herr, + arr_hfe=arr_hfe, + arr_hfs=arr_hfs, + arr_v=arr_v, + arr_vdir=arr_vdir, + arr_vmax=arr_vmax, + arr_fr=arr_fr, + dx=1.0, + dy=1.0, + dt=0.1, + g=9.81, + compute_vdir=compute_vdir, + compute_froude=compute_froude, + ) + + expected_v = sqrt(10.0**2 + 6.0**2) + expected_vdir = atan2(-6.0, 10.0) * 180.0 / pi % 360.0 + expected_froude = expected_v / sqrt(9.81 * 0.1) + tolerance = {"rtol": 1e-6, "atol": 1e-6} + + np.testing.assert_allclose(arr_v[1:-1, 1:-1], expected_v, **tolerance) + np.testing.assert_allclose(arr_vmax[1:-1, 1:-1], expected_v, **tolerance) + if compute_vdir: + np.testing.assert_allclose(arr_vdir[1:-1, 1:-1], expected_vdir, **tolerance) + else: + assert np.all(arr_vdir == dtype(-123.0)) + if compute_froude: + np.testing.assert_allclose(arr_fr[1:-1, 1:-1], expected_froude, **tolerance) + else: + assert np.all(arr_fr == dtype(-456.0)) + + def test_solve_h_uses_dx_and_dy_separately_in_flow_divergence(): """The water-depth update must use the x and y cell sizes independently.""" shape = (5, 5)