Rollout update - #2749
Open
sophie-xhonneux wants to merge 93 commits into
Open
Conversation
This reverts commit 973f3ef.
Co-authored-by: Tharmeekan <tsivalingam@ethz.ch>
…cmwf#2414) * Handle scalar datetimes in time encoders * Keep single target times one-dimensional --------- Co-authored-by: Christian Lessig <christian.lessig@ecmwf.int>
…2412) * ensure samples_per_mini_epoch >= world_size*batch_size * updated warning log
* Fix logging when NaNs occur * fix nan variables being included in terminal logged loss, replace nan by 0 * Removing incorrect reset --------- Co-authored-by: Sebastian Hickman <seb.hickman@gmail.com>
…fied projection logic (ecmwf#2401) * added arctic region for plotting in stereographic projection and modified projection logic * new projection logic in plotter * restored attention and positional_encoding * added icon, arome, uwc-west regions * removing accidental file --------- Co-authored-by: Matteo Broccoli <mbroccol@login01.leonardo.local> Co-authored-by: iluise <72020169+iluise@users.noreply.github.com>
* Align usage of stream configs in model with dict semantics. * parse stream configs into dictconfig instead of list * Align data sampling to dict semantics, wrap self.streams into self.streams_datasets * change list to dictconfig for streams * fixes * ruffed * fix backward compatibility. * address comments * linter * Adjust incoming fresh code. * fix unit tests, rephrase comment
* multiple ranks in evaluation * shorter version * shorter version * Minor correction to ranks PR (ecmwf#2420) * Fix minor bug * Minor correction * Linting --------- Co-authored-by: Savvas Melidonis <79579567+SavvasMel@users.noreply.github.com>
* remove dead code from train_logger * fix plot_loss_per_run * fix filter logic * address comments --------- Co-authored-by: Christian Lessig <christian.lessig@ecmwf.int>
* Changes to allow for score animations. Addition possibility for log-scale colorbar * Linting * Edit config files * Change to score map function name, move score function to new utils file, generalize single_animation function * Minor edits * Linting * Minor corrections to clarify the sa variable --------- Co-authored-by: iluise <72020169+iluise@users.noreply.github.com>
Co-authored-by: iluise <72020169+iluise@users.noreply.github.com>
* Delete unnecessary code * Linting
…text in legend entries (ecmwf#2433) * Made legend font size and columns configurable and removed redundant string from plot legend entries and moved this to title * Reverting unnecessary changes * Put losses per run behind command line arg (typically not used currently)
…data (ecmwf#2422) * Custom data reader for operational analysis that accounts for actual availability * Reduce mem footprint during training; add DataReaderAnemoiOperan * Config files * Fixed corner case: sample 0 that has no precedessor * Added additional check * Move OperAn data reader to extras readers * Cleaning up * Remove incorrect general setting for base permuation * Remove incorrect time_window_step settings * Linting * Improved encapsulation
* Fixes to forcing fields in io * Fixed linting
…on (ecmwf#2436) * time series scores * add offset for lead time without source * add fstep to ratio plots * fixed range for ratio plots * if the source is not available as a group in the zarr, instead use source interval attribute from target or prediction * add comments * lint * fix label in ratio plot * fix init time --------- Co-authored-by: Sebastian Hickman <seb.hickman@ecmwf.int>
* Move data reader fesom to extra package * ruff --------- Co-authored-by: Christian Lessig <christian.lessig@ecmwf.int>
* remove dead code from train_logger * fix plot_loss_per_run * fix filter logic * address comments * Smaller decoder for SYNOP * Custom data reader for operational analysis that accounts for actual availability * Reduce mem footprint during training; add DataReaderAnemoiOperan * Config files * Reduce spoofing to two tokens * Fixed corner case where there is no t_idx-1: very first sample * Fixes needed to run inference and write output with spoofed streams * Save soure_raw in val mode * Add additional sanity check * Reformatting * Updated config * Updated configs * Configurable plot_train * remove clim align when not necessary * Iluise/develop/init time plot (ecmwf#2457) * Sgrasse/develop/342 streams as dict (ecmwf#2242) * Align usage of stream configs in model with dict semantics. * parse stream configs into dictconfig instead of list * Align data sampling to dict semantics, wrap self.streams into self.streams_datasets * change list to dictconfig for streams * fixes * ruffed * fix backward compatibility. * address comments * linter * Adjust incoming fresh code. * fix unit tests, rephrase comment * combine inference from multiple ranks in evaluation (ecmwf#2400) * multiple ranks in evaluation * shorter version * shorter version * Minor correction to ranks PR (ecmwf#2420) * Fix minor bug * Minor correction * Linting --------- Co-authored-by: Savvas Melidonis <79579567+SavvasMel@users.noreply.github.com> * Sgrasse/develop/2353 plot train (ecmwf#2354) * remove dead code from train_logger * fix plot_loss_per_run * fix filter logic * address comments --------- Co-authored-by: Christian Lessig <christian.lessig@ecmwf.int> * Score animations with fixed colorbar (ecmwf#2402) * Changes to allow for score animations. Addition possibility for log-scale colorbar * Linting * Edit config files * Change to score map function name, move score function to new utils file, generalize single_animation function * Minor edits * Linting * Minor corrections to clarify the sa variable --------- Co-authored-by: iluise <72020169+iluise@users.noreply.github.com> * apply scale_z_channels (ecmwf#2417) Co-authored-by: iluise <72020169+iluise@users.noreply.github.com> * Fix bug in the evaluation wegen_reader.py (ecmwf#2431) * Delete unnecessary code * Linting * time series scores * add offset for lead time without source * add fstep to ratio plots * Made legend font size and columns configurable and removed redundant text in legend entries (ecmwf#2433) * Made legend font size and columns configurable and removed redundant string from plot legend entries and moved this to title * Reverting unnecessary changes * Put losses per run behind command line arg (typically not used currently) * Add data reader for Oper-an that accounts for actual availability of data (ecmwf#2422) * Custom data reader for operational analysis that accounts for actual availability * Reduce mem footprint during training; add DataReaderAnemoiOperan * Config files * Fixed corner case: sample 0 that has no precedessor * Added additional check * Move OperAn data reader to extras readers * Cleaning up * Remove incorrect general setting for base permuation * Remove incorrect time_window_step settings * Linting * Improved encapsulation * Fixes to forcing fields in io (ecmwf#2442) * Fixes to forcing fields in io * Fixed linting * Fix spoofing so that it only uses 2 tokens to avoid large memory footprint (ecmwf#2449) * fixed range for ratio plots * if the source is not available as a group in the zarr, instead use source interval attribute from target or prediction * add comments * lint * fix label in ratio plot * fix init time --------- Co-authored-by: Simon Grasse <161459968+grassesi@users.noreply.github.com> Co-authored-by: iluise <72020169+iluise@users.noreply.github.com> Co-authored-by: Savvas Melidonis <79579567+SavvasMel@users.noreply.github.com> Co-authored-by: Julius Polz <56866670+jpolz@users.noreply.github.com> Co-authored-by: iluise <luise.ilaria@gmail.com> Co-authored-by: Sebastian Hickman <seb.hickman@ecmwf.int> * lint * Revert "Merge branch 'clessig/develop/smaller_decoder_synop' into iluise.develop/fix-climatology" This reverts commit 9176522, reversing changes made to b7b6c9d. * julius comments --------- Co-authored-by: Simon Grasse <s.grasse@fz-juelich.de> Co-authored-by: Simon Grasse <161459968+grassesi@users.noreply.github.com> Co-authored-by: lessig2 <christian.lessig@ecmwf.int> Co-authored-by: Savvas Melidonis <79579567+SavvasMel@users.noreply.github.com> Co-authored-by: Julius Polz <56866670+jpolz@users.noreply.github.com> Co-authored-by: Sebastian Hickman <seb.hickman@ecmwf.int>
* first working implementation of spectra * fix zonal (ukmet) psd * remove old psd implementation * remove dead code * shorter functions * implement inverse transform and tests * lint * add var name to psd plot title * PR comments * add comments about limitations * lint * automatically detect grids * fix eval config * implement comments * lint
* initial casestudy implementation * initial linting * added some docstrings and comments * encapsulated dependencies, added instructions * changed path to relative, added README --------- Co-authored-by: buschow1 <s.buschow@fz-juelich.de>
Change default value of 'rank' from 0 to 'all'.
* Fix missing type of rnaks when specified as list * Fix import alias for omegaconf in wegen_reader.py --------- Co-authored-by: iluise <72020169+iluise@users.noreply.github.com>
* Implement regular subsampling * ruff * Fix handling of holes in dataset * Update stream config * Remove 0 as filler value --------- Co-authored-by: Savvas Melidonis <79579567+SavvasMel@users.noreply.github.com>
* Replace cf.rank==0 with utils.distributed.is_root * Restore streams_directory replacement semantics --------- Co-authored-by: wang85 <wang85@jwlogin22.juwels> Co-authored-by: wang85 <wang85@jwlogin24.juwels>
Co-authored-by: kasravi1 <j.kasravi@fz-jeulich.de> Co-authored-by: Christian Lessig <christian.lessig@ecmwf.int>
…mwf#2597) * add pdf merger * add pdf merge * feat(evaluate): filter PDF merge by run_ids from config - Rewrite pdf_merge.py to only merge PDFs whose filenames contain run_ids from the current evaluation config - Merged output files are named with the run_ids (e.g. merged_line_plots_runA_runB.pdf) - Different evaluations can coexist in the same directory without contaminating each other's merged output - Also produces per-metric grouped merges (e.g. merged_rmse_global_runA_runB.pdf) * feat(evaluate): granular plot directory structure Reorganize plot output into: <summary_dir>/<run_ids_str>/<plot_type>/<metric>/<region>/<files> - Prefix all output with run_ids_str (sorted concatenation of all run_ids from config) so different evaluations coexist cleanly - Add set_subdir(metric, region) to all plotter classes (LinePlots, ScoreCards, BarPlots, QuantilePlots) which creates and sets the metric/region subdirectory under each plot type - Update pdf_merge to recursively walk metric/region subdirs and produce per-leaf and per-plot-type merged PDFs Example final structure: plots/runA_runB/line_plots/rmse/global/compare_rmse_global_*.pdf plots/runA_runB/line_plots/rmse/global/merged_line_plots_runA_runB.pdf plots/runA_runB/line_plots/merged_line_plots_runA_runB.pdf plots/runA_runB/ratio_plots/rmse/global/... * restructure pdf merge * comments * lint
…mwf#2620) * new way of defining plots in the config * add git compat file * shorten logic * lint * Fix export bug, add IMERG N320 (ecmwf#2670) * added IMERG full res tp, fixed small bug * simplified if statement --------- Co-authored-by: buschow1 <s.buschow@fz-juelich.de> * merge PDFs and more granular folder structure for FastEval output (ecmwf#2597) * add pdf merger * add pdf merge * feat(evaluate): filter PDF merge by run_ids from config - Rewrite pdf_merge.py to only merge PDFs whose filenames contain run_ids from the current evaluation config - Merged output files are named with the run_ids (e.g. merged_line_plots_runA_runB.pdf) - Different evaluations can coexist in the same directory without contaminating each other's merged output - Also produces per-metric grouped merges (e.g. merged_rmse_global_runA_runB.pdf) * feat(evaluate): granular plot directory structure Reorganize plot output into: <summary_dir>/<run_ids_str>/<plot_type>/<metric>/<region>/<files> - Prefix all output with run_ids_str (sorted concatenation of all run_ids from config) so different evaluations coexist cleanly - Add set_subdir(metric, region) to all plotter classes (LinePlots, ScoreCards, BarPlots, QuantilePlots) which creates and sets the metric/region subdirectory under each plot type - Update pdf_merge to recursively walk metric/region subdirs and produce per-leaf and per-plot-type merged PDFs Example final structure: plots/runA_runB/line_plots/rmse/global/compare_rmse_global_*.pdf plots/runA_runB/line_plots/rmse/global/merged_line_plots_runA_runB.pdf plots/runA_runB/line_plots/merged_line_plots_runA_runB.pdf plots/runA_runB/ratio_plots/rmse/global/... * restructure pdf merge * comments * lint * address coments * add git compat file * shorten logic * solve conflict * lint --------- Co-authored-by: s6sebusc <49226935+s6sebusc@users.noreply.github.com> Co-authored-by: buschow1 <s.buschow@fz-juelich.de>
* original implementation * plumbing the seeps climatology through the plot orchestration * linting, bugfixes * remove specilaised seeps climatology file * switched to scores packages for seeps * linting * enabled user set masking parameters, updated documentation --------- Co-authored-by: buschow1 <s.buschow@fz-juelich.de> Co-authored-by: Christian Lessig <christian.lessig@ecmwf.int> Co-authored-by: iluise <72020169+iluise@users.noreply.github.com>
* new crps based on scores package and documentation * move import --------- Co-authored-by: iluise <72020169+iluise@users.noreply.github.com>
ecmwf#2694) * getting seeps scores instead of seeps error + doc string for the function * updating the seeps bits of documentation
* moving multi rank out * linting * fix type checking * fixing error message and linting * adjust for quaver parser * wrap around loop for verif * linting * channels -> stream_channels to prevent overwriting * remove typo * removing duplicate tp * add back cp --------- Co-authored-by: Ilaria Luise <luise.ilaria@gmail.com>
* initial fix with no real coordinate * hacked in the coordinate * coordinate ordering * sugggestions * linting * more changes * testing complete * ensemble members for verif, need to intuitively infer coordinate ordering * updating ensemble for verif * fixing help message * linting --------- Co-authored-by: buschow1 <s.buschow@fz-juelich.de> Co-authored-by: s6sebusc <49226935+s6sebusc@users.noreply.github.com>
* adding the experiment script * permission * adding touching files * fix the linter * update ruff everywhere
* Notes * Implementation of nc parser from pr 2500 * Notes update * Update of netcdf_parser.py for sub-6h models, and notes * Deleted print() * Deleted old implementations * Deleted old trials * Delete unecessary files * Fixing formatting issues * Rephrase comments * Formatting issue --------- Co-authored-by: Christian Lessig <christian.lessig@ecmwf.int>
* feat(output): add optional latent Zarr writer (enable via output.streams: ['..., latent']) Signed-off-by: evenmn <evenmn@mn.uio.no> * Fixed linting Signed-off-by: evenmn <evenmn@mn.uio.no> * Fixed unit tests Signed-off-by: evenmn <evenmn@mn.uio.no> * Using a names constant instead of 'latent' directly Signed-off-by: evenmn <evenmn@mn.uio.no> * latent_stream_name now defaults to a named constant Signed-off-by: evenmn <evenmn@mn.uio.no> * Moved logics related to getting latent state into separate function 'get_latent_output' Signed-off-by: evenmn <evenmn@mn.uio.no> * fix test * fix test * fix test * fix test * fix test * fix test * fix test * fix test * fix test * fix test * fix test * fix test * fix test * Pre-computing healpix coordinates for downstream use Signed-off-by: evenmn <evenmn@fys.uio.no> * Passing self.model_params to trainer object. This makes self.healpix_coords accessible from validation_io.py Signed-off-by: evenmn <evenmn@fys.uio.no> * Latent writer that splits up healpix features and extra features, and writes metadata correctly, was added Signed-off-by: evenmn <evenmn@fys.uio.no> * Moved healpix_coords calculations from ModelParams to validation_io.py. This avoids passing ModelParams to the Trainer Signed-off-by: evenmn <evenmn@fys.uio.no> * Split up extra_features into register and class tokens Signed-off-by: evenmn <evenmn@fys.uio.no> * Tracking config file Signed-off-by: evenmn <evenmn@fys.uio.no> * Split extra toksens into register toksens and class tokens correctly Signed-off-by: evenmn <evenmn@fys.uio.no> * Clean up Signed-off-by: evenmn <evenmn@fys.uio.no> * Using the same notation for latent data as for source, target and precition Signed-off-by: evenmn <evenmn@fys.uio.no> * Update src/weathergen/utils/validation_io.py Removed group initialization in latent writer Co-authored-by: Matthias Karlbauer <matthias.karlbauer@ecmwf.int> * Update src/weathergen/utils/validation_io.py Move and clean up group attributes dict Co-authored-by: Matthias Karlbauer <matthias.karlbauer@ecmwf.int> * Update src/weathergen/utils/validation_io.py Remove duplicated lines Co-authored-by: Matthias Karlbauer <matthias.karlbauer@ecmwf.int> * linter Signed-off-by: evenmn <evenmn@fys.uio.no> * Fix RoPE argument compatibility * Fix validation latent writer lint * Actually merging with teh right develop now, the last commit was with a stale devleop Signed-off-by: evenmn <evenmn@fys.uio.no> * Reverted some minor, unnecessary changes Signed-off-by: evenmn <evenmn@fys.uio.no> * Reverted changes in model.py and trainer.py to not touch upon unnecessary files Signed-off-by: evenmn <evenmn@fys.uio.no> * Going back to the original latent output structure Signed-off-by: Even Marius Nordhagen <even.nordhagen@gmail.com> * Clarify latent stream handling in validation IO Signed-off-by: Even Marius Nordhagen <evnor2743@met.no> * Removed silent overwrite behavior Signed-off-by: Even Marius Nordhagen <evnor2743@met.no> * Renamed latents_for_Step -> latents_in_step and latents_for_sample -> latents_in_sample Signed-off-by: Even Marius Nordhagen <evnor2743@met.no> * Renamed latent_state -> tokens, latent_state_class_token -> class_token and lantent_state_register_token -> register_token Signed-off-by: Even Marius Nordhagen <evnor2743@met.no> * Lint Signed-off-by: evenmn <evenmn@mn.uio.no> * Bumped up the indices of the latent states, such that 0 is reserved for the encoded initial states Signed-off-by: evenmn <evenmn@mn.uio.no> * Swtiched from getattr(cf...) to cf.get(...) to get rid of pylint warning Signed-off-by: evenmn <evenmn@mn.uio.no> --------- Signed-off-by: evenmn <evenmn@mn.uio.no> Signed-off-by: evenmn <evenmn@fys.uio.no> Signed-off-by: Even Marius Nordhagen <even.nordhagen@gmail.com> Signed-off-by: Even Marius Nordhagen <evnor2743@met.no> Co-authored-by: Tim Hunter <tim.hunter@ecmwf.int> Co-authored-by: Matthias Karlbauer <matthias.karlbauer@ecmwf.int> Co-authored-by: Even Marius Nordhagen <even.nordhagen@gmail.com>
Storing the steps and converting them via chunk_idx/batch_idx lets a partial rollout be indexed without callers tracking offsets. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The predictors now take the global forecast step and let ModelOutput resolve it, so they stay correct when the output covers only part of the rollout. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The context manager had been absorbed into the comment above it, so intermediate rollout steps still built a graph. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Sizing the zero arrays from the batch and stream config instead of indexing preds[0] lets a step that holds no predictions for a stream still be written. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The caller knows which steps its data covers, so the writer no longer has to reconstruct them from the target count and the offset. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Taking the steps from ModelOutput lets a partial rollout be written, at the cost of the first chunk emitting empty entries for its leading offset steps. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Collecting the model call and the write in _process_validation_chunks gives the rollout a single place to be split into chunks. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Each chunk continues from the previous one and is written as it completes; the predictions themselves are still retained for the loss computation. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
github-actions
Bot
added
data
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
20 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Issue Number
Is this PR a draft? Mark it as draft.
Checklist before asking for review
./scripts/actions.sh lint./scripts/actions.sh unit-test./scripts/actions.sh integration-testlaunch-slurm.py --time 60