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

82 statements  

« prev     ^ index     » next       coverage.py v7.13.4, created at 2026-03-16 12:31 +0000

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 __repr__(self) -> str: 

76 return str(self) 

77 

78 

79class Prim(DynasorAtoms): 

80 def __str__(self): 

81 strings = [f"""Primitive cell: 

82Number of atoms: {self.n_atoms} 

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

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

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

86Cell: 

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

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

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

90"""] 

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

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

93 atom_s = [] 

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

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

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

97 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}' 

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

99 

100 strings = strings + atom_s 

101 

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

103 

104 return string 

105 

106 

107class Supercell(DynasorAtoms): 

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

109 

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

111 calculated. 

112 

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

114 wrapped. 

115 

116 Parameters 

117 ---------- 

118 supercell 

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

120 prim 

121 Primitive structure. 

122 """ 

123 

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

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

126 super().__init__(supercell) 

127 

128 # determine P-matrix relating supercell to primitive cell 

129 from dynasor.tools.structures import get_P_matrix 

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

131 self._P_inv = inv(self.P) 

132 

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

134 from dynasor.tools.structures import get_offset_index 

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

136 

137 @property 

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

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

140 return self._P.copy() 

141 

142 @property 

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

144 """Inverse of `P`.""" 

145 return self._P_inv.copy() 

146 

147 @property 

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

149 """The offset of each atom.""" 

150 return self._offsets.copy() 

151 

152 @property 

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

154 """The basis index of each atom""" 

155 return self._indices.copy() 

156 

157 @property 

158 def n_cells(self) -> int: 

159 """Number of unit cells""" 

160 return self.n_atoms // self.prim.n_atoms 

161 

162 def __str__(self): 

163 

164 string = f"""Supercell: 

165Number of atoms: {self.n_atoms} 

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

167Number of unit cells: {self.n_cells} 

168Cell: 

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

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

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

172P-matrix: 

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

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

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

176{self.prim} 

177""" 

178 return string