diff --git a/src/struphy/fields_background/equils.py b/src/struphy/fields_background/equils.py index f1afacd35..888404d71 100644 --- a/src/struphy/fields_background/equils.py +++ b/src/struphy/fields_background/equils.py @@ -8,6 +8,8 @@ from time import time import cunumpy as xp +from psydac.ddm.mpi import MockMPI +from psydac.ddm.mpi import mpi as MPI from scipy.integrate import odeint, quad from scipy.interpolate import RectBivariateSpline, UnivariateSpline from scipy.optimize import fsolve, minimize @@ -31,6 +33,17 @@ from struphy.fields_background.mhd_equil.eqdsk import readeqdsk from struphy.utils.utils import read_state, subp_run +if isinstance(MPI, MockMPI): + comm = None + rank = 0 + size = 1 + Barrier = lambda: None +else: + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + size = comm.Get_size() + Barrier = comm.Barrier + class HomogenSlab(CartesianMHDequilibrium): r""" @@ -1748,7 +1761,8 @@ def __init__( # default input file if file is None: file = "AUGNLED_g031213.00830.high" - print(f"EQDSK: taking default file {file}.") + if rank == 0: + print(f"EQDSK: taking default file {file}.") # no rescaling if units are not provided if units is None: @@ -2131,9 +2145,11 @@ def __init__( import pytest with pytest.raises(SystemExit) as exc: - print("Simulation aborted, gvec must be installed (pip install gvec)!") + if rank == 0: + print("Simulation aborted, gvec must be installed (pip install gvec)!") sys.exit(1) - print(f"{exc.value.code =}") + if rank == 0: + print(f"{exc.value.code =}") import gvec @@ -2410,13 +2426,15 @@ def __init__( desc_spec = importlib.util.find_spec("desc") if desc_spec is None: - print("Simulation aborted, desc-opt must be installed!") - print("Install with:\npip install desc-opt") + if rank == 0: + print("Simulation aborted, desc-opt must be installed!") + print("Install with:\npip install desc-opt") sys.exit(1) import desc - print(f"DESC import: {time() - t} seconds") + if rank == 0: + print(f"DESC import: {time() - t} seconds") from struphy.geometry.domains import DESCunit # no rescaling if units are not provided @@ -2449,7 +2467,8 @@ def __init__( else: self._eq = desc.io.load(eq_name) - print(f"Eq. load: {time() - t} seconds") + if rank == 0: + print(f"Eq. load: {time() - t} seconds") self._rmin = self.params["rmin"] self._use_nfp = self.params["use_nfp"] @@ -2878,7 +2897,7 @@ def desc_eval( assert xp.all(theta == theta1) assert xp.all(zeta == zeta1) - if verbose: + if verbose and rank == 0: # import sys print(f"\n{nfp =}") print(f"{self.eq.axis =}") @@ -2900,7 +2919,8 @@ def desc_eval( # make c-contiguous out = xp.ascontiguousarray(out) - print(f"desc_eval for {var}: {time() - ttime} seconds") + if rank == 0: + print(f"desc_eval for {var}: {time() - ttime} seconds") return out diff --git a/src/struphy/initial/utilities.py b/src/struphy/initial/utilities.py index 7af042249..69303a3a1 100644 --- a/src/struphy/initial/utilities.py +++ b/src/struphy/initial/utilities.py @@ -54,16 +54,15 @@ def __init__(self, derham, name, species, **params): else: key = "restart/" + species + "_" + name - if isinstance(data.file[key], h5py.Dataset): - self._vector = data.file[key][-1] + with h5py.File(data.file_path, "a") as file: + if isinstance(file[key], h5py.Dataset): + self._vector = file[key][-1] - else: - self._vector = [] - - for n in range(3): - self._vector += [data.file[key + "/" + str(n + 1)][-1]] + else: + self._vector = [] - data.file.close() + for n in range(3): + self._vector += [file[key + "/" + str(n + 1)][-1]] @property def vector(self): diff --git a/src/struphy/io/output_handling.py b/src/struphy/io/output_handling.py index 808c2bcf3..4c6e757fc 100644 --- a/src/struphy/io/output_handling.py +++ b/src/struphy/io/output_handling.py @@ -38,13 +38,10 @@ def __init__(self, path_out, file_name=None, comm=None): self._file_name = file_name # file path - file_path = os.path.join(path_out, "data/", self._file_name) + self._file_path = os.path.join(path_out, "data/", self._file_name) # check if file already exists - file_exists = os.path.exists(file_path) - - # open/create file - self._file = h5py.File(file_path, "a") + file_exists = os.path.exists(self.file_path) # dictionary with pairs (dataset key : object ID) self._dset_dict = {} @@ -53,9 +50,10 @@ def __init__(self, path_out, file_name=None, comm=None): if file_exists: dataset_keys = [] - self._file.visit( - lambda key: dataset_keys.append(key) if isinstance(self._file[key], h5py.Dataset) else None, - ) + with h5py.File(self.file_path, "a") as file: + file.visit( + lambda key: dataset_keys.append(key) if isinstance(self._file[key], h5py.Dataset) else None, + ) for key in dataset_keys: self._dset_dict[key] = None @@ -66,9 +64,9 @@ def file_name(self): return self._file_name @property - def file(self): - """The hdf5 file.""" - return self._file + def file_path(self): + """The absolute path to the hdf5 file.""" + return self._file_path @property def dset_dict(self): @@ -103,20 +101,21 @@ def add_data(self, data_dict): # create new dataset otherwise and save array else: - # scalar values are saved as 1d arrays of size 1 - if val.size == 1: - assert val.ndim == 1 - self._file.create_dataset(key, (1,), maxshape=(None,), dtype=val.dtype, chunks=True) - self._file[key][0] = val[0] - else: - self._file.create_dataset( - key, - (1,) + val.shape, - maxshape=(None,) + val.shape, - dtype=val.dtype, - chunks=True, - ) - self._file[key][0] = val + with h5py.File(self.file_path, "a") as file: + # scalar values are saved as 1d arrays of size 1 + if val.size == 1: + assert val.ndim == 1 + file.create_dataset(key, (1,), maxshape=(None,), dtype=val.dtype, chunks=True) + file[key][0] = val[0] + else: + file.create_dataset( + key, + (1,) + val.shape, + maxshape=(None,) + val.shape, + dtype=val.dtype, + chunks=True, + ) + file[key][0] = val # set object ID self._dset_dict[key] = id(val) @@ -130,25 +129,26 @@ def save_data(self, keys=None): keys : list Keys to the data objects specified when using "add_data". Default saves all specified data objects. """ - - # loop over all keys - if keys is None: - for key in self._dset_dict: - self._file[key].resize(self._file[key].shape[0] + 1, axis=0) - self._file[key][-1] = ctypes.cast(self._dset_dict[key], ctypes.py_object).value - - # only loop over given keys - else: - for key in keys: - self._file[key].resize(self._file[key].shape[0] + 1, axis=0) - self._file[key][-1] = ctypes.cast(self._dset_dict[key], ctypes.py_object).value + with h5py.File(self.file_path, "a") as file: + # loop over all keys + if keys is None: + for key in self._dset_dict: + file[key].resize(file[key].shape[0] + 1, axis=0) + file[key][-1] = ctypes.cast(self._dset_dict[key], ctypes.py_object).value + + # only loop over given keys + else: + for key in keys: + file[key].resize(file[key].shape[0] + 1, axis=0) + file[key][-1] = ctypes.cast(self._dset_dict[key], ctypes.py_object).value def info(self): """Print info of data sets to screen.""" for key in self._dset_dict: - print(f"\nData set name: {key}") - print("Shape:", self._file[key].shape) - print("Attributes:") - for attr, val in self._file[key].attrs.items(): - print(attr, val) + with h5py.File(self.file_path, "a") as file: + print(f"\nData set name: {key}") + print("Shape:", file[key].shape) + print("Attributes:") + for attr, val in file[key].attrs.items(): + print(attr, val) diff --git a/src/struphy/main.py b/src/struphy/main.py index f6bf8a4e2..bae521086 100644 --- a/src/struphy/main.py +++ b/src/struphy/main.py @@ -280,9 +280,10 @@ def run( if restart: model.initialize_from_restart(data) - time_state["value"][0] = data.file["restart/time/value"][-1] - time_state["value_sec"][0] = data.file["restart/time/value_sec"][-1] - time_state["index"][0] = data.file["restart/time/index"][-1] + with h5py.File(data.file_path, "a") as file: + time_state["value"][0] = file["restart/time/value"][-1] + time_state["value_sec"][0] = file["restart/time/value_sec"][-1] + time_state["index"][0] = file["restart/time/index"][-1] total_steps = str(int(round((Tend - time_state["value"][0]) / dt))) else: @@ -317,7 +318,6 @@ def run( if break_cond_1 or break_cond_2: # save restart data (other data already saved below) data.save_data(keys=save_keys_end) - data.file.close() end_simulation = time.time() if rank == 0: print(f"\nTime steps done: {time_state['index'][0]}") @@ -474,37 +474,34 @@ def pproc( return # check for fields and kinetic data in hdf5 file that need post processing - file = h5py.File(os.path.join(path, "data/", "data_proc0.hdf5"), "r") + with h5py.File(os.path.join(path, "data/", "data_proc0.hdf5"), "r") as file: + # save time grid at which post-processing data is created + xp.save(os.path.join(path_pproc, "t_grid.npy"), file["time/value"][::step].copy()) - # save time grid at which post-processing data is created - xp.save(os.path.join(path_pproc, "t_grid.npy"), file["time/value"][::step].copy()) - - if "feec" in file.keys(): - exist_fields = True - else: - exist_fields = False - - if "kinetic" in file.keys(): - exist_kinetic = {"markers": False, "f": False, "n_sph": False} - kinetic_species = [] - kinetic_kinds = [] - for name in file["kinetic"].keys(): - kinetic_species += [name] - kinetic_kinds += [next(iter(model.species[name].variables.values())).space] - - # check for saved markers - if "markers" in file["kinetic"][name]: - exist_kinetic["markers"] = True - # check for saved distribution function - if "f" in file["kinetic"][name]: - exist_kinetic["f"] = True - # check for saved sph density - if "n_sph" in file["kinetic"][name]: - exist_kinetic["n_sph"] = True - else: - exist_kinetic = None - - file.close() + if "feec" in file.keys(): + exist_fields = True + else: + exist_fields = False + + if "kinetic" in file.keys(): + exist_kinetic = {"markers": False, "f": False, "n_sph": False} + kinetic_species = [] + kinetic_kinds = [] + for name in file["kinetic"].keys(): + kinetic_species += [name] + kinetic_kinds += [next(iter(model.species[name].variables.values())).space] + + # check for saved markers + if "markers" in file["kinetic"][name]: + exist_kinetic["markers"] = True + # check for saved distribution function + if "f" in file["kinetic"][name]: + exist_kinetic["f"] = True + # check for saved sph density + if "n_sph" in file["kinetic"][name]: + exist_kinetic["n_sph"] = True + else: + exist_kinetic = None # field post-processing if exist_fields: diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 4afcf7709..e75d281e0 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -6,6 +6,7 @@ from textwrap import indent import cunumpy as xp +import h5py import yaml from line_profiler import profile from psydac.ddm.mpi import MockMPI @@ -976,7 +977,7 @@ def print_scalar_quantities(self): # threshold=obj.weights_params["threshold"], # ) - def initialize_from_restart(self, data): + def initialize_from_restart(self, data: DataContainer): """ Set initial conditions for FE coefficients (electromagnetic and fluid) and markers from restart group in hdf5 files. @@ -986,41 +987,42 @@ def initialize_from_restart(self, data): The data object that links to the hdf5 files. """ - # initialize em fields - if len(self.em_fields) > 0: - for key, val in self.em_fields.items(): - if "params" in key: - continue - else: - obj = val["obj"] - assert isinstance(obj, SplineFunction) - obj.initialize_coeffs_from_restart_file(data.file) - - # initialize fields - if len(self.fluid) > 0: - for species, val in self.fluid.items(): - for variable, subval in val.items(): - if "params" in variable: + with h5py.File(data.file_path, "a") as file: + # initialize em fields + if len(self.em_fields) > 0: + for key, val in self.em_fields.items(): + if "params" in key: continue else: - obj = subval["obj"] + obj = val["obj"] assert isinstance(obj, SplineFunction) - obj.initialize_coeffs_from_restart_file( - data.file, - species, - ) + obj.initialize_coeffs_from_restart_file(file) + + # initialize fields + if len(self.fluid) > 0: + for species, val in self.fluid.items(): + for variable, subval in val.items(): + if "params" in variable: + continue + else: + obj = subval["obj"] + assert isinstance(obj, SplineFunction) + obj.initialize_coeffs_from_restart_file( + file, + species, + ) - # initialize particles - if len(self.kinetic) > 0: - for key, val in self.kinetic.items(): - obj = val["obj"] - assert isinstance(obj, Particles) - obj.draw_markers(verbose=self.verbose) - obj._markers[:, :] = data.file["restart/" + key][-1, :, :] + # initialize particles + if len(self.kinetic) > 0: + for key, val in self.kinetic.items(): + obj = val["obj"] + assert isinstance(obj, Particles) + obj.draw_markers(verbose=self.verbose) + obj._markers[:, :] = file["restart/" + key][-1, :, :] - # important: sets holes attribute of markers! - if self.comm_world is not None: - obj.mpi_sort_markers(do_test=True) + # important: sets holes attribute of markers! + if self.comm_world is not None: + obj.mpi_sort_markers(do_test=True) def initialize_data_output(self, data: DataContainer, size): """ @@ -1049,111 +1051,112 @@ def initialize_data_output(self, data: DataContainer, size): key_scalar = "scalar/" + key data.add_data({key_scalar: val}) - # store grid_info only for runs with 512 ranks or smaller - if self._scalar_quantities and self.derham is not None: - if size <= 512: - data.file["scalar"].attrs["grid_info"] = self.derham.domain_array + with h5py.File(data.file_path, "a") as file: + # store grid_info only for runs with 512 ranks or smaller + if self._scalar_quantities and self.derham is not None: + if size <= 512: + file["scalar"].attrs["grid_info"] = self.derham.domain_array + else: + file["scalar"].attrs["grid_info"] = self.derham.domain_array[0] else: - data.file["scalar"].attrs["grid_info"] = self.derham.domain_array[0] - else: - pass - - # save feec data in group 'feec/' - feec_species = self.field_species | self.fluid_species | self.diagnostic_species - for species, val in feec_species.items(): - assert isinstance(val, Species) + pass - species_path = os.path.join("feec", species) - species_path_restart = os.path.join("restart", species) + # save feec data in group 'feec/' + feec_species = self.field_species | self.fluid_species | self.diagnostic_species + for species, val in feec_species.items(): + assert isinstance(val, Species) - for variable, subval in val.variables.items(): - assert isinstance(subval, FEECVariable) - spline = subval.spline + species_path = os.path.join("feec", species) + species_path_restart = os.path.join("restart", species) - # in-place extraction of FEM coefficients from field.vector --> field.vector_stencil! - spline.extract_coeffs(update_ghost_regions=False) + for variable, subval in val.variables.items(): + assert isinstance(subval, FEECVariable) + spline = subval.spline - # save numpy array to be updated each time step. - if subval.save_data: - key_field = os.path.join(species_path, variable) + # in-place extraction of FEM coefficients from field.vector --> field.vector_stencil! + spline.extract_coeffs(update_ghost_regions=False) - if isinstance(spline.vector_stencil, StencilVector): - data.add_data( - {key_field: spline.vector_stencil._data}, - ) + # save numpy array to be updated each time step. + if subval.save_data: + key_field = os.path.join(species_path, variable) - else: - for n in range(3): - key_component = os.path.join(key_field, str(n + 1)) + if isinstance(spline.vector_stencil, StencilVector): data.add_data( - {key_component: spline.vector_stencil[n]._data}, + {key_field: spline.vector_stencil._data}, ) - # save field meta data - data.file[key_field].attrs["space_id"] = spline.space_id - data.file[key_field].attrs["starts"] = spline.starts - data.file[key_field].attrs["ends"] = spline.ends - data.file[key_field].attrs["pads"] = spline.pads + else: + for n in range(3): + key_component = os.path.join(key_field, str(n + 1)) + data.add_data( + {key_component: spline.vector_stencil[n]._data}, + ) - # save numpy array to be updated only at the end of the simulation for restart. - key_field_restart = os.path.join(species_path_restart, variable) + # save field meta data + file[key_field].attrs["space_id"] = spline.space_id + file[key_field].attrs["starts"] = spline.starts + file[key_field].attrs["ends"] = spline.ends + file[key_field].attrs["pads"] = spline.pads - if isinstance(spline.vector_stencil, StencilVector): - data.add_data( - {key_field_restart: spline.vector_stencil._data}, - ) - else: - for n in range(3): - key_component_restart = os.path.join(key_field_restart, str(n + 1)) + # save numpy array to be updated only at the end of the simulation for restart. + key_field_restart = os.path.join(species_path_restart, variable) + + if isinstance(spline.vector_stencil, StencilVector): data.add_data( - {key_component_restart: spline.vector_stencil[n]._data}, + {key_field_restart: spline.vector_stencil._data}, ) + else: + for n in range(3): + key_component_restart = os.path.join(key_field_restart, str(n + 1)) + data.add_data( + {key_component_restart: spline.vector_stencil[n]._data}, + ) - # save kinetic data in group 'kinetic/' - for name, species in self.particle_species.items(): - assert isinstance(species, ParticleSpecies) - assert len(species.variables) == 1, "More than 1 variable per kinetic species is not allowed." - for varname, var in species.variables.items(): - assert isinstance(var, PICVariable | SPHVariable) - obj = var.particles - assert isinstance(obj, Particles) - - key_spec = os.path.join("kinetic", name) - key_spec_restart = os.path.join("restart", name) - - # restart data - data.add_data({key_spec_restart: obj.markers}) - - # marker data - key_mks = os.path.join(key_spec, "markers") - data.add_data({key_mks: var.saved_markers}) - - # binning plot data - for bin_plot in species.binning_plots: - key_f = os.path.join(key_spec, "f", bin_plot.slice) - key_df = os.path.join(key_spec, "df", bin_plot.slice) - - data.add_data({key_f: bin_plot.f}) - data.add_data({key_df: bin_plot.df}) + # save kinetic data in group 'kinetic/' + for name, species in self.particle_species.items(): + assert isinstance(species, ParticleSpecies) + assert len(species.variables) == 1, "More than 1 variable per kinetic species is not allowed." + for varname, var in species.variables.items(): + assert isinstance(var, PICVariable | SPHVariable) + obj = var.particles + assert isinstance(obj, Particles) - for dim, be in enumerate(bin_plot.bin_edges): - data.file[key_f].attrs["bin_centers" + "_" + str(dim + 1)] = be[:-1] + (be[1] - be[0]) / 2 + key_spec = os.path.join("kinetic", name) + key_spec_restart = os.path.join("restart", name) - for i, kd_plot in enumerate(species.kernel_density_plots): - key_n = os.path.join(key_spec, "n_sph", f"view_{i}") + # restart data + data.add_data({key_spec_restart: obj.markers}) - data.add_data({key_n: kd_plot.n_sph}) - # save 1d point values, not meshgrids, because attrs size is limited - eta1 = kd_plot.plot_pts[0][:, 0, 0] - eta2 = kd_plot.plot_pts[1][0, :, 0] - eta3 = kd_plot.plot_pts[2][0, 0, :] - data.file[key_n].attrs["eta1"] = eta1 - data.file[key_n].attrs["eta2"] = eta2 - data.file[key_n].attrs["eta3"] = eta3 + # marker data + key_mks = os.path.join(key_spec, "markers") + data.add_data({key_mks: var.saved_markers}) - # TODO: maybe add other data - # else: - # data.add_data({key_dat: val1}) + # binning plot data + for bin_plot in species.binning_plots: + key_f = os.path.join(key_spec, "f", bin_plot.slice) + key_df = os.path.join(key_spec, "df", bin_plot.slice) + + data.add_data({key_f: bin_plot.f}) + data.add_data({key_df: bin_plot.df}) + + for dim, be in enumerate(bin_plot.bin_edges): + file[key_f].attrs["bin_centers" + "_" + str(dim + 1)] = be[:-1] + (be[1] - be[0]) / 2 + + for i, kd_plot in enumerate(species.kernel_density_plots): + key_n = os.path.join(key_spec, "n_sph", f"view_{i}") + + data.add_data({key_n: kd_plot.n_sph}) + # save 1d point values, not meshgrids, because attrs size is limited + eta1 = kd_plot.plot_pts[0][:, 0, 0] + eta2 = kd_plot.plot_pts[1][0, :, 0] + eta3 = kd_plot.plot_pts[2][0, 0, :] + file[key_n].attrs["eta1"] = eta1 + file[key_n].attrs["eta2"] = eta2 + file[key_n].attrs["eta3"] = eta3 + + # TODO: maybe add other data + # else: + # data.add_data({key_dat: val1}) # keys to be saved at each time step and only at end (restart) save_keys_all = [] diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 1dc148cdb..c49b3147e 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -1327,25 +1327,20 @@ def _load_external( Cumulative sum of number of markers on each process at loading stage. """ if self.mpi_rank == 0: - file = h5py.File( - self.loading_params.dir_external, - "r", - ) - print(f"\nLoading markers from file: {file}") - - self._markers[ - : n_mks_load_cum_sum[0], - :, - ] = file["markers"][: n_mks_load_cum_sum[0], :] - - for i in range(1, self._mpi_size): - self._mpi_comm.Send( - file["markers"][n_mks_load_cum_sum[i - 1] : n_mks_load_cum_sum[i], :], - dest=i, - tag=123, - ) + with h5py.File(self.loading_params.dir_external, "r") as file: + print(f"\nLoading markers from file: {file}") - file.close() + self._markers[ + : n_mks_load_cum_sum[0], + :, + ] = file["markers"][: n_mks_load_cum_sum[0], :] + + for i in range(1, self._mpi_size): + self._mpi_comm.Send( + file["markers"][n_mks_load_cum_sum[i - 1] : n_mks_load_cum_sum[i], :], + dest=i, + tag=123, + ) else: recvbuf = xp.zeros( (n_mks_load_loc, self.markers.shape[1]), @@ -1370,7 +1365,8 @@ def _load_restart(self): data_path = self.loading_params.dir_particles_abs data = DataContainer(data_path, comm=self.mpi_comm) - self._markers[:, :] = data.file["restart/" + self.loading_params.restart_key][-1, :, :] + with h5py.File(data.file_path, "a") as file: + self._markers[:, :] = file["restart/" + self.loading_params.restart_key][-1, :, :] def _load_tesselation(self, n_quad: int = 1): """ diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index e0759bb63..dfb418e8c 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -154,18 +154,17 @@ def create_femfields( ) # get fields names, space IDs and time grid from 0-th rank hdf5 file - file = h5py.File(os.path.join(path, "data/", "data_proc0.hdf5"), "r") - space_ids = {} - print("\nReading hdf5 data of following species:") - for species, dset in file["feec"].items(): - space_ids[species] = {} - print(f"{species}:") - for var, ddset in dset.items(): - space_ids[species][var] = ddset.attrs["space_id"] - print(f" {var}:", ddset) - - t_grid = file["time/value"][::step].copy() - file.close() + with h5py.File(os.path.join(path, "data/", "data_proc0.hdf5"), "r") as file: + space_ids = {} + print("\nReading hdf5 data of following species:") + for species, dset in file["feec"].items(): + space_ids[species] = {} + print(f"{species}:") + for var, ddset in dset.items(): + space_ids[species][var] = ddset.attrs["space_id"] + print(f" {var}:", ddset) + + t_grid = file["time/value"][::step].copy() # create one FemField for each snapshot fields = {} @@ -184,66 +183,57 @@ def create_femfields( print("") for rank in range(int(nproc)): # open hdf5 file - file = h5py.File( - os.path.join( - path, - "data/", - "data_proc" + str(rank) + ".hdf5", - ), - "r", - ) - - for species, dset in file["feec"].items(): - for var, ddset in tqdm(dset.items()): - # get global start indices, end indices and pads - gl_s = ddset.attrs["starts"] - gl_e = ddset.attrs["ends"] - pads = ddset.attrs["pads"] - - assert gl_s.shape == (3,) or gl_s.shape == (3, 3) - assert gl_e.shape == (3,) or gl_e.shape == (3, 3) - assert pads.shape == (3,) or pads.shape == (3, 3) - - # loop over time - for n, t in enumerate(t_grid): - # scalar field - if gl_s.shape == (3,): - s1, s2, s3 = gl_s - e1, e2, e3 = gl_e - p1, p2, p3 = pads - - data = ddset[n * step, p1:-p1, p2:-p2, p3:-p3].copy() - - fields[t][species][var].vector[ - s1 : e1 + 1, - s2 : e2 + 1, - s3 : e3 + 1, - ] = data - # update after each data addition, can be made more efficient - fields[t][species][var].vector.update_ghost_regions() - - # vector-valued field - else: - for comp in range(3): - s1, s2, s3 = gl_s[comp] - e1, e2, e3 = gl_e[comp] - p1, p2, p3 = pads[comp] - - data = ddset[str(comp + 1)][ - n * step, - p1:-p1, - p2:-p2, - p3:-p3, - ].copy() - - fields[t][species][var].vector[comp][ + with h5py.File(os.path.join(path, "data/", f"data_proc{rank}.hdf5"), "r") as file: + for species, dset in file["feec"].items(): + for var, ddset in tqdm(dset.items()): + # get global start indices, end indices and pads + gl_s = ddset.attrs["starts"] + gl_e = ddset.attrs["ends"] + pads = ddset.attrs["pads"] + + assert gl_s.shape == (3,) or gl_s.shape == (3, 3) + assert gl_e.shape == (3,) or gl_e.shape == (3, 3) + assert pads.shape == (3,) or pads.shape == (3, 3) + + # loop over time + for n, t in enumerate(t_grid): + # scalar field + if gl_s.shape == (3,): + s1, s2, s3 = gl_s + e1, e2, e3 = gl_e + p1, p2, p3 = pads + + data = ddset[n * step, p1:-p1, p2:-p2, p3:-p3].copy() + + fields[t][species][var].vector[ s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1, ] = data - # update after each data addition, can be made more efficient - fields[t][species][var].vector.update_ghost_regions() - file.close() + # update after each data addition, can be made more efficient + fields[t][species][var].vector.update_ghost_regions() + + # vector-valued field + else: + for comp in range(3): + s1, s2, s3 = gl_s[comp] + e1, e2, e3 = gl_e[comp] + p1, p2, p3 = pads[comp] + + data = ddset[str(comp + 1)][ + n * step, + p1:-p1, + p2:-p2, + p3:-p3, + ].copy() + + fields[t][species][var].vector[comp][ + s1 : e1 + 1, + s2 : e2 + 1, + s3 : e3 + 1, + ] = data + # update after each data addition, can be made more efficient + fields[t][species][var].vector.update_ghost_regions() print("Creation of Struphy Fields done.") @@ -527,20 +517,9 @@ def post_process_markers( nproc = meta["MPI processes"] # open hdf5 files and get names and number of saved markers of kinetic species - files = [ - h5py.File( - os.path.join( - path_in, - "data/", - f"data_proc{i}.hdf5", - ), - "r", - ) - for i in range(int(nproc)) - ] - - # get number of time steps and markers - nt, n_markers, n_cols = files[0]["kinetic/" + species + "/markers"].shape + with h5py.File(os.path.join(path_in, "data/data_proc0.hdf5"), "r") as file_0: + # get number of time steps and markers + nt, n_markers, n_cols = file_0["kinetic/" + species + "/markers"].shape log_nt = int(xp.log10(int(((nt - 1) / step)))) + 1 @@ -581,11 +560,12 @@ def post_process_markers( species + "_{0:0{1}d}.txt".format(n, log_nt), ) - for file in files: - markers = file["kinetic/" + species + "/markers"] - ids = markers[n * step, :, -1].astype("int") - ids = ids[ids != -1] # exclude holes - temp[ids] = markers[n * step, : ids.size, save_index] + for i in range(int(nproc)): + with h5py.File(os.path.join(path_in, "data/", f"data_proc{i}.hdf5"), "r") as file: + markers = file["kinetic/" + species + "/markers"] + ids = markers[n * step, :, -1].astype("int") + ids = ids[ids != -1] # exclude holes + temp[ids] = markers[n * step, : ids.size, save_index] # sorting out lost particles ids = temp[:, -1].astype("int") @@ -612,10 +592,6 @@ def post_process_markers( temp = xp.roll(temp, 1, axis=1) xp.savetxt(file_txt, temp[:, (0, 1, 2, 3, -1)], fmt="%12.6f", delimiter=", ") - # close hdf5 files - for file in files: - file.close() - def post_process_f( path_in, @@ -653,19 +629,6 @@ def post_process_f( meta = yaml.load(f, Loader=yaml.FullLoader) nproc = meta["MPI processes"] - # open hdf5 files and get names and number of saved markers of kinetic species - files = [ - h5py.File( - os.path.join( - path_in, - "data/", - f"data_proc{i}.hdf5", - ), - "r", - ) - for i in range(int(nproc)) - ] - # directory for .npy files path_distr = os.path.join(path_out, "distribution_function") @@ -678,126 +641,121 @@ def post_process_f( print("Evaluation of distribution functions for " + str(species)) # Create grids - for slice_name in tqdm(files[0]["kinetic/" + species + "/f"]): - # create a new folder for each slice - path_slice = os.path.join(path_distr, slice_name) - os.mkdir(path_slice) - - # Find out all names of slices - slice_names = slice_name.split("_") - - # save grid - for n_gr, (_, grid) in enumerate(files[0]["kinetic/" + species + "/f/" + slice_name].attrs.items()): - grid_path = os.path.join( - path_slice, - "grid_" + slice_names[n_gr] + ".npy", - ) - xp.save(grid_path, grid[:]) - - # compute distribution function - for slice_name in tqdm(files[0]["kinetic/" + species + "/f"]): - # path to folder of slice - path_slice = os.path.join(path_distr, slice_name) - - # Find out all names of slices - slice_names = slice_name.split("_") - - # load full-f data - data = files[0]["kinetic/" + species + "/f/" + slice_name][::step].copy() - for rank in range(1, int(nproc)): - data += files[rank]["kinetic/" + species + "/f/" + slice_name][::step] - - # load delta-f data - data_df = files[0]["kinetic/" + species + "/df/" + slice_name][::step].copy() - for rank in range(1, int(nproc)): - data_df += files[rank]["kinetic/" + species + "/df/" + slice_name][::step] - - # save distribution functions - xp.save(os.path.join(path_slice, "f_binned.npy"), data) - xp.save(os.path.join(path_slice, "delta_f_binned.npy"), data_df) - - if compute_bckgr: - # bckgr_params = params["kinetic"][species]["background"] - - # f_bckgr = None - # for fi, maxw_params in bckgr_params.items(): - # if fi[-2] == "_": - # fi_type = fi[:-2] - # else: - # fi_type = fi - - # if f_bckgr is None: - # f_bckgr = getattr(maxwellians, fi_type)( - # maxw_params=maxw_params, - # ) - # else: - # f_bckgr = f_bckgr + getattr(maxwellians, fi_type)( - # maxw_params=maxw_params, - # ) - - spec: ParticleSpecies = getattr(params_in.model, species) - var: PICVariable = spec.var - f_bckgr: KineticBackground = var.backgrounds - - # load all grids of the variables of f - grid_tot = [] - factor = 1.0 - - # eta-grid - for comp in range(1, 4): - current_slice = "e" + str(comp) - filename = os.path.join( + with h5py.File(os.path.join(path_in, "data/data_proc0.hdf5"), "r") as file_0: + for slice_name in tqdm(file_0["kinetic/" + species + "/f"]): + # create a new folder for each slice + path_slice = os.path.join(path_distr, slice_name) + os.mkdir(path_slice) + + # Find out all names of slices + slice_names = slice_name.split("_") + + # save grid + for n_gr, (_, grid) in enumerate(file_0["kinetic/" + species + "/f/" + slice_name].attrs.items()): + grid_path = os.path.join( path_slice, - "grid_" + current_slice + ".npy", - ) - - # check if file exists and is in slice_name - if os.path.exists(filename) and current_slice in slice_names: - grid_tot += [xp.load(filename)] - - # otherwise evaluate at zero - else: - grid_tot += [xp.zeros(1)] - - # v-grid - for comp in range(1, f_bckgr.vdim + 1): - current_slice = "v" + str(comp) - filename = os.path.join( - path_slice, - "grid_" + current_slice + ".npy", + "grid_" + slice_names[n_gr] + ".npy", ) + xp.save(grid_path, grid[:]) + + # compute distribution function + for slice_name in tqdm(file_0["kinetic/" + species + "/f"]): + # path to folder of slice + path_slice = os.path.join(path_distr, slice_name) + + # Find out all names of slices + slice_names = slice_name.split("_") + + # load full-f data + data = file_0["kinetic/" + species + "/f/" + slice_name][::step].copy() + data_df = file_0["kinetic/" + species + "/df/" + slice_name][::step].copy() + for rank in range(1, int(nproc)): + with h5py.File(os.path.join(path_in, "data/", f"data_proc{rank}.hdf5"), "r") as file: + data += file["kinetic/" + species + "/f/" + slice_name][::step] + data_df += file["kinetic/" + species + "/df/" + slice_name][::step] + + # save distribution functions + xp.save(os.path.join(path_slice, "f_binned.npy"), data) + xp.save(os.path.join(path_slice, "delta_f_binned.npy"), data_df) + + if compute_bckgr: + # bckgr_params = params["kinetic"][species]["background"] + + # f_bckgr = None + # for fi, maxw_params in bckgr_params.items(): + # if fi[-2] == "_": + # fi_type = fi[:-2] + # else: + # fi_type = fi + + # if f_bckgr is None: + # f_bckgr = getattr(maxwellians, fi_type)( + # maxw_params=maxw_params, + # ) + # else: + # f_bckgr = f_bckgr + getattr(maxwellians, fi_type)( + # maxw_params=maxw_params, + # ) + + spec: ParticleSpecies = getattr(params_in.model, species) + var: PICVariable = spec.var + f_bckgr: KineticBackground = var.backgrounds + + # load all grids of the variables of f + grid_tot = [] + factor = 1.0 + + # eta-grid + for comp in range(1, 4): + current_slice = "e" + str(comp) + filename = os.path.join( + path_slice, + "grid_" + current_slice + ".npy", + ) + + # check if file exists and is in slice_name + if os.path.exists(filename) and current_slice in slice_names: + grid_tot += [xp.load(filename)] + + # otherwise evaluate at zero + else: + grid_tot += [xp.zeros(1)] - # check if file exists and is in slice_name - if os.path.exists(filename) and current_slice in slice_names: - grid_tot += [xp.load(filename)] + # v-grid + for comp in range(1, f_bckgr.vdim + 1): + current_slice = "v" + str(comp) + filename = os.path.join( + path_slice, + "grid_" + current_slice + ".npy", + ) - # otherwise evaluate at zero - else: - grid_tot += [xp.zeros(1)] - # correct integrating out in v-direction, TODO: check for 5D Maxwellians - factor *= xp.sqrt(2 * xp.pi) + # check if file exists and is in slice_name + if os.path.exists(filename) and current_slice in slice_names: + grid_tot += [xp.load(filename)] - grid_eval = xp.meshgrid(*grid_tot, indexing="ij") + # otherwise evaluate at zero + else: + grid_tot += [xp.zeros(1)] + # correct integrating out in v-direction, TODO: check for 5D Maxwellians + factor *= xp.sqrt(2 * xp.pi) - data_bckgr = f_bckgr(*grid_eval).squeeze() + grid_eval = xp.meshgrid(*grid_tot, indexing="ij") - # correct integrating out in v-direction - data_bckgr *= factor + data_bckgr = f_bckgr(*grid_eval).squeeze() - # Now all data is just the data for delta_f - data_delta_f = data_df + # correct integrating out in v-direction + data_bckgr *= factor - # save distribution function - xp.save(os.path.join(path_slice, "delta_f_binned.npy"), data_delta_f) - # add extra axis for data_bckgr since data_delta_f has axis for time series - xp.save( - os.path.join(path_slice, "f_binned.npy"), - data_delta_f + data_bckgr[tuple([None])], - ) + # Now all data is just the data for delta_f + data_delta_f = data_df - # close hdf5 files - for file in files: - file.close() + # save distribution function + xp.save(os.path.join(path_slice, "delta_f_binned.npy"), data_delta_f) + # add extra axis for data_bckgr since data_delta_f has axis for time series + xp.save( + os.path.join(path_slice, "f_binned.npy"), + data_delta_f + data_bckgr[tuple([None])], + ) def post_process_n_sph( @@ -831,19 +789,6 @@ def post_process_n_sph( meta = yaml.load(f, Loader=yaml.FullLoader) nproc = meta["MPI processes"] - # open hdf5 files - files = [ - h5py.File( - os.path.join( - path_in, - "data/", - f"data_proc{i}.hdf5", - ), - "r", - ) - for i in range(int(nproc)) - ] - # directory for .npy files path_n_sph = os.path.join(path_out, "n_sph") @@ -855,34 +800,36 @@ def post_process_n_sph( print("Evaluation of sph density for " + str(species)) - # Create grids - for i, view in enumerate(files[0]["kinetic/" + species + "/n_sph"]): - # create a new folder for each view - path_view = os.path.join(path_n_sph, view) - os.mkdir(path_view) - - # build meshgrid and save - eta1 = files[0]["kinetic/" + species + "/n_sph/" + view].attrs["eta1"] - eta2 = files[0]["kinetic/" + species + "/n_sph/" + view].attrs["eta2"] - eta3 = files[0]["kinetic/" + species + "/n_sph/" + view].attrs["eta3"] - - ee1, ee2, ee3 = xp.meshgrid( - eta1, - eta2, - eta3, - indexing="ij", - ) + with h5py.File(os.path.join(path_in, "data/data_proc0.hdf5"), "r") as file_0: + # Create grids + for i, view in enumerate(file_0["kinetic/" + species + "/n_sph"]): + # create a new folder for each view + path_view = os.path.join(path_n_sph, view) + os.mkdir(path_view) + + # build meshgrid and save + eta1 = file_0["kinetic/" + species + "/n_sph/" + view].attrs["eta1"] + eta2 = file_0["kinetic/" + species + "/n_sph/" + view].attrs["eta2"] + eta3 = file_0["kinetic/" + species + "/n_sph/" + view].attrs["eta3"] + + ee1, ee2, ee3 = xp.meshgrid( + eta1, + eta2, + eta3, + indexing="ij", + ) - grid_path = os.path.join( - path_view, - "grid_n_sph.npy", - ) - xp.save(grid_path, (ee1, ee2, ee3)) + grid_path = os.path.join( + path_view, + "grid_n_sph.npy", + ) + xp.save(grid_path, (ee1, ee2, ee3)) - # load n_sph data - data = files[0]["kinetic/" + species + "/n_sph/" + view][::step].copy() - for rank in range(1, int(nproc)): - data += files[rank]["kinetic/" + species + "/n_sph/" + view][::step] + # load n_sph data + data = file_0["kinetic/" + species + "/n_sph/" + view][::step].copy() + for rank in range(1, int(nproc)): + with h5py.File(os.path.join(path_in, "data/", f"data_proc{rank}.hdf5"), "r") as file: + data += file["kinetic/" + species + "/n_sph/" + view][::step] - # save distribution functions - xp.save(os.path.join(path_view, "n_sph.npy"), data) + # save distribution functions + xp.save(os.path.join(path_view, "n_sph.npy"), data) diff --git a/src/struphy/post_processing/pproc_struphy.py b/src/struphy/post_processing/pproc_struphy.py index 940c94ab3..893b8052b 100644 --- a/src/struphy/post_processing/pproc_struphy.py +++ b/src/struphy/post_processing/pproc_struphy.py @@ -61,32 +61,29 @@ def main( os.mkdir(path_pproc) # check for fields and kinetic data in hdf5 file that need post processing - file = h5py.File(os.path.join(path, "data/", "data_proc0.hdf5"), "r") - - # save time grid at which post-processing data is created - xp.save(os.path.join(path_pproc, "t_grid.npy"), file["time/value"][::step].copy()) - - if "feec" in file.keys(): - exist_fields = True - else: - exist_fields = False - - if "kinetic" in file.keys(): - exist_kinetic = {"markers": False, "f": False, "n_sph": False} - for name in file["kinetic"].keys(): - # check for saved markers - if "markers" in file["kinetic"][name]: - exist_kinetic["markers"] = True - # check for saved distribution function - if "f" in file["kinetic"][name]: - exist_kinetic["f"] = True - # check for saved sph density - if "n_sph" in file["kinetic"][name]: - exist_kinetic["n_sph"] = True - else: - exist_kinetic = None - - file.close() + with h5py.File(os.path.join(path, "data/", "data_proc0.hdf5"), "r") as file: + # save time grid at which post-processing data is created + xp.save(os.path.join(path_pproc, "t_grid.npy"), file["time/value"][::step].copy()) + + if "feec" in file.keys(): + exist_fields = True + else: + exist_fields = False + + if "kinetic" in file.keys(): + exist_kinetic = {"markers": False, "f": False, "n_sph": False} + for name in file["kinetic"].keys(): + # check for saved markers + if "markers" in file["kinetic"][name]: + exist_kinetic["markers"] = True + # check for saved distribution function + if "f" in file["kinetic"][name]: + exist_kinetic["f"] = True + # check for saved sph density + if "n_sph" in file["kinetic"][name]: + exist_kinetic["n_sph"] = True + else: + exist_kinetic = None # import parameters params_in = import_parameters_py(os.path.join(path, "parameters.py"))