Coverage for sparc/sparc_parsers/ion.py: 94%
151 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-18 16:19 +0000
1"""
2Created on Thu Oct 18 14:16:21 2018
4Ben Comer (Georgia Tech)
6This file has been heavily modified since SPARC 0.1
8TODO: more descriptions about this file io parser
9"""
10import re
11import textwrap
12from warnings import warn
14import numpy as np
15from ase.units import Bohr
17# Safe wrappers for both string and fd
18from ase.utils import reader, writer
20from ..api import SparcAPI
21from .utils import (
22 bisect_and_strip,
23 make_reverse_mapping,
24 read_block_input,
25 strip_comments,
26)
29class InvalidSortingComment(ValueError):
30 def __init__(self, message):
31 self.message = message
34defaultAPI = SparcAPI()
37@reader
38def _read_ion(fileobj, validator=defaultAPI):
39 """
40 Read information from the .ion file. Note, this method does not return an atoms object,
41 but rather return a dict. Thus the label option is not necessary to keep
44 Reads an ion file. Because some of the information necessary to create
45 an atoms object is found in the .inpt file, this function also attemtps to read
46 that as a source of data. If the file is not found or the information is invalid,
47 it will look for it in the comments of the ion file, as written.
48 """
49 contents = fileobj.read()
50 # label = get_label(fileobj, ".ion")
51 data, comments = strip_comments(contents)
52 # We do not read the cell at this time!
53 sort, resort, new_comments = _read_sort_comment(comments)
55 # find the index for all atom type lines. They should be at the top of their block
56 # @TT 2025.06.04 HUBBARD block comes at the end of the
57 # ion file defines the last line in data
58 hubbard_bounds = [i for i, x in enumerate(data) if "HUBBARD" in x]
59 if len(hubbard_bounds) == 0:
60 atom_type_bounds_lim = len(data)
61 elif len(hubbard_bounds) == 1:
62 # The hubbard bounds
63 atom_type_bounds_lim = hubbard_bounds[0]
64 else:
65 # TODO: make it a format error
66 raise ValueError("Bad .ion file format, multiple HUBBARD sections exist.")
68 atom_type_bounds = [i for i, x in enumerate(data) if re.match("^ATOM_TYPE", x)]
69 atom_type_bounds += [atom_type_bounds_lim]
70 atom_blocks = [
71 read_block_input(data[start:end], validator=validator)
72 for start, end in zip(atom_type_bounds[:-1], atom_type_bounds[1:])
73 ]
75 extra_blocks = {}
76 # Now handle hubbard information --> extra
77 # the hubbard block is currently the last to appear in
78 # .ion file
79 if len(hubbard_bounds) == 1:
80 hubbard_settings = _parse_hubbard_block(
81 data[hubbard_bounds[0] :], validator=validator
82 )
83 extra_blocks["hubbard"] = hubbard_settings
85 return {
86 "ion": {
87 "atom_blocks": atom_blocks,
88 "comments": new_comments,
89 "extra": extra_blocks,
90 "sorting": {"sort": sort, "resort": resort},
91 }
92 }
95@writer
96def _write_ion(
97 fileobj,
98 data_dict,
99 validator=defaultAPI,
100):
101 """
102 Writes the ion file content from the atom_dict
104 Please note this is not a Atoms-compatible function!
106 The data_dict takes similar format as _read_ion
108 Basically, we want to ensure
109 data_dict = _read_ion("some.ion")
110 _write_ion("some.ion", data_dict)
111 shows the same format
112 """
113 ion_dict = data_dict.get("ion", None)
114 if ion_dict is None:
115 raise ValueError("No ion data provided in the input!")
116 if "atom_blocks" not in ion_dict:
117 raise ValueError(
118 "Must provide a data-section in the data_dict (blocks of atomic information)"
119 )
121 comments = ion_dict.get("comments", [])
122 banner = "Ion File Generated by SPARC ASE Calculator"
123 if len(comments) == 0:
124 comments = [banner]
125 elif "ASE" not in comments[0]:
126 comments = [banner] + comments
128 # Handle the sorting mapping
129 # the line wrap is 80 words
130 if "sorting" in ion_dict:
131 # print(ion_dict["sorting"])
132 resort = ion_dict["sorting"].get("resort", [])
133 # Write resort information only when it's actually useful
134 if len(resort) > 0:
135 comments.append("ASE-SORT:")
136 index_lines = textwrap.wrap(" ".join(map(str, resort)), width=80)
137 comments.extend(index_lines)
138 comments.append("END ASE-SORT")
140 for line in comments:
141 fileobj.write(f"# {line}\n")
143 fileobj.write("\n")
144 blocks = ion_dict["atom_blocks"]
145 for block in blocks:
146 for key in [
147 "ATOM_TYPE",
148 "N_TYPE_ATOM",
149 "PSEUDO_POT",
150 "COORD_FRAC",
151 "COORD",
152 "SPIN",
153 "RELAX",
154 ]:
155 val = block.get(key, None)
156 # print(key, val)
157 if (key not in ["RELAX", "COORD", "COORD_FRAC", "SPIN"]) and (val is None):
158 raise ValueError(f"Key {key} is not provided! Abort writing ion file")
159 # TODO: change the API version
160 if val is None:
161 continue
163 val_string = validator.convert_value_to_string(key, val)
164 # print(val_string)
165 # TODO: make sure 1 line is accepted
166 # TODO: write pads to vector lines
167 if (val_string.count("\n") > 0) or (
168 key in ["COORD_FRAC", "COORD", "RELAX", "SPIN"]
169 ):
170 output = f"{key}:\n{val_string}\n"
171 else:
172 output = f"{key}: {val_string}\n"
173 fileobj.write(output)
174 # TODO: check extra keys
175 # TODO: how to handle multiple psp files?
176 # Write a split line
177 # TODO: do we need to distinguish the last line?
178 fileobj.write("\n")
180 # @TT 2025.06.04 add support for HUBBARD parameters
181 extra_blocks = ion_dict.get("extra", {})
182 if "hubbard" in extra_blocks:
183 _check_hubbard_block(ion_dict, extra_blocks["hubbard"])
184 _write_hubbard_block(fileobj, extra_blocks["hubbard"], validator)
185 return
188def _ion_coord_to_ase_pos(data_dict, cell=None):
189 """Convert the COORD or COORD_FRAC from atom blocks to ASE's positions
191 Arguments:
192 cell: a unit cell in ASE-unit (i.e. parsed from inpt._inpt_cell_to_ase_cell)
194 This function modifies the data_dict in-place to add a field '_ase_positions'
195 to the atom_blocks
196 """
197 treated_blocks = []
198 can_have_coord_frac = cell is not None
199 ion_atom_blocks = data_dict["ion"]["atom_blocks"]
200 for i, block in enumerate(ion_atom_blocks):
201 if ("COORD" in block.keys()) and ("COORD_FRAC" in block.keys()):
202 raise KeyError("COORD and COORD_FRAC cannot co-exist!")
203 if (not can_have_coord_frac) and ("COORD_FRAC" in block.keys()):
204 raise KeyError("COORD_FRAC must be acompanied by a cell!")
205 coord = block.get("COORD", None)
206 if coord is not None:
207 coord = coord * Bohr
208 else:
209 coord_frac = block["COORD_FRAC"]
210 # Cell is already in Bohr
211 coord = np.dot(coord_frac, cell)
212 data_dict["ion"]["atom_blocks"][i]["_ase_positions"] = coord
213 return
216def _read_sort_comment(lines):
217 """Parse the atom sorting info from the comment lines
218 Format
220 ASE-SORT:
221 r_i r_j r_k ....
222 END ASE-SORT
223 where r_i etc are the indices in the original ASE atoms object
224 """
225 i = 0
226 resort = []
227 record = False
228 new_lines = []
229 while i < len(lines):
230 line = lines[i]
231 key, value = bisect_and_strip(line, ":")
232 i += 1
233 if key == "ASE-SORT":
234 record = True
235 elif key == "END ASE-SORT":
236 record = False
237 break
238 elif record is True:
239 resort += list(map(int, line.strip().split(" ")))
240 else:
241 # Put original lines in new_lines
242 new_lines.append(line)
243 # Put all remaining lines in new_lines
244 for j in range(i, len(lines)):
245 line = lines[j]
246 if "ASE-SORT" in line:
247 raise InvalidSortingComment(
248 "There appears to be multiple sorting information in the ion comment section!"
249 )
250 new_lines.append(line)
251 if record:
252 warn(
253 "ASE atoms resort comment block is not properly formatted, this may cause data loss!"
254 )
255 sort = make_reverse_mapping(resort)
256 assert set(sort) == set(resort), "Sort and resort info are of different length!"
257 return sort, resort, new_lines
260def _parse_hubbard_block(block, validator=defaultAPI):
261 """Parse the hubbard blocks into the following list
262 [{"U_ATOM_TYPE": <atom-name>,
263 "U_VAL": array},
264 ]
266 A hubbard block (after stripping the extra comments) may look like:
267 ['HUBBARD:'
268 'U_ATOM_TYPE: Ni',
269 'U_VAL: 0 0 0.05 0']
271 The U_ATOM_TYPE and U_VAL must come in pairs ordered
272 """
273 if "HUBBARD:" not in block[0]:
274 raise ValueError("Ill-formatted HUBBARD block in .ion file!")
275 if (len(block) - 1) % 2 != 0:
276 raise ValueError("U_ATOM_TYPE and U_VAL are not paired in the HUBBARD block!")
277 u_pairs = []
278 for i in range((len(block) - 1) // 2):
279 u_sub_block = block[i + 1 : i + 3]
280 u_dict = read_block_input(u_sub_block, validator)
281 u_pairs.append(u_dict)
282 return u_pairs
285def _check_hubbard_block(ion_dict, hubbard_u_pairs):
286 """Sanity check for hubbard parameters
287 1. U_ATOM_TYPE must match one existing element
288 2. No duplicated element of U_ATOM_TYPE
289 3. U value must be 4-tuples
290 """
291 structure_elements = set([entry["ATOM_TYPE"] for entry in ion_dict["atom_blocks"]])
292 hubbard_elements = set()
293 for pair in hubbard_u_pairs:
294 elem = pair["U_ATOM_TYPE"]
295 if elem not in structure_elements:
296 raise ValueError(
297 f"Element {elem} in the HUBBARD setting does not exist in the input structure!"
298 )
299 if elem in hubbard_elements:
300 raise ValueError(f"Element {elem} is duplicated in the HUBBARD setting!")
301 hubbard_elements.add(elem)
302 val = pair["U_VAL"]
303 if len(val) != 4:
304 raise ValueError(f"U_VAL for element {elem} must have length of 4!")
307def _write_hubbard_block(fileobj, u_pairs=[], validator=defaultAPI):
308 """Write the HUBBARD U-blocks at the end of the .ion file
309 format
311 HUBBARD:
312 U_ATOM_TYPE: Ni
313 U_VAL: 0 0 0.05 0
315 U_ATOM_TYPE: Cr
316 U_VAL: 0 0 0.05 0
317 """
318 if len(u_pairs) == 0:
319 return
320 fileobj.write("HUBBARD:\n")
321 for u_pair in u_pairs:
322 # TODO: add value checker
323 for key in ("U_ATOM_TYPE", "U_VAL"):
324 val = u_pair[key]
325 val_string = validator.convert_value_to_string(key, val)
326 fileobj.write(f"{key}: {val_string}\n")
327 fileobj.write("\n")
328 return