Allowing chain_file to be chain object to save reloading time (#3436)

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

openmc/deplete/abc.py

@@ -24,7 +24,7 @@ from openmc.mpi import comm
 from openmc.utility_funcs import change_directory
 from openmc import Material
 from .stepresult import StepResult
-from .chain import Chain
+from .chain import _get_chain
 from .results import Results, _SECONDS_PER_MINUTE, _SECONDS_PER_HOUR, \
     _SECONDS_PER_DAY, _SECONDS_PER_JULIAN_YEAR
 from .pool import deplete
@@ -126,8 +126,8 @@ class TransportOperator(ABC):
 
     Parameters
     ----------
-    chain_file : str
-        Path to the depletion chain XML file
+    chain_file : PathLike or Chain
+        Path to the depletion chain XML file or instance of openmc.deplete.Chain.
     fission_q : dict, optional
         Dictionary of nuclides and their fission Q values [eV]. If not given,
         values will be pulled from the ``chain_file``.
@@ -145,11 +145,12 @@ class TransportOperator(ABC):
         The depletion chain information necessary to form matrices and tallies.
 
     """
-    def __init__(self, chain_file, fission_q=None, prev_results=None):
+    def __init__(self, chain_file=None, fission_q=None, prev_results=None):
         self.output_dir = '.'
 
         # Read depletion chain
-        self.chain = Chain.from_xml(chain_file, fission_q)
+        self.chain = _get_chain(chain_file, fission_q)
+
         if prev_results is None:
             self.prev_res = None
         else:

openmc/deplete/chain.py

@@ -17,8 +17,9 @@ from typing import List
 import lxml.etree as ET
 import scipy.sparse as sp
 
-from openmc.checkvalue import check_type, check_greater_than
+from openmc.checkvalue import check_type, check_greater_than, PathLike
 from openmc.data import gnds_name, zam
+from openmc.exceptions import DataError
 from .nuclide import FissionYieldDistribution, Nuclide
 import openmc.data
 
@@ -1246,3 +1247,38 @@ class Chain:
         found.update(isotopes)
 
         return found
+
+
+def _get_chain(
+    chain_file: PathLike | Chain | None = None,
+    fission_q: dict | None = None
+) -> Chain:
+    """Get a depletion chain from a file or the runtime configuration.
+
+    Parameters
+    ----------
+    chain_file : PathLike or Chain, optional
+        Path to depletion chain XML file, a Chain instance, or None to use
+        the file specified in ``openmc.config['chain_file']``.
+    fission_q : dict, optional
+        Dictionary of nuclides and their fission Q values [eV]. If not given,
+        values will be pulled from the ``chain_file``.
+
+    Returns
+    -------
+    Chain
+        Depletion chain instance.
+    """
+    if isinstance(chain_file, Chain):
+        return chain_file
+    elif isinstance(chain_file, PathLike | None):
+        if chain_file is None:
+            chain_file = openmc.config.get('chain_file')
+            if 'chain_file' not in openmc.config:
+                raise DataError(
+                    "No depletion chain specified and could not find depletion "
+                    "chain in openmc.config['chain_file']"
+                )
+        return Chain.from_xml(chain_file, fission_q)
+    else:
+        raise TypeError("chain_file must be path-like, a Chain, or None")

openmc/deplete/coupled_operator.py

@@ -102,9 +102,9 @@ class CoupledOperator(OpenMCOperator):
     ----------
     model : openmc.model.Model
         OpenMC model object
-    chain_file : str, optional
-        Path to the depletion chain XML file. Defaults to
-        ``openmc.config['chain_file']``.
+    chain_file : PathLike or Chain, optional
+        Path to the depletion chain XML file or instance of openmc.deplete.Chain.
+        Defaults to ``openmc.config['chain_file']``.
     prev_results : Results, optional
         Results from a previous depletion calculation. If this argument is
         specified, the depletion calculation will start from the latest state

openmc/deplete/d1s.py

@@ -14,19 +14,20 @@ import numpy as np
 import openmc
 from openmc.data import half_life
 from .abc import _normalize_timesteps
-from .chain import Chain
+from .chain import Chain, _get_chain
+from ..checkvalue import PathLike
 
 
-def get_radionuclides(model: openmc.Model, chain_file: str | None = None) -> list[str]:
+def get_radionuclides(model: openmc.Model, chain_file: PathLike | Chain | None = None) -> list[str]:
     """Determine all radionuclides that can be produced during D1S.
 
     Parameters
     ----------
     model : openmc.Model
         Model that should be used for determining what nuclides are present
-    chain_file : str, optional
-        Which chain file to use for inspecting decay data. If None is passed,
-        defaults to ``openmc.config['chain_file']``
+    chain_file : PathLike | Chain
+        Path to the depletion chain XML file or instance of openmc.deplete.Chain.
+        Used for inspecting decay data. Defaults to ``openmc.config['chain_file']``
 
     Returns
     -------
@@ -39,9 +40,7 @@ def get_radionuclides(model: openmc.Model, chain_file: str | None = None) -> lis
                       for nuc in mat.get_nuclides()}
 
     # Load chain file
-    if chain_file is None:
-        chain_file = openmc.config['chain_file']
-    chain = Chain.from_xml(chain_file)
+    chain = _get_chain(chain_file)
 
     radionuclides = set()
     for nuclide in chain.nuclides:

openmc/deplete/independent_operator.py

@@ -50,9 +50,9 @@ class IndependentOperator(OpenMCOperator):
         Cross sections in [b] for each domain. If the
         :class:`~openmc.deplete.MicroXS` object is empty, a decay-only
         calculation will be run.
-    chain_file : str
-        Path to the depletion chain XML file. Defaults to
-        ``openmc.config['chain_file']``.
+    chain_file : PathLike or Chain, optional
+        Path to the depletion chain XML file or instance of openmc.deplete.Chain.
+        Defaults to ``openmc.config['chain_file']``.
     keff : 2-tuple of float, optional
        keff eigenvalue and uncertainty from transport calculation.
     prev_results : Results, optional
@@ -179,9 +179,9 @@ class IndependentOperator(OpenMCOperator):
         micro_xs : MicroXS
             Cross sections in [b]. If the :class:`~openmc.deplete.MicroXS`
             object is empty, a decay-only calculation will be run.
-        chain_file : str, optional
-            Path to the depletion chain XML file. Defaults to
-            ``openmc.config['chain_file']``.
+        chain_file : PathLike or Chain, optional
+            Path to the depletion chain XML file or instance of
+            openmc.deplete.Chain. Defaults to ``openmc.config['chain_file']``.
         nuc_units : {'atom/cm3', 'atom/b-cm'}, optional
             Units for nuclide concentration.
         keff : 2-tuple of float, optional

openmc/deplete/microxs.py

@@ -12,13 +12,12 @@ import pandas as pd
 import numpy as np
 
 from openmc.checkvalue import check_type, check_value, check_iterable_type, PathLike
-from openmc.exceptions import DataError
 from openmc.utility_funcs import change_directory
 from openmc import StatePoint
 from openmc.mgxs import GROUP_STRUCTURES
 from openmc.data import REACTION_MT
 import openmc
-from .chain import Chain, REACTIONS
+from .chain import Chain, REACTIONS, _get_chain
 from .coupled_operator import _find_cross_sections, _get_nuclides_with_data
 import openmc.lib
 from openmc.mpi import comm
@@ -28,24 +27,13 @@ _valid_rxns.append('fission')
 _valid_rxns.append('damage-energy')
 
 
-def _resolve_chain_file_path(chain_file: str | None):
-    if chain_file is None:
-        chain_file = openmc.config.get('chain_file')
-        if 'chain_file' not in openmc.config:
-            raise DataError(
-                "No depletion chain specified and could not find depletion "
-                "chain in openmc.config['chain_file']"
-            )
-    return chain_file
-
-
 def get_microxs_and_flux(
         model: openmc.Model,
         domains,
         nuclides: Iterable[str] | None = None,
         reactions: Iterable[str] | None = None,
         energies: Iterable[float] | str | None = None,
-        chain_file: PathLike | None = None,
+        chain_file: PathLike | Chain | None = None,
         run_kwargs=None
     ) -> tuple[list[np.ndarray], list[MicroXS]]:
     """Generate a microscopic cross sections and flux from a Model
@@ -66,9 +54,9 @@ def get_microxs_and_flux(
         reactions listed in the depletion chain file are used.
     energies : iterable of float or str
         Energy group boundaries in [eV] or the name of the group structure
-    chain_file : str, optional
-        Path to the depletion chain XML file that will be used in depletion
-        simulation. Used to determine cross sections for materials not
+    chain_file : PathLike or Chain, optional
+        Path to the depletion chain XML file or an instance of
+        openmc.deplete.Chain. Used to determine cross sections for materials not
         present in the inital composition. Defaults to
         ``openmc.config['chain_file']``.
     run_kwargs : dict, optional
@@ -86,8 +74,7 @@ def get_microxs_and_flux(
     original_tallies = model.tallies
 
     # Determine what reactions and nuclides are available in chain
-    chain_file = _resolve_chain_file_path(chain_file)
-    chain = Chain.from_xml(chain_file)
+    chain = _get_chain(chain_file)
     if reactions is None:
         reactions = chain.reactions
     if not nuclides:
@@ -245,9 +232,9 @@ class MicroXS:
             Energy group boundaries in [eV] or the name of the group structure
         multi_group_flux : iterable of float
             Energy-dependent multigroup flux values
-        chain_file : str, optional
-            Path to the depletion chain XML file that will be used in depletion
-            simulation.  Defaults to ``openmc.config['chain_file']``.
+        chain_file : PathLike or Chain, optional
+            Path to the depletion chain XML file or an instance of
+            openmc.deplete.Chain. Defaults to ``openmc.config['chain_file']``.
         temperature : int, optional
             Temperature for cross section evaluation in [K].
         nuclides : list of str, optional
@@ -278,9 +265,7 @@ class MicroXS:
         if len(multigroup_flux) != len(energies) - 1:
             raise ValueError('Length of flux array should be len(energies)-1')
 
-        chain_file_path = _resolve_chain_file_path(chain_file)
-        chain = Chain.from_xml(chain_file_path)
-
+        chain = _get_chain(chain_file)
         cross_sections = _find_cross_sections(model=None)
         nuclides_with_data = _get_nuclides_with_data(cross_sections)
 

openmc/deplete/openmc_operator.py

@@ -36,9 +36,9 @@ class OpenMCOperator(TransportOperator):
     cross_sections : str or list of MicroXS
         Path to continuous energy cross section library, or list of objects
         containing cross sections.
-    chain_file : str, optional
-        Path to the depletion chain XML file. Defaults to
-        openmc.config['chain_file'].
+    chain_file : PathLike or Chain, optional
+        Path to the depletion chain XML file or instance of openmc.deplete.Chain.
+        Defaults to ``openmc.config['chain_file']``.
     prev_results : Results, optional
         Results from a previous depletion calculation. If this argument is
         specified, the depletion calculation will start from the latest state

tests/unit_tests/test_d1s.py

@@ -22,7 +22,8 @@ def model():
 
 def test_get_radionuclides(model):
     # Check that radionuclides are correct and are unstable
-    nuclides = d1s.get_radionuclides(model, CHAIN_PATH)
+    chain = openmc.deplete.Chain.from_xml(CHAIN_PATH)
+    nuclides = d1s.get_radionuclides(model, chain)
     assert sorted(nuclides) == [
         'Co58', 'Co60', 'Co61', 'Co62', 'Co64',
         'Fe55', 'Fe59', 'Fe61', 'Ni57', 'Ni59', 'Ni63', 'Ni65'

tests/unit_tests/test_deplete_coupled_operator.py

@@ -5,7 +5,7 @@
 from pathlib import Path
 
 import pytest
-from openmc.deplete import CoupledOperator
+from openmc.deplete import CoupledOperator, Chain
 import openmc
 import numpy as np
 
@@ -106,11 +106,13 @@ def test_diff_volume_method_match_cell(model_with_volumes):
 def test_diff_volume_method_divide_equally(model_with_volumes):
     """Tests the volumes assigned to the materials are divided equally"""
 
+    chain = Chain.from_xml(CHAIN_PATH)
+
     operator = openmc.deplete.CoupledOperator(
         model=model_with_volumes,
         diff_burnable_mats=True,
         diff_volume_method='divide equally',
-        chain_file=CHAIN_PATH
+        chain_file=chain
     )
 
     all_cells = list(operator.model.geometry.get_all_cells().values())

tests/unit_tests/test_deplete_independent_operator.py

@@ -7,7 +7,7 @@ from pathlib import Path
 import pytest
 
 from openmc import Material
-from openmc.deplete import IndependentOperator, MicroXS
+from openmc.deplete import IndependentOperator, MicroXS, Chain
 
 CHAIN_PATH = Path(__file__).parents[1] / "chain_simple.xml"
 ONE_GROUP_XS = Path(__file__).parents[1] / "micro_xs_simple.csv"
@@ -25,8 +25,9 @@ def test_operator_init():
                 'O17': 1.7588724018066158e+19}
     flux = 1.0
     micro_xs = MicroXS.from_csv(ONE_GROUP_XS)
+    chain = Chain.from_xml(CHAIN_PATH)
     IndependentOperator.from_nuclides(
-        volume, nuclides, flux, micro_xs, CHAIN_PATH, nuc_units='atom/cm3')
+        volume, nuclides, flux, micro_xs, chain, nuc_units='atom/cm3')
 
     fuel = Material(name="uo2")
     fuel.add_element("U", 1, percent_type="ao", enrichment=4.25)