Coverage for dynasor/modes/atoms.py: 100%

1from typing import Union 

2import numpy as np 

3from ase.units import fs 

4from ase import Atoms 

5from numpy.typing import NDArray 

6from .tools import inv 

7from ..units import Dalton_to_dmu 

8 

9 

10class DynasorAtoms: 

11 """Dynasor's representation of a structure.""" 

12 def __init__(self, atoms: Atoms): 

13 """Initialized using an ASE Atoms object""" 

14 self._atoms = atoms 

15 

16 @property 

17 def pos(self) -> NDArray[float]: 

18 """Cartesian positions.""" 

19 return self._atoms.positions.copy() 

20 

21 @property 

22 def positions(self) -> NDArray[float]: 

23 """Cartesian positions.""" 

24 return self.pos 

25 

26 @property 

27 def spos(self) -> NDArray[float]: 

28 """Reduced (or scaled) positions of atoms.""" 

29 return self._atoms.get_scaled_positions() 

30 

31 @property 

32 def scaled_positions(self) -> NDArray[float]: 

33 """Reduced (or scaled) positions of atoms.""" 

34 return self.spos 

35 

36 @property 

37 def cell(self) -> NDArray[float]: 

38 """Cell of atoms with cell vectors as rows.""" 

39 return self._atoms.cell.array.copy() 

40 

41 @property 

42 def inv_cell(self) -> NDArray[float]: 

43 """The inverse cell transpose so the inverse cell vectors are rows, no 2pi.""" 

44 return np.linalg.inv(self._atoms.cell.array).T 

45 

46 @property 

47 def numbers(self) -> NDArray[int]: 

48 """Chemical number for each atom, e.g., 1 for H, 2 for He etc.""" 

49 return self._atoms.numbers.copy() 

50 

51 @property 

52 def masses(self) -> NDArray[float]: 

53 """Masses of atoms in dmu.""" 

54 return self._atoms.get_masses() / fs ** 2 # In eVfs²/Ų 

55 

56 @property 

57 def volume(self) -> float: 

58 """Volume of cell.""" 

59 return self._atoms.get_volume() 

60 

61 @property 

62 def n_atoms(self) -> int: 

63 """Number of atoms.""" 

64 return len(self._atoms) 

65 

66 @property 

67 def symbols(self) -> list[str]: 

68 """List of chemical symbol for each element.""" 

69 return list(self._atoms.symbols) 

70 

71 def to_ase(self) -> Atoms: 

72 """Converts the internal Atoms to ASE :class:`Atoms`.""" 

73 return Atoms(cell=self.cell, numbers=self.numbers, positions=self.positions, pbc=True) 

74 

75 def __str__(self) -> str: 

76 return f'{self.__class__.__name__} with {self.n_atoms} atoms' 

77 

78 def __repr__(self) -> str: 

79 return str(self) 

80 

81 

82class Prim(DynasorAtoms): 

83 def __str__(self): 

84 strings = [f"""Primitive cell: 

85Number of atoms: {self.n_atoms} 

86Volume: {self.volume:.3f} 

87Atomic species present: {set(self.symbols)} 

88Atomic numbers present: {set([int(n) for n in self.numbers])} 

89Cell: 

90[[{self.cell[0, 0]:<20}, {self.cell[0, 1]:<20}, {self.cell[0, 2]:<20}], 

91 [{self.cell[1, 0]:<20}, {self.cell[1, 1]:<20}, {self.cell[1, 2]:<20}], 

92 [{self.cell[2, 0]:<20}, {self.cell[2, 1]:<20}, {self.cell[2, 2]:<20}]] 

93"""] 

94 strings.append(f"{'Ind':<5}{'Sym':<5}{'Num':<5}{'Mass (Da)':<10}{'x':<10}{'y':<10}{'z':<10}" 

95 f"{'a':<10}{'b':<10}{'c':<10}") 

96 atom_s = [] 

97 for i, p, sp, m, n, s in zip( 

98 range(self.n_atoms), self.positions, self.spos, self.masses / Dalton_to_dmu, 

99 self.numbers, [a.symbol for a in self.to_ase()]): 

100 atom_s.append(f'{i:<5}{s:<5}{n:<5}{m:<10.2f}{p[0]:<10.3f}{p[1]:<10.3f}{p[2]:<10.3f}' 

101 f'{sp[0]:<10.3f}{sp[1]:<10.3f}{sp[2]:<10.3f}') 

102 

103 strings = strings + atom_s 

104 

105 string = '\n'.join(strings) 

106 

107 return string 

108 

109 

110class Supercell(DynasorAtoms): 

111 """The supercell takes care of some mappings between the primitive and repeated structure. 

112 

113 In particular the P-matrix connecting the cells as well as the offset-index of each atom is 

114 calculated. 

115 

116 Note that the positions cannot be revovered as `offset x cell + basis` since the atoms get 

117 wrapped. 

118 

119 Parameters 

120 ---------- 

121 supercell 

122 Some ideal repetition of the primitive structure and possible wrapping. 

123 prim 

124 Primitive structure. 

125 """ 

126 

127 def __init__(self, supercell: Union[Atoms, DynasorAtoms], prim: Union[Atoms, DynasorAtoms]): 

128 self.prim = Prim(prim.copy()) 

129 super().__init__(supercell) 

130 

131 # determine P-matrix relating supercell to primitive cell 

132 from dynasor.tools.structures import get_P_matrix 

133 self._P = get_P_matrix(self.prim.cell, self.cell) # P C = S 

134 self._P_inv = inv(self.P) 

135 

136 # find the index and offsets for supercell using primitive as base unit 

137 from dynasor.tools.structures import get_offset_index 

138 self._offsets, self._indices = get_offset_index(prim, supercell, wrap=True) 

139 

140 @property 

141 def P(self) -> NDArray[float]: 

142 """P-matrix is defined as dot(P, prim.cell) = supercell.cell""" 

143 return self._P.copy() 

144 

145 @property 

146 def P_inv(self) -> NDArray[float]: 

147 """Inverse of `P`.""" 

148 return self._P_inv.copy() 

149 

150 @property 

151 def offsets(self) -> NDArray[float]: 

152 """The offset of each atom.""" 

153 return self._offsets.copy() 

154 

155 @property 

156 def indices(self) -> NDArray[int]: 

157 """The basis index of each atom""" 

158 return self._indices.copy() 

159 

160 @property 

161 def n_cells(self) -> int: 

162 """Number of unit cells""" 

163 return self.n_atoms // self.prim.n_atoms 

164 

165 def __str__(self): 

166 

167 string = f"""Supercell: 

168Number of atoms: {self.n_atoms} 

169Volume: {self.volume:.3f} 

170Number of unit cells: {self.n_cells} 

171Cell: 

172[[{self.cell[0, 0]:<20}, {self.cell[0, 1]:<20}, {self.cell[0, 2]:<20}], 

173 [{self.cell[1, 0]:<20}, {self.cell[1, 1]:<20}, {self.cell[1, 2]:<20}], 

174 [{self.cell[2, 0]:<20}, {self.cell[2, 1]:<20}, {self.cell[2, 2]:<20}]] 

175P-matrix: 

176[[{self.P[0, 0]:<20}, {self.P[0, 1]:<20}, {self.P[0, 2]:<20}], 

177 [{self.P[1, 0]:<20}, {self.P[1, 1]:<20}, {self.P[1, 2]:<20}], 

178 [{self.P[2, 0]:<20}, {self.P[2, 1]:<20}, {self.P[2, 2]:<20}]] 

179{self.prim} 

180""" 

181 return string