feat/add-ensemble-metrics-crps-spread

[shauryam2807]

Describe your changes

This PR ports the sample-based CRPS (crps_ens) and ensemble variance (spread_squared) evaluation metrics from the prob_model_lam branch to main.

These metrics are essential for evaluating the probabilistic and ensemble models being developed (as tracked in Issue #62). I have also added a new test file tests/test_metrics.py with unit tests verifying the mathematical correctness and unbiased estimation of both metrics against expected values.

Dependencies: None

Issue Link

Addresses parts of #62

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
• 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
• I have requested a reviewer and an assignee

[joeloskarsson]

Hi, could you check how this relates to #229 ? Seems like these is potentially a lot of overlap, and it would be great if we could focus community efforts on one of these PRs.

[shauryam2807]

Hi @joeloskarsson, thanks for pointing this out! I’ve reviewed #229 by @Panchadip-128 and there is definitely a large overlap on the CRPS implementation and test infrastructure. PR #229 is very comprehensive and introduces excellent ensemble mode configurations.

To align with the community and avoid duplicated effort, I am more than happy to close this PR in favor of #229.

@Panchadip-128 I noticed my branch also ports the spread_squared (ensemble variance) metric, which might be a useful complement to CRPS for evaluating probabilistic spread. Let me know if you’d like me to open a PR against your branch to add it, or if it would be better for me to wait and open a follow-up PR adding spread_squared once #229 is merged into main!

[joeloskarsson]

I've been a bit reluctant to implement an ensemble CRPS metric before we have any model that can use it, but given that there seem to be so much interest in this I think we can start looking at it. Here are some quick scattered thoughts around CRPS metric:

• We should follow the shapes from [RFC/Design] Standardize probabilistic vs deterministic return contract to unblock evaluation integrations  #335
• It would be good to have both fair (unbiased), unfair (biased) and almost-fair (https://arxiv.org/html/2412.15832v1, with configurable alpha) CRPS. Maybe we could have a base function and then just wrappers of this with different names.
• We definitely need the sorting-based implementation for memory. Even with this, for long forecasts it is easy to run out of memory when computing the spread part. Would be good to be able to then internally batch this somehow. Unclear how to best control this.
• For clean implementation we can keep the N=2 members case as a special case. This will likely be used a lot for training, so good to have a clean and efficient implementation for.

Sharing a very general CRPS method that I have for another project (for inspiration):

def crps_ens(target, pred, ens_dim=1, estimator="unbiased", afc_alpha=None):
    """
    (Negative) Continuous Ranked Probability Score (CRPS)
    Estimator from samples. See e.g. Weatherbench 2.

    (..., M, ...,) is any number of batch dimensions, including ensemble
        dimension M
    target: (..., d_state), target
    pred: (..., M, ..., d_state), prediction
    ens_dim: batch dimension where ensemble members are laid out, to reduce over

    Returns:
    metric_val: (..., d_state)
    """
    num_ens = pred.shape[ens_dim]  # Number of ensemble members
    assert (
        num_ens > 1
    ), "CRPS can only be estimated for ensemble with more than 1 member"

    mean_mae = torch.mean(
        torch.abs(pred - target.unsqueeze(ens_dim)), dim=ens_dim
    )  # (..., d_state)

    if estimator == "biased":
        diff_factor = 1 / num_ens
    elif estimator == "unbiased":
        diff_factor = 1 / (num_ens - 1)
    elif estimator == "almost-fair":
        assert (
            afc_alpha is not None
        ), "afc_alpha must be provided for almost-fair CRPS estimator"
        diff_factor = (num_ens - 1 + afc_alpha) / (num_ens * (num_ens - 1))
    else:
        raise NotImplementedError(f"Unknown CRPS estimator: {estimator}")

    if num_ens == 2 and estimator == "unbiased":
        # Use simpler estimator
        pair_diffs_term = (
            -0.5
            * diff_factor
            * torch.abs(pred.select(ens_dim, 0) - pred.select(ens_dim, 1))
        )  # (..., d_state)
    else:
        # This is the rank-based implementation with O(M*log(M)) compute and
        # O(M) memory. See Zamo and Naveau and WB2 for explanation.
        # For smaller ensemble we can compute all of this directly in memory.

        # Ranks start at 1, two argsorts will compute entry ranks
        ranks = pred.argsort(dim=ens_dim).argsort(ens_dim) + 1

        # Note: We can batch this over any batch dimension if expensive
        pair_diffs_term = diff_factor * torch.mean(  # Note mean
            (num_ens + 1 - 2 * ranks) * pred,
            dim=ens_dim,
        )  # (..., d_state)

    crps_estimator = mean_mae + pair_diffs_term  # (..., d_state)
    return crps_estimator

Should naturally be adapted to the neural-lam metric API. But this gives an idea for how to get the fair/unfair/almost-fair versions in an easy way.

[Panchadip-128]

Hi @joeloskarsson and @shauryam2807,

Thanks for the detailed notes and for sharing the reference CRPS estimator — very helpful.

I’ve been looking into how this could integrate cleanly with the current metrics API and the forecast tensor shapes proposed in #335. A few thoughts that might help structure the implementation:

1. Shared base estimator + variants
   It seems clean to define a single core CRPS implementation and expose different estimators (biased / unbiased / almost-fair) as lightweight wrappers. This keeps the API simple while allowing flexibility depending on ensemble size and use case.

2. Memory-efficient spread term
   For larger ensembles, the pairwise formulation becomes expensive (O(S²)). The rank-based formulation you shared (O(S log S)) looks like the right default for scalability, especially for longer rollout horizons.

3. Special case for small ensembles (S = 2)
   Handling this explicitly could simplify the implementation and improve efficiency for small ensembles (which might be common during experimentation).

4. Shape alignment with [RFC/Design] Standardize probabilistic vs deterministic return contract to unblock evaluation integrations  #335
   To stay consistent with the proposed contract:

• deterministic: (B, T, N, F)
• ensemble: (B, S, T, N, F)

it may be useful for the metric to accept an explicit ens_dim argument (or assume a fixed position), so it can operate generically regardless of how the ensemble dimension is arranged internally.

5. Integration with existing metric reductions
   Since existing metrics use mask_and_reduce_metric, one option could be to have CRPS return entry-wise values (e.g., (..., N, F)) before applying reductions, to stay consistent with the current API.

I also sketched a small evaluation flow diagram based on #335 + CRPS integration to help ground the discussion — happy to refine or adjust if useful.

If helpful, I can follow up with a minimal implementation aligned with this structure (e.g., base estimator + rank-based spread term) and keep it tightly scoped.

Thanks!

[kshirajahere]

one thing that feels important to lock down early here is that crps_ens should probably stay entry-wise until after mask_and_reduce_metric, same as the deterministic metrics. otherwise latitude weighting / boundary masking / sum_vars can quietly diverge between deterministic and ensemble eval paths.

especially if there is any S=1 special casing, i’d be careful that grid_weights and reduction semantics do not get skipped there, because that kind of mismatch will be really hard to notice later when comparing deterministic vs ensemble evaluation.

[Panchadip-128]

Good point - keeping crps_ens entry-wise until after mask_and_reduce_metric ensures consistency with deterministic metrics, so masking, latitude weighting (grid_weights), and variable reductions apply uniformly across both paths. That kind of subtle divergence is very hard to catch later when comparing deterministic vs ensemble evaluation.

On S=1: since single-member CRPS is undefined as a proper scoring rule, an explicit assert S > 1 early in the function seems cleaner than silently short-circuiting - it also guarantees the reduction pipeline is never accidentally bypassed.
This is actually why I proposed in my earlier comment that crps_ens should return entry-wise (..., N, F) values before any reduction - so that mask_and_reduce_metric naturally becomes the single source of truth for all weighting and reductions across both paths.

Any future changes to reduction semantics (e.g., new masking strategies) then automatically apply to ensemble metrics without special-casing.

Happy to help align the implementation with this if useful.

[Sir-Sloth-The-Lazy]

Hi @joeloskarsson, @shauryam2807, @Panchadip-128, @kshirajahere —
I've opened #431 to add the spread_squared (ensemble variance) metric against main, incorporating the design consensus from this thread:

Entry-wise computation → mask_and_reduce_metric internally (same pattern as mse, nll, etc.), so grid_weights, boundary masking, and sum_vars apply uniformly across deterministic and ensemble eval paths
assert S > 1 — no silent S=1 fallback
Bessel's correction for unbiased estimation, matching the 1/(S-1) convention in the CRPS reference code shared in #226
ens_dim = -3 following the (..., S, N, d_state) convention
Full API signature match: (pred, target, pred_std, mask, average_grid, sum_vars)

[joeloskarsson]

@Sir-Sloth-The-Lazy can you explain why we need a separate PR for the ensemble spread? That is already added here, so #431 seems just like a duplicate of this now. Let me know if I am missing something.

[GiGiKoneti]

the N=1 branch has two issues when delegating to mae:

1. pred_std is hardcoded as None, which causes a TypeError
   (ones_like(): argument must be Tensor, not NoneType) if the caller
   passes a pred_std tensor
2. sum_vars is silently dropped, so sum_vars=False is ignored
   and output shape diverges from the N>1 path

verified both locally. fix:

if num_ens == 1:
    return mae(
        pred.squeeze(ens_dim),
        target,
        pred_std.squeeze(ens_dim) if pred_std is not None 
            else torch.ones_like(target),
        mask=mask,
        average_grid=average_grid,
        sum_vars=sum_vars,
    )

also no N=1 test case exists which is likely why this slipped
through ...best to add one to test_metrics.py

[GiGiKoneti]

also wrote a test that catches both issues — crashes on current
code, passes with the fix:

def test_crps_ens_single_member():
    # N=1 reduces to MAE — verify sum_vars and pred_std are correctly
    # passed through. Previously both were silently dropped causing a
    # TypeError crash and incorrect output shapes.
    pred = torch.tensor([[[[1.0, 2.0, 3.0]]]]) # shape (1, 1, 1, 3)
    target = torch.tensor([[[2.0, 2.0, 2.0]]]) # shape (1, 1, 3)
    pred_std = torch.ones_like(pred)

    result_summed = crps_ens(
        pred, target, pred_std,
        sum_vars=True, ens_dim=1
    )
    result_not_summed = crps_ens(
        pred, target, pred_std,
        sum_vars=False, ens_dim=1
    )

    assert result_not_summed.shape[-1] == 3, \
        "sum_vars=False should preserve feature dimension in N=1 path"
    assert result_summed.shape != result_not_summed.shape, \
        "sum_vars flag ignored in N=1 path"

[joeloskarsson]

Why would we need to pass any pred_std to mae, that is not used there? We could always pass it as None, no matter what was input to crps_ens, as the crps should not depend on predicted std.

[Sir-Sloth-The-Lazy]

@joeloskarsson To clarify my point of view, #226's spread_squared is bundled with crps_ens in a branch that currently has review issues (duplicate DEFINED_METRICS), and can't merge until the CRPS design questions are resolved (estimator variants, rank-based implementation, internal batching).
#431 decouples spread_squared so it can land independently. It also adds a dedicated test suite (9 cases covering mathematical correctness, torch.var equivalence, shape contract, mask behavior, and reduction flags) that #226 doesn't have for this metric.
If you'd prefer to keep them together in #226 and merge once CRPS is settled, happy to close #431, just let me know.😃

[shauryam2807]

Hello
Hi everyone, thanks for all the excellent feedback! Sorry for the delayed response.

@joeloskarsson — Thank you for sharing the reference CRPS implementation. I'll refactor my implementation to follow your design: rank-based O(M log M) approach with biased/unbiased/almost-fair estimator variants, aligned with the tensor shapes from #335.

@GiGiKoneti — Great catches on the N=1 path bugs. I'll fix both the pred_std=None crash and the dropped sum_vars flag, and add the test case you suggested.

I'll push the updated implementation shortly.

[Panchadip-128]

@GiGiKoneti One small addition to the test: it would be worth also verifying that the N=1 path and the N>1 path produce consistent results when S=1 reduces to a single-member forecast - specifically that mask is still applied correctly in both cases, since that's the subtle divergence that's hardest to catch.
Something like:

test_crps_ens_single_member_mask_consistency():
pred = torch.tensor([[[[1.0, 2.0, 3.0]]]])  # (1, 1, 1, 3)
target = torch.tensor([[[2.0, 2.0, 2.0]]])  # (1, 1, 3)
pred_std = torch.ones_like(pred)
mask = torch.tensor([[[1.0, 1.0, 0.0]]])    # mask last feature

result = crps_ens(pred, target, pred_std, mask=mask, sum_vars=False, ens_dim=1)
assert result.shape[-1] == 2, "masked feature should be excluded"

[GiGiKoneti]

verified locally on current PR #226 code (fix not yet pushed) —
the N=1 path crashes before mask logic even runs:

TypeError: ones_like(): argument must be Tensor, not NoneType

because mae() calls torch.ones_like(pred_std) but the N=1
branch passes None. current state:

1. pred_std=None crash still present at L273
2. sum_vars still dropped in N=1 path at L270
3. DEFINED_METRICS still defined twice
4. fair/unfair toggle not yet added
5. no N=1 test case

once the fix lands:

if num_ens == 1:
    return mae(
        pred.squeeze(ens_dim),
        target,
        pred_std.squeeze(ens_dim) if pred_std is not None
            else torch.ones_like(target),
        mask=mask,
        average_grid=average_grid,
        sum_vars=sum_vars,
    )

@Panchadip-128's mask consistency test should pass correctly
after this — mask_and_reduce_metric handles feature masking
in the mae() delegation. waiting to re-review once the
update lands

[GiGiKoneti]

reviewed current branch — fixes not yet pushed.
comments left inline on the N=1 path. will re-review
once the update lands.

[joeloskarsson]

Read over and added some comments. Some of these are more discussion points, that I think we need to make some decisions on the best way to handle.

[joeloskarsson]

These comments with all the dashed lines seem unnecessarily large

[joeloskarsson]

Unrelated changes

[shauryam2807]

Hello
Hi @joeloskarsson, thanks for the thorough review! I've addressed all your comments:

• Dead assert: Removed the unreachable assert num_ens > 1. The N=1 path now emits a warnings.warn() to alert users that CRPS is falling back to MAE.
• Large separator comments: Removed all dashed-line banners — comments are now clean and minimal.
• Unrelated file changes: Reverted all unrelated files. PR now only touches metrics.py and test_metrics.py.
• CRPS reference: Updated docstring to cite Lang et al., 2024 (AIFS) for the almost-fair estimator.
• assert → raise ValueError: Replaced all assert with explicit exceptions, following the convention from Replace assert with if...raise ValueError for runtime safety #279.
• N=2 generalization: The closed-form pair-difference shortcut now applies to all estimator types (biased, unbiased, almost-fair), not just unbiased.

New tests added:

• test_crps_ens_almost_fair_missing_alpha — verifies ValueError when afc_alpha is missing
• test_crps_ens_n2_all_estimators — verifies N=2 shortcut works for all estimator types
• test_spread_squared_single_member_raises — verifies ValueError for S=1

All 11/11 tests pass ✅

[Panchadip-128]

Naveau / AIFS formulation assumes α ∈ (0, 1) — passing afc_alpha=0 collapses to the unbiased estimator silently, and values outside (0, 1) produce mathematically invalid diff_factor. A simple guard would prevent silent misuse:

if not (0 < afc_alpha < 1):
raise ValueError(
f"afc_alpha must be in (0, 1) for almost-fair CRPS, got {afc_alpha}"
)

[joeloskarsson]

Good idea, yes checking that afc_alpha is in [0,1] would be good.

[GiGiKoneti]

@shauryam2807 Thanks for incorporating the fixes!

I've pulled the latest changes (b3978fa) and verified locally that the N=1 edge cases and sum_vars flags now behave correctly and avoid the previous crashes and shape divergences. The new test cases gracefully cover both issues.
Tests pass on my end. LGTM!

[Sir-Sloth-The-Lazy]

@shauryam2807 are you willing to work on this PR ? I would be really happy to help ! :) Great work here ! Just pinging @joeloskarsson if you have time to review this 😁

[joeloskarsson]

My previous comments are resolved, but this needs another look over as there are some conflicts to resolve and unrelated changes in the PR.

@shauryam2807 as @Sir-Sloth-The-Lazy mention above, he is happy to help as this relates to his overarching work of enabling probabilistic forecasting in the repo. Do let us know if you want to keep doing the work on this, get some help from @Sir-Sloth-The-Lazy, or hand over this fully to him.

[joeloskarsson]

Unrelated?

[joeloskarsson]

Unrelated?

[joeloskarsson]

Unrelated changes?

[joeloskarsson]

Unrelated changes?

[joeloskarsson]

Unrelated changes?

[joeloskarsson]

It would be good to add also specific biased, unbiased and almost-fair (with exactly afc_alpha=0.95) versions of crps_ens, so that they could be used in the same interface as other metrics, without having to deal with the afc_alpha argument in some special way. Suggestion to use functools.partial.

[joeloskarsson]

Do we need to allow for different ens_dim here? This is something we need to enforce throughout the codebase, so might be better to just specify clearly that we assume dim 1 to be the ens dim in all inputs.

[Sir-Sloth-The-Lazy]

I agree, Should we expose ens_dim as a configurable argument at all? Per the shape contract in #335 the ensemble axis is fixed at dim 1 ((B, S, T, N, F)), and we want that enforced consistently across the codebase rather than per-call. Suggest dropping the parameter and instead documenting/asserting that S is assumed at dim 1, that keeps the deterministic and ensemble eval paths from quietly diverging if a caller passes a different value. If we do keep it, it should at least be validated against the expected layout.

[shauryam2807]

Hello @joeloskarsson and @Sir-Sloth-The-Lazy,

Thank you for reviewing my PR. I will address all the conflicts and suggestions mentioned in the review.

I am currently occupied with my end-semester examinations and my responsibilities as a mentee in another project, so there may be some delay in my responses. However, I remain committed to this contribution and will ensure that the PR is properly maintained and updated.

I am also fully capable of maintaining this PR and will continue to actively work on it until all review comments and requested changes have been resolved.

Thank you for your understanding

[sadamov]

Closing this in favour of #649, which consolidates the probabilistic metrics + ensemble plumbing track for v0.8.0. The sample-based CRPS and spread_squared work here will be folded into the consolidated PR. Thanks for the contribution @shauryam2807, you'll be credited.

[sadamov]

oh yeah this one was a bit too eagerly closed for consolidation, sorry @shauryam2807 and thanks for your work here!

[shauryam2807]

Hello @joeloskarsson @sadamov @Sir-Sloth-The-Lazy

Yes I am exicted to work on this as I said you that I will mange it and work on it

Thank you for your time and believe

[joeloskarsson]

Hi @shauryam2807, I am trying to coordinate and push this through as quickly as possible. This PR is a high priority to us, since it directly relates to the work @Sir-Sloth-The-Lazy is doing with probabilistic models as part of his GSoC project.

Could you let us know your estimate for when you have time to look at this? We would really need to make some progress on this within a timeframe of ~ a week. If this does not seem feasible to you we will have to let @Sir-Sloth-The-Lazy take over the main work on this. I encourage you to instead collaborate on this PR, and let @Sir-Sloth-The-Lazy help with finalizing this work.

[shauryam2807]

"Hi @joeloskarsson, I am available and will prioritize this. I will resolve the conflicts and have this ready for you to review within 2-3 days.

[shauryam2807]

Hello @joeloskarsson, I have rebased the branch on latest main and cleaned up all unrelated file changes. The PR now only touches metrics.py and test_metrics.py. All previous review comments have been addressed. Ready for your review.

[Sir-Sloth-The-Lazy]

I went through the prob_model_lam branch this is ported from to confirm coverage. Good news: the function set matches (crps_ens + spread_squared), and this PR is actually ahead of the source by adding the biased / almost-fair estimators on top of the unbiased one. 👍

One functional gap though: the source's large-ensemble (M >= 10) memory-batched branch wasn't carried over, since it's the OOM case flagged earlier and the ensemble eval path will hit it. Everything else the ensemble models depend on (spread-skill ratio, ens_mse) lives in the model code, not here, so it's out of scope for this PR, spread_squared(..., sum_vars=False) already returns the shape that path needs.

[Sir-Sloth-The-Lazy]

This is an odd place to import 😅

[Sir-Sloth-The-Lazy]

If a caller passes pred_std of shape (F,) (a valid form elsewhere in the metrics API), .squeeze(ens_dim) raises IndexError: Dimension out of range. Since crps_ens doesn't use pred_std at all, we can follow Joel's earlier suggestion and just pass None unconditionally:

Suggested change

	return mae(
	pred.squeeze(ens_dim),
	target,
	pred_std.squeeze(ens_dim)
	if pred_std is not None
	else torch.ones_like(target),
	mask=mask,
	average_grid=average_grid,
	sum_vars=sum_vars,
	)
	return mae(
	pred.squeeze(ens_dim),
	target,
	torch.ones_like(target),
	mask=mask,
	average_grid=average_grid,
	sum_vars=sum_vars,
	)

[Sir-Sloth-The-Lazy]

I agree, Should we expose ens_dim as a configurable argument at all? Per the shape contract in #335 the ensemble axis is fixed at dim 1 ((B, S, T, N, F)), and we want that enforced consistently across the codebase rather than per-call. Suggest dropping the parameter and instead documenting/asserting that S is assumed at dim 1, that keeps the deterministic and ensemble eval paths from quietly diverging if a caller passes a different value. If we do keep it, it should at least be validated against the expected layout.

[Sir-Sloth-The-Lazy]

Hi @shauryam2807 , if you could look over the unresolved comments from Joel's review and solve them, I would be able to ping him for another review and we can get this PR merged sooner. Thank you for your hardwork 🤩

[shauryam2807]

Hi @joeloskarsson, regarding the large-ensemble OOM gap that @Sir-Sloth-The-Lazy flagged — I've been looking into a few approaches to handle memory efficiently for $M \ge 10$:

1. Spatial/Temporal Chunking (Internal Batching): Batching over the grid/time dimension (similar to what prob_model_lam originally did).
   • Pros: Exact same results, no approximations. Handles any ensemble/grid size.
   • Cons: Slightly slower due to the loop.
2. @torch.no_grad() Decorator: Since CRPS is primarily an evaluation metric, disabling gradient tracking saves ~50% memory.
   • Pros: One-line change, zero complexity.
   • Cons: Prevents CRPS from being used as a training loss function in the future.
3. PWM (Probability Weighted Moments) Formula: An alternative mathematical formulation that directly uses sorted values with a weighted sum, avoiding the need for double argsort ranks.
   • Pros: Most memory-efficient (one sort + one multiply).
   • Cons: Requires deriving weights for the biased/almost-fair estimators.
4. In-place Sort + Scatter: Replace argsort().argsort() with sort() + scatter_() to avoid creating intermediate rank tensors.
   • Pros: ~33% less intermediate memory allocation.
   • Cons: Slightly more complex code readability.

My thoughts:
My preference would be combining (1) Chunking + (2) no_grad as the safest default, as it guarantees we avoid OOMs regardless of scale.

However, before I implement this, I wanted to get your input on a few things:

• Chunking: Should chunk_size be configurable via the function signature, or hardcoded?
• Gradients: Do you foresee CRPS ever being used as a differentiable training loss in this repo? (If yes, we can't use no_grad).
• Formulation: Do you have any preference for exploring the PWM formula over the rank-based one?

Happy to implement whichever direction you think fits the project best!

[joeloskarsson]

Hi, weighing in on the OOM aspect. In general this is not super crucial to tackle optimally here, and something that can be optimized further later.

I also prefer approach 1 here, and configurable chunk size as keyword arguments, defaulting to no such chunking. In particular, I don't think this should be based on the number of ens members as in the prob_model_lam code, that was a bit arbitrary.

I would strongly disagree with adding in no_grad here, since it is up to the caller (and this is the case in e.g. ForecasterModule) to make sure that no gradients are tracked when computing metrics. We do also want CRPS as a training objective, e.g. for #648, so there we want gradients.

I don't understand the difference between the PWM approach (suggestion 3) and the current one, but this does not seem like it has potential to save that much. I thought we are following the PWM formulation as given in Zamo and Naveau. Doing in-place sorting is risky, and better to be avoided imo.

[shauryam2807]

Hello
Thanks @joeloskarsson for the clear direction! That makes a lot of sense.

To confirm, here's my plan based on your feedback:

1. Add chunk_size=None as a keyword argument to crps_ens — when None (default), the full grid is processed at once (current behavior). When set, the spread term computation will be batched over the grid dimension in chunks of the given size.

2. No @torch.no_grad() — agreed, this should be the caller's responsibility, especially since CRPS will be used as a training objective in feat: add latent encoder/decoder infrastructure for Graph-EFM port #648.

3. No in-place sorting — will keep the current safe approach.

4. PWM / rank-based formulation — you're right, the current implementation already follows Zamo and Naveau, so no change needed there.

I'll push the chunking update shortly. Thanks again for this

[joeloskarsson]

I would chunk over some other dim than the grid if possible, since it seems tricky to choose a reasonable chunk size for that. Or make the choice of dim also an option.

[Sir-Sloth-The-Lazy]

@shauryam2807, I hope you can get this work done a little faster, sorry for prying but this work is really important for visualisation PR #612

[shauryam2807]

Hy @Sir-Sloth-The-Lazy Yes, I know it's been delayed due to a lot of other work I currently have on my plate. I do have a plan for this and will get it done as soon as possible!

[Sir-Sloth-The-Lazy]

@shauryam2807 did you have sometime complete this work ? It would be really great if you did 😃 I am always here to help you have too much on your plate ! Hope you are doing well.