Add Atlas example

[pzharrington]

Earth2Studio Pull Request

Description

Adds an example (~5min runtime) demoing the two ways of doing Atlas inference. Also reorganizes some imports in nn/atlas.py for better protection against confusing errors. Previously physicsnemo.Module was imported in the same try block as natten, einops, timm, and torch_harmonics. If any one of those secondary deps fails to import, PhysicsNeMoModule falls back to object, so classes like NattenCombineDiT don't inherit Module._override_args and can't load a checkpoint (which points the user away from the actual issue). Fixed now with different import order.

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.
• Assess and address Greptile feedback (AI code review bot for guidance; use discretion, addressing all feedback is not required).

Dependencies

[greptile-apps]

Greptile Summary

This PR adds an Atlas inference example demonstrating both a manual forward-pass approach and the run.deterministic workflow, and fixes a subtle import-ordering bug in nn/atlas.py where a failure in any optional secondary dependency (natten, einops, etc.) would silently demote PhysicsNeMoModule to object and break checkpoint loading with a misleading error. The pyproject.toml also adds a <3.14 Python upper bound to guard against torch.compile/natten incompatibility.

Confidence Score: 5/5

Safe to merge; all remaining findings are P2 style/formatting issues that do not affect correctness.

The nn/atlas.py import fix is correct and unambiguously improves error clarity. The pyproject.toml change is a sensible guard. The example file has three minor issues (runtime length, one short RST underline, one trailing-whitespace character) but none affect runtime correctness or data integrity.

examples/02_medium_range/06_atlas_inference.py — trailing whitespace on line 75 and the short RST underline on line 150 should be fixed before the docs build to avoid pre-commit/Sphinx warnings.

Important Files Changed

Filename	Overview
earth2studio/models/nn/atlas.py	Separates the physicsnemo import into its own try/except block so that failures in the secondary optional dependencies (natten, einops, timm, torch_harmonics) no longer silently cause PhysicsNeMoModule to fall back to object, which previously broke checkpoint loading with a misleading error.
examples/02_medium_range/06_atlas_inference.py	New Atlas inference example demonstrating both a manual forward pass and the run.deterministic workflow. License header, structure, and plot are correct, but the stated ~14-minute runtime exceeds the 10-minute example rule, one RST section underline is too short, and a trailing-whitespace character on line 75 would fail the pre-commit hook.
pyproject.toml	Adds <3.14 upper bound to the Python version constraint to block Python 3.14 until torch.compile (a natten dependency) supports it; includes an explanatory inline comment.

_{Reviews (1): Last reviewed commit: "Add atlas example" | Re-trigger Greptile}

[pzharrington]

/blossom-ci

[pzharrington]

/blossom-ci

[NickGeneva]

imo maybe dont need this code block in the recipe to keep it more concise

[NickGeneva]

But if you think its really needed, can leave it.

[NickGeneva]

Consider making this a warning instead.

[NickGeneva]

Can't do this just for natten / atlas.

You should be able to restrict just natten to below a given python version I think...

But regardless, we suggest python 3.12 in the install docs.

[NickGeneva]

Haven't ran this yet, but whats the run time / GPU requirement?