Feature/visualize masking by wael-mika · Pull Request #2146 · ecmwf/WeatherGenerator

wael-mika · 2026-03-31T09:00:12Z

Description

Implementation of the masking visualization script

Issue Number

Closes #2145

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

…ntly fixing the white error on the maps

tjhunter · 2026-03-31T10:16:09Z

+This script loads a config, builds a MultiStreamDataSampler, extracts one batch,
+and plots a single variable for source and target with masking/cropping applied.
+This script can run on a cpu or logging node without GPUs.
+Please activate your .venv before running.


this is wrong. You should not have to do that. uv run should handle running without activating the environment.

I agree with you, but after dealing with different systems, there are different Python versions running on each of those, and I found out that using uv run is the easiest way to handle this situation. I can add a command using python as well if you agree

TillHae

Looks good, I just have a few concerns with your script.

I didn't check parsing, docstrings, plotting, meta info, instance checks and didn't test the script.

TillHae · 2026-03-31T11:20:44Z

+
+def _to_numpy(arr):
+    if isinstance(arr, torch.Tensor):
+        return arr.detach().cpu().numpy()


Why do we have to move this to cpu?

This script can be run completely on cpu, we do not need GPU for it

But if I have a gpu on my local machine I might want to use it for performance

I believe most of users will be using HPCs as the data is stored there. That said, it is only retrieving one sample so cpu is perfectly suited even on any machine. The performance is not important here as the script is used to generate one plot before training.

TillHae · 2026-03-31T11:30:41Z

+    np.ndarray
+        Boolean array of same length as lats/lons indicating visibility.
+    """
+    if len(lats) == 0:


Is len(lats) == len(lons) all the time?

Yes, always. Both lats and lons are derived from the same coords array via coords[:, 0] and coords[:, 1], so they are guaranteed to have equal length at every call site.

TillHae · 2026-03-31T11:56:37Z

+    if max_points is None or len(vals) <= max_points:
+        return lats, lons, vals
+    idxs = rng.choice(len(vals), size=max_points, replace=False)
+    return lats[idxs], lons[idxs], vals[idxs]


This might be dangerous to do as it generates a discontinuous array that looks continuous

Disagree. For scatter plots. Unlike line plots, scatter does not connect adjacent points. Each point is rendered independently, so random downsampling only reduces density. It doesn't create false continuity. The points=N label already shows how many points are displayed.
_downsample is only used when you set up "--max-points" for high resolution data and it is not used in individual cases if not specificaly sat up

TillHae · 2026-03-31T12:09:15Z

+        if len(vals_ref) == 0:
+            vals_ref = vals_src if len(vals_src) else vals_tgt
+        vmin = np.nanpercentile(vals_ref, 2)
+        vmax = np.nanpercentile(vals_ref, 98)


It doesn't make a big difference in most cases, but I think cutting the outliers can be unwanted sometimes, where we are interested in them and want to see it in our plot. Especially when evaluating masking this can be interesting.

I agree, but this is not an evaluation plot, so extreme events are prevented here to keep the colorscale consistent.

TillHae · 2026-03-31T12:19:47Z

+            )
+        else:
+            target_var_idx, target_var_name = _resolve_var_idx(
+                ds.source_channels,


Why do we have to use ds.source_channels when we have no target_tokens in our target stream?

When visualizing masking, I am choosing only one variable, which is enough for this purpose. Source and target here represent the student and teacher view which are different and can be set up in the config

TillHae · 2026-03-31T12:24:03Z

+
+        pairs_data.append(
+            {
+                "pair_idx": source_idx,


Shouldn't the key "pair_idx" point to the value of pair_indices defined in l.829

source_idx already is a value from pair_indices (the loop iterates for source_idx in pair_indices). So "pair_idx": source_idx stores the source sample index, which is meaningful (it identifies which batch sample this is). Storing the enumeration position (0, 1, 2…) would be less informative and would make the filename pair0_, pair1_ regardless of which samples were selected.
Though the naming pair_idx vs source_idx is admittedly redundant — pair_idx could be renamed to make it clearer.

TillHae · 2026-03-31T18:36:28Z

Thanks for clarifying. Now everything sounds logical to me.

wael-mika added 6 commits January 30, 2026 12:28

first draft of masking visualizer script

6146053

plotting improvement with shared colorbar, docstring, and point size

7b3f266

Merge branch 'ecmwf:develop' into feature/visualize-masking

f9023a6

Merge branch 'ecmwf:develop' into feature/visualize-masking

3ad7ff6

Clearing the wasted white space, fixing naming issue and most importa…

0ad0b64

…ntly fixing the white error on the maps

linting

4c8a144

github-project-automation bot added this to WeatherGen-dev Mar 31, 2026

github-actions bot added data Anything related to the datasets used in the project eval anything related to the model evaluation pipeline labels Mar 31, 2026

tjhunter reviewed Mar 31, 2026

View reviewed changes

TillHae reviewed Mar 31, 2026

View reviewed changes

Conversation

wael-mika commented Mar 31, 2026

Description

Issue Number

Checklist before asking for review

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

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!

TillHae commented Mar 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

TillHae commented Mar 31, 2026 •

edited

Loading