fennol.utils.io

  1import time
  2import numpy as np
  3from . import cell_vectors_to_lengths_angles, parse_cell
  4
  5def xyz_reader(
  6    filename,
  7    has_comment_line=True,
  8    indexed=False,
  9    start=1,
 10    stop=-1,
 11    step=1,
 12    max_frames=None,
 13    stream=False,
 14    interval=2.0,
 15    sleep=time.sleep,
 16):
 17    if stop > 0 and start > stop:
 18        return
 19    pf = open(filename, "r")
 20    inside_frame = False
 21    nat_read = False
 22    iframe = 0
 23    stride = step
 24    nframes = 0
 25    if start > 1:
 26        print(f"skipping {start-1} frames")
 27    if not has_comment_line:
 28        comment_line = ""
 29    while True:
 30        line = pf.readline()
 31
 32        if not line:
 33            if stream:
 34                sleep(interval)
 35                continue
 36            elif inside_frame or nat_read:
 37                raise Exception("Error: premature end of file!")
 38            else:
 39                break
 40
 41        line = line.strip()
 42
 43        if not inside_frame:
 44            if not nat_read:
 45                if line.startswith("#") or line == "":
 46                    continue
 47                nat = int(line.split()[0])
 48                nat_read = True
 49                iat = 0
 50                at_count = 0
 51                inside_frame = not has_comment_line
 52                xyz = np.zeros((nat, 3))
 53                symbols = []
 54                continue
 55
 56            # box = np.array(line.split(), dtype="float")
 57            comment_line = line
 58            inside_frame = True
 59            continue
 60
 61        if line.startswith("#") or line == "":
 62            raise Exception("Error: premature end of frame!")
 63
 64        if not nat_read:
 65            raise Exception("Error: nat not read!")
 66        ls = line.split()
 67        at_count += 1
 68        try:
 69            iat,s = (int(ls[0]),1)
 70            indexed = True
 71        except ValueError:
 72            if indexed:
 73                raise Exception("Error: line does not start with an integer index!")
 74            iat, s = (iat+1, 0)
 75        # iat, s = (int(ls[0]), 1) if indexed else (iat + 1, 0)
 76        symbols.append(ls[s])
 77        xyz[iat - 1, :] = np.array([ls[s + 1], ls[s + 2], ls[s + 3]], dtype="float")
 78        if at_count == nat:
 79            iframe += 1
 80            inside_frame = False
 81            nat_read = False
 82            if iframe < start:
 83                continue
 84            stride += 1
 85            if stride >= step:
 86                nframes += 1
 87                stride = 0
 88                yield symbols, xyz, comment_line
 89            if (max_frames is not None and nframes >= max_frames) or (
 90                stop > 0 and iframe >= stop
 91            ):
 92                break
 93
 94
 95def read_xyz(
 96    filename,
 97    has_comment_line=True,
 98    indexed=False,
 99    start=1,
100    stop=-1,
101    step=1,
102    max_frames=None,
103):
104    return [
105        frame
106        for frame in xyz_reader(
107            filename, has_comment_line, indexed, start, stop, step, max_frames
108        )
109    ]
110
111
112def last_xyz_frame(filename, has_comment_line=True, indexed=False):
113    last_frame = None
114    for frame in xyz_reader(
115        filename, has_comment_line, indexed, start=-1, stop=-1, step=1, max_frames=1
116    ):
117        last_frame = frame
118    return last_frame
119
120
121def write_arc_frame(
122    f,
123    symbols,
124    coordinates,
125    types=None,
126    nbonds=None,
127    connectivity=None,
128    cell=None,
129    **kwargs,
130):
131    nat = len(symbols)
132    f.write(f"{nat}\n")
133    if cell is not None:
134        box = cell_vectors_to_lengths_angles(parse_cell(cell))
135        f.write(" ".join([f"{x: 15.5f}" for x in box]) + "\n")
136    # f.write(f'{axis} {axis} {axis} 90.0 90.0 90.0 \n')
137    for i in range(nat):
138        line = f"{i+1} {symbols[i]:3} {coordinates[i,0]: 15.3f} {coordinates[i,1]: 15.3f} {coordinates[i,2]: 15.3f}"
139        if types is not None:
140            line += f"   {types[i]}"
141        if connectivity is not None and nbonds is not None:
142            line += "  " + " ".join([str(x + 1) for x in connectivity[i, : nbonds[i]]])
143        f.write(line + "\n")
144    f.flush()
145
146
147def write_extxyz_frame(
148    f, symbols, coordinates, cell=None, properties={}, forces=None, **kwargs
149):
150    nat = len(symbols)
151    f.write(f"{nat}\n")
152    comment_line = ""
153    if cell is not None:
154        cell = parse_cell(cell)
155        comment_line += (
156            'Lattice="' + " ".join([f"{x:.3f}" for x in cell.flatten()]) + '" '
157        )
158    comment_line += "Properties=species:S:1:pos:R:3"
159    if forces is not None:
160        comment_line += ":forces:R:3"
161    comment_line += " "
162    for k, v in properties.items():
163        comment_line += f"{k}={v} "
164    f.write(f"{comment_line}\n")
165    for i in range(nat):
166        line = f"{symbols[i]:3} {coordinates[i,0]: 15.5e} {coordinates[i,1]: 15.5e} {coordinates[i,2]: 15.5e}"
167        if forces is not None:
168            line += f" {forces[i,0]: 15.5e} {forces[i,1]: 15.5e} {forces[i,2]: 15.5e}"
169        f.write(f"{line}\n")
170    f.flush()
171
172
173def write_xyz_frame(f, symbols, coordinates,cell=None, **kwargs):
174    nat = len(symbols)
175    f.write(f"{nat}\n")
176    if cell is not None:
177        box = cell_vectors_to_lengths_angles(parse_cell(cell))
178        f.write(" ".join([f"{x:.3f}" for x in box]))
179    f.write("\n")
180    for i in range(nat):
181        f.write(
182            f"{symbols[i]:3} {coordinates[i,0]: 15.5e} {coordinates[i,1]: 15.5e} {coordinates[i,2]: 15.5e}\n"
183        )
184    f.flush()
185
186
187def human_time_duration(seconds: float):
188    """Convert seconds (duration) to human readable string
189
190    from https://gist.github.com/borgstrom/936ca741e885a1438c374824efb038b3
191    """
192
193    if seconds < 1.0:
194        return f"{seconds*1000:.3g} ms"
195    if seconds < 10.0:
196        return f"{seconds:.3g} s"
197
198    TIME_DURATION_UNITS = (
199        ("week", "s", 60 * 60 * 24 * 7),
200        ("day", "s", 60 * 60 * 24),
201        ("h", "", 60 * 60),
202        ("min", "", 60),
203        ("s", "", 1),
204    )
205    parts = []
206    for unit, plur, div in TIME_DURATION_UNITS:
207        amount, seconds = divmod(int(seconds), div)
208        if amount > 0:
209            parts.append(f"{amount} {unit}{plur if amount > 1 else ''}")
210    return " ".join(parts)