diff --git a/docs/notebooks/example_subset.ipynb b/docs/notebooks/example_subset.ipynb index a69445e..694b1cd 100644 --- a/docs/notebooks/example_subset.ipynb +++ b/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", diff --git a/suxarray/core/api.py b/suxarray/core/api.py index de63429..c4cf8c1 100644 --- a/suxarray/core/api.py +++ b/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 diff --git a/suxarray/core/dataarray.py b/suxarray/core/dataarray.py index 1fd725f..e3e1f09 100644 --- a/suxarray/core/dataarray.py +++ b/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) diff --git a/suxarray/core/dataset.py b/suxarray/core/dataset.py index 41b0008..faee36c 100644 --- a/suxarray/core/dataset.py +++ b/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} # }, diff --git a/suxarray/grid/area.py b/suxarray/grid/area.py index 6cb0662..d8735d7 100644 --- a/suxarray/grid/area.py +++ b/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,15 +71,13 @@ 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) @@ -84,15 +85,17 @@ def _integrate_quad(x, y, z): 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) diff --git a/suxarray/grid/integrate.py b/suxarray/grid/integrate.py index eaeec51..1ccb296 100644 --- a/suxarray/grid/integrate.py +++ b/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 \ No newline at end of file diff --git a/tests/test_area.py b/tests/test_area.py index 839ee1e..b843fc1 100644 --- a/tests/test_area.py +++ b/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) diff --git a/tests/test_suxarray.py b/tests/test_suxarray.py index d7f266a..f3f5f9b 100644 --- a/tests/test_suxarray.py +++ b/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): diff --git a/tests/test_volume.py b/tests/test_volume.py index 0792257..65859fa 100644 --- a/tests/test_volume.py +++ b/tests/test_volume.py @@ -65,4 +65,59 @@ def test_calculate_prism_volume_synthetic_shaved_cells(): # Spot-check a couple of entries with pytest.approx for readability. assert volumes[0, 0] == pytest.approx(2.0 / 3.0) - assert volumes[2, 2] == pytest.approx(2.0) \ No newline at end of file + assert volumes[2, 2] == pytest.approx(2.0) + + +def test_calculate_prism_volume_synthetic_float32(): + """Validate that prism volume calculation preserves float32 precision on a mixed tri-quad mesh. + + Mesh consists of: + - Face 0: Triangle with vertices (0,0), (2,0), (0,2). Area = 2 m^2. + - Face 1: Quadrilateral with vertices (2,0), (4,0), (4,2), (2,2). Area = 4 m^2. + """ + 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], # Triangle + [1, 3, 4, 5], # Quadrilateral + ], + dtype=np.int64, + ) + # Shaved cells: node 0 and node 3 have bottom interface masked + bottom_indices = np.array([1, 0, 0, 1, 0, 0], dtype=np.int64) + + # 4 layers -> 3 interfaces, all in float32 + z = np.tile( + np.array([-3.0, -2.0, -1.0, 0.0], dtype=np.float32), + (6, 1), + ) + + volumes = _calculate_prism_volumes( + node_x=node_x, + node_y=node_y, + z=z, + connectivity=connectivity, + bottom_indices=bottom_indices, + ) + + assert volumes.dtype == np.float32 + assert volumes.shape == (3, 2) + + # Expected volumes per layer and face: + # Layer 0: + # Tri [0,1,2]: thickness at nodes = [0, 1, 1], mean = 2/3 -> vol = 2.0 * (2/3) = 4/3 + # Quad [1,3,4,5]: thickness at nodes = [1, 0, 1, 1], mean = 3/4 -> vol = 4.0 * 0.75 = 3.0 + # Layers 1 and 2: + # Tri: thickness = 1.0 -> vol = 2.0 * 1.0 = 2.0 + # Quad: thickness = 1.0 -> vol = 4.0 * 1.0 = 4.0 + expected = np.array( + [ + [4.0 / 3.0, 3.0], + [2.0, 4.0], + [2.0, 4.0], + ], + dtype=np.float32, + ) + np.testing.assert_allclose(volumes, expected, rtol=1e-5, atol=1e-6) diff --git a/tests/testfixtures.py b/tests/testfixtures.py index 1423df0..6bb57ba 100644 --- a/tests/testfixtures.py +++ b/tests/testfixtures.py @@ -26,7 +26,7 @@ def sxds_test_dask(): path_var = [str(p_cur / "testdata/salinity_{}.nc".format(i)) for i in range(1, 3)] ds_salinity = xr.open_mfdataset( path_var, mask_and_scale=False, data_vars="minimal" - ).astype(np.float64) + ) sxds = sx.read_schism_nc(sxgrid, ds_salinity) return sxds @@ -45,6 +45,6 @@ def sxds_element_test(): for i in range(1, 3)] ds_vertVel = xr.open_mfdataset( path_var, mask_and_scale=False, data_vars="minimal" - ).astype(np.float64) + ) sxds = sx.read_schism_nc(sxgrid, ds_vertVel) return sxds \ No newline at end of file