1import re 

2 

3import numpy as np 

4 

5from ase import Atoms 

6from ase.geometry import cellpar_to_cell 

7from .parser import _define_pattern 

8 

9# Geometry block parser 

10_geom = _define_pattern( 

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

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

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

14 """\ 

15geometry units angstrom nocenter noautosym noautoz 

16 system crystal units angstrom 

17 lattice_vectors 

18 4.0000000000000000e+00 0.0000000000000000e+00 0.0000000000000000e+00 

19 0.0000000000000000e+00 5.5264780000000000e+00 0.0000000000000000e+00 

20 0.0000000000000000e+00 0.0000000000000000e+00 4.5963089999999998e+00 

21 end 

22 O 5.0000000000000000e-01 5.0000000000000011e-01 5.6486824536818558e-01 

23 H 5.0000000000000000e-01 6.3810586054988372e-01 4.3513175463181430e-01 

24 H 5.0000000000000000e-01 3.6189413945011639e-01 4.3513175463181430e-01 

25end 

26 

27""", re.M) 

28 

29# Finds crystal specification 

30_crystal = _define_pattern( 

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

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

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

34 """\ 

35 system crystal units angstrom 

36 lattice_vectors 

37 4.0000000000000000e+00 0.0000000000000000e+00 0.0000000000000000e+00 

38 0.0000000000000000e+00 5.5264780000000000e+00 0.0000000000000000e+00 

39 0.0000000000000000e+00 0.0000000000000000e+00 4.5963089999999998e+00 

40 end 

41""", re.M) 

42 

43# Finds 3d-periodic unit cell 

44_cell_3d = _define_pattern( 

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

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

47 """\ 

48 lattice_vectors 

49 4.0000000000000000e+00 0.0000000000000000e+00 0.0000000000000000e+00 

50 0.0000000000000000e+00 5.5264780000000000e+00 0.0000000000000000e+00 

51 0.0000000000000000e+00 0.0000000000000000e+00 4.5963089999999998e+00 

52""", re.M) 

53 

54# Extracts chemical species from a geometry block 

55_species = _define_pattern( 

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

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

58) 

59 

60 

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

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

63 atomslist = [] 

64 for match in _geom.findall(text): 

65 symbols = [] 

66 positions = [] 

67 for atom in _species.findall(match): 

68 atom = atom.split() 

69 symbols.append(atom[0]) 

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

71 positions = np.array(positions) 

72 atoms = Atoms(symbols) 

73 cell, pbc = _get_cell(text) 

74 pos = np.zeros_like(positions) 

75 for dim, ipbc in enumerate(pbc): 

76 if ipbc: 

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

78 else: 

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

80 atoms.set_cell(cell) 

81 atoms.pbc = pbc 

82 atoms.set_positions(pos) 

83 atomslist.append(atoms) 

84 return atomslist[index] 

85 

86 

87def _get_cell(text): 

88 # first check whether there is a lattice definition 

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

90 lattice = _cell_3d.findall(text) 

91 if lattice: 

92 pbc = [True, True, True] 

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

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

95 return cell, pbc 

96 pbc = [False, False, False] 

97 lengths = [None, None, None] 

98 angles = [None, None, None] 

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

100 row = row.strip().lower() 

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

102 if row.startswith('lat_{}'.format(vecname)): 

103 pbc[dim] = True 

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

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

106 if row.startswith(angle): 

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

108 

109 if not np.any(pbc): 

110 return None, pbc 

111 

112 for i in range(3): 

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

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

115 assert angles[c] is not None 

116 if angles[c] is not None: 

117 assert pbc[a] and pbc[b] 

118 

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

120 if np.all(pbc): 

121 return cellpar_to_cell(lengths + angles), pbc 

122 

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

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

125 dim = np.argmax(pbc) 

126 cell[dim, dim] = lengths[dim] 

127 return cell, pbc 

128 

129 # Hardest case: two lattice vectors are specified. 

130 dim1, dim2 = [dim for dim, ipbc in enumerate(pbc) if ipbc] 

131 angledim = np.argmin(pbc) 

132 cell[dim1, dim1] = lengths[dim1] 

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

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

135 return cell, pbc