Add values_dtype backend option to load values at full precision - ad…

…d tests
ecmwf · Dec 17, 2024 · 21ec8e8 · 21ec8e8
1 parent 5de59e8
commit 21ec8e8
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 1 deletion.
diff --git a/tests/test_30_dataset.py b/tests/test_30_dataset.py
@@ -380,3 +380,17 @@ def test_missing_field_values() -> None:
     t2 = res.variables["t2m"]
     assert np.isclose(np.nanmean(t2.data[0, :, :]), 268.375)
     assert np.isclose(np.nanmean(t2.data[1, :, :]), 270.716)
+
+
+def test_default_values_dtype() -> None:
+    res = dataset.open_file(TEST_DATA_MISSING_VALS)
+    assert res.variables["t2m"].data.dtype == np.dtype("float32")
+    assert res.variables["latitude"].data.dtype == np.dtype("float64")
+    assert res.variables["longitude"].data.dtype == np.dtype("float64")
+
+
+def test_float64_values_dtype() -> None:
+    res = dataset.open_file(TEST_DATA_MISSING_VALS, values_dtype=np.dtype("float64"))
+    assert res.variables["t2m"].data.dtype == np.dtype("float64")
+    assert res.variables["latitude"].data.dtype == np.dtype("float64")
+    assert res.variables["longitude"].data.dtype == np.dtype("float64")
diff --git a/tests/test_50_xarray_plugin.py b/tests/test_50_xarray_plugin.py
@@ -175,4 +175,18 @@ def test_xr_open_dataset_coords_to_attributes() -> None:
     assert "depthBelowLandLayer" not in ds.coords
 
     assert "GRIB_surface" in ds["t2m"].attrs
-    assert "GRIB_depthBelowLandLayer" in ds["stl1"].attrs
+    assert "GRIB_depthBelowLandLayer" in ds["stl1"].attrs
+
+
+def test_xr_open_dataset_default_values_dtype() -> None:
+    ds = xr.open_dataset(TEST_DATA_MISSING_VALS, engine="cfgrib")
+    assert ds["t2m"].data.dtype == np.dtype("float32")
+    assert ds["latitude"].data.dtype == np.dtype("float64")
+    assert ds["longitude"].data.dtype == np.dtype("float64")
+
+
+def test_xr_open_dataset_float64_values_dtype() -> None:
+    ds = xr.open_dataset(TEST_DATA_MISSING_VALS, engine="cfgrib", values_dtype=np.dtype("float64"))
+    assert ds["t2m"].data.dtype == np.dtype("float64")
+    assert ds["latitude"].data.dtype == np.dtype("float64")
+    assert ds["longitude"].data.dtype == np.dtype("float64")