def test_readThenWrite(self, suffix):
records = {
'test.txt': 'test string foo\n',
'blah.txt': 'another string',
'frames/0/otherblah.f32.ind': np.random.rand(5),
'frames/0/anotherblah.f64.ind': np.random.rand(7),
'frames/0/number.i32.uni': 14
}
convertedRecords = {
'frames/0/otherblah.f32.ind':
np.array(records['frames/0/otherblah.f32.ind'], dtype=np.float32)
}
with gtar.GTAR('test' + suffix, 'w') as arch:
for path in records:
arch.writePath(path, records[path])
for mode in ['r', 'a']:
with gtar.GTAR('test' + suffix, mode) as arch:
for path in records:
self.assertTrue(
np.all(
arch.readPath(path) == convertedRecords.get(
path, records[path])))
def test_write_readonly(self, suffix):
with gtar.GTAR('test' + suffix, 'w') as arch:
arch.writeStr('test.txt', 'good')
with gtar.GTAR('test' + suffix, 'r') as arch:
with self.assertRaises(RuntimeError):
arch.writeStr('test.txt', 'bad')
def test_overwrite(self, suffix):
with gtar.GTAR('test' + suffix, 'w') as arch:
arch.writeStr('test.txt', 'bad')
arch.writeStr('test.txt', 'good')
self.assertEqual(arch.readStr('test.txt'), 'good')
with gtar.GTAR('test' + suffix, 'r') as arch:
self.assertEqual(arch.readStr('test.txt'), 'good')
def test_group_selection(self, suffix):
if suffix == '/':
raise unittest.SkipTest(
'directory support is currently experimental'
' and more robust group parsing needs to be added')
with gtar.GTAR('test' + suffix, 'w') as arch:
arch.writeStr('prefix/test.txt', 'sample text')
with gtar.GTAR('test' + suffix, 'r') as arch:
self.assertEqual(1,
len(arch.framesWithRecordsNamed('test.txt')[1]))
self.assertEqual(
1,
len(
arch.framesWithRecordsNamed('test.txt',
group='prefix')[1]))
self.assertEqual(
1,
len(
arch.framesWithRecordsNamed('test.txt',
group_prefix='p')[1]))
self.assertEqual(
0,
len(
arch.framesWithRecordsNamed('test.txt',
group_prefix='t')[1]))
self.assertEqual(
0,
len(arch.framesWithRecordsNamed('test.txt', group='prefi')[1]))
self.assertEqual(1, len(list(arch.recordsNamed('test.txt'))))
self.assertEqual(
1, len(list(arch.recordsNamed('test.txt', group='prefix'))))
self.assertEqual(
1, len(list(arch.recordsNamed('test.txt', group_prefix='p'))))
self.assertEqual(
0, len(list(arch.recordsNamed('test.txt', group_prefix='t'))))
self.assertEqual(
0, len(list(arch.recordsNamed('test.txt', group='prefi'))))
self.assertEqual(
gtar.Record('prefix/test.txt'),
arch.staticRecordNamed('test.txt', group='prefix'))
self.assertEqual(
gtar.Record('prefix/test.txt'),
arch.staticRecordNamed('test.txt', group_prefix='p'))
def readData(id):
# read the zip file with gtar
fname = 'dump-eql-' + id + '.zip'
g = gtar.GTAR(fname, 'r')
types = g.staticRecordNamed('type')
#nParticles = len(list(g.recordsNamed('position'))[-1][1])
nFrames = len(list(g.recordsNamed('position')))
#begFr = 0
endFr = nFrames
# handles for file i/o
print('Unpacking frame ' + str(endFr))
# extract the data using gtar
pList = list(g.recordsNamed('position'))[endFr - 1][1]
vList = list(g.recordsNamed('velocity'))[endFr - 1][1]
# grab velocity by type, ignore walls
v_S1 = vList[types == 0]
#v_S2 = vList[types == 1]
# grab position by type, ignore walls
p_S1 = pList[types == 0]
#p_S2 = pList[types == 1]
v = list(map(list, list(zip(*v_S1)))) #transpose the list
p = list(map(list, list(zip(*p_S1))))
out = [p[2], v[0]]
out = list(map(list, list(zip(*out))))
index, value = max(enumerate(out[1]), key=operator.itemgetter(1))
v_max = [out[0][index], out[1][index]]
return out, v_max
def get_model(self, scope, storage):
model = scope['model']
filename = scope.get('filename', 'dump.zip')
optimizer_kwargs = scope.setdefault('optimizer_kwargs', {})
try:
with contextlib.ExitStack() as stack:
f = stack.enter_context(
storage.open(filename, 'rb', on_filesystem=True))
with gtar.GTAR(f.name, 'r') as traj:
if traj.readStr('end_reason.txt'):
scope['last_epoch'] = self.arguments['epochs']
return model
for (_, lr_array) in traj.recordsNamed('lr'):
optimizer_kwargs['lr'] = lr_array[-1]
traj = stack.enter_context(keras_gtar.Trajectory(f.name))
last_frame = int(traj.frames[-1])
new_model = traj.load()
# reuse the original model object to make sure the
# architecture did not change
model.set_weights(new_model.get_weights())
scope['last_epoch'] = last_frame
except FileNotFoundError:
scope['last_epoch'] = -1
return model
def test_readAndWrite(self, suffix):
records = {
'test.txt': 'test string foo\n',
'blah.txt': 'another string',
'frames/0/otherblah.f32.ind': np.random.rand(5),
'frames/0/anotherblah.f64.ind': np.random.rand(7),
'frames/0/number.i32.uni': 14
}
convertedRecords = {
'frames/0/otherblah.f32.ind':
np.array(records['frames/0/otherblah.f32.ind'], dtype=np.float32)
}
written = []
with gtar.GTAR('test' + suffix, 'w') as arch:
for path in records:
arch.writePath(path, records[path])
written.append(path)
for readpath in sorted(written):
self.assertTrue(
np.all(
arch.readPath(readpath) == convertedRecords.get(
readpath, records[readpath])))
def readData(id):
# read the zip file with gtar
fname = 'dump-eql-' + id + '.zip'
g = gtar.GTAR(fname, 'r')
types = g.staticRecordNamed('type')
#nParticles = len(list(g.recordsNamed('position'))[-1][1])
nFrames = len(list(g.recordsNamed('position')))
#begFr = 0
endFr = nFrames
# handles for file i/o
print('Unpacking frame ' + str(endFr))
# extract the data using gtar. Data in the last frame
pList = list(g.recordsNamed('position'))[endFr - 1][1]
vList = list(g.recordsNamed('velocity'))[endFr - 1][1]
# grab velocity by type, ignore walls
v_S1 = vList[types == 0]
v_S2 = vList[types == 1]
# grab position by type, ignore walls
p_S1 = pList[types == 0]
p_S2 = pList[types == 1]
v = map(list, zip(*v_S1)) #transpose the list
v1 = map(list, zip(*v_S2))
p = map(list, zip(*p_S1))
p1 = map(list, zip(*p_S2))
#out = [p[2],v[0]]
out = [p1[0], p1[1], p1[2], v1[0], v1[1], v1[2]]
out = map(list, zip(*out))
out1 = [p[0], p[1], p[2], v[0], v[1],
v[2]] #0 stands for x direction, 2 stands for z direction
out1 = map(list, zip(*out1))
#index, value = max(enumerate(out[2]), key=operator.itemgetter(1))
#v_max=[out[1][index],out[2][index]]
return out, out1 #,v_max
def setup_sample(self, N, dim=3):
self.position = np.random.rand(N, 3)
self.orientation = np.random.rand(N, 4)
self.velocity = np.random.rand(N, 3)
self.mass = np.random.rand(N)
self.charge = np.random.rand(N)
self.diameter = np.random.rand(N)
self.moment_inertia = np.random.rand(N, 3)
self.angmom = np.random.rand(N, 4)
self.image = np.random.randint(-1000,
1000,
size=(N, 3),
dtype=np.int32)
self.types = ['A', 'B']
self.typeid = N // 2 * [0] + (N - N // 2) * [1]
self.box = np.array([1.0, 1.0, 1.0, 0.0, 0.0, 0.0])
if dim == 2:
self.position[:, 2] = 0
self.velocity[:, 2] = 0
with gtar.GTAR(self.getar_file_fn, 'w') as traj:
traj.writePath('frames/0/position.f32.ind', self.position)
traj.writePath('frames/0/orientation.f32.ind', self.orientation)
traj.writePath('frames/0/velocity.f32.ind', self.velocity)
traj.writePath('frames/0/mass.f32.ind', self.mass)
traj.writePath('frames/0/charge.f32.ind', self.charge)
traj.writePath('frames/0/diameter.f32.ind', self.diameter)
traj.writePath('frames/0/moment_inertia.f32.ind',
self.moment_inertia)
traj.writePath('frames/0/angular_momentum_quat.f32.ind',
self.angmom)
traj.writePath('frames/0/box.f32.ind', self.box)
traj.writePath('frames/0/image.i32.ind', self.image)
traj.writePath('type.u32.ind', self.typeid)
traj.writePath('type_names.json', json.dumps(self.types))
def read(self, stream, default_type='A', default_box=None):
"""Read binary stream and return a trajectory instance.
:param stream: The stream, which contains the GeTarFile.
:type stream: A file-like binary stream.
:param default_type: The default particle type for
posfile dialects without type definition.
:type default_type: str
:param default_box: The default_box value is used if no
box is specified in the libgetar file.
Defaults to [Lx=Ly=Lz=1.0].
:type default_box: :class:`numpy.ndarray`
"""
if default_box is None:
default_box = np.diag([1.0] * 3)
try:
filename = stream.name
except AttributeError:
raise NotImplementedError(
"The current implementation of the GeTarFileReader requires "
"file objects with name attribute, such as NamedTemporaryFile "
"as the underlying library is reading the file by filename "
"and not directly from the stream.")
_trajectory = gtar.GTAR(filename, 'r')
_records = {rec.getName(): rec for rec in _trajectory.getRecordTypes() if not rec.getGroup()}
# assume that we care primarily about positions
try:
self._frames = _trajectory.queryFrames(_records['position'])
except KeyError:
raise RuntimeError("Given trajectory '{}' contained no "
"positions.".format(stream))
frames = [GetarFrame(_trajectory, _records, idx, default_type, default_box)
for idx in self._frames]
logger.info("Read {} frames.".format(len(frames)))
return Trajectory(frames)
def test_readAndWriteRecords(self, suffix):
with gtar.GTAR('test' + suffix, 'w') as arch:
reclen = len(arch.getRecordTypes())
for idx in range(10):
arch.writeStr(str(idx), '{}_contents'.format(idx))
self.assertEqual(reclen + 1, len(arch.getRecordTypes()))
reclen = len(arch.getRecordTypes())
def main(inputs, output):
"""Take all records from a set of getar-formatted files and output them to another.
:param inputs: Input filenames to concatenate
:param output: Output filename (can be the same as input)
"""
nameHalves = os.path.splitext(output)
tempName = output
while os.path.exists(tempName):
nameHalves = (nameHalves[0] + '_', nameHalves[1])
tempName = nameHalves[0] + nameHalves[1]
recordFiles = {}
for input in inputs:
with gtar.GTAR(input, 'r') as inpFile:
recs = {
rec: inpFile.queryFrames(rec)
for rec in inpFile.getRecordTypes()
}
for rec in recs:
frames = recs[rec]
for frame in frames:
rec.setIndex(frame)
recordFiles[rec.getPath()] = input
sourceRecords = defaultdict(list)
for path in recordFiles:
sourceRecords[recordFiles[path]].append(path)
try:
with gtar.GTAR(tempName, 'w') as out:
for source in sourceRecords:
allPaths = sorted(sourceRecords[source])
with gtar.GTAR(source, 'r') as inp:
for path in allPaths:
out.writePath(path, inp.readPath(path))
if not os.path.samefile(tempName, output):
os.rename(tempName, output)
finally:
if os.path.exists(tempName) and not os.path.samefile(tempName, output):
os.remove(tempName)
def test_widths():
with gtar.GTAR('test.zip', 'w') as arch:
arch.writePath('position.f32.ind', np.random.rand(6))
with gtar.GTAR('test.zip', 'r') as arch:
posx3 = arch.readPath('position.f32.ind')
gtar.widths['position'] = 2
posx2 = arch.readPath('position.f32.ind')
gtar.widths.clear()
posFlat = arch.readPath('position.f32.ind')
success = posx3.shape == (2, 3) and posx2.shape == (
3, 2) and posFlat.shape == (6, )
assert success
return success
def __init__(self, target, group=None):
try:
import gtar
except ImportError:
raise ImportError('libgetar must be installed to use GetarStorage')
self.target = target
self.group = group
self.gtar_file = gtar.GTAR(self.target, 'a')
def test_read_closed(self, suffix):
with gtar.GTAR('test' + suffix, 'w') as arch:
arch.writeStr('test.txt', 'good')
# throw when reading from a closed archive
with self.assertRaises(RuntimeError):
arch.readStr('test.txt')
# readPath goes through getRecord
with self.assertRaises(RuntimeError):
arch.readPath('test.txt')
def main(inputs):
archives = [gtar.GTAR(path, 'r') for path in inputs]
if len(inputs) == 1:
traj = archives[0]
definitions['traj'] = 'the opened input file'
records = {rec.getName(): rec for rec in traj.getRecordTypes()}
definitions['records'] = 'a dictionary of available records, indexed by property name'
recordFrames = {name: traj.queryFrames(records[name]) for name in records}
definitions['recordFrames'] = 'a dictionary of available frames for each property by name'
message = messageTemplate.format(
'\n'.join('- {} is {}.'.format(k, definitions[k]) for k in sorted(definitions)))
code.interact(banner=message, local=dict(globals(), **locals()))
def setUp(self):
self.tmp_dir = TemporaryDirectory(prefix='garnett_getar_tmp')
self.addCleanup(self.tmp_dir.cleanup)
self.getar_file_fn = os.path.join(self.tmp_dir.name, 'sample.tar')
with gtar.GTAR(self.getar_file_fn, 'w') as traj:
for frame in range(6):
traj.writePath('frames/{}/position.f32.ind'.format(frame),
[(0, 0, 0)])
if frame % 2:
traj.writePath('frames/{}/box.f32.uni'.format(frame),
[frame + 1, 1, 1, 0, 0, 0])
reader = garnett.reader.GetarFileReader()
self.getarfile = open(self.getar_file_fn, 'rb')
self.addCleanup(self.getarfile.close)
self.trajectory = reader.read(self.getarfile)
def save_tune_results(self, scope, storage, context):
import gtar
# step count before this stage
beginning_steps = scope['cumulative_steps'] - self.arguments['steps']
types = context.snapshot.particles.types
if scope['mpi_rank']:
return
integrator = scope['integrator']
translation, rotation = [], []
for t in types:
translation.append(integrator.get_d(t))
rotation.append(integrator.get_a(t))
updater = scope.get('box_updater', None)
box_distances = {}
if updater is not None:
for name in BOX_MOVE_NAMES:
box_distances[name] = getattr(updater, name)()['delta']
dump_filename = scope.get('dump_filename', 'dump.sqlite')
msg = 'Dumping tuned move distances to {}: translation {}, rotation {}'.format(
dump_filename, translation, rotation)
logger.debug(msg)
local_context = contextlib.ExitStack()
with local_context:
dump_file = local_context.enter_context(storage.open(
dump_filename, 'ab', on_filesystem=True, noop=scope['mpi_rank']))
getar_file = local_context.enter_context(gtar.GTAR(dump_file.name, 'a'))
path = 'hpmc/frames/{}/type_translation_distance.f32.uni'.format(beginning_steps)
getar_file.writePath(path, translation)
path = 'hpmc/frames/{}/type_rotation_distance.f32.uni'.format(beginning_steps)
getar_file.writePath(path, rotation)
for (key, value) in box_distances.items():
path = 'hpmc/frames/{}/box_{}_distance.f32.uni'.format(
beginning_steps, key)
getar_file.writePath(path, value)
def writeGetar(self, filename='test.zip'):
with gtar.GTAR(filename, 'w') as traj:
traj.writePath('type_names.json', json.dumps(self.typeNames))
traj.writePath('type.u32.ind', self.types)
for name in [
'position', 'velocity', 'acceleration', 'mass', 'charge',
'diameter'
]:
if getattr(self, name):
traj.writePath('{}.f32.ind'.format(name),
getattr(self, name))
traj.writePath('box.f32.uni', self.box)
if self.bonds:
traj.writePath('bond/type_names.json',
json.dumps(self.bondNames))
traj.writePath('bond/type.u32.ind',
[t for (t, _, _) in self.bonds])
traj.writePath('bond/tag.u32.ind',
[(l, r) for (_, l, r) in self.bonds])
def _get_traj(self, filename, storage):
for handle in (self._gtar_traj, self._gtar_file):
if handle is not None:
handle.close()
try:
gtar_file = gtar_traj = None
gtar_file = storage.open(self.arguments['filename'],
'rb',
on_filesystem=True)
gtar_traj = gtar.GTAR(gtar_file.name, 'r')
self._gtar_file, self._gtar_traj = gtar_file, gtar_traj
gtar_file = gtar_traj = None
finally:
for handle in (gtar_traj, gtar_file):
if handle is not None:
handle.close()
return self._gtar_traj
def run(self, scope, storage):
if 'seed' in self.arguments:
s = self.arguments['seed']
random.seed(s)
random.seed(random.randrange(2**32))
np.random.seed(random.randrange(2**32))
tf.random.set_seed(random.randrange(2**32))
model = keras.models.Model(scope['input_symbol'], scope['output'])
scope['model'] = model
for term in scope.get('extra_losses', []):
model.add_loss(term)
metrics = scope.get('metrics', [])
model.compile(self.arguments['optimizer'],
loss=scope['loss'],
metrics=metrics)
branches = self.arguments['branches']
with contextlib.ExitStack() as context_stack:
try:
handle = context_stack.enter_context(
storage.open(scope.get('dump_filename', 'dump.sqlite'),
'r',
on_filesystem=True))
with gtar.GTAR(handle.name, 'r') as traj:
queue_json = traj.readStr('job_queue.json')
if queue_json is None:
raise RuntimeError()
job_queue = list(map(tuple, json.loads(queue_json)))
(frames_completed, job_parent, job_grandparent, frame_index,
replica) = job_queue[0]
replica += 1
job_queue[0] = (frames_completed, job_parent, job_grandparent,
frame_index, replica)
# when we branch from this job, start at the second
# saved frame, rather than the first (which will be a
# duplicate of where we are branching from)
this_job_start = 1
if job_parent == NO_PARENT:
print(
'Initializing a replica from scratch, but data were found'
)
frame_index = 0
else:
print('Loading weights from parent ID {}, index {}'.format(
job_parent, frame_index))
traj = context_stack.enter_context(
keras_gtar.Trajectory(handle.name, 'r', job_parent))
model.set_weights(traj.get_weights(frame_index))
scope['last_epoch'] = int(traj.frames[frame_index])
print('restored frame index', scope['last_epoch'])
# file doesn't exist or json file isn't found
except (FileNotFoundError, RuntimeError):
initial_replica = branches - self.arguments.get(
'initial_states', branches)
job_queue = [(0, NO_PARENT, NO_PARENT, 0, initial_replica)]
job_parent, frame_index = NO_PARENT, 0
this_job_start = frames_completed = 0
print('Initializing a replica from scratch')
callbacks = scope.get('callbacks', [])
callbacks.append(
tfa.callbacks.TQDMProgressBar(show_epoch_progress=False,
update_per_second=1))
if 'early_stopping' in self.arguments:
callbacks.append(
keras.callbacks.EarlyStopping(
patience=self.arguments['early_stopping'],
monitor='val_loss'))
if 'reduce_lr' in self.arguments:
callbacks.append(
keras.callbacks.ReduceLROnPlateau(
patience=self.arguments['reduce_lr'],
monitor='val_loss',
factor=.5,
verbose=True))
job_hash = hashlib.sha1(
json.dumps(scope['workflow'].to_JSON()).encode()).hexdigest()[:32]
with contextlib.ExitStack() as context_stack:
dump_name = None
storage_handle = context_stack.enter_context(
storage.open(scope.get('dump_filename', 'dump.sqlite'),
'a',
on_filesystem=True))
if self.arguments.get('dump_period', None):
cbk = keras_gtar.GTARLogger(storage_handle.name,
self.arguments['dump_period'],
append=True,
when='pre_epoch',
group=job_hash)
callbacks.append(cbk)
model.fit(scope['x_train'],
scope['y_train'],
verbose=False,
epochs=self.arguments['epochs'],
batch_size=self.arguments['batch_size'],
validation_split=self.arguments['validation_split'],
callbacks=callbacks,
initial_epoch=scope.get('last_epoch', 0))
traj = context_stack.enter_context(
gtar.GTAR(storage_handle.name, 'a'))
metadata = dict(workflow=scope['workflow'].to_JSON())
metadata['parent'] = job_parent
metadata['frame_index'] = frame_index
metadata['job_hash'] = job_hash
traj.writeStr('{}/metadata.json'.format(job_hash),
json.dumps(metadata))
# make random initializations have themselves as parents
if job_parent == NO_PARENT:
job_parent = job_hash
this_job_start = 0
this_job_info = (frames_completed + this_job_start, job_hash,
job_parent, this_job_start, 0)
if self.arguments['exponential'] or this_job_info[0] == 0:
job_queue.append(this_job_info)
(frames_completed, job_parent, job_grandparent, frame_index,
replica) = job_queue[0]
if replica + 1 >= branches:
frame_index += 1
frames_completed += 1
replica = 0
job_queue.pop(0)
if job_parent != NO_PARENT:
job_queue.append((frames_completed, job_parent,
job_grandparent, frame_index, replica))
job_queue.sort()
traj.writeStr('job_queue.json', json.dumps(job_queue))
def run(self, scope, storage):
maybe_setup_tensorflow()
maybe_set_seed(self.arguments.get('seed', None))
model = self.get_model(scope, storage)
initial_epoch = scope['last_epoch'] + 1
if self.arguments['epochs'] <= initial_epoch:
return
metrics = []
for m in self.arguments['metrics']:
metrics.append(METRIC_MAP.get(m, m))
callbacks = list(scope.setdefault('callbacks', []))
callbacks.append(TimingCallback())
early_stopping_callback = None
if self.arguments.get('early_stopping', None):
patience = self.arguments['early_stopping']
early_stopping_callback = keras.callbacks.EarlyStopping(
patience=patience, restore_best_weights=True, verbose=True)
callbacks.append(early_stopping_callback)
time_callback = None
if self.arguments.get('time_limit', None):
time_callback = TimeLimitCallback(self.arguments['time_limit'])
callbacks.append(time_callback)
if self.arguments.get('reduce_lr', None):
reduce_lr_callback = keras.callbacks.ReduceLROnPlateau(
factor=self.arguments['reduce_lr_factor'],
patience=self.arguments['reduce_lr'])
callbacks.append(reduce_lr_callback)
optimizer_kwargs = dict(scope.get('optimizer_kwargs', {}))
optimizer_kwargs.update(
dict(self.arguments.get('optimizer_kwargs', {})))
optimizer_cls = getattr(keras.optimizers, self.arguments['optimizer'])
optimizer = optimizer_cls(**optimizer_kwargs)
model.compile(
optimizer,
loss=scope['loss'],
metrics=metrics,
)
args = []
kwargs = dict(scope.get('model_train_kwargs', {}))
kwargs.update(
dict(
callbacks=callbacks,
epochs=self.arguments['epochs'],
initial_epoch=initial_epoch,
verbose=False,
))
use_fit_generator = False
if 'training_data' in scope:
args.extend(scope['training_data'])
kwargs['batch_size'] = scope.get('batch_size', 32)
elif 'training_data_generator' in scope:
args.append(scope['training_data_generator'])
use_fit_generator = True
else:
raise NotImplementedError()
if 'validation_data' in scope:
kwargs['validation_data'] = scope['validation_data']
elif 'validation_data_generator' in scope:
kwargs['validation_data'] = scope['validation_data_generator']
elif self.arguments['validation_split'] > 0:
kwargs['validation_split'] = self.arguments['validation_split']
try:
if use_fit_generator:
history = model.fit_generator(*args, **kwargs)
else:
history = model.fit(*args, **kwargs)
except KeyboardInterrupt:
history = model.history
test_evals = {}
if 'test_data' in scope:
values = model.evaluate(*scope['test_data'],
verbose=False,
batch_size=kwargs['batch_size'])
test_evals.update(dict(zip(model.metrics_names, values)))
elif 'test_data_generator' in scope:
values = model.evaluate_generator(scope['test_data_generator'],
verbose=False,
steps=scope['test_steps'])
test_evals.update(dict(zip(model.metrics_names, values)))
metadata = scope.get('metadata', {})
filename = scope.get('filename', 'dump.zip')
final_epoch = history.epoch[-1]
end_reason = None
if final_epoch + 1 >= self.arguments['epochs']:
end_reason = 'completed'
elif time_callback is not None and time_callback.triggered:
pass # we should resume training later
elif early_stopping_callback is not None and early_stopping_callback.stopped_epoch:
end_reason = 'early_stopping'
with storage.open(filename, 'ab', on_filesystem=True) as f:
gtar_mode = 'a' if os.stat(f.name).st_size > 0 else 'w'
with keras_gtar.Trajectory(f.name, gtar_mode) as traj:
traj.save(model, str(final_epoch))
with gtar.GTAR(f.name, 'a') as traj:
for name, vals in history.history.items():
rec = gtar.Record('', name, str(final_epoch),
gtar.Behavior.Continuous,
gtar.Format.Float32,
gtar.Resolution.Uniform)
traj.writeRecord(rec, vals)
for name, vals in test_evals.items():
name = 'test_{}'.format(name)
rec = gtar.Record('', name, str(final_epoch),
gtar.Behavior.Discrete,
gtar.Format.Float32,
gtar.Resolution.Uniform)
traj.writeRecord(rec, vals)
traj.writeStr('metadata.json', json.dumps(metadata))
if end_reason is not None:
traj.writeStr('end_reason.txt', end_reason)
#!/usr/bin/env python
# Copyright (c) 2020 The Regents of the University of Michigan
# All rights reserved.
# This software is licensed under the BSD 3-Clause License.
import numpy as np
import gtar
N = 100
with gtar.GTAR('libgetar_sample.tar', 'w') as traj:
traj.writePath('frames/0/position.f32.ind', np.random.rand(N, 3))
traj.writePath('frames/0/orientation.f32.ind', np.random.rand(N, 4))
traj.writePath('type.u32.ind', N // 2 * [0] + (N - N // 2) * [1])
traj.writePath('type_names.json', '["type1", "type2"]')
def __init__(self, filename, mode='r', group=None):
self.filename = filename
self.mode = mode
self.handle = gtar.GTAR(filename, mode)
self.group = group
def load_move_distance(self, scope, storage, context, integrator, updater):
import gtar
# per-type move distance arrays for translation/rotation/box
distances = {}
frame = -1
dump_filename = scope.get('dump_filename', 'dump.sqlite')
local_context = contextlib.ExitStack()
with local_context:
try:
dump_file = local_context.enter_context(storage.open(
dump_filename, 'rb', on_filesystem=True, noop=scope['mpi_rank']))
traj = local_context.enter_context(gtar.GTAR(dump_file.name, 'r'))
except (FileNotFoundError, RuntimeError):
return -1, {}
# grab per-type distance arrays
for (frame, trans) in traj.recordsNamed('type_translation_distance'):
if int(frame) <= scope['cumulative_steps']:
distances['translation'] = trans
frame = int(frame)
else:
break
for (frame, rot) in traj.recordsNamed('type_rotation_distance'):
if int(frame) <= scope['cumulative_steps']:
distances['rotation'] = rot
frame = int(frame)
else:
break
for name in BOX_MOVE_NAMES:
box_name = 'box_{}'.format(name)
distance_name = '{}_distance'.format(box_name)
for (frame, rot) in traj.recordsNamed(distance_name):
if int(frame) <= scope['cumulative_steps']:
distances[box_name] = rot
frame = int(frame)
else:
break
if distances:
types = context.snapshot.particles.types
kwargs = dict()
if 'translation' in distances:
kwargs['d'] = dict(zip(types, distances['translation']))
if 'rotation' in distances:
kwargs['a'] = dict(zip(types, distances['rotation']))
integrator.set_params(**kwargs)
for (key, distance) in distances.items():
if key.startswith('box_') and updater is not None:
name = key[4:]
getattr(updater, name)(delta=distance)
msg = 'Loaded tuned move distances from {} for frame {}: {}'.format(
dump_filename, frame, distances)
logger.debug(msg)
return frame, distances
