Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 53 additions & 13 deletions src/access_moppy/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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(
Expand Down Expand Up @@ -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:
Expand Down
4 changes: 3 additions & 1 deletion src/access_moppy/driver.py
Original file line number Diff line number Diff line change
Expand Up @@ -482,7 +482,9 @@ def __init__(
"CFmon",
"CFday",
"3hr",
"3hrPt",
"6hrPlev",
"6hrPlevPt",
"E1hr",
"Eday",
"fx",
Expand Down Expand Up @@ -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: "
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions src/access_moppy/utilities.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down