Sophiex/dev ssl/seb temporal plus plus - #2656
Open
sophie-xhonneux wants to merge 62 commits into
Open
Conversation
* 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 <[email protected]>
…fied projection logic (#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 <[email protected]> Co-authored-by: iluise <[email protected]>
* 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 (#2420) * Fix minor bug * Minor correction * Linting --------- Co-authored-by: Savvas Melidonis <[email protected]>
* remove dead code from train_logger * fix plot_loss_per_run * fix filter logic * address comments --------- Co-authored-by: Christian Lessig <[email protected]>
* 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 <[email protected]>
Co-authored-by: iluise <[email protected]>
* Delete unnecessary code * Linting
…text in legend entries (#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 (#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 (#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 <[email protected]>
* Move data reader fesom to extra package * ruff --------- Co-authored-by: Christian Lessig <[email protected]>
* 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 (#2457) * Sgrasse/develop/342 streams as dict (#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 (#2400) * multiple ranks in evaluation * shorter version * shorter version * Minor correction to ranks PR (#2420) * Fix minor bug * Minor correction * Linting --------- Co-authored-by: Savvas Melidonis <[email protected]> * Sgrasse/develop/2353 plot train (#2354) * remove dead code from train_logger * fix plot_loss_per_run * fix filter logic * address comments --------- Co-authored-by: Christian Lessig <[email protected]> * Score animations with fixed colorbar (#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 <[email protected]> * apply scale_z_channels (#2417) Co-authored-by: iluise <[email protected]> * Fix bug in the evaluation wegen_reader.py (#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 (#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 (#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 (#2442) * Fixes to forcing fields in io * Fixed linting * Fix spoofing so that it only uses 2 tokens to avoid large memory footprint (#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 <[email protected]> Co-authored-by: iluise <[email protected]> Co-authored-by: Savvas Melidonis <[email protected]> Co-authored-by: Julius Polz <[email protected]> Co-authored-by: iluise <[email protected]> Co-authored-by: Sebastian Hickman <[email protected]> * 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 <[email protected]> Co-authored-by: Simon Grasse <[email protected]> Co-authored-by: lessig2 <[email protected]> Co-authored-by: Savvas Melidonis <[email protected]> Co-authored-by: Julius Polz <[email protected]> Co-authored-by: Sebastian Hickman <[email protected]>
* 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 <[email protected]>
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 <[email protected]>
* Implement regular subsampling * ruff * Fix handling of holes in dataset * Update stream config * Remove 0 as filler value --------- Co-authored-by: Savvas Melidonis <[email protected]>
* Replace cf.rank==0 with utils.distributed.is_root * Restore streams_directory replacement semantics --------- Co-authored-by: wang85 <[email protected]> Co-authored-by: wang85 <[email protected]>
Co-authored-by: kasravi1 <[email protected]> Co-authored-by: Christian Lessig <[email protected]>
* first implementatio regrid in eval * lint * workign version * remove unused functions * implement comments
* remove redundant plots * kint
* upgrade pytorch and flashattn wip * update download paths * Add Macos support * remove macos specific support * downgrade to cu126 again * Specify build dependency to avoid dependency issues on specific wheels * Revert "downgrade to cu126 again" This reverts commit 2334509. * Reset by hand * Trying pinned ABIfalse for x_86
* Implement dynamic loss * Linter * Move EMA to a separate class * Review * Linter
* feat: add timing metrics (startup, training, overall) - Track startup_time_seconds: time from run() start to training loop - Track total_training_time_seconds: time in training/validation cycles - Track overall_time_seconds: total wall-clock time from launch to finish - All metrics logged only on root rank to avoid file contention - Metrics written to metrics.json, automatically uploaded to MLflow - Console logs show timing summaries for quick monitoring * docs: add agent structure with skills, tasks, and docs - Created .hermes/ directory with skills/, tasks/, docs/ subfolders - Added skills overview (README.md) with task-type skills - Implemented 'planning' and 'metrics' skills - Documented timing metrics task in tasks/2026-04-17-timing-metrics/ - Added agent structure documentation - Updated .gitignore with optional .hermes/ entry * docs: add skills review cycle for periodic compactification - Added 2-3 month review cycle recommendation - Defined criteria for skill consolidation - Included usage frequency thresholds - Documented when to merge or remove skills * configs * Remove hermes tool tracking for now * Try duration metrics * Update metrics, store after each mini-epoch * Refactor configs/streams * Extract scaling data * Script to generate scaling plots * Script update * Repeat data in mini epoch * corrected time window length * Lower to 512 samples per mini epoch * Updated extraction script * Log time more often * Fix training start scope * Minimal validation * Increase samples_per_mini_epoch to 1024 * Final training duration and terminal/metric logging * log metrics after mini-epoch * Log metrics after mini-epoch, change schema * MEtric typo * Logging refactor * Update extraction script * NNode extraction * Logs path * Wait until all training complete and wait with validation until logs are done * Log seconds rather than hours * Measure dataset advancement time * LR scheduler lower bounds * At least two warmup steps * Len per rank at least 1 to avoid zero division error * Write csv for easier viewing * Extraction and plotting * Remove parent dir creation * more detailed extraction script * Remove overall time logging * Cleanup trainer * Metrics extraction and plot generation scripts * Add efficiency factor in plot * RM checkpoint and log metrics at last iteration * Detailed metrics * Remove barrier and extra logging on last batch * trainer code cleanup Co-authored-by: Copilot <[email protected]> * Lower bound beta2 in adam Co-authored-by: Copilot <[email protected]> * update script for scaling plots, loss as separate entry point * specify nodes in scaling data script * Update extract scaling data * Add pyarrow * Update script for scaling plots * Update to generating scaling plots * Move scaling scripts to package * init refactor * Setup and linting * Updated plot generation script * Update readme * Fewer diffs * no gitignore changes * Refactor logging and move time for mini epoch logging outside loop * Formatting and style fixes * Update config * Avoid duplicate metrics * Fix lint issues * t_training in __init__ * Renamed metric * mv performance package * Plot losses against elapsed training time via --x-axis flag - Replace --x_type ('step'/'reltime') with --x-axis column selector ('samples', 'elapsed_training_time') - Add x_axis param to plot_loss_avg (previously hardcoded num_samples) - Add friendly x-axis labels: 'elapsed training time [s]' when plotting against elapsed_training_time_seconds - plot_lr, plot_loss_avg, plot_loss_per_stream, plot_loss_per_run all now respect x_axis; xlabel is auto-derived from column name - Remove dead x_type parameter from plot_loss_per_stream * Fewer changes * rm configs * Remove startup time * Remove startup time * Formatting and removed time per epoch * Undo pyproject change * Undo pyproject changes --------- Co-authored-by: Copilot <[email protected]> Co-authored-by: Christian Lessig <[email protected]>
* NoOpGradScaler * NoOpGradScaler in trainer * Keep grad scaling by default for now * Linter fixes * False by default afteral * Revert "False by default afteral" This reverts commit 42d7d08. * Move grad_scaling totraining_config.optimizer in config --------- Co-authored-by: Christian Lessig <[email protected]>
* fix * activate * activate? * debug * link mode * copy * copy * disable cache * uv cache dir * trying other options * uv * change * fixues * more debug info * more debug info * link mode copy * fixes for anemoi * checks * check * changes * debug * debug * debug * debug * try * debug * debug * debug * debug * debug * debug * debug * try * trigger run * changes * cleanup * cleanup * change * change the uv cache in CI
* call BatchSample method instead of repeating content * improve object-oriented design of batch architecture * add explicit lenght checks for num source/target steps getter
* Add custom cmap support * Add glob support * Update range calculation with glob * Add colorbar scale options * Refactor _resolve_norm() * Add example for tp and log colorbars
* Add configurable coords column names to obs reader * Use count() to validate coordinate columns * Switch from raise to assert --------- Co-authored-by: Christian Lessig <[email protected]>
* reinstating necessary functions * linting
* Fix ema update with forecast offset * Remove commented code * Review --------- Co-authored-by: Christian Lessig <[email protected]>
* Add timeseries class, edit plot_orchestration * add timeseries.py * Take care of ensembles * Changes to reply to reviews. Remove functions for retrieving channels and samples * Include also timeseries per region. Change titles in timeseries to regions * Chnages to avoid Future Warnings related to concat * adapt configs * Linting
* add max num target warning * fix surfacecombined multi-rank * fix warning * fix warning * fix warning * address comment --------- Co-authored-by: Christian Lessig <[email protected]>
* better messages * revert matrix cashing * regridding class * generalise coordinates * fix logging * remove regrid opts * rename build+io_states
* Add Fast Evaluation Config Reference documentation This document serves as a comprehensive reference for all configuration options used in the command `uv run evaluate --config <your_config.yml>`, including templates and detailed descriptions of each configuration key. * Revise evaluate_config_reference.md addressing comments by iluise and clessig
add missing variables to export config to enable export of N320 data.
* score bug fixes * Align score-map ensemble handling with adjusted-scores branch Use the robust ens-detection variant in _plot_score_maps_per_stream so the block is identical (apart from the branch-specific tag string) to mh/full-pipeline-diffusion-adjusted-scores, minimising future merge conflicts: - restore ens labels via assign_coords(ens=preds.ens.values) (positional) instead of plot_metrics["ens"] = preds.ens (index-aligned) - gate has-ens detection on the ens *dimension* (all_ens) rather than the coordinate - compute per-metric ens iteration via metric_has_ens / ens_values Behaviour is unchanged. Co-Authored-By: Claude Opus 4.8 <[email protected]> * renames --------- Co-authored-by: Claude Opus 4.8 <[email protected]>
* score bug fixes * Align score-map ensemble handling with adjusted-scores branch Use the robust ens-detection variant in _plot_score_maps_per_stream so the block is identical (apart from the branch-specific tag string) to mh/full-pipeline-diffusion-adjusted-scores, minimising future merge conflicts: - restore ens labels via assign_coords(ens=preds.ens.values) (positional) instead of plot_metrics["ens"] = preds.ens (index-aligned) - gate has-ens detection on the ens *dimension* (all_ens) rather than the coordinate - compute per-metric ens iteration via metric_has_ens / ens_values Behaviour is unchanged. Co-Authored-By: Claude Opus 4.8 <[email protected]> * implemented spread adj and ssr adj * add latitude weighting * call calc_latitude_weights on p in ssr functions * Refactor latitude weight calculation to use static method and improve coordinate handling * move calc_latitude_weights to score utils * wire lat weights to config * add optional lat weighting to more scores * refactor: introduce _weighted_mean helper * unify ssr and ssr_adj via param switch * rename use_latitude_weights to latitude_weighting * move lat weights to verified data * add warning if no latitude was found * make adjusted ssr the default * lint * renames --------- Co-authored-by: Claude Opus 4.8 <[email protected]> Co-authored-by: jpolz <[email protected]> Co-authored-by: iluise <[email protected]>
adapted lr logging such that it logs lr for each of the optimizer used
…rameters for Muon
…b-temporal-plus-plus
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
Add Muon
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