Coverage for sparc/sparc_parsers/ion.py: 97%
100 statements
« prev ^ index » next coverage.py v7.6.9, created at 2024-12-12 01:13 +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 textwrap
11from warnings import warn
13import numpy as np
14from ase.units import Bohr
16# Safe wrappers for both string and fd
17from ase.utils import reader, writer
19from ..api import SparcAPI
20from .utils import (
21 bisect_and_strip,
22 make_reverse_mapping,
23 read_block_input,
24 strip_comments,
25)
28class InvalidSortingComment(ValueError):
29 def __init__(self, message):
30 self.message = message
33defaultAPI = SparcAPI()
36@reader
37def _read_ion(fileobj, validator=defaultAPI):
38 """
39 Read information from the .ion file. Note, this method does not return an atoms object,
40 but rather return a dict. Thus the label option is not necessary to keep
43 Reads an ion file. Because some of the information necessary to create
44 an atoms object is found in the .inpt file, this function also attemtps to read
45 that as a source of data. If the file is not found or the information is invalid,
46 it will look for it in the comments of the ion file, as written.
47 """
48 contents = fileobj.read()
49 # label = get_label(fileobj, ".ion")
50 data, comments = strip_comments(contents)
51 # We do not read the cell at this time!
52 sort, resort, new_comments = _read_sort_comment(comments)
54 # find the index for all atom type lines. They should be at the top of their block
55 atom_type_bounds = [i for i, x in enumerate(data) if "ATOM_TYPE" in x] + [len(data)]
56 atom_blocks = [
57 read_block_input(data[start:end], validator=validator)
58 for start, end in zip(atom_type_bounds[:-1], atom_type_bounds[1:])
59 ]
61 return {
62 "ion": {
63 "atom_blocks": atom_blocks,
64 "comments": new_comments,
65 "sorting": {"sort": sort, "resort": resort},
66 }
67 }
70@writer
71def _write_ion(
72 fileobj,
73 data_dict,
74 validator=defaultAPI,
75):
76 """
77 Writes the ion file content from the atom_dict
79 Please note this is not a Atoms-compatible function!
81 The data_dict takes similar format as _read_ion
83 Basically, we want to ensure
84 data_dict = _read_ion("some.ion")
85 _write_ion("some.ion", data_dict)
86 shows the same format
87 """
88 ion_dict = data_dict.get("ion", None)
89 if ion_dict is None:
90 raise ValueError("No ion data provided in the input!")
91 if "atom_blocks" not in ion_dict:
92 raise ValueError(
93 "Must provide a data-section in the data_dict (blocks of atomic information)"
94 )
96 comments = ion_dict.get("comments", [])
97 banner = "Ion File Generated by SPARC ASE Calculator"
98 if len(comments) == 0:
99 comments = [banner]
100 elif "ASE" not in comments[0]:
101 comments = [banner] + comments
103 # Handle the sorting mapping
104 # the line wrap is 80 words
105 if "sorting" in ion_dict:
106 # print(ion_dict["sorting"])
107 resort = ion_dict["sorting"].get("resort", [])
108 # Write resort information only when it's actually useful
109 if len(resort) > 0:
110 comments.append("ASE-SORT:")
111 index_lines = textwrap.wrap(" ".join(map(str, resort)), width=80)
112 comments.extend(index_lines)
113 comments.append("END ASE-SORT")
115 for line in comments:
116 fileobj.write(f"# {line}\n")
118 fileobj.write("\n")
119 blocks = ion_dict["atom_blocks"]
120 for block in blocks:
121 for key in [
122 "ATOM_TYPE",
123 "N_TYPE_ATOM",
124 "PSEUDO_POT",
125 "COORD_FRAC",
126 "COORD",
127 "SPIN",
128 "RELAX",
129 ]:
130 val = block.get(key, None)
131 # print(key, val)
132 if (key not in ["RELAX", "COORD", "COORD_FRAC", "SPIN"]) and (val is None):
133 raise ValueError(f"Key {key} is not provided! Abort writing ion file")
134 # TODO: change the API version
135 if val is None:
136 continue
138 val_string = validator.convert_value_to_string(key, val)
139 # print(val_string)
140 # TODO: make sure 1 line is accepted
141 # TODO: write pads to vector lines
142 if (val_string.count("\n") > 0) or (
143 key in ["COORD_FRAC", "COORD", "RELAX", "SPIN"]
144 ):
145 output = f"{key}:\n{val_string}\n"
146 else:
147 output = f"{key}: {val_string}\n"
148 fileobj.write(output)
149 # TODO: check extra keys
150 # TODO: how to handle multiple psp files?
151 # Write a split line
152 # TODO: do we need to distinguish the last line?
153 fileobj.write("\n")
154 return
157def _ion_coord_to_ase_pos(data_dict, cell=None):
158 """Convert the COORD or COORD_FRAC from atom blocks to ASE's positions
160 Arguments:
161 cell: a unit cell in ASE-unit (i.e. parsed from inpt._inpt_cell_to_ase_cell)
163 This function modifies the data_dict in-place to add a field '_ase_positions'
164 to the atom_blocks
165 """
166 treated_blocks = []
167 can_have_coord_frac = cell is not None
168 ion_atom_blocks = data_dict["ion"]["atom_blocks"]
169 for i, block in enumerate(ion_atom_blocks):
170 if ("COORD" in block.keys()) and ("COORD_FRAC" in block.keys()):
171 raise KeyError("COORD and COORD_FRAC cannot co-exist!")
172 if (not can_have_coord_frac) and ("COORD_FRAC" in block.keys()):
173 raise KeyError("COORD_FRAC must be acompanied by a cell!")
174 coord = block.get("COORD", None)
175 if coord is not None:
176 coord = coord * Bohr
177 else:
178 coord_frac = block["COORD_FRAC"]
179 # Cell is already in Bohr
180 coord = np.dot(coord_frac, cell)
181 data_dict["ion"]["atom_blocks"][i]["_ase_positions"] = coord
182 return
185def _read_sort_comment(lines):
186 """Parse the atom sorting info from the comment lines
187 Format
189 ASE-SORT:
190 r_i r_j r_k ....
191 END ASE-SORT
192 where r_i etc are the indices in the original ASE atoms object
193 """
194 i = 0
195 resort = []
196 record = False
197 new_lines = []
198 while i < len(lines):
199 line = lines[i]
200 key, value = bisect_and_strip(line, ":")
201 i += 1
202 if key == "ASE-SORT":
203 record = True
204 elif key == "END ASE-SORT":
205 record = False
206 break
207 elif record is True:
208 resort += list(map(int, line.strip().split(" ")))
209 else:
210 # Put original lines in new_lines
211 new_lines.append(line)
212 # Put all remaining lines in new_lines
213 for j in range(i, len(lines)):
214 line = lines[j]
215 if "ASE-SORT" in line:
216 raise InvalidSortingComment(
217 "There appears to be multiple sorting information in the ion comment section!"
218 )
219 new_lines.append(line)
220 if record:
221 warn(
222 "ASE atoms resort comment block is not properly formatted, this may cause data loss!"
223 )
224 sort = make_reverse_mapping(resort)
225 assert set(sort) == set(resort), "Sort and resort info are of different length!"
226 return sort, resort, new_lines