Merge pull request #2 from ESMCI/jpe/merge_branch

Jpe/merge branch

Author: jedwards4b

.github/workflows/pylint.yml

@@ -36,3 +36,8 @@ jobs:
         run: poetry install --no-interaction --no-ansi
       - name: Run pre-commit on all files
         run: poetry run pre-commit run --all-files
+#     the following can be used by developers to login to the github server in case of errors
+#     see https://github.com/marketplace/actions/debugging-with-tmate for further details
+#      - name: Setup tmate session
+#        if: ${{ failure() }}
+#        uses: mxschmitt/action-tmate@v3

cmip7_prep/cmor_writer.py

@@ -1,3 +1,4 @@
+# pylint: disable=too-many-lines
 """Thin CMOR wrapper used by cmip7_prep.
 
 This module centralizes CMOR session setup and writing so that the rest of the
@@ -122,6 +123,47 @@ def _encode_time_to_num(obj, units: str, calendar: str) -> np.ndarray:
     return nums.reshape(arr.shape)
 
 
+def _bounds_from_centers_1d(vals: np.ndarray, kind: str) -> np.ndarray:
+    """Compute [n,2] cell bounds from 1-D centers for 'lat' or 'lon'.
+
+    - For 'lat': clamps to [-90, 90]
+    - For 'lon': treats as periodic [0, 360)
+    - Works with non-uniform spacing (uses midpoints between neighbors)
+    """
+    v = np.asarray(vals, dtype="f8").reshape(-1)
+    n = v.size
+    if n < 2:
+        raise ValueError("Need at least 2 points to compute bounds")
+
+    # neighbor midpoints
+    mid = 0.5 * (v[1:] + v[:-1])  # length n-1
+    bounds = np.empty((n, 2), dtype="f8")
+    bounds[1:, 0] = mid
+    bounds[:-1, 1] = mid
+
+    # end caps: extrapolate by half-step at ends
+    first_step = v[1] - v[0]
+    last_step = v[-1] - v[-2]
+    bounds[0, 0] = v[0] - 0.5 * first_step
+    bounds[-1, 1] = v[-1] + 0.5 * last_step
+
+    if kind == "lat":
+        # clamp to physical limits
+        bounds[:, 0] = np.maximum(bounds[:, 0], -90.0)
+        bounds[:, 1] = np.minimum(bounds[:, 1], 90.0)
+    elif kind == "lon":
+        # wrap to [0, 360)
+        bounds = bounds % 360.0
+        # ensure each row is increasing in modulo arithmetic
+        wrap = bounds[:, 1] < bounds[:, 0]
+        if np.any(wrap):
+            bounds[wrap, 1] += 360.0
+    else:
+        raise ValueError("kind must be 'lat' or 'lon'")
+
+    return bounds
+
+
 def _encode_time_bounds_to_num(tb, units: str, calendar: str) -> np.ndarray:
     """
     Encode bounds array of shape (..., 2) to numeric CF time.
@@ -382,6 +424,31 @@ def _resolve_table_filename(tables_path: Path, key: str) -> str:
 DatasetJsonLike = Union[str, Path, AbstractContextManager]
 
 
+def _fx_glob_pattern(name: str) -> str:
+    # CMOR filenames vary; this finds most fx files for this var
+    # e.g., *_sftlf_fx_*.nc  or sftlf_fx_*.nc
+    return f"**/*_{name}_fx_*.nc"
+
+
+def _open_existing_fx(outdir: Path, name: str) -> xr.DataArray | None:
+    # Search recursively for an existing fx file for this var
+    for p in outdir.rglob(_fx_glob_pattern(name)):
+        try:
+            ds = xr.open_dataset(p, engine="netcdf4")
+            if name in ds:
+                return ds[name]
+        except FileNotFoundError:
+            return None
+        except (OSError, ValueError) as e:
+            # OSError: unreadable/corrupt file, low-level I/O; ValueError: engine/decoding issues
+            warnings.warn(f"[fx] failed to open {p} with netcdf4: {e}", RuntimeWarning)
+        except (ImportError, ModuleNotFoundError) as e:
+            # netCDF4 backend not installed
+            warnings.warn(f"[fx] netcdf4 backend unavailable: {e}", RuntimeWarning)
+
+    return None
+
+
 # ---------------------------------------------------------------------
 # CMOR session
 # ---------------------------------------------------------------------
@@ -422,8 +489,14 @@ def __init__(
         self._log_name = log_name
         self._log_path: Path | None = None
         self._pending_ps = None
-        self._outdir = Path(outdir) if outdir is not None else Path.cwd() / "CMIP7"
+        self._outdir = Path(outdir or "./CMIP7").resolve()
         self._outdir.mkdir(parents=True, exist_ok=True)
+        self._fx_written: set[str] = (
+            set()
+        )  # remembers which fx vars were written this run
+        self._fx_cache: dict[str, xr.DataArray] = (
+            {}
+        )  # regridded fx fields cached in-memory
 
     def __enter__(self) -> "CmorSession":
         # Resolve logfile path if requested
@@ -765,6 +838,100 @@ def _get_1d_with_bounds(dsi: xr.Dataset, name: str, units_default: str):
         axes_ids.extend([lat_id, lon_id])
         return axes_ids
 
+    def _write_fx_2d(self, ds: xr.Dataset, name: str, units: str) -> None:
+        if name not in ds:
+            return
+        table_filename = _resolve_table_filename(self.tables_path, "fx")
+        cmor.load_table(table_filename)
+
+        lat = ds["lat"].values
+        lon = ds["lon"].values
+        lat_b = ds.get("lat_bnds")
+        lon_b = ds.get("lon_bnds")
+        lat_b = (
+            lat_b.values
+            if isinstance(lat_b, xr.DataArray)
+            else _bounds_from_centers_1d(lat, "lat")
+        )
+        lon_b = (
+            lon_b.values
+            if isinstance(lon_b, xr.DataArray)
+            else _bounds_from_centers_1d(lon, "lon")
+        )
+
+        lat_id = cmor.axis(
+            "latitude", "degrees_north", coord_vals=lat, cell_bounds=lat_b
+        )
+        lon_id = cmor.axis(
+            "longitude", "degrees_east", coord_vals=lon, cell_bounds=lon_b
+        )
+        data_filled, fillv = _filled_for_cmor(ds[name])
+
+        var_id = cmor.variable(name, units, [lat_id, lon_id], missing_value=fillv)
+        print(f"write fx variable {name}")
+        cmor.write(
+            var_id,
+            np.asarray(data_filled),
+        )
+        cmor.close(var_id)
+
+    def ensure_fx_written_and_cached(self, ds_regr: xr.Dataset) -> xr.Dataset:
+        """Ensure sftlf and areacella exist in ds_regr and are written once as fx.
+        If not present in ds_regr, try to read from existing CMOR fx files in outdir.
+        If present in ds_regr but not yet written this run, write and cache them.
+        Returns ds_regr augmented with any missing fx fields.
+        """
+        need = [("sftlf", "%"), ("areacella", "m2")]
+        out = ds_regr
+
+        for name, units in need:
+            # 1) Already cached this run?
+            if name in self._fx_cache:
+                if name not in out:
+                    out = out.assign({name: self._fx_cache[name]})
+                continue
+
+            # 2) Present in regridded dataset? (best case)
+            if name in out:
+                self._fx_cache[name] = out[name]
+                if name not in self._fx_written:
+                    # Convert landfrac to % if needed
+                    if name == "sftlf":
+                        v = out[name]
+                        if (np.nanmax(v.values) <= 1.0) and v.attrs.get(
+                            "units", ""
+                        ) not in ("%", "percent"):
+                            out = out.assign(
+                                {
+                                    name: (v * 100.0).assign_attrs(
+                                        v.attrs | {"units": "%"}
+                                    )
+                                }
+                            )
+                            self._fx_cache[name] = out[name]
+                    self._write_fx_2d(out, name, units)
+                    self._fx_written.add(name)
+                    continue
+
+            # 3) Not present in ds_regr → try reading existing CMOR fx output
+            if self._outdir:
+                fx_da = _open_existing_fx(self._outdir, name)
+                if fx_da is not None:
+                    # Verify grid match (simple equality on lat/lon values)
+                    if (
+                        "lat" in out
+                        and "lon" in out
+                        and np.array_equal(out["lat"].values, fx_da["lat"].values)
+                        and np.array_equal(out["lon"].values, fx_da["lon"].values)
+                    ):
+                        out = out.assign({name: fx_da})
+                        self._fx_cache[name] = out[name]
+                        self._fx_written.add(name)  # already exists on disk
+                        continue
+                # If grid mismatch, you could regrid fx_da here; for now, skip.
+                # 4) Last resort: leave missing; caller may compute it later
+        return out
+
     # public API
     # -------------------------
     def write_variable(
@@ -801,6 +968,8 @@ def write_variable(
             time_da = ds.get("time")
         nt = 0
 
+        self.ensure_fx_written_and_cached(ds)
+
         # ---- Main variable write ----
 
         cmor.write(

cmip7_prep/data/cesm_to_cmip7.yaml

@@ -87,6 +87,24 @@ variables:
       - cesm_var: QFLX
         scale: -1.0     # flip sign if CESM positive upward
 
+  evspsblsoi:
+    table: Lmon
+    units: kg m-2 s-1
+    long_name: Water Evaporation from Soil
+    dims: [time, lat, lon]
+    positive: up
+    sources:
+      - cesm_var: QSOIL
+
+  evspsblveg:
+    table: Lmon
+    units: kg m-2 s-1
+    long_name: Water Evaporation from Soil
+    dims: [time, lat, lon]
+    positive: up
+    sources:
+      - cesm_var: QVEGE
+
   hfls:
     table: Amon
     units: "W m-2"
@@ -139,6 +157,51 @@ variables:
     sources:
       - cesm_var: QREFHT
 
+  lai:
+    table: Lmon
+    units: "1"
+    dims: [time, lat, lon]
+    sources:
+      - cesm_var: TLAI
+
+  mrfso:
+    table: Lmon
+    units: kg m-2
+    dims: [time, lat, lon]
+    formula: verticalsum(SOILICE, capped_at=5000)
+    sources:
+      - cesm_var: SOILICE
+
+  mrro:
+    table: Lmon
+    units: kg m-2 s-1
+    dims: [time, lat, lon]
+    sources:
+      - cesm_var: QRUNOFF
+
+  mrros:
+    table: Lmon
+    units: kg m-2 s-1
+    dims: [time, lat, lon]
+    sources:
+      - cesm_var: QOVER
+
+  mrso:
+    table: Lmon
+    units: kg m-2
+    dims: [time, lat, lon]
+    formula: verticalsum(SOILICE + SOILLIQ, capped_at=5000)
+    sources:
+      - cesm_var: SOILICE
+      - cesm_var: SOILLIQ
+
+  mrsos:
+    table: Lmon
+    units: kg m-2
+    dims: [time, lat, lon]
+    sources:
+      - cesm_var: SOILWATER_10CM
+
   pr:
     table: Amon
     units: kg m-2 s-1
@@ -432,10 +495,20 @@ variables:
     sources:
       - cesm_var: Z3
 
-fx:
-  areacella:
-    units: m2
-    source: cell_area
-  sftlf:
-    units: '%'
-    source: sftlf
+#  areacella:
+#    table: fx
+#    units: m2
+#    formula: area * 1.e6
+#    standard_name: cell_area
+#    dims: [lat, lon]
+#    sources:
+#      - cesm_var: ZBOT
+#
+#
+#  sftlf:
+#    table: fx
+#    units: '%'
+#    formula: landfrac * 100
+#    dims: [lat, lon]
+#    sources:
+#      - cesm_var: ZBOT

cmip7_prep/mapping_compat.py

@@ -102,7 +102,23 @@ def _safe_eval(expr: str, local_names: Dict[str, Any]) -> Any:
     2.0
     """
     safe_globals = {"__builtins__": {}}
-    locals_safe = {"np": np, "xr": xr}
+
+    # Add custom formula functions here
+    def verticalsum(arr, capped_at=None, dim="levsoi"):
+        # arr can be a DataArray or an expression
+        if isinstance(arr, xr.DataArray):
+            summed = arr.sum(dim=dim, skipna=True)
+        else:
+            summed = arr  # fallback, should be DataArray
+        if capped_at is not None:
+            summed = xr.where(summed > capped_at, capped_at, summed)
+        return summed
+
+    locals_safe = {
+        "np": np,
+        "xr": xr,
+        "verticalsum": verticalsum,
+    }
     locals_safe.update(local_names)
     # pylint: disable=eval-used
     return eval(expr, safe_globals, locals_safe)
@@ -274,6 +290,11 @@ def _realize_core(ds: xr.Dataset, vc: VarConfig) -> xr.DataArray:
             raise KeyError(f"source variable {vc.source!r} not found in dataset")
         return ds[vc.source]
 
+    if vc.name == "sftlf":
+        vc.raw_variables = "landfrac"
+    elif vc.name == "areacella":
+        vc.raw_variables = "area"
+
     # 2) identity mapping from a single raw variable
     if (
         vc.raw_variables