Coverage for ase / constraints / fix_com.py: 92.31%

1from ase.constraints.constraint import FixConstraint, IndexedConstraint 

2 

3 

4class FixCom(FixConstraint): 

5 """Constraint class for fixing the center of mass.""" 

6 

7 index = slice(None) # all atoms 

8 

9 def get_removed_dof(self, atoms): 

10 return 3 

11 

12 def adjust_positions(self, atoms, new): 

13 masses = atoms.get_masses()[self.index] 

14 old_cm = atoms.get_center_of_mass(indices=self.index) 

15 new_cm = masses @ new[self.index] / masses.sum() 

16 diff = old_cm - new_cm 

17 new += diff 

18 

19 def adjust_momenta(self, atoms, momenta): 

20 """Adjust momenta so that the center-of-mass velocity is zero.""" 

21 masses = atoms.get_masses()[self.index] 

22 velocity_com = momenta[self.index].sum(axis=0) / masses.sum() 

23 momenta[self.index] -= masses[:, None] * velocity_com 

24 

25 def adjust_forces(self, atoms, forces): 

26 # Eqs. (3) and (7) in https://doi.org/10.1021/jp9722824 

27 masses = atoms.get_masses()[self.index] 

28 lmd = masses @ forces[self.index] / sum(masses**2) 

29 forces[self.index] -= masses[:, None] * lmd 

30 

31 def todict(self): 

32 return {'name': 'FixCom', 'kwargs': {}} 

33 

34 

35class FixSubsetCom(FixCom, IndexedConstraint): 

36 """Constraint class for fixing the center of mass of a subset of atoms.""" 

37 

38 def __init__(self, indices): 

39 super().__init__(indices=indices) 

40 

41 def todict(self): 

42 return { 

43 'name': self.__class__.__name__, 

44 'kwargs': {'indices': self.index.tolist()}, 

45 }