Coverage for sparc/sparc_parsers/geopt.py: 88%

1""" 

2Created on Thu Oct 18 14:16:21 2018 

3 

4Ben Comer (Georgia Tech) 

5 

6This file has been heavily modified since SPARC 0.1 

7 

8TODO: more descriptions about this file io parser 

9""" 

10from warnings import warn 

11 

12import numpy as np 

13from ase.units import Bohr, GPa, Hartree 

14 

15# Safe wrappers for both string and fd 

16from ase.utils import reader, writer 

17 

18from ..api import SparcAPI 

19from .utils import strip_comments 

20 

21# TODO: should allow user to select the api 

22defaultAPI = SparcAPI() 

23 

24 

25@reader 

26def _read_geopt(fileobj): 

27 """ 

28 Parse the geopt information 

29 Each geopt is similar to the static block, except that the field name is started by 

30 ':' 

31 The relaxations are separated by ':RELAXSTEP:' seperators 

32 """ 

33 contents = fileobj.read() 

34 # label = get_label(fileobj, ".ion") 

35 # The geopt comments are simply discarded 

36 data, comments = strip_comments(contents) 

37 

38 # find the index for all atom type lines. They should be at the 

39 # top of their block 

40 step_bounds = [i for i, x in enumerate(data) if ":RELAXSTEP:" in x] + [len(data)] 

41 raw_geopt_blocks = [ 

42 data[start:end] for start, end in zip(step_bounds[:-1], step_bounds[1:]) 

43 ] 

44 geopt_steps = [_read_geopt_step(step) for step in raw_geopt_blocks] 

45 

46 return {"geopt": geopt_steps} 

47 

48 

49def _read_geopt_step(raw_step_text): 

50 """Parse a geopt step and compose the data dict 

51 

52 Arguments 

53 raw_step_text: list of lines within the step 

54 """ 

55 header, body = raw_step_text[0], raw_step_text[1:] 

56 if ":RELAXSTEP:" not in header: 

57 raise ValueError("Wrong geopt format! The :RELAXSTEP: label is missing.") 

58 # Geopt file uses 1-indexed step names, convert to 0-indexed 

59 step = int(header.split(":RELAXSTEP:")[-1]) - 1 

60 bounds = [i for i, x in enumerate(body) if ":" in x] + [len(body)] 

61 blocks = [body[start:end] for start, end in zip(bounds[:-1], bounds[1:])] 

62 data = {} 

63 for block in blocks: 

64 header_block, body_block = block[0], block[1:] 

65 header_name = header_block.split(":")[1] 

66 header_data = header_block.split(":")[-1].strip() 

67 if len(header_data) > 0: 

68 block_raw_data = [header_data] + body_block 

69 else: 

70 block_raw_data = body_block 

71 # import pdb; pdb.set_trace() 

72 raw_value = np.genfromtxt(block_raw_data, dtype=float) 

73 if "R(Bohr)" in header_name: 

74 name = "positions" 

75 value = raw_value.reshape((-1, 3)) * Bohr 

76 elif "E(Ha)" in header_name: 

77 name = "energy" 

78 value = float(raw_value) * Hartree 

79 elif "F(Ha/Bohr)" in header_name: 

80 name = "forces" 

81 value = raw_value.reshape((-1, 3)) * Hartree / Bohr 

82 elif "CELL" in header_name: 

83 name = "cell" 

84 value = raw_value * Bohr 

85 elif "VOLUME" in header_name: 

86 name = "volume" 

87 value = raw_value * Bohr**3 

88 elif "LATVEC" in header_name: 

89 # TODO: the LATVEC is ambiguous. Are the results a unit cell, or full cell? 

90 name = "latvec" 

91 value = raw_value * Bohr 

92 elif "STRESS" in header_name: 

93 # Stress handling in geopt output can be different 

94 # on low-dimensional systems. If the stress matrix is 3x3, 

95 # the unit is GPa, while lower dimensional stress matrices 

96 # are using Hartree / Bohr**2 or Hartree / Bohr 

97 dim = raw_value.shape[0] 

98 if dim == 3: 

99 name = "stress" 

100 stress_ev_a3 = raw_value * GPa 

101 # Standard stress value, use Voigt representation 

102 value = np.array( 

103 [ 

104 stress_ev_a3[0, 0], 

105 stress_ev_a3[1, 1], 

106 stress_ev_a3[2, 2], 

107 stress_ev_a3[1, 2], 

108 stress_ev_a3[0, 2], 

109 stress_ev_a3[0, 1], 

110 ] 

111 ) 

112 elif dim == 2: 

113 name = "stress_2d" 

114 value = raw_value * Hartree / Bohr**2 

115 elif dim == 1: 

116 name = "stress_1d" 

117 value = raw_value * Hartree / Bohr 

118 else: 

119 raise ValueError("Incorrect stress matrix dimension!") 

120 else: 

121 warn( 

122 f"Field {header_name} is not known to geopt! I'll use the results as is." 

123 ) 

124 name = header_name 

125 value = raw_value 

126 data[name] = value 

127 # Special treatment for latvec & cell 

128 if ("cell" in data) and ("latvec" in data): 

129 # TODO: check 

130 cell_, lat_ = data["cell"], data["latvec"] 

131 unit_lat = lat_ / np.linalg.norm(lat_, axis=1, keepdims=True) 

132 cell = (unit_lat.T * cell_).T 

133 data["ase_cell"] = cell 

134 data["step"] = step 

135 return data 

136 

137 

138@writer 

139def _write_geopt( 

140 fileobj, 

141 data_dict, 

142): 

143 raise NotImplementedError("Writing geopt file from SPARC-X-API not supported!")