Fix OOM in review_exf_conditions by computing stats and histograms in time chunks

fluidnumericsJoe · fluidnumericsJoe · commit b00ed19bf04d · 2026-03-11T08:57:35.000-04:00
Previously _var_stats and _make_histograms materialised the full time series
  as float64/float32 arrays (~30-60 GB per variable), and the non-negative and
  dewpoint checks created full boolean arrays of similar size. All four operations
  now iterate over TIME_CHUNK (744) time steps at a time, keeping peak memory
  to a few hundred MB per variable.
diff --git a/spectre_utils/animate_obc_conditions.py b/spectre_utils/animate_obc_conditions.py
@@ -0,0 +1,128 @@
+"""Animate ocean boundary condition (OBC) binary files.
+
+For each (variable, boundary) pair produces one MP4:
+  - 3D variables (U, V, T, S): pcolormesh of depth × boundary-position over time
+  - Eta: line plot of boundary-position over time
+
+Output goes to <simulation_directory>/animations/ocean_bcs/.
+"""
+
+import os
+import yaml
+import numpy as np
+import xarray as xr
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from matplotlib.animation import FFMpegWriter
+
+from spectre_utils import common
+
+VARS = ["U", "V", "T", "S", "Eta"]
+BOUNDARIES = ["south", "north", "west", "east"]
+
+PHYSICAL_BOUNDS = {
+    "U":   (-3.0,  3.0),
+    "V":   (-3.0,  3.0),
+    "T":   (-2.0, 35.0),
+    "S":   ( 0.0, 42.0),
+    "Eta": (-3.0,  3.0),
+}
+
+UNITS = {"U": "m/s", "V": "m/s", "T": "°C", "S": "PSU", "Eta": "m"}
+
+CMAPS = {
+    "U": "RdBu_r", "V": "RdBu_r", "T": "plasma", "S": "viridis", "Eta": "RdBu_r"
+}
+
+
+def animate_boundary(da, var, boundary, out_path, fps=4, dpi=100):
+    """Animate one (var, boundary) DataArray as an MP4."""
+    pos_dim = [d for d in da.dims if d not in ("time", "depth")][0]
+    pos     = da.coords[pos_dim].values
+    times   = da.coords["time"].values
+    nt      = da.sizes["time"]
+    xlabel  = "Longitude" if pos_dim == "lon" else "Latitude"
+
+    lo, hi  = PHYSICAL_BOUNDS.get(var, (None, None))
+    cmap    = CMAPS.get(var, "viridis")
+    unit    = UNITS.get(var, "")
+
+    is_3d = da.ndim == 3
+
+    if is_3d:
+        fig, ax = plt.subplots(figsize=(10, 5), dpi=dpi)
+        depth  = da.coords["depth"].values
+        frame0 = da.isel(time=0).values.astype(np.float32)
+        pcm = ax.pcolormesh(pos, depth, frame0, shading="auto",
+                            vmin=lo, vmax=hi, cmap=cmap)
+        plt.colorbar(pcm, ax=ax, label=f"{var} [{unit}]")
+        ax.invert_yaxis()
+        ax.set_xlabel(xlabel)
+        ax.set_ylabel("Depth (m)")
+
+        def update(i):
+            pcm.set_array(da.isel(time=i).values.astype(np.float32))
+
+    else:
+        fig, ax = plt.subplots(figsize=(10, 3), dpi=dpi)
+        frame0 = da.isel(time=0).values.astype(np.float32)
+        (line,) = ax.plot(pos, frame0, color="steelblue", lw=0.8)
+        if lo is not None:
+            ax.set_ylim(lo - abs(lo) * 0.1, hi + abs(hi) * 0.1)
+        ax.set_xlabel(xlabel)
+        ax.set_ylabel(f"{var} [{unit}]")
+        ax.grid(True, alpha=0.3)
+
+        def update(i):
+            line.set_ydata(da.isel(time=i).values.astype(np.float32))
+
+    title = ax.set_title(
+        f"{var} — {boundary}  |  {np.datetime_as_string(times[0], unit='D')}"
+    )
+
+    writer = FFMpegWriter(fps=fps, bitrate=-1)
+    with writer.saving(fig, out_path, dpi):
+        writer.grab_frame()
+        for i in range(1, nt):
+            update(i)
+            title.set_text(
+                f"{var} — {boundary}  |  {np.datetime_as_string(times[i], unit='D')}"
+            )
+            writer.grab_frame()
+
+    plt.close(fig)
+    print(f"  Saved {out_path}")
+
+
+def main():
+    args = common.cli()
+
+    with open(args.config_file, "r") as f:
+        config = yaml.safe_load(f)
+
+    simulation_directory = config["simulation_directory"]
+    working_directory    = config["working_directory"]
+    simulation_input_dir = os.path.join(simulation_directory, "input")
+
+    animations_dir = os.path.join(simulation_directory, "animations", "ocean_bcs")
+    os.makedirs(animations_dir, exist_ok=True)
+
+    print("Loading OBC binary files...")
+    data = common.load_obc_binaries(simulation_input_dir, working_directory, config)
+
+    for var in VARS:
+        for bnd in BOUNDARIES:
+            key = (var, bnd)
+            if key not in data:
+                print(f"  Skipping {var}.{bnd} (not found)")
+                continue
+            out_path = os.path.join(animations_dir, f"{var}_{bnd}.mp4")
+            print(f"Animating {var}.{bnd} -> {out_path}")
+            animate_boundary(data[key], var, bnd, out_path, fps=4)
+
+    print("Done.")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/spectre_utils/common.py b/spectre_utils/common.py
@@ -116,6 +116,103 @@ def load_exf_binaries(simulation_input_dir, var_names, working_directory, prefix
     return xr.Dataset(data_vars)
 
 
+def load_obc_binaries(simulation_input_dir, working_directory, config):
+    """Load ocean boundary condition binary files into a dict of labelled DataArrays.
+
+    Returns a mapping ``{(var, boundary): xr.DataArray}`` where:
+      - ``var``      ∈ {"U", "V", "T", "S", "Eta"}
+      - ``boundary`` ∈ {"south", "north", "west", "east"}
+
+    Coordinates (time, depth, and boundary-parallel position as longitude or
+    latitude) are extracted from the raw glorys12 NetCDF files.  The binary
+    data are memory-mapped so only accessed portions are read into RAM.
+    """
+    import xarray as xr
+    import numpy as np
+    import os
+    import glob as _glob
+
+    i0 = config.get("domain", {}).get("longitude", {}).get("start", 2)
+    i1 = config.get("domain", {}).get("longitude", {}).get("end", -2)
+    j0 = config.get("domain", {}).get("latitude",  {}).get("start", 1)
+    j1 = config.get("domain", {}).get("latitude",  {}).get("end",  -1)
+    prefix = config.get("ocean", {}).get("prefix", "glorysv12")
+
+    t_files = sorted(_glob.glob(f"{working_directory}/{prefix}_T_glorys12_raw.*.nc"))
+    if not t_files:
+        raise FileNotFoundError(f"No glorys12 T files found in {working_directory}")
+
+    # One file for static coords (depth, nav_lon, nav_lat, grid size).
+    one_ds = xr.open_dataset(t_files[0])
+    nz = one_ds.sizes["deptht"]
+    ny = one_ds.sizes["y"]
+    nx = one_ds.sizes["x"]
+    depth   = one_ds["deptht"].values
+    nav_lon = one_ds["nav_lon"].values   # (ny, nx)
+    nav_lat = one_ds["nav_lat"].values   # (ny, nx)
+    one_ds.close()
+
+    # Full time axis from all T files (lazy — no data read).
+    all_ds = xr.open_mfdataset(t_files, combine="by_coords")
+    times = all_ds["time_counter"].values
+    all_ds.close()
+    nt = len(times)
+
+    # Resolve negative indices to absolute positions.
+    i1_abs = nx + i1 if i1 < 0 else i1   # = nx - 2
+    j1_abs = ny + j1 if j1 < 0 else j1   # = ny - 1
+
+    # 1-D geographic coordinates along each boundary.
+    # South/north: position varies in x (longitude-like).
+    lon_sn   = nav_lon[j0, i0:i1_abs]       # V, T, S, Eta
+    lon_sn_u = nav_lon[j0, i0:i1_abs - 1]   # U (narrower by 1; C-grid stagger)
+    lat_sn   = nav_lat[j0, i0:i1_abs]
+    lat_sn_u = nav_lat[j0, i0:i1_abs - 1]
+    # West/east: position varies in y (latitude-like).
+    lon_we = nav_lon[j0:j1_abs, i0]
+    lat_we = nav_lat[j0:j1_abs, i0]
+
+    # (shape, pos_dim_name, pos_coord) keyed by (var, boundary).
+    _specs = {}
+    for bnd in ("south", "north"):
+        _specs[("U",   bnd)] = ((nt, nz, len(lon_sn_u)), "lon", lon_sn_u)
+        _specs[("V",   bnd)] = ((nt, nz, len(lon_sn)),   "lon", lon_sn)
+        _specs[("T",   bnd)] = ((nt, nz, len(lon_sn)),   "lon", lon_sn)
+        _specs[("S",   bnd)] = ((nt, nz, len(lon_sn)),   "lon", lon_sn)
+        _specs[("Eta", bnd)] = ((nt,     len(lon_sn)),   "lon", lon_sn)
+    for bnd in ("west", "east"):
+        _specs[("U",   bnd)] = ((nt, nz, len(lat_we)), "lat", lat_we)
+        _specs[("V",   bnd)] = ((nt, nz, len(lat_we)), "lat", lat_we)
+        _specs[("T",   bnd)] = ((nt, nz, len(lat_we)), "lat", lat_we)
+        _specs[("S",   bnd)] = ((nt, nz, len(lat_we)), "lat", lat_we)
+        _specs[("Eta", bnd)] = ((nt,     len(lat_we)), "lat", lat_we)
+
+    result = {}
+    for (var, bnd), (shape, pos_name, pos_coord) in _specs.items():
+        bin_path = os.path.join(simulation_input_dir, f"{var}.{bnd}.bin")
+        if not os.path.exists(bin_path):
+            print(f"Warning: binary file not found, skipping: {bin_path}", file=sys.stderr)
+            continue
+        arr = np.memmap(bin_path, dtype=">f4", mode="r", shape=shape)
+        if len(shape) == 3:   # time × depth × position
+            da = xr.DataArray(
+                arr,
+                dims=["time", "depth", pos_name],
+                coords={"time": times, "depth": depth, pos_name: pos_coord},
+                name=f"{var}_{bnd}",
+            )
+        else:                  # time × position (Eta)
+            da = xr.DataArray(
+                arr,
+                dims=["time", pos_name],
+                coords={"time": times, pos_name: pos_coord},
+                name=f"{var}_{bnd}",
+            )
+        result[(var, bnd)] = da
+
+    return result
+
+
 def load_atm_dataset(working_directory, prefix, years, atm_vars, t1, t2):
     """Load ERA5 atmospheric variables per MITgcm name, rename, and apply optional scale factors.
 
diff --git a/spectre_utils/review_exf_conditions.py b/spectre_utils/review_exf_conditions.py
@@ -46,24 +46,44 @@
 # Variables that should be non-negative after any scale factor
 NON_NEGATIVE = {"swdown", "lwdown", "precip", "runoff", "aqh"}
 
+# Number of time steps processed at once for memory-intensive operations.
+# ERA5 is hourly so 744 ≈ one month.
+TIME_CHUNK = 744
+
 
 # ---------------------------------------------------------------------------
 # Helpers
 # ---------------------------------------------------------------------------
 
 def _var_stats(da: xr.DataArray) -> dict:
-    """Compute global min/mean/max and count of non-finite values (lazy-safe)."""
-    arr = da.values.astype(np.float64)
-    n_bad = int(np.sum(~np.isfinite(arr)))
-    valid = arr[np.isfinite(arr)]
-    if valid.size == 0:
-        return dict(min=np.nan, mean=np.nan, max=np.nan, n_bad=n_bad, n=arr.size)
+    """Compute global min/mean/max and count of non-finite values in time chunks."""
+    nt = da.sizes["valid_time"]
+    n_bad = 0
+    n_total = 0
+    gmin = np.inf
+    gmax = -np.inf
+    gsum = 0.0
+
+    for i in range(0, nt, TIME_CHUNK):
+        chunk = da.isel(valid_time=slice(i, i + TIME_CHUNK)).values.astype(np.float64)
+        n_total += chunk.size
+        finite = np.isfinite(chunk)
+        n_bad += int((~finite).sum())
+        valid = chunk[finite]
+        if valid.size > 0:
+            gmin = min(gmin, float(valid.min()))
+            gmax = max(gmax, float(valid.max()))
+            gsum += float(valid.sum())
+
+    n_valid = n_total - n_bad
+    if n_valid == 0:
+        return dict(min=np.nan, mean=np.nan, max=np.nan, n_bad=n_bad, n=n_total)
     return dict(
-        min=float(valid.min()),
-        mean=float(valid.mean()),
-        max=float(valid.max()),
+        min=gmin if np.isfinite(gmin) else np.nan,
+        mean=gsum / n_valid,
+        max=gmax if np.isfinite(gmax) else np.nan,
         n_bad=n_bad,
-        n=int(arr.size),
+        n=n_total,
     )
 
 
@@ -166,9 +186,32 @@ def _make_histograms(ds: xr.Dataset, var_names: list[str], out_path: str) -> Non
 
     for idx, name in enumerate(present):
         ax = axes[idx // ncols][idx % ncols]
-        vals = ds[name].values.ravel()
-        vals = vals[np.isfinite(vals)]
-        ax.hist(vals, bins=120, color="steelblue", edgecolor="none", density=True)
+
+        # Determine bin range from known physical bounds or the first time chunk.
+        if name in PHYSICAL_BOUNDS:
+            lo, hi = PHYSICAL_BOUNDS[name]
+        else:
+            first = ds[name].isel(valid_time=slice(0, TIME_CHUNK)).values.ravel()
+            first = first[np.isfinite(first)]
+            lo = float(first.min()) if first.size > 0 else 0.0
+            hi = float(first.max()) if first.size > 0 else 1.0
+
+        bins = np.linspace(lo, hi, 121)
+        counts = np.zeros(120, dtype=np.float64)
+
+        nt = ds[name].sizes["valid_time"]
+        for i in range(0, nt, TIME_CHUNK):
+            chunk = ds[name].isel(valid_time=slice(i, i + TIME_CHUNK)).values.ravel()
+            chunk = chunk[np.isfinite(chunk)]
+            c, _ = np.histogram(chunk, bins=bins)
+            counts += c
+
+        bin_centers = (bins[:-1] + bins[1:]) / 2
+        bin_width = bins[1] - bins[0]
+        total = counts.sum()
+        density = counts / (total * bin_width) if total > 0 else counts
+        ax.bar(bin_centers, density, width=bin_width, color="steelblue", edgecolor="none")
+
         var_units = ds[name].attrs.get("units", "")
         ax.set_xlabel(f"{name} [{var_units}]" if var_units else name, fontsize=8)
         ax.set_ylabel("Density", fontsize=8)
@@ -264,7 +307,8 @@ def main():
             stat_rows.append((name, "—", "—", "—", "—", "—"))
             continue
 
-        st = _var_stats(ds[name])
+        da = ds[name]
+        st = _var_stats(da)
         checks: list[tuple[bool, str]] = []
 
         # 1. NaN / Inf
@@ -290,7 +334,11 @@ def main():
 
         # 3. Non-negative where required
         if name in NON_NEGATIVE:
-            n_neg = int((ds[name] < 0).sum().compute().item())
+            n_neg = 0
+            _nt = da.sizes["valid_time"]
+            for _i in range(0, _nt, TIME_CHUNK):
+                _chunk = da.isel(valid_time=slice(_i, _i + TIME_CHUNK)).values
+                n_neg += int((_chunk < 0).sum())
             if n_neg == 0:
                 checks.append((True, "All values ≥ 0"))
             else:
@@ -331,8 +379,14 @@ def main():
 
     # Dewpoint ≤ air temperature
     if "d2m" in ds and "atemp" in ds:
-        n_viol = int((ds["d2m"] > ds["atemp"]).sum().compute().item())
-        n_total = int(ds["d2m"].size)
+        n_viol = 0
+        n_total = 0
+        nt = ds["d2m"].sizes["valid_time"]
+        for _i in range(0, nt, TIME_CHUNK):
+            _d2m   = ds["d2m"].isel(valid_time=slice(_i, _i + TIME_CHUNK)).values
+            _atemp = ds["atemp"].isel(valid_time=slice(_i, _i + TIME_CHUNK)).values
+            n_viol += int((_d2m > _atemp).sum())
+            n_total += _d2m.size
         ok = n_viol == 0
         msg = (
             "Dewpoint ≤ air temperature everywhere"
diff --git a/spectre_utils/review_obc_conditions.py b/spectre_utils/review_obc_conditions.py
