Coverage for /builds/ase/ase/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 - f : Open filename. 

36 """ 

37 yield '\n' 

38 yield '#KPoint grid\n' 

39 yield '%block kgrid_Monkhorst_Pack\n' 

40 

41 for i in range(3): 

42 s = '' 

43 if i < len(kpts): 

44 number = kpts[i] 

45 displace = 0.0 

46 else: 

47 number = 1 

48 displace = 0 

49 for j in range(3): 

50 if j == i: 

51 write_this = number 

52 else: 

53 write_this = 0 

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

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

56 yield s 

57 yield '%endblock kgrid_Monkhorst_Pack\n' 

58 yield '\n' 

59 

60 @classmethod 

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

62 from ase.calculators.siesta.parameters import Species 

63 

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

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

66 tags = atoms.get_tags() 

67 default_species = [ 

68 s for s in species 

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

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

71 for symbol in atoms.symbols: 

72 if symbol not in default_symbols: 

73 spec = Species(symbol=symbol, 

74 basis_set=basis_set, 

75 tag=None) 

76 default_species.append(spec) 

77 default_symbols.append(symbol) 

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

79 

80 # Set default species as the first species. 

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

82 i = 1 

83 for spec in default_species: 

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

85 species_numbers[mask] = i 

86 i += 1 

87 

88 # Set up the non-default species. 

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

90 for spec in non_default_species: 

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

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

93 mask = np.logical_and(mask1, mask2) 

94 if sum(mask) > 0: 

95 species_numbers[mask] = i 

96 i += 1 

97 all_species = default_species + non_default_species 

98 

99 return all_species, species_numbers 

100 

101 @classmethod 

102 def make_xyz_constraints(cls, atoms: Atoms): 

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

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

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

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

107 """ 

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

109 for const in atoms.constraints: 

110 if isinstance(const, FixAtoms): 

111 moved[const.get_indices()] = 0 

112 elif isinstance(const, FixedLine): 

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

114 if not cls.is_along_cartesian(norm_dir): 

115 raise RuntimeError( 

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

117 ) 

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

119 moved[const.get_indices()] = norm_dir 

120 elif isinstance(const, FixedPlane): 

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

122 if not cls.is_along_cartesian(norm_dir): 

123 raise RuntimeError( 

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

125 ) 

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

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

128 elif isinstance(const, FixCartesian): 

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

130 else: 

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

132 return moved