Support flux collapse method in get_microxs_and_flux (#3466)

Co-authored-by: Jonathan Shimwell <drshimwell@gmail.com>

Author: paulromano

openmc/deplete/microxs.py

@@ -5,8 +5,10 @@
 """
 
 from __future__ import annotations
-from collections.abc import Iterable, Sequence
+from collections.abc import Sequence
+import shutil
 from tempfile import TemporaryDirectory
+from typing import Union, TypeAlias
 
 import pandas as pd
 import numpy as np
@@ -27,19 +29,39 @@
 _valid_rxns.append('damage-energy')
 
 
+# TODO: Replace with type statement when support is Python 3.12+
+DomainTypes: TypeAlias = Union[
+    Sequence[openmc.Material],
+    Sequence[openmc.Cell],
+    Sequence[openmc.Universe],
+    openmc.MeshBase,
+    openmc.Filter
+]
+
+
 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 | Chain | None = None,
-        run_kwargs=None
-    ) -> tuple[list[np.ndarray], list[MicroXS]]:
-    """Generate a microscopic cross sections and flux from a Model
+    model: openmc.Model,
+    domains: DomainTypes,
+    nuclides: Sequence[str] | None = None,
+    reactions: Sequence[str] | None = None,
+    energies: Sequence[float] | str | None = None,
+    reaction_rate_mode: str = 'direct',
+    chain_file: PathLike | Chain | None = None,
+    path_statepoint: PathLike | None = None,
+    run_kwargs=None
+) -> tuple[list[np.ndarray], list[MicroXS]]:
+    """Generate microscopic cross sections and fluxes for multiple domains.
+
+    This function runs a neutron transport solve to obtain the flux and reaction
+    rates in the specified domains and computes multigroup microscopic cross
+    sections that can be used in depletion calculations with the
+    :class:`~openmc.deplete.IndependentOperator` class.
 
     .. versionadded:: 0.14.0
 
+    .. versionchanged:: 0.15.3
+        Added `reaction_rate_mode` and `path_statepoint` arguments.
+
     Parameters
     ----------
     model : openmc.Model
@@ -53,12 +75,22 @@ def get_microxs_and_flux(
         Reactions to get cross sections for. If not specified, all neutron
         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
+        Energy group boundaries in [eV] or the name of the group structure.
+        If left as None energies will default to [0.0, 100e6]
+    reaction_rate_mode : {"direct", "flux"}, optional
+        Indicate how reaction rates should be calculated. The "direct" method
+        tallies reaction rates directly. The "flux" method tallies a multigroup
+        flux spectrum and then collapses multigroup reaction rates after a
+        transport solve (with an option to tally some reaction rates directly).
     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']``.
+    path_statepoint : path-like, optional
+        Path to write the statepoint file from the neutron transport solve to.
+        By default, The statepoint file is written to a temporary directory and
+        is not kept.
     run_kwargs : dict, optional
         Keyword arguments passed to :meth:`openmc.Model.run`
 
@@ -69,7 +101,13 @@ def get_microxs_and_flux(
     list of MicroXS
         Cross section data in [b] for each domain
 
+    See Also
+    --------
+    openmc.deplete.IndependentOperator
+
     """
+    check_value('reaction_rate_mode', reaction_rate_mode, {'direct', 'flux'})
+
     # Save any original tallies on the model
     original_tallies = model.tallies
 
@@ -85,8 +123,8 @@ def get_microxs_and_flux(
 
     # Set up the reaction rate and flux tallies
     if energies is None:
-        energy_filter = openmc.EnergyFilter([0.0, 100.0e6])
-    elif isinstance(energies, str):
+        energies = [0.0, 100.0e6]
+    if isinstance(energies, str):
         energy_filter = openmc.EnergyFilter.from_group_structure(energies)
     else:
         energy_filter = openmc.EnergyFilter(energies)
@@ -104,16 +142,18 @@ def get_microxs_and_flux(
     else:
         raise ValueError(f"Unsupported domain type: {type(domains[0])}")
 
-    rr_tally = openmc.Tally(name='MicroXS RR')
-    rr_tally.filters = [domain_filter, energy_filter]
-    rr_tally.nuclides = nuclides
-    rr_tally.multiply_density = False
-    rr_tally.scores = reactions
-
     flux_tally = openmc.Tally(name='MicroXS flux')
     flux_tally.filters = [domain_filter, energy_filter]
     flux_tally.scores = ['flux']
-    model.tallies = openmc.Tallies([rr_tally, flux_tally])
+    model.tallies = [flux_tally]
+
+    if reaction_rate_mode == 'direct':
+        rr_tally = openmc.Tally(name='MicroXS RR')
+        rr_tally.filters = [domain_filter, energy_filter]
+        rr_tally.nuclides = nuclides
+        rr_tally.multiply_density = False
+        rr_tally.scores = reactions
+        model.tallies.append(rr_tally)
 
     if openmc.lib.is_initialized:
         openmc.lib.finalize()
@@ -134,33 +174,55 @@ def get_microxs_and_flux(
         statepoint_path = model.run(**run_kwargs)
 
         if comm.rank == 0:
+            # Move the statepoint file if it is being saved to a specific path
+            if path_statepoint is not None:
+                shutil.move(statepoint_path, path_statepoint)
+                statepoint_path = path_statepoint
+
             with StatePoint(statepoint_path) as sp:
-                rr_tally = sp.tallies[rr_tally.id]
-                rr_tally._read_results()
+                if reaction_rate_mode == 'direct':
+                    rr_tally = sp.tallies[rr_tally.id]
+                    rr_tally._read_results()
                 flux_tally = sp.tallies[flux_tally.id]
                 flux_tally._read_results()
 
-    rr_tally = comm.bcast(rr_tally)
+    # Get flux values and make energy groups last dimension
     flux_tally = comm.bcast(flux_tally)
-    # Get reaction rates and flux values
-    reaction_rates = rr_tally.get_reshaped_data()  # (domains, groups, nuclides, reactions)
     flux = flux_tally.get_reshaped_data()  # (domains, groups, 1, 1)
-
-    # Make energy groups last dimension
-    reaction_rates = np.moveaxis(reaction_rates, 1, -1)  # (domains, nuclides, reactions, groups)
     flux = np.moveaxis(flux, 1, -1)  # (domains, 1, 1, groups)
 
-    # Divide RR by flux to get microscopic cross sections
-    xs = np.empty_like(reaction_rates) # (domains, nuclides, reactions, groups)
-    d, _, _, g = np.nonzero(flux)
-    xs[d, ..., g] = reaction_rates[d, ..., g] / flux[d, :, :, g]
+    # Create list where each item corresponds to one domain
+    fluxes = list(flux.squeeze((1, 2)))
+
+    if reaction_rate_mode == 'direct':
+        # Get reaction rates
+        rr_tally = comm.bcast(rr_tally)
+        reaction_rates = rr_tally.get_reshaped_data()  # (domains, groups, nuclides, reactions)
+
+        # Make energy groups last dimension
+        reaction_rates = np.moveaxis(reaction_rates, 1, -1)  # (domains, nuclides, reactions, groups)
+
+        # Divide RR by flux to get microscopic cross sections. The indexing
+        # ensures that only non-zero flux values are used, and broadcasting is
+        # applied to align the shapes of reaction_rates and flux for division.
+        xs = np.empty_like(reaction_rates) # (domains, nuclides, reactions, groups)
+        d, _, _, g = np.nonzero(flux)
+        xs[d, ..., g] = reaction_rates[d, ..., g] / flux[d, :, :, g]
+
+        # Create lists where each item corresponds to one domain
+        micros = [MicroXS(xs_i, nuclides, reactions) for xs_i in xs]
+    else:
+        micros = [MicroXS.from_multigroup_flux(
+            energies=energies,
+            multigroup_flux=flux_i,
+            chain_file=chain_file,
+            nuclides=nuclides,
+            reactions=reactions
+        ) for flux_i in fluxes]
 
     # Reset tallies
     model.tallies = original_tallies
 
-    # Create lists where each item corresponds to one domain
-    fluxes = list(flux.squeeze((1, 2)))
-    micros = [MicroXS(xs_i, nuclides, reactions) for xs_i in xs]
     return fluxes, micros
 
 
@@ -230,7 +292,7 @@ def from_multigroup_flux(
         ----------
         energies : iterable of float or str
             Energy group boundaries in [eV] or the name of the group structure
-        multi_group_flux : iterable of float
+        multigroup_flux : iterable of float
             Energy-dependent multigroup flux values
         chain_file : PathLike or Chain, optional
             Path to the depletion chain XML file or an instance of

openmc/filter_expansion.py

@@ -3,7 +3,7 @@
 import lxml.etree as ET
 
 import openmc.checkvalue as cv
-from . import Filter
+from .filter import Filter
 
 
 class ExpansionFilter(Filter):

tests/regression_tests/microxs/test.py

@@ -13,55 +13,56 @@
 @pytest.fixture(scope="module")
 def model():
     fuel = openmc.Material(name="uo2")
-    fuel.add_element("U", 1, percent_type="ao", enrichment=4.25)
-    fuel.add_element("O", 2)
+    fuel.add_nuclide("U235", 1.0)
+    fuel.add_nuclide("O16", 2.0)
     fuel.set_density("g/cc", 10.4)
 
-    clad = openmc.Material(name="clad")
-    clad.add_element("Zr", 1)
-    clad.set_density("g/cc", 6)
-
-    water = openmc.Material(name="water")
-    water.add_element("O", 1)
-    water.add_element("H", 2)
-    water.set_density("g/cc", 1.0)
-    water.add_s_alpha_beta("c_H_in_H2O")
-
-    radii = [0.42, 0.45]
-    fuel.volume = np.pi * radii[0] ** 2
-
-    materials = openmc.Materials([fuel, clad, water])
-
-    pin_surfaces = [openmc.ZCylinder(r=r) for r in radii]
-    pin_univ = openmc.model.pin(pin_surfaces, materials)
-    bound_box = openmc.model.RectangularPrism(1.24, 1.24, boundary_type="reflective")
-    root_cell = openmc.Cell(fill=pin_univ, region=-bound_box)
-    geometry = openmc.Geometry([root_cell])
+    sphere = openmc.Sphere(r=10.0, boundary_type='vacuum')
+    cell = openmc.Cell(region=-sphere, fill=fuel)
+    geometry = openmc.Geometry([cell])
 
     settings = openmc.Settings()
     settings.particles = 1000
     settings.inactive = 5
     settings.batches = 10
 
-    return openmc.Model(geometry, materials, settings)
+    return openmc.Model(geometry, settings=settings)
 
 
-@pytest.mark.parametrize("domain_type", ["materials", "mesh"])
-def test_from_model(model, domain_type):
+@pytest.mark.parametrize(
+    "domain_type, rr_mode",
+    [
+        ("materials", "direct"),
+        ("materials", "flux"),
+        ("mesh", "direct"),
+        ("mesh", "flux"),
+    ]
+)
+def test_from_model(model, domain_type, rr_mode):
     if domain_type == 'materials':
-        domains = model.materials[:1]
+        domains = list(model.geometry.get_all_materials().values())
     elif domain_type == 'mesh':
         mesh = openmc.RegularMesh()
-        mesh.lower_left = (-0.62, -0.62)
-        mesh.upper_right = (0.62, 0.62)
-        mesh.dimension = (3, 3)
+        mesh.lower_left = (-10., -10.)
+        mesh.upper_right = (10., 10.)
+        mesh.dimension = (1, 1)
         domains = mesh
-    nuclides = ['U234', 'U235', 'U238', 'U236', 'O16', 'O17', 'I135', 'Xe135',
-                'Xe136', 'Cs135', 'Gd157', 'Gd156']
-    _, test_xs = get_microxs_and_flux(model, domains, nuclides, chain_file=CHAIN_FILE)
+    nuclides = ['U235', 'O16', 'Xe135']
+    kwargs = {
+        'reaction_rate_mode': rr_mode,
+        'chain_file': CHAIN_FILE,
+        'path_statepoint': 'neutron_transport.h5',
+    }
+    if rr_mode == 'flux':
+        kwargs['energies'] = 'CASMO-40'
+    _, test_xs = get_microxs_and_flux(model, domains, nuclides, **kwargs)
     if config['update']:
-        test_xs[0].to_csv(f'test_reference_{domain_type}.csv')
-
-    ref_xs = MicroXS.from_csv(f'test_reference_{domain_type}.csv')
+        test_xs[0].to_csv(f'test_reference_{domain_type}_{rr_mode}.csv')
 
+    # Make sure results match reference results
+    ref_xs = MicroXS.from_csv(f'test_reference_{domain_type}_{rr_mode}.csv')
     np.testing.assert_allclose(test_xs[0].data, ref_xs.data, rtol=1e-11)
+
+    # Make sure statepoint file was saved
+    assert Path('neutron_transport.h5').exists()
+    Path('neutron_transport.h5').unlink()

tests/regression_tests/microxs/test_reference_materials.csv

@@ -0,0 +1,7 @@
+nuclides,reactions,groups,xs
+U235,"(n,gamma)",1,0.14765501510184456
+U235,fission,1,1.2517200956290817
+O16,"(n,gamma)",1,0.00010872314985710938
+O16,fission,1,0.0
+Xe135,"(n,gamma)",1,0.014333667335215764
+Xe135,fission,1,0.0

tests/regression_tests/microxs/test_reference_materials_direct.csv

@@ -0,0 +1,7 @@
+nuclides,reactions,groups,xs
+U235,"(n,gamma)",1,0.15018326758942505
+U235,fission,1,1.2603151141390958
+O16,"(n,gamma)",1,0.00012159621938463765
+O16,fission,1,0.0
+Xe135,"(n,gamma)",1,0.015180177095633546
+Xe135,fission,1,0.0

tests/regression_tests/microxs/test_reference_materials_flux.csv

@@ -0,0 +1,7 @@
+nuclides,reactions,groups,xs
+U235,"(n,gamma)",1,0.14765501510184456
+U235,fission,1,1.2517200956290815
+O16,"(n,gamma)",1,0.00010872314985710936
+O16,fission,1,0.0
+Xe135,"(n,gamma)",1,0.014333667335215761
+Xe135,fission,1,0.0