diff --git a/docs/source/pythonapi/deplete.rst b/docs/source/pythonapi/deplete.rst index b9970ffb11..b01de477ce 100644 --- a/docs/source/pythonapi/deplete.rst +++ b/docs/source/pythonapi/deplete.rst @@ -143,8 +143,7 @@ with :func:`cram.CRAM48` being the default. cram.CRAM16 cram.CRAM48 - cram.deplete - cram.timed_deplete + pool.deplete The following classes are used to help the :class:`openmc.deplete.Operator` compute quantities like effective fission yields, reaction rates, and diff --git a/openmc/deplete/abc.py b/openmc/deplete/abc.py index ddc48bc4de..8846906b44 100644 --- a/openmc/deplete/abc.py +++ b/openmc/deplete/abc.py @@ -4,15 +4,16 @@ This module contains the Operator class, which is then passed to an integrator to run a full depletion simulation. """ -from collections import namedtuple -from collections import defaultdict -from collections.abc import Iterable +from collections import namedtuple, defaultdict +from collections.abc import Iterable, Callable import os from pathlib import Path from abc import ABC, abstractmethod from copy import deepcopy from warnings import warn from numbers import Real, Integral +from inspect import signature +import time from numpy import nonzero, empty, asarray from uncertainties import ufloat @@ -23,6 +24,7 @@ from openmc.checkvalue import check_type, check_greater_than from .results import Results from .chain import Chain from .results_list import ResultsList +from .pool import deplete __all__ = [ @@ -595,7 +597,7 @@ class TalliedFissionYieldHelper(FissionYieldHelper): class Integrator(ABC): - """Abstract class for solving the time-integration for depletion + r"""Abstract class for solving the time-integration for depletion Parameters ---------- @@ -623,6 +625,17 @@ class Integrator(ABC): seconds, 'min' means minutes, 'h' means hours, and 'MWd/kg' indicates that the values are given in burnup (MW-d of energy deposited per kilogram of initial heavy metal). + solver : str or callable, optional + If a string, must be the name of the solver responsible for + solving the Bateman equations. Current options are: + + * ``cram16`` - 16th order IPF CRAM + * ``cram48`` - 48th order IPF CRAM [default] + + If a function or other callable, must adhere to the requirements in + :attr:`solver`. + + .. versionadded:: 0.12 Attributes ---------- @@ -634,10 +647,27 @@ class Integrator(ABC): Size of each depletion interval in [s] power : iterable of float Power of the reactor in [W] for each interval in :attr:`timesteps` + solver : callable + Function that will solve the Bateman equations + :math:`\frac{\partial}{\partial t}\vec{n} = A_i\vec{n}_i` with a step + size :math:`t_i`. Can be configured using the ``solver`` argument. + User-supplied functions are expected to have the following signature: + ``solver(A, n0, t) -> n1`` where + + * ``A`` is a :class:`scipy.sparse.csr_matrix` making up the + depletion matrix + * ``n0`` is a 1-D :class:`numpy.ndarray` of initial compositions + for a given material in atoms/cm3 + * ``t`` is a float of the time step size in seconds, and + * ``n1`` is a :class:`numpy.ndarray` of compositions at the + next time step. Expected to be of the same shape as ``n0`` + + .. versionadded:: 0.12 + """ def __init__(self, operator, timesteps, power=None, power_density=None, - timestep_units='s'): + timestep_units='s', solver="cram48"): # Check number of stages previously used if operator.prev_res is not None: res = operator.prev_res[-1] @@ -677,24 +707,24 @@ class Integrator(ABC): # Determine number of seconds for each timestep seconds = [] - for time, unit, watts in zip(times, units, power): + for timestep, unit, watts in zip(times, units, power): # Make sure values passed make sense - check_type('timestep', time, Real) - check_greater_than('timestep', time, 0.0, False) + check_type('timestep', timestep, Real) + check_greater_than('timestep', timestep, 0.0, False) check_type('timestep units', unit, str) check_type('power', watts, Real) check_greater_than('power', watts, 0.0, True) if unit in ('s', 'sec'): - seconds.append(time) + seconds.append(timestep) elif unit in ('min', 'minute'): - seconds.append(time*_SECONDS_PER_MINUTE) + seconds.append(timestep*_SECONDS_PER_MINUTE) elif unit in ('h', 'hr', 'hour'): - seconds.append(time*_SECONDS_PER_HOUR) + seconds.append(timestep*_SECONDS_PER_HOUR) elif unit in ('d', 'day'): - seconds.append(time*_SECONDS_PER_DAY) + seconds.append(timestep*_SECONDS_PER_DAY) elif unit.lower() == 'mwd/kg': - watt_days_per_kg = 1e6*time + watt_days_per_kg = 1e6*timestep kilograms = 1e-3*operator.heavy_metal days = watt_days_per_kg * kilograms / watts seconds.append(days*_SECONDS_PER_DAY) @@ -704,6 +734,59 @@ class Integrator(ABC): self.timesteps = asarray(seconds) self.power = asarray(power) + if isinstance(solver, str): + # Delay importing of cram module, which requires this file + if solver == "cram48": + from .cram import CRAM48 + self._solver = CRAM48 + elif solver == "cram16": + from .cram import CRAM16 + self._solver = CRAM16 + else: + raise ValueError( + "Solver {} not understood. Expected 'cram48' or " + "'cram16'".format(solver)) + else: + self.solver = solver + + @property + def solver(self): + return self._solver + + @solver.setter + def solver(self, func): + if not isinstance(func, Callable): + raise TypeError( + "Solver must be callable, not {}".format(type(func))) + try: + sig = signature(func) + except ValueError: + # Guard against callables that aren't introspectable, e.g. + # fortran functions wrapped by F2PY + warn("Could not determine arguments to {}. Proceeding " + "anyways".format(func)) + self._solver = func + return + + # Inspect arguments + if len(sig.parameters) != 3: + raise ValueError("Function {} does not support three arguments: " + "{!s}".format(func, sig)) + + for ix, param in enumerate(sig.parameters.values()): + if param.kind in {param.KEYWORD_ONLY, param.VAR_KEYWORD}: + raise ValueError( + "Keyword arguments like {} at position {} are not " + "allowed".format(ix, param)) + + self._solver = func + + def _timed_deplete(self, concs, rates, dt, matrix_func=None): + start = time.time() + results = deplete( + self._solver, self.chain, concs, rates, dt, matrix_func) + return time.time() - start, results + @abstractmethod def __call__(self, conc, rates, dt, power, i): """Perform the integration across one time step @@ -807,7 +890,7 @@ class Integrator(ABC): class SIIntegrator(Integrator): - """Abstract class for the Stochastic Implicit Euler integrators + r"""Abstract class for the Stochastic Implicit Euler integrators Does not provide a ``__call__`` method, but scales and resets the number of particles used in initial transport calculation @@ -841,6 +924,17 @@ class SIIntegrator(Integrator): n_steps : int, optional Number of stochastic iterations per depletion interval. Must be greater than zero. Default : 10 + solver : str or callable, optional + If a string, must be the name of the solver responsible for + solving the Bateman equations. Current options are: + + * ``cram16`` - 16th order IPF CRAM + * ``cram48`` - 48th order IPF CRAM [default] + + If a function or other callable, must adhere to the requirements in + :attr:`solver`. + + .. versionadded:: 0.12 Attributes ---------- @@ -854,12 +948,31 @@ class SIIntegrator(Integrator): Power of the reactor in [W] for each interval in :attr:`timesteps` n_steps : int Number of stochastic iterations per depletion interval + solver : callable + Function that will solve the Bateman equations + :math:`\frac{\partial}{\partial t}\vec{n} = A_i\vec{n}_i` with a step + size :math:`t_i`. Can be configured using the ``solver`` argument. + User-supplied functions are expected to have the following signature: + ``solver(A, n0, t) -> n1`` where + + * ``A`` is a :class:`scipy.sparse.csr_matrix` making up the + depletion matrix + * ``n0`` is a 1-D :class:`numpy.ndarray` of initial compositions + for a given material in atoms/cm3 + * ``t`` is a float of the time step size in seconds, and + * ``n1`` is a :class:`numpy.ndarray` of compositions at the + next time step. Expected to be of the same shape as ``n0`` + + .. versionadded:: 0.12 + """ def __init__(self, operator, timesteps, power=None, power_density=None, - timestep_units='s', n_steps=10): + timestep_units='s', n_steps=10, solver="cram48"): check_type("n_steps", n_steps, Integral) check_greater_than("n_steps", n_steps, 0) - super().__init__(operator, timesteps, power, power_density, timestep_units) + super().__init__( + operator, timesteps, power, power_density, timestep_units, + solver=solver) self.n_steps = n_steps def _get_bos_data_from_operator(self, step_index, step_power, bos_conc): diff --git a/openmc/deplete/cram.py b/openmc/deplete/cram.py index 7dad820a9c..6b6a909ec1 100644 --- a/openmc/deplete/cram.py +++ b/openmc/deplete/cram.py @@ -4,9 +4,6 @@ Implements two different forms of CRAM for use in openmc.deplete. """ import numbers -from itertools import repeat -from multiprocessing import Pool -import time import numpy as np import scipy.sparse as sp @@ -15,76 +12,7 @@ import scipy.sparse.linalg as sla from openmc.checkvalue import check_type, check_length from .abc import DepSystemSolver -__all__ = [ - "deplete", "timed_deplete", "CRAM16", "CRAM48", - "Cram16Solver", "Cram48Solver", "IPFCramSolver"] - - -def deplete(chain, x, rates, dt, matrix_func=None): - """Deplete materials using given reaction rates for a specified time - - Parameters - ---------- - chain : openmc.deplete.Chain - Depletion chain - x : list of numpy.ndarray - Atom number vectors for each material - rates : openmc.deplete.ReactionRates - Reaction rates (from transport operator) - dt : float - Time in [s] to deplete for - maxtrix_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:`CRAM48`. - - Returns - ------- - x_result : list of numpy.ndarray - Updated atom number vectors for each material - """ - - fission_yields = chain.fission_yields - if len(fission_yields) == 1: - fission_yields = repeat(fission_yields[0]) - elif len(fission_yields) != len(x): - raise ValueError( - "Number of material fission yield distributions {} is not equal " - "to the number of compositions {}".format(len(fission_yields), - len(x))) - - if matrix_func is None: - matrices = map(chain.form_matrix, rates, fission_yields) - else: - matrices = map(matrix_func, repeat(chain), rates, fission_yields) - - # Use multiprocessing pool to distribute work - with Pool() as pool: - inputs = zip(matrices, x, repeat(dt)) - x_result = list(pool.starmap(CRAM48, inputs)) - - return x_result - - -def timed_deplete(*args, **kwargs): - """Wrapper over :func:`deplete` that also returns process time - - All arguments and keyword arguments are passed onto - :func:`deplete` directly. - - Returns - ------- - proc_time: float - Process time required to return from deplete - results: list of numpy arrays - Output from :func:`deplete` call - """ - - start = time.time() - results = deplete(*args, **kwargs) - return time.time() - start, results +__all__ = ["CRAM16", "CRAM48", "Cram16Solver", "Cram48Solver", "IPFCramSolver"] class IPFCramSolver(DepSystemSolver): diff --git a/openmc/deplete/integrators.py b/openmc/deplete/integrators.py index ddb2130379..1fe6369003 100644 --- a/openmc/deplete/integrators.py +++ b/openmc/deplete/integrators.py @@ -2,7 +2,6 @@ import copy from itertools import repeat from .abc import Integrator, SIIntegrator, OperatorResult -from .cram import timed_deplete from ._matrix_funcs import ( cf4_f1, cf4_f2, cf4_f3, cf4_f4, celi_f1, celi_f2, leqi_f1, leqi_f2, leqi_f3, leqi_f4, rk4_f1, rk4_f4 @@ -54,6 +53,17 @@ class PredictorIntegrator(Integrator): that the values are given in burnup (MW-d of energy deposited per kilogram of initial heavy metal). + .. versionadded:: 0.12 + solver : str or callable, optional + If a string, must be the name of the solver responsible for + solving the Bateman equations. Current options are: + + * ``cram16`` - 16th order IPF CRAM + * ``cram48`` - 48th order IPF CRAM [default] + + If a function or other callable, must adhere to the requirements in + :attr:`solver`. + .. versionadded:: 0.12 Attributes @@ -66,6 +76,23 @@ class PredictorIntegrator(Integrator): Size of each depletion interval in [s] power : iterable of float Power of the reactor in [W] for each interval in :attr:`timesteps` + solver : callable + Function that will solve the Bateman equations + :math:`\frac{\partial}{\partial t}\vec{n} = A_i\vec{n}_i` with a step + size :math:`t_i`. Can be configured using the ``solver`` argument. + User-supplied functions are expected to have the following signature: + ``solver(A, n0, t) -> n1`` where + + * ``A`` is a :class:`scipy.sparse.csr_matrix` making up the + depletion matrix + * ``n0`` is a 1-D :class:`numpy.ndarray` of initial compositions + for a given material in atoms/cm3 + * ``t`` is a float of the time step size in seconds, and + * ``n1`` is a :class:`numpy.ndarray` of compositions at the + next time step. Expected to be of the same shape as ``n0`` + + .. versionadded:: 0.12 + """ _num_stages = 1 @@ -96,7 +123,7 @@ class PredictorIntegrator(Integrator): operator with predictor """ - proc_time, conc_end = timed_deplete(self.chain, conc, rates, dt) + proc_time, conc_end = self._timed_deplete(conc, rates, dt) return proc_time, [conc_end], [] @@ -145,6 +172,17 @@ class CECMIntegrator(Integrator): that the values are given in burnup (MW-d of energy deposited per kilogram of initial heavy metal). + .. versionadded:: 0.12 + solver : str or callable, optional + If a string, must be the name of the solver responsible for + solving the Bateman equations. Current options are: + + * ``cram16`` - 16th order IPF CRAM + * ``cram48`` - 48th order IPF CRAM [default] + + If a function or other callable, must adhere to the requirements in + :attr:`solver`. + .. versionadded:: 0.12 Attributes @@ -157,6 +195,23 @@ class CECMIntegrator(Integrator): Size of each depletion interval in [s] power : iterable of float Power of the reactor in [W] for each interval in :attr:`timesteps` + solver : callable + Function that will solve the Bateman equations + :math:`\frac{\partial}{\partial t}\vec{n} = A_i\vec{n}_i` with a step + size :math:`t_i`. Can be configured using the ``solver`` argument. + User-supplied functions are expected to have the following signature: + ``solver(A, n0, t) -> n1`` where + + * ``A`` is a :class:`scipy.sparse.csr_matrix` making up the + depletion matrix + * ``n0`` is a 1-D :class:`numpy.ndarray` of initial compositions + for a given material in atoms/cm3 + * ``t`` is a float of the time step size in seconds, and + * ``n1`` is a :class:`numpy.ndarray` of compositions at the + next time step. Expected to be of the same shape as ``n0`` + + .. versionadded:: 0.12 + """ _num_stages = 2 @@ -187,12 +242,12 @@ class CECMIntegrator(Integrator): Eigenvalue and reaction rates from transport simulations """ # deplete across first half of inteval - time0, x_middle = timed_deplete(self.chain, conc, rates, dt / 2) + time0, x_middle = self._timed_deplete(conc, rates, dt / 2) res_middle = self.operator(x_middle, power) # deplete across entire interval with BOS concentrations, # MOS reaction rates - time1, x_end = timed_deplete(self.chain, conc, res_middle.rates, dt) + time1, x_end = self._timed_deplete(conc, res_middle.rates, dt) return time0 + time1, [x_middle, x_end], [res_middle] @@ -244,6 +299,17 @@ class CF4Integrator(Integrator): that the values are given in burnup (MW-d of energy deposited per kilogram of initial heavy metal). + .. versionadded:: 0.12 + solver : str or callable, optional + If a string, must be the name of the solver responsible for + solving the Bateman equations. Current options are: + + * ``cram16`` - 16th order IPF CRAM + * ``cram48`` - 48th order IPF CRAM [default] + + If a function or other callable, must adhere to the requirements in + :attr:`solver`. + .. versionadded:: 0.12 Attributes @@ -256,6 +322,23 @@ class CF4Integrator(Integrator): Size of each depletion interval in [s] power : iterable of float Power of the reactor in [W] for each interval in :attr:`timesteps` + solver : callable + Function that will solve the Bateman equations + :math:`\frac{\partial}{\partial t}\vec{n} = A_i\vec{n}_i` with a step + size :math:`t_i`. Can be configured using the ``solver`` argument. + User-supplied functions are expected to have the following signature: + ``solver(A, n0, t) -> n1`` where + + * ``A`` is a :class:`scipy.sparse.csr_matrix` making up the + depletion matrix + * ``n0`` is a 1-D :class:`numpy.ndarray` of initial compositions + for a given material in atoms/cm3 + * ``t`` is a float of the time step size in seconds, and + * ``n1`` is a :class:`numpy.ndarray` of compositions at the + next time step. Expected to be of the same shape as ``n0`` + + .. versionadded:: 0.12 + """ _num_stages = 4 @@ -287,27 +370,27 @@ class CF4Integrator(Integrator): simulations """ # Step 1: deplete with matrix 1/2*A(y0) - time1, conc_eos1 = timed_deplete( - self.chain, bos_conc, bos_rates, dt, matrix_func=cf4_f1) + time1, conc_eos1 = self._timed_deplete( + bos_conc, bos_rates, dt, matrix_func=cf4_f1) res1 = self.operator(conc_eos1, power) # Step 2: deplete with matrix 1/2*A(y1) - time2, conc_eos2 = timed_deplete( - self.chain, bos_conc, res1.rates, dt, matrix_func=cf4_f1) + time2, conc_eos2 = self._timed_deplete( + bos_conc, res1.rates, dt, matrix_func=cf4_f1) res2 = self.operator(conc_eos2, power) # Step 3: deplete with matrix -1/2*A(y0)+A(y2) list_rates = list(zip(bos_rates, res2.rates)) - time3, conc_eos3 = timed_deplete( - self.chain, conc_eos1, list_rates, dt, matrix_func=cf4_f2) + time3, conc_eos3 = self._timed_deplete( + conc_eos1, list_rates, dt, matrix_func=cf4_f2) res3 = self.operator(conc_eos3, power) # Step 4: deplete with two matrix exponentials list_rates = list(zip(bos_rates, res1.rates, res2.rates, res3.rates)) - time4, conc_inter = timed_deplete( - self.chain, bos_conc, list_rates, dt, matrix_func=cf4_f3) - time5, conc_eos5 = timed_deplete( - self.chain, conc_inter, list_rates, dt, matrix_func=cf4_f4) + time4, conc_inter = self._timed_deplete( + bos_conc, list_rates, dt, matrix_func=cf4_f3) + time5, conc_eos5 = self._timed_deplete( + conc_inter, list_rates, dt, matrix_func=cf4_f4) return (time1 + time2 + time3 + time4 + time5, [conc_eos1, conc_eos2, conc_eos3, conc_eos5], @@ -360,6 +443,17 @@ class CELIIntegrator(Integrator): that the values are given in burnup (MW-d of energy deposited per kilogram of initial heavy metal). + .. versionadded:: 0.12 + solver : str or callable, optional + If a string, must be the name of the solver responsible for + solving the Bateman equations. Current options are: + + * ``cram16`` - 16th order IPF CRAM + * ``cram48`` - 48th order IPF CRAM [default] + + If a function or other callable, must adhere to the requirements in + :attr:`solver`. + .. versionadded:: 0.12 Attributes @@ -372,6 +466,23 @@ class CELIIntegrator(Integrator): Size of each depletion interval in [s] power : iterable of float Power of the reactor in [W] for each interval in :attr:`timesteps` + solver : callable + Function that will solve the Bateman equations + :math:`\frac{\partial}{\partial t}\vec{n} = A_i\vec{n}_i` with a step + size :math:`t_i`. Can be configured using the ``solver`` argument. + User-supplied functions are expected to have the following signature: + ``solver(A, n0, t) -> n1`` where + + * ``A`` is a :class:`scipy.sparse.csr_matrix` making up the + depletion matrix + * ``n0`` is a 1-D :class:`numpy.ndarray` of initial compositions + for a given material in atoms/cm3 + * ``t`` is a float of the time step size in seconds, and + * ``n1`` is a :class:`numpy.ndarray` of compositions at the + next time step. Expected to be of the same shape as ``n0`` + + .. versionadded:: 0.12 + """ _num_stages = 2 @@ -403,17 +514,17 @@ class CELIIntegrator(Integrator): simulation """ # deplete to end using BOS rates - proc_time, conc_ce = timed_deplete(self.chain, bos_conc, rates, dt) + proc_time, conc_ce = self._timed_deplete(bos_conc, rates, dt) res_ce = self.operator(conc_ce, power) # deplete using two matrix exponentials list_rates = list(zip(rates, res_ce.rates)) - time_le1, conc_inter = timed_deplete( - self.chain, bos_conc, list_rates, dt, matrix_func=celi_f1) + time_le1, conc_inter = self._timed_deplete( + bos_conc, list_rates, dt, matrix_func=celi_f1) - time_le2, conc_end = timed_deplete( - self.chain, conc_inter, list_rates, dt, matrix_func=celi_f2) + time_le2, conc_end = self._timed_deplete( + conc_inter, list_rates, dt, matrix_func=celi_f2) return proc_time + time_le1 + time_le1, [conc_ce, conc_end], [res_ce] @@ -463,6 +574,17 @@ class EPCRK4Integrator(Integrator): that the values are given in burnup (MW-d of energy deposited per kilogram of initial heavy metal). + .. versionadded:: 0.12 + solver : str or callable, optional + If a string, must be the name of the solver responsible for + solving the Bateman equations. Current options are: + + * ``cram16`` - 16th order IPF CRAM + * ``cram48`` - 48th order IPF CRAM [default] + + If a function or other callable, must adhere to the requirements in + :attr:`solver`. + .. versionadded:: 0.12 Attributes @@ -475,6 +597,18 @@ class EPCRK4Integrator(Integrator): Size of each depletion interval in [s] power : iterable of float Power of the reactor in [W] for each interval in :attr:`timesteps` + solver : str or callable, optional + If a string, must be the name of the solver responsible for + solving the Bateman equations. Current options are: + + * ``cram16`` - 16th order IPF CRAM + * ``cram48`` - 48th order IPF CRAM [default] + + If a function or other callable, must adhere to the requirements in + :attr:`solver`. + + .. versionadded:: 0.12 + """ _num_stages = 4 @@ -507,24 +641,23 @@ class EPCRK4Integrator(Integrator): """ # Step 1: deplete with matrix A(y0) / 2 - time1, conc1 = timed_deplete( - self.chain, conc, rates, dt, matrix_func=rk4_f1) + time1, conc1 = self._timed_deplete( + conc, rates, dt, matrix_func=rk4_f1) res1 = self.operator(conc1, power) # Step 2: deplete with matrix A(y1) / 2 - time2, conc2 = timed_deplete( - self.chain, conc, res1.rates, dt, matrix_func=rk4_f1) + time2, conc2 = self._timed_deplete( + conc, res1.rates, dt, matrix_func=rk4_f1) res2 = self.operator(conc2, power) # Step 3: deplete with matrix A(y2) - time3, conc3 = timed_deplete( - self.chain, conc, res2.rates, dt) + time3, conc3 = self._timed_deplete(conc, res2.rates, dt) res3 = self.operator(conc3, power) # Step 4: deplete with matrix built from weighted rates list_rates = list(zip(rates, res1.rates, res2.rates, res3.rates)) - time4, conc4 = timed_deplete( - self.chain, conc, list_rates, dt, matrix_func=rk4_f4) + time4, conc4 = self._timed_deplete( + conc, list_rates, dt, matrix_func=rk4_f4) return (time1 + time2 + time3 + time4, [conc1, conc2, conc3, conc4], [res1, res2, res3]) @@ -586,6 +719,17 @@ class LEQIIntegrator(Integrator): that the values are given in burnup (MW-d of energy deposited per kilogram of initial heavy metal). + .. versionadded:: 0.12 + solver : str or callable, optional + If a string, must be the name of the solver responsible for + solving the Bateman equations. Current options are: + + * ``cram16`` - 16th order IPF CRAM + * ``cram48`` - 48th order IPF CRAM [default] + + If a function or other callable, must adhere to the requirements in + :attr:`solver`. + .. versionadded:: 0.12 Attributes @@ -598,6 +742,23 @@ class LEQIIntegrator(Integrator): Size of each depletion interval in [s] power : iterable of float Power of the reactor in [W] for each interval in :attr:`timesteps` + solver : callable + Function that will solve the Bateman equations + :math:`\frac{\partial}{\partial t}\vec{n} = A_i\vec{n}_i` with a step + size :math:`t_i`. Can be configured using the ``solver`` argument. + User-supplied functions are expected to have the following signature: + ``solver(A, n0, t) -> n1`` where + + * ``A`` is a :class:`scipy.sparse.csr_matrix` making up the + depletion matrix + * ``n0`` is a 1-D :class:`numpy.ndarray` of initial compositions + for a given material in atoms/cm3 + * ``t`` is a float of the time step size in seconds, and + * ``n1`` is a :class:`numpy.ndarray` of compositions at the + next time step. Expected to be of the same shape as ``n0`` + + .. versionadded:: 0.12 + """ _num_stages = 2 @@ -645,10 +806,10 @@ class LEQIIntegrator(Integrator): le_inputs = list(zip( self._prev_rates, bos_res.rates, repeat(prev_dt), repeat(dt))) - time1, conc_inter = timed_deplete( - self.chain, bos_conc, le_inputs, dt, matrix_func=leqi_f1) - time2, conc_eos0 = timed_deplete( - self.chain, conc_inter, le_inputs, dt, matrix_func=leqi_f2) + time1, conc_inter = self._timed_deplete( + bos_conc, le_inputs, dt, matrix_func=leqi_f1) + time2, conc_eos0 = self._timed_deplete( + conc_inter, le_inputs, dt, matrix_func=leqi_f2) res_inter = self.operator(conc_eos0, power) @@ -656,10 +817,10 @@ class LEQIIntegrator(Integrator): self._prev_rates, bos_res.rates, res_inter.rates, repeat(prev_dt), repeat(dt))) - time3, conc_inter = timed_deplete( - self.chain, bos_conc, qi_inputs, dt, matrix_func=leqi_f3) - time4, conc_eos1 = timed_deplete( - self.chain, conc_inter, qi_inputs, dt, matrix_func=leqi_f4) + time3, conc_inter = self._timed_deplete( + bos_conc, qi_inputs, dt, matrix_func=leqi_f3) + time4, conc_eos1 = self._timed_deplete( + conc_inter, qi_inputs, dt, matrix_func=leqi_f4) # store updated rates self._prev_rates = copy.deepcopy(bos_res.rates) @@ -710,6 +871,17 @@ class SICELIIntegrator(SIIntegrator): n_steps : int, optional Number of stochastic iterations per depletion interval. Must be greater than zero. Default : 10 + solver : str or callable, optional + If a string, must be the name of the solver responsible for + solving the Bateman equations. Current options are: + + * ``cram16`` - 16th order IPF CRAM + * ``cram48`` - 48th order IPF CRAM [default] + + If a function or other callable, must adhere to the requirements in + :attr:`solver`. + + .. versionadded:: 0.12 Attributes ---------- @@ -723,6 +895,23 @@ class SICELIIntegrator(SIIntegrator): Power of the reactor in [W] for each interval in :attr:`timesteps` n_steps : int Number of stochastic iterations per depletion interval + solver : callable + Function that will solve the Bateman equations + :math:`\frac{\partial}{\partial t}\vec{n} = A_i\vec{n}_i` with a step + size :math:`t_i`. Can be configured using the ``solver`` argument. + User-supplied functions are expected to have the following signature: + ``solver(A, n0, t) -> n1`` where + + * ``A`` is a :class:`scipy.sparse.csr_matrix` making up the + depletion matrix + * ``n0`` is a 1-D :class:`numpy.ndarray` of initial compositions + for a given material in atoms/cm3 + * ``t`` is a float of the time step size in seconds, and + * ``n1`` is a :class:`numpy.ndarray` of compositions at the + next time step. Expected to be of the same shape as ``n0`` + + .. versionadded:: 0.12 + """ _num_stages = 2 @@ -753,8 +942,7 @@ class SICELIIntegrator(SIIntegrator): Eigenvalue and reaction rates from intermediate transport simulations """ - proc_time, eos_conc = timed_deplete( - self.chain, bos_conc, bos_rates, dt) + proc_time, eos_conc = self._timed_deplete(bos_conc, bos_rates, dt) inter_conc = copy.deepcopy(eos_conc) # Begin iteration @@ -769,10 +957,10 @@ class SICELIIntegrator(SIIntegrator): res_bar = OperatorResult(k, rates) list_rates = list(zip(bos_rates, res_bar.rates)) - time1, inter_conc = timed_deplete( - self.chain, bos_conc, list_rates, dt, matrix_func=celi_f1) - time2, inter_conc = timed_deplete( - self.chain, inter_conc, list_rates, dt, matrix_func=celi_f2) + time1, inter_conc = self._timed_deplete( + bos_conc, list_rates, dt, matrix_func=celi_f1) + time2, inter_conc = self._timed_deplete( + inter_conc, list_rates, dt, matrix_func=celi_f2) proc_time += time1 + time2 # end iteration @@ -820,6 +1008,17 @@ class SILEQIIntegrator(SIIntegrator): n_steps : int, optional Number of stochastic iterations per depletion interval. Must be greater than zero. Default : 10 + solver : str or callable, optional + If a string, must be the name of the solver responsible for + solving the Bateman equations. Current options are: + + * ``cram16`` - 16th order IPF CRAM + * ``cram48`` - 48th order IPF CRAM [default] + + If a function or other callable, must adhere to the requirements in + :attr:`solver`. + + .. versionadded:: 0.12 Attributes ---------- @@ -833,6 +1032,24 @@ class SILEQIIntegrator(SIIntegrator): Power of the reactor in [W] for each interval in :attr:`timesteps` n_steps : int Number of stochastic iterations per depletion interval + solver : callable + Function that will solve the Bateman equations + :math:`\frac{\partial}{\partial t}\vec{n} = A_i\vec{n}_i` with a step + size :math:`t_i`. Can be configured using the ``solver`` argument. + User-supplied functions are expected to have the following signature: + ``solver(A, n0, t) -> n1`` where + + * ``A`` is a :class:`scipy.sparse.csr_matrix` making up the + depletion matrix + * ``n0`` is a 1-D :class:`numpy.ndarray` of initial compositions + for a given material in atoms/cm3 + * ``t`` is a float of the time step size in seconds, and + * ``n1`` is a :class:`numpy.ndarray` of compositions at the + next time step. Expected to be of the same shape as ``n0`` + + .. versionadded:: 0.12 + + """ _num_stages = 2 @@ -879,10 +1096,10 @@ class SILEQIIntegrator(SIIntegrator): # Perform remaining LE/QI inputs = list(zip(self._prev_rates, bos_rates, repeat(prev_dt), repeat(dt))) - proc_time, inter_conc = timed_deplete( - self.chain, bos_conc, inputs, dt, matrix_func=leqi_f1) - time1, eos_conc = timed_deplete( - self.chain, inter_conc, inputs, dt, matrix_func=leqi_f2) + proc_time, inter_conc = self._timed_deplete( + bos_conc, inputs, dt, matrix_func=leqi_f1) + time1, eos_conc = self._timed_deplete( + inter_conc, inputs, dt, matrix_func=leqi_f2) proc_time += time1 inter_conc = copy.deepcopy(eos_conc) @@ -899,10 +1116,10 @@ class SILEQIIntegrator(SIIntegrator): inputs = list(zip(self._prev_rates, bos_rates, res_bar.rates, repeat(prev_dt), repeat(dt))) - time1, inter_conc = timed_deplete( - self.chain, bos_conc, inputs, dt, matrix_func=leqi_f3) - time2, inter_conc = timed_deplete( - self.chain, inter_conc, inputs, dt, matrix_func=leqi_f4) + time1, inter_conc = self._timed_deplete( + bos_conc, inputs, dt, matrix_func=leqi_f3) + time2, inter_conc = self._timed_deplete( + inter_conc, inputs, dt, matrix_func=leqi_f4) proc_time += time1 + time2 return proc_time, [eos_conc, inter_conc], [res_bar] diff --git a/openmc/deplete/pool.py b/openmc/deplete/pool.py new file mode 100644 index 0000000000..d9d3597751 --- /dev/null +++ b/openmc/deplete/pool.py @@ -0,0 +1,58 @@ +"""Dedicated module containing depletion function + +Provided to avoid some circular imports +""" +from itertools import repeat +from multiprocessing import Pool + + +def deplete(func, chain, x, rates, dt, matrix_func=None): + """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) -> n1`` + chain : openmc.deplete.Chain + Depletion chain + x : list of numpy.ndarray + Atom number vectors for each material + rates : openmc.deplete.ReactionRates + Reaction rates (from transport operator) + dt : float + Time in [s] to deplete for + maxtrix_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`` + + Returns + ------- + x_result : list of numpy.ndarray + Updated atom number vectors for each material + + """ + + fission_yields = chain.fission_yields + if len(fission_yields) == 1: + fission_yields = repeat(fission_yields[0]) + elif len(fission_yields) != len(x): + raise ValueError( + "Number of material fission yield distributions {} is not equal " + "to the number of compositions {}".format( + len(fission_yields), len(x))) + + if matrix_func is None: + matrices = map(chain.form_matrix, rates, fission_yields) + else: + matrices = map(matrix_func, repeat(chain), rates, fission_yields) + + # Use multiprocessing pool to distribute work + with Pool() as pool: + inputs = zip(matrices, x, repeat(dt)) + x_result = list(pool.starmap(func, inputs)) + + return x_result diff --git a/tests/unit_tests/test_deplete_chain.py b/tests/unit_tests/test_deplete_chain.py index 36791b0182..7162be6de9 100644 --- a/tests/unit_tests/test_deplete_chain.py +++ b/tests/unit_tests/test_deplete_chain.py @@ -7,7 +7,7 @@ import os from pathlib import Path import numpy as np -from openmc.deplete import comm, Chain, reaction_rates, nuclide, cram +from openmc.deplete import comm, Chain, reaction_rates, nuclide, cram, pool import pytest from tests import cdtemp @@ -414,7 +414,7 @@ def test_fission_yield_attribute(simple_chain): dummy_conc = [[1, 2]] * (len(empty_chain.fission_yields) + 1) with pytest.raises( ValueError, match="fission yield.*not equal.*compositions"): - cram.deplete(empty_chain, dummy_conc, None, 0.5) + pool.deplete(cram.CRAM48, empty_chain, dummy_conc, None, 0.5) def test_validate(simple_chain): diff --git a/tests/unit_tests/test_deplete_integrator.py b/tests/unit_tests/test_deplete_integrator.py index 6c09b3feac..0560cb9f73 100644 --- a/tests/unit_tests/test_deplete_integrator.py +++ b/tests/unit_tests/test_deplete_integrator.py @@ -17,7 +17,8 @@ import pytest from openmc.deplete import ( ReactionRates, Results, ResultsList, comm, OperatorResult, PredictorIntegrator, CECMIntegrator, CF4Integrator, CELIIntegrator, - EPCRK4Integrator, LEQIIntegrator, SICELIIntegrator, SILEQIIntegrator) + EPCRK4Integrator, LEQIIntegrator, SICELIIntegrator, SILEQIIntegrator, + cram) from tests import dummy_operator @@ -147,6 +148,23 @@ def test_bad_integrator_inputs(): with pytest.raises(ValueError, match="n_steps"): SICELIIntegrator(op, timesteps, [1], n_steps=0) + with pytest.raises(ValueError, match="Solver failure"): + PredictorIntegrator(op, timesteps, power=1, solver="failure") + + with pytest.raises(TypeError, match=".*callable.*NoneType"): + PredictorIntegrator(op, timesteps, power=1, solver=None) + + with pytest.raises(ValueError, match=".*arguments"): + PredictorIntegrator(op, timesteps, power=1, solver=mock_bad_solver_nargs) + + +def mock_good_solver(A, n, t): + pass + + +def mock_bad_solver_nargs(A, n): + pass + @pytest.mark.parametrize("scheme", dummy_operator.SCHEMES) def test_integrator(run_in_tmpdir, scheme): @@ -174,6 +192,19 @@ def test_integrator(run_in_tmpdir, scheme): assert dep_time.shape == (2, ) assert all(dep_time > 0) + integrator = bundle.solver(operator, [0.75], 1, solver=cram.CRAM48) + assert integrator.solver is cram.CRAM48 + + integrator = bundle.solver(operator, [0.75], 1, solver="cram16") + assert integrator.solver is cram.CRAM16 + + integrator.solver = mock_good_solver + assert integrator.solver is mock_good_solver + + lfunc = lambda A, n, t: mock_good_solver(A, n, t) + integrator.solver = lfunc + assert integrator.solver is lfunc + @pytest.mark.parametrize("integrator", INTEGRATORS) def test_timesteps(integrator):