diff --git a/benchmarks/benchmark_solve_q.py b/benchmarks/benchmark_solve_q.py index bac059c..b0cfa9a 100644 --- a/benchmarks/benchmark_solve_q.py +++ b/benchmarks/benchmark_solve_q.py @@ -69,7 +69,7 @@ def setup_solve_q_args(num_cells: int) -> tuple: arr_qs = zero_padded_array(shape) arr_hfe = zero_padded_array(shape) arr_hfs = zero_padded_array(shape) - arr_bctype = zero_padded_array(shape) + arr_bctype = np.zeros((shape[0] + 2, shape[1] + 2), dtype=np.uint8) arr_qe_new = zero_padded_array(shape) arr_qs_new = zero_padded_array(shape) diff --git a/benchmarks/benchmark_update_h.py b/benchmarks/benchmark_update_h.py index fad4465..0c8b59b 100644 --- a/benchmarks/benchmark_update_h.py +++ b/benchmarks/benchmark_update_h.py @@ -65,7 +65,7 @@ def setup_update_h_args(num_cells: int) -> tuple: arr_ext = zero_padded_array(shape) arr_qe = full_padded_array(shape, np.float32(0.01)) arr_qs = full_padded_array(shape, np.float32(0.01)) - arr_bct = zero_padded_array(shape) + arr_bct = np.zeros((shape[0] + 2, shape[1] + 2), dtype=np.uint8) arr_bcv = zero_padded_array(shape) arr_hfe = full_padded_array(shape, starting_depth) arr_hfs = full_padded_array(shape, starting_depth) diff --git a/src/itzi_core/array_definitions.py b/src/itzi_core/array_definitions.py index 3471909..b5b6860 100644 --- a/src/itzi_core/array_definitions.py +++ b/src/itzi_core/array_definitions.py @@ -16,6 +16,7 @@ from enum import Enum import numpy as np +from numpy.typing import DTypeLike class ArrayCategory(Enum): @@ -41,8 +42,9 @@ class ArrayDefinition: unit: str # Physical units of the array cf_unit: str # The unit expected by the CF convention var_loc: str # Location of the value. Either "face" or "edge" - fill_value: float = 0.0 # Fill value (replace NaN) + fill_value: float | int = 0.0 # Fill value (replace NaN) computes_from: str | None = None # For accumulation arrays + dtype: DTypeLike | None = None # Optional storage dtype override # Centralized array definitions - Single source of truth @@ -190,7 +192,8 @@ class ArrayDefinition: unit="1", cf_unit="", var_loc="face", - fill_value=0.0, + fill_value=0, + dtype=np.uint8, ), ] # ===== INTERNAL ARRAYS ===== diff --git a/src/itzi_core/compute/partial_inertia_h.pyx b/src/itzi_core/compute/partial_inertia_h.pyx index 208bb46..ada601b 100644 --- a/src/itzi_core/compute/partial_inertia_h.pyx +++ b/src/itzi_core/compute/partial_inertia_h.pyx @@ -18,6 +18,7 @@ from libc.math cimport atan2 as c_atan from libc.math cimport fmax ctypedef cython.floating DTYPE_t +ctypedef unsigned char BCTYPE_t cdef float PI = 3.1415926535898 cdef int solve_h_tile_rows = 64 cdef int solve_h_tile_cols = 128 @@ -47,7 +48,7 @@ cdef inline void solve_h_tile( DTYPE_t[:, ::1] arr_ext, DTYPE_t[:, ::1] arr_qe, DTYPE_t[:, ::1] arr_qs, - DTYPE_t[:, ::1] arr_bct, + BCTYPE_t[:, ::1] arr_bct, DTYPE_t[:, ::1] arr_bcv, DTYPE_t[:, ::1] arr_h, DTYPE_t[:, ::1] arr_hmax, @@ -70,7 +71,8 @@ cdef inline void solve_h_tile( ) noexcept nogil: """Update depth, velocity, and Froude values for one tile.""" cdef int r, c - cdef DTYPE_t qext, qe, qw, qn, qs, h, q_sum, h_new, hmax, bct, bcv + cdef DTYPE_t qext, qe, qw, qn, qs, h, q_sum, h_new, hmax, bcv + cdef BCTYPE_t bct cdef DTYPE_t hfe, hfs, hfw, hfn, ve, vw, vn, vs, vx, vy, v, vdir cdef DTYPE_t eps = 1e-12 # Small epsilon to avoid division by zero @@ -141,7 +143,7 @@ def solve_h( DTYPE_t[:, ::1] arr_ext, DTYPE_t[:, ::1] arr_qe, DTYPE_t[:, ::1] arr_qs, - DTYPE_t[:, ::1] arr_bct, + BCTYPE_t[:, ::1] arr_bct, DTYPE_t[:, ::1] arr_bcv, DTYPE_t[:, ::1] arr_h, DTYPE_t[:, ::1] arr_hmax, diff --git a/src/itzi_core/compute/partial_inertia_q.pyx b/src/itzi_core/compute/partial_inertia_q.pyx index bc1563d..1f7d98a 100644 --- a/src/itzi_core/compute/partial_inertia_q.pyx +++ b/src/itzi_core/compute/partial_inertia_q.pyx @@ -18,6 +18,7 @@ from libc.math cimport sqrt as c_sqrt from libc.math cimport fmin, copysign ctypedef cython.floating DTYPE_t +ctypedef unsigned char BCTYPE_t cdef int solve_q_tile_rows = 64 cdef int solve_q_tile_cols = 128 @@ -122,7 +123,7 @@ cdef inline void solve_qe_west_boundary_at( DTYPE_t[:, ::1] arr_h, DTYPE_t[:, ::1] arr_qe, DTYPE_t[:, ::1] arr_hfe, - DTYPE_t[:, ::1] arr_bctype, + BCTYPE_t[:, ::1] arr_bctype, DTYPE_t[:, ::1] arr_qe_new, int r, ) noexcept nogil: @@ -164,7 +165,7 @@ cdef inline void solve_qe_east_boundary_at( DTYPE_t[:, ::1] arr_h, DTYPE_t[:, ::1] arr_qe, DTYPE_t[:, ::1] arr_hfe, - DTYPE_t[:, ::1] arr_bctype, + BCTYPE_t[:, ::1] arr_bctype, DTYPE_t[:, ::1] arr_qe_new, int col_east_boundary, int r, @@ -308,7 +309,7 @@ cdef inline void solve_qs_north_boundary_at( DTYPE_t[:, ::1] arr_h, DTYPE_t[:, ::1] arr_qs, DTYPE_t[:, ::1] arr_hfs, - DTYPE_t[:, ::1] arr_bctype, + BCTYPE_t[:, ::1] arr_bctype, DTYPE_t[:, ::1] arr_qs_new, int c, ) noexcept nogil: @@ -350,7 +351,7 @@ cdef inline void solve_qs_south_boundary_at( DTYPE_t[:, ::1] arr_h, DTYPE_t[:, ::1] arr_qs, DTYPE_t[:, ::1] arr_hfs, - DTYPE_t[:, ::1] arr_bctype, + BCTYPE_t[:, ::1] arr_bctype, DTYPE_t[:, ::1] arr_qs_new, int row_south_boundary, int c, @@ -511,7 +512,7 @@ def solve_q( DTYPE_t[:, ::1] arr_qs, DTYPE_t[:, ::1] arr_hfe, DTYPE_t[:, ::1] arr_hfs, - DTYPE_t[:, ::1] arr_bctype, + BCTYPE_t[:, ::1] arr_bctype, DTYPE_t[:, ::1] arr_qe_new, DTYPE_t[:, ::1] arr_qs_new, DTYPE_t dt, @@ -805,7 +806,7 @@ cdef DTYPE_t flow_GMS( @cython.cdivision(True) # Don't check division by zero @cython.boundscheck(False) # turn off bounds-checking for entire function cdef DTYPE_t boundary_flow( - DTYPE_t bctype, + BCTYPE_t bctype, DTYPE_t q_domain, DTYPE_t flow_depth_domain, DTYPE_t flow_depth_boundary, diff --git a/src/itzi_core/rasterdomain.py b/src/itzi_core/rasterdomain.py index 246c4fe..5d58db1 100644 --- a/src/itzi_core/rasterdomain.py +++ b/src/itzi_core/rasterdomain.py @@ -126,6 +126,11 @@ def __init__(self, dtype, arr_mask: np.ndarray, cell_shape: tuple[float, float]) if ArrayCategory.ACCUMULATION in arr_def.category ] self.k_all = set(self.k_input + self.k_internal + self.k_accum) + self.dtypes = { + arr_def.key: np.dtype(self.dtype if arr_def.dtype is None else arr_def.dtype) + for arr_def in ARRAY_DEFINITIONS + if arr_def.key in self.k_all + } # Instantiate arrays and padded arrays filled with zeros self.arr = dict.fromkeys(self.k_all) self.arrp = dict.fromkeys(self.k_all) @@ -144,7 +149,7 @@ def _create_arrays(self) -> Self: the unpadded arrays are a slice of the padded ones """ for k in self.arr.keys(): - arr = np.full(fill_value=self.fill_values[k], shape=self.shape, dtype=self.dtype) + arr = np.full(fill_value=self.fill_values[k], shape=self.shape, dtype=self.dtypes[k]) self.arr[k], self.arrp[k] = self.pad_array(arr) return self @@ -164,6 +169,8 @@ def mask_array(self, arr: np.ndarray, default_value: float) -> Self: def unmask_array(self, arr: np.ndarray) -> np.ndarray: """Replace values in the input array by NULL values from mask""" unmasked_array = np.copy(arr) + if np.issubdtype(unmasked_array.dtype, np.integer): + unmasked_array = unmasked_array.astype(self.dtype) unmasked_array[self.mask] = np.nan return unmasked_array @@ -195,10 +202,25 @@ def update_array(self, arr_key: str, arr: np.ndarray) -> Self: # Calculate actual depth and update the internal depth array arr = rastermetrics.calculate_h_from_wse(arr_wse=arr, arr_dem=self.get_array("dem")) arr_key = "water_depth" + elif arr_key == "bctype": + arr = self._prepare_bctype(arr) self.mask_array(arr, self.fill_values[arr_key]) self.arr[arr_key][:], self.arrp[arr_key][:] = self.pad_array(arr) return self + def _prepare_bctype(self, arr: np.ndarray) -> np.ndarray: + """Mask and validate boundary codes before assigning them to uint8 storage.""" + if not (np.issubdtype(arr.dtype, np.integer) or np.issubdtype(arr.dtype, np.floating)): + raise ValueError("Invalid values for 'bctype': expected an integer or floating array.") + + candidate = np.array(arr, copy=True) + self.mask_array(candidate, self.fill_values["bctype"]) + valid = np.isin(candidate, (0, 1, 2, 3, 4)) + if not np.all(valid): + invalid_values = candidate[~valid].reshape(-1)[:5].tolist() + raise ValueError(f"Invalid values for 'bctype': {invalid_values}") + return candidate + def get_array(self, k: str) -> np.ndarray: """return the unpadded, masked array of key 'k'""" return self.arr[k] @@ -297,19 +319,44 @@ def load_state(self, npz_data: io.BytesIO) -> Self: f"domain expects padded shape {padded_shape}" ) - # Verify dtype compatibility (allow safe casting) + # Verify dtype compatibility (allow safe casting), while accepting valid + # floating-point bctype arrays from legacy state archives. + converted_arrays: dict[str, np.ndarray] = {} for key in expected_keys: stored_arr = npz[key] - if not np.can_cast(stored_arr.dtype, self.dtype, casting="safe"): + target_dtype = self.dtypes[key] + if key == "bctype": + if stored_arr.dtype != target_dtype and not np.issubdtype( + stored_arr.dtype, np.floating + ): + raise HotstartError( + f"Array '{key}' dtype mismatch: archive has {stored_arr.dtype}, " + f"domain expects {target_dtype} (or a safely castable type)" + ) + candidate = np.array(stored_arr, copy=True) + padded_mask = np.pad(self.mask, 1, mode="edge") + masked = padded_mask + if np.issubdtype(stored_arr.dtype, np.floating): + masked = np.logical_or(np.isnan(candidate), masked) + candidate[masked] = self.fill_values["bctype"] + valid = np.isin(candidate, (0, 1, 2, 3, 4)) + if not np.all(valid): + invalid_values = candidate[~valid].reshape(-1)[:5].tolist() + raise HotstartError( + f"Invalid values for 'bctype' in raster state: {invalid_values}" + ) + converted_arrays[key] = candidate.astype(target_dtype) + elif not np.can_cast(stored_arr.dtype, target_dtype, casting="safe"): raise HotstartError( f"Array '{key}' dtype mismatch: archive has {stored_arr.dtype}, " - f"domain expects {self.dtype} (or a safely castable type)" + f"domain expects {target_dtype} (or a safely castable type)" ) + else: + converted_arrays[key] = stored_arr.astype(target_dtype) # All validations passed - restore the arrays for key in expected_keys: - # Get the stored padded array and convert to domain dtype - arrp = npz[key].astype(self.dtype) + arrp = converted_arrays[key] # Store the padded array directly self.arrp[key][:] = arrp # Extract the interior (unpadded) slice for self.arr using simple_pad diff --git a/tests/conftest.py b/tests/conftest.py index 581c3bf..f596e4f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -167,7 +167,7 @@ def domain_5by5() -> Domain5by5Data: arr_mask = np.full(domain_data.shape, False, dtype=np.bool_) - arr_bctype = np.zeros(domain_data.shape, dtype=np.float32) + arr_bctype = np.zeros(domain_data.shape, dtype=np.uint8) arr_bctype[0, :] = 2 arr_bctype[4, :] = 2 arr_bctype[:, 0] = 2 diff --git a/tests/test_flow.py b/tests/test_flow.py index 20740ee..9c5838c 100644 --- a/tests/test_flow.py +++ b/tests/test_flow.py @@ -51,7 +51,7 @@ def _solve_q_at_face( arr_qs = np.zeros(shape, dtype=dtype) arr_hfe = np.zeros(shape, dtype=dtype) arr_hfs = np.zeros(shape, dtype=dtype) - arr_bctype = np.zeros(shape, dtype=dtype) + arr_bctype = np.zeros(shape, dtype=np.uint8) arr_qe_new = np.zeros(shape, dtype=dtype) arr_qs_new = np.zeros(shape, dtype=dtype) @@ -219,7 +219,7 @@ def test_solve_h_uses_dx_and_dy_separately_in_flow_divergence(): arr_ext = np.zeros(shape, dtype=dtype) arr_qe = np.zeros(shape, dtype=dtype) arr_qs = np.zeros(shape, dtype=dtype) - arr_bct = np.zeros(shape, dtype=dtype) + arr_bct = np.zeros(shape, dtype=np.uint8) arr_bcv = np.zeros(shape, dtype=dtype) arr_h = np.zeros(shape, dtype=dtype) arr_hmax = np.zeros(shape, dtype=dtype) @@ -292,7 +292,7 @@ def setup_method(self): self.arr_ext = np.zeros(self.shape, dtype=self.dtype) self.arr_qe = np.ones(self.shape, dtype=self.dtype) * 0.5 self.arr_qs = np.ones(self.shape, dtype=self.dtype) * 0.3 - self.arr_bct = np.zeros(self.shape, dtype=self.dtype) + self.arr_bct = np.zeros(self.shape, dtype=np.uint8) self.arr_bcv = np.zeros(self.shape, dtype=self.dtype) self.arr_h = np.ones(self.shape, dtype=self.dtype) * 0.1 self.arr_hmax = np.ones(self.shape, dtype=self.dtype) * 0.1 @@ -421,7 +421,7 @@ def setup_method(self): [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], ] - self.arr_bct = np.array(bct_values, dtype=self.dtype) + self.arr_bct = np.array(bct_values, dtype=np.uint8) assert self.shape == self.arr_bct.shape def test_adding_water(self): diff --git a/tests/test_rasterdomain.py b/tests/test_rasterdomain.py new file mode 100644 index 0000000..28be840 --- /dev/null +++ b/tests/test_rasterdomain.py @@ -0,0 +1,115 @@ +"""Tests for RasterDomain storage dtypes and boundary-type validation.""" + +import io + +import numpy as np +import pytest + +from itzi_core.itzi_error import HotstartError +from itzi_core.rasterdomain import RasterDomain + + +def make_domain( + dtype: type[np.floating] = np.float32, mask: np.ndarray | None = None +) -> RasterDomain: + if mask is None: + mask = np.zeros((3, 3), dtype=bool) + return RasterDomain(dtype=dtype, arr_mask=mask, cell_shape=(1.0, 1.0)) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_bctype_is_the_only_integer_storage_array(dtype): + domain = make_domain(dtype) + + assert domain.get_array("bctype").dtype == np.dtype(np.uint8) + assert domain.get_padded("bctype").dtype == np.dtype(np.uint8) + for key in domain.k_all - {"bctype"}: + assert domain.get_array(key).dtype == np.dtype(dtype) + + +@pytest.mark.parametrize("dtype", [np.uint8, np.int16, np.float32, np.float64]) +def test_bctype_accepts_exact_codes(dtype): + domain = make_domain() + values = np.array([[0, 1, 2], [3, 4, 0], [1, 2, 3]], dtype=dtype) + + domain.update_array("bctype", values) + + np.testing.assert_array_equal(domain.get_array("bctype"), values) + assert domain.get_array("bctype").dtype == np.dtype(np.uint8) + + +def test_bctype_masks_nan_and_domain_mask_before_validation(): + mask = np.zeros((3, 3), dtype=bool) + mask[0, 0] = True + domain = make_domain(mask=mask) + values = np.zeros((3, 3), dtype=np.float64) + values[0, 0] = 99 + values[1, 1] = np.nan + + domain.update_array("bctype", values) + + assert domain.get_array("bctype")[0, 0] == 0 + assert domain.get_array("bctype")[1, 1] == 0 + assert domain.get_unmasked("bctype").dtype == np.dtype(np.float32) + assert np.isnan(domain.get_unmasked("bctype")[0, 0]) + + +@pytest.mark.parametrize("invalid_value", [2.5, np.inf, -1, 5]) +def test_invalid_bctype_update_does_not_mutate_state(invalid_value): + domain = make_domain() + domain.update_array("bctype", np.ones(domain.shape, dtype=np.uint8)) + before = domain.get_array("bctype").copy() + before_padded = domain.get_padded("bctype").copy() + values = np.ones(domain.shape, dtype=np.float64) + values[1, 1] = invalid_value + + with pytest.raises(ValueError, match="bctype"): + domain.update_array("bctype", values) + + np.testing.assert_array_equal(domain.get_array("bctype"), before) + np.testing.assert_array_equal(domain.get_padded("bctype"), before_padded) + + +def make_archive(domain: RasterDomain, bctype: np.ndarray) -> io.BytesIO: + saved = domain.save_state() + saved.seek(0) + npz = np.load(saved, allow_pickle=False) + arrays = {key: npz[key] for key in npz.files} + arrays["bctype"] = bctype + buffer = io.BytesIO() + np.savez(buffer, allow_pickle=False, **arrays) + buffer.seek(0) + return buffer + + +def test_load_state_accepts_legacy_float_bctype_and_restores_uint8(): + source = make_domain(np.float32) + source.update_array("bctype", np.arange(9, dtype=np.uint8).reshape(3, 3) % 5) + saved = source.save_state() + saved.seek(0) + npz = np.load(saved, allow_pickle=False) + legacy_bctype = npz["bctype"].astype(np.float64) + + restored = make_domain(np.float32) + restored.load_state(make_archive(source, legacy_bctype)) + + np.testing.assert_array_equal(restored.get_padded("bctype"), npz["bctype"]) + assert restored.get_array("bctype").dtype == np.dtype(np.uint8) + + +def test_invalid_legacy_bctype_does_not_partially_restore_state(): + source = make_domain(np.float32) + source.update_array("bctype", np.zeros(source.shape, dtype=np.uint8)) + invalid = source.get_padded("bctype").astype(np.float32) + invalid[2, 2] = 2.5 + + restored = make_domain(np.float32) + restored.update_array("water_depth", np.full(restored.shape, 7, dtype=np.float32)) + before_bctype = restored.get_padded("bctype").copy() + before_depth = restored.get_padded("water_depth").copy() + + with pytest.raises(HotstartError, match="bctype"): + restored.load_state(make_archive(source, invalid)) + + np.testing.assert_array_equal(restored.get_padded("bctype"), before_bctype) + np.testing.assert_array_equal(restored.get_padded("water_depth"), before_depth)