From b35677e2051b38c69e689574e56985175ff56567 Mon Sep 17 00:00:00 2001 From: PavanSiligam Date: Tue, 28 Apr 2026 07:07:20 +0200 Subject: [PATCH 1/3] fix: cast lat/lon coordinates to float64 at CMOR writer level MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Geographic coordinate variables (lat, lon, and their bounds) must be float64 per CMOR3 and CF Conventions §2.2. Since pycmor writes via xarray rather than the CMOR C library, this must be enforced explicitly. Adds _cast_geo_coords_to_float64() called in save_dataset() so the requirement is met automatically for all pipelines. Closes #269 Co-Authored-By: Claude Sonnet 4.6 --- src/pycmor/std_lib/files.py | 43 +++++++++++++++++-- tests/unit/test_savedataset.py | 76 ++++++++++++++++++++++++++++++++++ 2 files changed, 116 insertions(+), 3 deletions(-) diff --git a/src/pycmor/std_lib/files.py b/src/pycmor/std_lib/files.py index fcc1b0be..2d2c1fd3 100644 --- a/src/pycmor/std_lib/files.py +++ b/src/pycmor/std_lib/files.py @@ -40,6 +40,7 @@ from pathlib import Path +import numpy as np import pandas as pd import xarray as xr from xarray.core.utils import is_scalar @@ -379,6 +380,41 @@ def _save_dataset_with_native_timespan( ) +_GEO_COORD_NAMES = frozenset( + { + "lat", "lon", "latitude", "longitude", + "lat_bnds", "lon_bnds", "lat_bounds", "lon_bounds", + "latitude_bnds", "longitude_bnds", "latitude_bounds", "longitude_bounds", + } +) + + +def _cast_geo_coords_to_float64(da): + """Cast geographic coordinate variables to float64 (CMOR3 / CF requirement). + + Accepts both xr.DataArray and xr.Dataset. + """ + coord_updates = { + name: da[name].astype(np.float64) + for name in da.coords + if name in _GEO_COORD_NAMES and da[name].dtype != np.float64 + } + if coord_updates: + for name in coord_updates: + logger.debug( + f"Casting {name!r} from {da[name].dtype} to float64 for CMIP compliance" + ) + da = da.assign_coords(coord_updates) + if isinstance(da, xr.Dataset): + for name in list(da.data_vars): + if name in _GEO_COORD_NAMES and da[name].dtype != np.float64: + logger.debug( + f"Casting {name!r} from {da[name].dtype} to float64 for CMIP compliance" + ) + da[name] = da[name].astype(np.float64) + return da + + def save_dataset(da: xr.DataArray, rule): """ Save dataset to one or more files. @@ -436,6 +472,10 @@ def save_dataset(da: xr.DataArray, rule): # Set default calendar if none is specified if time_encoding.get("calendar") is None: time_encoding["calendar"] = "standard" + + # CMOR3 / CF requirement: geographic coordinates must be float64 + da = _cast_geo_coords_to_float64(da) + if not has_time_axis(da): filepath = create_filepath(da, rule) return da.to_netcdf( @@ -455,7 +495,6 @@ def save_dataset(da: xr.DataArray, rule): ) if isinstance(da, xr.DataArray): da = da.to_dataset() - # Set time variable attributes if rule._pycmor_cfg("xarray_time_set_standard_name"): da[time_label].attrs["standard_name"] = "time" @@ -510,8 +549,6 @@ def save_dataset(da: xr.DataArray, rule): da[time_label].attrs["calendar"] = time_encoding["calendar"] # Ensure the encoding is set on the time variable itself - if isinstance(da, xr.DataArray): - da = da.to_dataset() da[time_label].encoding.update(time_encoding) if not has_time_axis(da): diff --git a/tests/unit/test_savedataset.py b/tests/unit/test_savedataset.py index a9582899..12ab2ce5 100644 --- a/tests/unit/test_savedataset.py +++ b/tests/unit/test_savedataset.py @@ -302,3 +302,79 @@ def test_save_dataset_with_custom_time_settings(tmp_path): assert ( time_var.encoding["calendar"] == custom_calendar ), f"XArray encoding calendar does not match. Expected {custom_calendar}, got {time_var.encoding['calendar']}" + + +@pytest.mark.parametrize("coord_name", ["lat", "lon", "latitude", "longitude"]) +@pytest.mark.parametrize("input_dtype", [np.float32, np.float16]) +def test_save_dataset_casts_geo_coords_to_float64(tmp_path, coord_name, input_dtype): + """Geographic coordinate variables must be written as float64 (CMOR3 / CF requirement).""" + dates = xr.cftime_range(start="2001", periods=2, freq="MS", calendar="noleap") + coords = { + "time": dates, + coord_name: np.array([-45.0, 0.0, 45.0], dtype=input_dtype), + } + da = xr.DataArray( + np.zeros((2, 3), dtype=np.float32), + coords=coords, + dims=["time", coord_name], + name="tos", + ) + rule = Mock() + rule._pycmor_cfg = PycmorConfigManager.from_pycmor_cfg({}) + rule.data_request_variable.frequency = "mon" + rule.data_request_variable.table_header.approx_interval = 30 + rule.cmor_variable = "tos" + rule.variant_label = "r1i1p1f1" + rule.source_id = "AWI-ESM-3" + rule.experiment_id = "historical" + rule.file_timespan = "2YS" + rule.output_directory = str(tmp_path) + + save_dataset(da, rule) + + saved = list(tmp_path.glob("*.nc")) + assert len(saved) == 1 + with xr.open_dataset(saved[0]) as ds: + assert ds[coord_name].dtype == np.float64, ( + f"{coord_name} should be float64, got {ds[coord_name].dtype}" + ) + + +@pytest.mark.parametrize("bounds_name", ["lat_bnds", "lon_bnds", "lat_bounds", "lon_bounds"]) +def test_save_dataset_casts_geo_bounds_to_float64(tmp_path, bounds_name): + """Geographic bounds variables must also be written as float64.""" + dates = xr.cftime_range(start="2001", periods=2, freq="MS", calendar="noleap") + coord_name = "lat" if "lat" in bounds_name else "lon" + coord_vals = np.array([-45.0, 0.0, 45.0], dtype=np.float32) + bounds_vals = np.array( + [[-67.5, -22.5], [-22.5, 22.5], [22.5, 67.5]], dtype=np.float32 + ) + ds = xr.Dataset( + { + "tos": xr.DataArray( + np.zeros((2, 3), dtype=np.float32), + coords={"time": dates, coord_name: coord_vals}, + dims=["time", coord_name], + ), + bounds_name: xr.DataArray(bounds_vals, dims=[coord_name, "bnds"]), + } + ) + rule = Mock() + rule._pycmor_cfg = PycmorConfigManager.from_pycmor_cfg({}) + rule.data_request_variable.frequency = "mon" + rule.data_request_variable.table_header.approx_interval = 30 + rule.cmor_variable = "tos" + rule.variant_label = "r1i1p1f1" + rule.source_id = "AWI-ESM-3" + rule.experiment_id = "historical" + rule.file_timespan = "2YS" + rule.output_directory = str(tmp_path) + + save_dataset(ds, rule) + + saved = list(tmp_path.glob("*.nc")) + assert len(saved) == 1 + with xr.open_dataset(saved[0]) as result: + assert result[bounds_name].dtype == np.float64, ( + f"{bounds_name} should be float64, got {result[bounds_name].dtype}" + ) From bbcf9b79a260a33495cad53fe89578f05948f82a Mon Sep 17 00:00:00 2001 From: PavanSiligam Date: Thu, 4 Jun 2026 11:08:17 +0200 Subject: [PATCH 2/3] style: fix black formatting in files.py and test_savedataset.py to pass CI lint check Co-Authored-By: Claude Sonnet 4.6 --- src/pycmor/std_lib/files.py | 15 ++++++++++++--- tests/unit/test_savedataset.py | 16 +++++++++------- 2 files changed, 21 insertions(+), 10 deletions(-) diff --git a/src/pycmor/std_lib/files.py b/src/pycmor/std_lib/files.py index 2d2c1fd3..5f7a3aec 100644 --- a/src/pycmor/std_lib/files.py +++ b/src/pycmor/std_lib/files.py @@ -382,9 +382,18 @@ def _save_dataset_with_native_timespan( _GEO_COORD_NAMES = frozenset( { - "lat", "lon", "latitude", "longitude", - "lat_bnds", "lon_bnds", "lat_bounds", "lon_bounds", - "latitude_bnds", "longitude_bnds", "latitude_bounds", "longitude_bounds", + "lat", + "lon", + "latitude", + "longitude", + "lat_bnds", + "lon_bnds", + "lat_bounds", + "lon_bounds", + "latitude_bnds", + "longitude_bnds", + "latitude_bounds", + "longitude_bounds", } ) diff --git a/tests/unit/test_savedataset.py b/tests/unit/test_savedataset.py index 12ab2ce5..12eaed5b 100644 --- a/tests/unit/test_savedataset.py +++ b/tests/unit/test_savedataset.py @@ -335,12 +335,14 @@ def test_save_dataset_casts_geo_coords_to_float64(tmp_path, coord_name, input_dt saved = list(tmp_path.glob("*.nc")) assert len(saved) == 1 with xr.open_dataset(saved[0]) as ds: - assert ds[coord_name].dtype == np.float64, ( - f"{coord_name} should be float64, got {ds[coord_name].dtype}" - ) + assert ( + ds[coord_name].dtype == np.float64 + ), f"{coord_name} should be float64, got {ds[coord_name].dtype}" -@pytest.mark.parametrize("bounds_name", ["lat_bnds", "lon_bnds", "lat_bounds", "lon_bounds"]) +@pytest.mark.parametrize( + "bounds_name", ["lat_bnds", "lon_bnds", "lat_bounds", "lon_bounds"] +) def test_save_dataset_casts_geo_bounds_to_float64(tmp_path, bounds_name): """Geographic bounds variables must also be written as float64.""" dates = xr.cftime_range(start="2001", periods=2, freq="MS", calendar="noleap") @@ -375,6 +377,6 @@ def test_save_dataset_casts_geo_bounds_to_float64(tmp_path, bounds_name): saved = list(tmp_path.glob("*.nc")) assert len(saved) == 1 with xr.open_dataset(saved[0]) as result: - assert result[bounds_name].dtype == np.float64, ( - f"{bounds_name} should be float64, got {result[bounds_name].dtype}" - ) + assert ( + result[bounds_name].dtype == np.float64 + ), f"{bounds_name} should be float64, got {result[bounds_name].dtype}" From a225a946be180c1ed30e0b5d6fadc65668f810de Mon Sep 17 00:00:00 2001 From: PavanSiligam Date: Thu, 4 Jun 2026 15:21:05 +0200 Subject: [PATCH 3/3] refactor: defer geo coord float64 cast to write time via xarray encoding Replace in-memory `.astype(np.float64)` + `assign_coords` with `variable.encoding["dtype"] = np.float64` so the cast happens at NetCDF write time rather than mutating the caller's DataArray. This uses the same xarray encoding mechanism already in place for time variables and covers all write paths (to_netcdf, save_mfdataset, split and resample code paths) without modifying any call sites. Tests renamed to `writes_geo_*_as_float64` and extended with an assertion that in-memory dtype is not mutated. Co-Authored-By: Claude Sonnet 4.6 --- src/pycmor/std_lib/files.py | 34 ++++++++++++++++------------------ tests/unit/test_savedataset.py | 18 +++++++++++++----- 2 files changed, 29 insertions(+), 23 deletions(-) diff --git a/src/pycmor/std_lib/files.py b/src/pycmor/std_lib/files.py index 5f7a3aec..6575a9e0 100644 --- a/src/pycmor/std_lib/files.py +++ b/src/pycmor/std_lib/files.py @@ -398,30 +398,28 @@ def _save_dataset_with_native_timespan( ) -def _cast_geo_coords_to_float64(da): - """Cast geographic coordinate variables to float64 (CMOR3 / CF requirement). +def _set_geo_coords_encoding_float64(da): + """Set float64 encoding on geographic coordinate variables (CMOR3 / CF requirement). - Accepts both xr.DataArray and xr.Dataset. + Defers the cast to write time via xarray's encoding mechanism — in-memory data + is unchanged; NetCDF output lands as float64. Accepts both xr.DataArray and + xr.Dataset. """ - coord_updates = { - name: da[name].astype(np.float64) - for name in da.coords - if name in _GEO_COORD_NAMES and da[name].dtype != np.float64 - } - if coord_updates: - for name in coord_updates: + for name in _GEO_COORD_NAMES: + if name in da.coords and da.coords[name].dtype != np.float64: logger.debug( - f"Casting {name!r} from {da[name].dtype} to float64 for CMIP compliance" + f"Setting encoding dtype=float64 for coord {name!r} " + f"(in-memory dtype: {da.coords[name].dtype})" ) - da = da.assign_coords(coord_updates) + da[name].encoding["dtype"] = np.float64 if isinstance(da, xr.Dataset): - for name in list(da.data_vars): + for name in da.data_vars: if name in _GEO_COORD_NAMES and da[name].dtype != np.float64: logger.debug( - f"Casting {name!r} from {da[name].dtype} to float64 for CMIP compliance" + f"Setting encoding dtype=float64 for data_var {name!r} " + f"(in-memory dtype: {da[name].dtype})" ) - da[name] = da[name].astype(np.float64) - return da + da[name].encoding["dtype"] = np.float64 def save_dataset(da: xr.DataArray, rule): @@ -482,8 +480,8 @@ def save_dataset(da: xr.DataArray, rule): if time_encoding.get("calendar") is None: time_encoding["calendar"] = "standard" - # CMOR3 / CF requirement: geographic coordinates must be float64 - da = _cast_geo_coords_to_float64(da) + # CMOR3 / CF requirement: geographic coordinates must be float64 on disk + _set_geo_coords_encoding_float64(da) if not has_time_axis(da): filepath = create_filepath(da, rule) diff --git a/tests/unit/test_savedataset.py b/tests/unit/test_savedataset.py index 12eaed5b..8b1ec395 100644 --- a/tests/unit/test_savedataset.py +++ b/tests/unit/test_savedataset.py @@ -306,8 +306,12 @@ def test_save_dataset_with_custom_time_settings(tmp_path): @pytest.mark.parametrize("coord_name", ["lat", "lon", "latitude", "longitude"]) @pytest.mark.parametrize("input_dtype", [np.float32, np.float16]) -def test_save_dataset_casts_geo_coords_to_float64(tmp_path, coord_name, input_dtype): - """Geographic coordinate variables must be written as float64 (CMOR3 / CF requirement).""" +def test_save_dataset_writes_geo_coords_as_float64(tmp_path, coord_name, input_dtype): + """Geographic coordinate variables must be written as float64 (CMOR3 / CF requirement). + + The cast is deferred to write time via xarray encoding — in-memory dtype is + preserved unchanged. + """ dates = xr.cftime_range(start="2001", periods=2, freq="MS", calendar="noleap") coords = { "time": dates, @@ -337,13 +341,15 @@ def test_save_dataset_casts_geo_coords_to_float64(tmp_path, coord_name, input_dt with xr.open_dataset(saved[0]) as ds: assert ( ds[coord_name].dtype == np.float64 - ), f"{coord_name} should be float64, got {ds[coord_name].dtype}" + ), f"{coord_name} should be float64 on disk, got {ds[coord_name].dtype}" + # encoding is deferred — in-memory dtype is not mutated + assert da.coords[coord_name].dtype == input_dtype @pytest.mark.parametrize( "bounds_name", ["lat_bnds", "lon_bnds", "lat_bounds", "lon_bounds"] ) -def test_save_dataset_casts_geo_bounds_to_float64(tmp_path, bounds_name): +def test_save_dataset_writes_geo_bounds_as_float64(tmp_path, bounds_name): """Geographic bounds variables must also be written as float64.""" dates = xr.cftime_range(start="2001", periods=2, freq="MS", calendar="noleap") coord_name = "lat" if "lat" in bounds_name else "lon" @@ -379,4 +385,6 @@ def test_save_dataset_casts_geo_bounds_to_float64(tmp_path, bounds_name): with xr.open_dataset(saved[0]) as result: assert ( result[bounds_name].dtype == np.float64 - ), f"{bounds_name} should be float64, got {result[bounds_name].dtype}" + ), f"{bounds_name} should be float64 on disk, got {result[bounds_name].dtype}" + # encoding is deferred — in-memory dtype is not mutated + assert ds[bounds_name].dtype == np.float32