[1849][model] Add optional latent Zarr writer by evenmn · Pull Request #1860 · ecmwf/WeatherGenerator

evenmn · 2026-02-17T13:55:40Z

Description

This PR introduces a writer for the latent vector. To avoid additional config options, I decided to add "latent" as a special case of a stream output, e.g:

uv run inference ... --options test_config.output.streams=[ERA5,latent]

The latent output contains the latent vector itself, metadata to map healpix cells to coordinates, as well as extra features which may be useful for explainable AI applications. The combined ERA5 and latent output file takes the form (here for healpix level 5 and 2048 latent channels):

├── 0                                                                                                                                                                                                              
│   ├── ERA5                                                                                                                                                                                                       
│   │   ├── 0                                                                                                                                                                                                      
│   │   │   └── source                                                                                                                                                                                             
│   │   │       ├── coords (40320, 2) float32
│   │   │       ├── data (40320, 2) float32
│   │   │       ├── geoinfo (40320, 0) float32
│   │   │       └── times (40320,) datetime64
│   │   ├── 1
│   │   │   ├── prediction
│   │   │   │   ├── coords (40320, 2) float32
│   │   │   │   ├── data (40320, 1, 1) float32
│   │   │   │   ├── geoinfo (40320, 0) float32
│   │   │   │   └── times (40320,) datetime64
│   │   │   └── target
│   │   │       ├── coords (40320, 2) float32
│   │   │       ├── data (40320, 1) float32
│   │   │       ├── geoinfo (40320, 0) float32
│   │   │       └── times (40320,) datetime64
...
│   └── latent
│          ├── 0                                                                                                                                                                                                      
│          │   ├── coords (12288, 2) float32                                                                                                                                                                          
│          │   ├── geoinfo (12288, 0) float32                                                                                                                                                                         
│          │   ├── tokens (12288, 2048) float32                                                                                                                                                                 
│          │   ├── class_token (0, 2048) float32                                                                                                                                                         
│          │   ├── register_tokens (0, 2048) float32                                                                                                                                                     
│          │   └── times (12288,) datetime64 
...

Going further, we can visualize the latent vector:

unzip validation_chkpt00000_rank0000.zip -d test.zarr

import zarr
import cartopy.crs as ccrs
import matplotlib.pyplot as plt

z = zarr.open("test.zarr", mode="r", zarr_version=3)

coords = z["0/latent/0/coords"][:]
latent = z["0/latent/0/latent_state"][:]

lat = coords[:,0]
lon = coords[:,1]
values = latent[:,0]

fig = plt.figure(figsize=(12,6))
ax = plt.axes(projection=ccrs.PlateCarree())

sc = ax.scatter(lon, lat, c=values, s=10, transform=ccrs.PlateCarree())

ax.coastlines()
plt.colorbar(sc, label="latent feature 0")

plt.savefig("latent_feature_0.png", dpi=300)

Healpix level 4:

Healpix level 5:

Issue Number

Closes #1849

Is this PR a draft? Mark it as draft.

Checklist before asking for review

I have performed a self-review of my code
My changes comply with basic sanity checks:
- I have fixed formatting issues with ./scripts/actions.sh lint
- I have run unit tests with ./scripts/actions.sh unit-test
- I have documented my code and I have updated the docstrings.
- I have added unit tests, if relevant
I have tried my changes with data and code:
- I have run the integration tests with ./scripts/actions.sh integration-test
- (bigger changes) I have run a full training and I have written in the comment the run_id(s): launch-slurm.py --time 60
- (bigger changes and experiments) I have shared a hegdedoc in the github issue with all the configurations and runs for this experiments
I have informed and aligned with people impacted by my change:
- for config changes: the MatterMost channels and/or a design doc
- for changes of dependencies: the MatterMost software development channel

…ams: ['..., latent']) Signed-off-by: evenmn <[email protected]>

clessig · 2026-02-18T18:57:36Z

@evenmn : can you please make sure the linter and unit tests pass.

clessig · 2026-02-18T18:58:23Z

@grassesi : can you have a look at the changes in io.py. How far are you with refactoring the output writing?

Signed-off-by: evenmn <[email protected]>

…Generator into feature/latent-zarr-writer

grassesi

Overall I like the idea of treating latent output just as another stream. One thing that needs to be addressed is that there is some masking logic that prevents the latent stream from trying to be processed during evaluation (esp. the to_xarray call will not work). Otherwise just some stylistic remarks.
I saw in the Issue that the latents should be eventually exposed via an JSON API? Would it make sense to already implement this and not try to piggyback on ZarrIO?

grassesi · 2026-02-19T09:46:08Z

+        # additionally yield latent output items if a latent stream name was provided
+        if self.latent_stream_name is not None and self.latents:
+            for s, fo_s in itertools.product(self.samples, self.forecast_steps):
+                key = ItemKey(int(s), int(fo_s), self.latent_stream_name)
+                latent_item = self._make_latent_item(key)
+                if latent_item is not None:
+                    yield latent_item
+


Please try to wrap this logic into the above for iteration: having two yielding for loops works but is confusing at best. I see two alternatives:

Mix the writing of latent items into the normal loop. Something like

for s, fo_s, fi_s in itertools.product( self.samples, self.forecast_steps, self.streams.keys() ): key = ItemKey(int(s), int(fo_s), fi_s) if fi_s == LATENT_STREAM: latent_item = self._make_latent_item(key) if latent_item is not None: yield latent_item else: yield self.extract(ItemKey(int(s), int(fo_s), fi_s))

have the writing of latent items in a separate method:

def latent_items(self): if self.latents: for s, fo_s in itertools.product(self.samples, self.forecast_steps): key = ItemKey(int(s), int(fo_s), LATENT_STREAM) latent_item = self._make_latent_item(key) if latent_item is not None: yield latent_item ... with zarrio_writer(config.get_path_results(cf, mini_epoch)) as zio: for subset in data.items(): zio.write_zarr(subset) for latent in data.latent_items(): zio.write_zarr(latent)

Option 2. is maybe a bit more clearer and more equivalent to the current solution and also provides more flexibility for the future. But Option 1 would be also fine with me.

grassesi · 2026-02-19T09:59:25Z

+    # optional name to use for latent pseudo-stream when yielding latent items
+    latent_stream_name: str | None = None


Please use a named constant for this.

evenmn · 2026-02-19T12:30:49Z

Overall I like the idea of treating latent output just as another stream. One thing that needs to be addressed is that there is some masking logic that prevents the latent stream from trying to be processed during evaluation (esp. the to_xarray call will not work). Otherwise just some stylistic remarks. I saw in the Issue that the latents should be eventually exposed via an JSON API? Would it make sense to already implement this and not try to piggyback on ZarrIO?

Thanks for your feedback. Exposing the latent space via an JSON API is useful when running the model operationally. However, I still think we should export the latent state as a Zarr file, since this is useful for other applications, for instance explanable AI.

clessig · 2026-02-19T16:30:10Z

Overall I like the idea of treating latent output just as another stream. One thing that needs to be addressed is that there is some masking logic that prevents the latent stream from trying to be processed during evaluation (esp. the to_xarray call will not work). Otherwise just some stylistic remarks. I saw in the Issue that the latents should be eventually exposed via an JSON API? Would it make sense to already implement this and not try to piggyback on ZarrIO?

Thanks for your feedback. Exposing the latent space via an JSON API is useful when running the model operationally. However, I still think we should export the latent state as a Zarr file, since this is useful for other applications, for instance explanable AI.

Yes, json API is a separate issue and this PR should address writing the latent space to disc as zarr.

clessig

Overall looks good. But validation_io will be refactored and it should be discussed how to best do the latent output going forward.

clessig · 2026-02-18T18:51:04Z

+    output_streams = {}
+    for name in output_stream_names:
+        if name == "latent":
+            latent_stream_name = name


I don't understand the logic here. Wouldn't it be enough to have

if "latent" in output_stream_names

in l 158? Do we expect to have have multiple latent states?

clessig · 2026-02-18T18:55:23Z

+                per_sample = {}
+                for lname, lval in latent_pred.items():
+                    if isinstance(lval, LatentState):
+                        for field in ("z_pre_norm", "patch_tokens", "register_tokens", "class_token"):


The latent state that should be relevant for the output are the patch_tokens. These are used for the decoder. To be fully future proof we could have an argument which part of LatentState is written although it might be over-engineering.

clessig · 2026-02-18T18:56:11Z


+    # collect latent outputs per forecast step and per sample (optional)
+    latents_all = []
+    if latent_stream_name is not None:


This should go to a separate function.

Signed-off-by: evenmn <[email protected]>

grassesi · 2026-02-23T10:55:09Z

-    latent_stream_name: str | None = None
+    latent_stream_name: str | None = LATENT_STREAM


Please remove: since you are always using the default here anyway it makes no difference if self.latent_stream_name or LATENT_STREAM is used. But using using latent_stream_name clutters up the namespace/interface of OutputBatchData

grassesi · 2026-02-23T10:58:21Z

self.latent_stream_name is currently an alias to LATENT_STREAM which is never None. So please just use if self.latents (This should not disregard my previous comment on these lines.)

grassesi · 2026-02-23T11:10:08Z

    stream_names = [stream.name for stream in cf.streams]
+    # include known pseudo-stream names (e.g. latent) so they are treated as known
+    if io.LATENT_STREAM not in stream_names:
+        stream_names.append(io.LATENT_STREAM)


Please remove and just use:

if io.LATENT_STREAM in output_stream_names: output_streams[io.LATENT_STREAM] = None

in ll. 136

grassesi · 2026-02-23T11:11:37Z

    for name in output_stream_names:
-        if name == "latent":
+        if name == io.LATENT_STREAM:
            latent_stream_name = name


Please remove this if clause, instead implement the suggestion I commented above

grassesi · 2026-02-23T11:12:35Z

Please use io.LATENT_STREAM here.

grassesi · 2026-02-23T12:41:57Z

Overall looks good. But validation_io will be refactored and it should be discussed how to best do the latent output going forward.

The choosen approach should translate relatively well into the refactored version.

…get_latent_output' Signed-off-by: evenmn <[email protected]>

evenmn · 2026-02-23T13:16:02Z

Thanks to both of you for the feedback, it truly improved this PR. I believe I have incorporated the suggested changes in my latest commit, but I still need to test the implementation before it gets merged in.

clessig · 2026-02-23T21:18:09Z

Thanks to both of you for the feedback, it truly improved this PR. I believe I have incorporated the suggested changes in my latest commit, but I still need to test the implementation before it gets merged in.

Can you please test it as far as you can, and then I have a final look.

Signed-off-by: evenmn <[email protected]>

MatKbauer · 2026-04-23T11:08:19Z

Do the written latents have a notion of space and time, i.e., are latent healpix tokens attributed with their lat/lon and time coordinates?
EDIT: Found coords and time here. Looks good to me. @evenmn, can you update to develop once more so that we can merge it?

evenmn · 2026-04-23T11:16:10Z

Do the written latents have a notion of space and time, i.e., are latent healpix tokens attributed with their lat/lon and time coordinates?

@MatKbauer It comes with lat/lon, you can see the output dimensions in the first comment above. The time dimension is wrong (it has coordinate dimensions), but so is the time dimension in the ERA5 outputs

…r-writer # Conflicts: # integration_tests/small1.yaml # integration_tests/small1_test.py # packages/common/src/weathergen/common/config.py # src/weathergen/train/trainer.py

MatKbauer

I successfully wrote some latent states to disk and appreciate they are contained seamlessly in the validation_chkpt00000_rank0000.zip.

Writing the latents with the 6x1h model resulted in some conflicts and I have tested and proposed a solution provided by Claude. There are some other small comments still, please have a look @evenmn.

With the new .zip instead of .zarr output, I'd suggest to update the PR description where the written latents can be plotted as follows:

store = ZipStore("validation_chkpt00000_rank0000.zip", mode="r")
z = zarr.open_group(store=store, zarr_format=3, mode="r")

That is, add the open_group explicitly to be compatible with the .zip format, which now is the default.

Also, please add a note to the PR description telling how latents can be written, i.e., when passing an according stream info to the inference command, such as

uv run inference ... --options test_config.output.streams=[ERA5,latent]

Edit: just realised that it is mentioned already in the description. Maybe make it more explicit with the uv run inference prefix.

@clessig, can you still have a look and merge if you agree?

…cition Signed-off-by: evenmn <[email protected]>

Removed group initialization in latent writer Co-authored-by: Matthias Karlbauer <[email protected]>

Move and clean up group attributes dict Co-authored-by: Matthias Karlbauer <[email protected]>

Remove duplicated lines Co-authored-by: Matthias Karlbauer <[email protected]>

Signed-off-by: evenmn <[email protected]>

…-writer # Conflicts: # src/weathergen/model/attention.py # src/weathergen/model/positional_encoding.py

… a stale devleop Signed-off-by: evenmn <[email protected]>

Signed-off-by: evenmn <[email protected]>

…sary files Signed-off-by: evenmn <[email protected]>

Signed-off-by: Even Marius Nordhagen <[email protected]>

MatKbauer

Thanks for updating this, @evenmn! I have tested it again successfully and found the latents in the validation.zarr. Before we can merge, please address the comments I've left.

MatKbauer · 2026-06-18T10:39:18Z

+            datasets.append(ds)
+
+        # TODO: missing forecast offset
+        return OutputItem(key=key, forecast_offset=None, latent=datasets)


Is there a reason for which we don't have forecast_offset included here?

The idea here was that the latent-only items do not need the offset for target/prediction inclusion. If forecast_offset is 1 and forecast_step is 0, the offset_key method will return -1, which can cause strange behavior. I may be wrong

I understand. The core question is: Do we want to write the initial condition t_0 as well when writing forecast latents, or do we start writing at $t_0+\Delta t$. The current setup would not write t_0, which is in line with running a forecast after all.

Yes, I think you're right. For a 40 step forecast the latent states are indexed 0 to 39, which might be confusing since index 0 is actually $t_0+\Delta t$. I think it would be useful to include the initial encoded state ($t_0$) as well, and have indices 0 to 40. What do you think?

Good idea, to include the IC. This might mess up definitions at other places, though, as in forecast mode with offset: 1, we only write predictions and targets starting at $t_0 + \Delta t$. @clessig suggested to only start writing the latent at $t_0 + \Delta t$ for consistency. If we want to write latents for the initial condition, we have to run inference differently, e.g., not in forecast mode.

That is, when running in forecast mode, we want to have everything in step 0 to be empty. No predictions, targets, and latents.

If I'm not mistaken, we could realise this by changing this line in the latent writer from

now

for t_idx, latents_for_step in enumerate(data.latents):

to

for t_idx, latents_for_step in enumerate(data.latents, start=data.forecast_offset):

Signed-off-by: Even Marius Nordhagen <[email protected]>

… latents_in_sample Signed-off-by: Even Marius Nordhagen <[email protected]>

…en and lantent_state_register_token -> register_token Signed-off-by: Even Marius Nordhagen <[email protected]>

Signed-off-by: evenmn <[email protected]>

feat(output): add optional latent Zarr writer (enable via output.stre…

0e798b7

…ams: ['..., latent']) Signed-off-by: evenmn <[email protected]>

github-project-automation Bot added this to WeatherGen-dev Feb 17, 2026

evenmn and others added 4 commits February 19, 2026 08:34

Merge branch 'develop' into feature/latent-zarr-writer

9fd2193

Fixed linting

972ba6b

Signed-off-by: evenmn <[email protected]>

Fixed unit tests

4c538c1

Signed-off-by: evenmn <[email protected]>

Merge branch 'feature/latent-zarr-writer' of github.com:metno/Weather…

62b8b92

…Generator into feature/latent-zarr-writer

grassesi requested changes Feb 19, 2026

View reviewed changes

github-project-automation Bot moved this to In Progress in WeatherGen-dev Feb 19, 2026

github-actions Bot added data Anything related to the datasets used in the project infra Issues related to infrastructure model Related to model training or definition (not generic infra) labels Feb 19, 2026

clessig reviewed Feb 23, 2026

View reviewed changes

evenmn added 2 commits February 23, 2026 11:22

Using a names constant instead of 'latent' directly

f614f86

Signed-off-by: evenmn <[email protected]>

latent_stream_name now defaults to a named constant

22a8462

Signed-off-by: evenmn <[email protected]>

grassesi requested changes Feb 23, 2026

View reviewed changes

Moved logics related to getting latent state into separate function '…

0e5b118

…get_latent_output' Signed-off-by: evenmn <[email protected]>

tjhunter added the app label Feb 26, 2026

tjhunter added 6 commits March 5, 2026 09:56

fix test

d6c0706

fix test

9b6eb0e

fix test

0315054

fix test

ce7ec67

fix test

b8407ed

fix test

21994ff

Clean up

31c8919

Signed-off-by: evenmn <[email protected]>

Merge remote-tracking branch 'origin/develop' into feature/latent-zar…

0599f57

…r-writer # Conflicts: # integration_tests/small1.yaml # integration_tests/small1_test.py # packages/common/src/weathergen/common/config.py # src/weathergen/train/trainer.py

MatKbauer requested changes Apr 24, 2026

View reviewed changes

Comment thread packages/common/src/weathergen/common/io.py Outdated

Comment thread src/weathergen/utils/validation_io.py Outdated

Comment thread src/weathergen/utils/validation_io.py

Comment thread src/weathergen/utils/validation_io.py Outdated

evenmn and others added 6 commits May 8, 2026 14:06

Using the same notation for latent data as for source, target and pre…

d0cb859

…cition Signed-off-by: evenmn <[email protected]>

Update src/weathergen/utils/validation_io.py

8d7769e

Removed group initialization in latent writer Co-authored-by: Matthias Karlbauer <[email protected]>

Update src/weathergen/utils/validation_io.py

a1d3c02

Move and clean up group attributes dict Co-authored-by: Matthias Karlbauer <[email protected]>

Update src/weathergen/utils/validation_io.py

f9d4e9f

Remove duplicated lines Co-authored-by: Matthias Karlbauer <[email protected]>

Merge branch 'develop' into feature/latent-zarr-writer

c52e27c

linter

6f5f8ca

Signed-off-by: evenmn <[email protected]>

grassesi mentioned this pull request Jun 8, 2026

Long rollout related issues #2465

Open

7 tasks

evenmn and others added 7 commits June 17, 2026 11:42

Merge remote-tracking branch 'ecmwf/develop' into feature/latent-zarr…

4bc6c7e

…-writer # Conflicts: # src/weathergen/model/attention.py # src/weathergen/model/positional_encoding.py

Fix RoPE argument compatibility

32fef4e

Fix validation latent writer lint

baf730e

Actually merging with teh right develop now, the last commit was with…

dfe40c1

… a stale devleop Signed-off-by: evenmn <[email protected]>

Reverted some minor, unnecessary changes

3889c75

Signed-off-by: evenmn <[email protected]>

Reverted changes in model.py and trainer.py to not touch upon unneces…

9684c72

…sary files Signed-off-by: evenmn <[email protected]>

Going back to the original latent output structure

3eb2d9e

Signed-off-by: Even Marius Nordhagen <[email protected]>

MatKbauer reviewed Jun 18, 2026

View reviewed changes

evenmn and others added 9 commits June 19, 2026 08:53

Merge branch 'develop' into feature/latent-zarr-writer

e547681

Clarify latent stream handling in validation IO

8c70058

Signed-off-by: Even Marius Nordhagen <[email protected]>

Removed silent overwrite behavior

25515ad

Signed-off-by: Even Marius Nordhagen <[email protected]>

Renamed latents_for_Step -> latents_in_step and latents_for_sample ->…

367fd92

… latents_in_sample Signed-off-by: Even Marius Nordhagen <[email protected]>

Renamed latent_state -> tokens, latent_state_class_token -> class_tok…

672e275

…en and lantent_state_register_token -> register_token Signed-off-by: Even Marius Nordhagen <[email protected]>

Merge branch 'develop' into feature/latent-zarr-writer

12b6d2c

Merge branch 'develop' into feature/latent-zarr-writer

729f7d6

Merge branch 'develop' into feature/latent-zarr-writer

dad70ba

Lint

4d26f9f

Signed-off-by: evenmn <[email protected]>

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

		# optional name to use for latent pseudo-stream when yielding latent items
		latent_stream_name: str \| None = None

		latent_stream_name: str \| None = None
		latent_stream_name: str \| None = LATENT_STREAM

Uh oh!

Conversation

evenmn commented Feb 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Issue Number

Checklist before asking for review

Uh oh!

clessig commented Feb 18, 2026

Uh oh!

clessig commented Feb 18, 2026

Uh oh!

grassesi left a comment

Choose a reason for hiding this comment

Uh oh!

grassesi Feb 19, 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!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

evenmn commented Feb 19, 2026

Uh oh!

clessig commented Feb 19, 2026

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!

grassesi commented Feb 23, 2026

Uh oh!

evenmn commented Feb 23, 2026

Uh oh!

clessig commented Feb 23, 2026

Uh oh!

MatKbauer commented Apr 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

evenmn commented Apr 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

MatKbauer left a comment • 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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

MatKbauer 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!

Uh oh!

evenmn commented Feb 17, 2026 •

edited

Loading

grassesi Feb 19, 2026 •

edited

Loading

MatKbauer commented Apr 23, 2026 •

edited

Loading

evenmn commented Apr 23, 2026 •

edited

Loading

MatKbauer left a comment •

edited

Loading