Coverage for /builds/debichem-team/python-ase/ase/io/siesta_input.py: 67.47%

1"""SiestaInput""" 

2import warnings 

3 

4import numpy as np 

5 

6from ase import Atoms 

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

8 

9 

10class SiestaInput: 

11 """SiestaInput""" 

12 @classmethod 

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

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

15 directions = [ 

16 [+1, 0, 0], 

17 [-1, 0, 0], 

18 [0, +1, 0], 

19 [0, -1, 0], 

20 [0, 0, +1], 

21 [0, 0, -1], 

22 ] 

23 for direction in directions: 

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

25 return True 

26 return False 

27 

28 @classmethod 

29 def generate_kpts(cls, kpts): 

30 """Write kpts. 

31 

32 Parameters: 

33 - f : Open filename. 

34 """ 

35 yield '\n' 

36 yield '#KPoint grid\n' 

37 yield '%block kgrid_Monkhorst_Pack\n' 

38 

39 for i in range(3): 

40 s = '' 

41 if i < len(kpts): 

42 number = kpts[i] 

43 displace = 0.0 

44 else: 

45 number = 1 

46 displace = 0 

47 for j in range(3): 

48 if j == i: 

49 write_this = number 

50 else: 

51 write_this = 0 

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

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

54 yield s 

55 yield '%endblock kgrid_Monkhorst_Pack\n' 

56 yield '\n' 

57 

58 @classmethod 

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

60 from ase.calculators.siesta.parameters import Species 

61 

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

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

64 tags = atoms.get_tags() 

65 default_species = [ 

66 s for s in species 

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

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

69 for symbol in atoms.symbols: 

70 if symbol not in default_symbols: 

71 spec = Species(symbol=symbol, 

72 basis_set=basis_set, 

73 tag=None) 

74 default_species.append(spec) 

75 default_symbols.append(symbol) 

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

77 

78 # Set default species as the first species. 

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

80 i = 1 

81 for spec in default_species: 

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

83 species_numbers[mask] = i 

84 i += 1 

85 

86 # Set up the non-default species. 

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

88 for spec in non_default_species: 

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

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

91 mask = np.logical_and(mask1, mask2) 

92 if sum(mask) > 0: 

93 species_numbers[mask] = i 

94 i += 1 

95 all_species = default_species + non_default_species 

96 

97 return all_species, species_numbers 

98 

99 @classmethod 

100 def make_xyz_constraints(cls, atoms: Atoms): 

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

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

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

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

105 """ 

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

107 for const in atoms.constraints: 

108 if isinstance(const, FixAtoms): 

109 moved[const.get_indices()] = 0 

110 elif isinstance(const, FixedLine): 

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

112 if not cls.is_along_cartesian(norm_dir): 

113 raise RuntimeError( 

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

115 ) 

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

117 moved[const.get_indices()] = norm_dir 

118 elif isinstance(const, FixedPlane): 

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

120 if not cls.is_along_cartesian(norm_dir): 

121 raise RuntimeError( 

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

123 ) 

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

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

126 elif isinstance(const, FixCartesian): 

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

128 else: 

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

130 return moved