Add LLC CPU data canonicalization by jder · Pull Request #670 · m2lines/Samudra

jder · 2026-04-08T19:36:20Z

Add LLC data support (preprocessing, etc). Brings time configuration down into the data source since LLC data has different time formats than OM4, and we likely want to use different time ranges even for different resolutions of OM4 data.

It looks big but a lot of it is mechanical config changes and tests; happy to walk through live and/or break up further if desired!

jder · 2026-05-07T14:13:51Z

@fomo-bot can you please rebase this PR into a single clean commit on top of the latest main?

fomo-bot · 2026-05-07T14:17:08Z

Done. PR #670 is now a single commit on top of latest main:

47d5ed9e Add LLC CPU data canonicalization

I force-pushed with lease to llc-stack-03-llc-cpu-canonicalization. Local validation passed:

uv run pytest tests/test_config.py tests/test_utils_data.py -m "not manual and not cuda"

37 passed, 1 warning.

GitHub checks have started and were still queued when I checked.

Codex work

Repository: Open-Athena/Ocean_Emulator
Trigger: https://github.com/Open-Athena/Ocean_Emulator/pull/670#issuecomment-4397887329
Request ID: ff1d0274-40c6-4df3-b5e7-655a457a6cfd
Accepted: true

Codex transcript

- I’ll inspect the current branch state, fetch the latest base branch, then squash the PR changes into one commit and update the remote branch with a lease-protected force push.
- The branch is clean and the PR targets `main`. I see the PR history includes several older stack commits plus merge commits, so I’m fetching `origin/main` now and will squash the actual branch diff against the fresh base rather than preserving that noisy ancestry.
- I’m going to keep a local backup ref at the current branch tip, then replay only the three LLC canonicalization commits onto `origin/main`. That should avoid carrying the old pre-merge dataset-spec history into the rewritten PR.
- The replay onto latest `origin/main` completed without conflicts. I’m checking the rewritten history and staged diff now before collapsing it to one commit.
- The rebased branch now has exactly the three LLC commits on top of current `origin/main`, with a five-file diff. I’m squashing those into a single commit named after the PR.
- The branch is now one local commit on top of current `origin/main`. I’m going to run a focused validation on the affected config/data tests before pushing the rewritten branch.
- `uv` had to create the local virtualenv, so the focused pytest run is still in progress. I’ll wait for that result before updating the remote branch.
- The focused tests passed: 37 tests in `tests/test_config.py` and `tests/test_utils_data.py`. I’m doing a final status check, then I’ll force-push with lease to update the PR branch.
- The force-with-lease push succeeded, updating the PR branch from `5c2d9008` to `47d5ed9e`. I’m verifying GitHub now reports the PR as a single commit on the rewritten branch.
- Done. PR #670 is now a single commit on top of latest `main`:

`47d5ed9e Add LLC CPU data canonicalization`

I force-pushed with lease to `llc-stack-03-llc-cpu-canonicalization`. Local validation passed:

`uv run pytest tests/test_config.py tests/test_utils_data.py -m "not manual...[truncated]

jder · 2026-05-29T18:41:24Z

@fomo-bot can we unify the two different mock LLC datasets in tests, please?

I also don't really understand all the complexity around the time stuff. Why do we need to change slice if we are already remapping the times in the canonicalize function?

fomo-bot · 2026-05-29T18:52:23Z

Done and pushed.

The PR is still a single commit, now rebased onto current origin/main:

0d2e5f36 Add LLC CPU data canonicalization

I unified the LLC test data in tests/llc_fixtures.py and removed the generic numeric-time DataSource.slice changes from data.py. The LLC path now relies on canonicalize_llc_datasets to remap raw LLC times to Julian cftime, so slice stays on the existing xarray label-slice path.

Local checks passed:

uv run pytest tests/test_config.py tests/test_utils_data.py -m "not manual and not cuda"

uvx pre-commit run --files ...

GitHub checks are running; a couple of data jobs had already passed when I checked, with the longer CI jobs still in progress.

Codex work

Repository: Open-Athena/Ocean_Emulator
Trigger: https://github.com/Open-Athena/Ocean_Emulator/pull/670#issuecomment-4578675510
Request ID: 08f257b8-3d05-449e-afed-671c6b4ca592
Accepted: true

Codex transcript

- I’ll inspect the current test helpers and the LLC canonicalization path first, then simplify the tests and the time handling where the code proves it’s redundant.
- The local branch is clean, but the local `origin/main` ref is stale enough that its comparison is misleading. I’m going to fetch current refs before editing so this cleanup lands on the actual current PR state.
- `main` has moved six commits since this PR was last rebased. I’m going to make the requested cleanup against the PR diff, then rebase the single commit onto current `origin/main` if the touched files still merge cleanly.
- I found the duplication: one LLC fixture writes zarr/netcdf in `test_config.py`, while another builds an in-memory copy in `test_utils_data.py`. I’m moving that to one shared test helper, and I’m reverting the generic `DataSource.slice` time-coordinate coercion so LLC time no...[truncated]
- The code now leaves `DataSource.slice` on the existing xarray label-slice implementation and keeps only the LLC-specific time conversion in `canonicalize_llc_datasets`. I’m checking formatting/import fallout before running the focused tests.
- `ruff` is not installed in this uv environment, so I’ll rely on `git diff --check` plus the focused pytest run for now. The cleanup is otherwise syntactically small: the new shared helper is untracked until I stage/amend.
- The focused tests pass with the shared fixture and without the generic numeric-time slice path. I’m amending the single PR commit now, then rebasing that one commit onto today’s `origin/main` because the PR is currently behind.
- The rebase onto current `origin/main` completed without conflicts. I’m rerunning the focused tests on the rebased commit and checking the final one-commit shape before pushing.
- The rebased focused tests pass. `pre-commit` is not installed in the project environment, so I’m trying it through `uvx` to catch the CI pre-commit checks before I push.
- The pre-commit run has passed ruff, formatting, YAML, and basic file checks; it is currently on mypy. I’ll let it finish because this is the same class of check that previously failed on the PR.
- Pre-commit passed on the touched files, including reuse/license lint. I’m force-pushing the rewritten one-commit branch with lease now.
- The branch was updated on GitHub from `47d5ed9e` to `0d2e5f36`. I’m verifying the PR metadata and check status now.
- GitHub checks are running now. I started a short watch to catch any immediate failures, but I won’t block on the longer EC2/GPU jobs unless they fail quickly.
- Done and pushed.

The PR is still a single commit, now rebased onto current `origin/main`:

`0d2e5f36 Add LLC CPU data canonicalization`

I unified the LLC test data in [tests/llc_fixtures.py](/home/fomo-bot/oa-jder-bot/worktrees/github.com/Open-Athena/Ocean_Emulator/670/tests/ll...[truncated]

jder · 2026-07-15T17:57:49Z

@fomo-bot can you rebase this on the latest main please

fomo-bot · 2026-07-15T18:05:56Z

Rebased PR branch llc-stack-03-llc-cpu-canonicalization onto latest origin/main (8fe8b080) and pushed it with --force-with-lease.

New branch head: d304b256

Validation:

uv run pytest for the LLC-specific tests: 5 passed
uv run --with ruff ruff check ...: passed
Broader focused pytest run had 156 passed and 12 errors, all from missing optional flash_perceiver dependency when perceiver_implementation=auto selected flash.

Debug info

Repository: m2lines/Samudra
Trigger: Add LLC CPU data canonicalization #670 (comment)
Request ID: c801366d-6049-4951-b8b6-bbea975a201b
Accepted: true

jder · 2026-07-16T20:11:21Z

@fomo-bot please fix precommit

fomo-bot · 2026-07-16T20:12:55Z

Fixed pre-commit and pushed b6df79a9.

The failure was ruff-format; it reformatted one line in src/samudra/utils/data.py.

Verification:

uv run --with pre-commit pre-commit run --all-files: passed
Branch is clean and synced with origin/llc-stack-03-llc-cpu-canonicalization.

Debug info

Repository: m2lines/Samudra
Trigger: Add LLC CPU data canonicalization #670 (comment)
Request ID: 757c93bf-d366-497c-a9f4-b6c72b490174
Accepted: true

alxmrs

Looks really great! I mostly have questions about factoring (where to put the new things), and then some conversations with myself to understand the PR.

Thanks for this change, this is really valuable.

alxmrs · 2026-07-21T22:09:09Z


 if TYPE_CHECKING:
-    from samudra.config import TimeConfig
+    from samudra.config import SourceTimeConfig


Is there also a non-source TimeConfig that we're keeping around? If not, then should we just keep the name "TimeConfig"? WDYT?

Yeah, makes sense.

alxmrs · 2026-07-21T22:10:47Z

        )


+def _rename_llc_level_index_vars(ds: xr.Dataset) -> xr.Dataset:


I hope there are roundtrip tests for these. Nice to see a function and its inverse.

alxmrs · 2026-07-21T22:13:06Z

+    return mask_vars
+
+
+def canonicalize_llc_datasets(


Idea: should we start to create data source specific utility files? e.g. keep data.py generic, and then separate out OM4 vs llc specific utilities? This might help us create better abstractions later on.

We could also do this after one more dataset gets added (to channel my inner YAGNI).

It seems like we're fairly close to creating new overridable abstract methods in DatasetSpec for these kinds of functions, but not quite yet.

Looks like this somewhat exists in config.py. Nice.

Yeah, I can break out the LLC-specific stuff.

Co-authored-by: Alex Merose <alex@openathena.ai>

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: ee56e6c977

ℹ️ About Codex in GitHub

Your team has set up Codex to 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 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

chatgpt-codex-connector · 2026-07-22T14:59:53Z

+            train=source.slice(self.train_time),
+            val=source.slice(self.val_time),


Avoid train/val slicing during eval-only builds

For eval-only runs whose data source contains only the configured inference period, Eval.__init__ calls cfg.data.build() and then only consumes data_container.inference_source, but this return path now slices the train and validation windows for every source first. Because DataSource.slice() raises when those windows are outside the store, standalone evaluation fails before reaching inference even though the train/val splits are unused; defer these slices to training or let eval build just the inference source plus unsliced source metadata.

Useful? React with 👍 / 👎.

jder · 2026-07-22T15:26:46Z

@alxmrs I think I did all the things above, RFAL!

alxmrs

LGTM! Thanks!

…canonicalization # Conflicts: # src/samudra/train.py # tests/conftest.py # tests/test_datasets.py

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 1c26183151

ℹ️ About Codex in GitHub

Your team has set up Codex to 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 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

chatgpt-codex-connector · 2026-07-23T00:21:31Z

+
+        n_levels = data_copy[name].sizes["lev"]
+        expected_levels = len(dataset_spec.depth_i_levels)
+        if n_levels != expected_levels:


Select requested LLC levels instead of rejecting full-depth stores

When a raw LLC store contains more k/lev levels than the model spec consumes (for example a full-depth root while single_1 only needs Theta_0), this equality check raises before any requested level can be selected. That makes otherwise usable LLC data/configs fail during canonicalization; flatten the configured level indices and only error when a needed level is missing.

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-07-23T00:21:31Z

+    requested_data_vars.update(
+        [dataset_spec.mask_all_levels_var, "mask_c", *dataset_spec.mask_vars]
+    )


Keep hFacC as a center wet-mask fallback

When an LLC store uses hFacC for tracer-cell wetness—the center-grid counterpart to the hFacW/hFacS fallbacks handled below—this whitelist drops it before the rename to wetmask, so flatten_masks() later cannot build wetmask_* and LLC loading fails. Include hFacC and rename it like mask_c for stores that do not already provide mask_c.

Useful? React with 👍 / 👎.

@fomo-bot can you please verify if our LLC mock data includes hFacC? If it does, please fix this comment. Either way let's then get this merged.

Verified: the LLC mock fixture includes mask_c, hFacW, and hFacS, but not hFacC, so I did not change the comment.

PR #670 was approved, checks passed, and I added it to the merge queue. It merged successfully at 2026-07-24T13:31:33Z: #670

Debug info

Repository: m2lines/Samudra

Trigger: Add LLC CPU data canonicalization #670 (comment)

Request ID: 01fa0742-b13e-4d3f-98a4-4a2b2b3b058a

Accepted: true

jder mentioned this pull request Apr 8, 2026

Enable LLC training configs #671

Draft

oa-jder-bot force-pushed the llc-stack-03-llc-cpu-canonicalization branch from 509b762 to 91815a4 Compare April 8, 2026 20:01

oa-jder-bot force-pushed the llc-stack-02-dataset-specs branch from d201ce9 to 8aa8e76 Compare April 8, 2026 21:07

oa-jder-bot force-pushed the llc-stack-03-llc-cpu-canonicalization branch from 91815a4 to 7440baa Compare April 8, 2026 21:07

fomo-bot force-pushed the llc-stack-02-dataset-specs branch from 8aa8e76 to d738b00 Compare April 13, 2026 15:22

fomo-bot force-pushed the llc-stack-03-llc-cpu-canonicalization branch 4 times, most recently from 6aa69de to 35caa19 Compare April 15, 2026 21:12

fomo-bot force-pushed the llc-stack-03-llc-cpu-canonicalization branch 2 times, most recently from a21e9c9 to 4276e4a Compare April 22, 2026 19:42

fomo-bot force-pushed the llc-stack-03-llc-cpu-canonicalization branch from 4276e4a to 5c2d900 Compare April 29, 2026 18:19

Base automatically changed from llc-stack-02-dataset-specs to main April 30, 2026 18:48

fomo-bot force-pushed the llc-stack-03-llc-cpu-canonicalization branch from 5c2d900 to 47d5ed9 Compare May 7, 2026 14:16

fomo-bot force-pushed the llc-stack-03-llc-cpu-canonicalization branch from 47d5ed9 to 0d2e5f3 Compare May 29, 2026 18:50

jder commented May 29, 2026

View reviewed changes

Comment thread src/ocean_emulators/utils/data.py Outdated

fomo-bot force-pushed the llc-stack-03-llc-cpu-canonicalization branch from 0d2e5f3 to dece93c Compare May 29, 2026 20:22

jder commented Jun 2, 2026

View reviewed changes

Comment thread src/ocean_emulators/utils/data.py Outdated

fomo-bot force-pushed the llc-stack-03-llc-cpu-canonicalization branch from dece93c to 0c477e4 Compare June 12, 2026 14:59

alxmrs mentioned this pull request Jul 9, 2026

Train a Samudra model to submit to OceanBench #786

Open

This was referenced Jul 10, 2026

implement core replay buffer logic #792

Open

add shardtensor-op-probe #793

Open

fomo-bot force-pushed the llc-stack-03-llc-cpu-canonicalization branch from fc087cd to d304b25 Compare July 15, 2026 18:05

jder marked this pull request as ready for review July 15, 2026 20:31

comment

50466b3

jder mentioned this pull request Jul 16, 2026

Untangle i/j/x/y/lat/lon given LLC datasets too #801

Open

Reject LLC static data vars

16a2638

fomo-bot force-pushed the llc-stack-03-llc-cpu-canonicalization branch from 379d3ab to 16a2638 Compare July 16, 2026 19:09

Use LLC staggered masks

7649a17

Format LLC mask helper

b6df79a

jder requested a review from alxmrs July 16, 2026 20:12

alxmrs reviewed Jul 21, 2026

View reviewed changes

Update src/samudra/config.py

ee56e6c

Co-authored-by: Alex Merose <alex@openathena.ai>

chatgpt-codex-connector Bot reviewed Jul 22, 2026

View reviewed changes

jder added 2 commits July 22, 2026 11:21

PR feedback

7799e54

move

56d2d13

jder mentioned this pull request Jul 22, 2026

Delete "static data" #816

Closed

alxmrs approved these changes Jul 22, 2026

View reviewed changes

jder enabled auto-merge July 22, 2026 22:52

jder disabled auto-merge July 22, 2026 23:46

Merge remote-tracking branch 'origin/main' into llc-stack-03-llc-cpu-…

1c26183

…canonicalization # Conflicts: # src/samudra/train.py # tests/conftest.py # tests/test_datasets.py

jder force-pushed the llc-stack-03-llc-cpu-canonicalization branch from ba58035 to 1c26183 Compare July 23, 2026 00:12

jder enabled auto-merge July 23, 2026 00:12

chatgpt-codex-connector Bot reviewed Jul 23, 2026

View reviewed changes

jder disabled auto-merge July 23, 2026 00:25

fomo-bot added this pull request to the merge queue Jul 24, 2026

Merged via the queue into main with commit e68ca1a Jul 24, 2026
20 checks passed

fomo-bot deleted the llc-stack-03-llc-cpu-canonicalization branch July 24, 2026 13:31

This was referenced Jul 24, 2026

Fix LLC center-mask fallback and fixture schema #820

Merged

Simplify the canonical data loading boundary #823

Draft

		)


		def _rename_llc_level_index_vars(ds: xr.Dataset) -> xr.Dataset:

		train=source.slice(self.train_time),
		val=source.slice(self.val_time),

Uh oh!

Conversation

jder commented Apr 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jder commented May 7, 2026

Uh oh!

fomo-bot commented May 7, 2026

Uh oh!

jder commented May 29, 2026

Uh oh!

fomo-bot commented May 29, 2026

Uh oh!

Uh oh!

Uh oh!

jder commented Jul 15, 2026

Uh oh!

fomo-bot commented Jul 15, 2026

Uh oh!

jder commented Jul 16, 2026

Uh oh!

fomo-bot commented Jul 16, 2026

Uh oh!

alxmrs left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Jul 22, 2026

Choose a reason for hiding this comment

Uh oh!

jder commented Jul 22, 2026

Uh oh!

alxmrs left a comment

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

jder commented Apr 8, 2026 •

edited

Loading