Coverage for ase / constraints / fix_scaled.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 FixScaled(IndexedConstraint): 

10 """Fix atoms in the directions of the unit vectors. 

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 Cell directions to be fixed. (False: unfixed, True: fixed) 

18 """ 

19 

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

21 # XXX The unused cell keyword is there for compatibility 

22 # with old trajectory files. 

23 super().__init__(indices=a) 

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

25 

26 def get_removed_dof(self, atoms: Atoms): 

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

28 

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

30 cell = atoms.cell 

31 scaled_old = cell.scaled_positions(atoms.positions[self.index]) 

32 scaled_new = cell.scaled_positions(new[self.index]) 

33 scaled_new[:, self.mask] = scaled_old[:, self.mask] 

34 new[self.index] = cell.cartesian_positions(scaled_new) 

35 

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

37 # Forces are covariant to the coordinate transformation, 

38 # use the inverse transformations 

39 cell = atoms.cell 

40 scaled_forces = cell.cartesian_positions(forces[self.index]) 

41 scaled_forces *= -(self.mask - 1) 

42 forces[self.index] = cell.scaled_positions(scaled_forces) 

43 

44 def todict(self): 

45 return { 

46 'name': 'FixScaled', 

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

48 } 

49 

50 def __repr__(self): 

51 name = type(self).__name__ 

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