Skip to content

fix(graphs): undo masking when torch-cluster is installed #375

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 43 commits into from
Jul 29, 2025
Merged

fix(graphs): undo masking when torch-cluster is installed #375

merged 43 commits into from
Jul 29, 2025

Conversation

@JPXKQX
Copy link
Member

@JPXKQX JPXKQX commented Jun 19, 2025 •
edited by github-actions bot
Loading

Description

The current workflow did not function as intended with torch-cluster when a source or target mask was specified. The torch-cluster library is used only when graph creation needs to be optimized. It is not listed as a dependency because it is difficult to install in the CI environment, and as a result, it is only lightly tested.

What problem does this change solve?

This PR fixes the previously reported bug and introduces a new base class shared by both KNNEdges and CutOffEdges to handle masking.

What issue or task does this change relate to?

Additional notes

As a contributor to the Anemoi framework, please ensure that your changes include unit tests, updates to any affected dependencies and documentation, and have been tested in a parallel setting (i.e., with multiple GPUs). As a reviewer, you are also responsible for verifying these aspects and requesting changes if they are not adequately addressed. For guidelines about those please refer to https://anemoi.readthedocs.io/en/latest/

📚 Documentation preview 📚: https://anemoi-training--375.org.readthedocs.build/en/375/

📚 Documentation preview 📚: https://anemoi-graphs--375.org.readthedocs.build/en/375/

📚 Documentation preview 📚: https://anemoi-models--375.org.readthedocs.build/en/375/

@JPXKQX JPXKQX self-assigned this Jun 19, 2025
@JPXKQX JPXKQX added bug Something isn't working graphs ATS Approval Not Needed No approval needed by ATS labels Jun 19, 2025
@github-project-automation github-project-automation bot moved this to Now In Progress in Anemoi-dev Jun 25, 2025
anaprietonem
anaprietonem reviewed Jul 8, 2025
View reviewed changes
anaprietonem
anaprietonem reviewed Jul 8, 2025
View reviewed changes
anaprietonem
anaprietonem reviewed Jul 8, 2025
View reviewed changes
anaprietonem
anaprietonem reviewed Jul 8, 2025
View reviewed changes
Copy link
Contributor

@anaprietonem anaprietonem left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Some minor comments and suggestion for docs, otherwise good to go

@JPXKQX JPXKQX requested a review from dietervdb-meteo July 11, 2025 11:28
JPXKQX and others added 11 commits July 17, 2025 11:28
dietervdb-meteo
dietervdb-meteo reviewed Jul 25, 2025
View reviewed changes
dietervdb-meteo
dietervdb-meteo reviewed Jul 25, 2025
View reviewed changes
dietervdb-meteo
dietervdb-meteo previously approved these changes Jul 25, 2025
View reviewed changes
Copy link
Contributor

@dietervdb-meteo dietervdb-meteo left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tested with and without torch-cluster, for various graph configs, including the new reversed classes. Once my suggestions above are taken into account this should be good to go.

@anaprietonem anaprietonem self-requested a review July 29, 2025 11:38
anaprietonem
anaprietonem approved these changes Jul 29, 2025
View reviewed changes
Copy link
Contributor

@anaprietonem anaprietonem left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @JPXKQX for addressing the comments and @dietervdb-meteo for the review, good to be merged!

@JPXKQX JPXKQX merged commit 9f75c06 into main Jul 29, 2025
22 checks passed
@JPXKQX JPXKQX deleted the fix/undo-masking-torchcluster branch July 29, 2025 11:39
@github-project-automation github-project-automation bot moved this from Now In Progress to Done in Anemoi-dev Jul 29, 2025
@DeployDuck DeployDuck mentioned this pull request Jul 28, 2025
Merged
anaprietonem pushed a commit that referenced this pull request Aug 4, 2025
@DeployDuck
chore: Release main (#380)
9e7f90d 
🤖 Automated Release PR

This PR was created by `release-please` to prepare the next release.
Once merged:

1. A new version tag will be created
2. A GitHub release will be published
3. The changelog will be updated

Changes to be included in the next release:
---


<details><summary>training: 0.6.0</summary>

##
[0.6.0](training-0.5.1...training-0.6.0)
(2025-08-01)


### ⚠ BREAKING CHANGES

* for schemas of data processors
([#433](#433))
* BaseGraphModule and tasks introduced in anemoi-core
([#399](#399))

### Features

* Add metadata back to pl checkpoint.
([#303](#303))
([0193b28](0193b28))
* BaseGraphModule and tasks introduced in anemoi-core
([#399](#399))
([f8ab962](f8ab962))
* **deps:** Use mlflow-skinny instead of mlflow
([#418](#418))
([6a8beb3](6a8beb3))
* Log FTT2 loss + Fourier Correlation loss
([#148](#148))
([345b0ab](345b0ab))
* **model:** Postprocessors for leaky boundings
([#315](#315))
([b54562b](b54562b))
* **models:** Checkpointed Mapper Chunking
([#406](#406))
([8577772](8577772))
* **models:** Mapper edge sharding
([#366](#366))
([326751d](326751d))
* Variable filtering
([#208](#208))
([fba5e47](fba5e47))


### Bug Fixes

* Dropping 3.9 ([#436](#436))
([f6c0214](f6c0214))
* For schemas of data processors
([#433](#433))
([539939b](539939b))
* Mlflow hp params limit
([#424](#424))
([138bc3a](138bc3a))
* Mlflowlogger duplicated key
([#414](#414))
([cb64a1c](cb64a1c))
* **models,traininig:** Hierarchical model + integration test
([#400](#400))
([71dfd89](71dfd89))
* **models:** Add removed sharded_input_key in PR400
([#425](#425))
([089fe6f](089fe6f))
* New checkpoint
([#445](#445))
([a25df93](a25df93))
* Plotting error when precip related params are not diagnostic
([#369](#369))
([010cfa3](010cfa3))
* **training:** Address issues with
[#208](#208)
([#417](#417))
([665f462](665f462))
* **training:** Scaler memory usage
([#391](#391))
([a9d30e1](a9d30e1))
* Update import mflow utils unit tests
([#427](#427))
([70ecdd9](70ecdd9))
* Update level retrieval logic
([#405](#405))
([f393bc3](f393bc3))
* Use transforms: Variable for ExtractVariableGroupAndLevel
([#321](#321))
([7649f4f](7649f4f))
* Warm restart ([#443](#443))
([ff96236](ff96236))


### Documentation

* **graphs:** Documenting some missing features
([#423](#423))
([8addbd8](8addbd8))
</details>

<details><summary>graphs: 0.6.3</summary>

##
[0.6.3](graphs-0.6.2...graphs-0.6.3)
(2025-08-01)


### Features

* **graphs:** Add lat weighted attribute
([#223](#223))
([5dd32ca](5dd32ca))
* **graphs:** Support to export edges to npz
([#395](#395))
([e21738f](e21738f))


### Bug Fixes

* Dropping 3.9 ([#436](#436))
([f6c0214](f6c0214))
* **graphs:** Revert PR
[#379](#379)
([#409](#409))
([d51219f](d51219f))
* **graphs:** Throw error instead of raising warning when graph exists.
([#379](#379))
([6ec6c18](6ec6c18))
* **graphs:** Undo masking when torch-cluster is installed
([#375](#375))
([9f75c06](9f75c06))


### Documentation

* **graphs:** Documenting some missing features
([#423](#423))
([8addbd8](8addbd8))
</details>

<details><summary>models: 0.9.0</summary>

##
[0.9.0](models-0.8.1...models-0.9.0)
(2025-08-01)


### ⚠ BREAKING CHANGES

* for schemas of data processors
([#433](#433))

### Features

* **model:** Postprocessors for leaky boundings
([#315](#315))
([b54562b](b54562b))
* **models:** Checkpointed Mapper Chunking
([#406](#406))
([8577772](8577772))
* **models:** Mapper edge sharding
([#366](#366))
([326751d](326751d))


### Bug Fixes

* Dropping 3.9 ([#436](#436))
([f6c0214](f6c0214))
* For schemas of data processors
([#433](#433))
([539939b](539939b))
* **models,traininig:** Hierarchical model + integration test
([#400](#400))
([71dfd89](71dfd89))
* **models:** Remove repeated lines
([#377](#377))
([1f0b861](1f0b861))
* **models:** Uneven channel sharding
([#385](#385))
([dd095c4](dd095c4))
* Pydantic model validator not working in transformer schema
([#422](#422))
([42f437a](42f437a))
* Remove dead code and fix typo
([#357](#357))
([8c615ba](8c615ba))
</details>

---
> [!IMPORTANT]
> Please do not change the PR title, manifest file, or any other
automatically generated content in this PR unless you understand the
implications. Changes here can break the release process.
> 
> ⚠️ Merging this PR will:
> - Create a new release
> - Trigger deployment pipelines
> - Update package versions

 **Before merging:**
 - Ensure all tests pass
 - Review the changelog carefully
 - Get required approvals

[Release-please
documentation](https://github.com/googleapis/release-please)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@anaprietonem anaprietonem anaprietonem approved these changes

@dietervdb-meteo dietervdb-meteo dietervdb-meteo left review comments

Assignees

@JPXKQX JPXKQX

Labels

ATS Approval Not Needed No approval needed by ATS bug Something isn't working graphs

Projects

Anemoi-dev
Status: Done

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@JPXKQX @anaprietonem @dietervdb-meteo