From 469909e1678009c839dacf4d6410f52ea56ad9dc Mon Sep 17 00:00:00 2001 From: SorooshMani-NOAA Date: Fri, 16 Oct 2026 07:29:43 -0400 Subject: [PATCH 1/4] Make ADT datum offset a config based parameter --- nos_utils/config.py | 11 +++++++++-- nos_utils/forcing/adt.py | 22 ++++++++++------------ 2 files changed, 19 insertions(+), 14 deletions(-) diff --git a/nos_utils/config.py b/nos_utils/config.py index fbc128a..f1eab09 100644 --- a/nos_utils/config.py +++ b/nos_utils/config.py @@ -231,6 +231,7 @@ class ForcingConfig: obc_sal_outside: float = 33.0 # ADT satellite SSH blending (STOFS-3D-ATL uses CMEMS ADT to correct boundary SSH) adt_enabled: bool = False + adt_datum_offset: float = -0.45 # Ops ESMF map for the ADT regrid (yaml ocean.adt.weight_file); a bare name resolves against # FIXofs. Unset or missing falls back to the bilinear regrid. MJ (10/02/26) adt_weight_file: Optional[Path] = None @@ -577,6 +578,7 @@ def for_stofs_3d_atl(cls, pdy: str, cyc: int, **overrides) -> "ForcingConfig": # carry only the 778 elevation nodes SCHISM expects. obc_elev_segments=[0, 1], adt_enabled=True, + adt_datum_offset=-0.45, adt_weight_file=Path("stofs_3d_atl_ufs.adt_weight.nc"), obc_ssh_hold_first_record=True, obc_ops_timeline=True, @@ -641,6 +643,7 @@ def for_stofs_3d_atl_ufs(cls, pdy: str, cyc: int, **overrides) -> "ForcingConfig # carry only the 778 elevation nodes SCHISM expects. obc_elev_segments=[0, 1], adt_enabled=True, + adt_datum_offset=-0.45, adt_weight_file=Path("stofs_3d_atl_ufs.adt_weight.nc"), obc_ssh_hold_first_record=True, obc_ops_timeline=True, @@ -701,7 +704,8 @@ def for_stofs_3d_pac(cls, pdy: str, cyc: int, **overrides) -> "ForcingConfig": # Segments 2-8 are flow-only river OBCs (excluded from *.th.nc). obc_elev_segments=[0, 1, 9], obc_ssh_offset=0.0, - adt_enabled=False, + adt_enabled=True, + adt_datum_offset=-0.4715517, nudging_enabled=True, nudging_timescale_seconds=86400.0, nwm_product="medium_range_mem1", @@ -730,7 +734,8 @@ def for_stofs_3d_pac_ufs(cls, pdy: str, cyc: int, **overrides) -> "ForcingConfig dynamic_adjust_enabled=False, obc_elev_segments=[0, 1, 9], obc_ssh_offset=0.0, - adt_enabled=False, + adt_enabled=True, + adt_datum_offset=-0.4715517, nudging_enabled=True, nudging_timescale_seconds=86400.0, nwm_product="medium_range_mem1", @@ -956,6 +961,7 @@ def from_yaml(cls, yaml_path, **overrides) -> "ForcingConfig": if _adt_wt is None and str(_sys.get("name", "")).startswith("stofs_3d_atl"): _adt_wt = "stofs_3d_atl_ufs.adt_weight.nc" _adt_wt = Path(_adt_wt) if _adt_wt else None + _adt_off = float(adt.get("datum_offset", 0.0)) if isinstance(adt, dict) else 0.0 # Ops river inputs (static FIX source_sink.in / msource.th / vsink.th, ops source # order): STOFS-3D-ATL by name (the coupled executable now has PREC_EVAP), yaml @@ -1011,6 +1017,7 @@ def from_yaml(cls, yaml_path, **overrides) -> "ForcingConfig": datm_hrrr_weight=_hrrr_weight, hrrr_small_grib=_hrrr_sg, adt_enabled=adt.get("enabled", False) if isinstance(adt, dict) else False, + adt_datum_offset=_adt_off, adt_weight_file=_adt_wt, nwm_product=nwm_product, river_ops_static=_river_ops, diff --git a/nos_utils/forcing/adt.py b/nos_utils/forcing/adt.py index 8b5b47a..de5e23d 100644 --- a/nos_utils/forcing/adt.py +++ b/nos_utils/forcing/adt.py @@ -5,12 +5,12 @@ boundary condition accuracy for STOFS-3D-ATL. The core formula (ops stofs_3d_atl_create_obc_3d_th_non_adjust.sh, SSH_1.nc block): - SSH_final = SSH_rtofs - SSH_rtofs(t=0) + (ADT - 0.45) + SSH_final = SSH_rtofs - SSH_rtofs(t=0) + (ADT + adt_datum_offset) ADT is the mean of the valid-day and previous-day CMEMS files (ncra), mapped onto the SSH_1 grid with the ops ESMF map (ncremap with stofs_3d_atl_adt_weight.nc, which is a nearest-source map with S=1 despite its "bilinear" global attribute). -The -0.45 m datum offset (stofs_3d_atl_adt_cvtz.nco) is applied per day and rounded +The datum offset (stofs_3d_atl_adt_cvtz.nco) is applied per day and rounded to float32 before the mean, as ops does. Without the map, a bilinear regrid is used and the result is not ops-exact. @@ -273,7 +273,7 @@ def _apply_esmf_map(S, row, col, n_b, src): return np.where(tally > 0, num, np.nan) def _regrid_esmf(self, adt_files, ssh_path: Path): - """Ops ADT step with the ESMF map: ROI subset, ncremap, float32(adt-0.45), two-day mean. + """Ops ADT step with the ESMF map: ROI subset, ncremap, float32(adt+datum_offset), two-day mean. Each day is rounded to float32 after the offset and the valid days are averaged (float64 accumulation, float32 result), as ncap2 and ncra do. Returns the float32 @@ -320,7 +320,7 @@ def dlon_diff(a, b): or np.abs(yca - LA.ravel()).max() > ADT_COORD_TOL): raise ValueError("ADT ROI coordinates do not match map xc_a/yc_a") out = self._apply_esmf_map(S, row, col, n_b, a.ravel()) - days.append((out - 0.45).astype(np.float32) if self.ops_numerics else out - 0.45) + days.append(out.astype(np.float32) if self.ops_numerics else out) if self.ops_numerics: stack = np.stack(days) ok = np.isfinite(stack) @@ -343,7 +343,7 @@ def dlon_diff(a, b): return field def _write_adt_on_rtofs(self, adt_dst, dst_lon, dst_lat, out: Path): - """adt_aft_cvtz_cln.nc analogue: float32 surf_el(time=1, ylat, xlon) in m after -0.45. + """adt_aft_cvtz_cln.nc analogue: float32 surf_el(time=1, ylat, xlon) in m after datum offset. Missing cells carry the ops fill (the CMEMS int fill held as float32). """ @@ -421,7 +421,7 @@ def _load_previous(self, ssh_path: Path): return None def _read_adt(self, adt_path: Path, bounds=None): - """Read and subset ADT to bounds, offset by -0.45 m. + """Read and subset ADT to bounds, offset by the user specified amount. bounds is (lon_min, lon_max, lat_min, lat_max) in -180/180; default is the ops ROI. Returns (adt[lat, lon], lons, lats) or None. Fill values become NaN. @@ -445,9 +445,7 @@ def _read_adt(self, adt_path: Path, bounds=None): return None adt_var = "adt" if "adt" in ds.variables else "surf_el" - # 0.45 from NCO file - # TODO: Should be different for different systems, e.g. PAC - adt_data = ds.variables[adt_var][...] - 0.45 + adt_data = ds.variables[adt_var][...] sl_y = slice(lat_idx[0], lat_idx[-1] + 1) sl_x = slice(lon_idx[0], lon_idx[-1] + 1) if adt_data.ndim == 3: @@ -564,8 +562,8 @@ def _apply_adt_blend(self, ssh_path: Path, adt_field, work_dir: Path, adt_dst=None) -> Optional[Path]: """Apply the ops ADT blend to SSH_1.nc. - Formula: SSH_final(t) = SSH_rtofs(t) - SSH_rtofs(t=0) + ADT - (ADT already carries the -0.45 offset), in float32 as ncap2. Where either + Formula: SSH_final(t) = SSH_rtofs(t) - SSH_rtofs(t=0) + ADT + datum_offset + (ADT datum offset, e.g. -0.45 is applied here), in float32 as ncap2. Where either side is invalid the point is left as the -30000 fill (dry in gen_3Dth). """ try: @@ -574,7 +572,7 @@ def _apply_adt_blend(self, ssh_path: Path, adt_field, work_dir: Path, shutil.copy2(ssh_path, output) ds = Dataset(str(output), "r+") - ssh = ds.variables["ssh"] + ssh = ds.variables["ssh"][...] + (self.config.adt_datum_offset or 0) nt, ny, nx = ssh.shape dst_lon = np.array(ds.variables["xlon"][:], dtype=np.float64) From c2d20d8d060d208d776a3277a587452e78b0d3ef Mon Sep 17 00:00:00 2001 From: SorooshMani-NOAA Date: Fri, 16 Oct 2026 09:04:25 -0400 Subject: [PATCH 2/4] Process RTOFS slice strides --- nos_utils/config.py | 6 +++--- nos_utils/forcing/nudging.py | 3 ++- nos_utils/forcing/rtofs.py | 15 +++++++++------ 3 files changed, 14 insertions(+), 10 deletions(-) diff --git a/nos_utils/config.py b/nos_utils/config.py index f1eab09..6b683ed 100644 --- a/nos_utils/config.py +++ b/nos_utils/config.py @@ -172,11 +172,11 @@ class ForcingConfig: # --- OBC (ocean boundary) settings --- # RTOFS ROI indices for 2D (ssh) extraction - obc_roi_2d: Optional[dict] = None # {x1, x2, y1, y2} + obc_roi_2d: Optional[dict] = None # {x1, x2, y1, y2, x_stride, y_stride} # RTOFS ROI indices for 3D (T,S,U,V) extraction - obc_roi_3d: Optional[dict] = None # {x1, x2, y1, y2} + obc_roi_3d: Optional[dict] = None # {x1, x2, y1, y2, x_stride, y_stride} # RTOFS ROI indices for nudging (slightly larger domain than OBC) - nudge_roi_3d: Optional[dict] = None # {x1, x2, y1, y2} + nudge_roi_3d: Optional[dict] = None # {x1, x2, y1, y2, x_stride, y_stride} # Which regional RTOFS 3dz tile to read for 3D T/S/U/V (OBC and nudging # share this discovery step). Ops stages "alaska", "US_east" and # "US_west" tiles side by side for every valid time with no region in diff --git a/nos_utils/forcing/nudging.py b/nos_utils/forcing/nudging.py index bc81078..6f7dd7b 100644 --- a/nos_utils/forcing/nudging.py +++ b/nos_utils/forcing/nudging.py @@ -323,7 +323,8 @@ def read_raw(path): # ncap2 honours only _FillValue; netCDF4's auto-mask would also drop values outside # valid_range. MJ (10/02/26) ds.set_auto_mask(False) - ys, xs = slice(roi["y1"], roi["y2"] + 1), slice(roi["x1"], roi["x2"] + 1) + ys = slice(roi["y1"], roi["y2"] + 1, roi.get("y_stride", 1)) + xs = slice(roi["x1"], roi["x2"] + 1, roi.get("x_stride", 1)) out = [] for nme in ("temperature", "salinity"): v = ds.variables[nme] diff --git a/nos_utils/forcing/rtofs.py b/nos_utils/forcing/rtofs.py index 2d28ed7..49c08dd 100644 --- a/nos_utils/forcing/rtofs.py +++ b/nos_utils/forcing/rtofs.py @@ -872,20 +872,23 @@ def _stofs_work_dir(self, keep: bool) -> Path: return work def _stofs_subset_roi(self, ds, roi: dict, variables: List[str]) -> dict: - """Subset RTOFS NetCDF by ROI indices (replaces NCO ncks -d X,x1,x2 -d Y,y1,y2).""" - x1, x2 = roi["x1"], roi["x2"] - y1, y2 = roi["y1"], roi["y2"] + """ + Subset RTOFS NetCDF by ROI indices & stride + (replaces NCO ncks -d X,x1,x2,i -d Y,y1,y2,j). + """ + x1, x2, x_stride = roi["x1"], roi["x2"], roi.get("x_stride", 1) + y1, y2, y_stride = roi["y1"], roi["y2"], roi.get("y_stride", 1) result = {} for var in variables: if var in ds.variables: data = ds.variables[var] ndim = len(data.dimensions) if ndim == 3: # (time, Y, X) - result[var] = data[:, y1:y2+1, x1:x2+1] + result[var] = data[:, y1:y2+1:y_stride, x1:x2+1:x_stride] elif ndim == 4: # (time, depth, Y, X) - result[var] = data[:, :, y1:y2+1, x1:x2+1] + result[var] = data[:, :, y1:y2+1:y_stride, x1:x2+1:x_stride] elif ndim == 2: # (Y, X) - result[var] = data[y1:y2+1, x1:x2+1] + result[var] = data[y1:y2+1:y_stride, x1:x2+1:x_stride] else: result[var] = data[:] return result From ef86e5d5fef48251f3f47d6078a6d447a45c3528 Mon Sep 17 00:00:00 2001 From: SorooshMani-NOAA Date: Thu, 22 Oct 2026 08:02:53 -0400 Subject: [PATCH 3/4] feat(forcing): add STOFS-3D-PAC ops features for RTOFS, ADT, and nudging This commit introduces the necessary structural changes to support the Pacific domain's operational processing requirements attempting to avoid breaking Atlantic parity. * nos_utils/config.py: - Decouple time steps by adding explicit `obc_dt` and `nudging_dt`. - Add config flags for `obc_time_smoothing` and `nudging_time_smoothing` - Add `rtofs_daily_3z` flag to support variable-split daily 3D files. - Expose `obc_min_timesteps` to be configurable via YAML. * nos_utils/coords.py: - Add `shift_grid_0360` utility to cleanly stitch and sort global [-180, 180] grids across the antimeridian for Pacific subsets. * nos_utils/forcing/adt.py: - Remove hardcoded Atlantic domain bounding boxes (`ADT_LON_MIN`, etc.). - Dynamically read destination bounds from `self.config.domain`. - Apply `shift_grid_0360` during `_subset_ops_roi` and `_read_adt` to seamlessly support 0-360 meshes. * nos_utils/forcing/rtofs.py & nos_utils/forcing/nudging.py: - Replace hardcoded 21600.0s time steps with `config.obc_dt` and `config.nudging_dt`. - Add `groupby` logic to process variable-split daily files (`t3z`, `s3z`) simultaneously so the time axis is not corrupted. - Implement a configurable 2-step temporal moving average `(data[:-1] + data[1:]) / 2` for `elev2D`, `TEM_3D`, `SAL_3D`, and nudging arrays to match ops Pacific time smoothing. - Use dynamically configured `obc_min_timesteps` for the phase plan and QC checks. --- nos_utils/config.py | 41 ++++++++++++++++-- nos_utils/coords.py | 21 ++++++++++ nos_utils/forcing/adt.py | 49 ++++++++++++---------- nos_utils/forcing/nudging.py | 56 ++++++++++++++++--------- nos_utils/forcing/rtofs.py | 80 ++++++++++++++++++++++-------------- 5 files changed, 173 insertions(+), 74 deletions(-) diff --git a/nos_utils/config.py b/nos_utils/config.py index 6b683ed..33add6b 100644 --- a/nos_utils/config.py +++ b/nos_utils/config.py @@ -192,6 +192,8 @@ class ForcingConfig: # for every system, not just Alaska, so correctness stops depending on # alphabetical luck. rtofs_3d_region: Optional[str] = None + # Whether to use global daily 3D files (split by variable) instead of regional 6-hourly. + rtofs_daily_3z: bool = False # SSH offset applied to boundary elevation (meters) # Geoid-to-MSL datum offset (meters). OFS-specific — verify from Fortran source. # SECOFS: 1.25 (confirmed). STOFS-3D-ATL: 0.04. Other OFS: 0.0 until verified. @@ -226,6 +228,10 @@ class ForcingConfig: # sliced per phase, non-zero uv3D. Off keeps the previous-cycle search, 3-hourly T/S and zero # uv3D that SECOFS, Pacific and Alaska use. MJ (10/02/26) obc_ops_timeline: bool = False + # Temporal resolution of the output boundary files (seconds). + obc_dt: float = 21600.0 + # Apply 2-step temporal moving average to OBC files (STOFS-3D-PAC ops parity). + obc_time_smoothing: bool = False # T,S for boundary nodes outside the RTOFS grid (stofs_3d_atl_obc_3dth_nc.in). MJ (10/01/26) obc_tem_outside: float = 20.0 obc_sal_outside: float = 33.0 @@ -238,6 +244,10 @@ class ForcingConfig: # Nudging enabled and timescale nudging_enabled: bool = False nudging_timescale_seconds: float = 86400.0 + # Temporal resolution of the nudging files (seconds). + nudging_dt: float = 21600.0 + # Apply 2-step temporal moving average to nudging fields. + nudging_time_smoothing: bool = False # Number of vertical levels n_levels: int = 51 # SCHISM model timestep (seconds) used to write the elev2D.th.nc @@ -709,7 +719,15 @@ def for_stofs_3d_pac(cls, pdy: str, cyc: int, **overrides) -> "ForcingConfig": nudging_enabled=True, nudging_timescale_seconds=86400.0, nwm_product="medium_range_mem1", - obc_min_timesteps=21, + obc_min_timesteps=4, + obc_dt=86400.0, + nudging_dt=86400.0, + obc_time_smoothing=True, + nudging_time_smoothing=True, + rtofs_daily_3z=True, + obc_roi_2d={"x1": 222, "x2": 2703, "y1": 1051, "y2": 2641, "x_stride": 2, "y_stride": 2}, + obc_roi_3d={"x1": 222, "x2": 2703, "y1": 1051, "y2": 2641, "x_stride": 2, "y_stride": 2}, + nudge_roi_3d={"x1": 222, "x2": 2703, "y1": 1051, "y2": 2641, "x_stride": 2, "y_stride": 2}, ) defaults.update(overrides) return cls(**defaults) @@ -739,7 +757,15 @@ def for_stofs_3d_pac_ufs(cls, pdy: str, cyc: int, **overrides) -> "ForcingConfig nudging_enabled=True, nudging_timescale_seconds=86400.0, nwm_product="medium_range_mem1", - obc_min_timesteps=21, + obc_min_timesteps=4, + obc_dt=86400.0, + nudging_dt=86400.0, + obc_time_smoothing=True, + nudging_time_smoothing=True, + rtofs_daily_3z=True, + obc_roi_2d={"x1": 222, "x2": 2703, "y1": 1051, "y2": 2641, "x_stride": 2, "y_stride": 2}, + obc_roi_3d={"x1": 222, "x2": 2703, "y1": 1051, "y2": 2641, "x_stride": 2, "y_stride": 2}, + nudge_roi_3d={"x1": 222, "x2": 2703, "y1": 1051, "y2": 2641, "x_stride": 2, "y_stride": 2}, # DATM grid — STOFS3D_PAC preset, verified from weight_gfs.nc # nx = round((290.0-92.5)/0.025)+1 = 7901; ny = round((67.0-(-33.5))/0.025)+1 = 4021 datm_lon_min=92.5, datm_lon_max=290.0, @@ -1039,6 +1065,12 @@ def from_yaml(cls, yaml_path, **overrides) -> "ForcingConfig": kwargs["nudge_roi_3d"] = {k: int(v) for k, v in nudge_roi.items()} if rtofs_3d_region: kwargs["rtofs_3d_region"] = str(rtofs_3d_region) + + kwargs["rtofs_daily_3z"] = _strict_bool(ocean.get("rtofs_daily_3z"), "ocean.rtofs_daily_3z", False) + kwargs["obc_dt"] = float(obc.get("dt", 21600.0)) + kwargs["obc_time_smoothing"] = _strict_bool(obc.get("time_smoothing"), "obc.time_smoothing", False) + kwargs["nudging_dt"] = float(nudge.get("dt", kwargs["obc_dt"])) + kwargs["nudging_time_smoothing"] = _strict_bool(nudge.get("time_smoothing"), "nudging.time_smoothing", False) # NWM target counts if isinstance(river, dict): @@ -1131,7 +1163,10 @@ def from_yaml(cls, yaml_path, **overrides) -> "ForcingConfig": # Non-STOFS OFSes have shorter OBC windows (typically ~10 records for # a 54h run) and would be wrongly flagged if this ran against them; # keep the dataclass default of 0 (disabled) for non-STOFS. - if is_stofs_like: + yaml_min_ts = obc.get("min_timesteps") + if yaml_min_ts is not None: + kwargs["obc_min_timesteps"] = int(yaml_min_ts) + elif is_stofs_like: kwargs["obc_min_timesteps"] = 21 # Declarative prep extras override. When a root-level `prep.extras` diff --git a/nos_utils/coords.py b/nos_utils/coords.py index 7f72afa..20943cf 100644 --- a/nos_utils/coords.py +++ b/nos_utils/coords.py @@ -71,6 +71,27 @@ def normalize_lon( return lon +def shift_grid_0360( + lon: np.ndarray, data: "Optional[np.ndarray]" = None, lon_axis: int = -1 +) -> "Tuple[np.ndarray, Optional[np.ndarray]]": + """Shift a [-180, 180] grid and optional data array to [0, 360]. + + Useful for stitching global [-180, 180] fields (like ADT) across the + antimeridian to match Pacific domain bounds smoothly. + """ + lon = np.asarray(lon) + lon_360 = np.where(lon < 0.0, lon + 360.0, lon) + + sort_idx = np.argsort(lon_360) + lon_out = lon_360[sort_idx] + + if data is not None: + data_out = np.take(data, sort_idx, axis=lon_axis) + return lon_out, data_out + + return lon_out, None + + def lon_convention(config: "object") -> str: """Return the longitude convention used by a ``ForcingConfig``. diff --git a/nos_utils/forcing/adt.py b/nos_utils/forcing/adt.py index de5e23d..d422365 100644 --- a/nos_utils/forcing/adt.py +++ b/nos_utils/forcing/adt.py @@ -50,10 +50,6 @@ # Ops ADT subset ROI, used only when the SSH_1 destination grid cannot be read. # MJ (10/01/26) -ADT_LON_MIN = -62.5 -ADT_LON_MAX = -51.5 -ADT_LAT_MIN = 7.0 -ADT_LAT_MAX = 54.0 ADT_PAD = 1.0 # Ops holds ADT-missing cells as the CMEMS int fill carried through float32. MJ (10/02/26) @@ -147,10 +143,9 @@ def blend_ssh(self, ssh_path: Path, work_dir: Path) -> Optional[Path]: log.warning(f"ADT blending failed: {e}") return None - @staticmethod - def _dest_bounds(ssh_path: Path): + def _dest_bounds(self, ssh_path: Path): """ADT subset (lon_min, lon_max, lat_min, lat_max): SSH_1 grid extent padded - by ADT_PAD (same -180/180 convention as ADT), else the ops ROI. MJ (10/01/26) + by ADT_PAD, else the config domain bounds. MJ (10/01/26) """ try: with Dataset(str(ssh_path)) as ds: @@ -159,8 +154,9 @@ def _dest_bounds(ssh_path: Path): return (float(np.nanmin(lon)) - ADT_PAD, float(np.nanmax(lon)) + ADT_PAD, float(np.nanmin(lat)) - ADT_PAD, float(np.nanmax(lat)) + ADT_PAD) except Exception as e: - log.warning(f"Cannot read SSH_1 grid extent ({e}) — using ops ADT ROI") - return ADT_LON_MIN, ADT_LON_MAX, ADT_LAT_MIN, ADT_LAT_MAX + log.warning(f"Cannot read SSH_1 grid extent ({e}) — using config domain bounds") + lon_min, lon_max, lat_min, lat_max = self.config.domain + return lon_min - ADT_PAD, lon_max + ADT_PAD, lat_min - ADT_PAD, lat_max + ADT_PAD def _find_adt_data(self) -> Optional[Path]: """Newest available CMEMS ADT file (valid day, else previous day).""" @@ -218,23 +214,29 @@ def _warn(self, msg): log.warning(msg) self.warnings.append(msg) - @staticmethod - def _subset_ops_roi(adt_path: Path): + def _subset_ops_roi(self, adt_path: Path, bounds): """ncks -d longitude,-62.5,-51.5 -d latitude,7.0,54.0: inclusive coordinate bounds. Returns (adt[lat, lon] in m with NaN for missing, lons, lats, n_time). MJ (10/02/26) """ + lon_min, lon_max, lat_min, lat_max = bounds with Dataset(str(adt_path)) as ds: lon_name = "longitude" if "longitude" in ds.variables else "lon" lat_name = "latitude" if "latitude" in ds.variables else "lat" lons = np.array(ds.variables[lon_name][:], dtype=np.float64) lats = np.array(ds.variables[lat_name][:], dtype=np.float64) - li = np.where((lons >= ADT_LON_MIN) & (lons <= ADT_LON_MAX))[0] - la = np.where((lats >= ADT_LAT_MIN) & (lats <= ADT_LAT_MAX))[0] - if li.size == 0 or la.size == 0: - raise ValueError("ADT file does not cover the ops ROI") var = ds.variables["adt" if "adt" in ds.variables else "surf_el"] - a = np.ma.filled(var[..., la[0]:la[-1] + 1, li[0]:li[-1] + 1], np.nan) + a = var[...] + + from ..coords import lon_convention, shift_grid_0360 + if lon_convention(self.config) == "0360": + lons, a = shift_grid_0360(lons, a, lon_axis=-1) + + li = np.where((lons >= lon_min) & (lons <= lon_max))[0] + la = np.where((lats >= lat_min) & (lats <= lat_max))[0] + if li.size == 0 or la.size == 0: + raise ValueError("ADT file does not cover the domain ROI") + a = np.ma.filled(a[..., la[0]:la[-1] + 1, li[0]:li[-1] + 1], np.nan) a = np.asarray(a, dtype=np.float64) if a.ndim == 3: a = a[0] @@ -307,9 +309,10 @@ def dlon_diff(a, b): if row.min() < 0 or row.max() >= n_b or col.min() < 0 or col.max() >= n_a: raise ValueError("map row/col out of range") + bounds = self._dest_bounds(ssh_path) days = [] for f in adt_files: - a, lons, lats = self._subset_ops_roi(f) + a, lons, lats = self._subset_ops_roi(f, bounds) if a.size != n_a: raise ValueError(f"ADT ROI subset has {a.size} cells, map n_a={n_a}") if src_dims and tuple(src_dims) != (a.shape[1], a.shape[0]): @@ -426,8 +429,7 @@ def _read_adt(self, adt_path: Path, bounds=None): bounds is (lon_min, lon_max, lat_min, lat_max) in -180/180; default is the ops ROI. Returns (adt[lat, lon], lons, lats) or None. Fill values become NaN. """ - lon_min, lon_max, lat_min, lat_max = bounds or ( - ADT_LON_MIN, ADT_LON_MAX, ADT_LAT_MIN, ADT_LAT_MAX) + lon_min, lon_max, lat_min, lat_max = bounds or self._dest_bounds(None) try: ds = Dataset(str(adt_path)) try: @@ -436,6 +438,12 @@ def _read_adt(self, adt_path: Path, bounds=None): lons = np.array(ds.variables[lon_name][:], dtype=np.float64) lats = np.array(ds.variables[lat_name][:], dtype=np.float64) + adt_var = "adt" if "adt" in ds.variables else "surf_el" + adt_data = ds.variables[adt_var][...] + + from ..coords import lon_convention, shift_grid_0360 + if lon_convention(self.config) == "0360": + lons, adt_data = shift_grid_0360(lons, adt_data, lon_axis=-1) lon_idx = np.where((lons >= lon_min) & (lons <= lon_max))[0] lat_idx = np.where((lats >= lat_min) & (lats <= lat_max))[0] @@ -443,9 +451,6 @@ def _read_adt(self, adt_path: Path, bounds=None): if len(lon_idx) == 0 or len(lat_idx) == 0: log.warning("ADT data doesn't cover target domain") return None - - adt_var = "adt" if "adt" in ds.variables else "surf_el" - adt_data = ds.variables[adt_var][...] sl_y = slice(lat_idx[0], lat_idx[-1] + 1) sl_x = slice(lon_idx[0], lon_idx[-1] + 1) if adt_data.ndim == 3: diff --git a/nos_utils/forcing/nudging.py b/nos_utils/forcing/nudging.py index 6f7dd7b..8e0f1bc 100644 --- a/nos_utils/forcing/nudging.py +++ b/nos_utils/forcing/nudging.py @@ -742,17 +742,26 @@ def _process_python(self) -> ForcingResult: # time records (rare, but the reader supports it), they are # treated as consecutive 6h steps starting at the file's stamped # forecast hour. - rtofs_dt_input = 21600.0 + rtofs_dt_input = self.config.nudging_dt - for fi, f in enumerate(files_3d): + is_split = getattr(self.config, "rtofs_daily_3z", False) + if is_split: + from itertools import groupby + def get_time_key(fp): return fp.name.split('_daily')[0] + file_groups = [list(g) for k, g in groupby(sorted(files_3d, key=get_time_key), key=get_time_key)] + else: + file_groups = [[f] for f in files_3d] + + for group in file_groups: try: - ds = Dataset(str(f)) - lon = ds.variables.get("Longitude") or ds.variables.get("lon") - lat = ds.variables.get("Latitude") or ds.variables.get("lat") - depth_var = ds.variables.get("Depth") or ds.variables.get("lev") + datasets = [Dataset(str(f)) for f in group] + ds_main = datasets[0] + lon = ds_main.variables.get("Longitude") or ds_main.variables.get("lon") + lat = ds_main.variables.get("Latitude") or ds_main.variables.get("lat") + depth_var = ds_main.variables.get("Depth") or ds_main.variables.get("lev") if lon is None or lat is None: - ds.close() + for ds in datasets: ds.close() continue lon_arr = lon[:] @@ -762,7 +771,7 @@ def _process_python(self) -> ForcingResult: # File valid time relative to sim_start (model t=0). # Parses the f### or n### tag from the filename. - file_t0 = (rtofs_proc.valid_time(f) - sim_start).total_seconds() + file_t0 = (rtofs_proc.valid_time(group[0]) - sim_start).total_seconds() # Validate precomputed weights grid on first file if use_precomputed and not grid_validated: @@ -781,14 +790,20 @@ def _process_python(self) -> ForcingResult: # stay aligned with ``all_temp`` / ``all_salt``. appended_per_var = {"temperature": 0, "salinity": 0} + available_vars = {} + for ds in datasets: + for v in ["temperature", "salinity"]: + if v in ds.variables: + available_vars[v] = ds.variables[v] + for var_name, target_list in [ ("temperature", all_temp), ("salinity", all_salt), ]: - if var_name not in ds.variables: + if var_name not in available_vars: continue - data = ds.variables[var_name][:] + data = available_vars[var_name][:] data = np.ma.filled(data, fill_value=np.nan) data[np.abs(data) >= 99.0] = np.nan @@ -868,9 +883,9 @@ def _process_python(self) -> ForcingResult: for k in range(n_appended): all_times.append(file_t0 + k * rtofs_dt_input) - ds.close() + for ds in datasets: ds.close() except Exception as e: - log.warning(f"Failed to read RTOFS 3D file {f.name}: {e}") + log.warning(f"Failed to read RTOFS 3D group {group[0].name}: {e}") if not all_temp and not all_salt: return ForcingResult( @@ -882,10 +897,10 @@ def _process_python(self) -> ForcingResult: # ``sim_start`` and ``sim_duration`` were computed above so the # per-file valid times in ``all_times`` are already relative to # the model t=0 anchor (= cycle - nowcast_hours). - target_dt = 10800.0 # 3-hourly output (matches COMF Fortran) + target_dt = getattr(self.config, "nudging_dt", 10800.0) if self.config.obc_ops_timeline: - # step_nu_tr = 21600 s: phase-anchored 6-hourly records, time 0 = phase start. MJ (10/02/26) - target_dt = 21600.0 + # Ops phase-anchored records, time 0 = phase start. + target_dt = self.config.nudging_dt # The output time axis covers [0, sim_duration] at ``target_dt`` # cadence — i.e. 0, 3h, 6h, ..., nowcast+forecast hours from the @@ -1040,10 +1055,13 @@ def _process_python(self) -> ForcingResult: # Subset to valid nodes stacked = stacked[:, valid_indices, :] - # Align the time axis with the data length. The two variables - # share ``output_times`` but defensive clipping keeps them - # honest if one variable has fewer raw timesteps. - times_for_var = output_times[: stacked.shape[0]] + # Apply 2-step temporal moving average if configured + if getattr(self.config, "nudging_time_smoothing", False) and stacked.shape[0] > 1: + stacked = (stacked[:-1] + stacked[1:]) / 2.0 + times_for_var = output_times[:stacked.shape[0]] + log.info(f"Nudging: applied 2-step temporal smoothing, resulting in {stacked.shape[0]} steps") + else: + times_for_var = output_times[:stacked.shape[0]] self._write_nudge_nc( out_path, stacked, valid_node_ids, times_for_var, diff --git a/nos_utils/forcing/rtofs.py b/nos_utils/forcing/rtofs.py index 49c08dd..623dd48 100644 --- a/nos_utils/forcing/rtofs.py +++ b/nos_utils/forcing/rtofs.py @@ -62,12 +62,6 @@ HAS_NETCDF4 = False -# Ops OBC record spacing (stofs_3d_atl_obc_3dth_nc.in dtout) and the ops record count: 23 6-hourly -# records from the nowcast start, n012 to f120. MJ (10/02/26) -OPS_OBC_DT = 21600.0 -OPS_MIN_SLOTS = 23 - - def _pack_for_fortran(real, scale, offset, fill, fill_mask=None): """ Pack a float field to float exactly as the Fortran gen_*_hycom expects. @@ -1015,16 +1009,28 @@ def _stofs_prepare_tsuv( all_lat = None all_depth = None - for f in files_3d: + # Group files by time if they are split daily files + is_split = getattr(self.config, "rtofs_daily_3z", False) + if is_split: + from itertools import groupby + def get_time_key(fp): return fp.name.split('_daily')[0] + file_groups = [list(g) for k, g in groupby(sorted(files_3d, key=get_time_key), key=get_time_key)] + else: + file_groups = [[f] for f in files_3d] + + for group in file_groups: try: - ds = Dataset(str(f)) + datasets = [Dataset(str(f)) for f in group] + # Merge variables from multiple files into one dictionary subset = self._stofs_subset_roi( - ds, roi, + datasets[0], roi, ["temperature", "salinity", "u", "v", "Longitude", "Latitude"], ) + for ds in datasets[1:]: + subset.update(self._stofs_subset_roi(ds, roi, ["temperature", "salinity", "u", "v"])) if "temperature" not in subset: - ds.close() + for ds in datasets: ds.close() continue for t in range(subset["temperature"].shape[0]): @@ -1040,9 +1046,9 @@ def _stofs_prepare_tsuv( if all_depth is None and "Depth" in ds.variables: all_depth = np.array(ds.variables["Depth"][:]) - ds.close() + for ds in datasets: ds.close() except Exception as e: - log.warning(f"Failed to read 3D file {f.name}: {e}") + log.warning(f"Failed to read 3D group {group[0].name}: {e}") if not all_temp: return None @@ -1580,8 +1586,8 @@ def _ops_slot_times(self) -> List[datetime]: t0 = cycle_dt - timedelta(hours=int(self.config.nowcast_hours)) span = (cycle_dt - t0).total_seconds() + ( int(self.config.forecast_hours) + self.buffer_hours) * 3600.0 - n = max(OPS_MIN_SLOTS, int(np.ceil(span / OPS_OBC_DT)) + 1) - return [t0 + timedelta(seconds=OPS_OBC_DT * k) for k in range(n)] + n = max(self.config.obc_min_timesteps, int(np.ceil(span / self.config.obc_dt)) + 1) + return [t0 + timedelta(seconds=self.config.obc_dt * k) for k in range(n)] def _ops_candidates(self, date: datetime, kind: str): """[(valid time, is_nowcast, path)] of the size-valid 2D or 3D files of one RTOFS cycle date.""" @@ -1590,8 +1596,11 @@ def _ops_candidates(self, date: datetime, kind: str): if kind == "2d": globs, floor = ["rtofs_glo_2ds_*_diag.nc"], self.MIN_FILE_SIZE_2D else: - tag = region if region else "*" - globs = [f"rtofs_glo_3dz_*_6hrly_hvr_{tag}.nc", f"rtofs_glo_3dz_*_6hrly_hvr_{tag}.nc4"] + if getattr(self.config, "rtofs_daily_3z", False): + globs = ["rtofs_glo_3dz_*_daily_3z*io.nc"] + else: + tag = region if region else "*" + globs = [f"rtofs_glo_3dz_*_6hrly_hvr_{tag}.nc", f"rtofs_glo_3dz_*_6hrly_hvr_{tag}.nc4"] floor = self._minimum_3d_file_size() for d in (self.input_path / f"rtofs.{date_str}", self.input_path / date_str, self.input_path): if not d.exists(): @@ -1646,10 +1655,10 @@ def _find_ops_files(self) -> Tuple[List[Path], List[Path]]: if have and n_prev[kind]: self._select_notes.append( f"RTOFS {kind}: {n_prev[kind]} of {len(have)} slots filled from the previous cycle") - miss = [vt for vt in slots[:OPS_MIN_SLOTS] if vt not in picked[kind]] if have else [] + miss = [vt for vt in slots[:self.config.obc_min_timesteps] if vt not in picked[kind]] if have else [] if miss: self._select_notes.append( - f"RTOFS {kind}: {len(miss)} of the first {OPS_MIN_SLOTS} 6-hourly slots missing " + f"RTOFS {kind}: {len(miss)} of the first {self.config.obc_min_timesteps} slots missing " f"in both cycles (first {miss[0]:%Y-%m-%d %Hz})") for m in self._select_notes: log.warning(m) @@ -2190,6 +2199,12 @@ def _process_2d(self, files_2d: List[Path], ssh_source: Optional[Path] = None) - ssh_array, time_axis, model_t0, ) + # Apply 2-step temporal moving average if configured + if getattr(self.config, "obc_time_smoothing", False) and ssh_array.shape[0] > 1: + ssh_array = (ssh_array[:-1] + ssh_array[1:]) / 2.0 + time_axis = time_axis[:ssh_array.shape[0]] + log.info(f"elev2D: applied 2-step temporal smoothing, resulting in {ssh_array.shape[0]} steps") + return self._write_elev2d(ssh_array, time_axis, dt_out) except Exception as e: @@ -2248,15 +2263,15 @@ def _write_elev2d_ops(self, files_2d: List[Path], ops_ssh: np.ndarray, hold: boo try: off = np.float32(self.config.obc_ssh_offset) model_t0, _, dur = self._get_output_window(self.phase) - n = int(np.ceil(dur / OPS_OBC_DT)) + 1 + n = int(np.ceil(dur / self.config.obc_dt)) + 1 if hold: rows = np.repeat(ops_ssh[:1] + off, n, axis=0) else: times = [(self.valid_time(f) - model_t0).total_seconds() for f in files_2d] - rows = resample_records(ops_ssh + off, times, OPS_OBC_DT, n) - time_axis = np.arange(n, dtype=np.float64) * OPS_OBC_DT + rows = resample_records(ops_ssh + off, times, self.config.obc_dt, n) + time_axis = np.arange(n, dtype=np.float64) * self.config.obc_dt rows, self._wl_bias_info = self._apply_wl_bias(rows, time_axis, model_t0) - return self._write_elev2d(rows, time_axis, OPS_OBC_DT) + return self._write_elev2d(rows, time_axis, self.config.obc_dt) except Exception as e: log.error(f"Failed to write ops elev2D: {e}") return None @@ -2778,6 +2793,11 @@ def _process_3d(self, files_3d: List[Path]) -> List[Path]: dt = target_dt_3d if len(all_temp) > 1 else rtofs_dt_3d + if getattr(self.config, "obc_time_smoothing", False) and len(all_temp) > 1: + for lst in [all_temp, all_salt]: + arr = np.stack(lst, axis=0) + lst[:] = list((arr[:-1] + arr[1:]) / 2.0) + if all_temp: fpath = self.output_path / "TEM_3D.th.nc" merged = np.stack(all_temp, axis=0) @@ -2815,26 +2835,26 @@ def _process_3d(self, files_3d: List[Path]) -> List[Path]: def _write_3d_ops(self, files_3d, all_temp, all_salt, all_u, all_v) -> List[Path]: """Ops TEM_3D, SAL_3D and uv3D: this phase's slice of the 6-hourly series, time 0 at the phase start.""" model_t0, _, dur = self._get_output_window(self.phase) - n = int(np.ceil(dur / OPS_OBC_DT)) + 1 + n = int(np.ceil(dur / self.config.obc_dt)) + 1 times = [(self.valid_time(f) - model_t0).total_seconds() for f in files_3d] - if n > 1 and max(times) < (n - 1) * OPS_OBC_DT: - self._ops_note(f"T/S/uv: the RTOFS series ends {(n - 1) * OPS_OBC_DT / 3600 - max(times) / 3600:g} h " + if n > 1 and max(times) < (n - 1) * self.config.obc_dt: + self._ops_note(f"T/S/uv: the RTOFS series ends {(n - 1) * self.config.obc_dt / 3600 - max(times) / 3600:g} h " f"before the {n}-record phase window; the last record is held") n_bnd = len(self._bnd_lons) out = [] series = {} for name, lst in (("t", all_temp), ("s", all_salt), ("u", all_u), ("v", all_v)): - series[name] = resample_records(np.stack(lst), times, OPS_OBC_DT, n) + series[name] = resample_records(np.stack(lst), times, self.config.obc_dt, n) for fname, key, var, units in (("TEM_3D.th.nc", "t", "temperature", "degC"), ("SAL_3D.th.nc", "s", "salinity", "PSU")): path = self.output_path / fname - self._write_3d_th(path, series[key], var, units, OPS_OBC_DT, n_bnd) + self._write_3d_th(path, series[key], var, units, self.config.obc_dt, n_bnd) out.append(path) - log.info(f"Created {fname}: {series[key].shape} at dt={OPS_OBC_DT:g}s") + log.info(f"Created {fname}: {series[key].shape} at dt={self.config.obc_dt:g}s") path = self.output_path / "uv3D.th.nc" - self._write_uv3d_th(path, series["u"], series["v"], OPS_OBC_DT, n_bnd) + self._write_uv3d_th(path, series["u"], series["v"], self.config.obc_dt, n_bnd) out.append(path) - log.info(f"Created uv3D.th.nc: {series['u'].shape} at dt={OPS_OBC_DT:g}s from RTOFS u/v") + log.info(f"Created uv3D.th.nc: {series['u'].shape} at dt={self.config.obc_dt:g}s from RTOFS u/v") return out def _interpolate_vertical( From 50343eb0fc49832f7eec47bd074475c0c2f3775c Mon Sep 17 00:00:00 2001 From: SorooshMani-NOAA Date: Thu, 22 Oct 2026 08:03:23 -0400 Subject: [PATCH 4/4] feat(forcing): implement surface restore processor and dynamic archiving Adds the Surface Restore processor to replicate the operational STOFS-3D-PAC surface temperature and salinity relaxation pipeline, and refactors orchestrator archiving to dynamically reflect configuration toggles. * nos_utils/forcing/surface_restore.py: - Create `SurfaceRestoreProcessor` to extract and interpolate RTOFS surface T/S to target nodes. - Implement a configurable 2-step temporal moving average for surface restore outputs. - Write results to standard `surface_restore.nc`. * nos_utils/config.py: - Add `surface_restore_enabled`, `surface_restore_time_smoothing`, and `surface_restore_weight_file` to `ForcingConfig`. * nos_utils/orchestrator.py: - Integrate `_run_surface_restore` into the main prep pipeline. - Add `surface_restore.nc` to the `_OPS_RUN_INPUTS` completeness check. - Refactor `_OBC_ARCHIVE_NAMES` into a dynamic `@property` to automatically include nudging and surface restore files in QA/QC gates and fallback staging when enabled. - Update tarball generation blocks (`payload_files`, `obc_files`) to dynamically read from `_OBC_ARCHIVE_NAMES.keys()`. --- nos_utils/config.py | 14 ++ nos_utils/forcing/surface_restore.py | 205 +++++++++++++++++++++++++++ nos_utils/orchestrator.py | 66 ++++++--- 3 files changed, 268 insertions(+), 17 deletions(-) create mode 100755 nos_utils/forcing/surface_restore.py diff --git a/nos_utils/config.py b/nos_utils/config.py index 33add6b..e9ef202 100644 --- a/nos_utils/config.py +++ b/nos_utils/config.py @@ -248,6 +248,10 @@ class ForcingConfig: nudging_dt: float = 21600.0 # Apply 2-step temporal moving average to nudging fields. nudging_time_smoothing: bool = False + # Surface restore enabled and smoothing + surface_restore_enabled: bool = False + surface_restore_time_smoothing: bool = False + surface_restore_weight_file: Optional[Path] = None # Number of vertical levels n_levels: int = 51 # SCHISM model timestep (seconds) used to write the elev2D.th.nc @@ -728,6 +732,8 @@ def for_stofs_3d_pac(cls, pdy: str, cyc: int, **overrides) -> "ForcingConfig": obc_roi_2d={"x1": 222, "x2": 2703, "y1": 1051, "y2": 2641, "x_stride": 2, "y_stride": 2}, obc_roi_3d={"x1": 222, "x2": 2703, "y1": 1051, "y2": 2641, "x_stride": 2, "y_stride": 2}, nudge_roi_3d={"x1": 222, "x2": 2703, "y1": 1051, "y2": 2641, "x_stride": 2, "y_stride": 2}, + surface_restore_enabled=True, + surface_restore_time_smoothing=True, ) defaults.update(overrides) return cls(**defaults) @@ -766,6 +772,8 @@ def for_stofs_3d_pac_ufs(cls, pdy: str, cyc: int, **overrides) -> "ForcingConfig obc_roi_2d={"x1": 222, "x2": 2703, "y1": 1051, "y2": 2641, "x_stride": 2, "y_stride": 2}, obc_roi_3d={"x1": 222, "x2": 2703, "y1": 1051, "y2": 2641, "x_stride": 2, "y_stride": 2}, nudge_roi_3d={"x1": 222, "x2": 2703, "y1": 1051, "y2": 2641, "x_stride": 2, "y_stride": 2}, + surface_restore_enabled=True, + surface_restore_time_smoothing=True, # DATM grid — STOFS3D_PAC preset, verified from weight_gfs.nc # nx = round((290.0-92.5)/0.025)+1 = 7901; ny = round((67.0-(-33.5))/0.025)+1 = 4021 datm_lon_min=92.5, datm_lon_max=290.0, @@ -1072,6 +1080,12 @@ def from_yaml(cls, yaml_path, **overrides) -> "ForcingConfig": kwargs["nudging_dt"] = float(nudge.get("dt", kwargs["obc_dt"])) kwargs["nudging_time_smoothing"] = _strict_bool(nudge.get("time_smoothing"), "nudging.time_smoothing", False) + surf_res = ocean.get("surface_restore", {}) if isinstance(ocean, dict) else {} + kwargs["surface_restore_enabled"] = _strict_bool(surf_res.get("enabled"), "ocean.surface_restore.enabled", False) + kwargs["surface_restore_time_smoothing"] = _strict_bool(surf_res.get("time_smoothing"), "ocean.surface_restore.time_smoothing", False) + if surf_res.get("weight_file"): + kwargs["surface_restore_weight_file"] = Path(surf_res["weight_file"]) + # NWM target counts if isinstance(river, dict): n_target = river.get("n_list_target") diff --git a/nos_utils/forcing/surface_restore.py b/nos_utils/forcing/surface_restore.py new file mode 100755 index 0000000..48869b9 --- /dev/null +++ b/nos_utils/forcing/surface_restore.py @@ -0,0 +1,205 @@ +""" +Surface Restore processor for SCHISM. + +Interpolates RTOFS surface temperature and salinity to boundary/restore nodes +and applies a 2-step temporal moving average, replicating the ops pipeline: + ncks -O -d time,0,3 ... ncks -O -d time,1,4 ... ncecat ... ncwa + +Output: + - surface_restore.nc +""" + +import logging +from pathlib import Path +from typing import List, Optional, Tuple +import numpy as np + +from ..config import ForcingConfig +from .base import ForcingProcessor, ForcingResult + +log = logging.getLogger(__name__) + +try: + from netCDF4 import Dataset + HAS_NETCDF4 = True +except ImportError: + HAS_NETCDF4 = False + + +class SurfaceRestoreProcessor(ForcingProcessor): + """Generate surface_restore.nc from RTOFS data.""" + + SOURCE_NAME = "SURFACE_RESTORE" + + def __init__( + self, + config: ForcingConfig, + input_path: Path, + output_path: Path, + weight_file: Optional[Path] = None, + phase: Optional[str] = None, + ): + super().__init__(config, input_path, output_path) + self.weight_file = weight_file + self.phase = phase + + def process(self) -> ForcingResult: + if not HAS_NETCDF4: + return ForcingResult( + success=False, source=self.SOURCE_NAME, + errors=["netCDF4 required for surface restore"], + ) + + if not getattr(self.config, "surface_restore_enabled", False): + return ForcingResult( + success=True, source=self.SOURCE_NAME, + warnings=["Surface restore disabled in config"], + ) + + self.create_output_dir() + + # Reuse RTOFS 3D discovery to get the raw daily/6-hourly files + from .rtofs import RTOFSProcessor + rtofs_proc = RTOFSProcessor( + self.config, self.input_path, self.output_path, phase=self.phase + ) + _, files_3d = rtofs_proc.find_input_files_by_type() + + if not files_3d: + return ForcingResult( + success=False, source=self.SOURCE_NAME, + errors=["No RTOFS 3D files found for surface restore"], + ) + + # Load target nodes from the gr3 weight file (same layout as nudging) + node_ids, lons, lats = self._load_restore_nodes() + if node_ids is None: + return ForcingResult( + success=False, source=self.SOURCE_NAME, + errors=["Cannot identify restore nodes: weight_file missing"], + ) + + n_nodes = len(node_ids) + all_temp, all_salt, times = [], [], [] + + # RTOFS Extraction (Surface Only) + for f in files_3d: + try: + ds = Dataset(str(f)) + lon_arr = ds.variables.get("Longitude") or ds.variables.get("lon") + lat_arr = ds.variables.get("Latitude") or ds.variables.get("lat") + + if lon_arr is None or lat_arr is None: + ds.close() + continue + + # Fetch surface layer (index 0) + temp = ds.variables.get("temperature") + salt = ds.variables.get("salinity") + + if temp is not None: + # [time, depth, lat, lon] -> grab surface depth (0) + t_surf = np.ma.filled(temp[:, 0, :, :], fill_value=np.nan) + t_surf[np.abs(t_surf) >= 99.0] = np.nan + + s_surf = np.ma.filled(salt[:, 0, :, :], fill_value=np.nan) + s_surf[np.abs(s_surf) >= 99.0] = np.nan + + for t in range(t_surf.shape[0]): + t_bnd = rtofs_proc._interpolate_2d_to_boundary(lon_arr[:], lat_arr[:], t_surf[t]) + s_bnd = rtofs_proc._interpolate_2d_to_boundary(lon_arr[:], lat_arr[:], s_surf[t]) + all_temp.append(t_bnd) + all_salt.append(s_bnd) + times.append(float(t)) # Placeholder for actual valid time + + ds.close() + except Exception as e: + log.warning(f"Failed to read 3D file for restore {f.name}: {e}") + + if not all_temp: + return ForcingResult( + success=False, source=self.SOURCE_NAME, + errors=["Failed to extract surface T/S"], + ) + + # Convert to numpy arrays: shape (n_times, n_nodes) + temp_arr = np.stack(all_temp, axis=0) + salt_arr = np.stack(all_salt, axis=0) + + dt_out = getattr(self.config, "obc_dt", 86400.0) + + # Apply 2-step temporal moving average if configured + if getattr(self.config, "surface_restore_time_smoothing", False) and temp_arr.shape[0] > 1: + temp_avg = (temp_arr[:-1] + temp_arr[1:]) / 2.0 + salt_avg = (salt_arr[:-1] + salt_arr[1:]) / 2.0 + time_axis = np.arange(temp_avg.shape[0]) * dt_out + log.info(f"Surface restore: applied 2-step temporal smoothing, resulting in {temp_avg.shape[0]} steps") + else: + temp_avg, salt_avg = temp_arr, salt_arr + time_axis = np.arange(temp_avg.shape[0]) * dt_out + + # Replace remaining NaNs with typical oceanic fill values + temp_avg = np.where(np.isnan(temp_avg), 15.0, temp_avg) + salt_avg = np.where(np.isnan(salt_avg), 35.0, salt_avg) + + # Write output file + out_path = self.output_path / "surface_restore.nc" + self._write_restore_nc(out_path, temp_avg, salt_avg, node_ids, time_axis) + + return ForcingResult( + success=True, + source=self.SOURCE_NAME, + output_files=[out_path], + metadata={ + "n_restore_nodes": n_nodes, + "n_timesteps": len(time_axis), + } + ) + + def _load_restore_nodes(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Load target nodes from the gr3 weight file.""" + from ..io.schism_grid import SchismGrid + weight_file = self.weight_file or self._find_restore_file() + + if not weight_file or not Path(weight_file).exists(): + return None, None, None + + node_ids, lons, lats, values = SchismGrid.read_gr3_values(weight_file) + mask = values > 0 + return node_ids[mask], lons[mask], lats[mask] + + def _find_restore_file(self) -> Optional[Path]: + import os + patterns = ["*.nudge.gr3", "*.TEM_nudge.gr3"] + for env_var in ["FIXofs", "FIXstofs3d"]: + fix_dir = os.environ.get(env_var) + if fix_dir: + for pattern in patterns: + matches = sorted(Path(fix_dir).glob(pattern)) + if matches: return matches[0] + return None + + def _write_restore_nc(self, output_path: Path, temp: np.ndarray, salt: np.ndarray, + node_ids: np.ndarray, time_axis: np.ndarray) -> None: + """Write surface_restore.nc using standard SCHISM formats.""" + nt, n_nodes = temp.shape + + nc = Dataset(str(output_path), "w", format="NETCDF4_CLASSIC") + nc.createDimension("time", nt) + nc.createDimension("node", n_nodes) + nc.createDimension("one", 1) + + time_var = nc.createVariable("time", "f8", ("time",)) + time_var[:] = time_axis + + map_var = nc.createVariable("map_to_global_node", "i4", ("node",)) + map_var[:] = node_ids.astype(np.int32) + + t_var = nc.createVariable("temperature", "f8", ("time", "node", "one")) + t_var[:, :, 0] = temp + + s_var = nc.createVariable("salinity", "f8", ("time", "node", "one")) + s_var[:, :, 0] = salt + + nc.close() + log.info(f"Created surface_restore.nc: {nt} times, {n_nodes} nodes.") diff --git a/nos_utils/orchestrator.py b/nos_utils/orchestrator.py index 10c8f38..4004bdd 100644 --- a/nos_utils/orchestrator.py +++ b/nos_utils/orchestrator.py @@ -309,6 +309,10 @@ def run(self, phase: str = "nowcast") -> PrepResult: if self.config.nudging_enabled and "rtofs" in self.paths: results.append(self._run_nudging(output_dir, phase, time_hotstart)) + # Surface restore needs RTOFS to have completed + if getattr(self.config, "surface_restore_enabled", False) and "rtofs" in self.paths: + results.append(self._run_surface_restore(output_dir, phase)) + # param.nml only needs time_hotstart if "fix" in self.paths: results.append(self._run_param_nml(output_dir, phase, time_hotstart)) @@ -675,16 +679,25 @@ def _prev_cycle_dirs(self, max_days: int = 3) -> List[Path]: dirs += [r / f"{self.run_name}.{day}", r / day] return list(dict.fromkeys(dirs)) - # Mapping from the active OBC filename (what SCHISM reads) to the - # standard archive basename under $COMOUT_PREV/rerun/ (see - # stofs_3d_atl_create_obc_3d_th_non_adjust.sh:641-642). The archive - # file is stored as `...nc`. - _OBC_ARCHIVE_NAMES: Dict[str, str] = { - "elev2D.th.nc": "elev2dth_non_adj", - "TEM_3D.th.nc": "tem3dth", - "SAL_3D.th.nc": "sal3dth", - "uv3D.th.nc": "uv3dth", - } + @property + def _OBC_ARCHIVE_NAMES(self) -> Dict[str, str]: + """Mapping of active OBC files to their archive basenames. + + Dynamically includes nudging and surface restore if enabled, and handles + the difference between ATL (elev2dth_non_adj) and PAC (elev2dth). + """ + names = { + "elev2D.th.nc": "elev2dth_non_adj" if getattr(self.config, "dynamic_adjust_enabled", False) else "elev2dth", + "TEM_3D.th.nc": "tem3dth", + "SAL_3D.th.nc": "sal3dth", + "uv3D.th.nc": "uv3dth", + } + if getattr(self.config, "nudging_enabled", False): + names["TEM_nu.nc"] = "temnu" + names["SAL_nu.nc"] = "salnu" + if getattr(self.config, "surface_restore_enabled", False): + names["surface_restore.nc"] = "surface_restore" + return names def _obc_archive_name(self, archive_base: str, phase: str) -> str: """``...nc``; the forecast-phase file gets ``.fcst`` before the suffix. @@ -834,6 +847,8 @@ def _check_ops_obc_inputs(self, output_dir: Path) -> ForcingResult: (nco_v315 scripts/stofs_3d_atl/exstofs_3d_atl_prep_processing.sh:398-746). MJ (10/05/26) """ wanted = list(self._OPS_RUN_INPUTS) + if getattr(self.config, "surface_restore_enabled", False): + wanted.append("surface_restore.nc") if self.config.nws == 4: wanted += ["gfs_forcing.nc"] + (["hrrr_forcing.nc"] if self.config.met_num >= 2 else []) else: @@ -1018,6 +1033,27 @@ def _run_nudging(self, output_dir: Path, phase: str = "nowcast", ) return proc.process() + def _run_surface_restore(self, output_dir: Path, phase: str = "nowcast") -> ForcingResult: + """Step 5c: Surface restore (STOFS-3D-PAC).""" + from .forcing.surface_restore import SurfaceRestoreProcessor + + weight_file = getattr(self.config, "surface_restore_weight_file", None) + if not weight_file: + fix_dir = self.paths.get("fix") + if fix_dir: + for pattern in ["*.nudge.gr3", "*.TEM_nudge.gr3"]: + matches = sorted(Path(fix_dir).glob(pattern)) + if matches: + weight_file = matches[0] + break + + proc = SurfaceRestoreProcessor( + self.config, self.paths["rtofs"], output_dir, + weight_file=weight_file, + phase=phase, + ) + return proc.process() + def _run_tidal(self, output_dir: Path, phase: str = "nowcast", time_hotstart=None) -> ForcingResult: """Step 6: Tidal forcing.""" @@ -1220,10 +1256,7 @@ def _build_archive_manifest( # --- COMMON: OBC (phase-specific, backward-compat) --- { "tar_name": f"{prefix}.{cycle}.{pdy}.obc.{phase}.tar", - "payload_files": [ - "elev2D.th.nc", "TEM_3D.th.nc", "SAL_3D.th.nc", - "uv3D.th.nc", "TEM_nu.nc", "SAL_nu.nc", - ], + "payload_files": list(self._OBC_ARCHIVE_NAMES.keys()), "condition": True, "label": "OBC", }, @@ -1555,9 +1588,8 @@ def archive_to_comout(self, result: PrepResult, comout: Path) -> List[Path]: # unchanged, executed verbatim when the manifest flag is OFF. # Tar OBC files (phase-specific): obc.nowcast.tar / obc.forecast.tar, - # each with all 6 files (boundary + nudging). - obc_files = ["elev2D.th.nc", "TEM_3D.th.nc", "SAL_3D.th.nc", "uv3D.th.nc", - "TEM_nu.nc", "SAL_nu.nc"] + # each with all active files (boundary + nudging + surface restore). + obc_files = list(self._OBC_ARCHIVE_NAMES.keys()) existing_obc = [work_dir / f for f in obc_files if (work_dir / f).exists()] if existing_obc: # Phase-specific tar (backward compat): obc.nowcast.tar / obc.forecast.tar