Skip to content

feat: add TorchSim integration#155

Merged
yanghan234 merged 19 commits into
mainfrom
feat/torchsim-integration
May 11, 2026
Merged

feat: add TorchSim integration#155
yanghan234 merged 19 commits into
mainfrom
feat/torchsim-integration

Conversation

@yanghan234
Copy link
Copy Markdown
Collaborator

@yanghan234 yanghan234 commented Apr 29, 2026 •
edited
Loading

Summary

Port TorchSim support from the internal repo, providing a GPU-accelerated model wrapper that adapts MatterSim potentials as a TorchSim ModelInterface.

Note: This PR depends on #154 and should only be merged after #154 is merged. Once that is done, this branch will be rebased onto main.

Breaking Change: Python >=3.12 Required

torch-sim-atomistic requires Python 3.12+, so requires-python has been bumped from >=3.10 to >=3.12. All CI pipelines, docs, README, MODEL_CARD, environment.yaml, and cibuildwheel config have been updated accordingly.

Changes

New package: mattersim.torchsim

  • TorchSimWrapper: adapts MatterSim Potential as a TorchSim ModelInterface
  • Graph construction bridge (build_graph_from_simstate)
  • Model loading with AOTI compilation support (get_torchsim_wrapper)

New documentation

  • docs/user_guide/torchsim.rst: user guide covering wrapper creation, structure relaxation with ts.optimize(), molecular dynamics with ts.integrate(), and trajectory saving

Dependencies

  • Added torch-sim-atomistic>=0.6.0 as required dependency
  • Removed unused emmet-core dependency
  • Bumped requires-python from >=3.10 to >=3.12

CI/CD

  • Test matrix: dropped Python 3.10, added 3.13
  • Release pipeline: wheel builds only for 3.12/3.13
  • GitHub Pages: bumped to Python 3.12

Tests

Package import smoke test and TorchSimWrapper forward pass test. GPU tests marked @slow and @requires_gpu.

@yanghan234 yanghan234 changed the base branch from feat/extract-internal-optimizations to main April 29, 2026 09:22
yanghan234 and others added 11 commits April 30, 2026 13:09
@yanghan234
feat: add load_mattersim helper to forcefield
bd8d94f 
Add a convenience function that wraps Potential.from_checkpoint() with
gradient checkpointing support and automatic version normalization from
checkpoint paths. Exported from mattersim.forcefield.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@yanghan234
feat: add torch-sim-atomistic as required dependency
a4a0114 
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@yanghan234
feat: add torchsim integration for batch MD and relaxation
1cf3171 
Port TorchSim support from the internal repo, providing GPU-accelerated
batch molecular dynamics and structure relaxation via torch-sim.

New package: mattersim.torchsim
- TorchSimWrapper: adapts MatterSim Potential as a TorchSim ModelInterface
- TorchSimBatchMD: batch MD with NVT/NPT integrators, temperature schedules
- TorchSimBatchRelaxer: batch structure optimization with FIRE/etc.
- Settings classes (IntegratorSettings, OptimizerSettings) with validation
- Trajectory loading/saving via TorchSim H5MD format
- Graph construction bridge (build_graph_from_simstate)
- Model loading with AOTI compilation support

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@yanghan234
test: add torchsim integration tests
baab0cf 
27 tests covering settings validation, TorchSimWrapper forward pass,
batch relaxation (file + in-memory), batch MD with multiple integrators,
trajectory continuation, per-system temperature schedules, and
trajectory loader convenience constructors.

GPU-heavy tests are marked @slow and @requires_gpu so they are skipped
on CPU-only machines.

Also adds mattersim_potential_best_device fixture to conftest.py.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@yanghan234
fix: set version in load_mattersim for autobatcher lookup
88d399b 
load_mattersim now always sets Potential.version so the pre-computed
autobatcher memory scaler table is used instead of falling back to
slow runtime memory estimation.

- Default checkpoint (load_path=None) gets 'mattersim-v1.0.0-1M'
- Explicit paths are normalized to canonical version strings

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@yanghan234
chore: clean up dead python_version < 3.10 conditionals
f1e5110 
Since requires-python is now >=3.12 (for torch-sim-atomistic), the
python_version < 3.10 dependency conditionals for emmet-core and numpy
are unreachable. Simplify to unconditional versions.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@yanghan234
ci: update test matrix to Python 3.12+ only
ba32da3 
Drop Python 3.10 from CI matrix since requires-python is now >=3.12.
Add Python 3.13 to the test matrix. Update macOS job to 3.12.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@yanghan234
ci: drop Python 3.10/3.11 wheel builds from release pipeline
08f307c 
Only build wheels for Python 3.12 and 3.13 since torch-sim-atomistic
requires Python >=3.12.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@yanghan234
ci: bump gh-pages docs build to Python 3.12
d28f581 
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@yanghan234
docs: update Python version requirement to 3.12+ in README
a7f5a83 
Update badge, prerequisite, and conda create example to reflect
the new minimum Python 3.12 requirement.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@yanghan234
chore: update Python version to 3.12+ in MODEL_CARD, environment.yaml…
3abacda 
…, cibuildwheel

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@yanghan234 yanghan234 force-pushed the feat/torchsim-integration branch from 23c4659 to 3abacda Compare April 30, 2026 13:09
@yanghan234 yanghan234 marked this pull request as ready for review April 30, 2026 13:12
yanghan234 and others added 3 commits May 6, 2026 16:03
@yanghan234
refactor: remove batch MD and relaxation runners from torchsim
b378ce0 
Remove the batch runner infrastructure (TorchSimBatchRunner, TorchSimBatchMD,
TorchSimBatchRelaxer), their settings classes (IntegratorSettings,
OptimizerSettings), trajectory loaders, and associated tests.

Keep the core TorchSim wrapper (TorchSimWrapper, graph_construction,
model_loading) that adapts MatterSim potentials as TorchSim ModelInterface.

Deleted files:
- src/mattersim/torchsim/base.py
- src/mattersim/torchsim/batch_relax.py
- src/mattersim/torchsim/md.py
- src/mattersim/torchsim/settings.py
- src/mattersim/torchsim/settings_base.py
- src/mattersim/torchsim/trajectory_loader.py

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@yanghan234
docs: add TorchSim simulation guide
2d631f8 
Add user guide for running structure relaxation and molecular dynamics
using the TorchSim backend with MatterSim potentials. Covers wrapper
creation, optimization, MD with various integrators, and trajectory
saving.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@yanghan234
docs: fix convergence function API in TorchSim guide
c2b7b96 
Use correct parameter name 'force_tol' (not 'fmax') and drop
'include_cell_forces' from the minimal example.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@yanghan234 yanghan234 changed the title feat: add TorchSim integration for batch MD and relaxation feat: add TorchSim integration May 6, 2026
flt17
flt17 reviewed May 6, 2026
View reviewed changes
Comment thread docs/user_guide/torchsim.rst Outdated
flt17
flt17 reviewed May 6, 2026
View reviewed changes
Comment thread src/mattersim/forcefield/potential.py Outdated
flt17
flt17 reviewed May 6, 2026
View reviewed changes
Comment thread pyproject.toml Outdated
flt17
flt17 approved these changes May 6, 2026
View reviewed changes
Copy link
Copy Markdown
Contributor

@flt17 flt17 left a comment

Choose a reason for hiding this comment

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

LGTM, a few minor things only.

danielzuegner
danielzuegner reviewed May 7, 2026
View reviewed changes
Comment thread src/mattersim/torchsim/torchsim_wrapper.py Outdated
danielzuegner
danielzuegner reviewed May 7, 2026
View reviewed changes
Comment thread pyproject.toml Outdated
yanghan234 and others added 3 commits May 11, 2026 12:20
@yanghan234
fix: remove unused emmet-core dependency
48f556c 
emmet-core is not imported anywhere in the codebase. Removing the
dead dependency per review feedback.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@yanghan234
bump torchsim >= 0.6.0
45a0c40
@yanghan234
fix: remove remaining load_mattersim references
673db37 
Replace load_mattersim() calls in model_loading.py with direct
Potential.from_checkpoint() + enable_gradient_checkpointing() calls.
Remove load_mattersim from forcefield __init__.py exports.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
yanghan234 and others added 2 commits May 11, 2026 14:23
@yanghan234
fix: remove _sanitize_outputs and sanitize_nan plumbing
621e7ce 
Remove internal-only NaN sanitization logic per review feedback.
Removed from TorchSimWrapper and get_torchsim_wrapper.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@yanghan234
fix: drop explicit float64, use default float32 dtype
c229010 
float64 support belongs to PR #156. For now use the wrapper's
default float32 in model_loading and doc examples.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@yanghan234 yanghan234 merged commit c6801fd into main May 11, 2026
7 checks passed
@yanghan234 yanghan234 deleted the feat/torchsim-integration branch May 11, 2026 14:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@danielzuegner danielzuegner danielzuegner approved these changes

@flt17 flt17 flt17 approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@yanghan234 @danielzuegner @flt17