From e0063e9fd24a40f66dff83ee8da9e306f332e38b Mon Sep 17 00:00:00 2001 From: Ethan Peterson Date: Tue, 28 Nov 2023 11:31:45 -0500 Subject: [PATCH] add splu cache to depletion solver --- openmc/deplete/abc.py | 20 ++++++++++++++------ openmc/deplete/cram.py | 18 +++++++++++++----- openmc/deplete/integrators.py | 19 ++++++++++--------- openmc/deplete/pool.py | 6 +++--- 4 files changed, 40 insertions(+), 23 deletions(-) diff --git a/openmc/deplete/abc.py b/openmc/deplete/abc.py index 5d906efd840..9c7aabb74c5 100644 --- a/openmc/deplete/abc.py +++ b/openmc/deplete/abc.py @@ -671,8 +671,8 @@ def solver(self, func): return # Inspect arguments - if len(sig.parameters) != 3: - raise ValueError("Function {} does not support three arguments: " + if len(sig.parameters) < 3: + raise ValueError("Function {} does not support less than three arguments: " "{!s}".format(func, sig)) for ix, param in enumerate(sig.parameters.values()): @@ -683,15 +683,16 @@ def solver(self, func): self._solver = func - def _timed_deplete(self, n, rates, dt, matrix_func=None): + def _timed_deplete(self, n, rates, dt, matrix_func=None, + use_cache=False): start = time.time() results = deplete( self._solver, self.chain, n, rates, dt, matrix_func, - self.transfer_rates) + self.transfer_rates, use_cache=use_cache) return time.time() - start, results @abstractmethod - def __call__(self, n, rates, dt, source_rate, i): + def __call__(self, n, rates, dt, source_rate, i, use_cache=False): """Perform the integration across one time step Parameters @@ -781,7 +782,10 @@ def integrate(self, final_step=True, output=True): n = self.operator.initial_condition() t, self._i_res = self._get_start_data() + prev_dt = None + prev_source_rate = None for i, (dt, source_rate) in enumerate(self): + use_cache = (prev_dt == dt) and (prev_source_rate == source_rate) if output and comm.rank == 0: print(f"[openmc.deplete] t={t} s, dt={dt} s, source={source_rate}") @@ -792,7 +796,9 @@ def integrate(self, final_step=True, output=True): n, res = self._get_bos_data_from_restart(i, source_rate, n) # Solve Bateman equations over time interval - proc_time, n_list, res_list = self(n, res.rates, dt, source_rate, i) + proc_time, n_list, res_list = self(n, res.rates, dt, + source_rate, i, + use_cache=use_cache) # Insert BOS concentration, transport results n_list.insert(0, n) @@ -804,6 +810,8 @@ def integrate(self, final_step=True, output=True): StepResult.save(self.operator, n_list, res_list, [t, t + dt], source_rate, self._i_res + i, proc_time) + prev_dt = dt + prev_source_rate = source_rate t += dt # Final simulation -- in the case that final_step is False, a zero diff --git a/openmc/deplete/cram.py b/openmc/deplete/cram.py index 9575e103580..49e367837c3 100644 --- a/openmc/deplete/cram.py +++ b/openmc/deplete/cram.py @@ -54,8 +54,9 @@ def __init__(self, alpha, theta, alpha0): self.alpha = alpha self.theta = theta self.alpha0 = alpha0 + self._splu_cache = [] - def __call__(self, A, n0, dt): + def __call__(self, A, n0, dt, use_cache=False): """Solve depletion equations using IPF CRAM Parameters @@ -75,11 +76,18 @@ def __call__(self, A, n0, dt): Final compositions after ``dt`` """ - A = dt * sp.csc_matrix(A, dtype=np.float64) y = n0.copy() - ident = sp.eye(A.shape[0], format='csc') - for alpha, theta in zip(self.alpha, self.theta): - y += 2*np.real(alpha*sla.splu(A - theta*ident).solve(y)) + if use_cache: + for alpha, splu in zip(self.alpha, self._splu_cache): + y += 2*np.real(alpha*splu.solve(y)) + else: + A = dt * sp.csc_matrix(A, dtype=np.float64) + ident = sp.eye(A.shape[0], format='csc') + self._splu_cache = [] + for alpha, theta in zip(self.alpha, self.theta): + splu = sla.splu(A - theta*ident) + self._splu_cache.append(splu) + y += 2*np.real(alpha*splu.solve(y)) return y * self.alpha0 diff --git a/openmc/deplete/integrators.py b/openmc/deplete/integrators.py index a877c4900f6..5aaff9f6832 100644 --- a/openmc/deplete/integrators.py +++ b/openmc/deplete/integrators.py @@ -26,7 +26,7 @@ class PredictorIntegrator(Integrator): """ _num_stages = 1 - def __call__(self, n, rates, dt, source_rate, _i=None): + def __call__(self, n, rates, dt, source_rate, _i=None, use_cache=False): """Perform the integration across one time step Parameters @@ -54,7 +54,8 @@ def __call__(self, n, rates, dt, source_rate, _i=None): with predictor """ - proc_time, n_end = self._timed_deplete(n, rates, dt) + proc_time, n_end = self._timed_deplete(n, rates, dt, + use_cache=use_cache) return proc_time, [n_end], [] @@ -78,7 +79,7 @@ class CECMIntegrator(Integrator): """ _num_stages = 2 - def __call__(self, n, rates, dt, source_rate, _i=None): + def __call__(self, n, rates, dt, source_rate, _i=None, use_cache=False): """Integrate using CE/CM Parameters @@ -142,7 +143,7 @@ class CF4Integrator(Integrator): """ _num_stages = 4 - def __call__(self, n_bos, bos_rates, dt, source_rate, _i=None): + def __call__(self, n_bos, bos_rates, dt, source_rate, _i=None, use_cache=False): """Perform the integration across one time step Parameters @@ -220,7 +221,7 @@ class CELIIntegrator(Integrator): """ _num_stages = 2 - def __call__(self, n_bos, rates, dt, source_rate, _i=None): + def __call__(self, n_bos, rates, dt, source_rate, _i=None, use_cache=False): """Perform the integration across one time step Parameters @@ -286,7 +287,7 @@ class EPCRK4Integrator(Integrator): """ _num_stages = 4 - def __call__(self, n, rates, dt, source_rate, _i=None): + def __call__(self, n, rates, dt, source_rate, _i=None, use_cache=False): """Perform the integration across one time step Parameters @@ -368,7 +369,7 @@ class LEQIIntegrator(Integrator): """ _num_stages = 2 - def __call__(self, n_bos, bos_rates, dt, source_rate, i): + def __call__(self, n_bos, bos_rates, dt, source_rate, i, use_cache=False): """Perform the integration across one time step Parameters @@ -450,7 +451,7 @@ class SICELIIntegrator(SIIntegrator): """ _num_stages = 2 - def __call__(self, n_bos, bos_rates, dt, source_rate, _i=None): + def __call__(self, n_bos, bos_rates, dt, source_rate, _i=None, use_cache=False): """Perform the integration across one time step Parameters @@ -516,7 +517,7 @@ class SILEQIIntegrator(SIIntegrator): """ _num_stages = 2 - def __call__(self, n_bos, bos_rates, dt, source_rate, i): + def __call__(self, n_bos, bos_rates, dt, source_rate, i, use_cache=False): """Perform the integration across one time step Parameters diff --git a/openmc/deplete/pool.py b/openmc/deplete/pool.py index 27ecaa4dd8b..38402004f0e 100644 --- a/openmc/deplete/pool.py +++ b/openmc/deplete/pool.py @@ -41,7 +41,7 @@ def _distribute(items): j += chunk_size def deplete(func, chain, n, rates, dt, matrix_func=None, transfer_rates=None, - *matrix_args): + *matrix_args, use_cache=False): """Deplete materials using given reaction rates for a specified time Parameters @@ -140,7 +140,7 @@ def deplete(func, chain, n, rates, dt, matrix_func=None, transfer_rates=None, # Concatenate vectors of nuclides in one n_multi = np.concatenate(n) - n_result = func(matrix, n_multi, dt) + n_result = func(matrix, n_multi, dt, use_cache=use_cache) # 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]) @@ -155,7 +155,7 @@ def deplete(func, chain, n, rates, dt, matrix_func=None, transfer_rates=None, return n_result - inputs = zip(matrices, n, repeat(dt)) + inputs = zip(matrices, n, repeat(dt), repeat(use_cache)) if USE_MULTIPROCESSING: with Pool(NUM_PROCESSES) as pool: