Add Gaussian perturbation checkpoint state

[NickGeneva]

Earth2Studio Pull Request

Description

Second PR in the checkpointing stack.

Adds checkpoint opt-in state to the Gaussian perturbation so restartable workflows can preserve its internal torch.Generator state without relying on global Torch or NumPy RNG state. The generator remains internal to the perturbation; checkpointing is controlled through the bound checkpoint state proxy and the selected state_policy.

Stack

1. Add checkpoint utilities and workflow support #912: checkpoint utilities and built-in workflow support
2. This PR: Gaussian perturbation checkpoint state
3. Add FCN3 checkpoint state #923: FCN3 checkpoint state
4. Add model checkpoint update developer skill #924: developer skill for adding model checkpoint support

Validation

• uv run ruff check earth2studio/perturbation/gaussian.py test/perturbation/test_gaussian.py
• uv run pytest test/perturbation/test_gaussian.py::test_gaussian_checkpoint_state_round_trip -q
• git diff --check

Full test/perturbation/test_gaussian.py was not clean in this local environment because torch_harmonics is not installed for the optional CorrelatedSphericalGaussian tests.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• An issue is linked to this pull request.

Dependencies

Depends on #912.

[greptile-apps]

Greptile Summary

This PR opts the Gaussian perturbation into the checkpoint framework introduced in #912 by adding a _GaussianCheckpointState dataclass that holds a torch.Tensor snapshot of an internal torch.Generator, and wiring save/restore logic through bind_checkpoint_state.

• _get_generator lazily creates a per-device generator and restores its state from the checkpoint when one is loaded, replacing the previous global-RNG dependency.
• _save_generator_state records the pre-call state (level 1, for exact replay) or the post-call state (level 2, for resuming from the next step), gated by checkpoint_enabled and checkpoint_level.
• A new test_gaussian_checkpoint test exercises both level-1 and level-2 round-trips with disk-backed checkpoint sessions.

Confidence Score: 5/5

The change is well-scoped and safe to merge; the generator save/restore paths are exercised by the new test and the core perturbation logic is unchanged.

All changes are additive: a new internal generator, a new checkpoint state dataclass, and two new private methods. The existing __call__ output is equivalent to the previous global-RNG call when checkpointing is disabled. No data paths, serialization formats, or public APIs are broken.

No files require special attention; both changed files are straightforward.

Important Files Changed

Filename	Overview
earth2studio/perturbation/gaussian.py	Adds _GaussianCheckpointState dataclass and wires an internal torch.Generator to the checkpoint framework; pre/post state capture logic and device-aware generator selection look correct.
test/perturbation/test_gaussian.py	New test_gaussian_checkpoint exercises level-1 and level-2 round-trips correctly; name in the PR description does not match the actual test function name.

_{Reviews (2): Last reviewed commit: "Rename Gaussian checkpoint test" | Re-trigger Greptile}

[NickGeneva]

@greptile-apps

[pzharrington]

LGTM! Love PRs like this 😅

[NickGeneva]

/blossom-ci

[NickGeneva]

/blossom-ci

[NickGeneva]

/blossom-ci

[NickGeneva]

/blossom-ci

[NickGeneva]

/blossom-ci