Skip to content

Conversation

@hholb
Copy link
Contributor

@hholb hholb commented Oct 20, 2025

Summary

Resolves: #291

This updates the NPTLangevinState to inherit from MDState instead of SimState. Updated the state calculations to use momenta instead of velocities to match other MD integrators.

Checklist

Before a pull request can be merged, the following items must be checked:

  • Doc strings have been added in the Google docstring format.
  • Run ruff on your code.
  • Tests have been added for any new functionality or bug fixes.

We highly recommended installing the prek hooks running in CI locally to speedup the development process. Simply run pip install prek && prek install to install the hooks which will check your code before each commit.

@hholb hholb marked this pull request as ready for review October 20, 2025 16:15
hholb added 3 commits October 20, 2025 11:02
When inheriting from MDState (which ultimately inherits from SimState
with a field that has a default value), Python's dataclass mechanism
requires all fields without defaults to come before fields with defaults.
Using kw_only=True resolves this ordering constraint while maintaining
backward compatibility since all existing instantiations already use
keyword arguments.
NPTLangevinState now follows the same pattern as MDState (its parent class):
- Store momenta as a field
- Compute velocities as a property (inherited from MDState)

This makes NPTLangevinState consistent with other MD integrator states and
properly leverages the MDState inheritance hierarchy.

Changes:
- Remove velocities field from NPTLangevinState dataclass
- Remove custom momenta property (now inherited from MDState)
- Update all step functions to work with momenta:
  - Convert momenta to velocities where needed: momenta / masses.unsqueeze(-1)
  - Convert velocities back to momenta for updates: velocities * masses.unsqueeze(-1)
- Update noise generation to match momenta shape
- Update npt_langevin_init to pass momenta instead of velocities
- Update test fixture to use momenta field
Copy link
Member

@CompRhys CompRhys left a comment

Choose a reason for hiding this comment

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

Thanks for the fixes! A few nit comments that should be fixed before merging.

@hholb
Copy link
Contributor Author

hholb commented Oct 20, 2025

Thanks for the review! Tweaks made!

@CompRhys CompRhys self-requested a review October 20, 2025 21:26
Signed-off-by: Rhys Goodall <rhys.goodall@outlook.com>
Copy link
Member

@CompRhys CompRhys left a comment

Choose a reason for hiding this comment

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

lgtm

@CompRhys CompRhys merged commit 728c29c into TorchSim:main Oct 20, 2025
95 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Convert NPTLangevinState to inherit from MDState and not StateSim

2 participants