Roll seam-crossing longitudes in the downscaling data layer by frodre · Pull Request #1236 · ai2cm/ace

frodre · 2026-06-06T21:42:33Z

PR 3 of 5 in the prime-meridian longitude stack. Applies the roll primitives (PR 2) in the data layer so a longitude interval that crosses the 0/360 seam can be subset instead of raising NotImplementedError. In-range intervals resolve to a zero roll and behave exactly as before.

Changes:

fme.downscaling.data.datasets.HorizontalSubsetDataset: roll data and coordinates into the requested interval's convention rather than raising on wraparound.
fme.downscaling.data.config: extract _build_aligned_subset_pair, which rolls coarse and fine lon coords into the extent's convention (_roll_lons_to_extent_convention) before adjust_fine_coord_range, so fine/coarse subselection stays aligned across the seam.
fme.downscaling.data.static.StaticInputs.roll: roll static fields and their lon coordinates to match.
fme.downscaling.data.test_config, fme.downscaling.data.test_datasets, fme.downscaling.data.test_static: tests for seam-crossing subsetting (negative and >360 conventions), fine/coarse scale-factor preservation across the seam (even and odd downscale factors), end-to-end paired loader with a seam-crossing extent, and StaticInputs.roll.

Note: surfacing the coarse grid convention on GriddedData/PairedGriddedData (coarse_latlon_coords) was deferred to the integration PR after review discussion.

Tests added
If dependencies changed, "deps only" image rebuilt and "latest_deps_only_image.txt" file updated

Base: feature/lon-roll-primitives (PR 2)

Stack

PR	Head → Base	Title
#1234	`refactor/moe-validate-experts-init` → `main`	Validate expert grid compatibility in `DenoisingMoEPredictor.__init__`
#1235	`feature/lon-roll-primitives` → PR1	Add longitude roll primitives
#1236	`feature/lon-roll-data-layer` → PR2	Roll seam-crossing longitudes in the data layer
#1237	`feature/lon-roll-model` → PR3	Add with_rolled_lon to models
#1238	`feature/lon-roll-integration` → PR4	Roll the model in inference/predict/evaluator

…1234) First in a 5-PR stack adding support for longitude domains that cross the 0/360 prime meridian in downscaling. This standalone hardening PR moves expert grid-compatibility validation into the predictor constructor so every construction path is protected, not just the config-build path: only the primary expert's coordinates are used for input prep and output coords, so an expert built on a mismatched grid would otherwise silently downscale onto the wrong grid. Changes: - `fme.downscaling.predictors.serial_denoising`: move `_validate_experts_compatible` from `DenoisingMoEConfig.build` into `DenoisingMoEPredictor.__init__`, so it holds for `build`, `from_state`, and future callers (e.g. `with_rolled_lon`). - `fme.downscaling.test_models`: add `test_denoising_moe_predictor_rejects_mismatched_expert_grids`, constructing the predictor directly with mismatched-grid experts and asserting it raises. - [x] Tests added - [ ] If dependencies changed, "deps only" image rebuilt and "latest_deps_only_image.txt" file updated Base: `main` ### Stack | PR | Head → Base | Title | |----|-------------|-------| | [#1234](#1234) | `refactor/moe-validate-experts-init` → `main` | Validate expert grid compatibility in `DenoisingMoEPredictor.__init__` | | [#1235](#1235) | `feature/lon-roll-primitives` → PR1 | Add longitude roll primitives | | [#1236](#1236) | `feature/lon-roll-data-layer` → PR2 | Roll seam-crossing longitudes in the data layer | | [#1237](#1237) | `feature/lon-roll-model` → PR3 | Add with_rolled_lon to models | | [#1238](#1238) | `feature/lon-roll-integration` → PR4 | Roll the model in inference/predict/evaluator |

) PR 2 of 5 in the prime-meridian longitude stack. Adds the pure coordinate/data rolling utilities needed to re-express a global grid in a seam-crossing domain's convention. These have no production callers yet — later PRs wire them into the data and model layers — so they are reviewable in isolation with full unit coverage. The interval-based roll only triggers when an interval actually crosses the seam (`start < 0` or `stop > 360`), so in-range intervals are a no-op and non-global grids are left untouched. Primitives overview (PR #1235) These primitives are always used as a pair: find_roll_anchor (or find_roll_anchor_from_interval) computes the roll amount once; callers pass it to all subsequent roll_lon_coords and roll_lon_data so coordinates and field tensors shift by the same amount. Two downstream pathways use them: - Dataset load — rolls each loaded grid into the user's configured lon_extent convention (PR #1236) - Model setup — rolls the model's fine grid to match the incoming coarse batch's convention (PR #1237) Changes: - `fme.downscaling.data.utils`: add `ClosedInterval.finite_values`, `_requires_lon_roll`, `coords_require_lon_roll`, `find_roll_anchor`, `find_roll_anchor_from_interval`, `roll_lon_coords`, `roll_lon_data`, and private helpers `_validate_rollable_lon` and `_validate_monotonic_lon`. - `roll_lon_coords` (1-D coordinate tensor) and `roll_lon_data` (N-D field tensor) form a parallel pair: both apply the same roll amount, but `roll_lon_coords` also remaps values to keep the result monotonically increasing, while `roll_lon_data` is a pure cyclic shift. Callers pre-compute the roll amount once via `find_roll_anchor` and pass it to both. - `roll_latlon_coords` is not included here; it operates on a `LatLonCoordinates` struct rather than a raw tensor and belongs in the PR that first uses it. - `fme.downscaling.data` (`__init__`): export the new roll helpers. - `fme.downscaling.data.test_utils`: unit tests for roll amounts, seam-crossing conventions, round-trip invertibility, non-global/non-uniform rejection, and invalid input validation. - [x] Tests added - [ ] If dependencies changed, "deps only" image rebuilt and "latest_deps_only_image.txt" file updated Base: `refactor/moe-validate-experts-init` (PR 1) ### Stack | PR | Head → Base | Title | |----|-------------|-------| | [#1234](#1234) | `refactor/moe-validate-experts-init` → `main` | Validate expert grid compatibility in `DenoisingMoEPredictor.__init__` | | [#1235](#1235) | `feature/lon-roll-primitives` → PR1 | Add longitude roll primitives | | [#1236](#1236) | `feature/lon-roll-data-layer` → PR2 | Roll seam-crossing longitudes in the data layer | | [#1237](#1237) | `feature/lon-roll-model` → PR3 | Add with_rolled_lon to models | | [#1238](#1238) | `feature/lon-roll-integration` → PR4 | Roll the model in inference/predict/evaluator |

Apply the roll primitives so the data layer can subset longitude domains that cross the 0/360 seam: - HorizontalSubsetDataset now rolls its data and coordinates to the requested interval's convention instead of raising NotImplementedError on wraparound; in-range intervals resolve to a zero roll and behave as before. - StaticInputs.roll rolls static fields and their lon coordinates to match. - BatchItemDatasetAdapter exposes latlon_coordinates, and GriddedData / PairedGriddedData carry coarse_latlon_coords (populated in config) so the coarse grid convention is available to consumers. Adds tests for seam-crossing subsetting (both negative and >360 conventions) and StaticInputs.roll.

adjust_fine_coord_range received unrolled (0-360) longitude tensors, so for an interval like (-16, 30) the coarse_min snapped to ~0 (first 0-360 coord >= -16) instead of -16. This made the computed fine extent too narrow (0-30° instead of -16-30°), causing a scale-factor mismatch error when the paired dataset validated fine vs coarse dimensions. Fix: roll both the coarse and fine lon tensors into the interval's convention before calling adjust_fine_coord_range. The fine anchor is placed one half-coarse-spacing before lon_start so that adjust_fine_coord_range can access the fine half-cells below the first coarse grid point. For non-crossing domains (coarse_roll=0) the roll is a no-op and behaviour is unchanged.

AnnaKwa · 2026-06-11T22:19:49Z

Claude flagged an edge case where longitude intervals crossing 180 should require a roll in the case where lon range is (-180, 180) but not if the max lon is 360. A check on max longitude coord could differentiate these scenarios.

AnnaKwa

Mostly LGTM, could you expand on the data loader test and if possible do a bit of refactoring so the tests of HorizontalSubsetDataset are more clearly associated with the code being tested?

AnnaKwa · 2026-06-11T21:50:39Z

+    assert (
+        batch.coarse.data["var0"].shape[-1] * scale_factor
+        == batch.fine.data["var0"].shape[-1]
+    )


Can you expand this test to have it also check that the coordinates and data values are correct for a rolled data batch and coords? If the data values are just the longitudes this should be straightforward.

AnnaKwa · 2026-06-11T22:10:41Z

+            full_fine_coord=rolled_fine_lon,
        )

        dataset_fine_subset = HorizontalSubsetDataset(


It looks like HorizontalSubsetDataset is what is tested in the new tests; is there a small refactor that can be done here to isolate the code that produces this object into a function or method? That would make it a lot clearer what parts of the code are being tested in the additions to test_datasets.py, and easier to debug if needed in the future.

AnnaKwa · 2026-06-11T22:17:41Z

            variable_metadata=variable_metadata,
            all_times=all_times,
            fine_coords=get_latlon_coords_from_properties(properties_fine),
+            coarse_extent_latlon_coords=dataset_coarse_subset.latlon_coordinates,


Is this used in a later PR?

Yes, I'll move that to the PR it's used in.

…to wt/roll-lon-data-layer

AnnaKwa · 2026-06-12T18:09:46Z

+    # mod 360 (e.g. original 337.5 -> coord -22.5); see the value-level check in
+    # test_build_aligned_subset_pair_preserves_scale_factor_across_seam.
+    for grid in (batch.coarse, batch.fine):
+        lon = grid.latlon_coordinates.lon[0].cpu()  # batch members are identical


Can you also check that the longitude values are consistent with the original lon_extent?

…to wt/roll-lon-data-layer

PR 4 of 5 in the prime-meridian longitude stack (PRs 1–3 now merged to main). Lets a model re-express its grid in a seam-crossing coarse domain's longitude convention while sharing the trained network weights, so a single checkpoint can generate over a domain expressed west of 0 or east of 360. Changes: - `fme.downscaling.models.DiffusionModel.with_rolled_lon`: rebuild the model through its constructor with `full_fine_coords` and `static_inputs` rolled to match the coarse grid, anchored on the western coarse-cell edge so the fine grid stays aligned to whole coarse cells; returns `self` when no roll is needed. Inference-only (rebuilding re-wraps the module under torch distributed). - `fme.downscaling.predictors.serial_denoising.DenoisingMoEPredictor.with_rolled_lon`: roll every expert (preserving the shared-grid invariant) and rebuild so the sigma dispatcher is reconstructed from the rolled experts. - `fme.downscaling.data` exports `roll_lon_coords` for the model layer. - `fme.downscaling.test_models`: tests for no-roll passthrough, coord shifting with shared weights (including value-level checks that coords and static data roll together, and that a double roll is a no-op), and coarse-cell alignment for a seam-crossing domain. MoE rolling tests live in `test_serial_denoising` next to the existing grid-validation test. - Test cleanup: shared `cell_centered_coordinate` helper in `test_utils` replaces per-file midpoint-coordinate constructions (`test_models`, `test_config`); removed a test and helper in `test_models`/`test_serial_denoising` duplicated from #1234. - [x] Tests added - [ ] If dependencies changed, "deps only" image rebuilt and "latest_deps_only_image.txt" file updated Base: `main` (PRs 1–3 of the stack merged) ### Stack | PR | Head → Base | Title | Status | |----|-------------|-------|--------| | [#1234](#1234) | `refactor/moe-validate-experts-init` → `main` | Validate expert grid compatibility in `DenoisingMoEPredictor.__init__` | merged | | [#1235](#1235) | `feature/lon-roll-primitives` → `main` | Add longitude roll primitives | merged | | [#1236](#1236) | `feature/lon-roll-data-layer` → `main` | Roll seam-crossing longitudes in the data layer | merged | | [#1237](#1237) | `feature/lon-roll-model` → `main` | Add with_rolled_lon to models | this PR | | [#1238](#1238) | `feature/lon-roll-integration` → PR4 | Roll the model in inference/predict/evaluator | open |

frodre force-pushed the feature/lon-roll-primitives branch from 12baad6 to 1df468b Compare June 8, 2026 21:57

Base automatically changed from feature/lon-roll-primitives to main June 9, 2026 20:54

frodre added 2 commits June 9, 2026 14:25

frodre force-pushed the feature/lon-roll-data-layer branch from 7455589 to 806b5cd Compare June 9, 2026 21:25

frodre added 10 commits June 9, 2026 15:52

Merge branch 'main' into feature/lon-roll-data-layer

dd4f6e9

Use renamed roll_data_along_lon_dim from merged primitives PR

3bfaaae

Extract lon roll alignment into _roll_lons_to_extent_convention helper

e82e94e

clean up test_config

a9aa95c

horizontal dataset helper

362f621

remove extraneous helper

a1ecf55

Module level imports

45ebf8f

Use helper for coords

86d07b2

Add edge case test for rederived fine coordinate anchor

9ba49bf

Merge branch 'main' into feature/lon-roll-data-layer

7955cc3

frodre marked this pull request as ready for review June 11, 2026 21:13

AnnaKwa requested changes Jun 11, 2026

View reviewed changes

frodre added 6 commits June 11, 2026 15:46

move coordinate pipelining to another PR

f64372f

Add data check to apired data loader test

2279b96

Separate out rolling for targeted test

8923889

Merge branch 'feature/lon-roll-data-layer' of github.com:ai2cm/ace in…

00a9138

…to wt/roll-lon-data-layer

Cleanup comments

7e2bc63

Add value check to helper method test

f33fccb

frodre added 5 commits June 12, 2026 09:45

Parametrize duplicate prime-meridian subset tests

7271c8c

Derive roll amount via find_roll_anchor in StaticInputs.roll test

b78b95a

Add 360 convention check to primitives

aa9f733

Clarify roll-anchor comments with concrete values

6f874fe

Explain why the roll anchor mismatch is harmless

d21e5e7

frodre mentioned this pull request Jun 12, 2026

ACE Aggregator Flaky Test? #1267

Closed

frodre requested a review from AnnaKwa June 12, 2026 17:50

AnnaKwa approved these changes Jun 12, 2026

View reviewed changes

frodre added 3 commits June 12, 2026 12:36

Merge branch 'main' into feature/lon-roll-data-layer

0ac853d

Test cleanup

52bd5ce

Merge branch 'feature/lon-roll-data-layer' of github.com:ai2cm/ace in…

e9d0e43

…to wt/roll-lon-data-layer

frodre enabled auto-merge (squash) June 12, 2026 20:05

frodre merged commit f83caaf into main Jun 12, 2026
7 checks passed

frodre deleted the feature/lon-roll-data-layer branch June 12, 2026 20:17

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Roll seam-crossing longitudes in the downscaling data layer#1236

Roll seam-crossing longitudes in the downscaling data layer#1236
frodre merged 26 commits into
mainfrom
feature/lon-roll-data-layer

frodre commented Jun 6, 2026 •

edited

Loading

Uh oh!

AnnaKwa commented Jun 11, 2026

Uh oh!

AnnaKwa left a comment

Uh oh!

AnnaKwa Jun 11, 2026

Uh oh!

AnnaKwa Jun 11, 2026

Uh oh!

AnnaKwa Jun 11, 2026

Uh oh!

frodre Jun 11, 2026

Uh oh!

AnnaKwa Jun 12, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

frodre commented Jun 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Stack

Uh oh!

AnnaKwa commented Jun 11, 2026

Uh oh!

AnnaKwa left a comment

Choose a reason for hiding this comment

Uh oh!

AnnaKwa Jun 11, 2026

Choose a reason for hiding this comment

Uh oh!

AnnaKwa Jun 11, 2026

Choose a reason for hiding this comment

Uh oh!

AnnaKwa Jun 11, 2026

Choose a reason for hiding this comment

Uh oh!

frodre Jun 11, 2026

Choose a reason for hiding this comment

Uh oh!

AnnaKwa Jun 12, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

frodre commented Jun 6, 2026 •

edited

Loading