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

72 statements  

« prev     ^ index     » next       coverage.py v7.5.3, created at 2025-08-02 00:12 +0000

1# fmt: off 

2 

3import re 

4 

5import numpy as np 

6 

7from ase import Atoms 

8from ase.geometry import cellpar_to_cell 

9 

10from .parser import _define_pattern 

11 

12# Geometry block parser 

13_geom = _define_pattern( 

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

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

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

17 """\ 

18geometry units angstrom nocenter noautosym noautoz 

19 system crystal units angstrom 

20 lattice_vectors 

21 4.0000000000000000e+00 0.0000000000000000e+00 0.0000000000000000e+00 

22 0.0000000000000000e+00 5.5264780000000000e+00 0.0000000000000000e+00 

23 0.0000000000000000e+00 0.0000000000000000e+00 4.5963089999999998e+00 

24 end 

25 O 5.0000000000000000e-01 5.0000000000000011e-01 5.6486824536818558e-01 

26 H 5.0000000000000000e-01 6.3810586054988372e-01 4.3513175463181430e-01 

27 H 5.0000000000000000e-01 3.6189413945011639e-01 4.3513175463181430e-01 

28end 

29 

30""", re.M) 

31 

32# Finds crystal specification 

33_crystal = _define_pattern( 

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

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

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

37 """\ 

38 system crystal units angstrom 

39 lattice_vectors 

40 4.0000000000000000e+00 0.0000000000000000e+00 0.0000000000000000e+00 

41 0.0000000000000000e+00 5.5264780000000000e+00 0.0000000000000000e+00 

42 0.0000000000000000e+00 0.0000000000000000e+00 4.5963089999999998e+00 

43 end 

44""", re.M) 

45 

46# Finds 3d-periodic unit cell 

47_cell_3d = _define_pattern( 

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

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

50 """\ 

51 lattice_vectors 

52 4.0000000000000000e+00 0.0000000000000000e+00 0.0000000000000000e+00 

53 0.0000000000000000e+00 5.5264780000000000e+00 0.0000000000000000e+00 

54 0.0000000000000000e+00 0.0000000000000000e+00 4.5963089999999998e+00 

55""", re.M) 

56 

57# Extracts chemical species from a geometry block 

58_species = _define_pattern( 

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

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

61) 

62 

63 

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

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

66 atomslist = [] 

67 for match in _geom.findall(text): 

68 symbols = [] 

69 positions = [] 

70 for atom in _species.findall(match): 

71 atom = atom.split() 

72 symbols.append(atom[0]) 

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

74 positions = np.array(positions) 

75 atoms = Atoms(symbols) 

76 cell, pbc = _get_cell(text) 

77 pos = np.zeros_like(positions) 

78 for dim, ipbc in enumerate(pbc): 

79 if ipbc: 

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

81 else: 

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

83 atoms.set_cell(cell) 

84 atoms.pbc = pbc 

85 atoms.set_positions(pos) 

86 atomslist.append(atoms) 

87 return atomslist[index] 

88 

89 

90def _get_cell(text): 

91 # first check whether there is a lattice definition 

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

93 lattice = _cell_3d.findall(text) 

94 if lattice: 

95 pbc = [True, True, True] 

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

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

98 return cell, pbc 

99 pbc = [False, False, False] 

100 lengths = [None, None, None] 

101 angles = [None, None, None] 

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

103 row = row.strip().lower() 

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

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

106 pbc[dim] = True 

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

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

109 if row.startswith(angle): 

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

111 

112 if not np.any(pbc): 

113 return None, pbc 

114 

115 for i in range(3): 

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

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

118 assert angles[c] is not None 

119 if angles[c] is not None: 

120 assert pbc[a] and pbc[b] 

121 

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

123 if np.all(pbc): 

124 return cellpar_to_cell(lengths + angles), pbc 

125 

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

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

128 dim = np.argmax(pbc) 

129 cell[dim, dim] = lengths[dim] 

130 return cell, pbc 

131 

132 # Hardest case: two lattice vectors are specified. 

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

134 angledim = np.argmin(pbc) 

135 cell[dim1, dim1] = lengths[dim1] 

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

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

138 return cell, pbc