Coverage for /builds/ase/ase/ase/outputs.py: 98.04%

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

2 

3from __future__ import annotations 

4 

5from abc import ABC, abstractmethod 

6from collections.abc import Mapping 

7from typing import Sequence, Union 

8 

9import numpy as np 

10 

11 

12class Properties(Mapping): 

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

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

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

16 self._setvalue(name, value) 

17 

18 def __len__(self) -> int: 

19 return len(self._dct) 

20 

21 def __iter__(self): 

22 return iter(self._dct) 

23 

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

25 return self._dct[name] 

26 

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

28 if name in self._dct: 

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

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

31 

32 prop = all_outputs[name] 

33 value = prop.normalize_type(value) 

34 shape = np.shape(value) 

35 

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

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

38 

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

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

41 continue 

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

43 

44 self._dct[name] = value 

45 

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

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

48 

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

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

51 """ 

52 shapespec = prop.shapespec 

53 shape = np.shape(value) 

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

55 return False 

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

57 if isinstance(dimspec, str): 

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

59 if dimspec != dim: 

60 return False 

61 return True 

62 

63 def __repr__(self) -> str: 

64 clsname = type(self).__name__ 

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

66 

67 

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

69 

70 

71class Property(ABC): 

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

73 self.name = name 

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

75 raise ValueError(dtype) 

76 self.dtype = dtype 

77 self.shapespec = shapespec 

78 

79 @abstractmethod 

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

81 

82 def __repr__(self) -> str: 

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

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

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

86 

87 

88class ScalarProperty(Property): 

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

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

91 

92 def normalize_type(self, value): 

93 if not np.isscalar(value): 

94 raise TypeError('Expected scalar') 

95 return self.dtype(value) 

96 

97 

98class ArrayProperty(Property): 

99 def normalize_type(self, value): 

100 if np.isscalar(value): 

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

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

103 

104 

105ShapeSpec = Union[str, int] 

106 

107 

108def _defineprop( 

109 name: str, 

110 dtype: type = float, 

111 shape: Union[ShapeSpec, Sequence[ShapeSpec]] = (), 

112) -> Property: 

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

114 

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

116 shape = (shape,) 

117 

118 shape = tuple(shape) 

119 prop: Property 

120 if len(shape) == 0: 

121 prop = ScalarProperty(name, dtype) 

122 else: 

123 prop = ArrayProperty(name, dtype, shape) 

124 

125 assert name not in all_outputs, name 

126 all_outputs[name] = prop 

127 return prop 

128 

129 

130# Atoms, energy, forces, stress: 

131_defineprop('natoms', int) 

132_defineprop('energy', float) 

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

134_defineprop('free_energy', float) 

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

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

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

138 

139# Electronic structure: 

140_defineprop('nbands', int) 

141_defineprop('nkpts', int) 

142_defineprop('nspins', int) 

143_defineprop('fermi_level', float) 

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

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

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

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

148 

149# Miscellaneous: 

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

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

152_defineprop('magmom', float) 

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

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

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

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

157 

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

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

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

161 

162 

163# def to_singlepoint(self, atoms): 

164# from ase.calculators.singlepoint import SinglePointDFTCalculator 

165# return SinglePointDFTCalculator(atoms, 

166# efermi=self.fermi_level, 

167 

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

169# 

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

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

172# the outputs. 

173# 

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

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

176# 

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

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

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

180# versus values provided within the output. 

181# 

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

183# band structure. 

184# 

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

186# 

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

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