hmax and vmax as normal output maps

benchmarks/benchmark_snippets.py

@@ -22,6 +22,11 @@
 num_cells_params = [10_000_000]
 num_cells_ids = ["10M"]
 
+maximum_shapes = {
+    "1M": (1_000, 1_000),
+    "10M": (2_000, 5_000),
+}
+
 
 def make_array_shape(num_cells: int) -> tuple[int, int]:
     side_length = int(math.sqrt(num_cells))
@@ -44,6 +49,59 @@ def annotate_math_benchmark(benchmark, *, num_cells: int, formula: str, math_pat
     benchmark.extra_info["math_path"] = math_path
 
 
+def maximum_numpy(arr_maximum: np.ndarray, arr_values: np.ndarray) -> None:
+    np.maximum(arr_maximum, arr_values, out=arr_maximum)
+
+
+def setup_maximum_args(
+    shape: tuple[int, int], padded_domain: bool
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    rng = np.random.default_rng(72)
+    if padded_domain:
+        storage_shape = (shape[0] + 2, shape[1] + 2)
+        initial = rng.random(size=storage_shape, dtype=np.float32)
+        values = rng.random(size=storage_shape, dtype=np.float32)
+        arr_maximum = np.empty(storage_shape, dtype=np.float32)
+    else:
+        initial = rng.random(size=shape, dtype=np.float32)
+        values = rng.random(size=shape, dtype=np.float32)
+        arr_maximum = np.empty(shape, dtype=np.float32)
+    return initial, values, arr_maximum
+
+
+def run_maximum_benchmark(
+    benchmark,
+    implementation,
+    shape: tuple[int, int],
+    padded_domain: bool,
+    use_interior_view: bool,
+) -> None:
+    initial, values, arr_maximum = setup_maximum_args(shape, padded_domain)
+    if padded_domain and use_interior_view:
+        initial = initial[1:-1, 1:-1]
+        values = values[1:-1, 1:-1]
+        arr_maximum = arr_maximum[1:-1, 1:-1]
+    expected = np.maximum(initial, values)
+
+    np.copyto(arr_maximum, initial)
+    implementation(arr_maximum, values)
+    np.testing.assert_array_equal(arr_maximum, expected)
+
+    def setup():
+        np.copyto(arr_maximum, initial)
+        return (arr_maximum, values), {}
+
+    benchmark.pedantic(implementation, setup=setup, rounds=7, iterations=1)
+    benchmark.extra_info["lattice_updates"] = shape[0] * shape[1]
+    benchmark.extra_info["domain_layout"] = "padded" if padded_domain else "packed"
+    if not padded_domain:
+        benchmark.extra_info["call_layout"] = "packed"
+    elif use_interior_view:
+        benchmark.extra_info["call_layout"] = "interior_view"
+    else:
+        benchmark.extra_info["call_layout"] = "padded_storage"
+
+
 def setup_almeida_args(num_cells: int) -> tuple[object, ...]:
     arr_shape = make_array_shape(num_cells)
     rng = np.random.default_rng(42)
@@ -85,6 +143,39 @@ def setup_velocity_diagnostics_args(num_cells: int) -> tuple[object, ...]:
     return arr_qx, arr_qy, arr_h, arr_v, arr_vdir, arr_fr, g
 
 
+## element-wise maximum ##
+
+
+@pytest.mark.parametrize("shape", maximum_shapes.values(), ids=maximum_shapes.keys())
+@pytest.mark.parametrize("padded_domain", [False, True], ids=["packed", "padded_domain"])
+def test_benchmark_maximum_numpy(benchmark, shape, padded_domain):
+    run_maximum_benchmark(benchmark, maximum_numpy, shape, padded_domain, use_interior_view=True)
+
+
+@pytest.mark.parametrize("shape", maximum_shapes.values(), ids=maximum_shapes.keys())
+@pytest.mark.parametrize("padded_domain", [False, True], ids=["packed", "padded_domain"])
+def test_benchmark_maximum_cython_serial(benchmark, shape, padded_domain):
+    run_maximum_benchmark(
+        benchmark,
+        snippets.arr_maximum_serial,
+        shape,
+        padded_domain,
+        use_interior_view=False,
+    )
+
+
+@pytest.mark.parametrize("shape", maximum_shapes.values(), ids=maximum_shapes.keys())
+@pytest.mark.parametrize("padded_domain", [False, True], ids=["packed", "padded_domain"])
+def test_benchmark_maximum_cython_parallel(benchmark, shape, padded_domain):
+    run_maximum_benchmark(
+        benchmark,
+        snippets.arr_maximum_parallel,
+        shape,
+        padded_domain,
+        use_interior_view=False,
+    )
+
+
 ## velocity ##
 
 

src/itzi_core/array_definitions.py

@@ -293,7 +293,7 @@ class ArrayDefinition:
         key="hmax",
         csdms_name="land_surface_water__max_of_depth",
         cf_name="",
-        category=[ArrayCategory.INTERNAL],  # output together with depth arrays
+        category=[ArrayCategory.INTERNAL, ArrayCategory.OUTPUT],
         description="Maximum water depth reached since the beginning of the simulation.",
         unit="m",
         cf_unit="",
@@ -323,7 +323,7 @@ class ArrayDefinition:
         key="vmax",
         csdms_name="land_surface_water_flow__max_of_speed",
         cf_name="",
-        category=[ArrayCategory.INTERNAL],  # output together with speed arrays
+        category=[ArrayCategory.INTERNAL, ArrayCategory.OUTPUT],
         description="Maximum water speed reached since the beginning of the simulation.",
         unit="m s-1",
         cf_unit="",

src/itzi_core/compute/snippets.pyx

@@ -44,6 +44,50 @@ cdef float DEG_360_F = 360.0
 cdef float EPS_F = 1e-12
 
 
+@cython.wraparound(False)
+@cython.boundscheck(False)
+@cython.initializedcheck(False)
+@cython.nonecheck(False)
+def arr_maximum_serial(
+    DTYPE_t[:, ::1] arr_maximum,
+    DTYPE_t[:, ::1] arr_values,
+):
+    """Update an element-wise maximum with a serial loop for benchmarking."""
+    cdef Py_ssize_t rows = arr_maximum.shape[0]
+    cdef Py_ssize_t cols = arr_maximum.shape[1]
+    cdef Py_ssize_t row_idx, col_idx
+
+    if rows != arr_values.shape[0] or cols != arr_values.shape[1]:
+        raise ValueError("Maximum arrays must have the same shape")
+
+    for row_idx in range(rows):
+        for col_idx in range(cols):
+            if arr_values[row_idx, col_idx] > arr_maximum[row_idx, col_idx]:
+                arr_maximum[row_idx, col_idx] = arr_values[row_idx, col_idx]
+
+
+@cython.wraparound(False)
+@cython.boundscheck(False)
+@cython.initializedcheck(False)
+@cython.nonecheck(False)
+def arr_maximum_parallel(
+    DTYPE_t[:, ::1] arr_maximum,
+    DTYPE_t[:, ::1] arr_values,
+):
+    """Update an element-wise maximum with OpenMP for benchmarking."""
+    cdef Py_ssize_t rows = arr_maximum.shape[0]
+    cdef Py_ssize_t cols = arr_maximum.shape[1]
+    cdef Py_ssize_t row_idx, col_idx
+
+    if rows != arr_values.shape[0] or cols != arr_values.shape[1]:
+        raise ValueError("Maximum arrays must have the same shape")
+
+    for row_idx in prange(rows, nogil=True, schedule="static"):
+        for col_idx in range(cols):
+            if arr_values[row_idx, col_idx] > arr_maximum[row_idx, col_idx]:
+                arr_maximum[row_idx, col_idx] = arr_values[row_idx, col_idx]
+
+
 @cython.wraparound(False)  # Disable negative index check
 @cython.cdivision(True)  # Don't check division by zero
 @cython.boundscheck(False)  # turn off bounds-checking for entire function

src/itzi_core/providers/base.py

@@ -23,7 +23,7 @@
 
     import numpy as np
 
-    from itzi_core.data_containers import DrainageNetworkData, MassBalanceData, SimulationData
+    from itzi_core.data_containers import DrainageNetworkData, MassBalanceData
     from itzi_core.providers.domain_data import DomainData
 
 
@@ -61,9 +61,8 @@ def write_arrays(
     ) -> None:
         """Write all arrays for the current time step."""
 
-    @abstractmethod
-    def finalize(self, final_data: SimulationData) -> None:
-        """Finalize outputs and cleanup."""
+    def finalize(self) -> None:
+        """Flush and close provider resources."""
 
 
 class VectorOutputProvider(ABC):
@@ -75,9 +74,8 @@ def write_vector(
     ) -> None:
         """Write simulation data for current time step."""
 
-    @abstractmethod
-    def finalize(self, drainage_data: DrainageNetworkData) -> None:
-        """Finalize outputs and cleanup."""
+    def finalize(self) -> None:
+        """Flush and close provider resources."""
 
 
 class MassBalanceOutputProvider(ABC):

src/itzi_core/providers/csv_output.py

@@ -14,17 +14,21 @@
 
 from __future__ import annotations
 
+import csv
 from datetime import datetime, timedelta
-from typing import TypedDict, TYPE_CHECKING, Tuple, List
 from io import StringIO
 from pathlib import PurePosixPath, PureWindowsPath
-import csv
+from typing import TYPE_CHECKING, TypedDict
 
 import pandas as pd
 
+from itzi_core.data_containers import (
+    DrainageLinkAttributes,
+    DrainageLinkData,
+    DrainageNodeAttributes,
+    DrainageNodeData,
+)
 from itzi_core.providers.base import VectorOutputProvider
-from itzi_core.data_containers import DrainageLinkData, DrainageLinkAttributes
-from itzi_core.data_containers import DrainageNodeData, DrainageNodeAttributes
 
 if TYPE_CHECKING:
     from itzi_core.data_containers import DrainageNetworkData
@@ -117,9 +121,6 @@ def write_vector(
         self._update_csv(sim_time_str, "link", drainage_data.links)
         self.number_of_writes["link"] += 1
 
-    def finalize(self, drainage_data: DrainageNetworkData) -> None:
-        """Finalize outputs and cleanup."""
-
     def _check_existing_csv(self, geom_type: str):
         """In order to be compatible, an existing CSV should have:
             - Same headers
@@ -206,18 +207,16 @@ def _update_csv(
         self,
         sim_time_str: str,
         geom_type: str,
-        drainage_elements: Tuple[DrainageNodeData | DrainageLinkData, ...],
+        drainage_elements: tuple[DrainageNodeData | DrainageLinkData, ...],
     ):
         """Update adequate CSV in object store"""
         # Check compatibility on first write
         if 0 == self.number_of_writes[geom_type] and self.existing_ids[geom_type]:
             # IDs must match
-            new_ids = set(
-                [
-                    drainage_elem.attributes.model_dump()[f"{geom_type}_id"]
-                    for drainage_elem in drainage_elements
-                ]
-            )
+            new_ids = {
+                drainage_elem.attributes.model_dump()[f"{geom_type}_id"]
+                for drainage_elem in drainage_elements
+            }
             if not new_ids == self.existing_ids[geom_type]:
                 raise ValueError(
                     f"Object ids mismatch for {geom_type}: "
@@ -236,7 +235,7 @@ def _update_csv(
         updated_csv = existing_csv + new_rows
         obstore.put(self.store, self.file_paths[geom_type], file=updated_csv.encode("utf-8"))
 
-    def _attrs_line(self, drainage_element: DrainageNodeData | DrainageLinkData) -> List[str, ...]:
+    def _attrs_line(self, drainage_element: DrainageNodeData | DrainageLinkData) -> list[str, ...]:
         """Return a list of attributes"""
         # Convert attributes to list
         attributes = [str(a) for a in drainage_element.attributes.model_dump().values()]