diff --git a/openmc/deplete/microxs.py b/openmc/deplete/microxs.py index e351c923df0..84497b616b4 100644 --- a/openmc/deplete/microxs.py +++ b/openmc/deplete/microxs.py @@ -170,13 +170,14 @@ def get_microxs_and_flux( # Reinitialize with tallies openmc.lib.init(intracomm=comm) - # create temporary run with TemporaryDirectory() as temp_dir: - if run_kwargs is None: - run_kwargs = {} - else: - run_kwargs = dict(run_kwargs) - run_kwargs.setdefault('cwd', temp_dir) + # Indicate to run in temporary directory unless being executed through + # openmc.lib, in which case we don't need to specify the cwd + run_kwargs = dict(run_kwargs) if run_kwargs else {} + if not openmc.lib.is_initialized: + run_kwargs.setdefault('cwd', temp_dir) + + # Run transport simulation statepoint_path = model.run(**run_kwargs) if comm.rank == 0: @@ -189,15 +190,18 @@ def get_microxs_and_flux( if path_input is not None: model.export_to_model_xml(path_input) - with StatePoint(statepoint_path) as sp: - if reaction_rate_mode == 'direct': - rr_tally = sp.tallies[rr_tally.id] - rr_tally._read_results() - flux_tally = sp.tallies[flux_tally.id] - flux_tally._read_results() + # Broadcast updated statepoint path to all ranks + statepoint_path = comm.bcast(statepoint_path) + + # Read in tally results (on all ranks) + with StatePoint(statepoint_path) as sp: + if reaction_rate_mode == 'direct': + rr_tally = sp.tallies[rr_tally.id] + rr_tally._read_results() + flux_tally = sp.tallies[flux_tally.id] + flux_tally._read_results() # Get flux values and make energy groups last dimension - flux_tally = comm.bcast(flux_tally) flux = flux_tally.get_reshaped_data() # (domains, groups, 1, 1) flux = np.moveaxis(flux, 1, -1) # (domains, 1, 1, groups) @@ -206,7 +210,6 @@ def get_microxs_and_flux( if reaction_rate_mode == 'direct': # Get reaction rates - rr_tally = comm.bcast(rr_tally) reaction_rates = rr_tally.get_reshaped_data() # (domains, groups, nuclides, reactions) # Make energy groups last dimension diff --git a/openmc/deplete/r2s.py b/openmc/deplete/r2s.py index 7b5deddc23c..7f3e94edd96 100644 --- a/openmc/deplete/r2s.py +++ b/openmc/deplete/r2s.py @@ -11,6 +11,9 @@ from .microxs import get_microxs_and_flux, write_microxs_hdf5, read_microxs_hdf5 from .results import Results from ..checkvalue import PathLike +from ..mpi import comm +from openmc.lib import TemporarySession +from openmc.utility_funcs import change_directory def get_activation_materials( @@ -199,8 +202,10 @@ def run( """ if output_dir is None: + # Create timestamped output directory and broadcast to all ranks for + # consistency (different ranks may have slightly different times) stamp = datetime.now().strftime('%Y-%m-%dT%H-%M-%S') - output_dir = Path(f'r2s_{stamp}') + output_dir = Path(comm.bcast(f'r2s_{stamp}')) # Set run_kwargs for the neutron transport step if micro_kwargs is None: @@ -257,18 +262,19 @@ def step1_neutron_transport( """ - output_dir = Path(output_dir) + output_dir = Path(output_dir).resolve() output_dir.mkdir(parents=True, exist_ok=True) if self.method == 'mesh-based': # Compute material volume fractions on the mesh if mat_vol_kwargs is None: mat_vol_kwargs = {} - self.results['mesh_material_volumes'] = mmv = \ - self.domains.material_volumes(self.neutron_model, **mat_vol_kwargs) + self.results['mesh_material_volumes'] = mmv = comm.bcast( + self.domains.material_volumes(self.neutron_model, **mat_vol_kwargs)) # Save results to file - mmv.save(output_dir / 'mesh_material_volumes.npz') + if comm.rank == 0: + mmv.save(output_dir / 'mesh_material_volumes.npz') # Create mesh-material filter based on what combos were found domains = openmc.MeshMaterialFilter.from_volumes(self.domains, mmv) @@ -299,13 +305,16 @@ def step1_neutron_transport( micro_kwargs.setdefault('path_statepoint', output_dir / 'statepoint.h5') micro_kwargs.setdefault('path_input', output_dir / 'model.xml') - # Run neutron transport and get fluxes and micros - self.results['fluxes'], self.results['micros'] = get_microxs_and_flux( - self.neutron_model, domains, **micro_kwargs) + # Run neutron transport and get fluxes and micros. Run via openmc.lib to + # maintain a consistent parallelism strategy with the activation step. + with TemporarySession(): + self.results['fluxes'], self.results['micros'] = get_microxs_and_flux( + self.neutron_model, domains, **micro_kwargs) # Save flux and micros to file - np.save(output_dir / 'fluxes.npy', self.results['fluxes']) - write_microxs_hdf5(self.results['micros'], output_dir / 'micros.h5') + if comm.rank == 0: + np.save(output_dir / 'fluxes.npy', self.results['fluxes']) + write_microxs_hdf5(self.results['micros'], output_dir / 'micros.h5') def step2_activation( self, @@ -457,15 +466,17 @@ def step3_photon_transport( # photon model if it is different from the neutron model to account for # potential material changes if self.method == 'mesh-based' and different_photon_model: - self.results['mesh_material_volumes_photon'] = photon_mmv = \ - self.domains.material_volumes(self.photon_model, **mat_vol_kwargs) + self.results['mesh_material_volumes_photon'] = photon_mmv = comm.bcast( + self.domains.material_volumes(self.photon_model, **mat_vol_kwargs)) # Save photon MMV results to file - photon_mmv.save(output_dir / 'mesh_material_volumes.npz') + if comm.rank == 0: + photon_mmv.save(output_dir / 'mesh_material_volumes.npz') - tally_ids = [tally.id for tally in self.photon_model.tallies] - with open(output_dir / 'tally_ids.json', 'w') as f: - json.dump(tally_ids, f) + if comm.rank == 0: + tally_ids = [tally.id for tally in self.photon_model.tallies] + with open(output_dir / 'tally_ids.json', 'w') as f: + json.dump(tally_ids, f) self.results['photon_tallies'] = {} @@ -514,8 +525,9 @@ def step3_photon_transport( time_index = len(self.results['depletion_results']) + time_index # Run photon transport calculation - run_kwargs['cwd'] = Path(output_dir) / f'time_{time_index}' - statepoint_path = self.photon_model.run(**run_kwargs) + photon_dir = Path(output_dir) / f'time_{time_index}' + with TemporarySession(self.photon_model, cwd=photon_dir): + statepoint_path = self.photon_model.run(**run_kwargs) # Store tally results with openmc.StatePoint(statepoint_path) as sp: diff --git a/openmc/lib/core.py b/openmc/lib/core.py index 9f8db69d577..8a7bcbe1926 100644 --- a/openmc/lib/core.py +++ b/openmc/lib/core.py @@ -13,6 +13,7 @@ from . import _dll from .error import _error_handler +from ..mpi import comm from openmc.checkvalue import PathLike import openmc.lib import openmc @@ -632,6 +633,9 @@ class TemporarySession: model : openmc.Model, optional OpenMC model to use for the session. If None, a minimal working model is created. + cwd : PathLike, optional + Working directory in which to run OpenMC. If None, a temporary directory + is created and deleted automatically. **init_kwargs Keyword arguments to pass to :func:`openmc.lib.init`. @@ -639,10 +643,13 @@ class TemporarySession: ---------- model : openmc.Model The OpenMC model used for the session. + comm : mpi4py.MPI.Intracomm + The MPI intracommunicator used for the session. """ - def __init__(self, model=None, **init_kwargs): - self.init_kwargs = init_kwargs + def __init__(self, model=None, cwd=None, **init_kwargs): + self.init_kwargs = dict(init_kwargs) + self.cwd = cwd if model is None: surf = openmc.Sphere(boundary_type="vacuum") cell = openmc.Cell(region=-surf) @@ -652,6 +659,10 @@ def __init__(self, model=None, **init_kwargs): particles=1, batches=1, output={'summary': False}) self.model = model + # Determine MPI intercommunicator + self.init_kwargs.setdefault('intracomm', comm) + self.comm = self.init_kwargs['intracomm'] + def __enter__(self): """Initialize the OpenMC library in a temporary directory.""" # If already initialized, the context manager is a no-op @@ -662,14 +673,28 @@ def __enter__(self): # Store original working directory self.orig_dir = Path.cwd() - # Set up temporary directory - self.tmp_dir = TemporaryDirectory() - working_dir = Path(self.tmp_dir.name) + if self.cwd is None: + # Set up temporary directory on rank 0 + if self.comm.rank == 0: + self.tmp_dir = TemporaryDirectory() + path_str = self.tmp_dir.name + else: + path_str = None + + # Broadcast the path so that all ranks use the same directory + path_str = self.comm.bcast(path_str) + working_dir = Path(path_str) + else: + working_dir = Path(self.cwd) + + # Create and change to specified directory working_dir.mkdir(parents=True, exist_ok=True) os.chdir(working_dir) - # Export model and initialize OpenMC - self.model.export_to_model_xml() + # Export model on first rank and initialize OpenMC + if self.comm.rank == 0: + self.model.export_to_model_xml() + self.comm.barrier() openmc.lib.init(**self.init_kwargs) return self @@ -683,7 +708,11 @@ def __exit__(self, exc_type, exc_value, traceback): finalize() finally: os.chdir(self.orig_dir) - self.tmp_dir.cleanup() + + # Make sure all ranks have finalized before deleting temporary dir + self.comm.barrier() + if hasattr(self, 'tmp_dir'): + self.tmp_dir.cleanup() class _DLLGlobal: