feat: preserve native dtype for SCHISM outputs

docs/notebooks/example_subset.ipynb

@@ -200,9 +200,9 @@
     "files_zcoords = [str(p) for p in [dir_data / 'zCoordinates_1.nc', dir_data / 'zCoordinates_2.nc']]\n",
     "grid = sx.open_grid(files_out2d, files_zcoords)\n",
     "\n",
-    "# We are converting the data to float64 for later calculation.\n",
+    "# Read salinity data directly in native float32 format\n",
     "files_salinity = [dir_data / 'salinity_1.nc', dir_data / 'salinity_2.nc']\n",
-    "ds_salinity = xr.open_mfdataset(files_salinity, parallel=True).astype(np.float64)\n",
+    "ds_salinity = xr.open_mfdataset(files_salinity, parallel=True)\n",
     "\n",
     "ds_salinity.attrs = {\n",
     "    \"Conventions\": \"CF-1.8\",\n",

suxarray/core/api.py

@@ -16,6 +16,23 @@
 )
 from suxarray.conventions.schism_grid import SCHISM_GRID_TIME_VARIABLES_IN_OUT2D
 
+def _safe_mask_dataset(ds: xr.Dataset) -> xr.Dataset:
+    """Mask missing/fill values on float variables while preserving native dtypes (e.g. float32).
+
+    Leaves integer variables untouched so connectivity arrays with _FillValue=-1
+    remain integer, and avoids xarray CFMaskCoder upcasting float32 to float64.
+    """
+    for varname in list(ds.data_vars.keys()):
+        da = ds[varname]
+        if np.issubdtype(da.dtype, np.floating):
+            fill_val = da.attrs.get("_FillValue", da.encoding.get("_FillValue", None))
+            missing_val = da.attrs.get("missing_value", da.encoding.get("missing_value", None))
+            for fv in (fill_val, missing_val):
+                if fv is not None and not np.isnan(fv):
+                    ds[varname] = da.where(da != fv)
+    return ds
+
+
 def open_ts_dataset(
     nc_file_path: Union[str, os.PathLike, Path],
     ref_time: Optional[pd.Timestamp] = None,
@@ -29,19 +46,20 @@ def open_ts_dataset(
     Returns:
     xarray.Dataset: The dataset read from the NetCDF file.
     """
-    ds = xr.open_dataset(nc_file_path, decode_times=False)
+    ds = xr.open_dataset(nc_file_path, decode_times=False, mask_and_scale=False)
     if "time" in ds:
         if "units" in ds.time.attrs:
             origin = parse(
                 ds.time.attrs["units"].split("since", 1)[1].strip(), ignoretz=True
             )  # removes the timezone info from the origin timestring
             ds.time.attrs["units"] = f"seconds since {origin.isoformat()}"
-            ds = xr.decode_cf(ds)
+            ds = xr.decode_cf(ds, mask_and_scale=False)
         elif ref_time is not None:
             # If the time variable is missing or doesn't have units, create a time variable based on ref_time
             ds.time.attrs["units"] = f"seconds since {ref_time.isoformat()}"
-            ds = xr.decode_cf(ds)
+            ds = xr.decode_cf(ds, mask_and_scale=False)
 
+    ds = _safe_mask_dataset(ds)
     return ds
 
 def _choose_netcdf_engine() -> str:
@@ -59,6 +77,7 @@ def _choose_netcdf_engine() -> str:
 
 
 def read_schism_nc(grid: Grid, ds_data: xr.Dataset) -> SxDataset:
+    ds_data = _safe_mask_dataset(ds_data)
     ds_data = ux.core.utils._map_dims_to_ugrid(ds_data, grid._source_dims_dict, grid)
     if "nSCHISM_vgrid_layers" in ds_data.dims:
         ds_data = ds_data.swap_dims({"nSCHISM_vgrid_layers": "n_layer"})
@@ -147,12 +166,14 @@ def open_grid(
         ds_zcoords = xr.open_mfdataset(
             files_zcoords,
             engine=engine,
+            mask_and_scale=False,
             chunks=chunks,
             join="override",
             coords="minimal",
             data_vars="minimal",
             compat="override",
         )
+        ds_zcoords = _safe_mask_dataset(ds_zcoords)
     else:
         ds_zcoords = None
 

suxarray/core/dataarray.py

@@ -78,7 +78,7 @@ def depth_average(self) -> SxDataArray:
                 [],
             ],
             dask="parallelized",
-            output_dtypes=[float],
+            output_dtypes=[self.dtype],
         )
         name = f"depth_averaged_{self.name}"
         return SxDataArray(da_da, sxgrid=self.sxgrid, name=name)

suxarray/core/dataset.py

@@ -74,7 +74,7 @@ def __getitem__(self, key):
     #                 "nMesh2_face",
     #             ],
     #         ],
-    #         output_dtypes=[float],
+    #         output_dtypes=[dataarray.dtype],
     #         dask_gufunc_kwargs={
     #             "output_sizes": {"nMesh2_face": self.uxgrid.nMesh2_face}
     #         },

suxarray/grid/area.py

@@ -26,11 +26,11 @@ def _integrate_triangle(x, y, z):
     """
     Integrate over a triangle using the three corner points.
     """
-    return (
-        np.abs((x[1] - x[0]) * (y[2] - y[0]) - (x[2] - x[0]) * (y[1] - y[0]))
-        * np.array([z[i, :].mean() for i in range(z.shape[0])])
-        * 0.5
-    )
+    area_2 = np.abs((x[1] - x[0]) * (y[2] - y[0]) - (x[2] - x[0]) * (y[1] - y[0]))
+    result = np.empty((z.shape[0],), dtype=z.dtype)
+    for k in range(z.shape[0]):
+        result[k] = 0.5 * area_2 * z[k, :].mean()
+    return result
 
 
 @njit(cache=True)
@@ -41,24 +41,27 @@ def _gauss_quadrature_points_order2():
 
 
 @njit(cache=True)
-def _shape_functions_quad_bilinear(xi, eta):
+def _shape_functions_quad_bilinear(xi, eta, z_dtype):
     return np.array(
         [
             0.25 * (1 - xi) * (1 - eta),
             0.25 * (1 + xi) * (1 - eta),
             0.25 * (1 + xi) * (1 + eta),
             0.25 * (1 - xi) * (1 + eta),
-        ]
+        ],
+        dtype=z_dtype,
     )
 
 
 @njit(cache=True)
-def _shape_functions_derivative_quad_bilinear(xi, eta):
+def _shape_functions_derivative_quad_bilinear(xi, eta, x_dtype):
     dxi = np.array(
         [-0.25 * (1 - eta), 0.25 * (1 - eta), 0.25 * (1 + eta), -0.25 * (1 + eta)],
+        dtype=x_dtype,
     )
     deta = np.array(
         [-0.25 * (1 - xi), -0.25 * (1 + xi), 0.25 * (1 + xi), 0.25 * (1 - xi)],
+        dtype=x_dtype,
     )
     return dxi, deta
 
@@ -68,31 +71,31 @@ def _integrate_quad(x, y, z):
     # points, weights = _gauss_quadrature_points_order2()
     points = np.array([[-0.577350269189626, 0.577350269189626]])
     weights = np.array([1.0, 1.0])
-    result = np.zeros((z.shape[0],))
+    result = np.zeros((z.shape[0],), dtype=z.dtype)
 
-    for i, wi in enumerate(weights):
-        for j, wj in enumerate(weights):
+    for i in range(2):
+        for j in range(2):
             xi, eta = points[0][i], points[0][j]
-            shape = _shape_functions_quad_bilinear(xi, eta)
-            dxi, deta = _shape_functions_derivative_quad_bilinear(xi, eta)
-
-            z_local = np.array([np.dot(shape, z[i, :]) for i in range(z.shape[0])])
+            shape = _shape_functions_quad_bilinear(xi, eta, z.dtype)
+            dxi, deta = _shape_functions_derivative_quad_bilinear(xi, eta, x.dtype)
 
             dx_dxi = np.dot(dxi, x)
             dy_dxi = np.dot(dxi, y)
             dx_deta = np.dot(deta, x)
             dy_deta = np.dot(deta, y)
 
             jacobian_det = np.abs(dx_dxi * dy_deta - dx_deta * dy_dxi)
-            result += wi * wj * z_local * jacobian_det
+            for k in range(z.shape[0]):
+                z_local = np.dot(shape, z[k, :])
+                result[k] += weights[i] * weights[j] * z_local * jacobian_det
 
     return result
 
 
 @njit(cache=True)
-def _integrate_nodal(node_x, node_y, values, connectivity):
+def _integrate_nodal_2d(node_x, node_y, values, connectivity):
     n_elements = connectivity.shape[0]
-    result = np.zeros((values.shape[0], n_elements))
+    result = np.zeros((values.shape[0], n_elements), dtype=values.dtype)
 
     for i in range(n_elements):
         # Select indices that are not negative
@@ -108,3 +111,14 @@ def _integrate_nodal(node_x, node_y, values, connectivity):
             result[:, i] = _integrate_quad(x, y, z)
 
     return result
+
+
+@njit(cache=True)
+def _integrate_nodal(node_x, node_y, values, connectivity):
+    if values.ndim == 1:
+        v_2d = np.empty((1, len(values)), dtype=values.dtype)
+        v_2d[0, :] = values
+        res_2d = _integrate_nodal_2d(node_x, node_y, v_2d, connectivity)
+        return res_2d[0]
+    else:
+        return _integrate_nodal_2d(node_x, node_y, values, connectivity)

suxarray/grid/integrate.py

@@ -67,7 +67,7 @@ def _face_mean_from_nodal(values, connectivity):
     """
     n_node, n_interface = values.shape
     n_face = connectivity.shape[0]
-    face_mean = np.zeros((n_interface, n_face), dtype=np.float64)
+    face_mean = np.zeros((n_interface, n_face), dtype=values.dtype)
 
     for face_i in range(n_face):
         node_ids = connectivity[face_i]
@@ -128,10 +128,9 @@ def _calculate_prism_volumes(node_x, node_y, z, connectivity, bottom_indices):
     ndarray, shape (n_layer-1, n_face).
         Prism volumes per layer and face.
     """
-    # Enforce float64 for numba
-    node_x = node_x.astype(np.float64)
-    node_y = node_y.astype(np.float64)
-    z = z.astype(np.float64)
+    # Ensure coordinate arrays are float64 for geometric accuracy
+    node_x = np.asarray(node_x, dtype=np.float64)
+    node_y = np.asarray(node_y, dtype=np.float64)
 
     z = _coerce_single_timestep(z, "z")
 
@@ -230,7 +229,7 @@ def integrate_prism_variable(sxda, return_per_layer=False):
 
     # Integrate!
     integrated = _integrate_nodal(
-        node_x, node_y, values.astype(np.float64), connectivity
+        node_x, node_y, values, connectivity
     )
 
     if return_per_layer:
@@ -305,7 +304,7 @@ def _calculate_volumes_uxarray(sxda, radius_m=None):
     grid = sxda.sxgrid
 
     z = _coerce_single_timestep(
-        grid.sgrid_info.zCoordinates.values.astype(np.float64), "z"
+        grid.sgrid_info.zCoordinates.values, "z"
     )
     connectivity = grid.face_node_connectivity.values
 
@@ -419,7 +418,7 @@ def _calculate_mass_in_ocean(sxda_salinity, sxda_temperature=None):
     values = rho_h.T  # (n_layer-1, n_node)
 
     integrated = _integrate_nodal(
-        node_x, node_y, values.astype(np.float64), connectivity
+        node_x, node_y, values, connectivity
     )
 
     return integrated

tests/test_area.py

@@ -8,15 +8,16 @@
 
 def test_integrate_nodal(sxds_test_dask):
     da = sxds_test_dask.salinity.depth_average()
+    assert da.dtype == np.float32
     result = suxarray.grid.area._integrate_nodal(
         da.sxgrid.node_x.values,
         da.sxgrid.node_y.values,
         da.values,
         da.sxgrid.face_node_connectivity.values,
     )
-    assert result[0, 64] == pytest.approx(41743.07605, abs=1e-3)
+    assert result.dtype == np.float32
+    assert result[0, 64] == pytest.approx(41743.07605, abs=1e-1)
 
-    # da = da.astype(np.float64)
     result = xr.apply_ufunc(
         suxarray.grid.area._integrate_nodal,
         da.sxgrid.node_x,
@@ -39,14 +40,93 @@ def test_integrate_nodal(sxds_test_dask):
             ],
         ],
         output_core_dims=[["n_face"]],
-        output_dtypes=np.float64,
+        output_dtypes=da.dtype,
         dask="parallelized",
     )
+    assert result.dtype == np.float32
     assert result.isel(time=0, n_face=64).values[()] == pytest.approx(
-        41743.07555287409, rel=1e-5
+        41743.07555287409, rel=1e-4
     )
 
     result = da.integrate()
+    assert result.dtype == np.float32
     assert result.isel(time=0, n_face=64, ignore_grid=True).values[()] == pytest.approx(
-        41743.07555287409, rel=1e-5
+        41743.07555287409, rel=1e-4
     )
+
+
+def test_integrate_nodal_float32_and_float64_shapes():
+    """Test 1D and 2D shapes and preservation of float32 and float64 dtypes."""
+    node_x = np.array([0.0, 2.0, 0.0], dtype=np.float64)
+    node_y = np.array([0.0, 0.0, 3.0], dtype=np.float64)
+    connectivity = np.array([[0, 1, 2]], dtype=np.int64)
+
+    # 1D float32 input -> 1D float32 output
+    z_1d_f32 = np.array([4.0, 4.0, 4.0], dtype=np.float32)
+    res_1d_f32 = suxarray.grid.area._integrate_nodal(node_x, node_y, z_1d_f32, connectivity)
+    assert res_1d_f32.dtype == np.float32
+    assert res_1d_f32.shape == (1,)
+    assert res_1d_f32[0] == pytest.approx(12.0, rel=1e-5)
+
+    # 2D float32 input -> 2D float32 output
+    z_2d_f32 = np.array([[4.0, 4.0, 4.0]], dtype=np.float32)
+    res_2d_f32 = suxarray.grid.area._integrate_nodal(node_x, node_y, z_2d_f32, connectivity)
+    assert res_2d_f32.dtype == np.float32
+    assert res_2d_f32.shape == (1, 1)
+    assert res_2d_f32[0, 0] == pytest.approx(12.0, rel=1e-5)
+
+    # 1D float64 input -> 1D float64 output
+    z_1d_f64 = np.array([4.0, 4.0, 4.0], dtype=np.float64)
+    res_1d_f64 = suxarray.grid.area._integrate_nodal(node_x, node_y, z_1d_f64, connectivity)
+    assert res_1d_f64.dtype == np.float64
+    assert res_1d_f64.shape == (1,)
+    assert res_1d_f64[0] == pytest.approx(12.0, rel=1e-12)
+
+
+def test_integrate_nodal_quad_float32():
+    """Test quad integration with float32 values."""
+    node_x = np.array([0.0, 2.0, 2.0, 0.0], dtype=np.float64)
+    node_y = np.array([0.0, 0.0, 3.0, 3.0], dtype=np.float64)
+    connectivity = np.array([[0, 1, 2, 3]], dtype=np.int64)
+
+    z_f32 = np.array([[4.0, 4.0, 4.0, 4.0]], dtype=np.float32)
+    res = suxarray.grid.area._integrate_nodal(node_x, node_y, z_f32, connectivity)
+    assert res.dtype == np.float32
+    assert res.shape == (1, 1)
+    assert res[0, 0] == pytest.approx(24.0, rel=1e-5)
+
+
+def test_integrate_nodal_mixed_tri_quad_float32():
+    """Test nodal integration on a mixed mesh containing both triangles and quads in float32."""
+    # Element 0: Triangle with nodes [0, 1, 2, -1], vertices (0,0), (2,0), (0,2). Area = 2.0
+    # Element 1: Quad with nodes [1, 3, 4, 5], vertices (2,0), (4,0), (4,2), (2,2). Area = 4.0
+    node_x = np.array([0.0, 2.0, 0.0, 4.0, 4.0, 2.0], dtype=np.float64)
+    node_y = np.array([0.0, 0.0, 2.0, 0.0, 2.0, 2.0], dtype=np.float64)
+    connectivity = np.array(
+        [
+            [0, 1, 2, -1],
+            [1, 3, 4, 5],
+        ],
+        dtype=np.int64,
+    )
+
+    # 1D float32 constant field z = 3.0 -> tri: 6.0, quad: 12.0
+    z_1d_f32 = np.full(6, 3.0, dtype=np.float32)
+    res_1d = suxarray.grid.area._integrate_nodal(node_x, node_y, z_1d_f32, connectivity)
+    assert res_1d.dtype == np.float32
+    assert res_1d.shape == (2,)
+    np.testing.assert_allclose(res_1d, [6.0, 12.0], rtol=1e-5)
+
+    # 2D float32 field with 2 timesteps -> timestep 0: [6.0, 12.0], timestep 1: [10.0, 20.0]
+    z_2d_f32 = np.array(
+        [
+            [3.0, 3.0, 3.0, 3.0, 3.0, 3.0],
+            [5.0, 5.0, 5.0, 5.0, 5.0, 5.0],
+        ],
+        dtype=np.float32,
+    )
+    res_2d = suxarray.grid.area._integrate_nodal(node_x, node_y, z_2d_f32, connectivity)
+    assert res_2d.dtype == np.float32
+    assert res_2d.shape == (2, 2)
+    np.testing.assert_allclose(res_2d[0], [6.0, 12.0], rtol=1e-5)
+    np.testing.assert_allclose(res_2d[1], [10.0, 20.0], rtol=1e-5)

tests/test_suxarray.py

@@ -144,9 +144,11 @@ def test_isel(sxds_test_dask):
 
 
 def test_depth_average(sxds_test_dask):
+    assert sxds_test_dask["salinity"].dtype == np.float32
     da = sxds_test_dask["salinity"].depth_average()
+    assert da.dtype == np.float32
     # The value is calculated from VisIt
-    assert da.sel(n_node=439).values[0] == pytest.approx(11.961535, abs=1e-6)
+    assert da.sel(n_node=439).values[0] == pytest.approx(11.961535, abs=1e-5)
 
 
 def test_slice_and_depth_average(sxds_test_dask):
@@ -159,9 +161,11 @@ def test_slice_and_depth_average(sxds_test_dask):
         ]
     )
     da_subset = sxds_test_dask["salinity"].subset.bounding_polygon(polygon)
+    assert da_subset.dtype == np.float32
     da_da = da_subset.depth_average()
+    assert da_da.dtype == np.float32
     # Node 186 is now 7
-    assert da_da.sel(n_node=7).values[0] == pytest.approx(10.334445, abs=1e-6)
+    assert da_da.sel(n_node=7).values[0] == pytest.approx(10.334445, abs=1e-5)
 
 
 # def test_interpolate_xy(sxds_test_dask):