Skip to content

Add with_rolled_lon to downscaling models#1237

Draft
frodre wants to merge 1 commit into
feature/lon-roll-data-layerfrom
feature/lon-roll-model
Draft

Add with_rolled_lon to downscaling models#1237
frodre wants to merge 1 commit into
feature/lon-roll-data-layerfrom
feature/lon-roll-model

Conversation

@frodre
Copy link
Copy Markdown
Collaborator

@frodre frodre commented Jun 6, 2026

PR 4 of 5 in the prime-meridian longitude stack. 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. Depends on the roll primitives (PR 2) and the constructor-time expert validation (PR 1).

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.
  • 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.test_models: tests for no-roll passthrough, coord shifting with shared weights, idempotency, coarse-cell alignment, and rolling all MoE experts.
  • Tests added
  • If dependencies changed, "deps only" image rebuilt and "latest_deps_only_image.txt" file updated

Base: feature/lon-roll-data-layer (PR 3)

Stack

PR Head → Base Title
#1234 refactor/moe-validate-experts-initmain 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

Let models re-express their grid in a seam-crossing coarse domain's
longitude convention while sharing network weights:

- DiffusionModel.with_rolled_lon rebuilds the model through its constructor
  with full_fine_coords and static_inputs rolled to match the coarse grid.
  The roll is anchored on the western coarse-cell edge so the fine grid
  stays aligned to whole coarse cells. Returns self when no roll is needed.
- DenoisingMoEPredictor.with_rolled_lon rolls every expert (preserving the
  shared-grid invariant) and rebuilds so the sigma dispatcher is
  reconstructed from the rolled experts.

Adds tests for no-roll passthrough, coord shifting with shared weights,
idempotency, coarse-cell alignment, and rolling all MoE experts.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant