CRPS decoder integrated in latest develop by saibr · Pull Request #2578 · ecmwf/WeatherGenerator

saibr · 2026-07-01T21:28:03Z

Description

This is the cleaned CRPS decoder code integrated in the latest develop on 1/7. It is based on the code from https://github.com/saibr/WeatherGenerator/tree/crps-decoder. This version should be easy to merge without conflicts.

In the config, the following lines should be added to use the CRPS decoder:

decoder_ens_latent_perturbation : 
  num_members: 2    # 0 / 1 = disabled; >1 = number of ensemble members
  sigma_init: 0.2  # initial noise standard deviation
  sigma_learnable: True  # true = nn.Parameter; false = fixed buffer

Besides the changes on the decoder, I made a modification on the loss function, to make sure the model can run with the kernel_crps loss without including the mse loss, by using the following setup:

losses : {
    "physical": {
        type: LossPhysical,
        loss_fcts: {"mse": null, "kernel_crps": {}, },
        },
    }

In case there is an alternative for this that I missed, these modifications can be discarded.

clessig

Thanks for the implementation. Left some comments: I think the code can be a bit tidier overall. Three questions:

Should we generate a warning if training is with multiple ens members and CRPS perturbations but without CRPS loss (or another ens loss)?
Do we know how the memory requirements change with the CRPS decoder?
Do we have any training results?

For the config, can we make it:

decoder_ens_latent_perturbation : 
  num_members: 2    # 0 / 1 = disabled; >1 = number of ensemble members
  sigma_init: 0.2  # initial noise standard deviation
  sigma_learnable: True  # true = nn.Parameter; false = fixed buffer

This is clearer and more composable (we still need to refactor the model config into dicts).

clessig · 2026-07-02T07:10:03Z

-            (s[0] * s[1] + 1,), fill_value=9, dtype=torch.int32, device=tokens_nbors.device
-        )
-        tokens_nbors_lens[0] = 0
+        # remove register and class tokens


Can we encapsulate this into a separate function

Which lines are you referring to for the separate function? This comment refers only to the single line below

The whole block that is green (which I think is a rewrite of some code that was there before and some new code).

I created _build_ensemble_latent_tokens for this

clessig · 2026-07-02T07:10:21Z

+        # remove register and class tokens
+        tokens = tokens[:, self.num_aux_tokens:]
+
+        B = len(batch)


The single letter vars make the code more difficult to read

clessig · 2026-07-02T07:11:44Z

+            tokens_tiled = tokens
+        assert tokens_tiled.shape == (M * B, H * Q, D)
+
+        MB = M * B


Is this variable really useful? Again, short var name is an issue

clessig · 2026-07-02T07:12:51Z

+        tokens_nbors = tokens_flat[idxs.flatten()].flatten(0, 1)  # [MB*H*9*Q, D]
+
+        tokens_nbors_lens = tokens_flat.new_zeros(MB * H + 1, dtype=torch.int32)
+        tokens_nbors_lens[1:] = 9 * Q


We should not hard-code 9 but have a class-level variable with a useful name, e.g. hp_num_nbors

clessig · 2026-07-02T07:13:12Z

            ]
            t_coords_lens = [len(t) for t in t_coords]
            t_coords = torch.cat(t_coords)
+            P = t_coords.shape[0]


Again, single letter var name.

clessig · 2026-07-02T07:25:16Z

+        # Latent-perturbation noise scale (learnable or fixed)
+        # torch.zeros() is used (not torch.tensor) so this respects the meta-device context used by FSDP. The actual value is filled in reset_parameters().
+        num_mem = cf.get("latent_perturbation_num_members", 0)
+        if num_mem > 1:


why not if cf.get("latent_perturbation_num_members", 0) > 1 :

With the changed config, we should use

if cf.get("decoder_ens_latent_perturbation") is not None

or we check in the constructor and have a class variable for this.

clessig · 2026-07-02T07:27:42Z

        Args:
            model_params : Query and embedding parameters
-            fstep : Number of forecast steps
+            step : Forecast step


Why not keep fstep?

predict_decoders takes step as input, not fstep. This is not something I changed, it is used in the latest develop. I believe the comment wasn't updated before.

clessig · 2026-07-02T07:29:05Z

+            eps = torch.randn(M, B, H * Q, D, device=tokens.device, dtype=tokens.dtype)
+            tokens_tiled = (tokens.unsqueeze(0) + sigma * eps).reshape(M * B, H * Q, D)
+        else:
+            M = 1


Avoid by M = self.cf.get("latent_perturbation_num_members", 1) above. But I think it should be valid to have self.cf.get("latent_perturbation_num_members", 1) and then we also add noise for decoding.

clessig · 2026-07-02T07:29:56Z

+        # fully vectorised. Layout convention: M is the OUTER axis, B is INNER.
+        #   tokens_tiled[m*B + b] belongs to member m, batch item b
+        M = self.cf.get("latent_perturbation_num_members", 0)
+        if M > 1 and self.latent_perturbation_log_sigma is not None:


Maybe really good to have a class variable if CRPS-perturbs are used or not

clessig · 2026-07-02T07:30:22Z

-        )
-        tokens_nbors_lens[0] = 0
+        # remove register and class tokens
+        tokens = tokens[:, self.num_aux_tokens:]


If we call this tokens_tiled (or a bit later) then we don't need the else branch in 803

clessig · 2026-07-02T12:49:48Z

@saibr : can you also lint (this might also complain about single letter variables ;)) and open a PR.

saibr · 2026-07-02T15:27:28Z

Thank you for all the comments @clessig , I incorporated them. I am not sure how to deal with the scenario of num_members: 1. I added noise to it, but the current implementation of the kernel crps loss does not allow num_members to be 1. Let me know if you have a specific use case for this in mind and I can adapt it accordingly.

shmh40 · 2026-07-02T17:26:38Z

+        # Latent-perturbation noise scale (learnable or fixed)
+        self.use_latent_perturbation = (
+            self.decoder_ens_latent_perturbation is not None
+            and self.decoder_ens_latent_perturbation.get("num_members", 0) >= 1


Should this be > 1? Do we only want this to work when the user has specified multiple members? What happens if they set it as 1? Do you get a perturbed single member?

Otherwise the implementation looks ok to me, thanks @saibr !

@shmh40 Christian mentioned in a comment that perturbation on 1 member would be a valid case, so I included it, but I'm not sure about the use case. The current implementation of the crps loss does not accept one member, so in my first implementation I discarded the scenario with 0 and 1 members.

…ation_log_sigma and removed mse:null

…erator into crps_decoder_v2

clessig

Thanks for the changes. Some more comments.

clessig · 2026-07-10T10:19:35Z

+                        f"expected {num_members * t_coords.shape[0]} pts in pred, "
+                        f"got {total_target_points}"
+                    )
+                    pred = (


This statement cuts across 12 lines and is extremely difficult to read for me. Can we precompute the shapes to get this statement in 1 or 2 lines

Changed in latest commit

clessig · 2026-07-10T10:26:24Z

                if self.cf.decoder_type == "Linear":
+                    # repeat target-coord tokens once per ensemble member,
+                    # matching tokens_tiled's ordering
+                    tc_tokens_in = tc_tokens.repeat(num_members, 1)


Can't we repeat the output of line 983 instead of repeating the input? There is no ensembling happening here. Otherwise, the if self.cf.decoder_type == "Linear": branch should be in line 1002, I think.

I believe we cannot remove the repeat output instead if the input, the ensemble component is in tokens_tiled which is create by _build_ensemble_latent_tokens. If we repeat the output num_members times, we get the same prediction for each member

clessig · 2026-07-10T10:27:34Z

+                    tc_tokens_outs = []
+                    for m in range(num_members):
+                        member_tokens = tokens_tiled[m * batch_size: (m + 1) * batch_size]
+                        tokens_nbors, tokens_nbors_lens = self._gather_neighbor_tokens(


clessig · 2026-07-10T10:30:54Z

-            (s[0] * s[1] + 1,), fill_value=9, dtype=torch.int32, device=tokens_nbors.device
-        )
-        tokens_nbors_lens[0] = 0
+        ens = self._build_ensemble_latent_tokens(model_params, tokens, batch)


It looks like EnsembleLatentTokens is only used as return type for the function. Then we should not introduce it and just return a tuple/list with the values.

I agree, changed

clessig · 2026-07-10T10:35:21Z

+                    )
+
+        if self.use_latent_perturbation:
+            if self.decoder_ens_latent_perturbation.get("sigma_learnable", True):


The branches differ just by requires_grad. We can save a few lines and make the code more readable by having one statement and just having the flag conditional (either using a temp variable or by directly having self.decoder_ens_latent_perturbation.get("sigma_learnable", True) there. (self.decoder_ens_latent_perturbation is very long and leads to a lot of line breaks that make the code difficult to read; maybe use a shorter name, e.g. self.ens_latent_perturb

Yes thanks, changed it

clessig · 2026-07-10T10:38:35Z

+        )
+
+        tokens_tiled = tokens
+        if self.use_latent_perturbation:


Could one return from this function immediately at the beginning if self.use_latent_perturbation == False

…erator into crps_decoder_v2

clessig · 2026-07-13T08:12:29Z

Closes via #2623

clean crps decoder code and compatible with latest develop 1/7

f086728

github-project-automation Bot added this to WeatherGen-dev Jul 1, 2026

github-actions Bot added the model Related to model training or definition (not generic infra) label Jul 1, 2026

clessig reviewed Jul 2, 2026

View reviewed changes

Sarah Ibrahimi added 2 commits July 2, 2026 16:58

incorporated suggested modifications from PR ecmwf#2578

374d882

after linting

99c38f2

shmh40 reviewed Jul 2, 2026

View reviewed changes

removed use of a buffer and used parameter instead for latent_perturb…

ab61f4c

…ation_log_sigma and removed mse:null

github-actions Bot added the eval anything related to the model evaluation pipeline label Jul 9, 2026

saibr added 4 commits July 10, 2026 09:12

Merge branch 'develop' into crps_decoder_v2

361569e

remove crps config file

f2ab1a6

Merge branch 'crps_decoder_v2' of https://github.com/saibr/WeatherGen…

62da414

…erator into crps_decoder_v2

Merge branch 'develop' into crps_decoder_v2

8c48624

clessig reviewed Jul 10, 2026

View reviewed changes

saibr added 2 commits July 10, 2026 13:24

incorporated comments from Christian on 10/7

4db1f0f

Merge branch 'crps_decoder_v2' of https://github.com/saibr/WeatherGen…

c96c20d

…erator into crps_decoder_v2

clessig mentioned this pull request Jul 13, 2026

Implementation CRPS decoder #2623

Open

4 tasks

clessig closed this Jul 13, 2026

github-project-automation Bot moved this to Done in WeatherGen-dev Jul 13, 2026

Uh oh!

Conversation

saibr commented Jul 1, 2026 • edited by clessig Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Uh oh!

clessig left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

clessig commented Jul 2, 2026

Uh oh!

saibr commented Jul 2, 2026

Uh oh!

shmh40 Jul 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

clessig left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

clessig commented Jul 13, 2026

Uh oh!

Reviewers

Assignees

saibr commented Jul 1, 2026 •

edited by clessig

Loading

shmh40 Jul 2, 2026 •

edited

Loading