Coverage for ase / constraints / fix_cartesian.py: 100.00%

1from __future__ import annotations 

2 

3import numpy as np 

4 

5from ase import Atoms 

6from ase.constraints.constraint import IndexedConstraint 

7 

8 

9class FixCartesian(IndexedConstraint): 

10 """Fix atoms in the directions of the cartesian coordinates. 

11 

12 Parameters 

13 ---------- 

14 a : Sequence[int] 

15 Indices of atoms to be fixed. 

16 mask : tuple[bool, bool, bool], default: (True, True, True) 

17 Cartesian directions to be fixed. (False: unfixed, True: fixed) 

18 """ 

19 

20 def __init__(self, a, mask=(True, True, True)): 

21 super().__init__(indices=a) 

22 self.mask = np.asarray(mask, bool) 

23 

24 def get_removed_dof(self, atoms: Atoms): 

25 return self.mask.sum() * len(self.index) 

26 

27 def adjust_positions(self, atoms: Atoms, new): 

28 new[self.index] = np.where( 

29 self.mask[None, :], 

30 atoms.positions[self.index], 

31 new[self.index], 

32 ) 

33 

34 def adjust_forces(self, atoms: Atoms, forces): 

35 forces[self.index] *= ~self.mask[None, :] 

36 

37 def todict(self): 

38 return { 

39 'name': 'FixCartesian', 

40 'kwargs': {'a': self.index.tolist(), 'mask': self.mask.tolist()}, 

41 } 

42 

43 def __repr__(self): 

44 name = type(self).__name__ 

45 return f'{name}(indices={self.index.tolist()}, {self.mask.tolist()})'