diff --git a/src/access_moppy/base.py b/src/access_moppy/base.py index b550fd3a..83a13731 100644 --- a/src/access_moppy/base.py +++ b/src/access_moppy/base.py @@ -296,6 +296,24 @@ def __setitem__(self, key, value): def __repr__(self): return repr(self.ds) + def _is_fx_variable(self) -> bool: + """Return True when compound_name corresponds to a fixed-field CMIP table.""" + if not self.compound_name: + return False + + table_id = self.compound_name.split(".", 1)[0].lower() + return table_id.endswith("fx") + + def _squeeze_fx_singleton_time(self) -> None: + """Drop UM-style singleton time axis for fixed (fx) variables.""" + if ( + self.ds is not None + and self._is_fx_variable() + and "time" in self.ds.dims + and self.ds.sizes.get("time") == 1 + ): + self.ds = self.ds.isel(time=0, drop=True) + def load_dataset(self, required_vars: Optional[List[str]] = None): """ Load dataset from input files or use provided xarray objects with optional frequency validation. @@ -363,11 +381,30 @@ def load_dataset(self, required_vars: Optional[List[str]] = None): # Original file-based loading logic def _preprocess(ds): ds = ds[list(required_vars & set(ds.data_vars))] - # Drop auxiliary UM time coordinates (time_0, time_1) that differ - # across files. Without this, xr.open_mfdataset's join='outer' - # unions every distinct value into a growing dimension, inflating - # the Dask task graph to several GiB before any computation starts. - aux_time_coords = [c for c in ("time_0", "time_1") if c in ds] + # Canonicalize UM auxiliary time dimensions (time_0/time_1) to + # a single "time" axis when the selected variables use exactly + # one such axis. Keep that primary axis and drop the unused one. + selected_data_vars = list(ds.data_vars) + used_time_dims = { + dim + for var in selected_data_vars + for dim in ds[var].dims + if dim.startswith("time") + } + primary_time_dim = "time" if "time" in used_time_dims else None + if primary_time_dim is None and len(used_time_dims) == 1: + primary_time_dim = next(iter(used_time_dims)) + + if primary_time_dim and primary_time_dim != "time": + ds = ds.rename({primary_time_dim: "time"}) + used_time_dims.discard(primary_time_dim) + used_time_dims.add("time") + + aux_time_coords = [ + c + for c in ("time_0", "time_1") + if c in ds.coords and c not in used_time_dims + ] if aux_time_coords: ds = ds.drop_vars(aux_time_coords) return ds @@ -383,8 +420,9 @@ def _preprocess(ds): if required_vars else list(_probe.data_vars) ) - _has_time = (required_vars is None or "time" in required_vars) and any( - "time" in _probe[v].dims for v in _probe_target_vars + _has_time = any( + any(dim.startswith("time") for dim in _probe[v].dims) + for v in _probe_target_vars ) # Validate frequency consistency and CMIP6 compatibility before concatenation @@ -394,7 +432,7 @@ def _preprocess(ds): # Time-independent variables typically have "fx" (fixed) in their table ID is_time_independent = ( self.compound_name and "fx" in self.compound_name.lower() - ) or "time" not in _probe_dims + ) or not any(dim.startswith("time") for dim in _probe_dims) if is_time_independent: logger.debug( @@ -496,11 +534,13 @@ def _preprocess(ds): if required_vars: vars_to_keep = [v for v in required_vars if v in self.ds.data_vars] self.ds = self.ds[vars_to_keep] - # UM source files always include a time dimension (size=1) even for - # static fields. Drop it so downstream CMOR processing sees the - # expected (lat, lon) shape rather than (time=1, lat, lon). - if "time" in self.ds.dims: - self.ds = self.ds.isel(time=0, drop=True) + # UM source files can include a time=1 axis for static fields. + # Keep squeeze behavior centralized in _squeeze_fx_singleton_time(). + + # UM source files can carry time=1 for fixed fields even when loaded + # through the time-aware branch. Squeeze once here before any downstream + # frequency handling, rechunking, or missing-value normalization. + self._squeeze_fx_singleton_time() # Apply temporal resampling if enabled and needed if self.enable_resampling and self.compound_name: @@ -790,6 +830,93 @@ def _numeric_delta_to_days(delta: float, time_units: Optional[str]) -> float: return delta / 86400.0 return delta + @staticmethod + def _days_to_numeric_units(days: float, time_units: Optional[str]) -> float: + """Convert day-length values back to the numeric coordinate units.""" + if not time_units: + return days + + interval = str(time_units).split("since", 1)[0].strip().lower() + if interval in {"day", "days", "d"}: + return days + if interval in {"hour", "hours", "hr", "hrs", "h"}: + return days * 24.0 + if interval in {"minute", "minutes", "min", "mins", "m"}: + return days * 1440.0 + if interval in {"second", "seconds", "sec", "secs", "s"}: + return days * 86400.0 + return days + + def _align_subdaily_point_time_to_square_grid(self) -> None: + """Align point-sampled sub-daily timestamps to day-boundary slots. + + Some ACCESS streams begin point-sampled sub-daily series at +3h/+6h. + WCRP TIME001 expects these frequencies to be square on the canonical + daily grid. Shift the whole time axis (and time bounds if present) by + that leading offset. + """ + if "time" not in self.ds.coords or self.cmor_name not in self.ds: + return + + table_id = (self.compound_name or "").split(".", 1)[0] + table_hours = { + "3hr": 3.0, + "3hrPt": 3.0, + "6hrPlevPt": 6.0, + } + target_hours = table_hours.get(table_id) + if target_hours is None: + return + + cell_methods = str(self.ds[self.cmor_name].attrs.get("cell_methods", "")) + if "time: point" not in cell_methods: + return + + time_da = self.ds["time"] + values = np.asarray(time_da.values) + if values.size == 0 or not np.issubdtype(values.dtype, np.number): + return + + time_units = time_da.attrs.get("units") + if not isinstance(time_units, str) or "since" not in time_units: + return + + first_value = float(values.flat[0]) + first_days = self._numeric_delta_to_days(first_value, time_units) + leading_offset_days = first_days - np.floor(first_days) + if np.isclose(leading_offset_days, 0.0, atol=1e-10): + return + + target_days = target_hours / 24.0 + if not np.isclose(leading_offset_days, target_days, atol=1e-8): + return + + shift_units = self._days_to_numeric_units(leading_offset_days, time_units) + self.ds["time"] = xr.DataArray( + values - shift_units, + dims=time_da.dims, + coords=time_da.coords, + attrs=time_da.attrs, + ) + + bounds_name = time_da.attrs.get("bounds") + if isinstance(bounds_name, str) and bounds_name in self.ds: + bnds = self.ds[bounds_name] + if np.issubdtype(np.asarray(bnds.values).dtype, np.number): + self.ds[bounds_name] = xr.DataArray( + bnds.values - shift_units, + dims=bnds.dims, + coords=bnds.coords, + attrs=bnds.attrs, + ) + + logger.debug( + "Shifted '%s' time axis by %.6f %s to align sub-daily point timestamps", + self.cmor_name, + shift_units, + time_units.split("since", 1)[0].strip(), + ) + def rechunk_dataset(self): """ Apply intelligent rechunking to the dataset. @@ -1146,6 +1273,10 @@ def write(self): self.ds[var].attrs["units"] = normalized logger.debug("Normalized '%s' units: %r -> %r", var, units, normalized) + # Align point-sampled sub-daily timestamps (e.g. 3hrPt/6hrPlevPt) + # to WCRP TIME001 square-grid expectations before writing. + self._align_subdaily_point_time_to_square_grid() + # ========== Prepare String Coordinates ========== # Detect and prepare all string/character coordinates before writing string_coords_info = self._prepare_string_coordinates() diff --git a/src/access_moppy/driver.py b/src/access_moppy/driver.py index 65ce29c1..014e0a88 100644 --- a/src/access_moppy/driver.py +++ b/src/access_moppy/driver.py @@ -482,7 +482,9 @@ def __init__( "CFmon", "CFday", "3hr", + "3hrPt", "6hrPlev", + "6hrPlevPt", "E1hr", "Eday", "fx", @@ -553,7 +555,7 @@ def __init__( raise ValueError( f"Unsupported CMIP table '{table}' in compound_name '{compound_name}'. " f"Supported legacy CMIP6 tables — " - f"atmosphere: ('Amon', 'Lmon', 'LImon', 'Emon', 'AERmon', 'AERday', 'day', 'CFmon', 'CFday', '3hr', '6hrPlev', 'E1hr', 'Eday', 'fx', 'Efx', 'atmos'), " + f"atmosphere: ('Amon', 'Lmon', 'LImon', 'Emon', 'AERmon', 'AERday', 'day', 'CFmon', 'CFday', '3hr', '3hrPt', '6hrPlev', '6hrPlevPt', 'E1hr', 'Eday', 'fx', 'Efx', 'atmos'), " f"ocean: ('Oyr', 'Oday', 'Omon', 'Ofx'), " f"sea-ice: ('SImon', 'SIday'). " f"MIP CMOR table prefixes are also supported: " diff --git a/src/access_moppy/examples/batch_config_esm1-6_cmip7_baseline.yml b/src/access_moppy/examples/batch_config_esm1-6_cmip7_baseline.yml index 8e63cc36..c2b55a73 100644 --- a/src/access_moppy/examples/batch_config_esm1-6_cmip7_baseline.yml +++ b/src/access_moppy/examples/batch_config_esm1-6_cmip7_baseline.yml @@ -109,7 +109,7 @@ variables: - seaIce.siconc.tavg-u-hxy-u.mon.glb # SImon.siconc - seaIce.simass.tavg-u-hxy-si.mon.glb # SImon.simass # - seaIce.snd.tavg-u-hxy-sn.mon.glb # SImon.sisnthick — NOT MAPPED - - seaIce.ts.tavg-u-hxy-si.mon.glb # SImon.sitemptop + # - seaIce.ts.tavg-u-hxy-si.mon.glb # SImon.sitemptop - seaIce.sithick.tavg-u-hxy-si.mon.glb # SImon.sithick - seaIce.sitimefrac.tavg-u-hxy-sea.mon.glb # SImon.sitimefrac - seaIce.siu.tavg-u-hxy-si.mon.glb # SImon.siu diff --git a/src/access_moppy/utilities.py b/src/access_moppy/utilities.py index b1765bd3..8dbf52ad 100644 --- a/src/access_moppy/utilities.py +++ b/src/access_moppy/utilities.py @@ -694,6 +694,7 @@ def parse_cmip6_table_frequency(compound_name: str) -> pd.Timedelta: "SIday": pd.Timedelta(days=1), # Additional frequency tables "3hr": pd.Timedelta(hours=3), + "3hrPt": pd.Timedelta(hours=3), "6hr": pd.Timedelta(hours=6), "day": pd.Timedelta(days=1), "mon": pd.Timedelta(days=30),