fennol.utils.io

  1import time
  2import numpy as np
  3
  4
  5def xyz_reader(
  6    filename,
  7    has_comment_line=False,
  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                inside_frame = not has_comment_line
 51                xyz = np.zeros((nat, 3))
 52                symbols = []
 53                continue
 54
 55            # box = np.array(line.split(), dtype="float")
 56            comment_line = line
 57            inside_frame = True
 58            continue
 59
 60        if line.startswith("#") or line == "":
 61            raise Exception("Error: premature end of frame!")
 62
 63        if not nat_read:
 64            raise Exception("Error: nat not read!")
 65        ls = line.split()
 66        iat, s = (int(ls[0]), 1) if indexed else (iat + 1, 0)
 67        symbols.append(ls[s])
 68        xyz[iat - 1, :] = np.array([ls[s + 1], ls[s + 2], ls[s + 3]], dtype="float")
 69        if iat == nat:
 70            iframe += 1
 71            inside_frame = False
 72            nat_read = False
 73            if iframe < start:
 74                continue
 75            stride += 1
 76            if stride >= step:
 77                nframes += 1
 78                stride = 0
 79                yield symbols, xyz, comment_line
 80            if (max_frames is not None and nframes >= max_frames) or (
 81                stop > 0 and iframe >= stop
 82            ):
 83                break
 84
 85
 86def read_xyz(
 87    filename,
 88    has_comment_line=False,
 89    indexed=False,
 90    start=1,
 91    stop=-1,
 92    step=1,
 93    max_frames=None,
 94):
 95    return [
 96        frame
 97        for frame in xyz_reader(
 98            filename, has_comment_line, indexed, start, stop, step, max_frames
 99        )
100    ]
101
102
103def last_xyz_frame(filename, has_comment_line=False, indexed=False):
104    last_frame = None
105    for frame in xyz_reader(
106        filename, has_comment_line, indexed, start=-1, stop=-1, step=1, max_frames=1
107    ):
108        last_frame = frame
109    return last_frame
110
111
112def write_arc_frame(
113    f,
114    symbols,
115    coordinates,
116    types=None,
117    nbonds=None,
118    connectivity=None,
119    cell=None,
120    **kwargs,
121):
122    nat = len(symbols)
123    f.write(f"{nat}\n")
124    if cell is not None:
125        f.write(" ".join([f"{x: 15.5f}" for x in cell.flatten()]) + "\n")
126    # f.write(f'{axis} {axis} {axis} 90.0 90.0 90.0 \n')
127    for i in range(nat):
128        line = f"{i+1} {symbols[i]:3} {coordinates[i,0]: 15.3f} {coordinates[i,1]: 15.3f} {coordinates[i,2]: 15.3f}"
129        if types is not None:
130            line += f"   {types[i]}"
131        if connectivity is not None and nbonds is not None:
132            line += "  " + " ".join([str(x + 1) for x in connectivity[i, : nbonds[i]]])
133        f.write(line + "\n")
134    f.flush()
135
136
137def write_extxyz_frame(
138    f, symbols, coordinates, cell=None, properties={}, forces=None, **kwargs
139):
140    nat = len(symbols)
141    f.write(f"{nat}\n")
142    comment_line = ""
143    if cell is not None:
144        comment_line += (
145            'Lattice="' + " ".join([f"{x:.3f}" for x in cell.flatten()]) + '" '
146        )
147    comment_line += "Properties=species:S:1:pos:R:3"
148    if forces is not None:
149        comment_line += ":forces:R:3"
150    comment_line += " "
151    for k, v in properties.items():
152        comment_line += f"{k}={v} "
153    f.write(f"{comment_line}\n")
154    for i in range(nat):
155        line = f"{symbols[i]:3} {coordinates[i,0]: 15.5e} {coordinates[i,1]: 15.5e} {coordinates[i,2]: 15.5e}"
156        if forces is not None:
157            line += f" {forces[i,0]: 15.5e} {forces[i,1]: 15.5e} {forces[i,2]: 15.5e}"
158        f.write(f"{line}\n")
159    f.flush()
160
161
162def write_xyz_frame(f, symbols, coordinates,cell=None, **kwargs):
163    nat = len(symbols)
164    f.write(f"{nat}\n")
165    if cell is not None:
166        f.write(" ".join([f"{x:.3f}" for x in cell.flatten()]))
167    f.write("\n")
168    for i in range(nat):
169        f.write(
170            f"{symbols[i]:3} {coordinates[i,0]: 15.5e} {coordinates[i,1]: 15.5e} {coordinates[i,2]: 15.5e}\n"
171        )
172    f.flush()
173
174
175def human_time_duration(seconds: float):
176    """Convert seconds (duration) to human readable string
177
178    from https://gist.github.com/borgstrom/936ca741e885a1438c374824efb038b3
179    """
180
181    if seconds < 1.0:
182        return f"{seconds*1000:.3g} ms"
183    if seconds < 10.0:
184        return f"{seconds:.3g} s"
185
186    TIME_DURATION_UNITS = (
187        ("week", "s", 60 * 60 * 24 * 7),
188        ("day", "s", 60 * 60 * 24),
189        ("h", "", 60 * 60),
190        ("min", "", 60),
191        ("s", "", 1),
192    )
193    parts = []
194    for unit, plur, div in TIME_DURATION_UNITS:
195        amount, seconds = divmod(int(seconds), div)
196        if amount > 0:
197            parts.append(f"{amount} {unit}{plur if amount > 1 else ''}")
198    return " ".join(parts)