Coverage for ase / io / siesta_input.py: 67.47%

1# fmt: off 

2 

3"""SiestaInput""" 

4import warnings 

5 

6import numpy as np 

7 

8from ase import Atoms 

9from ase.constraints import FixAtoms, FixCartesian, FixedLine, FixedPlane 

10 

11 

12class SiestaInput: 

13 """SiestaInput""" 

14 @classmethod 

15 def is_along_cartesian(cls, norm_dir: np.ndarray) -> bool: 

16 """Return whether `norm_dir` is along a Cartesian coordidate.""" 

17 directions = [ 

18 [+1, 0, 0], 

19 [-1, 0, 0], 

20 [0, +1, 0], 

21 [0, -1, 0], 

22 [0, 0, +1], 

23 [0, 0, -1], 

24 ] 

25 for direction in directions: 

26 if np.allclose(norm_dir, direction, rtol=0.0, atol=1e-6): 

27 return True 

28 return False 

29 

30 @classmethod 

31 def generate_kpts(cls, kpts): 

32 """Write kpts. 

33 

34 Parameters 

35 ---------- 

36 - f : Open filename. 

37 """ 

38 yield '\n' 

39 yield '#KPoint grid\n' 

40 yield '%block kgrid_Monkhorst_Pack\n' 

41 

42 for i in range(3): 

43 s = '' 

44 if i < len(kpts): 

45 number = kpts[i] 

46 displace = 0.0 

47 else: 

48 number = 1 

49 displace = 0 

50 for j in range(3): 

51 if j == i: 

52 write_this = number 

53 else: 

54 write_this = 0 

55 s += f' {write_this:d} ' 

56 s += f'{displace:1.1f}\n' 

57 yield s 

58 yield '%endblock kgrid_Monkhorst_Pack\n' 

59 yield '\n' 

60 

61 @classmethod 

62 def get_species(cls, atoms, species, basis_set): 

63 from ase.calculators.siesta.parameters import Species 

64 

65 # For each element use default species from the species input, or set 

66 # up a default species from the general default parameters. 

67 tags = atoms.get_tags() 

68 default_species = [ 

69 s for s in species 

70 if (s['tag'] is None) and s['symbol'] in atoms.symbols] 

71 default_symbols = [s['symbol'] for s in default_species] 

72 for symbol in atoms.symbols: 

73 if symbol not in default_symbols: 

74 spec = Species(symbol=symbol, 

75 basis_set=basis_set, 

76 tag=None) 

77 default_species.append(spec) 

78 default_symbols.append(symbol) 

79 assert len(default_species) == len(set(atoms.symbols)) 

80 

81 # Set default species as the first species. 

82 species_numbers = np.zeros(len(atoms), int) 

83 i = 1 

84 for spec in default_species: 

85 mask = atoms.symbols == spec['symbol'] 

86 species_numbers[mask] = i 

87 i += 1 

88 

89 # Set up the non-default species. 

90 non_default_species = [s for s in species if s['tag'] is not None] 

91 for spec in non_default_species: 

92 mask1 = tags == spec['tag'] 

93 mask2 = atoms.symbols == spec['symbol'] 

94 mask = np.logical_and(mask1, mask2) 

95 if sum(mask) > 0: 

96 species_numbers[mask] = i 

97 i += 1 

98 all_species = default_species + non_default_species 

99 

100 return all_species, species_numbers 

101 

102 @classmethod 

103 def make_xyz_constraints(cls, atoms: Atoms): 

104 """ Create coordinate-resolved list of constraints [natoms, 0:3] 

105 The elements of the list must be integers 0 or 1 

106 1 -- means that the coordinate will be updated during relaxation 

107 0 -- mains that the coordinate will be fixed during relaxation 

108 """ 

109 moved = np.ones((len(atoms), 3), dtype=int) # (0: fixed, 1: updated) 

110 for const in atoms.constraints: 

111 if isinstance(const, FixAtoms): 

112 moved[const.get_indices()] = 0 

113 elif isinstance(const, FixedLine): 

114 norm_dir = const.dir / np.linalg.norm(const.dir) 

115 if not cls.is_along_cartesian(norm_dir): 

116 raise RuntimeError( 

117 f'norm_dir {norm_dir} is not one of the Cartesian axes' 

118 ) 

119 norm_dir = norm_dir.round().astype(int) 

120 moved[const.get_indices()] = norm_dir 

121 elif isinstance(const, FixedPlane): 

122 norm_dir = const.dir / np.linalg.norm(const.dir) 

123 if not cls.is_along_cartesian(norm_dir): 

124 raise RuntimeError( 

125 f'norm_dir {norm_dir} is not one of the Cartesian axes' 

126 ) 

127 norm_dir = norm_dir.round().astype(int) 

128 moved[const.get_indices()] = abs(1 - norm_dir) 

129 elif isinstance(const, FixCartesian): 

130 moved[const.get_indices()] = 1 - const.mask.astype(int) 

131 else: 

132 warnings.warn(f'Constraint {const!s} is ignored') 

133 return moved