Features/constraints

CHANGELOG.md

@@ -1,6 +1,16 @@
 <!-- markdownlint-disable -->
 # Changelog
 
+## Unreleased
+
+### 🎉 New Features
+* Constraints support for molecular dynamics and optimization by @thomasloux in [#294](https://github.com/TorchSim/torch-sim/pull/294)
+  - Added `FixAtoms` constraint to fix specific atoms in place
+  - Added `FixCom` constraint to prevent center of mass drift
+  - Constraints automatically adjust degrees of freedom for accurate temperature calculations
+  - Full support across all integrators (NVE, NVT, NPT) and optimizers (FIRE, Gradient Descent)
+  - Constraints preserved during state manipulation (slicing, splitting, concatenation)
+
 ## v0.5.0
 
 This release focuses on improving batch processing capabilities across TorchSim. The neighbor list module has been completely refactored to support batched calculations with multiple backend implementations, elastic tensor calculations now leverage batched operations for improved performance, and a bug fix ensures Monte Carlo swaps work correctly with ragged (different-sized) systems.

examples/scripts/3_Dynamics/3.10_Hybrid_swap_mc.py

@@ -67,9 +67,8 @@ class HybridSwapMCState(ts.SwapMCState, MDState):
         last_swap: Last swap attempted
     """
 
-    last_permutation: torch.Tensor
     _atom_attributes = (
-        ts.SwapMCState._atom_attributes | MDState._atom_attributes | {"last_permutation"}  # noqa: SLF001
+        ts.SwapMCState._atom_attributes | MDState._atom_attributes  # noqa: SLF001
     )
     _system_attributes = (
         ts.SwapMCState._system_attributes | MDState._system_attributes  # noqa: SLF001

examples/tutorials/hybrid_swap_tutorial.py

@@ -100,9 +100,11 @@ class HybridSwapMCState(SwapMCState, MDState):
     from MDState.
     """
 
-    last_permutation: torch.Tensor
     _atom_attributes = (
-        MDState._atom_attributes | {"last_permutation"}  # noqa: SLF001
+        ts.SwapMCState._atom_attributes | MDState._atom_attributes  # noqa: SLF001
+    )
+    _system_attributes = (
+        ts.SwapMCState._system_attributes | MDState._system_attributes  # noqa: SLF001
     )
 
 

tests/conftest.py

@@ -339,7 +339,7 @@ def distorted_fcc_al_conventional_sim_state() -> ts.SimState:
     positions = atoms_fcc.get_positions()
     np_rng = np.random.default_rng(seed=42)
     positions += np_rng.normal(scale=0.01, size=positions.shape)
-    atoms_fcc.set_positions(positions)
+    atoms_fcc.positions = positions
 
     # Convert the ASE Atoms object to SimState (will be a single batch with 4 atoms)
     return ts.io.atoms_to_state(atoms_fcc, device=DEVICE, dtype=DTYPE)