Add fine coordinates to the model for easier inference handling by frodre · Pull Request #971 · ai2cm/ace

frodre · 2026-03-13T22:41:45Z

DiffusionModel previously relied on StaticInput.coords to store the fine-resolution lat/lon grid, coupling spatial metadata to individual topography fields. This made coordinate handling awkward since models without any static inputs had no coordinate information and would either fail or required a _downscale_coord bandaid to approximate fine-resolution coordinates from coarse ones.

Changes:

StaticInputs now carries a required coords: LatLonCoordinates field representing the full fine-resolution domain. coords is removed from individual StaticInput fields.
DiffusionModel always receives a non-optional StaticInputs (fields may be empty when no static data is needed). full_fine_coords is a property delegating to static_inputs.coords.
DiffusionModelConfig.build and DiffusionModel.__init__ no longer accept None for static_inputs.
StaticInputs.subset(lat_interval, lon_interval) replaces subset_latlon, computing slices internally and subsetting both fields and coords together.
ClosedInterval.slice_of renamed to slice_from; new subset_of convenience method returns the coordinate values within the interval.
All checkpoint backwards-compatibility loading logic is consolidated in StaticInputs.from_state_backwards_compatible in static.py. CheckpointModelConfig gains an optional fine_coordinates_path for old checkpoints with no stored coordinates.
load_fine_coords_from_path added to fme.downscaling.data and tested.
PairedGriddedData carries fine_coords, passed to the model at build time to use as a fallback if no static input's with coordinates in the dataset are specified.
Removed _downscale_coord from predict.py; model.get_fine_coords_for_batch is used instead.
Tests added

This PR finalizes the removal of the `StaticInput` handling by the data pipeline. The passing of static_input objects are removed from the data configuration, batch iteration, and model call signatures in favor of the direct model handling introduced in the previous downscaling PR (#954). Changes: - add `get_fine_coords_for_batch` to facilitate translation of an input batch domain to output coordinates via the models stored information. For now, this relies on the model's `static_inputs`, but will be switched to model's stored coordinates in (#971) - inference `Downscaler` now takes the batch `input_shape` instead of `static_inputs` to check the domain size and model type (regular `DiffusionModel` or `PatchPredictor` - downscaling `torch.datasets` generators for `BatchData` no longer include `StaticInputs` - removed `_apply_patch` and `_generate_from_patches` from `StaticInputs` - `config.py` no longer references static inputs as an argument - [x] Tests added

frodre · 2026-03-17T22:54:18Z

fme/downscaling/models.py

+        # Load fine_coords: new checkpoints store it directly; old checkpoints
+        # that had static_inputs with coords can auto-migrate from raw state.
+        fine_coords = state.get("fine_coords")
+        if fine_coords is not None:


This block is the pathway executed during training resumption and will fail if we try and resume training for a model without any static inputs or fine coordinates. Not totally sure if this sort of backwards compat is really necessary, since we'll likely just be training new models w/ fine_coords.

I'm ok with breaking backwards compatibility here. I think the oldest checkpoint we'd possibly want to continue to use would be the released precip-only model checkpoint, which has static inputs saved.

AnnaKwa

I agree with adding a fine coords attribute to the model and saving it as well. I do prefer to keep the the coordinates on the StaticInputs class though as the main use of the fine coords in the model is to subset the static inputs, and having them in that class this keeps the subsetting of that tensor cleaner while making it 100% clear that the coordinates are associated with that set of static inputs. Otherwise my worry in separating the coordinates from the static inputs is that down the line it may be easier to introduce a bug where the coordinates don't match.

Could the model attribute instead default to point to the static inputs fine coordinates and be set by an optional checkpoint config path if required for a model w/o static inputs?

AnnaKwa · 2026-03-18T18:28:29Z

fme/downscaling/models.py

                coarse to fine.
            sigma_data: The standard deviation of the data, used for diffusion
                model preconditioning.
+            fine_coords: the full-domain fine-resolution coordinates to use


Suggestion: call this full_fine_coords or something to that effect so it's obvious that this is the full domain and doesn't need to be updated if training is resumed on some different domain.

AnnaKwa · 2026-03-18T18:31:11Z

fme/downscaling/models.py

        self,
        coarse_shape: tuple[int, int],
        downscale_factor: int,
+        fine_coords: LatLonCoordinates,


For this and other places where this arg is added, see comment about naming to describe it's the full domain.

fme/downscaling/models.py

AnnaKwa · 2026-03-18T18:53:24Z

fme/downscaling/models.py

        coarse_shape: tuple[int, int],
        downscale_factor: int,
        sigma_data: float,
+        fine_coords: LatLonCoordinates,


Suggestion: make this arg optional, only to be used if building from a checkpoint that has an optional fine_coords_path arg to allow for models w/o static inputs to have correct coords in the saved predict/evaluate output. Otherwise for the standard case where the model uses static inputs, set the attribute self.fine_coords using the static inputs coords.

Co-authored-by: Anna Kwa <annak@allenai.org>

frodre · 2026-03-19T18:19:36Z

fme/core/coordinates.py

        return f"LatLonCoordinates(\n    lat={self.lat},\n    lon={self.lon}\n"

-    def to(self, device: str) -> "LatLonCoordinates":
+    def to(self, device: str | torch.device) -> "LatLonCoordinates":


To make my VSCode linter happy

frodre · 2026-03-19T23:18:28Z

fme/downscaling/data/static.py

 @dataclasses.dataclass
 class StaticInputs:
    fields: list[StaticInput]
+    coords: LatLonCoordinates


Not named full_coords because we do produce subsets with this class.

AnnaKwa · 2026-03-20T16:24:17Z

Ah, I'm very sorry for the confusion- when I said "I do prefer to keep the the coordinates on the StaticInputs" I meant I preferred them kept where there were currently were (within the static inputs but on individual StaticInput objects not the higher level StaticInputs). The potential mismatch in coords I was concerned about in the previous iteration was between the static inputs coords (post init ensures all StaticInput objects have same coords) and the DiffusionModel's coords attribute getting out of sync through mixing of saved checkpoints and updated configs. I think the static inputs coords are best kept in the lowest level object StaticInput so it's clear they are associated with the data information there.

PairedGriddedData carries fine_coords, passed to the model at build time.

Is this so that models without static inputs have the fine coords information? Could we instead set a _fine_coords_from_gridded_data attribute and fall back to using this if there are no static inputs (rather than setting the static inputs coords from the fine dataset coords)? See comment on the full_fine_coords property.
It should amount the same thing since the fine coords are usually the same from both sources, but this way it's 100% guaranteed that the static inputs coords are direct from their underlying datasets.

AnnaKwa · 2026-03-20T16:24:30Z

fme/downscaling/models.py

-        return self.static_inputs.subset_latlon(lat_interval, lon_interval)
+    @property
+    def full_fine_coords(self) -> LatLonCoordinates:
+        return self.static_inputs.coords


Could we do something like have an attribute self._full_fine_coords_from_gridded_data: None | LatLonCoordinates that is set in the build method, and then

if len(self.static_inputs.fields) > 0: # or maybe add a len method to the class return self.static_inputs.coords else: return self._full_fine_coords_from_gridded_data

Can the fallback be put here, rather than within the StaticInputs?

…tic inputs

AnnaKwa · 2026-03-20T20:06:25Z

fme/downscaling/models.py

-        return self.static_inputs.subset_latlon(lat_interval, lon_interval)
+    @property
+    def full_fine_coords(self) -> LatLonCoordinates:
+        return self.static_inputs.coords


Can the fallback be put here, rather than within the StaticInputs?

AnnaKwa · 2026-03-20T20:13:28Z

fme/downscaling/data/static.py

+        # no coords found with static inputs, use provided fallback
+        coords_to_use = fallback_coords
+    elif validate_coords:
+        _validate_coords("fallback", coords_to_use, fallback_coords)


Could this validation also get moved up to be done at the model level when it gets built?

AnnaKwa · 2026-03-20T20:34:08Z

fme/downscaling/data/static.py

+    try:
+        coords = _load_coords_from_ds(ds)
+    except ValueError:
+        # no coords available


Up to you, but in the data loading (get_horizontal_coordinates) it's assumed the last two dims are lat, lon so I think it's ok to assume this here as well. If for some reason there are no usable coords I think we should just raise the error here (i.e. if there are static inputs, we expect them to have valid coords).

frodre mentioned this pull request Mar 13, 2026

Remove static_input from data pipeline and model call signatures #956

Merged

1 task

frodre force-pushed the feature/downscaling-model-fine-coords branch 3 times, most recently from c7ff073 to a97e954 Compare March 16, 2026 22:59

frodre changed the base branch from refactor/remove-static-input-from-data-and-call-sigs to main March 16, 2026 23:07

frodre changed the base branch from main to refactor/remove-static-input-from-data-and-call-sigs March 16, 2026 23:08

frodre force-pushed the feature/downscaling-model-fine-coords branch 2 times, most recently from 35c2234 to 229e858 Compare March 16, 2026 23:36

frodre changed the base branch from refactor/remove-static-input-from-data-and-call-sigs to main March 16, 2026 23:44

frodre changed the base branch from main to refactor/remove-static-input-from-data-and-call-sigs March 16, 2026 23:44

Base automatically changed from refactor/remove-static-input-from-data-and-call-sigs to main March 17, 2026 21:01

frodre added 10 commits March 17, 2026 14:24

Initial shot

754f28a

Make fine coords required

7c5b4c5

Fine coords required for paired data

c7f58d9

Mesh with previous updates in refactor pr

563b57a

Simplify event downscaler coordinate in run()

f4218dd

use batch latlon coardinates for coarse

68cab61

Make fine coord loader public

83ca043

BatchLatLon coord access consistency

794a7d4

linting

b542080

Add no coords checkpoint with path test

5add727

frodre force-pushed the feature/downscaling-model-fine-coords branch from ae4b09e to 5add727 Compare March 17, 2026 21:24

frodre added 2 commits March 17, 2026 15:37

Small tweaks

36e80db

Add load_fine_coords_from_path test

b19c9d6

frodre commented Mar 17, 2026

View reviewed changes

frodre marked this pull request as ready for review March 17, 2026 22:58

AnnaKwa requested changes Mar 18, 2026

View reviewed changes

frodre and others added 2 commits March 18, 2026 13:20

Update fme/downscaling/models.py

7acf87c

Co-authored-by: Anna Kwa <annak@allenai.org>

Update fme/downscaling/models.py

eba2749

Co-authored-by: Anna Kwa <annak@allenai.org>

frodre and others added 3 commits March 18, 2026 13:20

Update fme/downscaling/models.py

2c25f1d

Co-authored-by: Anna Kwa <annak@allenai.org>

use latlon coords .to method for device fix

81e4eb8

Redo based on Anna's comments

325d324

frodre commented Mar 19, 2026

View reviewed changes

frodre added 9 commits March 19, 2026 11:28

Remove from_state docstring

1fac745

Move all state loading cases into static inputs code

3ef1613

Remove unused function from fme.downscaling.data

36e8101

fine_coords -> full_fine_coords

3f95457

Remove duplicated tests from models.py

6c9fcf7

Minor fixes

4b58b86

Fix imports

7e36f55

Final cleanup

a85fb05

Merge branch 'main' into feature/downscaling-model-fine-coords

b1ee3d4

frodre commented Mar 19, 2026

View reviewed changes

frodre requested a review from AnnaKwa March 19, 2026 23:22

AnnaKwa reviewed Mar 20, 2026

View reviewed changes

Add coordinate validation and use training data as a fallback vs. sta…

bedd276

…tic inputs

AnnaKwa reviewed Mar 20, 2026

View reviewed changes

frodre added 2 commits March 20, 2026 17:43

Updates based on discussion w/ Anna

a9dbae3

Merge branch 'main' into feature/downscaling-model-fine-coords

46670ff

This was referenced Mar 21, 2026

Add ClosedInterval subsetting and update StaticInput Subset #997

Draft

Update StaticInputs coordinate handling #998

Draft

Add fine_coordinates to model with validation against StaticInputs #999

Draft

Conversation

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

Uh oh!

Changes:

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AnnaKwa left a comment

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!

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!

AnnaKwa commented Mar 20, 2026

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

frodre commented Mar 13, 2026 •

edited

Loading