Data Assimilation Models¶
Data assimilation models in Earth2Studio combine observations with a model state or background estimate to produce an updated state. They are useful when a workflow needs to ingest observation-like data, correct a forecast or analysis state, or prepare an initial condition for a downstream model.
Unlike Prognostic Models, data assimilation models are not primarily responsible for rolling a forecast forward in time. Unlike Diagnostic Models, they are usually driven by observation batches or state-update logic rather than deriving a single physical field from an existing tensor state.
The list of data assimilation models that are already built into Earth2Studio can be found in the API documentation earth2studio.models.da.
Data Assimilation Interface¶
The full requirements for a standard data assimilation model are defined explicitly in
earth2studio/models/da/base.py.
@runtime_checkable
class AssimilationModel(Protocol):
"""Data assimilation model interface"""
def __call__(
self,
*args: pd.DataFrame | xr.DataArray | None,
) -> tuple[pd.DataFrame | xr.DataArray, ...]:
"""Stateless iteration for the data assimilation model.
Processes observations and returns assimilated data without maintaining
internal state between calls. This method is suitable for independent
processing of observation batches.
Parameters
----------
*args : pd.DataFrame | xr.DataArray | None
Variable number of observation arguments. Each argument can be a
DataFrame (pandas or cudf DataFrame) or xarray DataArray
containing observation data. None can be passed for optional
arguments when no input data is available.
Returns
-------
tuple[pd.DataFrame | xr.DataArray, ...]
Assimilated data output. Can return a combination of DataFrames or
xarray DataArrays depending on the particular model. Output is expect to be
on the same device as the model.
"""
def create_generator(
self,
*args: pd.DataFrame | xr.DataArray,
) -> Generator[
tuple[pd.DataFrame | xr.DataArray, ...],
tuple[pd.DataFrame | xr.DataArray | None, ...],
None,
]:
"""Creates a generator which accepts collection of input observations and
outputs a collection of assimilated data. Used for both stateless and stateful
iterations of the data assimilation model
The generator accepts observations (DataFrame or DataArray) via the send()
method and yields assimilated data (DataFrame or DataArray) as output.
Supports any number of arguments (variadic).
Parameters
----------
*args : pd.DataFrame | xr.DataArray
Variable number of initialization arguments, if any are required by
the model. Each argument can be a DataFrame (pandas or cudf
DataFrame) or xarray DataArray containing initial state data.
Yields
------
tuple[pd.DataFrame | xr.DataArray, ...]
Generator yields multiple arguments of assimilated data. Each argument
can be a DataFrame (PyArrow Table or cudf DataFrame) or xarray DataArray.
Supports any number of arguments.
Receives
--------
tuple[pd.DataFrame | xr.DataArray | None, ...]
Observations sent via generator.send() as multiple arguments. Each
argument can be a DataFrame (PyArrow Table or cudf DataFrame) or xarray
DataArray. None is sent initially to start the generator and can also be
sent for iterations where no input data is available. Supports any number
of arguments.
Examples
--------
>>> generator = model.create_generator()
>>> generator.send(None) # Prime the generator
>>> # Process observations over time
>>> for obs in observations:
... result = generator.send(obs) # Send observations, receive assimilated data
... # result is a tuple of DataFrames or DataArrays
>>> generator.close() # Clean up
"""
pass
def init_coords(self) -> tuple[FrameSchema | CoordSystem, ...] | None:
"""Initialization coordinate system required by the assimilation model.
Specifies the coordinate system(s) for initial state data that must be provided
before the model can process observations. The returned coordinate systems should
match the expected input format for the first argument(s) passed to ``__call__``
or sent to ``create_generator`` when initializing the model state.
Returns
-------
tuple[FrameSchema | CoordSystem, ...] | None
Tuple of coordinate systems or frame schemas defining the structure of
required initialization data. Returns ``None`` if the model does not require
initialization data (e.g., stateless models).
"""
pass
def input_coords(self) -> tuple[FrameSchema | CoordSystem, ...]:
"""Input coordinate system of assimilation model.
For DataFrame inputs, this should return a PyArrow schema (or a wrapper
containing schema and constraints). For tensor inputs, this should return
a CoordSystem.
Returns
-------
tuple[FrameSchema | CoordSystem, ...]
Tuple of coordinate systems or frame schemas, one for each input argument
that __call__ or create_generator accepts
"""
pass
def output_coords(
self,
input_coords: tuple[FrameSchema | CoordSystem, ...],
*args: Any,
**kwargs: Any,
) -> tuple[FrameSchema | CoordSystem, ...]:
"""Output coordinate system of the assimilation model given an input coordinate
system.
Parameters
----------
input_coords : tuple[FrameSchema | CoordSystem, ...]
Input coordinate system tuple. FrameSchema (OrderedDict mapping field names
to numpy arrays) for DataFrame inputs, or CoordSystem (OrderedDict mapping
dimension names to coordinate arrays) for tensor inputs
*args
Additional positional arguments
**kwargs
Additional keyword arguments, typically including request metadata such as
request_time and request_lead_time from DataFrame attrs
Returns
-------
tuple[FrameSchema | CoordSystem, ...]
Tuple of coordinate systems or frame schemas, one for each output argument
that __call__ or create_generator returns
"""
pass
def to(self, device: Any) -> AssimilationModel:
"""Moves assimilation model onto inference device, this is typically satisfied
via `torch.nn.Module`.
Parameters
----------
device : Any
Object representing the inference device, typically `torch.device` or str
Returns
-------
AssimilationModel
Returns instance of prognostic
"""
pass
Note
Data assimilation models do not need to inherit this protocol. The protocol defines the APIs that built-in workflows and utilities expect.
Data assimilation models can work with tensor data, tabular observation data, or both.
For tabular inputs, models commonly use a FrameSchema to describe fields and
constraints. For tensor inputs, models use the same CoordSystem convention described
in Data Movement.
Data Assimilation Usage¶
Loading a Pre-trained Data Assimilation Model¶
Use the concrete data assimilation model class you want to run. When that class supports automatic packages, the following pattern downloads and loads the pre-trained weights. More information on automatic downloading of checkpoints can be found in the AutoModels section.
from earth2studio.models.da import HealDA
model_package = HealDA.load_default_package()
model = HealDA.load_model(model_package)
Stateless Assimilation¶
The main work of a data assimilation model is the __call__ function. It accepts one
or more observation or state inputs and returns one or more assimilated outputs.
Stateful Assimilation¶
Some assimilation workflows need to process a sequence of observation batches while
maintaining internal state. For those workflows, use create_generator.
# Assume model is an instance of an AssimilationModel
generator = model.create_generator()
generator.send(None) # Prime the generator
for observations in observation_batches:
analysis, = generator.send(observations)
generator.close()
Custom Data Assimilation Models¶
To integrate your own data assimilation model, satisfy the interface above and keep the
input and output schemas explicit. The model should advertise initialization
requirements with init_coords(), accepted inputs with input_coords(), and produced
outputs with output_coords().
We recommend reviewing the extension examples, which show the style expected for adding custom Earth2Studio components.
Contributing a Data Assimilation Model¶
Want to add a data assimilation model to the package? We are happy to work with you. We expect the model to abide by the defined interface and meet the requirements set forth in our contribution guide. Typically, you are expected to provide any required weights or assets in a downloadable location that can be fetched.
Open an issue when you have an initial implementation you would like us to review. If you are aware of an existing model and want us to implement it, open a feature request and we will get it triaged.