feat: add support for output-only diagnostic variables

[palakbhati]

I have implemented support for the diagnostic data category as discussed. This allows users to designate certain variables as "output-only," preventing them from being fed back into the autoregressive loop as input for subsequent steps.

Summary of 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.

Issue Link

Closes #684

Type of change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation (Addition or improvements to documentation)

Checklist before requesting a review

• My branch is up-to-date with the target branch - if not update your fork with the changes from the target branch (use pull with --rebase option if possible).
• I have performed a self-review of my code
• For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
• I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
• I have updated the README to cover introduced code changes
• I have added tests that prove my fix is effective or that my feature works
• I have given the PR a name that clearly describes the change, written in imperative form (context).
• I have requested a reviewer and an assignee (assignee is responsible for merging). This applies only if you have write access to the repo, otherwise feel free to tag a maintainer to add a reviewer and assignee.

Checklist for reviewers

Each PR comes with its own improvements and flaws. The reviewer should check the following:

• the code is readable
• the code is well tested
• the code is documented (including return types and parameters)
• the code is easy to maintain

Author checklist after completed review

• I have added a line to the CHANGELOG describing this change, in a section
  reflecting type of change (add section where missing):

• added: when you have added new functionality
• changed: when default behaviour of the code has been changed
• fixes: when your contribution fixes a bug
• maintenance: when your contribution is relates to repo maintenance, e.g. CI/CD or documentation

Checklist for assignee

• PR is up to date with the base branch
• the tests pass
• (if the PR is not just maintenance/bugfix) the PR is assigned to the next milestone. If it is not, propose it for a future milestone.
• author has added an entry to the changelog (and designated the change as added, changed, fixed or maintenance)
  Once the PR is ready to be merged, squash commits and merge the PR.

[Sir-Sloth-The-Lazy]

Thanks for taking this on! The formatting/mypy points are quick fixes, but I think the core mechanism needs a rethink before we go further. Because boundary_states is the ground truth in every phase, overwriting the predicted diagnostic channels with it means the model never learns those variables (zero loss) and eval metrics leak the truth. Issue #684's intent is output-only variables, predicted and trained on, but not fed back as input which usually means the predictor's output dimension exceeds its input state dimension and the AR feedback slices the diagnostics off. Could you align on that design (probably with @joeloskarsson or @sadamov ) before iterating?

[Sir-Sloth-The-Lazy]

the comment does not describe the code, could you please write a comment which explain a future reader, what this line means ?

[Sir-Sloth-The-Lazy]

Should this live on TrainingConfig? Issue #684 frames diagnostic as a datastore data category (parallel to state/forcing/static), not a training hyperparameter. Worth confirming the design with @sadamov

[Sir-Sloth-The-Lazy]

(1)This will fail the mypy pre-commit hook,the annotation needs to be list[int] | None = None. (2) This new parameter isn't added to the init numpydoc Parameters section; It would be appreciated, if you do that.

[Sir-Sloth-The-Lazy]

This is a trailing whitespace, please remove it

[Sir-Sloth-The-Lazy]

Also, the indices are stored without validation, no bounds check against num_state_vars. This can cause a problem

[Sir-Sloth-The-Lazy]

This is the core concern. boundary_state is boundary_states[:, i], and ForecasterModule.common_step passes target_states as boundary_states in all phases (train/val/test). So this line sets the returned prediction's diagnostic channels equal to the ground truth. Consequences:

• Training: the loss on those channels is identically zero --> the model gets no gradient and never learns to predict diagnostics.
• Val/test: the "prediction" for diagnostics is the target --> data leakage / artificially perfect scores.
• new_state also becomes prev_state, so ground-truth diagnostics are fed back into the predictor input each step, the opposite of "output-only," and it leaks truth into the ordinary state predictions too.
  I think this needs a different design: diagnostics should be predicted (and included in the loss) but excluded from the AR feedback, which generally means the predictor outputs more channels than it takes as input, rather than overwriting predictions with truth.

[Sir-Sloth-The-Lazy]

The PR description says this validates that requested variables exist, but a missing variable just prints and continues, producing a silently-incomplete diagnostic_indices. I'd raise a config error here instead so misconfiguration fails loudly. Also, the codebase uses warnings.warn, not print.

[Sir-Sloth-The-Lazy]

Note this mapping logic only runs here in main(),load_forecaster_module_from_checkpoint constructs ARForecaster too, so a reloaded checkpoint silently loses its diagnostic config (gets None). Might be cleaner to resolve names --> indices inside the forecaster/datastore so both construction sites share it.

[Sir-Sloth-The-Lazy]

This datastore's state features are u100m, v100m, r2m, t2m there's no precipitation_rate in it. So with this config the lookup misses, prints the warning, and diagnostic_indices ends up empty --> the feature is silently disabled. The example meant to demonstrate the feature actually exercises the not-found path. It also highlights the underlying gap: the datastore has no diagnostic-category data, which is what #684 says needs to exist first.

[palakbhati]

@Sir-Sloth-The-Lazy Thanks for the thorough review agreed on all points, especially the core one.

I implemented diagnostics as a subset of existing state indices, masked with boundary_state (= target_states) at each AR step the same mechanism used for spatial boundary conditions, misapplied to feature indices. Since boundary_state is ground truth in every phase, this zeroes the training loss on those channels and leaks truth into eval metrics, exactly as you described. It also feeds true diagnostic values back into prev_state, which is the opposite of output-only.

Fix, not patch: rather than fix this mechanism in place, I'll rework it so the predictor's output width exceeds its input width: grid_output_dim = num_state_vars + num_diagnostic_vars. The AR loop will only feed the first num_state_vars channels back as prev_state/prev_prev_state; the remaining diagnostic channels are real predictions, included in the loss, never overwritten with truth.

This also means diagnostic needs to be a datastore-level category (parallel to state/forcing/static, per #684), not a TrainingConfig field agreed that's the wrong home for it. And you're right that today there's no diagnostic data in the datastore layer at all, which is why the example config's precipitation_rate just hits the not-found path.

@sadamov @joeloskarsson before I rebuild this, could you confirm the design above (extra output channels beyond state width, sliced off before AR feedback) matches what you had in mind for #684? Want to align once rather than iterate blind again.

If that's confirmed, I'll close this PR and split the rework into smaller, reviewable pieces:

Datastore support for a real diagnostic category (+ a diagnostic variable in tests/dummy_datastore.py so it's actually testable)
StepPredictor/BaseGraphModel output-width split (state-delta vs. diagnostic-absolute channels)
ARForecaster change: only state channels loop back, diagnostics collected separately
train_model.py wiring, with name→index resolution shared between main() and load_forecaster_module_from_checkpoint so checkpoint reloads don't silently lose the config

[sadamov]

Thanks @palakbhati, and thanks @Sir-Sloth-The-Lazy for the review. Agreed the current mechanism can't stand: overwriting the predicted channels with boundary_state at each AR step zeroes their loss and leaks truth. Let's go with your original proposal: a real diagnostic datastore category parallel to state/forcing/static, per #684.

Design choices:

• diagnostic a datastore category with its own std; output-only, so never in init_states, loaded only as targets.
• Predictor grid_output_dim = num_state + num_diagnostic, output head split: state slice is a residual delta, diagnostic slice absolute (no prev to add a delta to, since diagnostics aren't fed back).
• ARForecaster feeds only the state slice back; diagnostics collected separately, no boundary overwrite.
• Loss/metrics get a diagnostic term with its own weighting.

I'd close this PR and split into: (1) datastore category + a diagnostic var in dummy_datastore, (2) data module targets, (3) predictor output split, (4) forecaster feedback, (5) loss wiring. Start with the datastore piece so the rest is testable.

@joeloskarsson @leifdenby What do you think? Let's wait for their input as well.