feat: Add support for interpolator from files #335

dnerini · 2025-09-29T08:40:44Z

Description

This PR adds support for running the interpolator on (GRIB) files from a forecaster model.

What issue or task does this change relate to?

N/A

Additional notes

This draft includes all the changes introduced by @OpheliaMiralles in anemoi-inference in order to run interpolation. We should have a discussion on what changes should be factored out in a separate PR (e.g. support for model sharding).

As a contributor to the Anemoi framework, please ensure that your changes include unit tests, updates to any affected dependencies and documentation, and have been tested in a parallel setting (i.e., with multiple GPUs). As a reviewer, you are also responsible for verifying these aspects and requesting changes if they are not adequately addressed. For guidelines about those please refer to https://anemoi.readthedocs.io/en/latest/

By opening this pull request, I affirm that all authors agree to the Contributor License Agreement.

Copilot

Pull Request Overview

This PR adds support for running the interpolator on GRIB files from forecaster models, enabling interpolation on file-based inputs rather than just analysis data.

Key changes:

Added support for interpolation from GRIB files with different reference date handling
Implemented distributed computing support with process group initialization and cleanup
Enhanced pre-processors to work with State objects instead of just FieldList objects

Reviewed Changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
src/anemoi/inference/runners/interpolator.py	Added file-based interpolation support and device management
src/anemoi/inference/runners/default.py	Added distributed process group cleanup and reference date handling
src/anemoi/inference/runner.py	Enhanced with distributed computing, improved tensor operations, and better device handling
src/anemoi/inference/pre_processors/no_missing_values.py	Updated to process State objects instead of FieldList
src/anemoi/inference/pre_processors/forward_transform_filter.py	Updated to process State objects instead of FieldList
src/anemoi/inference/pre_processors/extract.py	New processor for extracting subsets of data using masks or slices
src/anemoi/inference/inputs/gribfile.py	Added support for reference date index parameter
src/anemoi/inference/inputs/ekd.py	Enhanced state creation with reference date support and improved pre-processing
src/anemoi/inference/inputs/cutout.py	Added kwargs support for input state creation
src/anemoi/inference/device.py	Enhanced device selection with distributed computing support
src/anemoi/inference/config/init.py	Relaxed configuration validation to allow extra fields
pyproject.toml	Updated omegaconf version constraint

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

Copilot · 2025-09-29T11:03:27Z

src/anemoi/inference/runners/interpolator.py


        super().__init__(config)
-
+        self.from_analysis = True if hasattr(config.input, "test") else False


Use direct boolean assignment instead of conditional expression. This can be simplified to self.from_analysis = hasattr(config.input, "test").

Suggested change

self.from_analysis = True if hasattr(config.input, "test") else False

self.from_analysis = hasattr(config.input, "test")

Copilot · 2025-09-29T11:03:27Z

src/anemoi/inference/runners/interpolator.py

+                input_state = input.create_input_state(
+                    date=window_start_date, ref_date_index=0
+                )  # for interpolator, the first date is present and the last is future. For AR models with multiple input states, the last date is the current date. This is why the distinction is made here.


The inline comment is too long and detailed for a single line. Consider moving this explanation to a docstring or a separate comment block above the code for better readability.

Suggested change

input_state = input.create_input_state(

date=window_start_date, ref_date_index=0

) # for interpolator, the first date is present and the last is future. For AR models with multiple input states, the last date is the current date. This is why the distinction is made here.

# For interpolator, the first date is present and the last is future.

# For AR models with multiple input states, the last date is the current date.

# This is why the distinction is made here.

input_state = input.create_input_state(

date=window_start_date, ref_date_index=0

)

Copilot · 2025-09-29T11:03:28Z

src/anemoi/inference/runner.py

                # No need to prepare next input tensor if we are at the last step
                if is_last_step:
-                    break
+                    break  # noqa: F701


The noqa comment for F701 (break outside loop) is incorrect here since this break statement is properly inside a for loop. This noqa directive should be removed.

Suggested change

break # noqa: F701

break

Copilot · 2025-09-29T11:03:28Z

src/anemoi/inference/runner.py

+        if check[pmask_in_np].any():
+            # Report which ones are conflicting
+            conflicting = [self._input_tensor_by_name[i] for i in pmask_in_np[check[pmask_in_np]]]
+            raise AssertionError(f"Overwriting prognostic input slots for {conflicting}")


The error message could be more descriptive. Consider adding context about what this means and potential solutions, such as 'Attempting to overwrite prognostic input slots that have already been set for variables: {conflicting}. This may indicate a configuration issue with variable mappings.'

Suggested change

raise AssertionError(f"Overwriting prognostic input slots for {conflicting}")

raise AssertionError(

f"Attempting to overwrite prognostic input slots that have already been set for variables: {conflicting}. "

"This may indicate a configuration issue with variable mappings. "

"Please check your variable configuration to ensure that each prognostic input slot is only set once."

)

Copilot · 2025-09-29T11:03:28Z

src/anemoi/inference/inputs/ekd.py

+        dates = sorted([to_datetime(d) for d in dates])
+        date_to_index = {d.isoformat(): i for i, d in enumerate(dates)}


This code duplicates the sorting and date_to_index creation that was already done at lines 235-237. The duplicated logic should be removed to avoid redundancy and potential inconsistencies.

Suggested change

dates = sorted([to_datetime(d) for d in dates])

date_to_index = {d.isoformat(): i for i, d in enumerate(dates)}

# dates and date_to_index are already sorted and created earlier

gmertes · 2025-09-30T10:55:21Z

Thanks for this! As you say, I think we first need to discuss how to bring this functionality into the codebase. I think this would be a good topic for in two weeks when Harrison and I are in Zurich ?

OpheliaMiralles added 19 commits September 26, 2025 13:26

Changes to pyproject

fdf767e

WIP

922a80a

Set reference date to the first date

7fbee60

Precommit

a385c38

Replace str with torch.device

53605a7

Cleanup

6391edd

Precommit

f40e077

Reset permissions

cb44806

Add extract pre-processors

09135af

Precommit

add4ed6

WIP

4b0e2cc

Reset permissions

176a5a7

Remove debugging changes

02df7df

Fix cutout for interpolator

06728e6

Default date change

fabf63f

Allow interpolator to run from files

79651de

Sharding

7749194

Re-add mps support

737e536

Rebase on main/remove include_forcings

bc963bd

github-project-automation bot added this to Anemoi-dev Sep 29, 2025

github-project-automation bot moved this to To be triaged in Anemoi-dev Sep 29, 2025

github-actions bot added dependencies Pull requests that update a dependency file config labels Sep 29, 2025

dnerini requested a review from frazane September 29, 2025 08:41

dnerini changed the title ~~Add support for interpolator from files~~ feat: Add support for interpolator from files Sep 29, 2025

dnerini assigned dnerini and OpheliaMiralles Sep 29, 2025

dnerini requested a review from Copilot September 29, 2025 11:01

Copilot AI reviewed Sep 29, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat: Add support for interpolator from files #335

feat: Add support for interpolator from files #335

Uh oh!

dnerini commented Sep 29, 2025

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Sep 29, 2025

Uh oh!

Copilot AI Sep 29, 2025

Uh oh!

Copilot AI Sep 29, 2025

Uh oh!

Copilot AI Sep 29, 2025

Uh oh!

Copilot AI Sep 29, 2025

Uh oh!

gmertes commented Sep 30, 2025

Uh oh!

Uh oh!


		super().__init__(config)

		self.from_analysis = True if hasattr(config.input, "test") else False

	self.from_analysis = True if hasattr(config.input, "test") else False
	self.from_analysis = hasattr(config.input, "test")

-            raise AssertionError(f"Overwriting prognostic input slots for {conflicting}")
+            raise AssertionError(
+                f"Attempting to overwrite prognostic input slots that have already been set for variables: {conflicting}. "
+                "This may indicate a configuration issue with variable mappings. "
+                "Please check your variable configuration to ensure that each prognostic input slot is only set once."
+            )

		dates = sorted([to_datetime(d) for d in dates])
		date_to_index = {d.isoformat(): i for i, d in enumerate(dates)}

	dates = sorted([to_datetime(d) for d in dates])
	date_to_index = {d.isoformat(): i for i, d in enumerate(dates)}
	# dates and date_to_index are already sorted and created earlier

feat: Add support for interpolator from files #335

Are you sure you want to change the base?

feat: Add support for interpolator from files #335

Uh oh!

Conversation

dnerini commented Sep 29, 2025

Description

What issue or task does this change relate to?

Additional notes

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Copilot AI Sep 29, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Sep 29, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Sep 29, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Sep 29, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Sep 29, 2025

Choose a reason for hiding this comment

Uh oh!

gmertes commented Sep 30, 2025

Uh oh!

Uh oh!