1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
|
"""Functions for in-place manipulation of bundletrajectories.
This module defines a number of functions that can be used to
extract and delete data from BundleTrajectories directly on
disk. The functions are intended for large-scale MD output,
so they avoid copying the potentially large amounts of data.
Instead, data is either directly deleted in-place; or copies
are made by creating a new directory structure, but hardlinking
the data files. Hard links makes it possible to delete the
original data without invalidating the copy.
Usage from command line:
python -m ase.io.bundlemanipulate inbundle outbundle [start [end [step]]]
"""
import os
import json
from typing import Optional
import numpy as np
from ase.io.bundletrajectory import UlmBundleBackend
def copy_frames(inbundle, outbundle, start=0, end=None, step=1,
verbose=False):
"""Copies selected frame from one bundle to the next."""
if not (isinstance(start, int) and
(isinstance(end, int) or end is None) and
isinstance(step, int)):
raise TypeError("copy_frames: start, end and step must be integers.")
metadata, nframes = read_bundle_info(inbundle)
if metadata['backend'] == 'ulm':
backend = UlmBundleBackend(True, metadata['ulm.singleprecision'])
elif metadata['backend'] == 'pickle':
raise IOError("Input BundleTrajectory uses the 'pickle' backend. " +
"This is not longer supported for security reasons")
else:
raise IOError("Unknown backend type '{}'".format(metadata['backend']))
if start < 0:
start += nframes
if end is None:
end = nframes
if end < 0:
end += nframes
if start < 0 or (start > nframes - 1 and end > 0):
raise ValueError("copy_frames: Invalid start value.")
if end < 0 or (end > nframes - 1 and end < 0):
raise ValueError("copy_frames: Invalid end value.")
if step == 0:
raise ValueError("copy_frames: Invalid step value (zero)")
frames = list(range(start, end, step))
if verbose:
print("Copying the frames", frames)
# Make the new bundle directory
os.mkdir(outbundle)
with open(os.path.join(outbundle, 'metadata.json'), 'w') as f:
json.dump(metadata, f, indent=2)
for nout, nin in enumerate(frames):
if verbose:
print("F%i -> F%i" % (nin, nout))
indir = os.path.join(inbundle, "F" + str(nin))
outdir = os.path.join(outbundle, "F" + str(nout))
os.mkdir(outdir)
names = os.listdir(indir)
for name in names:
fromfile = os.path.join(indir, name)
tofile = os.path.join(outdir, name)
os.link(fromfile, tofile)
if nout == 0 and nin != 0:
if verbose:
print("F0 -> F0 (supplemental)")
# The smalldata.ulm stuff must be updated.
# At the same time, check if the number of fragments
# has not changed.
data0 = backend.read_small(os.path.join(inbundle, "F0"))
data1 = backend.read_small(indir)
split_data = (metadata['subtype'] == 'split')
if split_data:
fragments0 = data0['fragments']
fragments1 = data1['fragments']
data0.update(data1) # Data in frame overrides data from frame 0.
backend.write_small(outdir, data0)
# If data is written in split mode, it must be reordered
firstnames = os.listdir(os.path.join(inbundle, "F0"))
if not split_data:
# Simple linking
for name in firstnames:
if name not in names:
if verbose:
print(" ", name, " (linking)")
fromfile = os.path.join(inbundle, "F0", name)
tofile = os.path.join(outdir, name)
os.link(fromfile, tofile)
else:
# Must read and rewrite data
# First we read the ID's from frame 0 and N
assert 'ID_0.ulm' in firstnames and 'ID_0.ulm' in names
backend.nfrag = fragments0
f0_id, dummy = backend.read_split(
os.path.join(inbundle, "F0"), "ID"
)
backend.nfrag = fragments1
fn_id, fn_sizes = backend.read_split(indir, "ID")
for name in firstnames:
# Only look at each array, not each file
if '_0.' not in name:
continue
if name not in names:
# We need to load this array
arrayname = name.split('_')[0]
print(" Reading", arrayname)
backend.nfrag = fragments0
f0_data, dummy = backend.read_split(
os.path.join(inbundle, "F0"), arrayname
)
# Sort data
f0_data[f0_id] = np.array(f0_data)
# Unsort with new ordering
f0_data = f0_data[fn_id]
# Write it
print(" Writing reshuffled", arrayname)
pointer = 0
backend.nfrag = fragments1
for i, s in enumerate(fn_sizes):
segment = f0_data[pointer:pointer + s]
pointer += s
backend.write(outdir, '{}_{}'.format(arrayname, i),
segment)
# Finally, write the number of frames
with open(os.path.join(outbundle, 'frames'), 'w') as fd:
fd.write(str(len(frames)) + '\n')
# Helper functions
def read_bundle_info(name):
"""Read global info about a bundle.
Returns (metadata, nframes)
"""
if not os.path.isdir(name):
raise IOError("No directory (bundle) named '%s' found." % (name,))
metaname = os.path.join(name, 'metadata.json')
if not os.path.isfile(metaname):
if os.path.isfile(os.path.join(name, 'metadata')):
raise IOError("Found obsolete metadata in unsecure Pickle format. Refusing to load.")
else:
raise IOError("'{}' does not appear to be a BundleTrajectory "
"(no {})".format(name, metaname))
with open(metaname) as fd:
mdata = json.load(fd)
if 'format' not in mdata or mdata['format'] != 'BundleTrajectory':
raise IOError("'%s' does not appear to be a BundleTrajectory" %
(name,))
if mdata['version'] != 1:
raise IOError("Cannot manipulate BundleTrajectories with version "
"number %s" % (mdata['version'],))
with open(os.path.join(name, "frames")) as fd:
nframes = int(fd.read())
if nframes == 0:
raise IOError("'%s' is an empty BundleTrajectory" % (name,))
return mdata, nframes
if __name__ == '__main__':
import sys
if len(sys.argv) < 3:
print(__doc__)
sys.exit()
inname, outname = sys.argv[1:3]
if len(sys.argv) > 3:
start = int(sys.argv[3])
else:
start = 0
if len(sys.argv) > 4:
end: Optional[int] = int(sys.argv[4])
else:
end = None
if len(sys.argv) > 5:
step = int(sys.argv[5])
else:
step = 1
copy_frames(inname, outname, start, end, step, verbose=1)
|
