save work

Author: jedwards4b

.pylintrc

@@ -14,5 +14,5 @@ disable=
 max-line-length=100
 
 [TYPECHECK]
-ignored-modules=xarray,xesmf,cmor,numpy,geocat.comp,yaml
+ignored-modules=xarray,xesmf,cmor,numpy,geocat.comp,yaml,cftime
 ignored-classes=cmor

cmip7_prep/cli.py

@@ -74,7 +74,11 @@ def main(argv: list[str] | None = None) -> int:
         "--cmor-tables", required=True, type=Path, help="Path to CMOR Tables directory."
     )
     p_prep.add_argument(
-        "--dataset-json", required=True, type=Path, help="Path to cmor_dataset.json."
+        "--dataset-json",
+        required=True,
+        type=Path,
+        help="Path to cmor_dataset.json.",
+        default=Path("data/cmor_dataset.json"),
     )
     p_prep.add_argument(
         "--outdir",

cmip7_prep/cmor_writer.py

@@ -3,20 +3,72 @@
 This module centralizes CMOR session setup and writing so that the rest of the
 pipeline can stay xarray-first. It supports either a dataset JSON file (preferred)
 or directly injected global attributes, and creates axes based on the coordinates
-present in the provided dataset.
+present in the provided dataset. It also supports a packaged default
+`cmor_dataset.json` living under `cmip7_prep/data/`.
 """
 
 from __future__ import annotations
 
 from contextlib import AbstractContextManager
+from importlib.resources import as_file, files
 from pathlib import Path
 from typing import Any, Dict, Optional
 
-import cmor  # type: ignore
 import numpy as np
 import xarray as xr
+from xarray.coding.times import encode_cf_datetime
+
+import cmor  # type: ignore
+
+
+# ---------------------------------------------------------------------
+# Packaged resource helpers
+# ---------------------------------------------------------------------
+def packaged_dataset_json(filename: str = "cmor_dataset.json") -> Any:
+    """Return a context manager yielding a real FS path to a packaged dataset JSON.
+
+    Looks under cmip7_prep/data/.
+    Usage:
+        with packaged_dataset_json() as p:
+            cmor.dataset_json(str(p))
+    """
+    res = files("cmip7_prep").joinpath(f"data/{filename}")
+    return as_file(res)
 
 
+# ---------------------------------------------------------------------
+# Time encoding
+# ---------------------------------------------------------------------
+def _encode_time_to_num(time_da: xr.DataArray, units: str, calendar: str) -> np.ndarray:
+    """Return numeric CF time values (float64) acceptable to CMOR.
+
+    Tries xarray's encoder first; if that fails and cftime is available,
+    falls back to cftime.date2num. Raises a ValueError with details otherwise.
+    """
+    # 1) xarray encoder (handles numpy datetime64 and cftime objects if cftime present)
+    try:
+        out = encode_cf_datetime(time_da.values, units=units, calendar=calendar)
+        return np.asarray(out, dtype="f8")
+    except (ValueError, TypeError) as exc_xr:
+        last_err = exc_xr
+
+    # 2) Optional cftime path (lazy import to keep lint/typecheck happy)
+    try:
+        import cftime  # type: ignore # pylint: disable=import-outside-toplevel
+
+        seq = time_da.values.tolist()  # handles object arrays / cftime arrays
+        out = cftime.date2num(seq, units=units, calendar=calendar)
+        return np.asarray(out, dtype="f8")
+    except Exception as exc_cf:  # noqa: BLE001 - we surface both causes together
+        raise ValueError(
+            f"Could not encode time to numeric CF values with units={units!r}, "
+            f"calendar={calendar!r}. xarray error: {last_err}; cftime error: {exc_cf}"
+        ) from exc_cf
+
+
+# ---------------------------------------------------------------------
+# CMOR session
+# ---------------------------------------------------------------------
 class CmorSession(AbstractContextManager):
     """Context manager for CMOR sessions.
 
@@ -43,74 +95,72 @@ def __init__(
     def __enter__(self) -> "CmorSession":
         """Initialize CMOR and register dataset metadata."""
         cmor.setup(inpath=self.tables_path, netcdf_file_action=cmor.CMOR_REPLACE_3)
+
         if self.dataset_json_path:
             cmor.dataset_json(self.dataset_json_path)
         elif self.dataset_attrs:
             for key, value in self.dataset_attrs.items():
                 cmor.setGblAttr(key, value)
         else:
-            raise ValueError(
-                "CmorSession requires either dataset_json path or dataset_attrs."
-            )
+            # Fallback to packaged cmor_dataset.json if available
+            with packaged_dataset_json() as p:
+                cmor.dataset_json(str(p))
+
         return self
 
-    def __exit__(self, exc_type, exc, tb) -> None:
+    def __exit__(self, exc_type, exc, tb) -> None:  # noqa: D401
         """Finalize CMOR, closing any open handles."""
         cmor.close()
 
     # -------------------------
     # internal helpers
     # -------------------------
-
     def _define_axes(self, ds: xr.Dataset, vdef: Any) -> list[int]:
-        """Create CMOR axis IDs based on dataset coordinates and optional vdef levels.
-
-        Parameters
-        ----------
-        ds : xr.Dataset
-            Dataset containing coordinates like time, lat, lon, and optionally plev/lev.
-        vdef : Any
-            Variable definition object. If it has a ``levels`` mapping (e.g.,
-            ``{'axis_entry': 'plev19', 'name': 'plev', 'units': 'Pa'}``), that will
-            be used to define the vertical axis.
-
-        Returns
-        -------
-        list[int]
-            List of CMOR axis IDs in the order they should be used for the variable.
-        """
+        """Create CMOR axis IDs based on dataset coordinates and optional vdef levels."""
         axes: list[int] = []
 
-        # Time axis
+        # ---- time axis ----
         if "time" in ds.coords or "time" in ds:
             time = ds["time"]
             t_units = time.attrs.get("units", "days since 1850-01-01")
-            cal = time.attrs.get("calendar", "standard")
-            if np.issubdtype(time.dtype, np.datetime64):
-                # convert to numeric time using CF units if needed
-                tvals = xr.conventions.times.encode_cf_datetime(
-                    time.values, t_units, calendar=cal
-                )
-            else:
-                tvals = time.values
+            cal = time.attrs.get("calendar", time.encoding.get("calendar", "noleap"))
+            tvals = _encode_time_to_num(time, t_units, cal)
+
+            # Optional bounds: try common names or CF 'bounds' attribute
+            tb = ds.get("time_bnds") or ds.get("time_bounds")
+            if tb is None and isinstance(time.attrs.get("bounds"), str):
+                bname = time.attrs["bounds"]
+                tb = ds.get(bname)
+
+            t_bnds = None
+            if tb is not None:
+                try:
+                    t_bnds = _encode_time_to_num(tb, t_units, cal)
+                except ValueError:
+                    t_bnds = None
+
             axes.append(
-                cmor.axis(table_entry="time", units=str(t_units), coord_vals=tvals)
+                cmor.axis(
+                    table_entry="time",
+                    units=str(t_units),
+                    coord_vals=tvals,
+                    cell_bounds=t_bnds if t_bnds is not None else None,
+                )
             )
 
-        # Vertical axis (pressure or model levels)
+        # ---- vertical axis (plev or lev) ----
         levels_info = getattr(vdef, "levels", None)
-        # Prefer an explicit plev axis in the data
         if "plev" in ds.coords:
             p = ds["plev"]
             p_units = p.attrs.get("units", "Pa")
-            # If vdef specifies an axis entry (e.g., plev19), use that; else let CMOR infer
+            # If vdef specifies an axis entry (e.g., plev19), use that; else infer
             table_entry = None
             if isinstance(levels_info, dict):
                 table_entry = levels_info.get("axis_entry") or levels_info.get("name")
             if table_entry is None:
-                # Heuristic: choose plev19 if there are 19 levels, else generic plev
                 table_entry = "plev19" if p.size == 19 else "plev"
-            p_bnds = ds.get("plev_bnds")
+
+            p_bnds = ds.get("plev_bnds") or ds.get("plev_bounds")
             axes.append(
                 cmor.axis(
                     table_entry=table_entry,
@@ -120,7 +170,6 @@ def _define_axes(self, ds: xr.Dataset, vdef: Any) -> list[int]:
                 )
             )
         elif "lev" in ds.coords:
-            # Generic hybrid "lev" axis; rely on table entry provided via vdef or default to "alev"
             lev = ds["lev"]
             table_entry = "alev"
             if isinstance(levels_info, dict):
@@ -129,12 +178,13 @@ def _define_axes(self, ds: xr.Dataset, vdef: Any) -> list[int]:
                 cmor.axis(table_entry=table_entry, units="1", coord_vals=lev.values)
             )
 
-        # Latitude / Longitude
+        # ---- horizontal axes ----
         if "lat" in ds.coords and "lon" in ds.coords:
             lat = ds["lat"]
             lon = ds["lon"]
-            lat_b = ds.get("lat_bnds")
-            lon_b = ds.get("lon_bnds")
+            lat_b = ds.get("lat_bnds") or ds.get("lat_bounds")
+            lon_b = ds.get("lon_bnds") or ds.get("lon_bounds")
+
             axes.append(
                 cmor.axis(
                     table_entry="lat",
@@ -157,58 +207,48 @@ def _define_axes(self, ds: xr.Dataset, vdef: Any) -> list[int]:
     # -------------------------
     # public API
     # -------------------------
-
     def write_variable(
         self, ds: xr.Dataset, varname: str, vdef: Any, outdir: Path
     ) -> None:
-        """Write one variable from `ds` to a CMOR-compliant NetCDF file.
-
-        Parameters
-        ----------
-        ds : xr.Dataset
-            Dataset containing the variable and its coordinates.
-        varname : str
-            Name of the variable in `ds` to CMORize.
-        vdef : Any
-            An object with fields: ``name``, ``realm``,
-            optional ``units``, ``positive``, and optional
-            ``levels`` dict (see :meth:`_define_axes`).
-            This is typically a light-weight holder.
-        outdir : Path
-            Output directory for the CMORized NetCDF file.
-        """
-        # Load the appropriate CMOR table without relying on exceptions.
-        realm = getattr(vdef, "realm", "Amon")
-        candidate7 = Path(self.tables_path) / f"CMIP7_{realm}.json"
-        candidate6 = Path(self.tables_path) / f"CMIP6_{realm}.json"
+        """Write one variable from `ds` to a CMOR-compliant NetCDF file."""
+        # Pick CMOR table: prefer vdef.table, else vdef.realm (default Amon)
+        table_key = getattr(vdef, "table", None) or getattr(vdef, "realm", "Amon")
+        table_key = str(table_key)
+        candidate7 = Path(self.tables_path) / f"CMIP7_{table_key}.json"
+        candidate6 = Path(self.tables_path) / f"CMIP6_{table_key}.json"
+
         if candidate7.exists():
             cmor.load_table(str(candidate7))
         elif candidate6.exists():
             cmor.load_table(str(candidate6))
         else:
-            # Fall back to passing the bare table name; CMOR will search its inpath.
-            # This branch avoids broad exception handling while still being flexible.
-            cmor.load_table(f"CMIP7_{realm}.json")
+            # Let CMOR search inpath; will raise if not found
+            cmor.load_table(f"CMIP7_{table_key}.json")
 
         axes_ids = self._define_axes(ds, vdef)
-        units = getattr(vdef, "units", "")
+
+        units = getattr(vdef, "units", "") or ""
         var_id = cmor.variable(
             getattr(vdef, "name", varname),
             units,
             axes_ids,
             positive=getattr(vdef, "positive", None),
         )
 
-        # Optional variable attributes (e.g., cell_methods)
+        # Optional variable attributes (e.g., cell_methods, long_name, standard_name)
         if getattr(vdef, "cell_methods", None):
             cmor.set_variable_attribute(var_id, "cell_methods", vdef.cell_methods)
+        if getattr(vdef, "long_name", None):
+            cmor.set_variable_attribute(var_id, "long_name", vdef.long_name)
+        if getattr(vdef, "standard_name", None):
+            cmor.set_variable_attribute(var_id, "standard_name", vdef.standard_name)
 
-        # Ensure time is the leading dimension if present
         data = ds[varname]
         if "time" in data.dims:
             data = data.transpose("time", ...)
 
-        cmor.write(var_id, data.values, ntimes_passed=data.sizes.get("time", 1))
+        # CMOR expects a NumPy array; this will materialize data as needed.
+        cmor.write(var_id, np.asarray(data), ntimes_passed=data.sizes.get("time", 1))
 
         outdir = Path(outdir)
         outdir.mkdir(parents=True, exist_ok=True)

cmip7_prep/mapping_compat.py

@@ -20,7 +20,7 @@
 
 def packaged_mapping_resource(filename: str = "cesm_to_cmip7.yaml"):
     """Context manager yielding a real filesystem path to the packaged mapping file."""
-    res = files("cmip7_prep").joinpath(f"mapping/{filename}")
+    res = files("cmip7_prep").joinpath(f"data/{filename}")
     return as_file(res)
 
 
@@ -119,7 +119,11 @@ def _load_yaml(path: Path) -> Dict[str, VarConfig]:
             data = yaml.safe_load(f)
 
         # Accept top-level "variables:" wrapper
-        if isinstance(data, dict) and "variables" in data and isinstance(data["variables"], dict):
+        if (
+            isinstance(data, dict)
+            and "variables" in data
+            and isinstance(data["variables"], dict)
+        ):
             data = data["variables"]
 
         result: Dict[str, VarConfig] = {}

pyproject.toml

@@ -8,10 +8,10 @@ readme = "README.md"
 packages = [{ include = "cmip7_prep" }]
 # Include non-Python data files in both sdist and wheel
 include = [
-  { path = "cmip7_prep/mapping/*.yaml", format = "sdist" },
-  { path = "cmip7_prep/mapping/*.yaml", format = "wheel" },
-  { path = "cmip7_prep/mapping/*.json", format = "sdist" },
-  { path = "cmip7_prep/mapping/*.json", format = "wheel" },
+  { path = "cmip7_prep/data/*.yaml", format = "sdist" },
+  { path = "cmip7_prep/data/*.yaml", format = "wheel" },
+  { path = "cmip7_prep/data/*.json", format = "sdist" },
+  { path = "cmip7_prep/data/*.json", format = "wheel" },
 ]
 
 [tool.poetry.dependencies]

scripts/testone.py

@@ -50,7 +50,9 @@
     xr.Dataset({var: da}), var, output_time_chunk=time_chunk, dtype="float32"
 )
 
-with CmorSession(tables_path="Tables", dataset_json="cmor_dataset.json") as cm:
+with CmorSession(
+    tables_path="/glade/work/cmip7/e3sm_to_cmip/cmip6-cmor-tables/Tables"
+) as cm:
     vdef = type(
         "VDef", (), {"name": var, "realm": "Amon", "units": cfg.get("units", "")}
     )()