From 3d71ff506b60887cd4c4a79239bcab21aaa4ca21 Mon Sep 17 00:00:00 2001 From: SanggyuChong Date: Mon, 26 May 2025 13:56:18 +0200 Subject: [PATCH 1/5] ase random rotation --- src/flashmd/ase/velocity_verlet.py | 36 ++++++++++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/src/flashmd/ase/velocity_verlet.py b/src/flashmd/ase/velocity_verlet.py index b7f7452..7ee38e3 100644 --- a/src/flashmd/ase/velocity_verlet.py +++ b/src/flashmd/ase/velocity_verlet.py @@ -8,9 +8,9 @@ from metatensor.torch.atomistic import System import ase from ..stepper import FlashMDStepper -import numpy as np - +import numpy as np +from scipy.spatial.transform import Rotation class VelocityVerlet(MolecularDynamics): def __init__( self, @@ -19,6 +19,7 @@ def __init__( model: MetatensorAtomisticModel | List[MetatensorAtomisticModel], device: str | torch.device = "auto", rescale_energy: bool = True, + random_rotation: bool = False, **kwargs, ): super().__init__(atoms, timestep, **kwargs) @@ -47,15 +48,39 @@ def __init__( self.stepper = FlashMDStepper(models, n_time_steps, self.device) self.rescale_energy = rescale_energy + self.random_rotation = random_rotation def step(self): + if self.rescale_energy: old_energy = self.atoms.get_total_energy() system = _convert_atoms_to_system( self.atoms, device=self.device, dtype=self.dtype ) + + if self.random_rotation: + # generate a random rotation matrix (with i-PI utils for consistency) + R = torch.tensor( + _random_R(), + device=system.positions.device, + dtype=system.positions.dtype, + ) + # apply the random rotation + system.cell = system.cell @ R.T + system.positions = system.positions @ R.T + momenta = system.get_data("momenta").block(0).values.squeeze() + momenta[:] = momenta @ R.T # does the change in place + new_system = self.stepper.step(system) + + if self.random_rotation: + # revert q, p, and cell to the original reference frame + new_system.cell = system.cell @ R + new_system.positions = system.positions @ R + new_momenta = new_system.get_data("momenta").block(0).values.squeeze() + new_momenta[:] = new_momenta @ R + self.atoms.set_positions(new_system.positions.detach().cpu().numpy()) self.atoms.set_momenta( new_system.get_data("momenta") @@ -126,3 +151,10 @@ def _convert_atoms_to_system( ), ) return system + + +def _random_R(): + R = Rotation.random().as_matrix() + if np.random.rand() < 0.5: + R[:, 0] *= -1 + return R From 1d10d05aad7a0dd04db874b8d355306bcfb3ae83 Mon Sep 17 00:00:00 2001 From: SanggyuChong Date: Mon, 26 May 2025 14:20:49 +0200 Subject: [PATCH 2/5] update dependencies --- src/flashmd/ase/bussi.py | 3 ++- src/flashmd/ase/langevin.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/flashmd/ase/bussi.py b/src/flashmd/ase/bussi.py index e773e07..63cbb12 100644 --- a/src/flashmd/ase/bussi.py +++ b/src/flashmd/ase/bussi.py @@ -17,9 +17,10 @@ def __init__( time_constant: float = 10.0 * ase.units.fs, device: str | torch.device = "auto", rescale_energy: bool = True, + random_rotation: bool = False, **kwargs, ): - super().__init__(atoms, timestep, model, device, rescale_energy, **kwargs) + super().__init__(atoms, timestep, model, device, rescale_energy, random_rotation, **kwargs) self.temperature_K = temperature_K self.time_constant = time_constant diff --git a/src/flashmd/ase/langevin.py b/src/flashmd/ase/langevin.py index ee71890..2bcd231 100644 --- a/src/flashmd/ase/langevin.py +++ b/src/flashmd/ase/langevin.py @@ -19,9 +19,10 @@ def __init__( time_constant: float = 100.0 * ase.units.fs, device: str | torch.device = "auto", rescale_energy: bool = True, + random_rotation: bool = False, **kwargs, ): - super().__init__(atoms, timestep, model, device, rescale_energy, **kwargs) + super().__init__(atoms, timestep, model, device, rescale_energy, random_rotation, **kwargs) self.temperature_K = temperature_K self.friction = 1.0 / time_constant From 1c7f49c5463e0fbc75f3ff8aacae964e69ffb41d Mon Sep 17 00:00:00 2001 From: SanggyuChong Date: Mon, 26 May 2025 14:24:02 +0200 Subject: [PATCH 3/5] prevent potential error build up in cell transformation --- src/flashmd/ase/velocity_verlet.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/flashmd/ase/velocity_verlet.py b/src/flashmd/ase/velocity_verlet.py index 7ee38e3..43cb4d3 100644 --- a/src/flashmd/ase/velocity_verlet.py +++ b/src/flashmd/ase/velocity_verlet.py @@ -67,6 +67,7 @@ def step(self): dtype=system.positions.dtype, ) # apply the random rotation + old_cell = system.cell system.cell = system.cell @ R.T system.positions = system.positions @ R.T momenta = system.get_data("momenta").block(0).values.squeeze() @@ -75,8 +76,8 @@ def step(self): new_system = self.stepper.step(system) if self.random_rotation: - # revert q, p, and cell to the original reference frame - new_system.cell = system.cell @ R + # revert q, p to the original reference frame, load old cell + new_system.cell = old_cell new_system.positions = system.positions @ R new_momenta = new_system.get_data("momenta").block(0).values.squeeze() new_momenta[:] = new_momenta @ R From f1b55d8465cf6da009fe4294b663a794e93dff97 Mon Sep 17 00:00:00 2001 From: SanggyuChong Date: Mon, 26 May 2025 14:27:36 +0200 Subject: [PATCH 4/5] update comment --- src/flashmd/ase/velocity_verlet.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/flashmd/ase/velocity_verlet.py b/src/flashmd/ase/velocity_verlet.py index 43cb4d3..c9f3433 100644 --- a/src/flashmd/ase/velocity_verlet.py +++ b/src/flashmd/ase/velocity_verlet.py @@ -60,7 +60,7 @@ def step(self): ) if self.random_rotation: - # generate a random rotation matrix (with i-PI utils for consistency) + # generate a random rotation matrix with SciPy R = torch.tensor( _random_R(), device=system.positions.device, From 58c47b4af07077e3a9ac5f413afa281e80c72a20 Mon Sep 17 00:00:00 2001 From: "Sanggyu \"Raymond\" Chong" <87842409+SanggyuChong@users.noreply.github.com> Date: Mon, 26 May 2025 14:33:41 +0200 Subject: [PATCH 5/5] Update README.md to warn about ASE --- README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index a1eb3eb..caabccc 100644 --- a/README.md +++ b/README.md @@ -6,4 +6,5 @@ This repository contains model architectures and helper functions to use transfo This is experimental software and should only be used if you know what you're doing. See the [cookbook recipe](http://atomistic-cookbook.org) for a usage example, and [this preprint](http://arxiv.org) for a discussion of the theory, benchmarks, and -of the potential limitations. +of the potential limitations. Note that the ASE implementation is even more experimental +than i-PI, and we strongly recommend the i-PI implementation over ASE.