Add a `diagnostic` data category (output-only)

[sadamov]

Right now a datastore variable falls into one of three categories: state (model input and output), forcing (input only) and static (input only). There is no output-only category for variables the model predicts but does not feed back as input for the next autoregressive step, e.g. precipitation or radiation fluxes.

A diagnostic category fits naturally as the output-only counterpart to forcing:

Category	Model input	Model output
state	✓	✓
forcing	✓
static	✓
diagnostic		✓

This came up while documenting the data categories in #635 (@joeloskarsson suggested tracking it as its own issue rather than mentioning it in the README before it exists).

[palakbhati]

I am interested in working on this. I've been looking at the current data category schema, and I plan to add 'diagnostic' as a fourth category. I will ensure the autoregressive loop logic is updated to exclude these variables from the next time step. Does this alignment match the maintainers' vision for the implementation?

[sadamov]

Hi @palakbhati thanks! One big question will be how the user can specify these types of variables (since they are not really in a datastore). Most likely the current config yaml file on the model side could be a good place for that. There are several PRs open that affect the datastores yaml-configs and similar. It makes sense to consider these before we settle on an implementation

[palakbhati]

Hi @sadamov thanks for the feedback! That makes sense keeping the datastore definition separate from the model’s use of diagnostic variables definitely avoids unnecessary complexity in the Zarr loading pipeline.

If I move the specification of 'diagnostic' variables to the model-side config.yaml, I see two primary benefits:

It allows users to define diagnostic outputs without needing to rebuild their entire datastore.
It keeps the core autoregressive loop clean, as I can just filter these out during the predict_step or loss calculation based on the model config.

I’ll take a look at the other open PRs you mentioned to ensure my approach to the config.yaml doesn't conflict with those changes. Does this seem like the right direction before I start drafting a PoC?

[palakbhati]

Hi @sadamov,

I have implemented support for the diagnostic data category as discussed. Here is a summary of the changes:

Configuration Schema: Updated NeuralLAMConfig in config.py to include a diagnostic_vars field, allowing users to specify output-only variables directly in the model-side config.yaml.

Autoregressive Logic: Modified ARForecaster.forward in autoregressive.py to filter these variables during the unrolling process. By setting the diagnostic indices to the ground truth (boundary_state) at each step, we ensure the model treats them as output-only and does not propagate them autoregressively.

Dynamic Mapping: Updated train_model.py to dynamically map variable names provided in the config to their corresponding indices in the datastore, with validation to ensure the requested variables exist.

I have verified the implementation using the current training pipeline, and it successfully handles the configuration loading and variable filtering.

I am ready to open a Pull Request for this. Should I proceed, or are there any additional tests/checks you would like me to include in the PR?