From 792579513197d14d9f2f1f9461cbeaad6f551c23 Mon Sep 17 00:00:00 2001 From: "Keaton J. Burns" Date: Mon, 16 Dec 2024 10:56:58 -0500 Subject: [PATCH] Add loading from coeff hdf5 data --- dedalus/core/field.py | 27 ++++++++++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/dedalus/core/field.py b/dedalus/core/field.py index 875c0561..c7359fb6 100644 --- a/dedalus/core/field.py +++ b/dedalus/core/field.py @@ -724,9 +724,30 @@ def load_from_hdf5(self, file, index, task=None, func=None): if task is None: task = self.name dset = file['tasks'][task] - if not np.all(dset.attrs['grid_space']): - raise ValueError("Can only load data from grid space") - self.load_from_global_grid_data(dset, pre_slices=(index,), func=func) + if np.all(dset.attrs['grid_space']): + self.load_from_global_grid_data(dset, pre_slices=(index,), func=func) + elif np.all(~dset.attrs['grid_space']): + self.load_from_global_coeff_data(dset, pre_slices=(index,), func=func) + else: + raise ValueError("Can only load global data from pure grid or coeff space") + + def load_from_global_coeff_data(self, global_data, pre_slices=tuple(), func=None): + """Load local coeff data from array-like global coeff data.""" + dim = self.dist.dim + layout = self.dist.coeff_layout + # Check shapes + data_shape = global_data.shape[-dim:] + self_shape = layout.global_shape(self.domain, scales=1) + if data_shape != self_shape: + raise ValueError("Cannot change global shape when loading coeff data.") + # Extract local data from global data + component_slices = tuple(slice(None) for cs in self.tensorsig) + spatial_slices = layout.slices(self.domain, scales=1) + local_slices = pre_slices + component_slices + spatial_slices + if func is None: + self[layout] = global_data[local_slices] + else: + self[layout] = func(global_data[local_slices]) def load_from_global_grid_data(self, global_data, pre_slices=tuple(), func=None): """Load local grid data from array-like global grid data."""