crp/OpenMC
1
0
Fork 0
mirror of https://github.com/openmc-dev/openmc.git synced 2026-07-28 14:15:42 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions 5

Support MPI parallelism in R2SManager (#3632)

Browse source
This commit is contained in:
Paul Romano 2025-11-19 11:19:27 -06:00 • committed by GitHub
parent 5c63e0df21
commit 028f440448
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 82 additions and 42 deletions
Download patch file Download diff file Expand all files Collapse all files

31
openmc/deplete/microxs.py
View file

@ -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

48
openmc/deplete/r2s.py
View file

@ -11,6 +11,9 @@ from . import IndependentOperator, PredictorIntegrator
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 @@ class R2SManager:
"""
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 @@ class R2SManager:
"""
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 @@ class R2SManager:
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 @@ class R2SManager:
# 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 @@ class R2SManager:
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:

45
openmc/lib/core.py
View file

@ -13,6 +13,7 @@ from numpy.ctypeslib import as_array
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 @@ class TemporarySession:
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,24 @@ class TemporarySession:
# Store original working directory
self.orig_dir = Path.cwd()
# Set up temporary directory
self.tmp_dir = TemporaryDirectory()
working_dir = Path(self.tmp_dir.name)
working_dir.mkdir(parents=True, exist_ok=True)
os.chdir(working_dir)
if self.cwd is None:
# Set up temporary directory on rank 0
if self.comm.rank == 0:
self._tmp_dir = TemporaryDirectory()
self.cwd = self._tmp_dir.name
# Export model and initialize OpenMC
self.model.export_to_model_xml()
# Broadcast the path so that all ranks use the same directory
self.cwd = self.comm.bcast(self.cwd)
# Create and change to specified directory
self.cwd = Path(self.cwd)
self.cwd.mkdir(parents=True, exist_ok=True)
os.chdir(self.cwd)
# 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 +704,11 @@ class TemporarySession:
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: