Coverage for ase / outputs.py: 97.98%

1"""Module for ``Property`` and ``Properties``.""" 

2 

3from __future__ import annotations 

4 

5from abc import ABC, abstractmethod 

6from collections.abc import Mapping, Sequence 

7 

8import numpy as np 

9 

10 

11class Properties(Mapping): 

12 def __init__(self, dct: dict) -> None: 

13 self._dct: dict[str, Property] = {} 

14 for name, value in dct.items(): 

15 self._setvalue(name, value) 

16 

17 def __len__(self) -> int: 

18 return len(self._dct) 

19 

20 def __iter__(self): 

21 return iter(self._dct) 

22 

23 def __getitem__(self, name) -> Property: 

24 return self._dct[name] 

25 

26 def _setvalue(self, name: str, value) -> None: 

27 if name in self._dct: 

28 # Which error should we raise for already existing property? 

29 raise ValueError(f'{name} already set') 

30 

31 prop = all_outputs[name] 

32 value = prop.normalize_type(value) 

33 shape = np.shape(value) 

34 

35 if not self.shape_is_consistent(prop, value): 

36 raise ValueError(f'{name} has bad shape: {shape}') 

37 

38 for i, spec in enumerate(prop.shapespec): 

39 if not isinstance(spec, str) or spec in self._dct: 

40 continue 

41 self._setvalue(spec, shape[i]) 

42 

43 self._dct[name] = value 

44 

45 def shape_is_consistent(self, prop: Property, value) -> bool: 

46 """Return whether shape of values is consistent with properties. 

47 

48 For example, forces of shape (7, 3) are consistent 

49 unless properties already have "natoms" with non-7 value. 

50 """ 

51 shapespec = prop.shapespec 

52 shape = np.shape(value) 

53 if len(shapespec) != len(shape): 

54 return False 

55 for dimspec, dim in zip(shapespec, shape): 

56 if isinstance(dimspec, str): 

57 dimspec = self._dct.get(dimspec, dim) 

58 if dimspec != dim: 

59 return False 

60 return True 

61 

62 def __repr__(self) -> str: 

63 clsname = type(self).__name__ 

64 return f'({clsname}({self._dct})' 

65 

66 

67all_outputs: dict[str, Property] = {} 

68 

69 

70class Property(ABC): 

71 def __init__(self, name: str, dtype: type, shapespec: tuple) -> None: 

72 self.name = name 

73 if dtype not in {float, int}: # Others? 

74 raise ValueError(dtype) 

75 self.dtype = dtype 

76 self.shapespec = shapespec 

77 

78 @abstractmethod 

79 def normalize_type(self, value): ... 

80 

81 def __repr__(self) -> str: 

82 typename = self.dtype.__name__ # Extend to other than float/int? 

83 shape = ', '.join(str(dim) for dim in self.shapespec) 

84 return f'Property({self.name!r}, dtype={typename}, shape=[{shape}])' 

85 

86 

87class ScalarProperty(Property): 

88 def __init__(self, name: str, dtype: type) -> None: 

89 super().__init__(name, dtype, ()) 

90 

91 def normalize_type(self, value): 

92 if not np.isscalar(value): 

93 raise TypeError('Expected scalar') 

94 return self.dtype(value) 

95 

96 

97class ArrayProperty(Property): 

98 def normalize_type(self, value): 

99 if np.isscalar(value): 

100 raise TypeError('Expected array, got scalar') 

101 return np.asarray(value, dtype=self.dtype) 

102 

103 

104ShapeSpec = str | int 

105 

106 

107def _defineprop( 

108 name: str, 

109 dtype: type = float, 

110 shape: ShapeSpec | Sequence[ShapeSpec] = (), 

111) -> Property: 

112 """Create, register, and return a property.""" 

113 

114 if isinstance(shape, (int, str)): 

115 shape = (shape,) 

116 

117 shape = tuple(shape) 

118 prop: Property 

119 if len(shape) == 0: 

120 prop = ScalarProperty(name, dtype) 

121 else: 

122 prop = ArrayProperty(name, dtype, shape) 

123 

124 assert name not in all_outputs, name 

125 all_outputs[name] = prop 

126 return prop 

127 

128 

129# Atoms, energy, forces, stress: 

130_defineprop('natoms', int) 

131_defineprop('energy', float) 

132_defineprop('energies', float, shape='natoms') 

133_defineprop('free_energy', float) 

134_defineprop('forces', float, shape=('natoms', 3)) 

135_defineprop('stress', float, shape=6) 

136_defineprop('stresses', float, shape=('natoms', 6)) 

137 

138# Electronic structure: 

139_defineprop('nbands', int) 

140_defineprop('nkpts', int) 

141_defineprop('nspins', int) 

142_defineprop('fermi_level', float) 

143_defineprop('kpoint_weights', float, shape='nkpts') 

144_defineprop('ibz_kpoints', float, shape=('nkpts', 3)) 

145_defineprop('eigenvalues', float, shape=('nspins', 'nkpts', 'nbands')) 

146_defineprop('occupations', float, shape=('nspins', 'nkpts', 'nbands')) 

147 

148# Miscellaneous: 

149_defineprop('dipole', float, shape=3) 

150_defineprop('charges', float, shape='natoms') 

151_defineprop('magmom', float) 

152_defineprop('magmoms', float, shape='natoms') # XXX spinors? 

153_defineprop('polarization', float, shape=3) 

154_defineprop('dielectric_tensor', float, shape=(3, 3)) 

155_defineprop('born_effective_charges', float, shape=('natoms', 3, 3)) 

156 

157# We might want to allow properties that are part of Atoms, such as 

158# positions, numbers, pbc, cell. It would be reasonable for those 

159# concepts to have a formalization outside the Atoms class. 

160 

161 

162# def to_singlepoint(self, atoms): 

163# from ase.calculators.singlepoint import SinglePointDFTCalculator 

164# return SinglePointDFTCalculator(atoms, 

165# efermi=self.fermi_level, 

166 

167# We can also retrieve (P)DOS and band structure. However: 

168# 

169# * Band structure may require bandpath, which is an input, and 

170# may not necessarily be easy or possible to reconstruct from 

171# the outputs. 

172# 

173# * Some calculators may produce the whole BandStructure object in 

174# one go (e.g. while parsing) 

175# 

176# * What about HOMO/LUMO? Can be obtained from 

177# eigenvalues/occupations, but some codes provide real data. We 

178# probably need to distinguish between HOMO/LUMO inferred by us 

179# versus values provided within the output. 

180# 

181# * HOMO is sometimes used as alternative reference energy for 

182# band structure. 

183# 

184# * What about spin-dependent (double) Fermi level? 

185# 

186# * What about 3D arrays? We will almost certainly want to be 

187# connected to an object that can load dynamically from a file.