feat: add enhanced probabilistic forecasting visualizations - #612
feat: add enhanced probabilistic forecasting visualizations#612Sir-Sloth-The-Lazy wants to merge 29 commits into
Conversation
- Update test_datasets.py to use ForecasterModule instead of GraphLAM - Update test_plotting.py to use ForecasterModule instead of GraphLAM - Fix interior_mask_bool property shape (1,) -> (N,) for correct loss masking - Fix all_gather_cat to handle single-device runs without incorrect dim collapse
…r hierarchy - Replace opaque argparse.Namespace with explicit keyword arguments in StepPredictor, BaseGraphModel, BaseHiGraphModel, GraphLAM, HiLAM, and HiLAMParallel __init__ methods - Reorder methods in step_predictor.py: forward/expand_to_batch now appear before clamping methods - Update all instantiation sites (train_model.py, test_training.py, test_prediction_model_classes.py) to pass explicit kwargs - HiLAM helper methods (make_same/up/down_gnns) now use self.hidden_dim and self.hidden_layers instead of args parameter Addresses review comments on PR mllam#208.
- Rename border to boundary in Forecaster - Pass Forecaster object to ForecasterModule init instead of Predictor - Remove inline imports in ForecasterModule - Move loss-related pred_std logic fully into ForecasterModule - Delete obsolete test_refactored_hierarchy.py
Co-authored-by: Joel Oskarsson <joel.oskarsson@outlook.com>
- Add predicts_std property to StepPredictor, Forecaster and ARForecaster so ForecasterModule can query the forecaster instead of taking output_std as a separate constructor argument - Remove output_std parameter from ForecasterModule; use self._forecaster.predicts_std throughout - Move fallback per_var_std logic out of forecast_for_batch into each step method so pred_std is None before fallback, enabling direct None checks instead of hparam checks - Replace len(datastore.boundary_mask) with datastore.num_grid_points in StepPredictor to avoid relying on boundary_mask - Move get_state_feature_weighting and ARForecaster inline imports to module-level imports in forecaster_module.py and train_model.py - Fix statement ordering in StepPredictor.__init__ so register_buffer for grid_static_features appears directly after building the tensor - Replace dict+loop pattern for registering state_mean/state_std buffers with two direct register_buffer calls - Remove all internal Item N checklist references from comments - Remove TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD env var hack; pass weights_only=False explicitly to load_from_checkpoint calls and weights_only=True to torch.load in test_graph_creation.py - Add test_step_predictor_no_static_features to verify models initialise and run correctly when the datastore returns None for static features - Fix graph= -> graph_name= and model.forecaster -> model._forecaster in tests to match current API
…r_batch Makes the forecasting path tolerant to batch-folded execution so that future ensemble generation can fold (S, B) into (S*B) before calling ARForecaster, without any changes to ARForecaster or StepPredictor. Prediction is kept folded through the existing deterministic logging and aggregation paths so all dim assumptions in training_step, validation_step, and test_step remain correct. Unfolding to (*leading, T, N, F) is deferred to ensemble-specific subclasses (e.g. EnsForecasterModule). Adds test_fold_unfold_equivalence to confirm ARForecaster's rollout is rank-transparent under a pre-entry fold.
…ecast_for_batch" This reverts commit e659f7a.
Remove the thin forecast_for_batch wrapper, inlining the batch unpacking and forecaster call into each step method. Also fix _forecaster -> forecaster attribute references in checkpoint loading and tests to match the actual attribute name.
…sue-49 Resolved conflicts: - Removed ar_model.py (replaced by ForecasterModule/ARForecaster hierarchy) - Updated test_training.py: use ForecasterModule imports, dynamic device allocation, and wandb.init(mode="disabled") - Updated test_plotting.py: replace ARModel.all_gather_cat with ForecasterModule.all_gather_cat, fix model.args.create_gif reference Ported features from ar_model.py: - GIF export (create_gif param + plot_examples logic) into ForecasterModule - Added common_step to ForecasterModule; refactored training/validation/ test steps to delegate to it, eliminating duplicated forecaster calls
Ports ensemble/latent plotting functions from prob_model_lam into the new ForecasterModule hierarchy and adds a suite of probabilistic forecast-verification diagnostics described in docs/visualization_proposal.md. Fixes to existing plots: - plot_error_map_absolute: unnormalized companion heatmap with a shared colour scale so errors can be compared across variables - get_var_cmap helper: picks RdBu_r for anomaly/error fields, YlOrRd for uncertainty/spread, viridis for raw physical fields - plot_ensemble_prediction std panel uses YlOrRd New plotting functions: - plot_spread_skill: ensemble spread vs RMSE with 1:1 diagonal - plot_rank_histogram: Talagrand diagram with flat reference - plot_reliability_diagram: calibration curve vs 45 degree diagonal - plot_crps_leadtime: CRPS vs lead-time line plot - plot_spaghetti: contour-bundle for ensemble spatial coherence - plot_plume: per-location fan diagram with 10/25/50/75/90 pct. bands - plot_exceedance_prob: probability-of-exceedance spatial map - save_forecast_animation: GIF from a list of figures - plot_power_spectrum: radially-averaged 2D FFT comparison - plot_latent_pca: NumPy-SVD PCA scatter of prior vs VI samples Ported from prob_model_lam: - plot_ensemble_prediction adapted to datastore + pcolormesh - plot_latent_samples Wiring: - ForecasterModule.create_metric_log_dict now logs the absolute error map alongside the normalized one, and adds a CRPS-vs-lead-time line plot whenever the metric name contains "crps".
Wrap three overlong docstring/comment lines and remove dead thresholds_np assignment left over from reliability-diagram fix.
|
@joeloskarsson @sadamov looking forward for your dig on this PR whenever you have time. ;D |
|
Nice work on this ! The
The current implementation computes a single global_max across all variables: If one variable dominates (e.g. geopotential height having much larger absolute errors than 2m temperature), the scatter for smaller-error variables gets compressed into a corner. Per-subplot scaling (
The auto-detection in is nice forward-looking design for when CRPS gets added to the metrics module. |
|
With more things merged, could you get this up to date with main so it is easier to discuss the diff :) |
| return s | ||
|
|
||
|
|
||
| def get_var_cmap(var_name: str) -> str: |
There was a problem hiding this comment.
Is this used? Don't fully get when it would be used. If I know I am plotting an error, I could just as well hard-code the colormap as call this with a hard-coded string? It would however be nice to use different cmaps different colors for different atmospheric variables, but this does not seem to be the case now?
| else: | ||
| std_vmin, std_vmax = vrange_std | ||
|
|
||
| n_members = min(samples_np.shape[0], 6) |
There was a problem hiding this comment.
Are we sure we want to limit this to 6? I guess this will need to be configured somewhere, but could be left as responsibility of caller (through number of entries in samples).
| img_side_size = int(np.sqrt(num_mesh_nodes)) | ||
| assert img_side_size**2 == num_mesh_nodes, ( | ||
| "Number of mesh nodes is not a square number, " | ||
| "cannot plot latent samples as images" | ||
| ) |
There was a problem hiding this comment.
I think we should figure out a better way to do these plots, rather than using imshow for a square. In a model like Graph-EFM, the latent vectors actually have a spatial position associated with them. A better approach here is to actually use those positions. Maybe an option would be to give this method also the coordinates associated with the N_mesh dimension?
|
|
||
|
|
||
| @matplotlib.rc_context(utils.fractional_plot_bundle(1)) | ||
| def plot_spread_skill(spread, rmse, var_names, title=None): |
There was a problem hiding this comment.
Is this intended for one forecast? Or the average spread and skill across the entire dataset? How do we show which lead time each point corresponds to?
|
It would here be very useful if you would put in here also one example plot from each new method, as we then make sure we reason about the same visual throughout the review of this :) |
|
While I am a big fan of verification&vizualization, and thanks for this PR @Sir-Sloth-The-Lazy, we also need to be aware how costly probabilistic plots can be. There is a real risk that validation steps take longer for plotting than actual inference. So setting sensible defaults and also moving some plots to post-training evaluation with a separate library, might be reasonable here. Just wanted to drop this note and maybe ask for some performance benchmarks once the details are settled. |
|
@sadamov , Thank you for the review. I would be reading more about how to make plotting efficient, duely noted :). @joeloskarsson I would be now attending to this PR and completing this ! I will need sometime to attend to all of your concerns. Sorry for the delay 😅 |
|
This is a very good point @sadamov, we should probably not implement a bunch of ensemble-related plots just because we can. This is something that would be good to get more community input on, to see what people want/need. @Sir-Sloth-The-Lazy a suggestion would be to 1. Mark in your list which plots are mainly for validation (during training) and which are for evaluation runs from a saved checkpoint. 2. Also add a figure of each plot, as the look now. 3. Either post this here, and direct people from slack to have a look, or post this directly as a thread on slack (in #neural-lam). |
|
@joeloskarsson I have started the work keeping your and @sadamov guidance in mind :) I would pushing the changes and updation soon. |
4 participants
Describe your changes
Ports the ensemble / latent-space plotting functions from the
prob_model_lambranch into the refactoredForecasterModulehierarchy and adds a suite of probabilistic forecast-verification diagnostics.To view only this PR's additions, compare against the post-refactor branch:
👉 View the clean diff on my fork
The goal is to land the visualization layer ahead of the full Graph-EFM port so that when probabilistic models arrive they already have the calibration, reliability, and uncertainty-quality plots that ensemble forecasting requires. Every function is self-contained and accepts pre-computed tensors/arrays, so model wiring can be added incrementally.
Fixes to existing plots
plot_error_map_absolute— unnormalized companion heatmap with a shared colour scale so errors can be compared directly across variablesget_var_cmap(var_name)helper — routes anomaly/error fields toRdBu_r, uncertainty/spread fields toYlOrRd, and raw physical fields toviridisplot_ensemble_predictionusesYlOrRd(previouslyplasma)New plotting functions
plot_spread_skillplot_rank_histogramplot_reliability_diagramplot_crps_leadtimeplot_spaghettiplot_plumeplot_exceedance_probsave_forecast_animationplot_power_spectrumplot_latent_pcaPorted from
prob_model_lamplot_ensemble_prediction— adapted to usedatastore+pcolormeshviaplot_on_axis(no moreimshow+GRID_LIMITS)plot_latent_samples— ported largely unchanged, with numpy-safe tensor handlingWiring into
ForecasterModulecreate_metric_log_dictnow logs the absolute error map alongside the normalized one for every metric"crps", the lead-time line plot is logged tooDependencies: no new dependencies. Uses only existing stack (cartopy, matplotlib, numpy, torch, xarray, Pillow).
Issue Link
prob_model_lambranch #62 — ports theprob_model_lam"Add plotting functions for ensemble forecasts" item and goes further with calibration/reliability diagnostics that were not present in either branch.ForecasterModule) with visualization hooks that are ready for probabilistic/ensembleForecastersubclasses.Type of change
Checklist before requesting a review
pullwith--rebaseoption if possible).Checklist for reviewers
Each PR comes with its own improvements and flaws. The reviewer should check the following:
Author checklist after completed review
reflecting type of change (add section where missing):
Checklist for assignee