type hinting openmc.deplete.abc.py (openmc-dev#2866)

Co-authored-by: Paul Romano <[email protected]>

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__(