NCEP obs refactor

[aayushg55]

Earth2Studio Pull Request

Description

NNJAObsConv combined three concerns in one class: request assembly, NNJA
S3/cache transport, and conventional/GPSRO product decoding. That made the
request lifecycle impossible to reuse for other NCEP observation products or
archives without copying it.

This PR is a behavior-preserving refactor that separates those concerns behind
a private seam:

• add _NCEPObsSourceBase for the shared DataFrame request lifecycle (input
  normalization, field projection, cycle-window construction, URI dedup,
  ordered decode/concat, source attrs, sync wrapping, cleanup);
• add the three-operation _NCEPObsStore protocol (fetch_files,
  local_path, cleanup);
• move NNJA anonymous-S3 session/retry/concurrency/raw-cache behavior into
  _NNJAObsStore;
• add _NNJAObsSourceBase for NNJA store wiring, six-hour validation,
  availability, and cache access;
• keep URI construction, product start dates, missing-file policy, lexicon
  routing, and PrepBUFR/GPSRO decoding on NNJAObsConv.

No public API, schema, variable, dtype, or runtime behavior changes.

Why This Boundary

NCEP observation sources vary along two independent axes:

• the archive transport determines how raw files are located, fetched, and
  cached;
• the observation product determines available times, file routes, variables,
  decoding, and missing-file policy.

This PR separates those axes without introducing a broad public framework:

_NNJAObsStore --implements--> _NCEPObsStore

_NNJAObsSourceBase --inherits--> _NCEPObsSourceBase
_NNJAObsSourceBase --composes--> _NNJAObsStore
NNJAObsConv --inherits---------> _NNJAObsSourceBase

The private seam is intended to support this follow-up shape (not included
here):

_NCEPObsSourceBase
  |-- _NNJAObsSourceBase
  |     |-- NNJAObsConv
  |     `-- NNJAObsSat         (follow-up)
  `-- NomadsGDASObsConv        (follow-up)

_NCEPObsStore
  |-- _NNJAObsStore
  `-- _NomadsObsStore          (follow-up)

The goal is only to establish and validate the minimal reusable ownership
boundary so those sources do not copy the request/download/assembly lifecycle
again.

Compatibility

• NNJAObsConv constructor and call signatures are unchanged.
• Variables, schema names/order/types, pandas dtypes, nullability, and attrs
  are unchanged.
• Cycle selection, tolerance handling, and row populations are unchanged.
• NNJA cache roots, raw cache filenames, retry/concurrency behavior, temporary
  cleanup, and missing-file behavior are unchanged.
• The existing optional BUFR dependency gate remains on the public source.

Validation

Origin/main and this branch independently decoded the same bounded real
PrepBUFR bytes. Exact positional DataFrame equality passed for 173,152 rows
across t, q, u, v, and pres, including row order, every value,
dtypes, null masks, column order, and attrs. The ordered row SHA-256 is
7becb34bc3bcec2211e046ca5f01ac9b2be389111914c4046c115ff10d80a31d.

Focused NNJA tests (19 passed, 2 slow deselected), Ruff, Black, mypy, and
git diff --check pass on the changed files.

Follow-ups

1. NNJAObsSat as a second _NNJAObsSourceBase consumer.
2. _NomadsObsStore + NomadsGDASObsConv migration deleting its duplicate
   lifecycle.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• Assess and address Greptile feedback (AI code review bot for guidance; use discretion, addressing all feedback is not required).

Dependencies

[greptile-apps]

Greptile Summary

This PR refactors NNJAObsConv by extracting the shared DataFrame request lifecycle into a new _NCEPObsSourceBase class and the NNJA S3 transport into _NNJAObsStore, wiring them together through a _NCEPObsStore protocol and a new _NNJAObsSourceBase intermediate layer. No public API, schema, variable, dtype, or runtime behavior changes are introduced.

• New ncep_obs.py introduces the _NCEPObsStore Protocol (three-method contract), _NCEPObsSourceBase (input normalization, URI dedup, ordered decode/concat, schema resolution), and a public resolve_fields alias — all private surface.
• Refactored nnja.py moves NNJA anonymous-S3 session/retry/cache behavior into _NNJAObsStore, adds _NNJAObsSourceBase for six-hour validation and store wiring, and narrows NNJAObsConv to URI construction, product dates, missing-file policy, and PrepBUFR/GPSRO decoding.
• Test updates correct _sync_async patch paths to the new module and add an asyncio test verifying the store-contract delegation (fetch, local_path, URI dedup).

Confidence Score: 4/5

Safe to merge — the refactor is behavior-preserving; transport, validation, and decoding paths produce identical outputs, verified by the author against 173k rows of real PrepBUFR data.

The decomposition is clean and the public-facing NNJAObsConv interface is unchanged. The _cycle_windows method brings a 6-hour-specific assumption into the generic _NCEPObsSourceBase layer, which could silently misdirect a future non-6-hourly NCEP source. The new store-contract test exercises fetch but not the call → finally → store.cleanup() delegation, leaving a small gap in the teardown coverage path.

earth2studio/data/ncep_obs.py — the _cycle_windows placement; test/data/test_nnja.py — cleanup delegation not exercised in the new asyncio test.

Important Files Changed

Filename	Overview
earth2studio/data/ncep_obs.py	New module introducing _NCEPObsStore Protocol and _NCEPObsSourceBase; clean separation of concerns, well-structured with appropriate abstract hooks and schema resolution logic
earth2studio/data/nnja.py	Decomposes _NNJAObsBase into _NNJAObsStore (transport) and _NNJAObsSourceBase (lifecycle wiring); NNJAObsConv public API and behavior preserved with correct MRO delegation and warn-only missing-file handler
test/data/test_nnja.py	Patch targets updated to ncep_obs module; new asyncio test validates store contract and lifecycle, covering fetch delegation, URI dedup, and attrs propagation

_{Reviews (1): Last reviewed commit: "rename to ncep_obs" | Re-trigger Greptile}

[aayushg55]

Note: this retains the original time request semantics of NNJAObsConv. Can patch in future PR: #967

[NickGeneva]

lets not create other interfaces... the interfaces are in base.py

focus on mixins, not protocals and base classes.

[NickGeneva]

Too many base classes.

Try to make just one. Avoiding some crazy inheritance stack is good. Because hard to look where code is at

Some repeat code is okay, this doesnt look like its adding much?

[aayushg55]

I can switch to mixin. The thing is it’s literally duplicate functionality across gdas and nnja (and gfs if we add that). We were also seeing divergence in functionality because of separation.

[aayushg55]

@NickGeneva Reworked this per your feedback. Quick rundown of what changed and where it's headed.

Addressing your notes

• Killed the _NCEPObsStore Protocol and the
  _NCEPObsSourceBase → _NNJAObsSourceBase → NNJAObsConv base-class
  stack — no new interfaces
• Deleted the separate _ncep_obs.py module; all shared NCEP-obs code
  lives in utils_ncep.py now
• It's flat mixins, not an inheritance stack.

New shape (2 files)

utils_ncep.py — the shared pieces:

• NCEPObsRequestMixin: the request lifecycle every NCEP obs source
  shares (__call__ / fetch / compile / 6h cycle-window planning /
  schema resolution).
• _NCEPObsTask (minimal base: uri + time window) + _NCEPConvTask
  (adds route + var_plan), and a _plan_conv_tasks(...) free
  function that routes prepbufr/gpsro variables into tasks. Plus the
  existing decode adapters.

nnja.py — NNJAObsMixin (NNJA S3 transport/cache/validation) + the
concrete NNJAObsConv(NNJAObsMixin, NCEPObsRequestMixin).

So a source = compose {a transport mixin} + {the request-lifecycle mixin}
and supply its lexicon / URI builder / decode. No chain, no protocol.

Why it's shaped this way

The next stacked PRs drop in additively:

• NNJAObsSat: (NNJAObsMixin, NCEPObsRequestMixin) + its own
  _NCEPSatTask + per-sensor planner + decode. Reuses NNJA transport + lifecycle;
  touches nothing shared.
• NomadsGDAS conv: (NomadsObsMixin, NCEPObsRequestMixin) +
  reuse _plan_conv_tasks with its own URLs; deletes its duplicated
  lifecycle.