| """NaN-in-mask check for a pair of initial-condition files. |
| |
| usage: python check_ic.py ATM_IC.nc OCEAN_IC.nc [OCEAN_FORCING.nc] |
| or: from check_ic import check; check(atm_ic, ocean_ic, ocean_forcing) # True on PASS |
| |
| The masks come from an ocean-forcing file of this package (default: |
| ../forcing_data/ocean-forcing-<first year found>.nc). They are the masks the checkpoints |
| apply. Rules: a 3-D ocean variable `<name>_<k>` uses mask_<k>; every other ocean variable |
| uses mask_2d. Passing means: |
| - no non-finite value inside the wet mask, and no finite value outside it (ocean); |
| - no non-finite value anywhere (atmosphere); |
| - sst NaN fraction == 1 - mean(mask_2d) (0.2958 for this grid). |
| Exit code 0 on pass, 1 on failure. |
| """ |
|
|
| import glob |
| import os |
| import sys |
|
|
| import numpy as np |
| import xarray as xr |
|
|
|
|
| def check(atm_path, ocn_path, mask_path=None): |
| """Check an IC pair; prints the details and returns True on PASS.""" |
| if mask_path is None: |
| here = os.path.dirname(os.path.abspath(__file__)) |
| found = sorted(glob.glob(os.path.join(here, "..", "forcing_data", "ocean-forcing-*.nc"))) |
| if not found: |
| raise FileNotFoundError("no ocean-forcing-*.nc found; pass the mask file") |
| mask_path = found[0] |
| masks = xr.open_dataset(mask_path) |
| atm = xr.open_dataset(atm_path, decode_times=False) |
| ocn = xr.open_dataset(ocn_path, decode_times=False) |
| ok = True |
| print(f"masks from {mask_path}") |
| print(f"times: atmosphere {atm.time.values}, ocean {ocn.time.values} ({ocn.time.attrs.get('units')})") |
| if not np.array_equal(atm.time.values, ocn.time.values): |
| print("FAIL: atmosphere and ocean IC times differ") |
| ok = False |
|
|
| def mask_for(name): |
| suffix = name.rsplit("_", 1)[-1] |
| if suffix.isdigit() and f"mask_{suffix}" in masks: |
| return masks[f"mask_{suffix}"].values > 0 |
| return masks["mask_2d"].values > 0 |
|
|
| bad_in, wet_out = {}, {} |
| for name in ocn.data_vars: |
| if "lat" not in ocn[name].dims: |
| continue |
| m = mask_for(name) |
| for i in range(ocn.sizes["time"]): |
| v = ocn[name].isel(time=i).values |
| b = int((m & ~np.isfinite(v)).sum()) |
| o = int((~m & np.isfinite(v)).sum()) |
| if b: |
| bad_in[f"{name}[t={i}]"] = b |
| if o: |
| wet_out[f"{name}[t={i}]"] = o |
| atm_bad = { |
| n: int((~np.isfinite(atm[n].values)).sum()) |
| for n in atm.data_vars |
| if "lat" in atm[n].dims and (~np.isfinite(atm[n].values)).any() |
| } |
| print(f"ocean: {len(ocn.data_vars)} variables; non-finite inside wet mask: {bad_in or 'NONE'}") |
| print(f"ocean: finite outside wet mask: {wet_out or 'NONE'}") |
| print(f"atmosphere: {len(atm.data_vars)} variables; non-finite anywhere: {atm_bad or 'NONE'}") |
| expect = 1 - float(masks["mask_2d"].mean()) |
| for i in range(ocn.sizes["time"]): |
| frac = float(np.isnan(ocn["sst"].isel(time=i).values).mean()) |
| good = abs(frac - expect) < 1e-6 |
| print(f"sst NaN fraction, IC {i}: {frac:.4f} (expected {expect:.4f}) {'ok' if good else 'FAIL'}") |
| ok &= good |
| ok &= not bad_in and not wet_out and not atm_bad |
| print("PASS" if ok else "FAIL") |
| return bool(ok) |
|
|
|
|
| def main(): |
| if len(sys.argv) < 3: |
| sys.exit(__doc__) |
| sys.exit(0 if check(*sys.argv[1:4]) else 1) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
