Skip to content

Commit

Permalink
run microxs with mpi (#3028)
Browse files Browse the repository at this point in the history
  • Loading branch information
church89 authored Aug 16, 2024
1 parent b22656e commit 54c28b7
Showing 1 changed file with 18 additions and 6 deletions.
24 changes: 18 additions & 6 deletions openmc/deplete/microxs.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
from .chain import Chain, REACTIONS
from .coupled_operator import _find_cross_sections, _get_nuclides_with_data
import openmc.lib
from openmc.mpi import comm

_valid_rxns = list(REACTIONS)
_valid_rxns.append('fission')
Expand Down Expand Up @@ -124,6 +125,15 @@ def get_microxs_and_flux(
flux_tally.scores = ['flux']
model.tallies = openmc.Tallies([rr_tally, flux_tally])

if openmc.lib.is_initialized:
openmc.lib.finalize()

if comm.rank == 0:
model.export_to_model_xml()
comm.barrier()
# Reinitialize with tallies
openmc.lib.init(intracomm=comm)

# create temporary run
with TemporaryDirectory() as temp_dir:
if run_kwargs is None:
Expand All @@ -133,12 +143,15 @@ def get_microxs_and_flux(
run_kwargs.setdefault('cwd', temp_dir)
statepoint_path = model.run(**run_kwargs)

with StatePoint(statepoint_path) as sp:
rr_tally = sp.tallies[rr_tally.id]
rr_tally._read_results()
flux_tally = sp.tallies[flux_tally.id]
flux_tally._read_results()
if comm.rank == 0:
with StatePoint(statepoint_path) as sp:
rr_tally = sp.tallies[rr_tally.id]
rr_tally._read_results()
flux_tally = sp.tallies[flux_tally.id]
flux_tally._read_results()

rr_tally = comm.bcast(rr_tally)
flux_tally = comm.bcast(flux_tally)
# Get reaction rates and flux values
reaction_rates = rr_tally.get_reshaped_data() # (domains, groups, nuclides, reactions)
flux = flux_tally.get_reshaped_data() # (domains, groups, 1, 1)
Expand Down Expand Up @@ -371,4 +384,3 @@ def to_csv(self, *args, **kwargs):
)
df = pd.DataFrame({'xs': self.data.flatten()}, index=multi_index)
df.to_csv(*args, **kwargs)

0 comments on commit 54c28b7

Please sign in to comment.