Skip to content

Add CRPS metric scaffold for ensemble probabilistic evaluation - #229

Closed
Panchadip-128 wants to merge 4 commits into
mllam:mainfrom
Panchadip-128:crps-metric-pr
Closed

Add CRPS metric scaffold for ensemble probabilistic evaluation#229
Panchadip-128 wants to merge 4 commits into
mllam:mainfrom
Panchadip-128:crps-metric-pr

Conversation

@Panchadip-128

@Panchadip-128 Panchadip-128 commented Feb 22, 2026

Copy link
Copy Markdown

Describe your changes

Following @leifdenby's review feedback, this PR has been expanded beyond
the original scaffold scope. CRPS and ensemble generation are now
implemented together as a coherent unit.

Changes include:

  • Added three config-driven ensemble modes (sar, lagged_ic, hybrid)
    in ar_model.py
  • Added lagged_ic: IC perturbation along recent tendency vectors
    (structured IC spread, not random jitter)
  • Added hybrid: lagged_ic at t=0 + per-step SAR noise (IC + dynamical
    uncertainty)
  • Added vectorized crps_ensemble implementation in crps.py
    (energy-form CRPS, proper scoring rule)
  • Registered crps_ensemble under key "crps" in DEFINED_METRICS
    (metrics.py and __init__.py)
  • Fixed neural_lam.metrics import failures via importlib dynamic load
    (get_metric, crps_ensemble now importable cleanly)
  • Added test_crps.py: shape, dtype, perfect ensemble → CRPS ≈ 0
  • Added test_ar_model_ensemble_modes.py: lagged_ic + hybrid shape and
    diversity tests

This PR does not modify existing training or loss pipelines. Fully
backward compatible with deterministic workflows.

Motivation and context

Probabilistic forecasting is a core research direction for neural weather
models. CRPS is the NWP standard for evaluating ensemble forecast quality —
a proper scoring rule that penalises both bias and spread errors.

The previous ensemble path added Gaussian jitter x + ε with no clear
ensemble definition. lagged_ic addresses this by perturbing initial
conditions along recent temporal tendency vectors — systematic IC variants
consistent with lagged-average forecast ensembles in NWP. hybrid adds
per-step dynamical noise on top, covering both IC and process uncertainty.

This PR establishes foundational probabilistic output support and prepares
the codebase for deep ensembles, calibration analysis, and evaluation
pipeline integration.

What this PR does now

Three named ensemble modes (config-driven)

Mode Mechanism Status
sar Per-step interior Gaussian noise, BCs overwritten each step Existing, preserved
lagged_ic IC perturbation along recent tendency vectors NEW
hybrid lagged_ic at t=0 + SAR noise at t=1…T NEW
training:
  ensemble_mode: hybrid          # sar | lagged_ic | hybrid
  ensemble_size: 10
  perturbation_scale: 0.01       # SAR per-step σ
  ic_perturbation_scale: 0.05    # IC perturbation σ

CRPS metric

Energy-form proper scoring rule, vectorised PyTorch, GPU-friendly.
Registered in DEFINED_METRICS under key "crps":

CRPS(F, y) = E[|X − y|] − ½ · E[|X − X′|]
import neural_lam.metrics as m
m.get_metric("crps")   # returns crps_ensemble function

Metrics API fix

neural_lam/metrics/__init__.py performs a pragmatic dynamic import
(via importlib) to avoid a package/module name collision — the directory
neural_lam/metrics/ and the implementation file metrics.py share a
name that causes import-time failures. This loader is a minimal bridge
that keeps metrics.py as the single source of truth.

Note: happy to follow up with a rename refactor if maintainers prefer
that style over the dynamic loader.

lagged_ic perturbs initial conditions along recent temporal tendency
vectors
— structured, interpretable IC spread, not random jitter. Members
are systematic IC variants consistent with IC-ensemble approaches in NWP
(cf. lagged-average forecasts).

hybrid adds per-step dynamical noise on top, covering both IC uncertainty
(dominant at short lead times) and process uncertainty (longer lead times).

image

Issue Link

N/A — foundational probabilistic evaluation support

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.

@Panchadip-128 Panchadip-128 changed the title This PR introduces initial Continuous Ranked Probability Score (CRPS) metric support for ensemble probabilistic forecasting within neural-lam. Key contributions: Added a vectorized crps_ensemble implementation for evaluating ensemble predictions Registered CRPS within the metrics module for seamless metric selection Added a lightweight unit test validating shape and functional correctness Maintained full backward compatibility with deterministic workflows This implementation provides a modular evaluation primitive that enables future probabilistic forecasting extensions without affecting existing training pipelines. Add CRPS metric scaffold for ensemble probabilistic evaluation Feb 22, 2026
@Panchadip-128

Copy link
Copy Markdown
Author

Hi @j6k4m8 @leifdenby @khintz ,
Would you mind taking a look at this PR when you have time?
It adds a minimal CRPS metric scaffold for ensemble probabilistic evaluation and is designed to be non-breaking and incremental.
Thank you very much!

@leifdenby leifdenby left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for this @Panchadip-128! Can you explain the motivation for this PR a bit more please? As I read it this PR doesn't add any ensemble architectures, but just adds some random noise to the deterministic model output, is that right?

https://github.com/mllam/neural-lam/pull/229/changes#diff-043440a2d7a2cd62bb349e74c9fe4f55e69a8d0e4801f6ab2e150d381b11c74dR71

@Panchadip-128

Copy link
Copy Markdown
Author

Hi @leifdenby sir - thanks for the feedback and for reviewing the PR.

You’re correct that this PR does not introduce a full ensemble architecture. The intent here though is to add CRPS as a modular probabilistic evaluation primitive, independent of the ensemble generation mechanism. The temporary stochastic sampling is only used to exercise the interface contract and validate tensor shapes, not as a final ensemble approach.

The motivation is to establish ensemble-aware evaluation early while keeping the current deterministic training pipeline unchanged and enabling incremental probabilistic extensions.

As follow-up work if you allow, I plan to propose PRs introducing:
• a proper ensemble generation pathway (e.g., multi-sample autoregressive rollout / decoder stochasticity)
• integration of CRPS with evaluation pipelines once ensemble outputs exist
• exploration of architectural ensemble variants and calibration analysis

I’ll also update the PR description to clarify that this PR focuses on evaluation scaffolding rather than ensemble modelling accordingly.

Thanks again for the guidance - happy to align with any preferred ordering of probabilistic components in the roadmap.

@Panchadip-128
Panchadip-128 marked this pull request as draft February 28, 2026 19:08
@Panchadip-128

Panchadip-128 commented Feb 28, 2026

Copy link
Copy Markdown
Author

Hi @leifdenby - thanks again for the previous feedbacks, based on that i have tried to improved this PR.

I've expanded the scope so CRPS and ensemble generation can be reviewed
together as a coherent unit and also converted to draft PR as per previous instructions.

Key changes:

  • lagged_ic: ICs perturbed along recent tendency vectors — structured
    IC spread, not random jitter. Consistent with lagged-average forecast
    ensembles in NWP.
  • hybrid: lagged_ic at t=0 + per-step SAR noise, covering both IC and
    dynamical uncertainty.
  • sar: retained as comparison baseline.
  • CRPS registered as "crps"metrics.get_metric("crps") works cleanly.
  • Metrics AttributeError on import fixed.
  • Tests passing: shape, diversity, CRPS correctness (perfect ensemble → 0).

I've also updated the PR description to reflect the expanded scope.

Happy to discuss the IC perturbation approach, adjust scope, or remove
the dynamic import bridge if you'd prefer a rename refactor instead.

Thanks so much for your time reviewing this!

Regards.

@Panchadip-128

Copy link
Copy Markdown
Author

Hi , please let me know if any changes are required for this PR. If everything looks good I’ll mark it ready for review and apply any further suggestions. Thanks for your review - any guidance to improve this is appreciated.

@sadamov
sadamov requested a review from leifdenby March 12, 2026 19:20
@joeloskarsson

Copy link
Copy Markdown
Collaborator

The name of this PR does not seem to match the scope at all now. If I understand correctly you are here proposing the addition of ensemble generation through IC perturbations? In general I would be very hesitant to merge in any ensemble method that does not build on a better structured model architecture such as discussed in #49 .

My personal opinion is that ensembling through IC perturbations is generally a bad idea, and if you do it there is a lot of care needed with magnitudes and correlations.

@Panchadip-128

Copy link
Copy Markdown
Author

Hi @joeloskarsson

Thanks for the clarification - that helps a lot.

Based on the discussion around #49, it makes sense that ensemble mechanisms should build on top of the upcoming model hierarchy refactor rather than being introduced directly in the current ARModel structure.

To keep this PR aligned with that direction, I’ll simplify it back to its original goal: introducing the CRPS metric as a standalone probabilistic evaluation primitive within the metrics module, without modifying ARModel or adding ensemble generation logic.

This keeps the change small and independent from the ongoing architecture work while still establishing the evaluation component that probabilistic models can use later.

Would you prefer the CRPS metric to remain purely metric-side for now, or should it already support the expected ensemble tensor shapes used by future probabilistic models?

Thanks again for the guidance - I’ll update the PR accordingly as you prefer and work on thereby.

@joeloskarsson

Copy link
Copy Markdown
Collaborator

To avoid duplicate PRs that try to implement the CRPS metric, could you help out in #226 for now? Bring good ideas for the CRPS implementation there.

@Panchadip-128

Copy link
Copy Markdown
Author

Thanks for the guidance!

That makes sense - consolidating the CRPS work into a single effort will definitely avoid duplicated implementations.

I’ll look at #226 for now and contribute ideas and improvements so the CRPS implementation can evolve collaboratively.

Happy to help refine the metric implementation and tests there.

@joeloskarsson

Copy link
Copy Markdown
Collaborator

Thanks, then I'll close this.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants