Coverage for /builds/kinetik161/ase/ase/io/nwchem/nwreader_in.py: 80.56%

1import re 

2 

3import numpy as np 

4 

5from ase import Atoms 

6from ase.geometry import cellpar_to_cell 

7 

8from .parser import _define_pattern 

9 

10# Geometry block parser 

11_geom = _define_pattern( 

12 r'^[ \t]*geometry[ \t\S]*\n' 

13 r'((?:^[ \t]*[\S]+[ \t\S]*\n)+)' 

14 r'^[ \t]*end\n\n', 

15 """\ 

16geometry units angstrom nocenter noautosym noautoz 

17 system crystal units angstrom 

18 lattice_vectors 

19 4.0000000000000000e+00 0.0000000000000000e+00 0.0000000000000000e+00 

20 0.0000000000000000e+00 5.5264780000000000e+00 0.0000000000000000e+00 

21 0.0000000000000000e+00 0.0000000000000000e+00 4.5963089999999998e+00 

22 end 

23 O 5.0000000000000000e-01 5.0000000000000011e-01 5.6486824536818558e-01 

24 H 5.0000000000000000e-01 6.3810586054988372e-01 4.3513175463181430e-01 

25 H 5.0000000000000000e-01 3.6189413945011639e-01 4.3513175463181430e-01 

26end 

27 

28""", re.M) 

29 

30# Finds crystal specification 

31_crystal = _define_pattern( 

32 r'^[ \t]*system crystal[ \t\S]*\n' 

33 r'((?:[ \t]*[\S]+[ \t\S]*\n)+?)' 

34 r'^[ \t]*end[ \t]*\n', 

35 """\ 

36 system crystal units angstrom 

37 lattice_vectors 

38 4.0000000000000000e+00 0.0000000000000000e+00 0.0000000000000000e+00 

39 0.0000000000000000e+00 5.5264780000000000e+00 0.0000000000000000e+00 

40 0.0000000000000000e+00 0.0000000000000000e+00 4.5963089999999998e+00 

41 end 

42""", re.M) 

43 

44# Finds 3d-periodic unit cell 

45_cell_3d = _define_pattern( 

46 r'^[ \t]*lattice_vectors[ \t]*\n' 

47 r'^((?:(?:[ \t]+[\S]+){3}\n){3})', 

48 """\ 

49 lattice_vectors 

50 4.0000000000000000e+00 0.0000000000000000e+00 0.0000000000000000e+00 

51 0.0000000000000000e+00 5.5264780000000000e+00 0.0000000000000000e+00 

52 0.0000000000000000e+00 0.0000000000000000e+00 4.5963089999999998e+00 

53""", re.M) 

54 

55# Extracts chemical species from a geometry block 

56_species = _define_pattern( 

57 r'^[ \t]*[A-Z][a-z]?(?:[ \t]+[\S]+){3}\n', 

58 " O 0.0 0.0 0.0\n", re.M, 

59) 

60 

61 

62def read_nwchem_in(fobj, index=-1): 

63 text = ''.join(fobj.readlines()) 

64 atomslist = [] 

65 for match in _geom.findall(text): 

66 symbols = [] 

67 positions = [] 

68 for atom in _species.findall(match): 

69 atom = atom.split() 

70 symbols.append(atom[0]) 

71 positions.append([float(x) for x in atom[1:]]) 

72 positions = np.array(positions) 

73 atoms = Atoms(symbols) 

74 cell, pbc = _get_cell(text) 

75 pos = np.zeros_like(positions) 

76 for dim, ipbc in enumerate(pbc): 

77 if ipbc: 

78 pos += np.outer(positions[:, dim], cell[dim, :]) 

79 else: 

80 pos[:, dim] = positions[:, dim] 

81 atoms.set_cell(cell) 

82 atoms.pbc = pbc 

83 atoms.set_positions(pos) 

84 atomslist.append(atoms) 

85 return atomslist[index] 

86 

87 

88def _get_cell(text): 

89 # first check whether there is a lattice definition 

90 cell = np.zeros((3, 3)) 

91 lattice = _cell_3d.findall(text) 

92 if lattice: 

93 pbc = [True, True, True] 

94 for i, row in enumerate(lattice[0].strip().split('\n')): 

95 cell[i] = [float(x) for x in row.split()] 

96 return cell, pbc 

97 pbc = [False, False, False] 

98 lengths = [None, None, None] 

99 angles = [None, None, None] 

100 for row in text.strip().split('\n'): 

101 row = row.strip().lower() 

102 for dim, vecname in enumerate(['a', 'b', 'c']): 

103 if row.startswith(f'lat_{vecname}'): 

104 pbc[dim] = True 

105 lengths[dim] = float(row.split()[1]) 

106 for i, angle in enumerate(['alpha', 'beta', 'gamma']): 

107 if row.startswith(angle): 

108 angles[i] = float(row.split()[1]) 

109 

110 if not np.any(pbc): 

111 return None, pbc 

112 

113 for i in range(3): 

114 a, b, c = np.roll(np.array([0, 1, 2]), i) 

115 if pbc[a] and pbc[b]: 

116 assert angles[c] is not None 

117 if angles[c] is not None: 

118 assert pbc[a] and pbc[b] 

119 

120 # The easiest case: all three lattice vectors and angles are specified 

121 if np.all(pbc): 

122 return cellpar_to_cell(lengths + angles), pbc 

123 

124 # Next easiest case: exactly one lattice vector has been specified 

125 if np.sum(pbc) == 1: 

126 dim = np.argmax(pbc) 

127 cell[dim, dim] = lengths[dim] 

128 return cell, pbc 

129 

130 # Hardest case: two lattice vectors are specified. 

131 dim1, dim2 = (dim for dim, ipbc in enumerate(pbc) if ipbc) 

132 angledim = np.argmin(pbc) 

133 cell[dim1, dim1] = lengths[dim1] 

134 cell[dim2, dim2] = lengths[dim2] * np.sin(angles[angledim]) 

135 cell[dim2, dim1] = lengths[dim2] * np.cos(angles[angledim]) 

136 return cell, pbc