From c5bc2302eb91745c810ddbaa97073ca8be0c3486 Mon Sep 17 00:00:00 2001 From: Lorenzo Chierici Date: Fri, 16 Jun 2023 06:00:20 +0200 Subject: [PATCH] Transfer rates mpi (#2551) Co-authored-by: Paul Romano --- openmc/deplete/coupled_operator.py | 3 +- openmc/deplete/independent_operator.py | 3 +- openmc/deplete/openmc_operator.py | 24 +----- openmc/deplete/pool.py | 100 +++++++++++++++++-------- openmc/deplete/transfer_rates.py | 3 + 5 files changed, 77 insertions(+), 56 deletions(-) diff --git a/openmc/deplete/coupled_operator.py b/openmc/deplete/coupled_operator.py index 669e98ddac..d8d4838a4a 100644 --- a/openmc/deplete/coupled_operator.py +++ b/openmc/deplete/coupled_operator.py @@ -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, diff --git a/openmc/deplete/independent_operator.py b/openmc/deplete/independent_operator.py index 6caef26352..8d7882ac9b 100644 --- a/openmc/deplete/independent_operator.py +++ b/openmc/deplete/independent_operator.py @@ -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 diff --git a/openmc/deplete/openmc_operator.py b/openmc/deplete/openmc_operator.py index adf1c2eb7b..55a7d72469 100644 --- a/openmc/deplete/openmc_operator.py +++ b/openmc/deplete/openmc_operator.py @@ -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. diff --git a/openmc/deplete/pool.py b/openmc/deplete/pool.py index fc006a8842..1fa5e26ac8 100644 --- a/openmc/deplete/pool.py +++ b/openmc/deplete/pool.py @@ -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 diff --git a/openmc/deplete/transfer_rates.py b/openmc/deplete/transfer_rates.py index c6f0cd94df..fe21d7d6f4 100644 --- a/openmc/deplete/transfer_rates.py +++ b/openmc/deplete/transfer_rates.py @@ -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}