Simplify the canonical data loading boundary

[jder]

Summary

• rebase the canonical-reader refactor needed by Add opt-in Rust OM4 data loader #800 onto the LLC support landed in Add LLC CPU data canonicalization #670
• replace the overlapping DatasetSpec / CanonicalDataset / OceanData abstractions with DataLayout, CanonicalSource, BatchPreprocessor, HostBatch, and ModelBatch
• make the storage-independent reader accept explicit time indices and canonical channel names
• push prognostic-versus-boundary selection into read requests instead of maintaining sliced source objects
• move OM4-specific canonicalization into Om4DataSourceConfig.canonicalize_datasets
• make OM4 and LLC canonicalizers derive their channel selections from variable keys and return the resulting DataLayout
• keep raw source conventions such as LLC staggered masks and OM4 mask names out of DataLayout
• preserve the grid-geometry, analysis-ready writer output, and training-progress changes currently on main

This intentionally does not add an LLC Rust loader. It establishes the narrower CanonicalReader seam that the Rust loader can implement later without exposing xarray or source-specific naming to the training pipeline.

Impact

Callers now request ordered canonical channels directly. Canonicalization owns source-specific naming, dimension, mask, and variable-selection rules, while DataLayout describes only the canonical/model-facing result. This removes repeated variable filtering and normalization wrappers from the hot data-loading path.

Validation

• uvx pre-commit run --all-files
• CUDA_VISIBLE_DEVICES='' uv run pytest -q -m 'not manual and not cuda'
  • 323 passed, 2 skipped, 10 xfailed

[alxmrs]

@fomo-bot will you make your best attempt to rebase in the main branch and address all outstanding merge conflicts?

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[fomo-bot]

I fetched origin and attempted the requested rebase onto origin/main. Git reported codex/canonical-reader-llc is already up to date, so there were no merge conflicts to resolve and no commit/push was needed.

Validation completed:

• No conflict markers found.
• git diff --check origin/main...HEAD passed.
• CUDA_VISIBLE_DEVICES= uv run pytest -m "not manual and not cuda" passed: 459 passed, 2 skipped, 68 deselected, 10 xfailed.
• GitHub checks on head 8a39ecb5a944 are complete and passing/skipped as expected.

PR #823 is still MERGEABLE, but GitHub reports mergeStateStatus: BLOCKED because it is draft and still review-required, not because of merge conflicts.

Debug info

• Repository: m2lines/Samudra
• Trigger: Simplify the canonical data loading boundary #823 (comment)
• Request ID: 20bd747d-2a5f-4b96-a8f4-d1004e4cdee3
• Accepted: true

[alxmrs]

I have reviewed this change, focusing on the core parts which are data.py and dataset.py. I'm really excited to get it merged. It LGTM.

[alxmrs]

@codex may I have your review?

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 8a39ecb5a9

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[alxmrs]

@codex, my agent has attempted to address your concerns from a previous review. Will you please review this PR once again?

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: b3d363ed92

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Preserve curvilinear coordinates while expanding levels

When an OM4 source has y/x dimensions with 2-D lat/lon coordinates and grid_type: tripolar, with_lat_lon_coords preserves the real geometry as lat_2d/lon_2d, but this whitelist copies only time, lat, and lon into the new dataset. The resulting CanonicalSource.coordinates() therefore omits the real 2-D coordinates, and the save_zarr path aborts in ZarrWriter._output_coords because tripolar geometry cannot be reconstructed by broadcasting the 1-D axes. Copy the preserved curvilinear coordinates into the canonical dataset.

Useful? React with 👍 / 👎.

[alxmrs]

@codex, may I have one more review?

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Keep it up!

Reviewed commit: 4399ced28c

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[alxmrs]

@fomo-bot Do you think this PR is ready to be merged? The GPU CI failures are because we can't request GPUs machines at this time, let's assume they would pass if we could. If you think this is ready, will you approve the PR?

[fomo-bot]

Reviewed current head 4399ced. I did not find merge-blocking issues; CPU/pre-commit/data/container checks pass, and the GPU failure is EC2 capacity before tests run.

[fomo-bot]

I think this PR is ready to merge under the stated assumption that the GPU jobs would pass if runners could be provisioned.



I reviewed current head 4399ced28cc13d00e95b40a5fcadf2c817627c77, checked the full PR diff, review threads, current CI, and the local worktree. The remaining GPU failure is EC2 InsufficientInstanceCapacity before tests run; CPU, pre-commit, data tests, and container smoke are passing. Local validation also passed with 463 passed, 2 skipped, 68 deselected, 10 xfailed.



I approved PR #823. GitHub now reports reviewDecision: APPROVED and mergeable: MERGEABLE. mergeStateStatus is still BLOCKED because of the known GPU-capacity status failures.

Debug info

• Repository: m2lines/Samudra
• Trigger: Simplify the canonical data loading boundary #823 (comment)
• Request ID: 98c7fbe5-7b34-48a0-a2bd-a0bc48ed066d
• Accepted: true