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)