1# -*- coding: utf-8 -*- 

2# 

3# Copyright (c) 2017, the cclib development team 

4# 

5# This file is part of cclib (http://cclib.github.io) and is distributed under 

6# the terms of the BSD 3-Clause License. 

7 

8"""A writer for chemical JSON (CJSON) files.""" 

9 

10import os.path 

11import json 

12import numpy as np 

13 

14from cclib.io import filewriter 

15from cclib.parser.data import ccData 

16from cclib.parser.utils import find_package 

17 

18_has_openbabel = find_package("openbabel") 

19 

20 

21class CJSON(filewriter.Writer): 

22 """A writer for chemical JSON (CJSON) files.""" 

23 

24 def __init__(self, ccdata, terse=False, *args, **kwargs): 

25 """Initialize the chemical JSON writer object. 

26 

27 Inputs: 

28 ccdata - An instance of ccData, parsed from a logfile. 

29 """ 

30 

31 super(CJSON, self).__init__(ccdata, terse=terse, *args, **kwargs) 

32 

33 def pathname(self, path): 

34 """Return filename without extension to be used as name.""" 

35 name = os.path.basename(os.path.splitext(path)[0]) 

36 return name 

37 

38 def as_dict(self): 

39 """ Build a Python dict with the CJSON data""" 

40 cjson_dict = dict() 

41 # Need to decide on a number format. 

42 cjson_dict['chemical json'] = 0 

43 if self.jobfilename is not None: 

44 cjson_dict['name'] = self.pathname(self.jobfilename) 

45 

46 # These are properties that can be collected using Open Babel. 

47 if _has_openbabel: 

48 cjson_dict['smiles'] = self.pbmol.write('smiles') 

49 cjson_dict['inchi'] = self.pbmol.write('inchi') 

50 cjson_dict['inchikey'] = self.pbmol.write('inchikey') 

51 cjson_dict['formula'] = self.pbmol.formula 

52 # TODO Incorporate unit cell information. 

53 

54 # Iterate through the attribute list present in ccData. Depending on 

55 # the availability of the attribute add it at the right 'level'. 

56 for attribute_name, v in ccData._attributes.items(): 

57 if not hasattr(self.ccdata, attribute_name): 

58 continue 

59 

60 attribute_path = v.attribute_path.split(":") 

61 

62 # Depth of the attribute in the CJSON. 

63 levels = len(attribute_path) 

64 

65 # The attributes which haven't been included in the CJSON format. 

66 if attribute_path[0] == 'N/A': 

67 continue 

68 

69 if attribute_path[0] not in cjson_dict: 

70 cjson_dict[attribute_path[0]] = dict() 

71 l1_data_object = cjson_dict[attribute_path[0]] 

72 

73 # 'moments' and 'atomcoords' key will contain processed data 

74 # obtained from the output file. TODO rewrite this 

75 if attribute_name in ('moments', 'atomcoords'): 

76 if attribute_name == 'moments': 

77 dipole_moment = self._calculate_total_dipole_moment() 

78 if dipole_moment is not None: 

79 cjson_dict['properties'][ccData._attributes['moments'].json_key] = dipole_moment 

80 else: 

81 cjson_dict['atoms']['coords'] = dict() 

82 cjson_dict['atoms']['coords']['3d'] = self.ccdata.atomcoords[-1].flatten().tolist() 

83 continue 

84 

85 if levels == 1: 

86 self.set_JSON_attribute(l1_data_object, attribute_name) 

87 elif levels >= 2: 

88 if attribute_path[1] not in l1_data_object: 

89 l1_data_object[attribute_path[1]] = dict() 

90 l2_data_object = l1_data_object[attribute_path[1]] 

91 

92 if levels == 2: 

93 self.set_JSON_attribute(l2_data_object, attribute_name) 

94 elif levels == 3: 

95 if attribute_path[2] not in l2_data_object: 

96 l2_data_object[attribute_path[2]] = dict() 

97 l3_data_object = l2_data_object[attribute_path[2]] 

98 self.set_JSON_attribute(l3_data_object, attribute_name) 

99 

100 # Attributes which are not directly obtained from the output files. 

101 if hasattr(self.ccdata, 'moenergies') and hasattr(self.ccdata, 'homos'): 

102 if 'energy' not in cjson_dict['properties']: 

103 cjson_dict['properties']['energy'] = dict() 

104 

105 cjson_dict['properties']['energy']['alpha'] = dict() 

106 cjson_dict['properties']['energy']['beta'] = dict() 

107 

108 homo_idx_alpha = int(self.ccdata.homos[0]) 

109 homo_idx_beta = int(self.ccdata.homos[-1]) 

110 energy_alpha_homo = self.ccdata.moenergies[0][homo_idx_alpha] 

111 energy_alpha_lumo = self.ccdata.moenergies[0][homo_idx_alpha + 1] 

112 energy_alpha_gap = energy_alpha_lumo - energy_alpha_homo 

113 energy_beta_homo = self.ccdata.moenergies[-1][homo_idx_beta] 

114 energy_beta_lumo = self.ccdata.moenergies[-1][homo_idx_beta + 1] 

115 energy_beta_gap = energy_beta_lumo - energy_beta_homo 

116 

117 cjson_dict['properties']['energy']['alpha']['homo'] = energy_alpha_homo 

118 cjson_dict['properties']['energy']['alpha']['gap'] = energy_alpha_gap 

119 cjson_dict['properties']['energy']['beta']['homo'] = energy_beta_homo 

120 cjson_dict['properties']['energy']['beta']['gap'] = energy_beta_gap 

121 cjson_dict['properties']['energy']['total'] = self.ccdata.scfenergies[-1] 

122 

123 if hasattr(self.ccdata, 'atomnos'): 

124 cjson_dict['atoms']['elements']['atom count'] = len(self.ccdata.atomnos) 

125 cjson_dict['atoms']['elements']['heavy atom count'] = len([x for x in self.ccdata.atomnos if x > 1]) 

126 

127 # Bond attributes: 

128 if _has_openbabel and (len(self.ccdata.atomnos) > 1): 

129 cjson_dict['bonds'] = dict() 

130 cjson_dict['bonds']['connections'] = dict() 

131 cjson_dict['bonds']['connections']['index'] = [] 

132 for bond in self.bond_connectivities: 

133 cjson_dict['bonds']['connections']['index'].append(bond[0]) 

134 cjson_dict['bonds']['connections']['index'].append(bond[1]) 

135 cjson_dict['bonds']['order'] = [bond[2] for bond in self.bond_connectivities] 

136 

137 if _has_openbabel: 

138 cjson_dict['properties']['molecular mass'] = self.pbmol.molwt 

139 cjson_dict['diagram'] = self.pbmol.write(format='svg') 

140 return cjson_dict 

141 

142 def generate_repr(self): 

143 """Generate the CJSON representation of the logfile data.""" 

144 cjson_dict = self.as_dict() 

145 if self.terse: 

146 return json.dumps(cjson_dict, cls=NumpyAwareJSONEncoder) 

147 else: 

148 return json.dumps(cjson_dict, cls=JSONIndentEncoder, sort_keys=True, indent=4) 

149 

150 def set_JSON_attribute(self, object, key): 

151 """ 

152 Args: 

153 object: Python dictionary which is being appended with the key value. 

154 key: cclib attribute name. 

155 

156 Returns: 

157 None. The dictionary is modified to contain the attribute with the 

158 cclib keyname as key 

159 """ 

160 if hasattr(self.ccdata, key): 

161 object[ccData._attributes[key].json_key] = getattr(self.ccdata, key) 

162 

163 

164class NumpyAwareJSONEncoder(json.JSONEncoder): 

165 """A encoder for numpy.ndarray's obtained from the cclib attributes. 

166 For all other types the json default encoder is called. 

167 Do Not rename the 'default' method as it is required to be implemented 

168 by any subclass of the json.JSONEncoder 

169 """ 

170 def default(self, obj): 

171 if isinstance(obj, np.ndarray): 

172 if obj.ndim == 1: 

173 nan_list = obj.tolist() 

174 return [None if np.isnan(x) else x for x in nan_list] 

175 else: 

176 return [self.default(obj[i]) for i in range(obj.shape[0])] 

177 return json.JSONEncoder.default(self, obj) 

178 

179 

180class JSONIndentEncoder(json.JSONEncoder): 

181 

182 def __init__(self, *args, **kwargs): 

183 super(JSONIndentEncoder, self).__init__(*args, **kwargs) 

184 self.current_indent = 0 

185 self.current_indent_str = "" 

186 

187 def encode(self, o): 

188 # Special Processing for lists 

189 if isinstance(o, (list, tuple)): 

190 primitives_only = True 

191 for item in o: 

192 if isinstance(item, (list, tuple, dict)): 

193 primitives_only = False 

194 break 

195 output = [] 

196 if primitives_only: 

197 for item in o: 

198 output.append(json.dumps(item, cls=NumpyAwareJSONEncoder)) 

199 return "[ " + ", ".join(output) + " ]" 

200 else: 

201 self.current_indent += self.indent 

202 self.current_indent_str = "".join([" " for x in range(self.current_indent)]) 

203 for item in o: 

204 output.append(self.current_indent_str + self.encode(item)) 

205 self.current_indent -= self.indent 

206 self.current_indent_str = "".join([" " for x in range(self.current_indent)]) 

207 return "[\n" + ",\n".join(output) + "\n" + self.current_indent_str + "]" 

208 elif isinstance(o, dict): 

209 output = [] 

210 self.current_indent += self.indent 

211 self.current_indent_str = "".join([" " for x in range(self.current_indent)]) 

212 for key, value in o.items(): 

213 output.append(self.current_indent_str + json.dumps(key, cls=NumpyAwareJSONEncoder) + ": " + 

214 str(self.encode(value))) 

215 self.current_indent -= self.indent 

216 self.current_indent_str = "".join([" " for x in range(self.current_indent)]) 

217 return "{\n" + ",\n".join(output) + "\n" + self.current_indent_str + "}" 

218 elif isinstance(o, np.generic): 

219 return json.dumps(o.item(), cls=NumpyAwareJSONEncoder) 

220 else: 

221 return json.dumps(o, cls=NumpyAwareJSONEncoder) 

222 

223 

224del find_package