Compute Froude and vmax only when needed

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)

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

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,

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:

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