Skip to content

[big merge][1300][model] Assemble diffusion model #1310

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Draft
MatKbauer wants to merge 222 commits into
base: develop
Choose a base branch
from
Draft

[big merge][1300][model] Assemble diffusion model #1310

MatKbauer wants to merge 222 commits into from

Conversation

@MatKbauer
Copy link
Contributor

@MatKbauer MatKbauer commented Nov 20, 2025

Description

We assemble the different PRs of the latent diffusion forecast engine together to test their compatibility and perform first explorations.

Issue Number

Closes #1300

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

sophie-xhonneux and others added 30 commits October 30, 2025 17:27
@sophie-xhonneux
Abstract class for target/aux computation
3f1bb7d 
Implemented Identity class

TODO: implement EMATeacher
@sophie-xhonneux
Start implementing the EMA Teacher
03ed148 
The big question on the EMA teacher side to me is how to allow for a
fleixble teacher and student architecture that can differ

We updated some APIs of the abstract base class to allow the ema_model
forward, subject to change given the loss calculator, which is imho the
second big question mark
@Jubeku
adding loss calculator base class
28d9b22
@sophie-xhonneux
Option for constructing teacher model flexibly
192beb6
@sophie-xhonneux
Extract get batch size util function
aac7e29 
Easier to read and as batchsize gets more complicated in SSL this will
be a useful abstraction
@sophie-xhonneux
Fix mismatched dtypes in the target computation
145d18a 
It runs so far. Next steps:
 - Route all the config options
 - Start writing the loss functions to understand the state requirements
@Jubeku
abstract loss calc structure
f1e7132
@Jubeku
add abstract method to loss calculator base class
e822e12
@Jubeku
add latent loss class
d24ef48
@Jubeku
update loss calc config and rename files
c259c20
@Jubeku
restructure loss modules
a19ee16
@Jubeku
add ModelOutput dataclass
bf3e128
@MatKbauer
First draft of diffusion model
711f29b
@clessig
NOT WORKING: initial draft for index-based masking. Implemented for r…
81bd6eb 
…andom and healpix masking. Open issues with _coords_local, centroids and probably other things.
@MatKbauer
Minor modifications
f367bb4
@MatKbauer
Linter
1cc168c
@MatKbauer
Copyright attribution to EDM
48934c2
@clessig
NOT WORKING: Finished src, target still to be done.
51f437f
@MatKbauer
Adapt diffusion model to expected data structure
6046694
@MatKbauer
Corrected data retrieval to only access model_samples and not target_…
f66c9fa 
…samples
@MatKbauer
Minor correction
7e48c39
@clessig
Masking target is working in principle but errors when feeding data t…
e4a9cc0 
…o the model.
@clessig
Working version for ERA5, NPP-ATMS. Problems with SYNOP with empty ce…
a581405 
…ll handling
@clessig
Minor cleanup
9229e48
@clessig
Fixed linting
db6f285
@MatKbauer
Restructuring and correcting forward pass during inference
7866ff7
@clessig
Fixed remaining problems that occured for NPP-ATMS and SYNOP.
ec38123 
TODO:
- Forecast still needs to be adapted
- Some more cleanup of variable naming, return values etc
@clessig
Enabled support for forecast. Cleaned up some bits and pieces.
0634105
@Jubeku
merge develop
0fa60db
@Jubeku
mv streams_data declaration under if condition
cab9fbe
clessig and others added 30 commits December 9, 2025 13:54
@clessig
Missing update to validation()
e0f6cc4
@clessig
Improved robustness through sanity checking of arguments
8f097ec
@clessig
Improved handling of corner cases
6b64511
@clessig
Merge branch 'develop' of github.com:ecmwf/WeatherGenerator into shmh…
ed886e2 
…40/dev/1270-idx-global-local
@clessig
- Fixed incorrect call to get_forecast_steps() in validation
303f48a 
- Fixed interface of target_aux_calculator
@MatKbauer
[NOT WORKING] Merged current data-branch. TargetAuxCalculator argumen…
9638de8 
…ts incorrect
@clessig
More fixed to validation
7299106
@clessig
Adding stream_id
45189a4
@MatKbauer
[NOT WORKING] Added modifications from data branch
50b0a89
@clessig
Cleaned up ModelOutput class to have proper access functions and a be…
5bed792 
…tter structure
@clessig
Switched to use dict to internally represent streams_datasets
06f2e06
@clessig
Improving robustness of interface of ModelOutput class
ad5a19c
@clessig
Re-enabling model output
4f8abbb
@clessig
Ruff
d36716c
@clessig
Minor clean-ups and additional comments
b8d95b2
@clessig
Minor cleanups
081d90a
@clessig
Cleaned up handling of masks and masking metadata
6b8fe83
@MatKbauer
Resolved bugs when updating data structure
5a8ad49
@MatKbauer
Updated to new data output structure
eedaa8a
@MatKbauer
Linter
f768046
@clessig
Current working version of default_config
ca9e605
@clessig
Fixed problem with branches with old code and incomplete cleanup
f8b1ca6
@clessig
Updated to test convergence of integration test.
003b0cf
@clessig
Updated settings
f38e6d2
@clessig
Clessig/ypd/dev/1353 add tokens latent state finalization (#1452)
7e7ff8e 
* Add LatentState

* Add class and register tokens for LatentState, adjust everything accordingly

* Add option in config file + minor changes

* Add pos.emb. for register tokens + remove class tokens + minor fixes

* Minor fix

* Changed empty to zeros pe_register

* Ruffed

* Clean-up and fixed positional encoding

* Fixing things that got lost during last merge

---------

Co-authored-by: Yura Perugachi Diaz <yperugac@santis-ln001.cscs.ch>
Co-authored-by: Yura Perugachi Diaz <yperugac@santis-ln002.cscs.ch>
@clessig
Ruffed
31a0b96
@clessig
Adding sanity check for register tokens
4fe90d7
@clessig
Merge branch 'develop' of github.com:ecmwf/WeatherGenerator into shmh…
46bd7a2 
…40/dev/1270-idx-global-local
@MatKbauer
Update to latest data branch: latent_state dataclass
48dee1e
@MatKbauer
Update to latest develop with new data structure
e2c09f2
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@tjhunter tjhunter tjhunter left review comments

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

model:rollout model Related to model training or definition (not generic infra)

Projects

WeatherGen-dev
Status: In Progress

Milestone

latent diffusion model

Development

Successfully merging this pull request may close these issues.

Assemble latent diffusion forecast engine

9 participants

@MatKbauer @tjhunter @sophie-xhonneux @Jubeku @clessig @shmh40 @moritzhauschulz @kctezcan