Cropping overlap control using cone distance #1554

wael-mika · 2026-01-06T15:19:44Z

Description

This PR implements cone distance for spatial sampling. The core logic uses angular distance (in degrees/radii) to determine the spacing between the centers of the Teacher and Student cones/crops.

Key Features:
Geodesic Support: Currently optimized to work specifically with geodesic_disk.
Multi-Level Sampling: Allows Student and Teacher crops to be sampled from different hl_level masks.
Performance Optimizations: I’ve implemented caching for the Healpix object initialization to improve efficiency.
Geometric Logic: Uses spherical trigonometry to calculate destination points given a starting (lon, lat), distance, and azimuth (bearing).

Configuration:
Usage Example: Refer to default_config for the necessary parameters and an explanation of how center_distance_degrees is used for overlap control.
Documentation: I have included long/detailed docstrings for all new functions to make the review process easier.

Issue Number

Closes #1515

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

…ve to extra readers (ecmwf#1231) * initial commit for ICON data reader - new version * merging ICON Zarr and ICON CMIP6 NetCDF4 data readers * fixed inconsistencies + added descriptive comments * removed unecessary attribute * added attribute cf.multiprocessing_method * changed self.colnames self.cols_idx from parent to children classes * config file for CMIP6 * ICON config * code ruffing * fixed lines that were too long * temporary fix for multiprocedding method * unecessary code removed and added more channels * title: changes to dataloading method description: - fixed datetime, lat, lon reglated bugs - added more memory efficient way of loading channels * change in len_hrs and step_hrs * config for CMIP6 monthly data * title: handling built-in ICON levels for CMIP6 data (details below) description: - added method select_by_level() to construct selected colnames using pre-selected levels in data config file - changes in _get() method to handle built-in plevs from CMIP6 ICON - added method get_cols() to construct general colnames - clipping lat/lon - changed data config attributes from source_channels/target_channels to source/target - fixed bug in asserting non positive stds * removed breaking change * committing ruffing changes * added ICON CMIP6 monthly data to the stream * removed the previous stream * changed per-level reading for CMIP6 only * dtype missmatch fix - training on ICON CMIP6 day + Amon channels * changed len_hrs value to accommodate monthly data * check in merged reader for art * remove cmip * revert icon cmip specific changes * ruff * speed up loading 500x by selecting time slice first * correct type * update config * linting * ruff * current working loader param setup * reset to standard setting * current settings * adjustments to multistream data sampler to be tested * reset to original since it works * update registry * current config * remove legacy code and local files * reset default config * adjust stream config * remove icon data reader from main codebase * move scheduler setting inside datareader --------- Co-authored-by: sbAsma <[email protected]> Co-authored-by: Timothy Hunter <[email protected]>

* Update to develop, prepare for new experiment series * Setting o48 as default in era5 config Committer: Matthias Karlbauer <[email protected]> On branch mk/develop/fe_experiments Your branch is ahead of 'origin/mk/develop/fe_experiments' by 57 commits. (use "git push" to publish your local commits) Changes to be committed: modified: config/streams/era5_1deg/era5.yml * Updated default config to 256 dim latent size On branch mk/develop/fe_experiments Your branch is ahead of 'origin/mk/develop/fe_experiments' by 58 commits. (use "git push" to publish your local commits) Changes to be committed: modified: config/default_config.yml * Update branch to latest develop * Change epochs from 64 to 32 * ICON ESM (CMIP6) data reader * hooking-up the ICON ESM data reader * ICON ESM historical data config * added config multiprocessing_method as param * ICON ESM custom config params * fixed multiprocessing method needed for ICON ESM * fixed data type + cleaned code * added param streams_output * changed stream config names * ruffing changes * restored era5 file * removed experiment specific config * removed unecessary file * restored default config * restore era5 file * removed wrong eval config * another attempt to fix era5 filename * ruff requested fixes * Add configurable multiprocessing method with fork as default * Another add configurable multiprocessing method with fork as default * fixing forgotten conflicts in default_config --------- Co-authored-by: Matthias Karlbauer <[email protected]> Co-authored-by: Matthias Karlbauer <[email protected]>

* added integ test for jepa * added weight > 0.0 checks * lint * added jepa test command * i_obs -> i_stream * added integration-test-all * using .get()

…tion

* modified ratio plot * heat-map * fix cosmetics * lint * change config * add metric name to heatmap * multiply forecasts by step_hrs * remove breakpoints * fix variable order * lint * implement lead time plotting * lint * implement comments

…wf#1551)

* Revert "Make missing fstep 0 produce a more meaningfull error (ecmwf#1444)" This reverts commit e3b29be. * add function to truncate runs * lint * add warning

* mae loss function * implement generalized lp-norm * docstring updates

shmh40 · 2026-01-13T10:25:19Z

config/default_config.yml

+        hl_mask: 3                     # Can be different from teacher
+        rate: 0.4                      # Student cone: 40% of sphere (~72° radius)
+        method: "geodesic_disk"        # Must be geodesic_disk
+        center_distance_degrees: 180    # REQUIRED: Student center is 45° from teacher


This comment is a little unclear, the center_distance_degrees is required, but the center_distance_degrees can be any chosen angle right?

Perhaps you could also include a range as below

Indeed, the comment is confusing as I had to set center_distance_degrees to 45 before submitting the PR ( I was trying different values to see if they work)
Any angle you choose should work as demonstrated in the example below.

shmh40 · 2026-01-13T10:26:18Z

src/weathergen/datasets/masking.py

        # number of healpix cells
        self.healpix_level_data = cf.healpix_level
        self.healpix_num_cells = 12 * (4**cf.healpix_level)
+        self._hp_cache = {}


Good, I think this is helpful

True, it helped speeding up the computation (we do not have to repeat hp calculation every single time)

shmh40 · 2026-01-13T10:28:08Z

src/weathergen/datasets/masking.py

            instance default if None.
        masking_strategy_config : dict | None
            Optional override of strategy config (e.g., {'hl_mask': 3}).
+        target_metadata : dict | None


Perhaps state here that this is required when we do overlap with the cones? Otherwise it is unused right?

Agree, this is only used when doing cone overlap. I am not sure if there is a more elegant way to do it without adding target_metadata

shmh40 · 2026-01-13T12:13:37Z

src/weathergen/datasets/masking.py

+                "relationship 'cone_distance' requires 'center_cell' in target_metadata"
+            )
+            # Get teacher's hl_mask level to properly convert coordinates
+            teacher_hl_mask = target_metadata.get("hl_mask", masking_strategy_config.get("hl_mask", 0))


Should this default to 0? Perhaps should check that hl_mask is specified with an assert like the above lines? Do we default to 0 in other parts of the code?

wael-mika and others added 8 commits December 25, 2025 00:35

First prototype of cone and deterministic overlp control, TO BE TESTED

67c0568

more doc (ecmwf#1477)

00ac394

Improve argument check readability of run_train.py (ecmwf#1485)

287d819

Adding an integration test for JEPA (ecmwf#1519)

08c9d90

* added integ test for jepa * added weight > 0.0 checks * lint * added jepa test command * i_obs -> i_stream * added integration-test-all * using .get()

Implementing the cone_distance overlap control with mathematical nota…

67784d3

…tion

Finalize the config for the cone overlap control

5826b45

wael-mika added science Scientific questions data:processing model:pretrain labels Jan 6, 2026

github-project-automation bot added this to WeatherGen-dev Jan 6, 2026

iluise and others added 5 commits January 7, 2026 13:59

rename mse_channel_location_weighted to mse and extend docstring (ecm…

54d49e8

…wf#1551)

truncate filename in plotting tool (ecmwf#1556)

5608a92

* Revert "Make missing fstep 0 produce a more meaningfull error (ecmwf#1444)" This reverts commit e3b29be. * add function to truncate runs * lint * add warning

Implement MAE loss function (ecmwf#1334)

0d0e6a5

* mae loss function * implement generalized lp-norm * docstring updates

Merge branch 'ecmwf:develop' into wm/dev/cropping_overlap

4fb73fd

shmh40 reviewed Jan 13, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Cropping overlap control using cone distance #1554

Cropping overlap control using cone distance #1554

wael-mika commented Jan 6, 2026

Uh oh!

shmh40 Jan 13, 2026

Uh oh!

shmh40 Jan 13, 2026

Uh oh!

wael-mika Jan 13, 2026

Uh oh!

shmh40 Jan 13, 2026

Uh oh!

wael-mika Jan 13, 2026

Uh oh!

shmh40 Jan 13, 2026

Uh oh!

wael-mika Jan 13, 2026

Uh oh!

shmh40 Jan 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants

Cropping overlap control using cone distance #1554

Are you sure you want to change the base?

Cropping overlap control using cone distance #1554

Conversation

wael-mika commented Jan 6, 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!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants