"""Dedicated module containing depletion function Provided to avoid some circular imports """ from itertools import repeat, starmap from multiprocessing import Pool import numpy as np from scipy.sparse import hstack, vstack from openmc.mpi import comm from .._sparse_compat import block_array, csc_array # Configurable switch that enables / disables the use of # multiprocessing routines during depletion USE_MULTIPROCESSING = True # Allow user to override the number of worker processes to use for depletion # 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 _add_external_source( matrices, n, chain, external_source_rates, current_timestep ): """Augment depletion matrices and nuclide vectors with external sources.""" sources = map(chain.form_ext_source_term, repeat(external_source_rates), repeat(current_timestep), external_source_rates.local_mats) matrices = [ hstack([matrix, source]) for matrix, source in zip(matrices, sources) ] n_solve = [arr.copy() for arr in n] # Homogenize the augmented matrices and nuclide vectors for i, matrix in enumerate(matrices): if matrix.shape[0] + 1 == matrix.shape[1]: matrices[i] = vstack( [matrix, csc_array((1, matrix.shape[1]))]) n_solve[i] = np.append(n_solve[i], 1.0) return matrices, n_solve def deplete(func, chain, n, rates, dt, current_timestep=None, matrix_func=None, transfer_rates=None, external_source_rates=None, substeps=1, *matrix_args): """Deplete materials using given reaction rates for a specified time Parameters ---------- func : callable Function to use to get new compositions. Expected to have the signature ``func(A, n0, t, substeps=1) -> n1``. chain : openmc.deplete.Chain Depletion chain n : list of numpy.ndarray List of atom number arrays for each material. Each array in the list contains the number of [atom] of each nuclide. rates : openmc.deplete.ReactionRates Reaction rates (from transport operator) dt : float Time in [s] to deplete for current_timestep : int Current timestep index matrix_func : callable, optional Function to form the depletion matrix after calling ``matrix_func(chain, rates, fission_yields)``, where ``fission_yields = {parent: {product: yield_frac}}`` Expected to return the depletion matrix required by ``func`` transfer_rates : openmc.deplete.TransferRates, Optional Transfer rates for continuous removal/feed. .. versionadded:: 0.14.0 external_source_rates : openmc.deplete.ExternalSourceRates, Optional External source rates for continuous removal/feed. .. versionadded:: 0.15.3 substeps : int, optional Number of substeps to pass to solvers that support substepping. matrix_args: Any, optional Additional arguments passed to matrix_func Returns ------- n_result : list of numpy.ndarray Updated list of atom number arrays for each material. Each array in the list contains the number of [atom] of each nuclide. """ fission_yields = chain.fission_yields if len(fission_yields) == 1: fission_yields = repeat(fission_yields[0]) elif len(fission_yields) != len(n): raise ValueError( "Number of material fission yield distributions {} is not " "equal to the number of compositions {}".format( len(fission_yields), len(n))) if matrix_func is None: matrices = map(chain.form_matrix, rates, fission_yields) else: matrices = map(matrix_func, repeat(chain), rates, fission_yields, *matrix_args) # Determine if transfer rates or external source rates are active transfer_active = transfer_rates is not None and \ current_timestep in transfer_rates.external_timesteps external_active = external_source_rates is not None and \ current_timestep in external_source_rates.external_timesteps n_solve = n if transfer_active: # Calculate transfer rate terms as diagonal matrices transfers = map(chain.form_rr_term, repeat(transfer_rates), repeat(current_timestep), transfer_rates.local_mats) # Subtract transfer rate terms from Bateman matrices matrices = [matrix - transfer for (matrix, transfer) in zip(matrices, transfers)] if transfer_rates.redox: for mat_idx, mat_id in enumerate(transfer_rates.local_mats): if mat_id in transfer_rates.redox: matrices[mat_idx] = chain.add_redox_term(matrices[mat_idx], transfer_rates.redox[mat_id][0], transfer_rates.redox[mat_id][1]) # Add external sources if present if external_active: matrices, n_solve = _add_external_source( matrices, n, chain, external_source_rates, current_timestep) # Set transfer rate terms with destination material if present if current_timestep in transfer_rates.index_transfer: # Gather all on comm.rank 0 matrices = comm.gather(matrices) n = comm.gather(n_solve) if comm.rank == 0: # Expand lists matrices = [elm for matrix in matrices for elm in matrix] n = [n_elm for n_mat in n for n_elm in n_mat] # Calculate transfer rate terms as diagonal matrices transfer_pair = {} for mat_pair in dict.fromkeys(transfer_rates.index_transfer[current_timestep]): transfer_matrix = chain.form_rr_term(transfer_rates, current_timestep, mat_pair) # check if destination material has a redox control if mat_pair[0] in transfer_rates.redox: transfer_matrix = chain.add_redox_term(transfer_matrix, transfer_rates.redox[mat_pair[0]][0], transfer_rates.redox[mat_pair[0]][1]) # Add external source rates if present if external_active: if len(external_source_rates.get_components(mat_pair[0], current_timestep)) > 0: transfer_matrix = vstack([transfer_matrix, csc_array((1, transfer_matrix.shape[1]))]) if len(external_source_rates.get_components(mat_pair[1], current_timestep)) > 0: transfer_matrix = hstack([transfer_matrix, csc_array((transfer_matrix.shape[0], 1))]) transfer_pair[mat_pair] = transfer_matrix # Combine all matrices together in a single block matrix of matrices # to be solved on one rank 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[current_timestep]: # Fill the off-diagonals with the transfer pair matrices cols.append(transfer_pair[mat_pair]) else: cols.append(None) rows.append(cols) matrix = block_array(rows) # Concatenate vectors of nuclides in one n_multi = np.concatenate(n) n_result = func(matrix, n_multi, dt, substeps) # Split back the nuclide vector result into the original form n_result = np.split(n_result, np.cumsum([len(i) for i in n])[:-1]) else: n_result = None # Broadcast result to other MPI ranks and then distribute n_result = comm.bcast(n_result) n_result = _distribute(n_result) # Remove extra values based on the materials local to each rank if external_active: external_source_rates.reformat_nuclide_vectors(n_result) return n_result # If only external source rates are present elif external_active: matrices, n_solve = _add_external_source( matrices, n, chain, external_source_rates, current_timestep) inputs = zip(matrices, n_solve, repeat(dt), repeat(substeps)) if USE_MULTIPROCESSING: with Pool(NUM_PROCESSES) as pool: n_result = list(pool.starmap(func, inputs)) else: n_result = list(starmap(func, inputs)) # Remove extra value at the end of the nuclide vectors if external source rates are present if external_active: external_source_rates.reformat_nuclide_vectors(n_result) return n_result