microstructure-tpw 0.3.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (119) hide show
  1. microstructure_tpw-0.3.0.dist-info/METADATA +234 -0
  2. microstructure_tpw-0.3.0.dist-info/RECORD +119 -0
  3. microstructure_tpw-0.3.0.dist-info/WHEEL +5 -0
  4. microstructure_tpw-0.3.0.dist-info/entry_points.txt +7 -0
  5. microstructure_tpw-0.3.0.dist-info/licenses/LICENSE +674 -0
  6. microstructure_tpw-0.3.0.dist-info/top_level.txt +1 -0
  7. odas_tpw/__init__.py +18 -0
  8. odas_tpw/chi/__init__.py +48 -0
  9. odas_tpw/chi/batchelor.py +189 -0
  10. odas_tpw/chi/chi.py +857 -0
  11. odas_tpw/chi/fp07.py +649 -0
  12. odas_tpw/chi/l2_chi.py +153 -0
  13. odas_tpw/chi/l3_chi.py +388 -0
  14. odas_tpw/chi/l4_chi.py +391 -0
  15. odas_tpw/config_base.py +423 -0
  16. odas_tpw/perturb/__init__.py +18 -0
  17. odas_tpw/perturb/adcp.py +274 -0
  18. odas_tpw/perturb/atomic_io.py +47 -0
  19. odas_tpw/perturb/binning.py +761 -0
  20. odas_tpw/perturb/cli.py +591 -0
  21. odas_tpw/perturb/combo.py +346 -0
  22. odas_tpw/perturb/config.py +1046 -0
  23. odas_tpw/perturb/ctd.py +424 -0
  24. odas_tpw/perturb/diag/__init__.py +12 -0
  25. odas_tpw/perturb/diag/_common.py +140 -0
  26. odas_tpw/perturb/diag/chi.py +42 -0
  27. odas_tpw/perturb/diag/cli.py +67 -0
  28. odas_tpw/perturb/diag/data.py +662 -0
  29. odas_tpw/perturb/diag/epsilon.py +45 -0
  30. odas_tpw/perturb/diag/inspector.py +342 -0
  31. odas_tpw/perturb/diag/mixing.py +94 -0
  32. odas_tpw/perturb/diag/render.py +543 -0
  33. odas_tpw/perturb/discover.py +40 -0
  34. odas_tpw/perturb/fp07_cal.py +354 -0
  35. odas_tpw/perturb/gps.py +458 -0
  36. odas_tpw/perturb/hotel.py +572 -0
  37. odas_tpw/perturb/logging_setup.py +181 -0
  38. odas_tpw/perturb/merge.py +320 -0
  39. odas_tpw/perturb/mk_sections.py +310 -0
  40. odas_tpw/perturb/netcdf_schema.py +520 -0
  41. odas_tpw/perturb/pipeline.py +2654 -0
  42. odas_tpw/perturb/plot/__init__.py +15 -0
  43. odas_tpw/perturb/plot/cli.py +107 -0
  44. odas_tpw/perturb/plot/diagnostics.py +195 -0
  45. odas_tpw/perturb/plot/eps_chi.py +524 -0
  46. odas_tpw/perturb/plot/figure.py +679 -0
  47. odas_tpw/perturb/plot/gamma_scaling.py +1097 -0
  48. odas_tpw/perturb/plot/grid.py +112 -0
  49. odas_tpw/perturb/plot/layout.py +361 -0
  50. odas_tpw/perturb/plot/overview.py +547 -0
  51. odas_tpw/perturb/plot/presets/eps-chi.yaml +16 -0
  52. odas_tpw/perturb/plot/presets/profiles.yaml +32 -0
  53. odas_tpw/perturb/plot/presets/scalar.yaml +21 -0
  54. odas_tpw/perturb/plot/profiles.py +652 -0
  55. odas_tpw/perturb/plot/scalar.py +391 -0
  56. odas_tpw/perturb/plot/sections.py +506 -0
  57. odas_tpw/perturb/plot/xaxis.py +396 -0
  58. odas_tpw/perturb/qc_gate.py +177 -0
  59. odas_tpw/perturb/qc_rules.py +373 -0
  60. odas_tpw/perturb/resolve.py +306 -0
  61. odas_tpw/perturb/seawater.py +76 -0
  62. odas_tpw/perturb/speed.py +14 -0
  63. odas_tpw/perturb/trim.py +313 -0
  64. odas_tpw/processing/__init__.py +62 -0
  65. odas_tpw/processing/bottom.py +215 -0
  66. odas_tpw/processing/chi_combine.py +300 -0
  67. odas_tpw/processing/ct_align.py +172 -0
  68. odas_tpw/processing/epsilon_combine.py +210 -0
  69. odas_tpw/processing/mixing.py +630 -0
  70. odas_tpw/processing/thorpe.py +458 -0
  71. odas_tpw/processing/top_trim.py +274 -0
  72. odas_tpw/pyturb/__init__.py +3 -0
  73. odas_tpw/pyturb/__main__.py +5 -0
  74. odas_tpw/pyturb/_compat.py +204 -0
  75. odas_tpw/pyturb/_profind.py +124 -0
  76. odas_tpw/pyturb/bin.py +113 -0
  77. odas_tpw/pyturb/cli.py +175 -0
  78. odas_tpw/pyturb/eps.py +353 -0
  79. odas_tpw/pyturb/merge.py +101 -0
  80. odas_tpw/pyturb/p2nc.py +76 -0
  81. odas_tpw/rsi/__init__.py +31 -0
  82. odas_tpw/rsi/adapter.py +314 -0
  83. odas_tpw/rsi/bench.py +904 -0
  84. odas_tpw/rsi/binning.py +193 -0
  85. odas_tpw/rsi/channels.py +376 -0
  86. odas_tpw/rsi/chi_io.py +980 -0
  87. odas_tpw/rsi/cli.py +1411 -0
  88. odas_tpw/rsi/combine.py +120 -0
  89. odas_tpw/rsi/config.py +126 -0
  90. odas_tpw/rsi/config_patch.py +790 -0
  91. odas_tpw/rsi/convert.py +651 -0
  92. odas_tpw/rsi/deconvolve.py +119 -0
  93. odas_tpw/rsi/diss_look.py +343 -0
  94. odas_tpw/rsi/dissipation.py +669 -0
  95. odas_tpw/rsi/helpers.py +527 -0
  96. odas_tpw/rsi/mixing_look.py +509 -0
  97. odas_tpw/rsi/p_file.py +771 -0
  98. odas_tpw/rsi/pipeline.py +924 -0
  99. odas_tpw/rsi/profile.py +624 -0
  100. odas_tpw/rsi/quick_look.py +758 -0
  101. odas_tpw/rsi/sensor_inventory.py +684 -0
  102. odas_tpw/rsi/shear_noise.py +128 -0
  103. odas_tpw/rsi/speed.py +221 -0
  104. odas_tpw/rsi/vehicle.py +74 -0
  105. odas_tpw/rsi/viewer_base.py +1071 -0
  106. odas_tpw/rsi/window.py +459 -0
  107. odas_tpw/scor160/__init__.py +36 -0
  108. odas_tpw/scor160/cli.py +242 -0
  109. odas_tpw/scor160/compare.py +484 -0
  110. odas_tpw/scor160/despike.py +227 -0
  111. odas_tpw/scor160/goodman.py +288 -0
  112. odas_tpw/scor160/io.py +690 -0
  113. odas_tpw/scor160/l2.py +311 -0
  114. odas_tpw/scor160/l3.py +279 -0
  115. odas_tpw/scor160/l4.py +850 -0
  116. odas_tpw/scor160/nasmyth.py +208 -0
  117. odas_tpw/scor160/ocean.py +287 -0
  118. odas_tpw/scor160/profile.py +231 -0
  119. odas_tpw/scor160/spectral.py +482 -0
@@ -0,0 +1,2654 @@
1
+ # Mar-2026, Claude and Pat Welch, pat@mousebrains.com
2
+ """Orchestration — per-file processing and full pipeline.
3
+
4
+ Reference: Code/process_P_files.m (233 lines), Code/mat2profile.m (170 lines)
5
+ """
6
+
7
+ import contextlib
8
+ import hashlib
9
+ import json
10
+ import logging
11
+ import os
12
+ import re
13
+ from concurrent.futures import ProcessPoolExecutor, as_completed
14
+ from contextlib import suppress
15
+ from pathlib import Path
16
+ from typing import TYPE_CHECKING, Any
17
+
18
+ if TYPE_CHECKING:
19
+ import xarray as xr
20
+
21
+ from odas_tpw.perturb.config import (
22
+ config_dir_of,
23
+ expand_config_dir,
24
+ merge_config,
25
+ resolve_output_dir,
26
+ resolve_window_config,
27
+ )
28
+ from odas_tpw.perturb.config import (
29
+ upstream_for as _upstream_for,
30
+ )
31
+ from odas_tpw.perturb.logging_setup import (
32
+ current_run_stamp,
33
+ init_worker_logging,
34
+ stage_log,
35
+ )
36
+
37
+ logger = logging.getLogger(__name__)
38
+
39
+
40
+ _SAFE_STEM_RE = re.compile(r"[^A-Za-z0-9_.-]+")
41
+
42
+
43
+ def _safe_stem_part(part: str) -> str:
44
+ """Return a filesystem-friendly stem segment."""
45
+ safe = _SAFE_STEM_RE.sub("_", part).strip("._")
46
+ return safe or "unnamed"
47
+
48
+
49
+ def _source_output_stem(path: Path, root: Path | str) -> str:
50
+ """Return a unique, stable output stem for *path* relative to *root*."""
51
+ path = Path(path)
52
+ root = Path(root)
53
+ try:
54
+ rel = path.resolve().relative_to(root.resolve()).with_suffix("")
55
+ parts = rel.parts
56
+ except ValueError:
57
+ # File is not under root. Use parent directory name to provide likely
58
+ # uniqueness without relying on an absolute path that breaks cache portability.
59
+ parts = (path.parent.name, path.stem)
60
+ return "__".join(_safe_stem_part(str(part)) for part in parts)
61
+
62
+
63
+ # {stem}_prof{NNN}.nc — greedy stem so a stem containing "_prof" (or a current
64
+ # stem that is a prefix of an orphan's stem) still anchors on the final
65
+ # ``_prof{digits}.nc`` suffix, recovering the exact source stem.
66
+ _PROF_NC_RE = re.compile(r"^(?P<stem>.+)_prof\d+\.nc$")
67
+
68
+
69
+ def _prune_orphan_profile_ncs(stage_dir: Path, valid_stems: set[str]) -> int:
70
+ """Delete per-profile NetCDFs whose source .p file is no longer present.
71
+
72
+ The per-stage output dir is content-hashed on *config* and reused across
73
+ runs, so its ``{stem}_prof{NNN}.nc`` files accumulate. The signature does
74
+ not include the resolved input-file set, so re-running the same config with
75
+ a changed set of .p files (one dropped, or an explicit reduced list) would
76
+ leave orphaned NCs that the binning glob silently folds into the combos.
77
+
78
+ Match the EXACT source stem (parse ``{stem}_prof{NNN}.nc`` and require
79
+ ``stem in valid_stems``) rather than a name-prefix: a prefix test wrongly
80
+ KEEPS an orphan whose stem has a current stem as a prefix (e.g. current
81
+ ``a`` would protect a ``a_extra``-derived file). Only recognized profile
82
+ NetCDFs are pruned, so files not matching the ``_prof{NNN}`` shape are left
83
+ untouched and a file produced this run is never deleted. Returns the count.
84
+ """
85
+ if not stage_dir.exists():
86
+ return 0
87
+ pruned = 0
88
+ for nc in stage_dir.glob("*_prof*.nc"):
89
+ m = _PROF_NC_RE.match(nc.name)
90
+ if m is None or m.group("stem") in valid_stems:
91
+ continue # not a recognized profile NC, or a current stem -> keep
92
+ try:
93
+ nc.unlink()
94
+ pruned += 1
95
+ logger.info("pruned orphaned %s (no source .p in current input set)", nc.name)
96
+ except OSError as exc:
97
+ logger.warning("could not prune %s: %s", nc.name, exc)
98
+ return pruned
99
+
100
+
101
+ def _write_binned_or_clear(ds: "xr.Dataset", out_dir: Path, manifest: str | None = None) -> None:
102
+ """Write *ds* to ``out_dir/binned.nc``, or remove a stale one if *ds* is empty.
103
+
104
+ An empty re-run (same config but a reduced/changed input set) must not leave
105
+ a prior run's ``binned.nc`` in place: ``resolve_output_dir`` hashes only the
106
+ config, not the input file set, so the same dir is reused across input sets,
107
+ and the combo glob would otherwise republish the stale file as current data
108
+ (#56).
109
+
110
+ *manifest* (a hash of the contributing per-file cache keys) is stored as the
111
+ ``_input_manifest`` attribute so a later run can skip re-binning when the
112
+ inputs are unchanged — keyed on content identity, NOT mtime (the data volume
113
+ may be exFAT, whose 2 s mtime granularity makes mtime ordering unsafe).
114
+ """
115
+ out = out_dir / "binned.nc"
116
+ if ds.data_vars:
117
+ if manifest is not None:
118
+ ds = ds.assign_attrs(_input_manifest=manifest)
119
+ # Atomic write: to a temp file in the same dir, then os.replace into
120
+ # place. A direct ds.to_netcdf(out) interrupted mid-payload (ENOSPC /
121
+ # SMB drop / Ctrl-C) leaves a truncated binned.nc whose _input_manifest
122
+ # header already validates as cache-current, permanently publishing
123
+ # fill data with no error (audit r1-5). os.replace is atomic within a
124
+ # filesystem, so a partial write never becomes the live file. (Mirrors
125
+ # _write_marker and trim.py; the temp file shares out_dir's filesystem.)
126
+ tmp = out.with_name(f".binned.nc.{os.getpid()}.tmp")
127
+ try:
128
+ ds.to_netcdf(tmp)
129
+ os.replace(tmp, out)
130
+ except BaseException:
131
+ with contextlib.suppress(OSError):
132
+ os.unlink(tmp)
133
+ raise
134
+ elif out.exists():
135
+ out.unlink()
136
+
137
+
138
+ def _atomic_to_netcdf(ds: "xr.Dataset", out_path: Path) -> None:
139
+ """Write *ds* to *out_path* atomically (temp file + os.replace).
140
+
141
+ Thin wrapper over the shared :func:`odas_tpw.perturb.atomic_io.atomic_to_netcdf`
142
+ (one canonical implementation, also used by ctd.py; kept under this name for
143
+ the diss/chi call sites). A direct ds.to_netcdf(out_path) interrupted
144
+ mid-payload (ENOSPC / SMB drop / Ctrl-C) leaves a readable PARTIAL NetCDF the
145
+ bin/combo manifest then locks in on a clean retry; os.replace never lets a
146
+ partial write become the live file. (audit 2026-07-01; #104 U5-2.)
147
+ """
148
+ from odas_tpw.perturb.atomic_io import atomic_to_netcdf
149
+
150
+ atomic_to_netcdf(ds, out_path)
151
+
152
+
153
+ _MANIFEST_ATTR = "_input_manifest"
154
+
155
+
156
+ def _inputs_manifest(
157
+ nc_files: "list[Path]", file_cachekeys: dict[str, str], volatile_cfg: dict | None = None
158
+ ) -> str:
159
+ """A content-identity hash of a bin/combo's inputs: each input file paired
160
+ with the per-file cache key of its source ``.p`` (or, for a binned input,
161
+ that file's own manifest). mtime-immune, so it is safe on exFAT."""
162
+ items: list[Any] = []
163
+ for f in sorted(nc_files):
164
+ m = _PROF_NC_RE.match(f.name)
165
+ stem = m.group("stem") if m else f.stem
166
+ items.append([f.name, file_cachekeys.get(stem, "?")])
167
+ if volatile_cfg:
168
+ items.append(["_volatile_cfg", volatile_cfg])
169
+ return hashlib.sha256(
170
+ json.dumps(items, sort_keys=True, separators=(",", ":")).encode()
171
+ ).hexdigest()
172
+
173
+
174
+ def _output_is_current(nc_path: Path, manifest: str | None) -> bool:
175
+ """True iff *manifest* is set, *nc_path* exists, and its ``_input_manifest``
176
+ matches (None manifest → never current → always rebuild)."""
177
+ return manifest is not None and _read_manifest_attr(nc_path) == manifest
178
+
179
+
180
+ def _read_manifest_attr(nc_path: Path) -> str | None:
181
+ """The ``_input_manifest`` global attribute of *nc_path*, or None."""
182
+ if not nc_path.exists():
183
+ return None
184
+ try:
185
+ import xarray as xr
186
+
187
+ with xr.open_dataset(nc_path) as ds:
188
+ val = ds.attrs.get(_MANIFEST_ATTR)
189
+ return str(val) if val is not None else None
190
+ except (OSError, ValueError, KeyError):
191
+ return None
192
+
193
+
194
+ def _stamp_manifest(nc_path: Path, manifest: str) -> None:
195
+ """Add the ``_input_manifest`` global attribute to an already-written NC
196
+ (header-only append — no data rewrite), so a combo can be skipped next run
197
+ when its inputs are unchanged."""
198
+ try:
199
+ import netCDF4
200
+
201
+ with netCDF4.Dataset(nc_path, "a") as nc:
202
+ nc.setncattr(_MANIFEST_ATTR, manifest)
203
+ except OSError as exc:
204
+ logger.warning("could not stamp manifest on %s: %s", nc_path, exc)
205
+
206
+
207
+ def _prune_orphan_named_ncs(stage_dir: Path, valid_stems: set[str]) -> int:
208
+ """Delete ``{stem}.nc`` per-file outputs whose source .p is no longer present.
209
+
210
+ Sibling of :func:`_prune_orphan_profile_ncs` for products named exactly
211
+ ``{stem}.nc`` (the CTD per-file outputs, ctd.py). Those carry no ``_prof``
212
+ suffix, so the prefix-based pruner skips them and orphaned casts from a
213
+ dropped .p file would leak into the CTD combo. Prune any ``*.nc`` whose stem
214
+ is not in *valid_stems* (exact membership); a file produced this run always
215
+ matches a current stem, so live data is never dropped. Returns the count.
216
+ """
217
+ if not stage_dir.exists():
218
+ return 0
219
+ pruned = 0
220
+ for nc in stage_dir.glob("*.nc"):
221
+ if nc.stem not in valid_stems:
222
+ try:
223
+ nc.unlink()
224
+ pruned += 1
225
+ logger.info(
226
+ "pruned orphaned CTD %s (no source .p in current input set)",
227
+ nc.name,
228
+ )
229
+ except OSError as exc:
230
+ logger.warning("could not prune %s: %s", nc.name, exc)
231
+ return pruned
232
+
233
+
234
+ # ---------------------------------------------------------------------------
235
+ # Per-file processing cache (incremental re-runs)
236
+ # ---------------------------------------------------------------------------
237
+ # A file's outputs are skipped on re-run when nothing that could change them has
238
+ # changed. The cache key folds: the input file's identity (size + mtime_ns of
239
+ # the trimmed/merged .p), external inputs referenced by path (hotel/GPS files),
240
+ # and the SIGNATURE HASHES of the output dirs the file targets — which already
241
+ # encode config + the engine-code fingerprint (see perturb.config). Keying on
242
+ # the signature hash, not the reusable ``{stage}_NN`` basename, prevents a stale
243
+ # marker matching a *different* config that happens to reuse a sequence number.
244
+ # A skip ALSO requires the recorded outputs to still exist on disk: a marker is
245
+ # a claim, the files are the truth — so a pruned/deleted output forces recompute.
246
+ _CACHE_DIRNAME = ".cache"
247
+ _SIG_PREFIX = ".params_sha256_"
248
+
249
+
250
+ def _file_fingerprint(p: Path) -> dict:
251
+ """Cheap content-change proxy for a file: size + mtime binned to 2 s.
252
+
253
+ mtime is floored to a 2 s boundary so the cache survives a copy to exFAT
254
+ (2 s mtime granularity); the trade-off is that a same-size edit landing in
255
+ the same 2 s bin is not re-detected (accepted file-identity model). A missing
256
+ file yields ``{"missing": True}`` (no skip will ever match it, so it is
257
+ reprocessed — and process_file then surfaces the real error)."""
258
+ try:
259
+ st = p.stat()
260
+ except OSError:
261
+ return {"missing": True}
262
+ return {"size": st.st_size, "mtime_2s": int(st.st_mtime) // 2 * 2}
263
+
264
+
265
+ def _external_input_fingerprints(config: dict, hotel_cfg: dict | None) -> dict:
266
+ """Fingerprint hotel/GPS files referenced *by path* in the config.
267
+
268
+ Their *contents* can change ε/χ/position while the config (and thus every
269
+ stage-dir hash) stays identical, so they must enter the per-file key or an
270
+ edited hotel/GPS file would be silently ignored on re-run.
271
+ """
272
+ out: dict[str, dict] = {}
273
+ cdir = config_dir_of(config)
274
+ gps_file = merge_config("gps", config.get("gps")).get("file")
275
+ for name, path in (("hotel", (hotel_cfg or {}).get("file")), ("gps", gps_file)):
276
+ if path:
277
+ fp = Path(expand_config_dir(path, cdir))
278
+ out[name] = _file_fingerprint(fp) if fp.exists() else {"missing": True}
279
+ return out
280
+
281
+
282
+ def _stage_signature_hashes(output_dirs: dict[str, Path]) -> dict[str, str]:
283
+ """Each output dir's stored ``.params_sha256_<hash>`` value — the config +
284
+ engine identity of that stage (collision-proof, unlike its ``{stage}_NN``
285
+ basename)."""
286
+ out: dict[str, str] = {}
287
+ for stage, d in output_dirs.items():
288
+ sigs = sorted(d.glob(f"{_SIG_PREFIX}*"))
289
+ if sigs:
290
+ out[stage] = sigs[0].name[len(_SIG_PREFIX) :]
291
+ return out
292
+
293
+
294
+ def _file_cachekey(
295
+ p: Path, ext_fp: dict, stage_hashes: dict[str, str], volatile_cfg: dict | None = None
296
+ ) -> str:
297
+ payload: dict[str, Any] = {
298
+ "input": _file_fingerprint(p),
299
+ "external": ext_fp,
300
+ "stages": stage_hashes,
301
+ }
302
+ if volatile_cfg:
303
+ payload["volatile"] = volatile_cfg
304
+ return hashlib.sha256(
305
+ json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
306
+ ).hexdigest()
307
+
308
+
309
+ def _marker_path(output_root: Path, output_stem: str) -> Path:
310
+ return output_root / _CACHE_DIRNAME / f"{output_stem}.json"
311
+
312
+
313
+ def _outputs_for_stem(
314
+ output_dirs: dict[str, Path], output_stem: str, output_root: Path
315
+ ) -> list[str]:
316
+ """Relative paths of the NetCDFs a file produced: per-profile
317
+ ``{stem}_prof*.nc`` (profiles/diss/chi) and ``{stem}.nc`` (ctd)."""
318
+ rels: list[str] = []
319
+ for stage, d in output_dirs.items():
320
+ pattern = f"{output_stem}.nc" if stage == "ctd" else f"{output_stem}_prof*.nc"
321
+ rels.extend(nc.relative_to(output_root).as_posix() for nc in d.glob(pattern))
322
+ return sorted(rels)
323
+
324
+
325
+ def _list_present_outputs(output_dirs: dict[str, Path], output_root: Path) -> set[str]:
326
+ """All output-NetCDF relpaths currently on disk, gathered with ONE directory
327
+ listing per stage. The per-file output-exists check is then set membership
328
+ rather than a ``stat`` per NC — which on a slow/network mount (SMB) turns
329
+ thousands of round-trips into a handful (one ``glob`` per stage dir)."""
330
+ present: set[str] = set()
331
+ for d in output_dirs.values():
332
+ if d.exists():
333
+ present.update(nc.relative_to(output_root).as_posix() for nc in d.glob("*.nc"))
334
+ return present
335
+
336
+
337
+ def _marker_is_current(marker: Path, cachekey: str, present_outputs: set[str]) -> bool:
338
+ """True iff the marker exists, its cachekey matches, AND every recorded
339
+ output is present (membership in the pre-listed *present_outputs* set, so a
340
+ pruned/deleted/foreign output still forces recompute — no per-output stat)."""
341
+ try:
342
+ data = json.loads(marker.read_text())
343
+ except (OSError, ValueError):
344
+ return False
345
+ if not isinstance(data, dict) or data.get("cachekey") != cachekey:
346
+ return False # unreadable/malformed marker -> miss, never a crash
347
+ outputs = data.get("outputs")
348
+ if not isinstance(outputs, list):
349
+ return False
350
+ return all(rel in present_outputs for rel in outputs)
351
+
352
+
353
+ def _write_marker(
354
+ output_root: Path, output_stem: str, cachekey: str, output_dirs: dict[str, Path]
355
+ ) -> None:
356
+ """Atomically record a file's cache marker after a successful process_file."""
357
+ marker = _marker_path(output_root, output_stem)
358
+ marker.parent.mkdir(parents=True, exist_ok=True)
359
+ payload = {
360
+ "cachekey": cachekey,
361
+ "outputs": _outputs_for_stem(output_dirs, output_stem, output_root),
362
+ }
363
+ tmp = marker.with_name(f"{marker.name}.{os.getpid()}.tmp")
364
+ tmp.write_text(json.dumps(payload))
365
+ os.replace(tmp, marker) # atomic; concurrent workers write distinct stems
366
+
367
+
368
+ _TIME_SCALAR_ATTRS: dict[str, dict[str, str]] = {
369
+ "stime": {
370
+ "standard_name": "time",
371
+ "long_name": "profile start time",
372
+ "units_metadata": "leap_seconds: utc",
373
+ "axis": "T",
374
+ },
375
+ "etime": {
376
+ "standard_name": "time",
377
+ "long_name": "profile end time",
378
+ "units_metadata": "leap_seconds: utc",
379
+ },
380
+ }
381
+ _LATLON_SCALAR_ATTRS: dict[str, dict[str, str]] = {
382
+ "lat": {
383
+ "units": "degrees_north",
384
+ "standard_name": "latitude",
385
+ "long_name": "profile latitude",
386
+ "axis": "Y",
387
+ },
388
+ "lon": {
389
+ "units": "degrees_east",
390
+ "standard_name": "longitude",
391
+ "long_name": "profile longitude",
392
+ "axis": "X",
393
+ },
394
+ }
395
+
396
+
397
+ def _scalars_to_dataarrays(scalars: dict[str, float]) -> dict:
398
+ """Build xr.DataArray scalars matching what xr.open_dataset would yield.
399
+
400
+ Returns ``{name: xr.DataArray}`` for whichever of ``lat/lon/stime/etime``
401
+ are present in *scalars*. Time scalars are decoded to ``datetime64[ns]``
402
+ with the encoding hints needed for round-trip consistency, exactly
403
+ matching the open_dataset path that previously fed _copy_profile_scalars.
404
+ """
405
+ import numpy as np
406
+ import xarray as xr
407
+
408
+ out: dict = {}
409
+ for name in ("stime", "etime"):
410
+ if name in scalars:
411
+ secs = float(scalars[name])
412
+ # np.datetime64 rejects np.int64 — must be a Python int.
413
+ # round(float) already returns a Python int, no extra cast needed.
414
+ ns = round(secs * 1e9)
415
+ arr = xr.DataArray(np.datetime64(ns, "ns"), attrs=_TIME_SCALAR_ATTRS[name])
416
+ arr.encoding = {
417
+ "units": "seconds since 1970-01-01",
418
+ "calendar": "standard",
419
+ "dtype": "float64",
420
+ }
421
+ out[name] = arr
422
+ for name in ("lat", "lon"):
423
+ if name in scalars:
424
+ arr = xr.DataArray(np.float64(scalars[name]), attrs=_LATLON_SCALAR_ATTRS[name])
425
+ out[name] = arr
426
+ return out
427
+
428
+
429
+ def _copy_profile_scalars(
430
+ prof_path: str | Path,
431
+ target_ds,
432
+ scalars_cache: dict[str, dict[str, float]] | None = None,
433
+ ) -> None:
434
+ """Copy CF §9 profile scalars (lat, lon, stime, etime) onto *target_ds*.
435
+
436
+ The diss/chi per-profile NetCDFs share the same profile bounds as the
437
+ source profile NetCDF, so they can carry the same lat/lon/time scalars.
438
+ Downstream :func:`bin_by_depth` picks them up to populate per-profile
439
+ 1-D variables in the binned/combo output, which in turn unlocks ACDD
440
+ ``geospatial_lat/lon_min/max`` and ``time_coverage_start/end``.
441
+
442
+ *scalars_cache* maps profile path → raw scalar dict (as produced by
443
+ :func:`extract_profiles` with ``return_scalars=True``). When supplied
444
+ and the path is present, the scalars are reconstructed in memory and
445
+ no NetCDF re-open is needed — saves ~10 ms * 166 profiles per file.
446
+ """
447
+ if scalars_cache is not None:
448
+ scalars = scalars_cache.get(str(prof_path))
449
+ if scalars is not None:
450
+ for name, da in _scalars_to_dataarrays(scalars).items():
451
+ target_ds[name] = da
452
+ return
453
+
454
+ import xarray as xr
455
+
456
+ try:
457
+ with xr.open_dataset(prof_path) as src:
458
+ for sname in ("lat", "lon", "stime", "etime"):
459
+ if sname in src.data_vars and src[sname].ndim == 0:
460
+ target_ds[sname] = src[sname]
461
+ except Exception:
462
+ # Profile file missing scalars — older runs / external NCs, fine.
463
+ pass
464
+
465
+
466
+ def _adjust_profile_bounds(
467
+ profiles: list[tuple[int, int]],
468
+ pf,
469
+ top_trim_cfg: dict,
470
+ bottom_cfg: dict,
471
+ file_label: str,
472
+ ) -> list[tuple[int, int]]:
473
+ """Push each profile's start forward (top_trim) and/or end backward (bottom).
474
+
475
+ Both stages are skipped when their ``enable`` flag is False, so this
476
+ returns the input list unchanged in that case.
477
+
478
+ Top-trim removes propeller-wash / surface-instability samples by
479
+ measuring shear/acceleration variance binned by depth and finding the
480
+ depth where variance settles below a quantile threshold. Bottom-crash
481
+ detection scans accelerometer variance from the deepest bin upward
482
+ and flags a vibration spike near the seafloor.
483
+
484
+ Both adjustments are clipped so the returned bounds always lie inside
485
+ the original profile (no shrinking past the original end / before the
486
+ original start). Any per-profile failure is logged and that profile
487
+ is left unchanged — the rest of the file still processes.
488
+ """
489
+ import numpy as np
490
+
491
+ do_top = top_trim_cfg.get("enable", False)
492
+ do_bottom = bottom_cfg.get("enable", False)
493
+ if not (do_top or do_bottom):
494
+ return profiles
495
+
496
+ P_slow = pf.channels.get("P")
497
+ if P_slow is None:
498
+ return profiles
499
+
500
+ ratio = round(pf.fs_fast / pf.fs_slow)
501
+ P_fast = pf.channels.get("P_fast")
502
+ # Many .p files only carry P on the slow channel — interpolate to fast
503
+ # rate when needed by repeating each slow sample ``ratio`` times.
504
+ if P_fast is None or len(P_fast) != len(pf.t_fast):
505
+ P_fast = np.repeat(P_slow, ratio)
506
+ if len(P_fast) > len(pf.t_fast):
507
+ P_fast = P_fast[: len(pf.t_fast)]
508
+ elif len(P_fast) < len(pf.t_fast):
509
+ P_fast = np.concatenate([P_fast, np.full(len(pf.t_fast) - len(P_fast), P_fast[-1])])
510
+
511
+ top_kwargs = {k: v for k, v in top_trim_cfg.items() if k != "enable"}
512
+ bottom_kwargs = {k: v for k, v in bottom_cfg.items() if k != "enable"}
513
+
514
+ if do_top:
515
+ from odas_tpw.processing.top_trim import compute_trim_depth
516
+ if do_bottom:
517
+ from odas_tpw.processing.bottom import detect_bottom_crash
518
+
519
+ adjusted: list[tuple[int, int]] = []
520
+ for pi, (s_slow, e_slow) in enumerate(profiles, 1):
521
+ new_s = s_slow
522
+ new_e = e_slow
523
+ s_fast = s_slow * ratio
524
+ e_fast = (e_slow + 1) * ratio
525
+ depth_seg = P_fast[s_fast:e_fast]
526
+
527
+ if do_top:
528
+ try:
529
+ # Use only the accelerometers (Ax, Ay) to find the prop-wash /
530
+ # instrument-startup exit. On ARCTERX VMP data the accelerometers
531
+ # are the only channels that settle at the true exit (~2-6 m): the
532
+ # shear probes (sh1/sh2), inclinometers, and fall rate all stay
533
+ # "elevated" deep because they respond to *ocean* turbulence the
534
+ # instrument falls through, dragging the trim to 30-50 m. The
535
+ # accelerometers capture the mechanical entry transient, which
536
+ # ocean turbulence does not reproduce. (VMP only; MRs are trimmed
537
+ # by a separate operation.)
538
+ fast_channels: dict = {}
539
+ for name in ("Ax", "Ay"):
540
+ if name in pf.channels and pf.is_fast(name):
541
+ fast_channels[name] = pf.channels[name][s_fast:e_fast]
542
+ trim_depth = compute_trim_depth(depth_seg, fast_channels, **top_kwargs)
543
+ if trim_depth is not None and len(depth_seg) > 0:
544
+ # Find first slow index where pressure exceeds trim_depth.
545
+ P_seg = P_slow[s_slow : e_slow + 1]
546
+ above = np.where(P_seg >= trim_depth)[0]
547
+ if len(above) > 0:
548
+ push = int(above[0])
549
+ candidate = s_slow + push
550
+ if candidate < e_slow:
551
+ new_s = candidate
552
+ logger.info(
553
+ "%s prof%d: top_trim pushed start %.2f -> %.2f dbar",
554
+ file_label,
555
+ pi,
556
+ float(P_slow[s_slow]),
557
+ float(P_slow[new_s]),
558
+ )
559
+ except Exception as exc:
560
+ logger.warning("%s prof%d top_trim failed: %s", file_label, pi, exc)
561
+
562
+ if do_bottom:
563
+ try:
564
+ # Build the vibration-channels dict from whatever fast accel/gyro
565
+ # axes this instrument has. VMP-250 has Ax+Ay; a 3-axis IMU
566
+ # would also contribute Az; a different platform might supply a
567
+ # pre-aggregated ``vibration_rms`` channel.
568
+ vibration: dict = {}
569
+ for name in ("Ax", "Ay", "Az"):
570
+ ch = pf.channels.get(name)
571
+ if ch is not None and pf.is_fast(name):
572
+ vibration[name] = ch[s_fast:e_fast]
573
+ if vibration:
574
+ bottom_depth = detect_bottom_crash(
575
+ depth_seg,
576
+ vibration,
577
+ pf.fs_fast,
578
+ **bottom_kwargs,
579
+ )
580
+ if bottom_depth is not None:
581
+ P_seg = P_slow[new_s : e_slow + 1]
582
+ below = np.where(P_seg >= bottom_depth)[0]
583
+ if len(below) > 0:
584
+ candidate = new_s + int(below[0]) - 1
585
+ if candidate > new_s:
586
+ new_e = candidate
587
+ logger.info(
588
+ "%s prof%d: bottom pushed end %.2f -> %.2f dbar",
589
+ file_label,
590
+ pi,
591
+ float(P_slow[e_slow]),
592
+ float(P_slow[new_e]),
593
+ )
594
+ except Exception as exc:
595
+ logger.warning("%s prof%d bottom failed: %s", file_label, pi, exc)
596
+
597
+ adjusted.append((new_s, new_e))
598
+
599
+ return adjusted
600
+
601
+
602
+ def _apply_fom_cut(ds, fom_max: float, file_label: str, two_sided: bool = False) -> int:
603
+ """NaN per-probe per-segment values where the fom fails its cut.
604
+
605
+ With ``two_sided=False`` (epsilon): ``fom >= fom_max``.
606
+ With ``two_sided=True`` (chi): ``fom >= fom_max`` OR ``fom <= 1/fom_max``.
607
+ The chi fom is an obs/model VARIANCE RATIO that is bad far from 1.0 in either
608
+ direction, so a model that overestimates observed variance (fom << 1) must be
609
+ cut too, not only a high fom.
610
+
611
+ Operates on the (probe, time) array set returned by the diss / chi
612
+ stages. Mutates *ds* in place; returns the count of (probe, segment)
613
+ pairs that were NaN'd.
614
+
615
+ The FOM cut runs *before* ``mk_epsilon_mean`` / ``mk_chi_mean`` so
616
+ the per-probe drops are reflected in the geometric-mean and
617
+ log-sigma aggregates (a probe with bad FOM at one time step contributes
618
+ nothing to the combined estimate at that step, while the other
619
+ probe still does).
620
+
621
+ Variables masked:
622
+ - ``epsilon`` / ``chi`` (probe, time) — set NaN at flagged cells.
623
+ - ``e_N`` / ``chi_N`` (time,) — set NaN at the time slots where
624
+ probe ``N`` was flagged. ``mk_*_mean`` prefers these companions
625
+ so they need to agree.
626
+
627
+ Other per-probe metadata vars (``fom``, ``K_max``, ``kB`` etc.) are
628
+ left untouched — they're useful even at flagged segments for QC
629
+ and post-hoc diagnostics.
630
+ """
631
+ import numpy as np
632
+
633
+ if "fom" not in ds or "probe" not in ds.dims:
634
+ return 0
635
+ fom = ds["fom"].values
636
+ bad = np.isfinite(fom) & (fom >= fom_max)
637
+ if two_sided and fom_max > 0:
638
+ bad = bad | (np.isfinite(fom) & (fom <= 1.0 / fom_max))
639
+ if not bad.any():
640
+ return 0
641
+ n_bad = int(bad.sum())
642
+ n_probe = bad.shape[0]
643
+
644
+ # 2-D (probe, time) value vars
645
+ for v in ("epsilon", "chi", "epsilon_T"):
646
+ if v in ds.data_vars and ds[v].shape == fom.shape:
647
+ arr = ds[v].values.copy()
648
+ arr[bad] = np.nan
649
+ ds[v].values[...] = arr
650
+
651
+ # 1-D companions used by mk_*_mean
652
+ for i in range(n_probe):
653
+ for prefix in ("e_", "chi_"):
654
+ name = f"{prefix}{i + 1}"
655
+ if name in ds.data_vars:
656
+ arr = ds[name].values.copy()
657
+ arr[bad[i, :]] = np.nan
658
+ ds[name].values[...] = arr
659
+
660
+ logger.info(
661
+ "%s: fom_max=%g cut %d (probe,segment) cells",
662
+ file_label,
663
+ fom_max,
664
+ n_bad,
665
+ )
666
+ return n_bad
667
+
668
+
669
+ def _time_epoch_seconds(da) -> Any:
670
+ """Convert a time variable to float64 epoch seconds (ndarray).
671
+
672
+ Handles both decoded (datetime64) and undecoded (float + CF
673
+ "seconds since ..." units) representations, so arrays from different
674
+ open modes can be paired on one time base.
675
+ """
676
+ import numpy as np
677
+
678
+ vals = np.asarray(da.values)
679
+ if np.issubdtype(vals.dtype, np.datetime64):
680
+ return vals.astype("datetime64[ns]").astype("int64") / 1e9
681
+ units = str(da.attrs.get("units", ""))
682
+ if " since " in units:
683
+ # np.datetime64 rejects timezone-aware strings; these stamps are UTC.
684
+ origin = units.split(" since ", 1)[1].strip()
685
+ origin = origin.replace("Z", "").replace("+00:00", "")
686
+ origin_s = np.datetime64(origin).astype("datetime64[ns]").astype("int64") / 1e9
687
+ return vals.astype(np.float64) + origin_s
688
+ return vals.astype(np.float64)
689
+
690
+
691
+ def _compute_slow_stratification(pf, profiles, T_name, C_name, window):
692
+ """Per-cast sorted N2/dT/dz on the full slow grid (NaN outside casts).
693
+
694
+ Computes background stratification once per detected cast with
695
+ :func:`profile_stratification` on a *window*-dbar scale, then interpolates
696
+ each cast's coarse result onto the slow pressure grid. Returns
697
+ ``(N2_full, dTdz_full)`` aligned with the slow channels, or ``(None, None)``
698
+ when pressure/temperature are unavailable. Practical salinity comes from
699
+ conductivity via TEOS-10 when present. Position defaults to 0/0 (a <0.3%
700
+ gravity and negligible salinity-anomaly effect on N2).
701
+ """
702
+ import gsw
703
+ import numpy as np
704
+
705
+ from odas_tpw.processing.mixing import profile_stratification
706
+
707
+ P = pf.channels.get("P")
708
+ T = pf.channels.get(T_name)
709
+ if P is None or T is None:
710
+ return None, None
711
+ P = np.asarray(P, dtype=np.float64)
712
+ T = np.asarray(T, dtype=np.float64)
713
+ C = pf.channels.get(C_name)
714
+ # Only usable if conductivity is a slow channel aligned with P/T; a fast C
715
+ # would be silently misaligned by the slow-index slice below.
716
+ if C is not None and getattr(pf, "is_fast", lambda _n: False)(C_name):
717
+ C = None
718
+ C = np.asarray(C, dtype=np.float64) if C is not None else None
719
+
720
+ n = len(P)
721
+ N2_full = np.full(n, np.nan)
722
+ dTdz_full = np.full(n, np.nan)
723
+ for s, e in profiles:
724
+ sl = slice(s, e + 1)
725
+ P_c, T_c = P[sl], T[sl]
726
+ S_c = gsw.SP_from_C(C[sl], T_c, P_c) if C is not None else None
727
+ target_P, N2, dTdz = profile_stratification(P_c, T_c, S=S_c, window=window)
728
+ good = np.isfinite(N2)
729
+ if good.sum() >= 2:
730
+ N2_full[sl] = np.interp(P_c, target_P[good], N2[good])
731
+ good_d = np.isfinite(dTdz)
732
+ if good_d.sum() >= 2:
733
+ dTdz_full[sl] = np.interp(P_c, target_P[good_d], dTdz[good_d])
734
+ return N2_full, dTdz_full
735
+
736
+
737
+ def _window_stratification_for_profile(
738
+ prof_path, win_t, half_w, T_name: str, C_name: str, file_label: str
739
+ ):
740
+ """Sorted N2/dTdz at window times *win_t* from a profile's CTD channels.
741
+
742
+ Returns ``(N2, dTdz, sal_note)`` evaluated on the *win_t* grid, or ``None``
743
+ when the profile cannot be read or lacks t_slow/P/T. Practical salinity
744
+ comes from the profile's own conductivity via TEOS-10 when present, else
745
+ 35 PSU is assumed (recorded in *sal_note*). Shared by the diss and chi
746
+ products, each evaluating stratification at its own window scale (the diss
747
+ path passes diss_length_seconds/2, the chi path 0.5*chi diss_length/fs), so
748
+ the two coincide only when their diss_length/fft_length settings match.
749
+ """
750
+ import gsw
751
+ import numpy as np
752
+ import xarray as xr
753
+
754
+ from odas_tpw.processing.mixing import sorted_stratification
755
+
756
+ try:
757
+ prof = xr.open_dataset(prof_path, decode_times=False)
758
+ except Exception as exc:
759
+ logger.info(
760
+ "stratification skipped for %s: cannot read profile (%s)",
761
+ file_label,
762
+ exc,
763
+ )
764
+ return None
765
+ with prof:
766
+ if any(v not in prof for v in ("t_slow", "P", T_name)):
767
+ logger.info(
768
+ "stratification skipped for %s: profile lacks t_slow/P/%s",
769
+ file_label,
770
+ T_name,
771
+ )
772
+ return None
773
+ t_slow = _time_epoch_seconds(prof["t_slow"])
774
+ P = prof["P"].values.astype(np.float64)
775
+ T = prof[T_name].values.astype(np.float64)
776
+ lat = float(prof["lat"].values) if "lat" in prof else np.nan
777
+ lon = float(prof["lon"].values) if "lon" in prof else np.nan
778
+ lat = 0.0 if not np.isfinite(lat) else lat
779
+ lon = 0.0 if not np.isfinite(lon) else lon
780
+ if C_name in prof:
781
+ C = prof[C_name].values.astype(np.float64)
782
+ S = gsw.SP_from_C(C, T, P)
783
+ sal_note = f"practical salinity from {C_name}/{T_name}/P (TEOS-10)"
784
+ else:
785
+ S = None
786
+ sal_note = (
787
+ "salinity assumed 35 PSU (no conductivity channel); N2 "
788
+ "reflects temperature stratification only"
789
+ )
790
+ strat = sorted_stratification(win_t, half_w, t_slow, P, T, S=S, lat=lat, lon=lon)
791
+ return strat.N2, strat.dTdz, sal_note
792
+
793
+
794
+ def _attach_window_stratification(ds, prof_path, half_w, T_name, C_name, file_label, scale_name):
795
+ """Attach sorted N2/dTdz to a window-grid dataset (e.g. diss). No-op on failure.
796
+
797
+ *half_w* is the window half-width [s]; *scale_name* names the window in the
798
+ variable comments (e.g. "dissipation").
799
+ """
800
+ import xarray as xr
801
+
802
+ if "t" not in ds:
803
+ return
804
+ out = _window_stratification_for_profile(
805
+ prof_path, _time_epoch_seconds(ds["t"]), half_w, T_name, C_name, file_label
806
+ )
807
+ if out is None:
808
+ return
809
+ N2, dTdz, sal_note = out
810
+ ds["N2"] = xr.DataArray(
811
+ N2,
812
+ dims=["time"],
813
+ attrs={
814
+ "units": "s-2",
815
+ "long_name": "buoyancy frequency squared (window scale)",
816
+ "comment": (
817
+ "TEOS-10 (gsw.Nsquared) between the shallow- and deep-half means "
818
+ f"of each {scale_name} window after Thorpe-sorting to a stable "
819
+ f"profile (overturns removed); {sal_note}."
820
+ ),
821
+ },
822
+ )
823
+ ds["dTdz"] = xr.DataArray(
824
+ dTdz,
825
+ dims=["time"],
826
+ attrs={
827
+ "units": "K m-1",
828
+ "long_name": "background temperature gradient (positive down)",
829
+ "comment": (
830
+ "Least-squares slope of the Thorpe-sorted in-situ temperature vs "
831
+ f"depth over each {scale_name} window."
832
+ ),
833
+ },
834
+ )
835
+
836
+
837
+ def _add_mixing_quantities(
838
+ chi_ds,
839
+ diss_ds,
840
+ prof_path,
841
+ T_name: str = "JAC_T",
842
+ C_name: str = "JAC_C",
843
+ file_label: str = "",
844
+ epsilon_provenance: str = "",
845
+ ):
846
+ """Append stratification and (when available) mixing quantities to chi.
847
+
848
+ Always computes window-scale N2 and dT/dz from the profile's own slow CTD
849
+ channels via the Thorpe-sorted (adiabatically leveled) method — with
850
+ practical salinity from C/T/P via TEOS-10 when the conductivity channel
851
+ exists, so N2 is fully constrained (unlike the rsi path, which may assume
852
+ 35 PSU). N2 and dT/dz depend only on the CTD profile, so they are written
853
+ regardless of epsilon/chi.
854
+
855
+ When epsilonMean (from the matching diss dataset) and chiMean are both
856
+ present, additionally computes K_T (Osborn-Cox), Gamma (Oakey 1982), and
857
+ K_rho (Osborn 1980, Gamma_0 = 0.2), pairing epsilonMean onto the chi grid.
858
+
859
+ Mutates and returns *chi_ds*. Returns it unchanged (with a log line) only
860
+ when the profile lacks the t_slow/P/T channels or chi window attributes
861
+ needed to compute stratification at all.
862
+ """
863
+ import numpy as np
864
+ import xarray as xr
865
+
866
+ from odas_tpw.processing.mixing import mixing_coefficients, pair_nearest
867
+
868
+ # N2/dTdz are stratification quantities — they need only the CTD profile,
869
+ # not epsilon or chi — so they are always computed and attached. Only the
870
+ # derived coefficients (K_T, Gamma, K_rho) require epsilon and/or chi.
871
+ have_mix = diss_ds is not None and "epsilonMean" in diss_ds and "chiMean" in chi_ds
872
+
873
+ try:
874
+ chi_t = _time_epoch_seconds(chi_ds["t"])
875
+ half_w = 0.5 * float(chi_ds.attrs["diss_length"]) / float(chi_ds.attrs["fs_fast"])
876
+ except (KeyError, TypeError, ValueError):
877
+ logger.info("mixing quantities skipped for %s: no chi time/window attrs", file_label)
878
+ return chi_ds
879
+
880
+ strat_out = _window_stratification_for_profile(
881
+ prof_path, chi_t, half_w, T_name, C_name, file_label
882
+ )
883
+ if strat_out is None:
884
+ return chi_ds
885
+ N2, dTdz, sal_note = strat_out
886
+
887
+ var_specs = {
888
+ "N2": (
889
+ N2,
890
+ {
891
+ "units": "s-2",
892
+ "long_name": "buoyancy frequency squared (window scale)",
893
+ "comment": (
894
+ "TEOS-10 (gsw.Nsquared) between the shallow- and deep-half "
895
+ "means of each chi window after Thorpe-sorting to a stable "
896
+ f"profile (overturns removed); {sal_note}."
897
+ ),
898
+ },
899
+ ),
900
+ "dTdz": (
901
+ dTdz,
902
+ {
903
+ "units": "K m-1",
904
+ "long_name": "background temperature gradient (positive down)",
905
+ "comment": (
906
+ "Least-squares slope of the Thorpe-sorted in-situ "
907
+ "temperature vs depth over each chi window."
908
+ ),
909
+ },
910
+ ),
911
+ }
912
+
913
+ # K_T / Gamma / K_rho need epsilon and/or chi; only add them when available.
914
+ if have_mix:
915
+ eps_on_chi = pair_nearest(
916
+ _time_epoch_seconds(diss_ds["t"]), diss_ds["epsilonMean"].values, chi_t
917
+ )
918
+ mix = mixing_coefficients(eps_on_chi, chi_ds["chiMean"].values, N2, dTdz)
919
+ var_specs.update(
920
+ {
921
+ # The exact epsilon that entered Gamma/K_rho, stored for
922
+ # traceability: consumers (e.g. perturb-plot gamma-scaling)
923
+ # can recompute ratio quantities from the same pairing the
924
+ # pipeline used instead of re-pairing against the diss
925
+ # product (and possibly choosing different windows).
926
+ "epsilon_paired": (
927
+ eps_on_chi,
928
+ {
929
+ "units": "W/kg",
930
+ "long_name": "epsilonMean paired onto the chi window grid",
931
+ "comment": (
932
+ "Nearest-window epsilonMean from the matching diss "
933
+ "dataset (processing.mixing.pair_nearest); the "
934
+ "epsilon used in this dataset's Gamma and K_rho. "
935
+ "NaN where no diss window paired." + epsilon_provenance
936
+ ),
937
+ },
938
+ ),
939
+ "K_T": (
940
+ mix.K_T,
941
+ {
942
+ "units": "m2 s-1",
943
+ "long_name": "Osborn-Cox eddy diffusivity of heat",
944
+ "comment": (
945
+ "K_T = chi / (2*(dT/dz)^2), Osborn & Cox (1972), "
946
+ "doi:10.1080/03091927208236085. NaN where "
947
+ "|dT/dz| < 1e-4 K/m (well-mixed)."
948
+ ),
949
+ },
950
+ ),
951
+ "Gamma": (
952
+ mix.Gamma,
953
+ {
954
+ "units": "1",
955
+ "long_name": "mixing coefficient (measured)",
956
+ "comment": (
957
+ "Gamma = N2*chi / (2*epsilon*(dT/dz)^2), Oakey (1982), "
958
+ "doi:10.1175/1520-0485(1982)012<0256:DOTROD>2.0.CO;2. "
959
+ "epsilonMean paired from the nearest dissipation window. "
960
+ "NaN where N2 < 1e-9 s-2 or |dT/dz| < 1e-4 K/m. "
961
+ "Canonical value ~0.2." + epsilon_provenance
962
+ ),
963
+ },
964
+ ),
965
+ "K_rho": (
966
+ mix.K_rho,
967
+ {
968
+ "units": "m2 s-1",
969
+ "long_name": "Osborn diapycnal diffusivity (Gamma_0 = 0.2)",
970
+ "comment": (
971
+ "K_rho = 0.2*epsilon/N2, Osborn (1980), "
972
+ "doi:10.1175/1520-0485(1980)010<0083:EOTLRO>2.0.CO;2. "
973
+ "Compare with K_T: agreement implies the measured Gamma "
974
+ "is near the canonical 0.2. NaN where N2 < 1e-9 s-2 or "
975
+ "K_rho > 10 m2 s-1 (physically implausible diffusivity: "
976
+ "the unbounded near-floor-N2 artifact, or contaminated "
977
+ "near-surface windows where epsilon is itself spurious)."
978
+ + epsilon_provenance
979
+ ),
980
+ },
981
+ ),
982
+ }
983
+ )
984
+
985
+ for name, (arr, attrs) in var_specs.items():
986
+ chi_ds[name] = xr.DataArray(arr, dims=["time"], attrs=attrs)
987
+ if have_mix:
988
+ n_gamma = int(np.sum(np.isfinite(var_specs["Gamma"][0])))
989
+ logger.info("mixing quantities for %s: %d valid Gamma estimates", file_label, n_gamma)
990
+ else:
991
+ n_n2 = int(np.sum(np.isfinite(N2)))
992
+ logger.info(
993
+ "stratification for %s: %d valid N2 windows (no epsilon/chi coefficients)",
994
+ file_label,
995
+ n_n2,
996
+ )
997
+ return chi_ds
998
+
999
+
1000
+ def _nan_excluded_probes(ds, excluded_probes: list[str], file_label: str) -> None:
1001
+ """Set per-probe epsilon to NaN for shear probes named in *excluded_probes*.
1002
+
1003
+ Mutates *ds* in place, so the subsequent mk_epsilon_mean call sees NaN for
1004
+ those probes and excludes them from the geometric mean (and, via the
1005
+ fallback path in chi, from chi Method 1). Both the 2-D ``epsilon``
1006
+ (probe x time) and any pre-existing 1-D ``e_N`` companion variables are
1007
+ masked. Unknown probe names are warned and ignored — typos in the
1008
+ config should be visible but not fatal.
1009
+ """
1010
+ import numpy as np
1011
+
1012
+ if "epsilon" not in ds or "probe" not in ds.dims:
1013
+ return
1014
+ probe_names = [str(p) for p in ds["probe"].values]
1015
+ for probe in excluded_probes:
1016
+ if probe not in probe_names:
1017
+ logger.warning(
1018
+ "instruments override for %s: probe %r not present (have %s)",
1019
+ file_label,
1020
+ probe,
1021
+ probe_names,
1022
+ )
1023
+ continue
1024
+ idx = probe_names.index(probe)
1025
+ ds["epsilon"].values[idx, :] = np.nan
1026
+ # mk_epsilon_mean prefers e_N variables when present; mask them too
1027
+ # so any caller that round-trips through e_N sees the same result.
1028
+ e_name = f"e_{idx + 1}"
1029
+ if e_name in ds.data_vars:
1030
+ ds[e_name].values[:] = np.nan
1031
+
1032
+
1033
+ def _profile_practical_salinity(prof_path, T_name: str, C_name: str):
1034
+ """Slow-time practical salinity for a profile NetCDF, or None.
1035
+
1036
+ Derives per-slow-sample practical salinity [PSU] from the profile's own
1037
+ conductivity, in-situ temperature, and pressure via TEOS-10
1038
+ (``gsw.SP_from_C`` — the same conversion used for the derived mixing
1039
+ quantities), returning ``None`` when any of the C/T/P channels are
1040
+ absent. Used to resolve ``chi.salinity: measured`` to an array the chi
1041
+ viscosity path consumes instead of the fixed 35-PSU default.
1042
+ """
1043
+ import gsw
1044
+ import numpy as np
1045
+ import xarray as xr
1046
+
1047
+ with xr.open_dataset(prof_path, decode_times=False) as ds:
1048
+ if any(v not in ds for v in ("P", T_name, C_name)):
1049
+ return None
1050
+ P = ds["P"].values.astype(np.float64)
1051
+ T = ds[T_name].values.astype(np.float64)
1052
+ C = ds[C_name].values.astype(np.float64)
1053
+ return gsw.SP_from_C(C, T, P)
1054
+
1055
+
1056
+ _SEAWATER_UNITS = {
1057
+ "SP": "1",
1058
+ "SA": "g/kg",
1059
+ "CT": "degree_Celsius",
1060
+ "sigma0": "kg m-3",
1061
+ "rho": "kg m-3",
1062
+ }
1063
+
1064
+
1065
+ def _inject_seawater_properties(pf, gps, T_name: str, C_name: str) -> tuple[str, ...]:
1066
+ """Compute SP/SA/CT/sigma0/rho on the slow grid and inject as pf channels.
1067
+
1068
+ For the profile product: the full-rate slow ``T_name``/``C_name``/``P`` feed
1069
+ TEOS-10 (:func:`add_seawater_properties`) *before* depth-binning, so the
1070
+ nonlinear conversions see unaveraged inputs; binning then averages the
1071
+ derived properties like any other channel. ``T_name`` is the in-situ CTD
1072
+ reference (e.g. JAC_T), untouched by FP07 recalibration. Per-slow-sample
1073
+ lat/lon come from *gps* (SA's small position dependence). Long names /
1074
+ canonical units are layered from ``COMBO_SCHEMA`` at write time, so only the
1075
+ minimal ``units``/``type`` channel_info is set here.
1076
+
1077
+ Returns the names injected, or ``()`` when the C/T/P slow channels are
1078
+ missing or length-mismatched (e.g. a fast conductivity).
1079
+ """
1080
+ import numpy as np
1081
+
1082
+ if not all(n in pf.channels for n in (T_name, C_name, "P")):
1083
+ return ()
1084
+ T = np.asarray(pf.channels[T_name], dtype=np.float64)
1085
+ C = np.asarray(pf.channels[C_name], dtype=np.float64)
1086
+ P = np.asarray(pf.channels["P"], dtype=np.float64)
1087
+ n_slow = len(pf.t_slow)
1088
+ if not (len(T) == len(C) == len(P) == n_slow):
1089
+ return ()
1090
+
1091
+ from odas_tpw.perturb.seawater import add_seawater_properties
1092
+
1093
+ epoch = pf.start_time.timestamp() + np.asarray(pf.t_slow, dtype=np.float64)
1094
+ sw = add_seawater_properties(T, C, P, gps.lat(epoch), gps.lon(epoch))
1095
+ for name in ("SP", "SA", "CT", "sigma0", "rho"):
1096
+ pf.channels[name] = sw[name]
1097
+ pf.channel_info[name] = {
1098
+ "units": _SEAWATER_UNITS[name],
1099
+ "type": "derived",
1100
+ "name": name,
1101
+ }
1102
+ return ("SP", "SA", "CT", "sigma0", "rho")
1103
+
1104
+
1105
+ def process_file(
1106
+ p_path: Path,
1107
+ config: dict,
1108
+ gps,
1109
+ output_dirs: dict,
1110
+ hotel_data=None,
1111
+ hotel_cfg=None,
1112
+ instrument_key: str | None = None,
1113
+ output_stem: str | None = None,
1114
+ ) -> dict:
1115
+ """Process a single .p file through the full enhancement chain.
1116
+
1117
+ Parameters
1118
+ ----------
1119
+ p_path : Path
1120
+ Path to the .p file.
1121
+ config : dict
1122
+ Full merged config (all sections).
1123
+ gps : GPSProvider
1124
+ GPS provider.
1125
+ output_dirs : dict
1126
+ Maps stage names to output directory Paths.
1127
+ instrument_key : str, optional
1128
+ Key into ``config["instruments"]`` for per-instrument overrides
1129
+ (e.g. ``"SN465"``). Falls back to ``p_path.parent.name`` if not
1130
+ supplied — but the caller normally passes this explicitly because
1131
+ the trim stage flattens the original ``<root>/<SN>/<file>.p``
1132
+ layout into a single ``trimmed/`` directory.
1133
+ output_stem : str, optional
1134
+ Unique stem for per-file outputs. Defaults to ``p_path.stem``.
1135
+
1136
+ Returns
1137
+ -------
1138
+ dict with output paths per stage.
1139
+ """
1140
+ from odas_tpw.perturb.fp07_cal import fp07_calibrate
1141
+ from odas_tpw.perturb.qc_gate import apply_qc_to_dataset
1142
+ from odas_tpw.processing.chi_combine import mk_chi_mean
1143
+ from odas_tpw.processing.ct_align import ct_align
1144
+ from odas_tpw.processing.epsilon_combine import mk_epsilon_mean
1145
+ from odas_tpw.rsi.dissipation import _compute_epsilon
1146
+ from odas_tpw.rsi.p_file import PFile
1147
+ from odas_tpw.rsi.profile import _smooth_fall_rate, get_profiles
1148
+
1149
+ result: dict[str, Any] = {
1150
+ "source": str(p_path),
1151
+ "profiles": [],
1152
+ "diss": [],
1153
+ "chi": [],
1154
+ }
1155
+
1156
+ try:
1157
+ pf = PFile(p_path)
1158
+ except Exception as exc:
1159
+ logger.error("loading %s: %s", p_path.name, exc)
1160
+ result.setdefault("errors", []).append(f"load: {exc}")
1161
+ return result
1162
+
1163
+ # ---- Hotel data injection ----
1164
+ if hotel_data is not None:
1165
+ from odas_tpw.perturb.hotel import merge_hotel_into_pfile
1166
+
1167
+ merge_hotel_into_pfile(hotel_data, pf, hotel_cfg or {})
1168
+
1169
+ # ---- Speed channel (downstream of .p load + hotel merge) ----
1170
+ # Inject ``speed_fast`` and ``W_slow`` so extract_profiles writes them
1171
+ # to every per-profile NetCDF, and downstream prepare_profiles uses
1172
+ # them instead of recomputing from P. Default method is "pressure",
1173
+ # which reproduces the historical VMP behavior exactly.
1174
+ speed_cfg = merge_config("speed", config.get("speed"))
1175
+ if "P" in pf.channels:
1176
+ try:
1177
+ from odas_tpw.perturb.speed import compute_speed_for_pfile
1178
+
1179
+ vehicle = pf.config.get("instrument_info", {}).get("vehicle", "").lower()
1180
+ speed_fast, W_slow = compute_speed_for_pfile(pf, speed_cfg, vehicle)
1181
+
1182
+ pf.channels["speed_fast"] = speed_fast
1183
+ pf._fast_channels.add("speed_fast")
1184
+ pf.channel_info["speed_fast"] = {
1185
+ "units": "m s-1",
1186
+ "type": "computed",
1187
+ "name": "speed_fast",
1188
+ }
1189
+ pf.channels["W_slow"] = W_slow
1190
+ pf.channel_info["W_slow"] = {
1191
+ "units": "dbar s-1",
1192
+ "type": "computed",
1193
+ "name": "W_slow",
1194
+ "long_name": "profiling rate (smoothed |dP/dt|)",
1195
+ }
1196
+ except Exception as exc:
1197
+ logger.warning(
1198
+ "speed channel computation failed for %s (method=%s): %s",
1199
+ p_path.name,
1200
+ speed_cfg.get("method", "pressure"),
1201
+ exc,
1202
+ )
1203
+
1204
+ # ---- Internal QC rules (after hotel + speed so all channels exist) ----
1205
+ qc_rules_cfg = (config.get("qc") or {}).get("rules") or {}
1206
+ if qc_rules_cfg:
1207
+ try:
1208
+ from odas_tpw.perturb.qc_rules import (
1209
+ evaluate_rules,
1210
+ register_rule_channels,
1211
+ )
1212
+
1213
+ evaluated = evaluate_rules(pf, qc_rules_cfg)
1214
+ register_rule_channels(pf, evaluated, qc_rules_cfg)
1215
+ except Exception as exc:
1216
+ logger.warning(
1217
+ "qc.rules evaluation failed for %s: %s",
1218
+ p_path.name,
1219
+ exc,
1220
+ )
1221
+
1222
+ # Per-stage log files: each ``with stage_log(...)`` adds a FileHandler
1223
+ # for that stage's output dir so e.g. all CTD-binning records for ``a.p``
1224
+ # land in ``ctd_NN/a.log`` *and* the worker/run logs. basename = stem so
1225
+ # ``a.p`` produces ``a.log``.
1226
+ output_stem = output_stem or p_path.stem
1227
+ log_basename = output_stem
1228
+
1229
+ profiles_cfg = config.get("profiles", {})
1230
+ fp07_cfg = config.get("fp07", {})
1231
+ ct_cfg = config.get("ct", {})
1232
+
1233
+ # ---- Profile detection + CT alignment ----
1234
+ # Runs BEFORE the CTD fork so the CTD product's salinity/density are
1235
+ # computed from time-aligned conductivity whenever profiles can be
1236
+ # detected (alignment needs the per-profile lag estimates). Files
1237
+ # with no pressure or no detectable profiles still get a CTD product
1238
+ # from the unaligned conductivity.
1239
+ profiles: list[tuple[int, int]] = []
1240
+ direction = "down" # resolved below when a pressure channel is present
1241
+ P_slow = pf.channels.get("P")
1242
+ with stage_log(output_dirs.get("profiles"), log_basename):
1243
+ if P_slow is None:
1244
+ logger.warning("No pressure channel in %s", p_path.name)
1245
+ elif len(P_slow) != len(pf.t_slow):
1246
+ # Defense-in-depth for U5-1: profile detection drives P on the SLOW
1247
+ # grid (fs_slow), so a P whose length != the slow-grid length — e.g.
1248
+ # a hotel-injected fast-rate P — would make dP/dt (hence W) wrong and
1249
+ # silently drop the file. Skip loudly instead. (merge_hotel_into_pfile
1250
+ # normally prevents this upstream by refusing to clobber native P.)
1251
+ logger.warning(
1252
+ "Pressure channel length %d != slow-grid length %d in %s; "
1253
+ "skipping profile detection (a fast-rate P cannot drive it)",
1254
+ len(P_slow),
1255
+ len(pf.t_slow),
1256
+ p_path.name,
1257
+ )
1258
+ else:
1259
+ # Resolve "auto" → vehicle default (e.g. slocum_glider → "glide").
1260
+ # ``scor160.profile.get_profiles`` doesn't know "auto" itself, so
1261
+ # without this it silently falls through to the "down" branch and
1262
+ # we lose every up-profile -- on a glider with MR-on-during-climb
1263
+ # only, that drops *all* the real flight data.
1264
+ from odas_tpw.rsi.vehicle import resolve_direction, resolve_tau
1265
+
1266
+ vehicle = pf.config.get("instrument_info", {}).get("vehicle", "").lower()
1267
+ # Smooth the detection fall rate with the VEHICLE tau, matching ODAS
1268
+ # (odas_p2mat.m). The default 1.5 s is correct for a VMP but 2-40x
1269
+ # too fast for gliders/floats (slocum 3.0 s, argo 60.0 s), which made
1270
+ # W noisy at the W_min threshold and fragmented profile boundaries.
1271
+ W = _smooth_fall_rate(P_slow, pf.fs_slow, tau=resolve_tau(vehicle))
1272
+ direction = resolve_direction(
1273
+ profiles_cfg.get("direction", "auto"),
1274
+ vehicle,
1275
+ )
1276
+ profiles = get_profiles(
1277
+ P_slow,
1278
+ W,
1279
+ pf.fs_slow,
1280
+ P_min=profiles_cfg.get("P_min", 0.5),
1281
+ W_min=profiles_cfg.get("W_min", 0.3),
1282
+ direction=direction,
1283
+ min_duration=profiles_cfg.get("min_duration", 7.0),
1284
+ )
1285
+ if not profiles:
1286
+ logger.warning("No profiles in %s", p_path.name)
1287
+
1288
+ # CT alignment (before the CTD fork; needs detected profiles)
1289
+ if profiles and ct_cfg.get("align", True):
1290
+ T_name = ct_cfg.get("T_name", "JAC_T")
1291
+ C_name = ct_cfg.get("C_name", "JAC_C")
1292
+ if T_name in pf.channels and C_name in pf.channels:
1293
+ try:
1294
+ C_aligned, _lag = ct_align(
1295
+ pf.channels[T_name],
1296
+ pf.channels[C_name],
1297
+ pf.fs_slow,
1298
+ profiles,
1299
+ )
1300
+ pf.channels[C_name] = C_aligned
1301
+ except Exception as exc:
1302
+ logger.warning("CT align failed for %s: %s", p_path.name, exc)
1303
+
1304
+ # ---- Background stratification (N2, dT/dz) on the slow grid ----
1305
+ # Injected as slow pf channels once, before the CTD and profile forks, so
1306
+ # extract_profiles writes them into every per-profile NetCDF and
1307
+ # ctd_bin_file bins them into the CTD product — all independent of eps/chi.
1308
+ strat_cfg = merge_config("stratification", config.get("stratification"))
1309
+ if profiles and bool(strat_cfg.get("enable", True)):
1310
+ try:
1311
+ N2_full, dTdz_full = _compute_slow_stratification(
1312
+ pf,
1313
+ profiles,
1314
+ ct_cfg.get("T_name", "JAC_T"),
1315
+ ct_cfg.get("C_name", "JAC_C"),
1316
+ float(strat_cfg.get("window", 2.0)),
1317
+ )
1318
+ if N2_full is not None:
1319
+ win = float(strat_cfg.get("window", 2.0))
1320
+ pf.channels["N2"] = N2_full
1321
+ pf.channel_info["N2"] = {
1322
+ "units": "s-2",
1323
+ "type": "derived",
1324
+ "name": "N2",
1325
+ "long_name": (
1326
+ f"buoyancy frequency squared (background, {win:g}-dbar Thorpe-sorted)"
1327
+ ),
1328
+ "comment": (
1329
+ "TEOS-10 N2 from the profile's own C/T/P over a "
1330
+ f"{win:g}-dbar pressure window, Thorpe-sorted to a stable "
1331
+ "profile. Background (profile/CTD) scale — distinct from "
1332
+ "the dissipation/chi-window N2 in the diss/chi products."
1333
+ ),
1334
+ }
1335
+ pf.channels["dTdz"] = dTdz_full
1336
+ pf.channel_info["dTdz"] = {
1337
+ "units": "K m-1",
1338
+ "type": "derived",
1339
+ "name": "dTdz",
1340
+ "long_name": (
1341
+ "background temperature gradient (positive down, "
1342
+ f"{win:g}-dbar Thorpe-sorted)"
1343
+ ),
1344
+ "comment": (
1345
+ "Least-squares slope of the Thorpe-sorted in-situ "
1346
+ f"temperature vs depth over a {win:g}-dbar pressure window."
1347
+ ),
1348
+ }
1349
+ except Exception as exc:
1350
+ logger.warning("stratification failed for %s: %s", p_path.name, exc)
1351
+
1352
+ # ---- CTD fork (full file, both up and down) ----
1353
+ ctd_cfg = config.get("ctd", {})
1354
+ if ctd_cfg.get("enable", True) and "ctd" in output_dirs:
1355
+ with stage_log(output_dirs.get("ctd"), log_basename):
1356
+ try:
1357
+ from odas_tpw.perturb.ctd import ctd_bin_file
1358
+
1359
+ ctd_bin_file(
1360
+ pf,
1361
+ gps,
1362
+ output_dirs["ctd"],
1363
+ bin_width=ctd_cfg.get("bin_width", 0.5),
1364
+ T_name=ctd_cfg.get("T_name", "JAC_T"),
1365
+ C_name=ctd_cfg.get("C_name", "JAC_C"),
1366
+ variables=ctd_cfg.get("variables"),
1367
+ method=ctd_cfg.get("method", "mean"),
1368
+ diagnostics=ctd_cfg.get("diagnostics", False),
1369
+ output_stem=output_stem,
1370
+ profiles=profiles,
1371
+ direction=direction,
1372
+ )
1373
+ except Exception as exc:
1374
+ logger.error("CTD binning %s: %s", p_path.name, exc)
1375
+ result.setdefault("errors", []).append(f"ctd: {exc}")
1376
+
1377
+ if P_slow is None or not profiles:
1378
+ return result
1379
+
1380
+ # ---- Seawater properties (SP/SA/CT/sigma0/rho) for the profile product ----
1381
+ # Injected AFTER the CTD fork, which derives its own from post-bin T/C/P;
1382
+ # here they are computed on the full slow-rate T/C/P and depth-binned with
1383
+ # the profile, so both products stay self-consistent without conflicting.
1384
+ try:
1385
+ _inject_seawater_properties(
1386
+ pf, gps, ct_cfg.get("T_name", "JAC_T"), ct_cfg.get("C_name", "JAC_C")
1387
+ )
1388
+ except Exception as exc:
1389
+ logger.warning("seawater properties failed for %s: %s", p_path.name, exc)
1390
+
1391
+ # ---- Profile fork ----
1392
+ # Wraps profile-bound adjustment, FP07 calibration, and the
1393
+ # extract_profiles write. These all conceptually feed profiles_NN/.
1394
+ with stage_log(output_dirs.get("profiles"), log_basename):
1395
+ # Adjust profile bounds (top-trim removes prop-wash, bottom detects seafloor crash)
1396
+ profiles = _adjust_profile_bounds(
1397
+ profiles,
1398
+ pf,
1399
+ config.get("top_trim", {}),
1400
+ config.get("bottom", {}),
1401
+ p_path.name,
1402
+ )
1403
+
1404
+ # FP07 calibration
1405
+ if fp07_cfg.get("calibrate", True):
1406
+ try:
1407
+ cal_result = fp07_calibrate(
1408
+ pf,
1409
+ profiles,
1410
+ reference=fp07_cfg.get("reference", "JAC_T"),
1411
+ order=fp07_cfg.get("order", 2),
1412
+ max_lag_seconds=fp07_cfg.get("max_lag_seconds", 10.0),
1413
+ must_be_negative=fp07_cfg.get("must_be_negative", True),
1414
+ )
1415
+ # Apply calibrated temperatures back to pf.channels —
1416
+ # both the base channels (T1, T2) and the recalibrated
1417
+ # deconvolved fast channels (T1_dT1, T2_dT2) that the
1418
+ # chi pipeline consumes. Without the latter, chi keeps
1419
+ # the factory calibration and scales with the square of
1420
+ # the calibration slope error.
1421
+ # Record the in-situ calibration as per-channel provenance so
1422
+ # the products are self-describing (and plots can label it).
1423
+ # Only channels actually recalibrated are tagged — if the
1424
+ # reference was missing, cal_result is empty and the channels
1425
+ # keep the factory calibration with no tag.
1426
+ cal_tag = (
1427
+ f"in-situ (Steinhart-Hart order {fp07_cfg.get('order', 2)} "
1428
+ f"vs {fp07_cfg.get('reference', 'JAC_T')})"
1429
+ )
1430
+ for ch_name, cal_data in cal_result.get("channels", {}).items():
1431
+ pf.channels[ch_name] = cal_data
1432
+ pf.channel_info.setdefault(ch_name, {})["calibration"] = cal_tag
1433
+ for ch_name, cal_data in cal_result.get("fast_channels", {}).items():
1434
+ pf.channels[ch_name] = cal_data
1435
+ pf.channel_info.setdefault(ch_name, {})["calibration"] = cal_tag
1436
+ except Exception as exc:
1437
+ logger.warning("FP07 cal failed for %s: %s", p_path.name, exc)
1438
+
1439
+ # Write per-profile NetCDFs
1440
+ from odas_tpw.rsi.profile import extract_profiles
1441
+
1442
+ prof_scalars_cache: dict[str, dict[str, float]] = {}
1443
+ if "profiles" in output_dirs:
1444
+ try:
1445
+ prof_paths, prof_scalars = extract_profiles(
1446
+ pf,
1447
+ output_dirs["profiles"],
1448
+ profiles=profiles,
1449
+ gps=gps,
1450
+ return_scalars=True,
1451
+ output_stem=output_stem,
1452
+ )
1453
+ result["profiles"] = [str(p) for p in prof_paths]
1454
+ prof_scalars_cache = {str(p): s for p, s in zip(prof_paths, prof_scalars)}
1455
+ except Exception as exc:
1456
+ logger.error("extracting profiles %s: %s", p_path.name, exc)
1457
+ result.setdefault("errors", []).append(f"profiles: {exc}")
1458
+ return result
1459
+
1460
+ # Per-profile dissipation. Resolve via merge_config so the DEFAULTS (e.g.
1461
+ # epsilon.fft_length=256) — and any None values, which merge_config drops —
1462
+ # are the values actually handed to _compute_epsilon/_compute_chi and the
1463
+ # diss_length_seconds below. config.get() raw would let an omitted/null
1464
+ # fft_length fall to _compute_epsilon's own 1024 default (mis-sizing windows
1465
+ # and making the diss-dir provenance hash, which IS merge_config-based, lie).
1466
+ eps_cfg = merge_config("epsilon", config.get("epsilon"))
1467
+ # Durations -> integer sample counts at THIS file's sampling rate, so one
1468
+ # duration-based config serves instruments at 512 Hz and 1-2 kHz alike.
1469
+ eps_cfg = resolve_window_config(eps_cfg, float(pf.fs_fast), section="epsilon")
1470
+ inst_lookup = instrument_key if instrument_key is not None else p_path.parent.name
1471
+ instrument_cfg = config.get("instruments", {}).get(inst_lookup, {})
1472
+ excluded_probes = list(instrument_cfg.get("exclude_shear_probes", []))
1473
+ chi_cfg = merge_config("chi", config.get("chi"))
1474
+ chi_enabled = bool(chi_cfg.get("enable", False)) and "chi" in output_dirs
1475
+ if chi_enabled:
1476
+ # Only resolve windows for a stage that will run: a mis-set chi
1477
+ # window in a chi-disabled config must not abort epsilon processing.
1478
+ chi_cfg = resolve_window_config(chi_cfg, float(pf.fs_fast), section="chi")
1479
+
1480
+ # Per-profile channels cache: filled here, consumed by the chi loop
1481
+ # below. Holds ~1 MB / profile of fast-time numpy arrays — bounded by
1482
+ # the per-file profile count (~80 for SN465), so worker-RSS impact is
1483
+ # well under 100 MB. Skipped entirely when chi is disabled.
1484
+ prof_data_cache: dict[str, dict[str, Any]] = {}
1485
+ diss_by_profile: dict[str, str] = {}
1486
+
1487
+ # QC gate config — resolved once for both diss and chi loops below.
1488
+ qc_cfg = merge_config("qc", config.get("qc"))
1489
+ qc_enabled = bool(qc_cfg.get("enable", True))
1490
+ qc_drop_action = qc_cfg.get("drop_action", "nan")
1491
+ qc_eps_drop_from = list(qc_cfg.get("epsilon_drop_from") or [])
1492
+ qc_chi_drop_from = list(qc_cfg.get("chi_drop_from") or [])
1493
+ strat_cfg = merge_config("stratification", config.get("stratification"))
1494
+ strat_enabled = bool(strat_cfg.get("enable", True))
1495
+ ct_T_name = ct_cfg.get("T_name", "JAC_T")
1496
+ ct_C_name = ct_cfg.get("C_name", "JAC_C")
1497
+ # resolve_window_config guarantees concrete integer diss_length values.
1498
+ diss_length_seconds = float(eps_cfg["diss_length"]) / float(pf.fs_fast)
1499
+ chi_diss_length_seconds = (
1500
+ float(chi_cfg["diss_length"]) / float(pf.fs_fast) if chi_enabled else None
1501
+ )
1502
+
1503
+ if "diss" in output_dirs and result["profiles"]:
1504
+ with stage_log(output_dirs.get("diss"), log_basename):
1505
+ for prof_path in result["profiles"]:
1506
+ try:
1507
+ pre_loaded: dict[str, Any] | None
1508
+ if chi_enabled:
1509
+ # Single NC pass that produces both channels and
1510
+ # therm-gradient channels; reused below for chi.
1511
+ from odas_tpw.rsi.chi_io import _load_therm_channels
1512
+
1513
+ pre_loaded = _load_therm_channels(prof_path)
1514
+ prof_data_cache[prof_path] = pre_loaded
1515
+ else:
1516
+ pre_loaded = None
1517
+ diss_results = _compute_epsilon(
1518
+ prof_path,
1519
+ **{
1520
+ k: v
1521
+ for k, v in eps_cfg.items()
1522
+ if k
1523
+ not in (
1524
+ "epsilon_minimum",
1525
+ "T_source",
1526
+ "T1_norm",
1527
+ "T2_norm",
1528
+ "fom_max",
1529
+ "diagnostics",
1530
+ )
1531
+ },
1532
+ _pre_loaded=pre_loaded,
1533
+ )
1534
+ eps_fom_max = eps_cfg.get("fom_max")
1535
+ for ds in diss_results:
1536
+ if excluded_probes:
1537
+ _nan_excluded_probes(ds, excluded_probes, p_path.name)
1538
+ if eps_fom_max is not None:
1539
+ _apply_fom_cut(ds, float(eps_fom_max), p_path.name)
1540
+ ds = mk_epsilon_mean(ds, eps_cfg.get("epsilon_minimum", 1e-13))
1541
+ if qc_enabled:
1542
+ n_probe = int(ds["probe"].size) if "probe" in ds.dims else 2
1543
+ apply_qc_to_dataset(
1544
+ ds,
1545
+ pf,
1546
+ qc_eps_drop_from,
1547
+ diss_length_seconds,
1548
+ flag_var_name="qc_drop_epsilon",
1549
+ value_vars=[
1550
+ "epsilonMean",
1551
+ "epsilonLnSigma",
1552
+ *[f"e_{i + 1}" for i in range(n_probe)],
1553
+ "epsilon",
1554
+ ],
1555
+ drop_action=qc_drop_action,
1556
+ )
1557
+ _copy_profile_scalars(prof_path, ds, prof_scalars_cache)
1558
+ if strat_enabled:
1559
+ _attach_window_stratification(
1560
+ ds,
1561
+ prof_path,
1562
+ diss_length_seconds / 2,
1563
+ ct_T_name,
1564
+ ct_C_name,
1565
+ p_path.name,
1566
+ "dissipation",
1567
+ )
1568
+ out_name = Path(prof_path).name
1569
+ out_path = output_dirs["diss"] / out_name
1570
+ _atomic_to_netcdf(ds, out_path)
1571
+ out_path_str = str(out_path)
1572
+ result["diss"].append(out_path_str)
1573
+ diss_by_profile[prof_path] = out_path_str
1574
+ except Exception as exc:
1575
+ logger.error("diss for %s: %s", Path(prof_path).name, exc)
1576
+ result.setdefault("errors", []).append(f"diss {Path(prof_path).name}: {exc}")
1577
+
1578
+ # Per-profile chi (if enabled)
1579
+ if chi_enabled and result["diss"]:
1580
+ # chi_enabled guarantees the chi windows were resolved above.
1581
+ assert chi_diss_length_seconds is not None
1582
+ try:
1583
+ from odas_tpw.rsi.chi_io import _compute_chi
1584
+ except ImportError:
1585
+ pass
1586
+ else:
1587
+ chi_use_epsilon = bool(chi_cfg.get("use_epsilon", True))
1588
+ chi_fom_max = chi_cfg.get("fom_max")
1589
+ chi_spectral_qc = bool(chi_cfg.get("spectral_qc", True))
1590
+ chi_kwargs = {
1591
+ k: v
1592
+ for k, v in chi_cfg.items()
1593
+ if k
1594
+ not in (
1595
+ "enable",
1596
+ "chi_minimum",
1597
+ "spectral_qc",
1598
+ "diagnostics",
1599
+ "use_epsilon",
1600
+ "fom_max",
1601
+ "mixing",
1602
+ "salinity",
1603
+ )
1604
+ }
1605
+ # Soft spectral-QC thresholds for mk_chi_mean — the SAME limits the rsi
1606
+ # pipeline's chi_final uses, so both pipelines filter chiMean/K_T/Gamma
1607
+ # identically (issue #104 U3-C2).
1608
+ if chi_spectral_qc:
1609
+ from odas_tpw.chi.l4_chi import _CHI_FOM_LIMIT, _CHI_K_MAX_RATIO_MIN
1610
+
1611
+ chi_qc_kwargs = {
1612
+ "fom_limit": _CHI_FOM_LIMIT,
1613
+ "k_max_ratio_min": _CHI_K_MAX_RATIO_MIN,
1614
+ }
1615
+ else:
1616
+ chi_qc_kwargs = {}
1617
+ # Resolve chi.salinity: "measured" -> per-profile practical salinity
1618
+ # from the profile's own C/T/P (TEOS-10); a number or None is passed
1619
+ # through (None -> fixed 35 PSU viscosity in process_l3_chi).
1620
+ chi_sal_cfg = chi_cfg.get("salinity")
1621
+ chi_use_measured_sal = (
1622
+ isinstance(chi_sal_cfg, str) and chi_sal_cfg.strip().lower() == "measured"
1623
+ )
1624
+ with stage_log(output_dirs.get("chi"), log_basename):
1625
+ for prof_path in result["profiles"]:
1626
+ diss_path = diss_by_profile.get(prof_path)
1627
+ if diss_path is None:
1628
+ logger.warning(
1629
+ "chi skipped for %s: no matching diss output",
1630
+ Path(prof_path).name,
1631
+ )
1632
+ prof_data_cache.pop(prof_path, None)
1633
+ continue
1634
+ diss_ds = None
1635
+ try:
1636
+ import xarray as xr
1637
+
1638
+ # Method 1 uses shear-probe epsilon to seed kB.
1639
+ # Method 2 (use_epsilon=False) does a pure spectral
1640
+ # fit and is appropriate when shear epsilon is
1641
+ # contaminated (e.g. MR on a vibrating glider).
1642
+ diss_ds = xr.open_dataset(diss_path) if chi_use_epsilon else None
1643
+ # Pop so the cache is released as we consume it.
1644
+ pre_loaded = prof_data_cache.pop(prof_path, None)
1645
+ if chi_use_measured_sal:
1646
+ chi_sal = _profile_practical_salinity(
1647
+ prof_path,
1648
+ ct_cfg.get("T_name", "JAC_T"),
1649
+ ct_cfg.get("C_name", "JAC_C"),
1650
+ )
1651
+ if chi_sal is None:
1652
+ logger.warning(
1653
+ "chi salinity='measured' but %s lacks C/T/P; "
1654
+ "using fixed 35 PSU",
1655
+ Path(prof_path).name,
1656
+ )
1657
+ else:
1658
+ chi_sal = chi_sal_cfg
1659
+ chi_results = _compute_chi(
1660
+ prof_path,
1661
+ epsilon_ds=diss_ds,
1662
+ salinity=chi_sal,
1663
+ **chi_kwargs,
1664
+ _pre_loaded=pre_loaded,
1665
+ )
1666
+ for chi_ds in chi_results:
1667
+ if chi_fom_max is not None:
1668
+ _apply_fom_cut(
1669
+ chi_ds,
1670
+ float(chi_fom_max),
1671
+ p_path.name,
1672
+ two_sided=True,
1673
+ )
1674
+ chi_ds = mk_chi_mean(
1675
+ chi_ds,
1676
+ chi_cfg.get("chi_minimum", 1e-13),
1677
+ **chi_qc_kwargs,
1678
+ )
1679
+ # Apply chi QC BEFORE deriving mixing quantities so a
1680
+ # QC-dropped chi does not leak into K_T/Gamma (audit
1681
+ # r1-4). Mixing then consumes the NaN'd chiMean, so
1682
+ # K_T and Gamma (chi-derived) are NaN wherever chi was
1683
+ # dropped. K_rho is epsilon-derived (0.2*eps/N2) and is
1684
+ # correctly gated instead by the epsilon-side QC, via
1685
+ # the NaN epsilonMean paired in by _add_mixing_quantities.
1686
+ if qc_enabled:
1687
+ n_cprobe = (
1688
+ int(chi_ds["probe"].size) if "probe" in chi_ds.dims else 2
1689
+ )
1690
+ apply_qc_to_dataset(
1691
+ chi_ds,
1692
+ pf,
1693
+ qc_chi_drop_from,
1694
+ chi_diss_length_seconds,
1695
+ flag_var_name="qc_drop_chi",
1696
+ value_vars=[
1697
+ "chiMean",
1698
+ "chiLnSigma",
1699
+ *[f"chi_{i + 1}" for i in range(n_cprobe)],
1700
+ "chi",
1701
+ "epsilon_T",
1702
+ ],
1703
+ drop_action=qc_drop_action,
1704
+ )
1705
+ # Derived mixing quantities (Gamma, K_T, K_rho)
1706
+ # with real salinity from the profile's own C/T/P
1707
+ if chi_cfg.get("mixing", True):
1708
+ mix_eps = (
1709
+ diss_ds if diss_ds is not None else xr.open_dataset(diss_path)
1710
+ )
1711
+ # Gamma and K_rho are built from the shear-probe
1712
+ # epsilonMean even under Method 2 (use_epsilon=False),
1713
+ # the setting chosen precisely when shear epsilon is
1714
+ # distrusted. Record that provenance in the attrs and
1715
+ # warn so a downstream user is not misled.
1716
+ if not chi_use_epsilon:
1717
+ eps_prov = (
1718
+ " PROVENANCE: chi.use_epsilon=false (Method 2), "
1719
+ "so chi is a pure spectral fit, but this epsilon "
1720
+ "is still the shear-probe epsilonMean (the source "
1721
+ "Method 2 is chosen to distrust); treat Gamma/K_rho "
1722
+ "accordingly."
1723
+ )
1724
+ logger.warning(
1725
+ "%s: chi.use_epsilon=false but mixing Gamma/K_rho "
1726
+ "still use shear epsilonMean; see variable attrs",
1727
+ Path(prof_path).name,
1728
+ )
1729
+ else:
1730
+ eps_prov = ""
1731
+ try:
1732
+ chi_ds = _add_mixing_quantities(
1733
+ chi_ds,
1734
+ mix_eps,
1735
+ prof_path,
1736
+ T_name=ct_cfg.get("T_name", "JAC_T"),
1737
+ C_name=ct_cfg.get("C_name", "JAC_C"),
1738
+ file_label=Path(prof_path).name,
1739
+ epsilon_provenance=eps_prov,
1740
+ )
1741
+ finally:
1742
+ if mix_eps is not diss_ds:
1743
+ mix_eps.close()
1744
+ _copy_profile_scalars(prof_path, chi_ds, prof_scalars_cache)
1745
+ out_name = Path(prof_path).name
1746
+ out_path = output_dirs["chi"] / out_name
1747
+ _atomic_to_netcdf(chi_ds, out_path)
1748
+ result["chi"].append(str(out_path))
1749
+ except Exception as exc:
1750
+ logger.error("chi for %s: %s", Path(prof_path).name, exc)
1751
+ result.setdefault("errors", []).append(f"chi {Path(prof_path).name}: {exc}")
1752
+ finally:
1753
+ if diss_ds is not None:
1754
+ diss_ds.close()
1755
+
1756
+ return result
1757
+
1758
+
1759
+ # ---------------------------------------------------------------------------
1760
+ # Stage runners
1761
+ # ---------------------------------------------------------------------------
1762
+
1763
+
1764
+ def _configured_input_root(config: dict) -> Path:
1765
+ """Return the configured input root for preserving relative paths."""
1766
+ files_cfg = config.get("files", {})
1767
+ root = expand_config_dir(files_cfg.get("p_file_root", "VMP/"), config_dir_of(config))
1768
+ return Path(root)
1769
+
1770
+
1771
+ def _configured_output_root(config: dict) -> Path:
1772
+ """The configured output root, with any ``<CONFIG_DIR>`` token resolved."""
1773
+ files_cfg = config.get("files", {})
1774
+ root = expand_config_dir(files_cfg.get("output_root", "results/"), config_dir_of(config))
1775
+ return Path(root)
1776
+
1777
+
1778
+ def _check_unique_outputs(destinations: list[Path]) -> None:
1779
+ """Raise if multiple inputs would write the same output path."""
1780
+ seen: dict[Path, int] = {}
1781
+ for dest in destinations:
1782
+ key = dest.resolve()
1783
+ seen[key] = seen.get(key, 0) + 1
1784
+ duplicates = [str(path) for path, count in seen.items() if count > 1]
1785
+ if duplicates:
1786
+ raise ValueError(
1787
+ "multiple input files map to the same output path: " + ", ".join(sorted(duplicates))
1788
+ )
1789
+
1790
+
1791
+ def _check_unique_output_stems(stems: list[str]) -> None:
1792
+ """Raise if multiple processed inputs would share one output stem."""
1793
+ seen: dict[str, int] = {}
1794
+ for stem in stems:
1795
+ seen[stem] = seen.get(stem, 0) + 1
1796
+ duplicates = [stem for stem, count in seen.items() if count > 1]
1797
+ if duplicates:
1798
+ raise ValueError(
1799
+ "multiple input files map to the same output stem: " + ", ".join(sorted(duplicates))
1800
+ )
1801
+
1802
+
1803
+ def _trim_one(args: tuple):
1804
+ """Worker: trim one .p file and return (trim_result, source_path, error).
1805
+
1806
+ Designed to be pickled and dispatched to a ProcessPoolExecutor — must
1807
+ take a single argument and import its dependencies inside the call.
1808
+ The first element is a :class:`~odas_tpw.perturb.trim.TrimResult` on
1809
+ success (or ``None`` with an error string on failure).
1810
+ """
1811
+ p, trim_dir, root, force, cache_dir = args
1812
+ try:
1813
+ from odas_tpw.perturb.trim import trim_p_file
1814
+
1815
+ return (
1816
+ trim_p_file(p, trim_dir, root=root, force=force, cache_dir=cache_dir),
1817
+ p,
1818
+ None,
1819
+ )
1820
+ except Exception as exc:
1821
+ return None, p, str(exc)
1822
+
1823
+
1824
+ def _trim_log_message(result) -> str:
1825
+ """Render the per-file trim log line, reporting what happened to the file.
1826
+
1827
+ Distinguishes the three outcomes so a re-run is legible: a real trim
1828
+ reports the bytes dropped, a complete file reports that it was referenced
1829
+ in place (not copied), and an up-to-date trimmed output reports skipped.
1830
+ """
1831
+ name = result.dest.name
1832
+ if result.action == "trimmed":
1833
+ return (
1834
+ f"Trimmed: {name} (removed {result.bytes_removed} B — "
1835
+ f"incomplete final record of {result.record_size} B)"
1836
+ )
1837
+ if result.action == "skipped":
1838
+ return f"Trim: {name} (skipped — trimmed output already up to date)"
1839
+ return f"Trim: {name} (complete — referenced original in place, not copied)"
1840
+
1841
+
1842
+ def _log_trim_summary(results: list, n_failed: int) -> None:
1843
+ """Emit one roll-up line with the per-action counts."""
1844
+ counts = {"trimmed": 0, "referenced": 0, "skipped": 0}
1845
+ for r in results:
1846
+ counts[r.action] = counts.get(r.action, 0) + 1
1847
+ logger.info(
1848
+ "Trim summary: %d trimmed, %d referenced in place, %d skipped (up to date), %d failed",
1849
+ counts["trimmed"],
1850
+ counts["referenced"],
1851
+ counts["skipped"],
1852
+ n_failed,
1853
+ )
1854
+
1855
+
1856
+ def run_trim(
1857
+ config: dict,
1858
+ p_files: list[Path] | None = None,
1859
+ *,
1860
+ jobs: int = 1,
1861
+ ) -> list[Path]:
1862
+ """Trim corrupt final records from .p files; reference complete ones in place.
1863
+
1864
+ *jobs* > 1 dispatches :func:`trim_p_file` calls onto a process pool
1865
+ so the per-file work overlaps. Each call is independent (different
1866
+ output path, no shared state), so process-level parallelism gives
1867
+ near-linear speedup until disk I/O saturates.
1868
+
1869
+ Complete files are not copied — their original path is returned so
1870
+ downstream stages read them in place. Only genuinely-truncated files
1871
+ are written under ``<output_root>/trimmed/``. A trimmed file is skipped
1872
+ when an up-to-date output already exists (``files.force_trim: true``
1873
+ re-trims unconditionally). The returned list is the per-file path to
1874
+ use downstream (original for referenced files, trimmed path otherwise).
1875
+ """
1876
+ from odas_tpw.perturb.discover import find_p_files
1877
+ from odas_tpw.perturb.trim import trim_destination
1878
+
1879
+ files_cfg = config.get("files", {})
1880
+ output_root = _configured_output_root(config)
1881
+ trim_dir = output_root / "trimmed"
1882
+ cache_dir = output_root / _CACHE_DIRNAME
1883
+ root = _configured_input_root(config)
1884
+ force = bool(files_cfg.get("force_trim", False))
1885
+
1886
+ if p_files is None:
1887
+ root = _configured_input_root(config)
1888
+ pattern = files_cfg.get("p_file_pattern", "**/*.p")
1889
+ p_files = find_p_files(root, pattern)
1890
+ if not p_files:
1891
+ return []
1892
+
1893
+ destinations = [trim_destination(p, trim_dir, root=root) for p in p_files]
1894
+ _check_unique_outputs(destinations)
1895
+
1896
+ trim_results: list = []
1897
+ n_failed = 0
1898
+ if jobs <= 1 or len(p_files) <= 1:
1899
+ from odas_tpw.perturb.trim import trim_p_file
1900
+
1901
+ for p in p_files:
1902
+ try:
1903
+ result = trim_p_file(p, trim_dir, root=root, force=force, cache_dir=cache_dir)
1904
+ trim_results.append(result)
1905
+ logger.info("%s", _trim_log_message(result))
1906
+ except Exception as exc:
1907
+ n_failed += 1
1908
+ logger.error("trimming %s: %s", p.name, exc)
1909
+ else:
1910
+ with ProcessPoolExecutor(max_workers=jobs) as executor:
1911
+ futures = {
1912
+ executor.submit(_trim_one, (p, trim_dir, root, force, cache_dir)): p
1913
+ for p in p_files
1914
+ }
1915
+ for future in as_completed(futures):
1916
+ result, src, err = future.result()
1917
+ if err is not None:
1918
+ n_failed += 1
1919
+ logger.error("trimming %s: %s", src.name, err)
1920
+ continue
1921
+ assert result is not None # narrow for the type checker
1922
+ trim_results.append(result)
1923
+ logger.info("%s", _trim_log_message(result))
1924
+
1925
+ _log_trim_summary(trim_results, n_failed)
1926
+ return [r.dest for r in trim_results]
1927
+
1928
+
1929
+ def run_merge(
1930
+ config: dict,
1931
+ p_files: list[Path] | None = None,
1932
+ *,
1933
+ include_singletons: bool = False,
1934
+ input_root: Path | str | None = None,
1935
+ merge_plan: list[tuple[Path, list[Path]]] | None = None,
1936
+ ) -> list[Path]:
1937
+ """Merge split .p files.
1938
+
1939
+ By default, returns only newly merged outputs for CLI compatibility.
1940
+ When *include_singletons* is true, returns the complete post-merge file
1941
+ list: merged outputs plus untouched non-mergeable inputs.
1942
+ """
1943
+ from odas_tpw.perturb.discover import find_p_files
1944
+ from odas_tpw.perturb.merge import merge_p_files, plan_merge_outputs
1945
+
1946
+ files_cfg = config.get("files", {})
1947
+ root = Path(input_root) if input_root is not None else _configured_input_root(config)
1948
+ output_root = _configured_output_root(config)
1949
+ merge_dir = output_root / "merged"
1950
+
1951
+ if p_files is None:
1952
+ pattern = files_cfg.get("p_file_pattern", "**/*.p")
1953
+ p_files = find_p_files(root, pattern)
1954
+ if not p_files:
1955
+ return []
1956
+
1957
+ plan = (
1958
+ merge_plan if merge_plan is not None else plan_merge_outputs(p_files, merge_dir, root=root)
1959
+ )
1960
+ planned_outputs = [
1961
+ output_path for output_path, chain in plan if include_singletons or len(chain) > 1
1962
+ ]
1963
+ _check_unique_outputs(planned_outputs)
1964
+
1965
+ results: list[Path] = []
1966
+ for output_path, chain in plan:
1967
+ if len(chain) == 1:
1968
+ if include_singletons:
1969
+ results.append(chain[0])
1970
+ continue
1971
+ try:
1972
+ merged = merge_p_files(chain, merge_dir, root=root)
1973
+ results.append(merged)
1974
+ logger.info("Merged %d files -> %s", len(chain), merged.name)
1975
+ except Exception as exc:
1976
+ logger.error("merging: %s", exc)
1977
+ if include_singletons:
1978
+ results.extend(chain)
1979
+ return results
1980
+
1981
+
1982
+ def _setup_output_dirs(config: dict) -> dict[str, Path]:
1983
+ """Set up versioned output directories based on config."""
1984
+ output_root = _configured_output_root(config)
1985
+
1986
+ dirs = {}
1987
+
1988
+ # Profiles directory — hash includes profiles + fp07 + ct + gps + bottom + top_trim
1989
+ profiles_params = merge_config("profiles", config.get("profiles"))
1990
+ dirs["profiles"] = resolve_output_dir(
1991
+ output_root,
1992
+ "profiles",
1993
+ "profiles",
1994
+ profiles_params,
1995
+ upstream=_upstream_for("profiles", config),
1996
+ )
1997
+
1998
+ # Diss directory — hash includes epsilon params + profiles upstream
1999
+ eps_params = merge_config("epsilon", config.get("epsilon"))
2000
+ dirs["diss"] = resolve_output_dir(
2001
+ output_root,
2002
+ "diss",
2003
+ "epsilon",
2004
+ eps_params,
2005
+ upstream=_upstream_for("diss", config),
2006
+ )
2007
+
2008
+ # Chi directory (if enabled). merge_config so the enable decision matches
2009
+ # process_file's (which now resolves chi via merge_config).
2010
+ chi_cfg = merge_config("chi", config.get("chi"))
2011
+ if chi_cfg.get("enable", False):
2012
+ chi_params = merge_config("chi", chi_cfg)
2013
+ dirs["chi"] = resolve_output_dir(
2014
+ output_root,
2015
+ "chi",
2016
+ "chi",
2017
+ chi_params,
2018
+ upstream=_upstream_for("chi", config),
2019
+ )
2020
+
2021
+ # CTD directory
2022
+ ctd_cfg = config.get("ctd", {})
2023
+ if ctd_cfg.get("enable", True):
2024
+ ctd_params = merge_config("ctd", ctd_cfg)
2025
+ dirs["ctd"] = resolve_output_dir(
2026
+ output_root,
2027
+ "ctd",
2028
+ "ctd",
2029
+ ctd_params,
2030
+ upstream=_upstream_for("ctd", config),
2031
+ )
2032
+
2033
+ return dirs
2034
+
2035
+
2036
+ def _format_elapsed(seconds: float) -> str:
2037
+ """Human-readable elapsed time: ``12.3s`` under a minute, else ``1m 23s``."""
2038
+ if seconds >= 60:
2039
+ return f"{int(seconds // 60)}m {seconds % 60:02.0f}s"
2040
+ return f"{seconds:.1f}s"
2041
+
2042
+
2043
+ def _done_message(name: str, result: Any) -> str:
2044
+ """Per-file completion line reporting products written and elapsed time.
2045
+
2046
+ *result* is the dict returned by :func:`process_file` (``profiles`` /
2047
+ ``diss`` / ``chi`` lists, plus ``elapsed_s`` when run through
2048
+ :func:`_process_file_timed`). Reports the profile count headline plus the
2049
+ dissipation and chi counts (``0`` chi usually means chi was disabled, but
2050
+ also when no chi outputs were produced — e.g. every profile errored or had
2051
+ no matching diss output) and the per-file wall-clock when available.
2052
+ """
2053
+ if isinstance(result, dict):
2054
+ n_prof = len(result.get("profiles", []))
2055
+ n_diss = len(result.get("diss", []))
2056
+ n_chi = len(result.get("chi", []))
2057
+ elapsed = result.get("elapsed_s")
2058
+ else:
2059
+ n_prof = n_diss = n_chi = 0
2060
+ elapsed = None
2061
+ elapsed_str = f" in {_format_elapsed(elapsed)}" if elapsed is not None else ""
2062
+ return (
2063
+ f"Done processing {name}: {n_prof} profiles "
2064
+ f"({n_diss} dissipation, {n_chi} chi){elapsed_str}"
2065
+ )
2066
+
2067
+
2068
+ def _process_file_timed(*args, cachekey: str | None = None, **kwargs) -> dict:
2069
+ """Run :func:`process_file`, recording its wall-clock as ``elapsed_s``.
2070
+
2071
+ Wrapping rather than instrumenting process_file's several return points
2072
+ keeps the timing in one place. The elapsed seconds are measured inside the
2073
+ worker process, so they reflect true per-file processing time, not the
2074
+ time a file spent queued waiting for a free worker.
2075
+
2076
+ On success, writes the per-file cache marker (in the worker, so the real
2077
+ outputs are on disk) when *cachekey* is supplied — keeping marker-writing
2078
+ out of ``process_file`` itself so the mocked-``process_file`` tests are
2079
+ unaffected.
2080
+ """
2081
+ import time
2082
+
2083
+ if cachekey is not None:
2084
+ output_dirs = args[3] if len(args) > 3 else kwargs.get("output_dirs")
2085
+ output_stem = kwargs.get("output_stem")
2086
+ if output_dirs and output_stem:
2087
+ output_root = Path(next(iter(output_dirs.values()))).parent
2088
+ # Invalidate the old marker immediately. If processing is interrupted mid-write,
2089
+ # we don't want a corrupted NetCDF to falsely validate against the old marker
2090
+ # on a subsequent run.
2091
+ with suppress(OSError):
2092
+ _marker_path(output_root, output_stem).unlink(missing_ok=True)
2093
+
2094
+ # Unlink previous outputs for this stem so a failure doesn't leave
2095
+ # stale NetCDFs that get silently swept up by the combo binning glob.
2096
+ for stage, d in output_dirs.items():
2097
+ pattern = f"{output_stem}.nc" if stage == "ctd" else f"{output_stem}_prof*.nc"
2098
+ for nc_file in Path(d).glob(pattern):
2099
+ with suppress(OSError):
2100
+ nc_file.unlink()
2101
+
2102
+ t0 = time.monotonic()
2103
+ result = process_file(*args, **kwargs)
2104
+ if isinstance(result, dict):
2105
+ result["elapsed_s"] = time.monotonic() - t0
2106
+ # Only cache a CLEAN run. process_file catches and swallows real failures
2107
+ # (PFile load, profile extraction, per-profile diss/chi, CTD binning),
2108
+ # returning an empty/partial result with ``errors`` set. Writing a marker
2109
+ # then would lock that incomplete output in as cache-valid until --force /
2110
+ # an input or config change. A legitimately-empty file (no profiles found)
2111
+ # has no ``errors`` and IS cached. The marker stays invalidated (unlinked
2112
+ # above) so the next run retries.
2113
+ if cachekey is not None and isinstance(result, dict) and not result.get("errors"):
2114
+ output_dirs = args[3] if len(args) > 3 else kwargs.get("output_dirs")
2115
+ output_stem = kwargs.get("output_stem")
2116
+ if output_dirs and output_stem:
2117
+ output_root = Path(next(iter(output_dirs.values()))).parent
2118
+ try:
2119
+ _write_marker(output_root, output_stem, cachekey, output_dirs)
2120
+ except OSError as exc:
2121
+ logger.warning("could not write cache marker for %s: %s", output_stem, exc)
2122
+ return result
2123
+
2124
+
2125
+ def run_pipeline(config: dict, p_files: list[Path] | None = None) -> None:
2126
+ """Run the full pipeline: discover -> trim -> merge -> process -> bin -> combo.
2127
+
2128
+ Parameters
2129
+ ----------
2130
+ config : dict
2131
+ Full merged config (all sections).
2132
+ p_files : list of Path, optional
2133
+ Override file discovery with explicit file list.
2134
+ """
2135
+ from odas_tpw.perturb.discover import find_p_files
2136
+ from odas_tpw.perturb.gps import create_gps
2137
+
2138
+ files_cfg = config.get("files", {})
2139
+
2140
+ # Discover files
2141
+ if p_files is None:
2142
+ root = _configured_input_root(config)
2143
+ pattern = files_cfg.get("p_file_pattern", "**/*.p")
2144
+ p_files = find_p_files(root, pattern)
2145
+
2146
+ if not p_files:
2147
+ logger.warning("No .p files found")
2148
+ return
2149
+
2150
+ logger.info("Found %d .p files", len(p_files))
2151
+
2152
+ input_root = _configured_input_root(config)
2153
+ current_root = input_root
2154
+ instrument_key_by_path = {p.resolve(): p.parent.name for p in p_files}
2155
+ output_stem_by_path = {p.resolve(): _source_output_stem(p, current_root) for p in p_files}
2156
+
2157
+ # Parallel processing
2158
+ parallel_cfg = config.get("parallel", {})
2159
+ jobs = parallel_cfg.get("jobs", 1)
2160
+ if jobs == 0:
2161
+ jobs = os.cpu_count() or 1
2162
+
2163
+ # Trim
2164
+ if files_cfg.get("trim", True):
2165
+ logger.info("Trimming...")
2166
+ trimmed = run_trim(config, p_files, jobs=jobs)
2167
+ if trimmed:
2168
+ from odas_tpw.perturb.trim import trim_destination
2169
+
2170
+ trim_dir = _configured_output_root(config) / "trimmed"
2171
+ # run_trim returns the trim-dir path for trimmed/skipped files and
2172
+ # the *original* path for complete (referenced) files. Recover the
2173
+ # per-file physical path and keep stems/instrument keys canonical
2174
+ # (derived from the original source under the input root) so a
2175
+ # referenced file and a trimmed file share one stable identity.
2176
+ trimmed_keys = {p.resolve() for p in trimmed}
2177
+ next_p_files: list[Path] = []
2178
+ next_instruments: dict[Path, str] = {}
2179
+ next_stems: dict[Path, str] = {}
2180
+ for source in p_files:
2181
+ src_key = source.resolve()
2182
+ dest = trim_destination(source, trim_dir, root=current_root)
2183
+ if dest.resolve() in trimmed_keys:
2184
+ physical = dest # rewritten (trimmed) or reused (skipped)
2185
+ elif src_key in trimmed_keys:
2186
+ physical = source # complete — referenced in place
2187
+ else:
2188
+ continue # trim failed for this file; drop it
2189
+ key = physical.resolve()
2190
+ next_p_files.append(physical)
2191
+ next_instruments[key] = instrument_key_by_path.get(src_key, source.parent.name)
2192
+ next_stems[key] = _source_output_stem(source, current_root)
2193
+ p_files = next_p_files
2194
+ instrument_key_by_path = next_instruments
2195
+ output_stem_by_path = next_stems
2196
+ # current_root stays the input root: referenced files (the common
2197
+ # case) live there, and trimmed files resolve via the stem map
2198
+ # above rather than a relative-path computation.
2199
+
2200
+ # Merge
2201
+ if files_cfg.get("merge", False):
2202
+ logger.info("Merging...")
2203
+ from odas_tpw.perturb.merge import plan_merge_outputs
2204
+
2205
+ merge_dir = _configured_output_root(config) / "merged"
2206
+ merge_plan = plan_merge_outputs(p_files, merge_dir, root=current_root)
2207
+ merged_files = run_merge(
2208
+ config,
2209
+ p_files,
2210
+ include_singletons=True,
2211
+ input_root=current_root,
2212
+ merge_plan=merge_plan,
2213
+ )
2214
+ if merged_files:
2215
+ merged_keys = {p.resolve() for p in merged_files}
2216
+ next_instruments = {}
2217
+ next_output_stems = {}
2218
+ for output_path, chain in merge_plan:
2219
+ output_key = output_path.resolve()
2220
+ if output_key in merged_keys:
2221
+ source = chain[0]
2222
+ next_instruments[output_key] = instrument_key_by_path.get(
2223
+ source.resolve(), source.parent.name
2224
+ )
2225
+ if len(chain) > 1:
2226
+ # Genuinely merged → new file under merge_dir.
2227
+ next_output_stems[output_key] = _source_output_stem(output_path, merge_dir)
2228
+ else:
2229
+ # Passthrough singleton — keep the canonical stem the
2230
+ # trim stage assigned (the physical file may sit under
2231
+ # the input root or the trim dir).
2232
+ next_output_stems[output_key] = output_stem_by_path.get(
2233
+ output_key, _source_output_stem(output_path, current_root)
2234
+ )
2235
+ continue
2236
+ for source in chain:
2237
+ source_key = source.resolve()
2238
+ if source_key in merged_keys:
2239
+ next_instruments[source_key] = instrument_key_by_path.get(
2240
+ source_key, source.parent.name
2241
+ )
2242
+ next_output_stems[source_key] = output_stem_by_path.get(
2243
+ source_key, _source_output_stem(source, current_root)
2244
+ )
2245
+ p_files = merged_files
2246
+ instrument_key_by_path = next_instruments
2247
+ output_stem_by_path = next_output_stems
2248
+ current_root = merge_dir
2249
+
2250
+ _check_unique_output_stems(
2251
+ [
2252
+ output_stem_by_path.get(p.resolve(), _source_output_stem(p, current_root))
2253
+ for p in p_files
2254
+ ]
2255
+ )
2256
+
2257
+ # Setup output directories
2258
+ output_dirs = _setup_output_dirs(config)
2259
+ output_root = _configured_output_root(config)
2260
+
2261
+ # GPS provider (resolve any <CONFIG_DIR> in gps.file before it hits disk)
2262
+ cdir = config_dir_of(config)
2263
+ gps_cfg = merge_config("gps", config.get("gps"))
2264
+ if gps_cfg.get("file"):
2265
+ gps_cfg = {**gps_cfg, "file": expand_config_dir(gps_cfg["file"], cdir)}
2266
+ gps = create_gps(gps_cfg)
2267
+
2268
+ # Hotel data
2269
+ hotel_cfg = merge_config("hotel", config.get("hotel"))
2270
+ hotel_data = None
2271
+ if hotel_cfg.get("enable", False) and hotel_cfg.get("file"):
2272
+ from odas_tpw.perturb.hotel import load_hotel
2273
+
2274
+ hotel_data = load_hotel(
2275
+ expand_config_dir(hotel_cfg["file"], cdir),
2276
+ hotel_cfg.get("time_column", "time"),
2277
+ hotel_cfg.get("time_format", "auto"),
2278
+ hotel_cfg.get("channels", {}),
2279
+ )
2280
+ logger.info("Loaded hotel file: %d channels", len(hotel_data.channels))
2281
+
2282
+ def _instrument_key(p):
2283
+ return instrument_key_by_path.get(p.resolve(), p.parent.name)
2284
+
2285
+ def _output_stem(p):
2286
+ return output_stem_by_path.get(p.resolve(), _source_output_stem(p, current_root))
2287
+
2288
+ # Per-file incremental cache: compute each file's cache key (input identity +
2289
+ # external inputs + the config/engine-hashed output dirs it targets) and skip
2290
+ # files whose marker still matches AND whose outputs still exist on disk.
2291
+ # ``--force`` (files.force) reprocesses everything. Every file's key — skipped
2292
+ # or not — also seeds the bin/combo manifest (``file_cachekeys``).
2293
+ force = bool(files_cfg.get("force", False))
2294
+ ext_fp = _external_input_fingerprints(config, hotel_cfg)
2295
+ stage_hashes = _stage_signature_hashes(output_dirs)
2296
+
2297
+ # Fold in config keys that affect per-file output but are excluded from
2298
+ # directory signature hashes (so they cause an in-place rebuild).
2299
+ per_file_volatile: dict[str, Any] = {}
2300
+ if "hotel" in config:
2301
+ per_file_volatile["hotel"] = config["hotel"]
2302
+ if config.get("ctd", {}).get("diagnostics"):
2303
+ per_file_volatile["ctd_diagnostics"] = True
2304
+
2305
+ # One directory listing per stage up front, so each file's output-exists
2306
+ # check is set membership instead of a stat per NetCDF (SMB-friendly).
2307
+ present_outputs = _list_present_outputs(output_dirs, output_root)
2308
+ file_cachekeys: dict[str, str] = {}
2309
+ to_process: list[tuple[Path, str]] = []
2310
+ n_skipped = 0
2311
+ for p in p_files:
2312
+ stem = _output_stem(p)
2313
+ cachekey = _file_cachekey(p, ext_fp, stage_hashes, volatile_cfg=per_file_volatile)
2314
+ file_cachekeys[stem] = cachekey
2315
+ if not force and _marker_is_current(
2316
+ _marker_path(output_root, stem), cachekey, present_outputs
2317
+ ):
2318
+ n_skipped += 1
2319
+ else:
2320
+ to_process.append((p, cachekey))
2321
+
2322
+ if n_skipped:
2323
+ logger.info(
2324
+ "Processing %d file(s) (jobs=%d); %d up to date (skipped)",
2325
+ len(to_process),
2326
+ jobs,
2327
+ n_skipped,
2328
+ )
2329
+ else:
2330
+ logger.info("Processing %d files (jobs=%d)...", len(to_process), jobs)
2331
+
2332
+ if jobs == 1:
2333
+ for p_path, cachekey in to_process:
2334
+ logger.info("Processing %s...", p_path.name)
2335
+ result = _process_file_timed(
2336
+ p_path,
2337
+ config,
2338
+ gps,
2339
+ output_dirs,
2340
+ hotel_data=hotel_data,
2341
+ hotel_cfg=hotel_cfg,
2342
+ instrument_key=_instrument_key(p_path),
2343
+ output_stem=_output_stem(p_path),
2344
+ cachekey=cachekey,
2345
+ )
2346
+ logger.info("%s", _done_message(p_path.name, result))
2347
+ else:
2348
+ # Spawn workers each get a per-pid log file inside <output_root>/logs/
2349
+ # so multi-process runs are diagnosable. ``run_stamp`` is the parent
2350
+ # CLI's invocation timestamp so all workers share a prefix.
2351
+ worker_log_dir = output_root / "logs"
2352
+ run_stamp = current_run_stamp()
2353
+ with ProcessPoolExecutor(
2354
+ max_workers=jobs,
2355
+ initializer=init_worker_logging,
2356
+ initargs=(worker_log_dir, run_stamp),
2357
+ ) as executor:
2358
+ futures = {
2359
+ executor.submit(
2360
+ _process_file_timed,
2361
+ p,
2362
+ config,
2363
+ gps,
2364
+ output_dirs,
2365
+ hotel_data=hotel_data,
2366
+ hotel_cfg=hotel_cfg,
2367
+ instrument_key=_instrument_key(p),
2368
+ output_stem=_output_stem(p),
2369
+ cachekey=cachekey,
2370
+ ): p
2371
+ for p, cachekey in to_process
2372
+ }
2373
+ for future in as_completed(futures):
2374
+ p = futures[future]
2375
+ try:
2376
+ result = future.result()
2377
+ logger.info("%s", _done_message(p.name, result))
2378
+ except Exception as exc:
2379
+ logger.error("processing %s: %s", p.name, exc)
2380
+
2381
+ # Binning — three independent stages (profiles / diss / chi) that
2382
+ # share no data, so we run them concurrently on a thread pool. Each
2383
+ # stage's heavy lifting is numpy + bincount + xarray-IO, which release
2384
+ # the GIL, so threads give us most of the wall-clock win without
2385
+ # nesting process pools or paying the ProcessPoolExecutor IPC cost on
2386
+ # whole xarray Datasets.
2387
+ binning_cfg = config.get("binning", {})
2388
+ bin_method = binning_cfg.get("method", "depth")
2389
+ bin_width = binning_cfg.get("width", 1.0)
2390
+ aggregation = binning_cfg.get("aggregation", "mean")
2391
+ diagnostics = binning_cfg.get("diagnostics", False)
2392
+
2393
+ from odas_tpw.perturb.binning import bin_by_depth, bin_by_time, bin_chi, bin_diss
2394
+
2395
+ prof_binned_dir: Path | None = None
2396
+ diss_binned_dir: Path | None = None
2397
+ chi_binned_dir: Path | None = None
2398
+
2399
+ # Drop per-profile NCs orphaned by a prior run on a different input set
2400
+ # before globbing, so the combos reflect exactly the current .p files
2401
+ # (new files are still picked up incrementally; only files no longer
2402
+ # present are removed). See _prune_orphan_profile_ncs.
2403
+ valid_stems = {_output_stem(p) for p in p_files}
2404
+ for _stage in ("profiles", "diss", "chi"):
2405
+ if _stage in output_dirs:
2406
+ _prune_orphan_profile_ncs(output_dirs[_stage], valid_stems)
2407
+ # CTD per-file outputs are named {stem}.nc (no _prof suffix), so prune them
2408
+ # with exact-stem matching before the CTD combo globs the dir.
2409
+ if "ctd" in output_dirs:
2410
+ _prune_orphan_named_ncs(output_dirs["ctd"], valid_stems)
2411
+
2412
+ # Resolve binned-output dirs *before* the bin call so we can pass them
2413
+ # as ``log_dir`` and each input .p file's records land in
2414
+ # ``<binned_dir>/<stem>.log``. This also means the binned dir is
2415
+ # created (with its signature) even if the bin call ends up emitting
2416
+ # no data — a small price for symmetric per-file logs.
2417
+ prof_ncs = sorted(output_dirs["profiles"].glob("*.nc")) if "profiles" in output_dirs else []
2418
+ if prof_ncs:
2419
+ binning_params = merge_config("binning", binning_cfg)
2420
+ prof_binned_dir = resolve_output_dir(
2421
+ output_root,
2422
+ "profiles_binned",
2423
+ "binning",
2424
+ binning_params,
2425
+ upstream=_upstream_for("profiles_binned", config),
2426
+ )
2427
+
2428
+ diss_ncs = sorted(output_dirs["diss"].glob("*.nc")) if "diss" in output_dirs else []
2429
+ diss_width = binning_cfg.get("diss_width") or bin_width
2430
+ diss_agg = binning_cfg.get("diss_aggregation") or aggregation
2431
+ if diss_ncs:
2432
+ diss_binned_dir = resolve_output_dir(
2433
+ output_root,
2434
+ "diss_binned",
2435
+ "binning",
2436
+ merge_config("binning", binning_cfg),
2437
+ upstream=_upstream_for("diss_binned", config),
2438
+ )
2439
+
2440
+ chi_ncs: list[Path] = []
2441
+ chi_width = binning_cfg.get("chi_width") or binning_cfg.get("diss_width") or bin_width
2442
+ chi_agg = (
2443
+ binning_cfg.get("chi_aggregation") or binning_cfg.get("diss_aggregation") or aggregation
2444
+ )
2445
+ if "chi" in output_dirs:
2446
+ chi_ncs = sorted(output_dirs["chi"].glob("*.nc"))
2447
+ if chi_ncs:
2448
+ chi_binned_dir = resolve_output_dir(
2449
+ output_root,
2450
+ "chi_binned",
2451
+ "binning",
2452
+ merge_config("binning", binning_cfg),
2453
+ upstream=_upstream_for("chi_binned", config),
2454
+ )
2455
+
2456
+ # Stages run sequentially because each one now drives a per-profile
2457
+ # ``ProcessPoolExecutor`` of size *jobs*. Running the 3 stages
2458
+ # concurrently on a thread pool would over-commit (3 * jobs workers
2459
+ # competing for the same cores) and cost more than it saves.
2460
+
2461
+ binning_volatile = {"diagnostics": diagnostics} if diagnostics else None
2462
+
2463
+ if prof_ncs and prof_binned_dir is not None:
2464
+ manifest = _inputs_manifest(prof_ncs, file_cachekeys, volatile_cfg=binning_volatile)
2465
+ if not force and _output_is_current(prof_binned_dir / "binned.nc", manifest):
2466
+ logger.info("Binning profiles... up to date (skipped)")
2467
+ else:
2468
+ logger.info("Binning profiles...")
2469
+ if bin_method == "depth":
2470
+ ds = bin_by_depth(
2471
+ prof_ncs,
2472
+ bin_width,
2473
+ aggregation,
2474
+ diagnostics,
2475
+ log_dir=prof_binned_dir,
2476
+ jobs=jobs,
2477
+ )
2478
+ else:
2479
+ ds = bin_by_time(
2480
+ prof_ncs, bin_width, aggregation, diagnostics, log_dir=prof_binned_dir
2481
+ )
2482
+ _write_binned_or_clear(ds, prof_binned_dir, manifest)
2483
+
2484
+ if diss_ncs and diss_binned_dir is not None:
2485
+ manifest = _inputs_manifest(diss_ncs, file_cachekeys, volatile_cfg=binning_volatile)
2486
+ if not force and _output_is_current(diss_binned_dir / "binned.nc", manifest):
2487
+ logger.info("Binning dissipation... up to date (skipped)")
2488
+ else:
2489
+ logger.info("Binning dissipation...")
2490
+ ds = bin_diss(
2491
+ diss_ncs,
2492
+ diss_width,
2493
+ diss_agg,
2494
+ bin_method,
2495
+ diagnostics,
2496
+ log_dir=diss_binned_dir,
2497
+ jobs=jobs,
2498
+ )
2499
+ _write_binned_or_clear(ds, diss_binned_dir, manifest)
2500
+
2501
+ if chi_ncs and chi_binned_dir is not None:
2502
+ manifest = _inputs_manifest(chi_ncs, file_cachekeys, volatile_cfg=binning_volatile)
2503
+ if not force and _output_is_current(chi_binned_dir / "binned.nc", manifest):
2504
+ logger.info("Binning chi... up to date (skipped)")
2505
+ else:
2506
+ logger.info("Binning chi...")
2507
+ ds = bin_chi(
2508
+ chi_ncs,
2509
+ chi_width,
2510
+ chi_agg,
2511
+ bin_method,
2512
+ diagnostics,
2513
+ log_dir=chi_binned_dir,
2514
+ jobs=jobs,
2515
+ )
2516
+ _write_binned_or_clear(ds, chi_binned_dir, manifest)
2517
+
2518
+ # Combo assembly
2519
+ _run_combo(
2520
+ output_root,
2521
+ prof_binned_dir,
2522
+ diss_binned_dir,
2523
+ chi_binned_dir,
2524
+ output_dirs.get("ctd"),
2525
+ config.get("netcdf", {}),
2526
+ config=config,
2527
+ bin_method=bin_method,
2528
+ force=force,
2529
+ file_cachekeys=file_cachekeys,
2530
+ )
2531
+
2532
+ logger.info("Pipeline complete.")
2533
+
2534
+
2535
+ def _run_combo(
2536
+ output_root: Path,
2537
+ prof_binned_dir: Path | None,
2538
+ diss_binned_dir: Path | None,
2539
+ chi_binned_dir: Path | None,
2540
+ ctd_dir: Path | None,
2541
+ netcdf_attrs: dict,
2542
+ config: dict | None = None,
2543
+ bin_method: str = "depth",
2544
+ force: bool = False,
2545
+ file_cachekeys: dict[str, str] | None = None,
2546
+ ) -> None:
2547
+ """Assemble combo NetCDFs from each populated binned directory.
2548
+
2549
+ Each combo writes ``output_root/<name>_combo/combo.nc`` with CF/ACDD
2550
+ metadata applied via :mod:`odas_tpw.perturb.netcdf_schema`. Missing
2551
+ or empty input dirs are silently skipped — this is a best-effort
2552
+ publish step run after binning completes.
2553
+
2554
+ *bin_method* selects the gluing strategy for profiles/diss/chi combos:
2555
+
2556
+ * ``"depth"`` — vertical profilers (cast VMP, glider-mounted MR).
2557
+ Each per-file binned NetCDF has dims ``(bin, profile)``; the combo
2558
+ stacks them widthwise so the result is a profile-set indexed by
2559
+ time. Produces CF ``featureType=profile``.
2560
+ * ``"time"`` — non-profiling instruments (mooring-mounted MR, AUV).
2561
+ Each per-file binned NetCDF has a ``time`` dimension; the combo
2562
+ concatenates them lengthwise and sorts by time, like the CTD/scalar
2563
+ combo. Produces CF ``featureType=trajectory``.
2564
+
2565
+ The ctd_combo path always uses ``time`` regardless — scalar sensor
2566
+ streams are inherently per-time, never per-depth.
2567
+
2568
+ When *config* is provided, each combo dir also receives a
2569
+ ``.params_sha256_<hash>`` signature capturing the full upstream chain
2570
+ so the hash matches the corresponding binned/source dir.
2571
+ """
2572
+ from odas_tpw.perturb.combo import make_combo, make_ctd_combo
2573
+ from odas_tpw.perturb.netcdf_schema import CHI_SCHEMA, COMBO_SCHEMA, CTD_SCHEMA
2574
+
2575
+ logger.info("Assembling combo files (method=%s)...", bin_method)
2576
+
2577
+ binning_p = (
2578
+ merge_config("binning", config.get("binning") if config else None)
2579
+ if config is not None
2580
+ else None
2581
+ )
2582
+
2583
+ def _resolve_dst(stage: str, section: str, params: dict | None) -> Path:
2584
+ """Versioned ``stage_NN/`` when config available; legacy fixed path otherwise."""
2585
+ if config is None or params is None:
2586
+ d = output_root / stage
2587
+ d.mkdir(parents=True, exist_ok=True)
2588
+ return d
2589
+ return resolve_output_dir(
2590
+ output_root,
2591
+ stage,
2592
+ section,
2593
+ params,
2594
+ upstream=_upstream_for(stage, config),
2595
+ )
2596
+
2597
+ targets = [
2598
+ (prof_binned_dir, "combo", COMBO_SCHEMA, bin_method, make_combo),
2599
+ (diss_binned_dir, "diss_combo", COMBO_SCHEMA, bin_method, make_combo),
2600
+ (chi_binned_dir, "chi_combo", CHI_SCHEMA, bin_method, make_combo),
2601
+ ]
2602
+ for src, stage, schema, method, func in targets:
2603
+ if src is None or not src.exists():
2604
+ continue
2605
+ dst = _resolve_dst(stage, "binning", binning_p)
2606
+ # Propagate the source binned.nc's input manifest: skip re-assembling
2607
+ # the combo when its (unchanged) inputs already produced it. None when
2608
+ # the binned dir carries no manifest (e.g. the standalone `perturb bin`
2609
+ # path) → always rebuild, safely.
2610
+ manifest = _read_manifest_attr(src / "binned.nc")
2611
+ if not force and _output_is_current(dst / "combo.nc", manifest):
2612
+ logger.info("%s up to date (skipped)", dst.name)
2613
+ continue
2614
+ with stage_log(dst, "combo"):
2615
+ try:
2616
+ out = func(src, dst, schema, netcdf_attrs=netcdf_attrs, method=method)
2617
+ if out is not None:
2618
+ if manifest is not None:
2619
+ _stamp_manifest(Path(out), manifest)
2620
+ logger.info("Wrote %s", out)
2621
+ else:
2622
+ # Empty binned source -> no combo produced. Remove any stale
2623
+ # combo.nc so a shrunk/zeroed input set does not leave the
2624
+ # previous run's combo as the apparently-current product
2625
+ # (mirrors _write_binned_or_clear on the binning side).
2626
+ stale = dst / "combo.nc"
2627
+ if stale.exists():
2628
+ stale.unlink()
2629
+ logger.info("Removed stale %s (no binned input)", stale)
2630
+ except Exception as exc:
2631
+ logger.error("combo %s: %s", dst.name, exc)
2632
+
2633
+ if ctd_dir is not None and ctd_dir.exists():
2634
+ ctd_p = merge_config("ctd", config.get("ctd")) if config is not None else None
2635
+ ctd_combo_dir = _resolve_dst("ctd_combo", "ctd", ctd_p)
2636
+ ctd_manifest = (
2637
+ _inputs_manifest(sorted(ctd_dir.glob("*.nc")), file_cachekeys)
2638
+ if file_cachekeys is not None
2639
+ else None
2640
+ )
2641
+ if not force and _output_is_current(ctd_combo_dir / "combo.nc", ctd_manifest):
2642
+ logger.info("%s up to date (skipped)", ctd_combo_dir.name)
2643
+ else:
2644
+ with stage_log(ctd_combo_dir, "combo"):
2645
+ try:
2646
+ out = make_ctd_combo(
2647
+ ctd_dir, ctd_combo_dir, CTD_SCHEMA, netcdf_attrs=netcdf_attrs
2648
+ )
2649
+ if out is not None:
2650
+ if ctd_manifest is not None:
2651
+ _stamp_manifest(Path(out), ctd_manifest)
2652
+ logger.info("Wrote %s", out)
2653
+ except Exception as exc:
2654
+ logger.error("ctd combo: %s", exc)