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

Transfer rates mpi (#2551)

Browse source 
Co-authored-by: Paul Romano <paul.k.romano@gmail.com>
This commit is contained in:
Lorenzo Chierici 2023-06-16 06:00:20 +02:00 • committed by GitHub
parent 787b3e0308
commit c5bc2302eb
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 77 additions and 56 deletions
Download patch file Download diff file Expand all files Collapse all files

3
openmc/deplete/coupled_operator.py
View file

@ -21,7 +21,8 @@ from openmc.exceptions import DataError
import openmc.lib
from openmc.mpi import comm
from .abc import OperatorResult
from .openmc_operator import OpenMCOperator, _distribute
from .openmc_operator import OpenMCOperator
from .pool import _distribute
from .results import Results
from .helpers import (
DirectReactionRateHelper, ChainFissionHelper, ConstantFissionYieldHelper,

3
openmc/deplete/independent_operator.py
View file

@ -15,7 +15,8 @@ import openmc
from openmc.checkvalue import check_type
from openmc.mpi import comm
from .abc import ReactionRateHelper, OperatorResult
from .openmc_operator import OpenMCOperator, _distribute
from .openmc_operator import OpenMCOperator
from .pool import _distribute
from .microxs import MicroXS
from .results import Results
from .helpers import ChainFissionHelper, ConstantFissionYieldHelper, SourceRateHelper

24
openmc/deplete/openmc_operator.py
View file

@ -16,33 +16,11 @@ from openmc.mpi import comm
from .abc import TransportOperator, OperatorResult
from .atom_number import AtomNumber
from .reaction_rates import ReactionRates
from .pool import _distribute
__all__ = ["OpenMCOperator", "OperatorResult"]
def _distribute(items):
"""Distribute items across MPI communicator
Parameters
----------
items : list
List of items of distribute
Returns
-------
list
Items assigned to process that called
"""
min_size, extra = divmod(len(items), comm.size)
j = 0
for i in range(comm.size):
chunk_size = min_size + int(i < extra)
if comm.rank == i:
return items[j:j + chunk_size]
j += chunk_size
class OpenMCOperator(TransportOperator):
"""Abstract class holding OpenMC-specific functions for running
depletion calculations.

100
openmc/deplete/pool.py
View file

@ -6,7 +6,7 @@ from itertools import repeat, starmap
from multiprocessing import Pool
from scipy.sparse import bmat
import numpy as np
from openmc.mpi import comm
# Configurable switch that enables / disables the use of
# multiprocessing routines during depletion
@ -16,6 +16,27 @@ USE_MULTIPROCESSING = True
# calculations
NUM_PROCESSES = None
def _distribute(items):
"""Distribute items across MPI communicator
Parameters
----------
items : list
List of items of distribute
Returns
-------
list
Items assigned to process that called
"""
min_size, extra = divmod(len(items), comm.size)
j = 0
for i in range(comm.size):
chunk_size = min_size + int(i < extra)
if comm.rank == i:
return items[j:j + chunk_size]
j += chunk_size
def deplete(func, chain, x, rates, dt, matrix_func=None, transfer_rates=None,
*matrix_args):
@ -72,45 +93,62 @@ def deplete(func, chain, x, rates, dt, matrix_func=None, transfer_rates=None,
if transfer_rates is not None:
# Calculate transfer rate terms as diagonal matrices
transfers = map(chain.form_rr_term, repeat(transfer_rates),
transfer_rates.burnable_mats)
transfer_rates.local_mats)
# Subtract transfer rate terms from Bateman matrices
matrices = [matrix - transfer for (matrix, transfer) in zip(matrices,
transfers)]
if len(transfer_rates.index_transfer) > 0:
# Calculate transfer rate terms as diagonal matrices
transfer_pair = {
mat_pair: chain.form_rr_term(transfer_rates, mat_pair)
for mat_pair in transfer_rates.index_transfer
}
# Gather all on comm.rank 0
matrices = comm.gather(matrices)
x = comm.gather(x)
# Combine all matrices together in a single matrix of matrices
# to be solved in one go
n_rows = n_cols = len(transfer_rates.burnable_mats)
rows = []
for row in range(n_rows):
cols = []
for col in range(n_cols):
mat_pair = (transfer_rates.burnable_mats[row],
transfer_rates.burnable_mats[col])
if row == col:
# Fill the diagonals with the Bateman matrices
cols.append(matrices[row])
elif mat_pair in transfer_rates.index_transfer:
# Fill the off-diagonals with the transfer pair matrices
cols.append(transfer_pair[mat_pair])
else:
cols.append(None)
if comm.rank == 0:
# Expand lists
matrices = [elm for matrix in matrices for elm in matrix]
x = [x_elm for x_mat in x for x_elm in x_mat]
rows.append(cols)
matrix = bmat(rows)
# Calculate transfer rate terms as diagonal matrices
transfer_pair = {
mat_pair: chain.form_rr_term(transfer_rates, mat_pair)
for mat_pair in transfer_rates.index_transfer
}
# Concatenate vectors of nuclides in one
x_multi = np.concatenate([xx for xx in x])
x_result = func(matrix, x_multi, dt)
# Combine all matrices together in a single matrix of matrices
# to be solved in one go
n_rows = n_cols = len(transfer_rates.burnable_mats)
rows = []
for row in range(n_rows):
cols = []
for col in range(n_cols):
mat_pair = (transfer_rates.burnable_mats[row],
transfer_rates.burnable_mats[col])
if row == col:
# Fill the diagonals with the Bateman matrices
cols.append(matrices[row])
elif mat_pair in transfer_rates.index_transfer:
# Fill the off-diagonals with the transfer pair matrices
cols.append(transfer_pair[mat_pair])
else:
cols.append(None)
# Split back the nuclide vector result into the original form
x_result = np.split(x_result, np.cumsum([len(i) for i in x])[:-1])
rows.append(cols)
matrix = bmat(rows)
# Concatenate vectors of nuclides in one
x_multi = np.concatenate([xx for xx in x])
x_result = func(matrix, x_multi, dt)
# Split back the nuclide vector result into the original form
x_result = np.split(x_result, np.cumsum([len(i) for i in x])[:-1])
else:
x_result = None
# Braodcast result to other ranks
x_result = comm.bcast(x_result)
# Distribute results across MPI
x_result = _distribute(x_result)
return x_result

3
openmc/deplete/transfer_rates.py
View file

@ -30,6 +30,8 @@ class TransferRates:
----------
burnable_mats : list of str
All burnable material IDs.
local_mats : list of str
All burnable material IDs being managed by a single process
transfer_rates : dict of str to dict
Container of transfer rates, elements and destination material
index_transfer : Set of pair of str
@ -40,6 +42,7 @@ class TransferRates:
self.materials = model.materials
self.burnable_mats = operator.burnable_mats
self.local_mats = operator.local_mats
#initialize transfer rates container dict
self.transfer_rates = {mat: {} for mat in self.burnable_mats}