Transformer UNet for data assimilation and forecasting

[SOHAMPAL23]

This PR introduces a modular data assimilation (DA) interface designed to integrate cleanly with existing and future transformer-based models in the repository.

The motivation comes from recent research exploring how data assimilation can be incorporated into neural models to cheaply generate large ensembles of predictions, improving uncertainty handling and forecast stability. Rather than tightly coupling DA logic to a specific architecture, this PR focuses on modularity and interchangeability, ensuring that DA components can be swapped or extended without affecting core model implementations.

The added components are intentionally lightweight and model-agnostic, operating on either a graph representation or a collection of nodes, and returning outputs in the same structural format. No existing model behavior, training loops, or APIs have been modified.

Fixes #21

How Has This Been Tested?

Manual validation by integrating the new data assimilation interface with existing model outputs.

Verified that the module:

• Accepts both graph-based and node-based inputs

• Produces structurally consistent outputs

• Does not alter existing training or inference workflows

Reproducibility

• Run an existing model forward pass to generate a state or ensemble.

• Pass the output through the new DA module.

• Confirm the returned structure matches the input format.

• Yes

Sanity Checks (Data Processing)

• Performed basic consistency checks on input/output shapes and data flow.

• No numerical transformations were applied that require plotting at this stage.

• Yes

Checklist

• My code follows OCF's coding style guidelines

• I have performed a self-review of my own code

• I have made corresponding changes to the documentation

• I have added tests that prove my feature works (to be added in a follow-up PR)

• I have checked my code and corrected any misspellings

[jacobbieker]

Thanks for this. It is very large PR, so can you split it up? Maybe one on the base, and then one for each of the Kalman, Particle, and Variational ones? Also, all of the newly added ones need to have tests in the tests folder.

[SOHAMPAL23]

@jacobbieker sure working on it