From c471c9a418e38fa1a8357073e099ed3827410d9c Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 30 Jul 2026 08:42:53 +0200 Subject: [PATCH 1/6] do not reset particel weights after initial Poisson solve in Vlasov Ampere --- src/struphy/models/vlasov_ampere_one_species.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/struphy/models/vlasov_ampere_one_species.py b/src/struphy/models/vlasov_ampere_one_species.py index 8e84912ea..3f09461ae 100644 --- a/src/struphy/models/vlasov_ampere_one_species.py +++ b/src/struphy/models/vlasov_ampere_one_species.py @@ -185,7 +185,7 @@ def allocate_helpers(self): logger.info("... Done.") # reset particle weights - particles.weights = particles.weights_at_t0.copy() + # particles.weights = particles.weights_at_t0.copy() ## default parameters def generate_default_parameter_file(self, path=None, prompt=True): From cd3890733fd7164e292515af529793c71b0a1c2b Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 30 Jul 2026 13:11:56 +0200 Subject: [PATCH 2/6] adavnce parallel pproc further, not yet done --- .../post_processing/post_processing_tools.py | 394 +++++++++++------- src/struphy/simulation/sim.py | 6 +- 2 files changed, 248 insertions(+), 152 deletions(-) diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index 5063c5566..df73ecaac 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -11,6 +11,7 @@ from feectools.ddm.mpi import mpi as MPI from pyevtk.hl import gridToVTK from tqdm import tqdm +from feectools.ddm.mpi import MockComm from struphy.feec.psydac_derham import Derham, SplineFunction from struphy.fields_background import equils @@ -133,6 +134,7 @@ def __init__( grid = params_in.grid derham_opts = params_in.derham_opts model = params_in.model + sim = params_in.sim elif os.path.exists(bin_path): with open(os.path.join(path, "env.bin"), "rb") as f: @@ -157,6 +159,7 @@ def __init__( with open(os.path.join(path, "model_class.bin"), "rb") as f: model_class: StruphyModel = pickle.load(f) model = model_class() + sim = None else: raise FileNotFoundError(f"Neither of the paths {params_path} or {bin_path} exists.") @@ -170,6 +173,7 @@ def __init__( self.grid = grid self.derham_opts = derham_opts self.model = model + self.sim = sim class PostProcessor: @@ -185,6 +189,8 @@ class PostProcessor: Simulation object of a finished run. If provided, its metadata and output paths are used. path_out : str, optional Path to the Struphy output folder. Required if ``sim`` is not given. + parallel_pproc : bool, optional + Whether to run post-processing in parallel using MPI. Default is False (serial post-processing). Attributes ---------- @@ -206,9 +212,10 @@ def __init__( self, sim: "Simulation" = None, path_out: str = None, + parallel_pproc: bool = False, ): - # create post-processing folder + # import simulation parameters from sim object or from path_out if sim is None: assert path_out is not None, ( "If no sim object is provided, a path_out must be given to retrieve the parameters of the run to post-process." @@ -218,37 +225,62 @@ def __init__( derham_opts = params_in.derham_opts domain = params_in.domain model = params_in.model + imported_sim = params_in.sim else: path_out = sim.env.path_out grid = sim.grid derham_opts = sim.derham_opts domain = sim.domain model = sim.model - + imported_sim = sim + + # get number of MPI ranks used in the simulation from meta.yml with open(os.path.join(path_out, "meta.yml"), "r") as f: meta = yaml.load(f, Loader=yaml.FullLoader) comm_size = meta["MPI processes"] - + self.comm_size = comm_size + + # create post-processing folder self.path_out = path_out self.path_pproc = os.path.join(path_out, "post_processing") - if grid is None or derham_opts is None: - self.derham = None - else: - self.derham = Derham( - grid, - derham_opts, - comm=None, - domain=domain, - ) + + # parallel post-processing (default: False) + self.parallel_pproc = parallel_pproc + + # struphy objects needed for post-processing self.domain = domain self.model = model - self.comm_size = comm_size - try: - os.mkdir(self.path_pproc) - except: - shutil.rmtree(self.path_pproc) - os.mkdir(self.path_pproc) + if self.parallel_pproc: + assert imported_sim is not None, ( + "Parallel post-processing only supported when the sim object is provided." + ) + self.derham = sim.derham + self.comm = self.derham.comm + self.rank = self.comm.Get_rank() + self.range_ranks = range(self.rank, self.rank + 1) + else: + if grid is None or derham_opts is None: + self.derham = None + else: + self.derham = Derham( + grid, + derham_opts, + comm=None, + domain=domain, + ) + self.comm = MockComm() + self.rank = 0 + self.range_ranks = range(int(self.comm_size)) + + # create or remove output paths + if self.rank == 0: + try: + os.mkdir(self.path_pproc) + except: + shutil.rmtree(self.path_pproc) + os.mkdir(self.path_pproc) + self.comm.Barrier() def plot_time_traces(self): path_time_trace = os.path.join(self.path_out, "profiling_time_trace.pkl") @@ -288,8 +320,10 @@ def process( # check for fields and kinetic data in hdf5 file that need post processing with h5py.File(os.path.join(self.path_out, "data/", "data_proc0.hdf5"), "r") as file: - # save time grid at which post-processing data is created - xp.save(os.path.join(self.path_pproc, "t_grid.npy"), file["time/value"][::step].copy()) + + if self.rank == 0: + # save time grid at which post-processing data is created + xp.save(os.path.join(self.path_pproc, "t_grid.npy"), file["time/value"][::step].copy()) if "feec" in file.keys(): self.exist_fields = True @@ -354,40 +388,42 @@ def process_fields( # directory for field data path_fields = os.path.join(self.path_pproc, "fields_data") - try: - os.mkdir(path_fields) - except: - shutil.rmtree(path_fields) - os.mkdir(path_fields) - - # save data dicts for each field - for species, vars in point_data.items(): - for name, val in vars.items(): - try: - os.mkdir(os.path.join(path_fields, species)) - except: - pass + if self.rank == 0: + try: + os.mkdir(path_fields) + except: + shutil.rmtree(path_fields) + os.mkdir(path_fields) - with open(os.path.join(path_fields, species, name + "_log.bin"), "wb") as handle: - pickle.dump(val, handle, protocol=pickle.HIGHEST_PROTOCOL) + # save data dicts for each field + for species, vars in point_data.items(): + for name, val in vars.items(): + try: + os.mkdir(os.path.join(path_fields, species)) + except: + pass - if physical: - with open(os.path.join(path_fields, species, name + "_phy.bin"), "wb") as handle: - pickle.dump(point_data_phy[species][name], handle, protocol=pickle.HIGHEST_PROTOCOL) + with open(os.path.join(path_fields, species, name + "_log.bin"), "wb") as handle: + pickle.dump(val, handle, protocol=pickle.HIGHEST_PROTOCOL) - # save grids - with open(os.path.join(path_fields, "grids_log.bin"), "wb") as handle: - pickle.dump(grids_log, handle, protocol=pickle.HIGHEST_PROTOCOL) + if physical: + with open(os.path.join(path_fields, species, name + "_phy.bin"), "wb") as handle: + pickle.dump(point_data_phy[species][name], handle, protocol=pickle.HIGHEST_PROTOCOL) - with open(os.path.join(path_fields, "grids_phy.bin"), "wb") as handle: - pickle.dump(grids_phy, handle, protocol=pickle.HIGHEST_PROTOCOL) + # save grids + with open(os.path.join(path_fields, "grids_log.bin"), "wb") as handle: + pickle.dump(grids_log, handle, protocol=pickle.HIGHEST_PROTOCOL) - # create vtk files - if create_vtk: - self._create_vtk(path_fields, t_grid, grids_phy, point_data) - if physical: - self._create_vtk(path_fields, t_grid, grids_phy, point_data_phy, physical=True) + with open(os.path.join(path_fields, "grids_phy.bin"), "wb") as handle: + pickle.dump(grids_phy, handle, protocol=pickle.HIGHEST_PROTOCOL) + # create vtk files + if create_vtk: + self._create_vtk(path_fields, t_grid, grids_phy, point_data) + if physical: + self._create_vtk(path_fields, t_grid, grids_phy, point_data_phy, physical=True) + self.comm.Barrier() + def process_particles( self, step: int = 1, @@ -402,22 +438,26 @@ def process_particles( # directory for kinetic data path_kinetics = os.path.join(self.path_pproc, "kinetic_data") - try: - os.mkdir(path_kinetics) - except: - shutil.rmtree(path_kinetics) - os.mkdir(path_kinetics) + if self.rank == 0: + try: + os.mkdir(path_kinetics) + except: + shutil.rmtree(path_kinetics) + os.mkdir(path_kinetics) + self.comm.Barrier() # kinetic post-processing for each species for n, species in enumerate(self.kinetic_species): # directory for each species path_kinetics_species = os.path.join(path_kinetics, species) - try: - os.mkdir(path_kinetics_species) - except: - shutil.rmtree(path_kinetics_species) - os.mkdir(path_kinetics_species) + if self.rank == 0: + try: + os.mkdir(path_kinetics_species) + except: + shutil.rmtree(path_kinetics_species) + os.mkdir(path_kinetics_species) + self.comm.Barrier() # markers if self.exist_particles["markers"]: @@ -446,12 +486,12 @@ def process_particles( compute_bckgr=compute_bckgr, ) - # sph density - if self.exist_particles["n_sph"]: - self._post_process_n_sph( - path_kinetics_species, - step, - ) + # # sph density + # if self.exist_particles["n_sph"]: + # self._post_process_n_sph( + # path_kinetics_species, + # step, + # ) def _create_femfields(self, step: int = 1): """Reconstruct FEEC spline field objects from HDF5 output files. @@ -498,8 +538,8 @@ def _create_femfields(self, step: int = 1): ) # get hdf5 data - logger.warning("") - for rank in range(int(self.comm_size)): + logger.warning("") + for rank in self.range_ranks: # open hdf5 file with h5py.File(os.path.join(self.path_out, "data/", f"data_proc{rank}.hdf5"), "r") as file: for species, dset in file["feec"].items(): @@ -615,67 +655,75 @@ def _eval_femfields( assert isinstance(field, SplineFunction) space_id = field.space_id - # field evaluation - temp_val = field(*grids_log) + # evaluate field locally on rank and send to rank 0 for collection + temp_val = field(*grids_log, local=True) + if self.parallel_pproc: + self.comm.reduce( + temp_val, + op=MPI.SUM, + root=0, + ) + # point_data will stay an empty list on all ranks except rank 0 point_data[species][name][t] = [] - # scalar spaces - if isinstance(temp_val, xp.ndarray): - if physical: - # push-forward - if space_id == "H1": - point_data[species][name][t].append( - self.domain.push( - temp_val, - *grids_log, - kind="0", - ), - ) - elif space_id == "L2": - point_data[species][name][t].append( - self.domain.push( - temp_val, - *grids_log, - kind="3", - ), - ) - - else: - point_data[species][name][t].append(temp_val) - - # vector-valued spaces - else: - for j in range(3): + if self.rank == 0: + # scalar spaces + if isinstance(temp_val, xp.ndarray): if physical: # push-forward - if space_id == "Hcurl": + if space_id == "H1": point_data[species][name][t].append( self.domain.push( temp_val, *grids_log, - kind="1", - )[j], + kind="0", + ), ) - elif space_id == "Hdiv": + elif space_id == "L2": point_data[species][name][t].append( self.domain.push( temp_val, *grids_log, - kind="2", - )[j], - ) - elif space_id == "H1vec": - point_data[species][name][t].append( - self.domain.push( - temp_val, - *grids_log, - kind="v", - )[j], + kind="3", + ), ) else: - point_data[species][name][t].append(temp_val[j]) + point_data[species][name][t].append(temp_val) + + # vector-valued spaces + else: + for j in range(3): + if physical: + # push-forward + if space_id == "Hcurl": + point_data[species][name][t].append( + self.domain.push( + temp_val, + *grids_log, + kind="1", + )[j], + ) + elif space_id == "Hdiv": + point_data[species][name][t].append( + self.domain.push( + temp_val, + *grids_log, + kind="2", + )[j], + ) + elif space_id == "H1vec": + point_data[species][name][t].append( + self.domain.push( + temp_val, + *grids_log, + kind="v", + )[j], + ) + + else: + point_data[species][name][t].append(temp_val[j]) return point_data, grids_log, grids_phy @@ -788,11 +836,13 @@ def _post_process_markers( else: save_index = list(range(0, 4)) + [-1] - try: - os.mkdir(path_orbits) - except: - shutil.rmtree(path_orbits) - os.mkdir(path_orbits) + if self.rank == 0: + try: + os.mkdir(path_orbits) + except: + shutil.rmtree(path_orbits) + os.mkdir(path_orbits) + self.comm.Barrier() # temporary array temp = xp.empty((n_markers, len(save_index)), order="C") @@ -815,12 +865,19 @@ def _post_process_markers( species + "_{0:0{1}d}.txt".format(n, log_nt), ) - for i in range(int(self.comm_size)): - with h5py.File(os.path.join(self.path_out, "data/", f"data_proc{i}.hdf5"), "r") as file: + for rank in self.range_ranks: + with h5py.File(os.path.join(self.path_out, "data/", f"data_proc{rank}.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] + + if self.parallel_pproc: + self.comm.reduce( + temp, + op=MPI.SUM, + root=0, + ) # sorting out lost particles ids = temp[:, -1].astype("int") @@ -841,11 +898,13 @@ def _post_process_markers( pos_phys = self.domain(xp.array(temp[~lost_particles_mask, :3]), change_out_order=True) temp[~lost_particles_mask, :3] = pos_phys - # save numpy - xp.save(file_npy, temp) - # move ids to first column and save txt - temp = xp.roll(temp, 1, axis=1) - xp.savetxt(file_txt, temp[:, (0, 1, 2, 3, -1)], fmt="%12.6f", delimiter=", ") + if self.rank == 0: + # save numpy + xp.save(file_npy, temp) + # move ids to first column and save txt + temp = xp.roll(temp, 1, axis=1) + xp.savetxt(file_txt, temp[:, (0, 1, 2, 3, -1)], fmt="%12.6f", delimiter=", ") + self.comm.Barrier() def _post_process_f( self, @@ -869,54 +928,91 @@ def _post_process_f( compute_bckgr : bool, optional If True, compute and add background contribution to the saved binned data. """ + print(f"{self.rank} starting post-processing of distribution functions for {path_kinetic_species} ...") + species = path_kinetic_species.split("/")[-1] species_obj: ParticleSpecies = self.model.particle_species[species] # directory for .npy files path_distr = os.path.join(path_kinetic_species, "distribution_function") - try: - os.mkdir(path_distr) - except: - shutil.rmtree(path_distr) - os.mkdir(path_distr) + if self.rank == 0: + try: + os.mkdir(path_distr) + except: + shutil.rmtree(path_distr) + os.mkdir(path_distr) + self.comm.Barrier() logger.warning("Evaluation of distribution functions for " + str(species)) # Create grids with h5py.File(os.path.join(self.path_out, "data/data_proc0.hdf5"), "r") as file_0: + slice_names = [] for slice_name in tqdm(file_0["kinetic/" + species + "/f"]): + slice_names += [slice_name] # create a new folder for each slice path_slice = os.path.join(path_distr, slice_name) - os.mkdir(path_slice) + if self.rank == 0: + os.mkdir(path_slice) + self.comm.Barrier() # Find out all names of slices - slice_names = slice_name.split("_") + slice_splits = 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_" + slice_names[n_gr] + ".npy", + "grid_" + slice_splits[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(self.comm_size)): - with h5py.File(os.path.join(self.path_out, "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] + if self.rank == 0: + xp.save(grid_path, grid[:]) + self.comm.Barrier() + + # compute distribution function + for slice_name in tqdm(slice_names): + logger.info(f"Processing slice {slice_name} for species {species}") + # path to folder of slice + path_slice = os.path.join(path_distr, slice_name) + + # Find out all names of slices + slice_splits = slice_name.split("_") + + for rank in self.range_ranks: + print(f"{rank = } ----------------------------") + with h5py.File(os.path.join(self.path_out, "data/", f"data_proc{rank}.hdf5"), "r") as file: + if self.parallel_pproc: + data = file["kinetic/" + species + "/f/" + slice_name][::step] + data_df = file["kinetic/" + species + "/df/" + slice_name][::step] + else: + if rank == 0: + data = file["kinetic/" + species + "/f/" + slice_name][::step].copy() + data_df = file["kinetic/" + species + "/df/" + slice_name][::step].copy() + self.comm.Barrier() + + if rank != 0: + data += file["kinetic/" + species + "/f/" + slice_name][::step] + data_df += file["kinetic/" + species + "/df/" + slice_name][::step] + + print(f"{self.rank =} with {xp.sum(data) =} and {xp.sum(data_df) =}") + + if self.parallel_pproc: + self.comm.reduce( + data, + op=MPI.SUM, + root=0, + ) + self.comm.reduce( + data_df, + op=MPI.SUM, + root=0, + ) + + print(f"{self.rank =} with {xp.sum(data) =} and {xp.sum(data_df) =}") + print(f"{self.rank =} done.") + if self.rank == 0: # 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) @@ -958,7 +1054,7 @@ def _post_process_f( ) # check if file exists and is in slice_name - if os.path.exists(filename) and current_slice in slice_names: + if os.path.exists(filename) and current_slice in slice_splits: grid_tot += [xp.load(filename)] # otherwise evaluate at zero @@ -974,7 +1070,7 @@ def _post_process_f( ) # check if file exists and is in slice_name - if os.path.exists(filename) and current_slice in slice_names: + if os.path.exists(filename) and current_slice in slice_splits: grid_tot += [xp.load(filename)] # otherwise evaluate at zero diff --git a/src/struphy/simulation/sim.py b/src/struphy/simulation/sim.py index ad524c0a3..e5504328b 100644 --- a/src/struphy/simulation/sim.py +++ b/src/struphy/simulation/sim.py @@ -474,7 +474,7 @@ def run(self, one_time_step: bool = False): If True, only perform one time step (useful for testing). """ - logger.warning(f"\nStarting run for model {self.model_name} ...") + logger.warning(f"\nStarting run for model {self.model_name} on {self.comm_size} ranks ...") if self.name != "": logger.info(f"Simulation name: {self.name}") if self.description != "": @@ -694,6 +694,7 @@ def pproc( classify: bool = False, create_vtk: bool = True, time_trace: bool = False, + parallel_pproc: bool = False, ): """Run post-processing on saved simulation data. @@ -702,8 +703,7 @@ def pproc( """ # setup post processor and plotting - if not hasattr(self, "_post_processor") and self.rank == 0: - self._post_processor = PostProcessor(sim=self) + self._post_processor = PostProcessor(sim=self, parallel_pproc=parallel_pproc) if time_trace: self.post_processor.plot_time_traces() From 86ba7cfe68d63b8ae5b4013aa4dee6d66789fbd7 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 30 Jul 2026 15:58:45 +0200 Subject: [PATCH 3/6] parallel pproc works for feec and f, not yet n_sph --- .../post_processing/post_processing_tools.py | 100 +++++++++++++----- 1 file changed, 74 insertions(+), 26 deletions(-) diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index df73ecaac..194c5b8e6 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -234,12 +234,6 @@ def __init__( model = sim.model imported_sim = sim - # get number of MPI ranks used in the simulation from meta.yml - with open(os.path.join(path_out, "meta.yml"), "r") as f: - meta = yaml.load(f, Loader=yaml.FullLoader) - comm_size = meta["MPI processes"] - self.comm_size = comm_size - # create post-processing folder self.path_out = path_out self.path_pproc = os.path.join(path_out, "post_processing") @@ -257,6 +251,7 @@ def __init__( ) self.derham = sim.derham self.comm = self.derham.comm + self.comm_size = self.comm.Get_size() self.rank = self.comm.Get_rank() self.range_ranks = range(self.rank, self.rank + 1) else: @@ -270,6 +265,10 @@ def __init__( domain=domain, ) self.comm = MockComm() + # get number of MPI ranks used in the simulation from meta.yml + with open(os.path.join(path_out, "meta.yml"), "r") as f: + meta = yaml.load(f, Loader=yaml.FullLoader) + self.comm_size = meta["MPI processes"] self.rank = 0 self.range_ranks = range(int(self.comm_size)) @@ -657,12 +656,43 @@ def _eval_femfields( # evaluate field locally on rank and send to rank 0 for collection temp_val = field(*grids_log, local=True) + print(f"{type(temp_val) = }") if self.parallel_pproc: - self.comm.reduce( - temp_val, - op=MPI.SUM, - root=0, - ) + if isinstance(temp_val, xp.ndarray): + print(f"{self.rank = } before \n{temp_val.shape = }\n{temp_val = }") + if self.rank == 0: + self.comm.Reduce( + MPI.IN_PLACE, + temp_val, + op=MPI.SUM, + root=0, + ) + else: + self.comm.Reduce( + temp_val, + None, + op=MPI.SUM, + root=0, + ) + print(f"{self.rank = } after\n {temp_val = }") + else: + for j in range(3): + print(f"{self.rank = } before \n{temp_val[j].shape = }\n{temp_val[j] = }") + if self.rank == 0: + self.comm.Reduce( + MPI.IN_PLACE, + temp_val[j], + op=MPI.SUM, + root=0, + ) + else: + self.comm.Reduce( + temp_val[j], + None, + op=MPI.SUM, + root=0, + ) + print(f"{self.rank = } after\n {temp_val[j] = }") # point_data will stay an empty list on all ranks except rank 0 point_data[species][name][t] = [] @@ -998,16 +1028,32 @@ def _post_process_f( print(f"{self.rank =} with {xp.sum(data) =} and {xp.sum(data_df) =}") if self.parallel_pproc: - self.comm.reduce( - data, - op=MPI.SUM, - root=0, - ) - self.comm.reduce( - data_df, - op=MPI.SUM, - root=0, - ) + if self.rank == 0: + self.comm.Reduce( + MPI.IN_PLACE, + data, + op=MPI.SUM, + root=0, + ) + self.comm.Reduce( + MPI.IN_PLACE, + data_df, + op=MPI.SUM, + root=0, + ) + else: + self.comm.Reduce( + data, + None, + op=MPI.SUM, + root=0, + ) + self.comm.Reduce( + data_df, + None, + op=MPI.SUM, + root=0, + ) print(f"{self.rank =} with {xp.sum(data) =} and {xp.sum(data_df) =}") @@ -1116,11 +1162,13 @@ def _post_process_n_sph( # directory for .npy files path_n_sph = os.path.join(path_kinetic_species, "n_sph") - try: - os.mkdir(path_n_sph) - except: - shutil.rmtree(path_n_sph) - os.mkdir(path_n_sph) + if self.rank == 0: + try: + os.mkdir(path_n_sph) + except: + shutil.rmtree(path_n_sph) + os.mkdir(path_n_sph) + self.comm.Barrier() logger.warning("Evaluation of sph density for " + str(species)) From 0ada74da39fb6605c96589bb9885b64e00f817a7 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 30 Jul 2026 16:03:19 +0200 Subject: [PATCH 4/6] add n_sph parallel --- .../post_processing/post_processing_tools.py | 72 +++++++++++++------ 1 file changed, 50 insertions(+), 22 deletions(-) diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index 194c5b8e6..dae5e9592 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -485,12 +485,12 @@ def process_particles( compute_bckgr=compute_bckgr, ) - # # sph density - # if self.exist_particles["n_sph"]: - # self._post_process_n_sph( - # path_kinetics_species, - # step, - # ) + # sph density + if self.exist_particles["n_sph"]: + self._post_process_n_sph( + path_kinetics_species, + step, + ) def _create_femfields(self, step: int = 1): """Reconstruct FEEC spline field objects from HDF5 output files. @@ -1019,9 +1019,7 @@ def _post_process_f( if rank == 0: data = file["kinetic/" + species + "/f/" + slice_name][::step].copy() data_df = file["kinetic/" + species + "/df/" + slice_name][::step].copy() - self.comm.Barrier() - - if rank != 0: + else: data += file["kinetic/" + species + "/f/" + slice_name][::step] data_df += file["kinetic/" + species + "/df/" + slice_name][::step] @@ -1173,11 +1171,15 @@ def _post_process_n_sph( logger.warning("Evaluation of sph density for " + str(species)) with h5py.File(os.path.join(self.path_out, "data/data_proc0.hdf5"), "r") as file_0: + views = list(file_0["kinetic/" + species + "/n_sph"]) + # Create grids - for i, view in enumerate(file_0["kinetic/" + species + "/n_sph"]): + for view in views: # create a new folder for each view path_view = os.path.join(path_n_sph, view) - os.mkdir(path_view) + if self.rank == 0: + os.mkdir(path_view) + self.comm.Barrier() # build meshgrid and save eta1 = file_0["kinetic/" + species + "/n_sph/" + view].attrs["eta1"] @@ -1191,19 +1193,45 @@ def _post_process_n_sph( indexing="ij", ) - grid_path = os.path.join( - path_view, - "grid_n_sph.npy", - ) - xp.save(grid_path, (ee1, ee2, ee3)) + if self.rank == 0: + grid_path = os.path.join( + path_view, + "grid_n_sph.npy", + ) + xp.save(grid_path, (ee1, ee2, ee3)) - # load n_sph data - data = file_0["kinetic/" + species + "/n_sph/" + view][::step].copy() - for rank in range(1, int(self.comm_size)): - with h5py.File(os.path.join(self.path_out, "data/", f"data_proc{rank}.hdf5"), "r") as file: - data += file["kinetic/" + species + "/n_sph/" + view][::step] + # compute sph density + for view in tqdm(views): + path_view = os.path.join(path_n_sph, view) - # save distribution functions + for rank in self.range_ranks: + with h5py.File(os.path.join(self.path_out, "data/", f"data_proc{rank}.hdf5"), "r") as file: + if self.parallel_pproc: + data = file["kinetic/" + species + "/n_sph/" + view][::step] + else: + if rank == 0: + data = file["kinetic/" + species + "/n_sph/" + view][::step].copy() + else: + data += file["kinetic/" + species + "/n_sph/" + view][::step] + + if self.parallel_pproc: + if self.rank == 0: + self.comm.Reduce( + MPI.IN_PLACE, + data, + op=MPI.SUM, + root=0, + ) + else: + self.comm.Reduce( + data, + None, + op=MPI.SUM, + root=0, + ) + + if self.rank == 0: + # save sph density xp.save(os.path.join(path_view, "n_sph.npy"), data) From 756027b8db249de350d2fa1d6856bcab4e073e34 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 30 Jul 2026 16:04:33 +0200 Subject: [PATCH 5/6] add unit test --- .../post_processing/post_processing_tools.py | 43 ++++--- .../post_processing/tests/test_pproc.py | 113 ++++++++++++++++++ 2 files changed, 133 insertions(+), 23 deletions(-) create mode 100644 src/struphy/post_processing/tests/test_pproc.py diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index dae5e9592..319f18a07 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -8,10 +8,10 @@ import cunumpy as xp import h5py import yaml +from feectools.ddm.mpi import MockComm from feectools.ddm.mpi import mpi as MPI from pyevtk.hl import gridToVTK from tqdm import tqdm -from feectools.ddm.mpi import MockComm from struphy.feec.psydac_derham import Derham, SplineFunction from struphy.fields_background import equils @@ -159,7 +159,7 @@ def __init__( with open(os.path.join(path, "model_class.bin"), "rb") as f: model_class: StruphyModel = pickle.load(f) model = model_class() - sim = None + sim = None else: raise FileNotFoundError(f"Neither of the paths {params_path} or {bin_path} exists.") @@ -233,22 +233,20 @@ def __init__( domain = sim.domain model = sim.model imported_sim = sim - + # create post-processing folder self.path_out = path_out self.path_pproc = os.path.join(path_out, "post_processing") - + # parallel post-processing (default: False) self.parallel_pproc = parallel_pproc - + # struphy objects needed for post-processing self.domain = domain self.model = model if self.parallel_pproc: - assert imported_sim is not None, ( - "Parallel post-processing only supported when the sim object is provided." - ) + assert imported_sim is not None, "Parallel post-processing only supported when the sim object is provided." self.derham = sim.derham self.comm = self.derham.comm self.comm_size = self.comm.Get_size() @@ -259,11 +257,11 @@ def __init__( self.derham = None else: self.derham = Derham( - grid, - derham_opts, - comm=None, - domain=domain, - ) + grid, + derham_opts, + comm=None, + domain=domain, + ) self.comm = MockComm() # get number of MPI ranks used in the simulation from meta.yml with open(os.path.join(path_out, "meta.yml"), "r") as f: @@ -319,7 +317,6 @@ def process( # check for fields and kinetic data in hdf5 file that need post processing with h5py.File(os.path.join(self.path_out, "data/", "data_proc0.hdf5"), "r") as file: - if self.rank == 0: # save time grid at which post-processing data is created xp.save(os.path.join(self.path_pproc, "t_grid.npy"), file["time/value"][::step].copy()) @@ -421,8 +418,8 @@ def process_fields( self._create_vtk(path_fields, t_grid, grids_phy, point_data) if physical: self._create_vtk(path_fields, t_grid, grids_phy, point_data_phy, physical=True) - self.comm.Barrier() - + self.comm.Barrier() + def process_particles( self, step: int = 1, @@ -537,7 +534,7 @@ def _create_femfields(self, step: int = 1): ) # get hdf5 data - logger.warning("") + logger.warning("") for rank in self.range_ranks: # open hdf5 file with h5py.File(os.path.join(self.path_out, "data/", f"data_proc{rank}.hdf5"), "r") as file: @@ -901,7 +898,7 @@ def _post_process_markers( ids = markers[n * step, :, -1].astype("int") ids = ids[ids != -1] # exclude holes temp[ids] = markers[n * step, : ids.size, save_index] - + if self.parallel_pproc: self.comm.reduce( temp, @@ -959,7 +956,7 @@ def _post_process_f( If True, compute and add background contribution to the saved binned data. """ print(f"{self.rank} starting post-processing of distribution functions for {path_kinetic_species} ...") - + species = path_kinetic_species.split("/")[-1] species_obj: ParticleSpecies = self.model.particle_species[species] @@ -999,7 +996,7 @@ def _post_process_f( if self.rank == 0: xp.save(grid_path, grid[:]) self.comm.Barrier() - + # compute distribution function for slice_name in tqdm(slice_names): logger.info(f"Processing slice {slice_name} for species {species}") @@ -1022,9 +1019,9 @@ def _post_process_f( else: data += file["kinetic/" + species + "/f/" + slice_name][::step] data_df += file["kinetic/" + species + "/df/" + slice_name][::step] - + print(f"{self.rank =} with {xp.sum(data) =} and {xp.sum(data_df) =}") - + if self.parallel_pproc: if self.rank == 0: self.comm.Reduce( @@ -1052,7 +1049,7 @@ def _post_process_f( op=MPI.SUM, root=0, ) - + print(f"{self.rank =} with {xp.sum(data) =} and {xp.sum(data_df) =}") print(f"{self.rank =} done.") diff --git a/src/struphy/post_processing/tests/test_pproc.py b/src/struphy/post_processing/tests/test_pproc.py new file mode 100644 index 000000000..af1184468 --- /dev/null +++ b/src/struphy/post_processing/tests/test_pproc.py @@ -0,0 +1,113 @@ +import logging +from pathlib import Path + +import numpy as np +import pytest +from feectools.ddm.mpi import MockComm +from feectools.ddm.mpi import mpi as MPI +from matplotlib import pyplot as plt + +from struphy import Simulation, set_logging_level +from struphy.io.setup import import_parameters_py + +set_logging_level(logging.WARNING) +logger = logging.getLogger("struphy") + +PARAMS_PATH = ( + Path(__file__).resolve().parents[4] + / "examples" + / "VlasovAmpereOneSpecies" + / "weak_Landau_damping" + / "params_weak_Landau_damping.py" +) + + +@pytest.mark.mpi(min_size=2) +def test_pproc_mpi(): + + def do_plotting(sim: Simulation, from_parallel=False): + sim.load_plotting_data() + + t_grid = sim.t_grid + eta1 = sim.grids_log[0] + e_field = sim.spline_values.em_fields.e_field_log + phi = sim.spline_values.em_fields.phi_log + + f = sim.f.kinetic_ions.e1_v1_density + print(f.__dict__.keys()) + bins_e1 = f.grid_e1 + bins_v1 = f.grid_v1 + f_binned = f.f_binned + df_binned = f.delta_f_binned + print(f"{f_binned.shape=}") + + if from_parallel: + extra = " (from parallel pproc)" + else: + extra = "" + + n = 0 # time index + + plt.figure(figsize=(12, 12)) + plt.subplot(2, 2, 1) + plt.plot(eta1, e_field.data[t_grid[n]][0][:, 0, 0], label="Ex") + plt.title(f"Ex at t={t_grid[n]} on rank 0{extra}") + plt.xlabel("$\\eta1$") + plt.ylabel("Ex") + plt.legend() + + plt.subplot(2, 2, 2) + plt.plot(eta1, phi.data[t_grid[n]][0][:, 0, 0], label="phi") + plt.title(f"phi at t={t_grid[n]} on rank 0{extra}") + plt.xlabel("$\\eta1$") + plt.ylabel("phi") + plt.legend() + + plt.subplot(2, 2, 3) + plt.pcolor(bins_e1, bins_v1, f_binned[n].T, shading="auto") + plt.title(f"full f at t={t_grid[n]} on rank 0{extra}") + plt.xlabel("$\\eta1$") + plt.ylabel("$v_x$") + + plt.subplot(2, 2, 4) + plt.pcolor(bins_e1, bins_v1, df_binned[n].T, shading="auto") + plt.title(f"delta f at t={t_grid[n]} on rank 0{extra}") + plt.xlabel("$\\eta1$") + plt.ylabel("$v_x$") + + return ( + e_field.data[t_grid[n]][0][:, 0, 0], + phi.data[t_grid[n]][0][:, 0, 0], + f_binned[n].T, + df_binned[n].T, + ) + + params = import_parameters_py(str(PARAMS_PATH), name="weak_Landau_damping") + + sim: Simulation = params.sim + + sim.run(one_time_step=True) + + if MPI.COMM_WORLD.Get_rank() == 0: + # serial pproc + sim.pproc() + r1, r2, r3, r4 = do_plotting(sim) + MPI.COMM_WORLD.Barrier() + + # parallel pproc + sim.pproc(parallel_pproc=True) + + # plot and compare results from serial and parallel pproc + if MPI.COMM_WORLD.Get_rank() == 0: + r1_mpi, r2_mpi, r3_mpi, r4_mpi = do_plotting(sim, from_parallel=True) + plt.show() + + assert np.allclose(r1, r1_mpi) + assert np.allclose(r2, r2_mpi) + assert np.allclose(r3, r3_mpi) + assert np.allclose(r4, r4_mpi) + print("All checks passed for parallel pproc vs serial pproc.") + + +if __name__ == "__main__": + test_pproc_mpi() From ec8f0f7178430f50529d00983f3d19d72a8e6ca5 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 31 Jul 2026 07:42:41 +0200 Subject: [PATCH 6/6] apply Copilot suggestions --- .../models/vlasov_ampere_one_species.py | 2 +- .../post_processing/post_processing_tools.py | 16 +++---- .../post_processing/tests/test_pproc.py | 1 - src/struphy/simulation/sim.py | 42 +++++++++++++------ 4 files changed, 35 insertions(+), 26 deletions(-) diff --git a/src/struphy/models/vlasov_ampere_one_species.py b/src/struphy/models/vlasov_ampere_one_species.py index 3f09461ae..8e84912ea 100644 --- a/src/struphy/models/vlasov_ampere_one_species.py +++ b/src/struphy/models/vlasov_ampere_one_species.py @@ -185,7 +185,7 @@ def allocate_helpers(self): logger.info("... Done.") # reset particle weights - # particles.weights = particles.weights_at_t0.copy() + particles.weights = particles.weights_at_t0.copy() ## default parameters def generate_default_parameter_file(self, path=None, prompt=True): diff --git a/src/struphy/post_processing/post_processing_tools.py b/src/struphy/post_processing/post_processing_tools.py index 319f18a07..1a2e274ab 100644 --- a/src/struphy/post_processing/post_processing_tools.py +++ b/src/struphy/post_processing/post_processing_tools.py @@ -247,7 +247,7 @@ def __init__( if self.parallel_pproc: assert imported_sim is not None, "Parallel post-processing only supported when the sim object is provided." - self.derham = sim.derham + self.derham = imported_sim.derham self.comm = self.derham.comm self.comm_size = self.comm.Get_size() self.rank = self.comm.Get_rank() @@ -653,10 +653,8 @@ def _eval_femfields( # evaluate field locally on rank and send to rank 0 for collection temp_val = field(*grids_log, local=True) - print(f"{type(temp_val) = }") if self.parallel_pproc: if isinstance(temp_val, xp.ndarray): - print(f"{self.rank = } before \n{temp_val.shape = }\n{temp_val = }") if self.rank == 0: self.comm.Reduce( MPI.IN_PLACE, @@ -671,10 +669,8 @@ def _eval_femfields( op=MPI.SUM, root=0, ) - print(f"{self.rank = } after\n {temp_val = }") else: for j in range(3): - print(f"{self.rank = } before \n{temp_val[j].shape = }\n{temp_val[j] = }") if self.rank == 0: self.comm.Reduce( MPI.IN_PLACE, @@ -689,7 +685,6 @@ def _eval_femfields( op=MPI.SUM, root=0, ) - print(f"{self.rank = } after\n {temp_val[j] = }") # point_data will stay an empty list on all ranks except rank 0 point_data[species][name][t] = [] @@ -900,11 +895,10 @@ def _post_process_markers( temp[ids] = markers[n * step, : ids.size, save_index] if self.parallel_pproc: - self.comm.reduce( - temp, - op=MPI.SUM, - root=0, - ) + if self.rank == 0: + self.comm.Reduce(MPI.IN_PLACE, temp, op=MPI.SUM, root=0) + else: + self.comm.Reduce(temp, None, op=MPI.SUM, root=0) # sorting out lost particles ids = temp[:, -1].astype("int") diff --git a/src/struphy/post_processing/tests/test_pproc.py b/src/struphy/post_processing/tests/test_pproc.py index af1184468..93809a4c2 100644 --- a/src/struphy/post_processing/tests/test_pproc.py +++ b/src/struphy/post_processing/tests/test_pproc.py @@ -3,7 +3,6 @@ import numpy as np import pytest -from feectools.ddm.mpi import MockComm from feectools.ddm.mpi import mpi as MPI from matplotlib import pyplot as plt diff --git a/src/struphy/simulation/sim.py b/src/struphy/simulation/sim.py index e5504328b..bb25f9362 100644 --- a/src/struphy/simulation/sim.py +++ b/src/struphy/simulation/sim.py @@ -703,19 +703,35 @@ def pproc( """ # setup post processor and plotting - self._post_processor = PostProcessor(sim=self, parallel_pproc=parallel_pproc) - - if time_trace: - self.post_processor.plot_time_traces() - - self.post_processor.process( - step=step, - celldivide=celldivide, - physical=physical, - guiding_center=guiding_center, - classify=classify, - create_vtk=create_vtk, - ) + if parallel_pproc: + self._post_processor = PostProcessor(sim=self, parallel_pproc=True) + + if time_trace: + self.post_processor.plot_time_traces() + + self.post_processor.process( + step=step, + celldivide=celldivide, + physical=physical, + guiding_center=guiding_center, + classify=classify, + create_vtk=create_vtk, + ) + else: + if self.rank == 0: + self._post_processor = PostProcessor(sim=self, parallel_pproc=False) + + if time_trace: + self.post_processor.plot_time_traces() + + self.post_processor.process( + step=step, + celldivide=celldivide, + physical=physical, + guiding_center=guiding_center, + classify=classify, + create_vtk=create_vtk, + ) def load_plotting_data(self): """Load plotting datasets produced by post-processing.