def append(self, snapshot):
snapshot_data = snapshot
if not isinstance(snapshot, np.ndarray):
snapshot_data = snapshot.to_array()
if self.in_memory:
if self._snapshots is not None:
if snapshot_data.shape != self.snapshot_shape:
raise ValueError("This snapshot doesn't match the initialized shape.")
else:
self.snapshot_shape = snapshot_data.shape
self._init_snapshot_storage()
if self._current_index >= self._reserved_length:
if self.max_snapshots:
raise IndexError("Appending exceeded the set maximum snapshot count.")
else:
self._expand_snapshot_storage()
self._snapshots[self._current_index, :] = snapshot_data
if self.output_path:
output_file_name = util.construct_snapshot_name(self.output_path, self._current_index)
util.save_snapshot(snapshot_data, output_file_name)
self._current_index += 1
self.snapshot_count += 1
def save(self, file_name):
""" Saves all appended snapshots to file or directory
Args:
file_name: Can be either a directory name, in which case .csv files
will be saved for each snapshot, or a .pkl file, in which case
all snapshots will be stored in a numpy pickle.
"""
if not self.in_memory:
raise IOError("Can't save the storage since it's not stored in memory.")
if file_name.endswith(".pkl"):
self._snapshots[:self._current_index, :].dump(file_name)
else:
for i in range(self._current_index):
output_file_name = util.construct_snapshot_name(file_name, i)
util.save_snapshot(self._snapshots[i, :], output_file_name)
