Geotransolver + GP by ktangsali · Pull Request #1557 · NVIDIA/physicsnemo

ktangsali · 2026-04-09T19:10:26Z

PhysicsNeMo Pull Request

Description

Adds a GP head to Geotransolver recipe.

Checklist

I am familiar with the Contributing Guidelines.
New or existing tests cover these changes.
The documentation is up to date with these changes.
The CHANGELOG.md is up to date with these changes.
An issue is linked to this pull request.
If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

* Bug fixes for ShardTensor+SongUNet * Handle dtensor spec in sharded view * Fix SongUNet with ShardTensor when using zero embedding * Use buffer for zero embed --------- Co-authored-by: Peter Harrington <48932392+pzharrington@users.noreply.github.com> Co-authored-by: Peter Harrington <pharrington@nvidia.com>

* comment out e2grid and makani installs * fix dtype * update sfno test * update version * some fixes for healpix tests, update to nc install, fix test dir path mismatch * revert changes to healpix code * don't change the default path * fix graphcast doctest, update doctest command to ignore onnx module because of conflict with onnx.utils * skip pytest for debugging, use floating point comparison for gumbel softmax * make the test more robust * bring back pytests

…DIA#1433) * Bugfix with num_steps parameters in CorrDiff generate.py Signed-off-by: Charlelie Laurent <claurent@nvidia.com> * Update examples/weather/corrdiff/conf/base/generation/sampler/stochastic.yaml Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --------- Signed-off-by: Charlelie Laurent <claurent@nvidia.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

NVIDIA#1430) * DSMLoss implementation + minor bugfixes in noise_schedulers.py and preconditioners.py Signed-off-by: Charlelie Laurent <claurent@nvidia.com> * Bugfix for missing abstract class in noise_schedulers.py Signed-off-by: Charlelie Laurent <claurent@nvidia.com> * Grammar improvements Signed-off-by: Charlelie Laurent <claurent@nvidia.com> * Added 'reduction' argument to MSEDSMLoss + added new WeightedMSEDSMLoss Signed-off-by: Charlelie Laurent <claurent@nvidia.com> --------- Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

* Fix dlesym compat map * modafno fix * format

@greptile-apps

* add fixes related to dask, dask is needed for healpix datapipes * update vtk install * add missing dependency * add dask dependency * Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

Add zarr to datapipe example requirements after it was dropped from core reqs.

The FigConvNet Example needs this preprocessing script, which is how I noticed pyvista and vtk are not listed here as requirements.

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

* Remove 'perf' extra from physicsnemo installation because NGC containers already include transformer-engine * update deterministic settings

* Add wandb to requirements * Modify requirements for trimesh and add wandb Updated trimesh version constraint and added wandb.

…IA#1415) * Fix attention and turn off feature broadcasting. * Fix scalar loading shapes * Update volume.yaml Ensure the volume example works out of the box. * Fix Geotransolver inference tests

* Adds mesh docs on top of RC branch * Update docs/mesh/boundaries.rst Co-authored-by: megnvidia <mmiranda@nvidia.com> * Apply suggestions from code review Co-authored-by: megnvidia <mmiranda@nvidia.com> * Refine documentation for mesh geometry functions to clarify usage and API exposure for advanced cases. * better subdivision descriptions * clearer docs --------- Co-authored-by: megnvidia <mmiranda@nvidia.com>

* Update API docs and structure. * clean-up and re-organization of docs * fix based on new api * remove unused sections * update image paths * Update docs/api/models/diffusion.rst Co-authored-by: megnvidia <mmiranda@nvidia.com> * Update docs/api/models/operators.rst Co-authored-by: megnvidia <mmiranda@nvidia.com> * Update docs/api/models/weather.rst Co-authored-by: megnvidia <mmiranda@nvidia.com> * Update docs/api/physicsnemo.core.rst Co-authored-by: megnvidia <mmiranda@nvidia.com> * Update docs/api/physicsnemo.diffusion.rst Co-authored-by: megnvidia <mmiranda@nvidia.com> * Update docs/api/physicsnemo.diffusion.rst Co-authored-by: megnvidia <mmiranda@nvidia.com> * Update docs/api/physicsnemo.utils.rst Co-authored-by: megnvidia <mmiranda@nvidia.com> * Update docs/api/physicsnemo.utils.rst Co-authored-by: megnvidia <mmiranda@nvidia.com> * Update docs/api/physicsnemo.utils.rst Co-authored-by: megnvidia <mmiranda@nvidia.com> * Update docs/api/physicsnemo.utils.rst Co-authored-by: megnvidia <mmiranda@nvidia.com> * Update docs/api/physicsnemo.utils.rst Co-authored-by: megnvidia <mmiranda@nvidia.com> * Update docs/api/physicsnemo.utils.rst Co-authored-by: megnvidia <mmiranda@nvidia.com> * Update docs/api/physicsnemo.utils.rst Co-authored-by: megnvidia <mmiranda@nvidia.com> * fix formatting --------- Co-authored-by: Kaustubh Tangsali <71059996+ktangsali@users.noreply.github.com> Co-authored-by: Kaustubh Tangsali <ktangsali@nvidia.com> Co-authored-by: megnvidia <mmiranda@nvidia.com>

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

* minor edits to the install guide * add more details * minor doc fix * add transolver to the api index

* New API docs for diffusion Signed-off-by: Charlelie Laurent <claurent@nvidia.com> * Some fixes in nested API references Signed-off-by: Charlelie Laurent <claurent@nvidia.com> * Revert some changes Signed-off-by: Charlelie Laurent <claurent@nvidia.com> * Some fixes Signed-off-by: Charlelie Laurent <claurent@nvidia.com> * Some clarifications in introduction.rst Signed-off-by: Charlelie Laurent <claurent@nvidia.com> * Some clarification in diffusion models.rst Signed-off-by: Charlelie Laurent <claurent@nvidia.com> * Fixed note sections in preconditioners.py Signed-off-by: Charlelie Laurent <claurent@nvidia.com> * Fix some broken short-form refs in losses.py Signed-off-by: Charlelie Laurent <claurent@nvidia.com> * Updated DPSScorePredictor class name in samplers.rst Signed-off-by: Charlelie Laurent <claurent@nvidia.com> * Enhance clarity and structure of introduction section Refactor introduction section for clarity and readability. Improve formatting and organization of key concepts related to the PhysicsNeMo diffusion framework. * Refactor metrics.rst for improved clarity and formatting Reformatted the description of the module to use bullet points for clarity. Adjusted wording for consistency and readability. * Fix punctuation and enhance clarity in models.rst Corrected punctuation and improved clarity in the documentation. * Addressed PR comments Signed-off-by: Charlelie Laurent <claurent@nvidia.com> --------- Signed-off-by: Charlelie Laurent <claurent@nvidia.com> Co-authored-by: megnvidia <mmiranda@nvidia.com>

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

* Trying again with datapipes check in * Update docs/api/datapipes/physicsnemo.datapipes.cae.rst Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Update docs/api/datapipes/physicsnemo.datapipes.cae.rst Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Update docs/api/datapipes/physicsnemo.datapipes.rst Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Update docs/api/datapipes/physicsnemo.datapipes.rst Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Update docs/api/datapipes/physicsnemo.datapipes.rst Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Update docs/api/datapipes/physicsnemo.datapipes.rst Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Update docs/api/datapipes/physicsnemo.datapipes.rst Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Update docs/api/datapipes/physicsnemo.datapipes.transforms.rst Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Resolve api commits for datapipes * Remove old datapipes api * Add link to the datapipe docs. --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* Improved docs for module.py Signed-off-by: Charlelie Laurent <claurent@nvidia.com> * Fix in save_checkpoint and load_checkpoint docstrings Signed-off-by: Charlelie Laurent <claurent@nvidia.com> * Addressed PR comments Signed-off-by: Charlelie Laurent <claurent@nvidia.com> * Improvements in docs Signed-off-by: Charlelie Laurent <claurent@nvidia.com> * Moved down section about static capture in physicsnemo.utils.rst Signed-off-by: Charlelie Laurent <claurent@nvidia.com> --------- Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

…with different learning rates for different params, matern kernel, test on ood, etc

laserkelvin

Just some minor stuff

laserkelvin

Talked pretty extensively offline - will approve, and some things can be amended

ktangsali · 2026-04-15T17:52:11Z

/blossom-ci

coreyjadams

I know Kelvin has reviewed this extensively, and while I'm sure there is always more to do, I'm approving to unblock. UQ is great to pursue and I think we should get this going!

ktangsali · 2026-04-15T21:11:12Z

/blossom-ci

ktangsali · 2026-04-15T21:24:56Z

/blossom-ci

jleinonen and others added 30 commits February 19, 2026 13:37

Fix dlesym compat map (NVIDIA#1439)

d323b89

* Fix dlesym compat map * modafno fix * format

Update requirements.txt (NVIDIA#1448)

d23f0aa

Add zarr to datapipe example requirements after it was dropped from core reqs.

Update requirements.txt (NVIDIA#1447)

ff17020

The FigConvNet Example needs this preprocessing script, which is how I noticed pyvista and vtk are not listed here as requirements.

Fix broken cross-ref links in docstrings (NVIDIA#1454)

8784aa7

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

Fixes from SFB builds / testing (NVIDIA#1459)

5a6f923

* Remove 'perf' extra from physicsnemo installation because NGC containers already include transformer-engine * update deterministic settings

Wandb fixes (NVIDIA#1458)

474e2c0

* Add wandb to requirements * Modify requirements for trimesh and add wandb Updated trimesh version constraint and added wandb.

GeoTransolver: Fix attention and turn off feature broadcasting. (NVID…

6194f04

…IA#1415) * Fix attention and turn off feature broadcasting. * Fix scalar loading shapes * Update volume.yaml Ensure the volume example works out of the box. * Fix Geotransolver inference tests

Fixes and renaming in dps_guidance.py (NVIDIA#1471)

2841666

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

Minor edits to the install guide (NVIDIA#1470)

eda7e9a

* minor edits to the install guide * add more details * minor doc fix * add transolver to the api index

Fix unresolved conflict (NVIDIA#1477)

eb05732

Signed-off-by: Charlelie Laurent <claurent@nvidia.com>

add wip code 1

7c6505b

first, kinda working setup

6d6d384

make embeddings configurable

fb4626e

add eos configs

51bf68d

add eos configs

f16243a

add scripts after succesful 2 stage training. includes updates to gp …

1a06913

…with different learning rates for different params, matern kernel, test on ood, etc

add combined transolver and GP run

1fb2c7b

add histogram plotting

6ba6b32

add spectral normalization, better consistency, enhanced plotting

92df748

add output scale prior, narrower length scale, embedding normalization

0138f6c

ktangsali and others added 6 commits April 9, 2026 12:18

updates after merging main

20cae52

add plots and experiment results

b5af99e

add plots and experiment results

e4b0713

fix license headers

1ff5313

remove unwanted files, fix docstring issues

68b610d

fix issues with importlinter

0e506f1

ktangsali removed request for CharlelieLrt, NickGeneva and pzharrington April 9, 2026 21:21

ktangsali and others added 4 commits April 9, 2026 14:29

remove changes to base files

54154c1

Merge branch 'main' into gpgeotransolver

f7c5807

revert some changes to upstream code

8b0e6b6

restore normalization.npz from main

a2d5071

ktangsali requested a review from laserkelvin April 13, 2026 23:46

laserkelvin requested changes Apr 14, 2026

View reviewed changes

ktangsali and others added 3 commits April 14, 2026 23:21

address review comments

9607b83

Merge branch 'main' into gpgeotransolver

b22e9ed

update docstrings and unify the readme

234f005

laserkelvin approved these changes Apr 15, 2026

View reviewed changes

ktangsali enabled auto-merge April 15, 2026 17:52

coreyjadams approved these changes Apr 15, 2026

View reviewed changes

RishikeshRanade approved these changes Apr 15, 2026

View reviewed changes

address rishi's comments

8bd3dc2

loliverhennigh approved these changes Apr 15, 2026

View reviewed changes

ktangsali added this pull request to the merge queue Apr 15, 2026

Merged via the queue into NVIDIA:main with commit 06f9d22 Apr 15, 2026
4 checks passed

Conversation

ktangsali commented Apr 9, 2026

PhysicsNeMo Pull Request

Description

Checklist

Dependencies

Review Process

Uh oh!

laserkelvin left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

laserkelvin left a comment

Choose a reason for hiding this comment

Uh oh!

ktangsali commented Apr 15, 2026

Uh oh!

coreyjadams left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ktangsali commented Apr 15, 2026

Uh oh!

ktangsali commented Apr 15, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants