diff --git a/openmc/deplete/abc.py b/openmc/deplete/abc.py index 4a92337e03a..6389ba4b52d 100644 --- a/openmc/deplete/abc.py +++ b/openmc/deplete/abc.py @@ -3,6 +3,7 @@ This module contains Abstract Base Classes for implementing operator, integrator, depletion system solver, and operator helper classes """ +from __future__ import annotations from abc import ABC, abstractmethod from collections import namedtuple, defaultdict from collections.abc import Iterable, Callable @@ -13,17 +14,21 @@ import os from pathlib import Path import time +from typing import Optional, Union, Sequence from warnings import warn -from numpy import nonzero, empty, asarray, take +import numpy as np + from uncertainties import ufloat from openmc.checkvalue import check_type, check_greater_than, PathLike, check_value from openmc.mpi import comm +from openmc import Material from .stepresult import StepResult from .chain import Chain from .results import Results from .pool import deplete +from .reaction_rates import ReactionRates from .transfer_rates import TransferRates from openmc import Material, Cell from .batchwise import (BatchwisePure, BatchwiseCellGeometrical, BatchwiseCellTemperature, @@ -179,7 +184,7 @@ def finalize(self): pass @abstractmethod - def write_bos_data(self, step): + def write_bos_data(self, step: int): """Document beginning of step data for a given step Called at the beginning of a depletion step and at @@ -218,7 +223,7 @@ class ReactionRateHelper(ABC): def __init__(self, n_nucs, n_react): self._nuclides = None - self._results_cache = empty((n_nucs, n_react)) + self._results_cache = np.empty((n_nucs, n_react)) @abstractmethod def generate_tallies(self, materials, scores): @@ -235,7 +240,12 @@ def nuclides(self, nuclides): self._nuclides = nuclides @abstractmethod - def get_material_rates(self, mat_id, nuc_index, react_index): + def get_material_rates( + self, + mat_id: int, + nuc_index: Sequence[str], + react_index: Sequence[str] + ): """Return 2D array of [nuclide, reaction] reaction rates Parameters @@ -248,7 +258,7 @@ def get_material_rates(self, mat_id, nuc_index, react_index): Ordering of reactions """ - def divide_by_atoms(self, number): + def divide_by_atoms(self, number: Sequence[float]): """Normalize reaction rates by number of atoms Acts on the current material examined by :meth:`get_material_rates` @@ -265,7 +275,7 @@ def divide_by_atoms(self, number): normalized by the number of nuclides """ - mask = nonzero(number) + mask = np.nonzero(number) results = self._results_cache for col in range(results.shape[1]): results[mask, col] /= number[mask] @@ -297,7 +307,7 @@ def reset(self): """Reset state for normalization""" @abstractmethod - def prepare(self, chain_nucs, rate_index): + def prepare(self, chain_nucs: Sequence[str], rate_index: dict): """Perform work needed to obtain energy produced This method is called prior to calculating the reaction rates @@ -336,7 +346,7 @@ def nuclides(self, nuclides): self._nuclides = nuclides @abstractmethod - def factor(self, source_rate): + def factor(self, source_rate: float): """Return normalization factor Parameters @@ -439,7 +449,7 @@ def generate_tallies(materials, mat_indexes): in parallel mode. """ - def update_tally_nuclides(self, nuclides): + def update_tally_nuclides(self, nuclides: Sequence[str]) -> list: """Return nuclides with non-zero densities and yield data Parameters @@ -565,8 +575,16 @@ class Integrator(ABC): """ - def __init__(self, operator, timesteps, power=None, power_density=None, - source_rates=None, timestep_units='s', solver="cram48"): + def __init__( + self, + operator: TransportOperator, + timesteps: Sequence[float], + power: Optional[Union[float, Sequence[float]]] = None, + power_density: Optional[Union[float, Sequence[float]]] = None, + source_rates: Optional[Sequence[float]] = None, + timestep_units: str = 's', + solver: str = "cram48" + ): # Check number of stages previously used if operator.prev_res is not None: res = operator.prev_res[-1] @@ -638,8 +656,8 @@ def __init__(self, operator, timesteps, power=None, power_density=None, else: raise ValueError("Invalid timestep unit '{}'".format(unit)) - self.timesteps = asarray(seconds) - self.source_rates = asarray(source_rates) + self.timesteps = np.asarray(seconds) + self.source_rates = np.asarray(source_rates) self.transfer_rates = None self.batchwise = None @@ -699,7 +717,14 @@ def _timed_deplete(self, n, rates, dt, matrix_func=None): return time.time() - start, results @abstractmethod - def __call__(self, n, rates, dt, source_rate, i): + def __call__( + self, + n: Sequence[np.ndarray], + rates: ReactionRates, + dt: float, + source_rate: float, + i: int + ): """Perform the integration across one time step Parameters @@ -980,8 +1005,8 @@ def integrate_adaptive(self, material_id, nuclides, final_step=True, def _adapt_timestep(self, n_pred, n_corr, mat_idx, nuc_ids, tol, err, rho1, rho2, f ): - filt_pred = take(n_pred[mat_idx], nuc_ids) - filt_corr = take(n_corr[mat_idx], nuc_ids) + filt_pred = np.take(n_pred[mat_idx], nuc_ids) + filt_corr = np.take(n_corr[mat_idx], nuc_ids) err_vec = abs(filt_corr - filt_pred) x = min((tol + err * filt_corr) / err_vec) @@ -990,8 +1015,14 @@ def _adapt_timestep(self, n_pred, n_corr, mat_idx, nuc_ids, tol, err, rho1, print(f'Timestep adapting factor: {adapt}') return adapt - def add_transfer_rate(self, material, components, transfer_rate, - transfer_rate_units='1/s', destination_material=None): + def add_transfer_rate( + self, + material: Union[str, int, Material], + components: Sequence[str], + transfer_rate: float, + transfer_rate_units: str = '1/s', + destination_material: Optional[Union[str, int, Material]] = None + ): """Add transfer rates to depletable material. Parameters @@ -1172,9 +1203,17 @@ class SIIntegrator(Integrator): """ - def __init__(self, operator, timesteps, power=None, power_density=None, - source_rates=None, timestep_units='s', n_steps=10, - solver="cram48"): + def __init__( + self, + operator: TransportOperator, + timesteps: Sequence[float], + power: Optional[Union[float, Sequence[float]]] = None, + power_density: Optional[Union[float, Sequence[float]]] = None, + source_rates: Optional[Sequence[float]] = None, + timestep_units: str = 's', + n_steps: int = 10, + solver: str = "cram48" + ): check_type("n_steps", n_steps, Integral) check_greater_than("n_steps", n_steps, 0) super().__init__(