Providing plugin hook getAxialExpansionChanger (#1870)

armi/plugins.py

@@ -127,6 +127,7 @@
 
 if TYPE_CHECKING:
     from armi.reactor.composites import Composite
+    from armi.reactor.converters.axialExpansionChanger import AxialExpansionChanger
 
 
 HOOKSPEC = pluggy.HookspecMarker("armi")
@@ -662,6 +663,45 @@ def defineSystemBuilders() -> Dict[str, Callable[[str], "Composite"]]:
         and a ``"sfp"`` lookup, triggered to run after all other hooks have been run.
         """
 
+    @staticmethod
+    @HOOKSPEC(firstresult=True)
+    def getAxialExpansionChanger() -> type["AxialExpansionChanger"]:
+        """Produce the class responsible for performing axial expansion.
+
+        Plugins can provide this hook to override or negate axial expansion.
+        Will be used during initial construction of the core and assemblies, and
+        can be a class to perform custom axial expansion routines.
+
+        The first object returned that is not ``None`` will be used.
+        Plugins are encouraged to add the ``tryfirst=True`` arguments to their
+        ``HOOKIMPL`` invocations to make sure their specific are earlier in the
+        hook call sequence.
+
+        Returns
+        -------
+        type of :class:`armi.reactor.converters.axialExpansionChanger.AxialExpansionChanger`
+
+        Notes
+        -----
+        This hook **should not** provide an instance of the class. The construction
+        of the changer will be handled by applications and plugins that need it.
+
+        This hook should only be provided by one additional plugin in your application. Otherwise
+        the `order of hook execution <https://pluggy.readthedocs.io/en/stable/index.html#call-time-order>`_
+        may not provide the behavior you expect.
+
+        Examples
+        --------
+        >>> class MyPlugin(ArmiPlugin):
+        ...     @staticmethod
+        ...     @HOOKIMPL(tryfirst=True)
+        ...     def getAxialExpansionChanger():
+        ...         from myproject.physics import BespokeAxialExpansion
+        ...
+        ...         return BespokeAxialExpansion
+
+        """
+
 
 class UserPlugin(ArmiPlugin):
     """

armi/reactor/__init__.py

@@ -90,13 +90,20 @@ def defineAssemblyTypes():
 
     @staticmethod
     @plugins.HOOKIMPL(trylast=True)
-    def defineSystemBuilders() -> Dict[
-        str, Callable[[str], Union["Core", "SpentFuelPool"]]
-    ]:
+    def defineSystemBuilders() -> (
+        Dict[str, Callable[[str], Union["Core", "SpentFuelPool"]]]
+    ):
         from armi.reactor.reactors import Core
         from armi.reactor.assemblyLists import SpentFuelPool
 
         return {
             "core": Core,
             "sfp": SpentFuelPool,
         }
+
+    @staticmethod
+    @plugins.HOOKIMPL(trylast=True)
+    def getAxialExpansionChanger():
+        from armi.reactor.converters.axialExpansionChanger import AxialExpansionChanger
+
+        return AxialExpansionChanger

armi/reactor/blueprints/__init__.py

@@ -151,7 +151,7 @@ def __new__(mcs, name, bases, attrs):
         else:
             pluginSections = pm.hook.defineBlueprintsSections()
             for plug in pluginSections:
-                for (attrName, section, resolver) in plug:
+                for attrName, section, resolver in plug:
                     assert isinstance(section, yamlize.Attribute)
                     if attrName in attrs:
                         raise plugins.PluginError(
@@ -326,10 +326,12 @@ def _prepConstruction(self, cs):
                     for a in list(self.assemblies.values())
                     if not any(a.hasFlags(f) for f in assemsToSkip)
                 )
-                axialExpansionChanger.axialExpansionChanger.expandColdDimsToHot(
-                    assemsToExpand,
-                    cs[CONF_DETAILED_AXIAL_EXPANSION],
-                )
+                axialExpander = getPluginManagerOrFail().hook.getAxialExpansionChanger()
+                if axialExpander is not None:
+                    axialExpander.expandColdDimsToHot(
+                        assemsToExpand,
+                        cs[CONF_DETAILED_AXIAL_EXPANSION],
+                    )
 
             getPluginManagerOrFail().hook.afterConstructionOfAssemblies(
                 assemblies=self.assemblies.values(), cs=cs

armi/reactor/converters/axialExpansionChanger/axialExpansionChanger.py

@@ -52,71 +52,6 @@ def makeAssemsAbleToSnapToUniformMesh(
         a.makeAxialSnapList(referenceAssembly)
 
 
-def expandColdDimsToHot(
-    assems: list,
-    isDetailedAxialExpansion: bool,
-    referenceAssembly=None,
-):
-    """
-    Expand BOL assemblies, resolve disjoint axial mesh (if needed), and update block BOL heights.
-
-    .. impl:: Perform expansion during core construction based on block heights at a specified temperature.
-        :id: I_ARMI_INP_COLD_HEIGHT
-        :implements: R_ARMI_INP_COLD_HEIGHT
-
-        This method is designed to be used during core construction to axially thermally expand the
-        assemblies to their "hot" temperatures (as determined by ``Thot`` values in blueprints).
-        First, The Assembly is prepared for axial expansion via ``setAssembly``. In
-        ``applyColdHeightMassIncrease``, the number densities on each Component is adjusted to
-        reflect that Assembly inputs are at cold (i.e., ``Tinput``) temperatures. To expand to
-        the requested hot temperatures, thermal expansion factors are then computed in
-        ``computeThermalExpansionFactors``. Finally, the Assembly is axially thermally expanded in
-        ``axiallyExpandAssembly``.
-
-        If the setting ``detailedAxialExpansion`` is ``False``, then each Assembly gets its Block mesh
-        set to match that of the "reference" Assembly (see ``getDefaultReferenceAssem`` and ``setBlockMesh``).
-
-        Once the Assemblies are axially expanded, the Block BOL heights are updated. To account for the change in
-        Block volume from axial expansion, ``completeInitialLoading`` is called to update any volume-dependent
-        Block information.
-
-    Parameters
-    ----------
-    assems: list[:py:class:`Assembly <armi.reactor.assemblies.Assembly>`]
-        list of assemblies to be thermally expanded
-    isDetailedAxialExpansion: bool
-        If False, assemblies will be forced to conform to the reference mesh after expansion
-    referenceAssembly: :py:class:`Assembly <armi.reactor.assemblies.Assembly>`, optional
-        Assembly whose mesh other meshes will conform to if isDetailedAxialExpansion is False.
-        If not provided, will assume the finest mesh assembly which is typically fuel.
-
-    Notes
-    -----
-    Calling this method will result in an increase in mass via applyColdHeightMassIncrease!
-
-    See Also
-    --------
-    :py:meth:`armi.reactor.converters.axialExpansionChanger.axialExpansionChanger.AxialExpansionChanger.applyColdHeightMassIncrease`
-    """
-    assems = list(assems)
-    if not referenceAssembly:
-        referenceAssembly = getDefaultReferenceAssem(assems)
-    axialExpChanger = AxialExpansionChanger(isDetailedAxialExpansion)
-    for a in assems:
-        axialExpChanger.setAssembly(a, expandFromTinputToThot=True)
-        axialExpChanger.applyColdHeightMassIncrease()
-        axialExpChanger.expansionData.computeThermalExpansionFactors()
-        axialExpChanger.axiallyExpandAssembly()
-    if not isDetailedAxialExpansion:
-        for a in assems:
-            a.setBlockMesh(referenceAssembly.getAxialMesh())
-    # update block BOL heights to reflect hot heights
-    for a in assems:
-        for b in a:
-            b.p.heightBOL = b.getHeight()
-            b.completeInitialLoading()
-
-
 class AxialExpansionChanger:
     """
     Axially expand or contract assemblies or an entire core.
@@ -148,6 +83,71 @@ def __init__(self, detailedAxialExpansion: bool = False):
         self.linked = None
         self.expansionData = None
 
+    @classmethod
+    def expandColdDimsToHot(
+        cls,
+        assems: list,
+        isDetailedAxialExpansion: bool,
+        referenceAssembly=None,
+    ):
+        """Expand BOL assemblies, resolve disjoint axial mesh (if needed), and update block BOL heights.
+
+        .. impl:: Perform expansion during core construction based on block heights at a specified temperature.
+            :id: I_ARMI_INP_COLD_HEIGHT
+            :implements: R_ARMI_INP_COLD_HEIGHT
+
+            This method is designed to be used during core construction to axially thermally expand the
+            assemblies to their "hot" temperatures (as determined by ``Thot`` values in blueprints).
+            First, The Assembly is prepared for axial expansion via ``setAssembly``. In
+            ``applyColdHeightMassIncrease``, the number densities on each Component is adjusted to
+            reflect that Assembly inputs are at cold (i.e., ``Tinput``) temperatures. To expand to
+            the requested hot temperatures, thermal expansion factors are then computed in
+            ``computeThermalExpansionFactors``. Finally, the Assembly is axially thermally expanded in
+            ``axiallyExpandAssembly``.
+
+            If the setting ``detailedAxialExpansion`` is ``False``, then each Assembly gets its Block mesh
+            set to match that of the "reference" Assembly (see ``getDefaultReferenceAssem`` and ``setBlockMesh``).
+
+            Once the Assemblies are axially expanded, the Block BOL heights are updated. To account for the change in
+            Block volume from axial expansion, ``completeInitialLoading`` is called to update any volume-dependent
+            Block information.
+
+        Parameters
+        ----------
+        assems: list[:py:class:`Assembly <armi.reactor.assemblies.Assembly>`]
+            list of assemblies to be thermally expanded
+        isDetailedAxialExpansion: bool
+            If False, assemblies will be forced to conform to the reference mesh after expansion
+        referenceAssembly: :py:class:`Assembly <armi.reactor.assemblies.Assembly>`, optional
+            Assembly whose mesh other meshes will conform to if isDetailedAxialExpansion is False.
+            If not provided, will assume the finest mesh assembly which is typically fuel.
+
+        Notes
+        -----
+        Calling this method will result in an increase in mass via applyColdHeightMassIncrease!
+
+        See Also
+        --------
+        :py:meth:`armi.reactor.converters.axialExpansionChanger.axialExpansionChanger.AxialExpansionChanger.applyColdHeightMassIncrease`
+        """
+        assems = list(assems)
+        if not referenceAssembly:
+            referenceAssembly = getDefaultReferenceAssem(assems)
+        axialExpChanger = cls(isDetailedAxialExpansion)
+        for a in assems:
+            axialExpChanger.setAssembly(a, expandFromTinputToThot=True)
+            axialExpChanger.applyColdHeightMassIncrease()
+            axialExpChanger.expansionData.computeThermalExpansionFactors()
+            axialExpChanger.axiallyExpandAssembly()
+        if not isDetailedAxialExpansion:
+            for a in assems:
+                a.setBlockMesh(referenceAssembly.getAxialMesh())
+        # update block BOL heights to reflect hot heights
+        for a in assems:
+            for b in a:
+                b.p.heightBOL = b.getHeight()
+                b.completeInitialLoading()
+
     def performPrescribedAxialExpansion(
         self, a, componentLst: list, percents: list, setFuel=True
     ):

armi/tests/test_plugins.py

@@ -26,6 +26,7 @@
 from armi import plugins
 from armi import settings
 from armi import utils
+from armi.reactor.converters.axialExpansionChanger import AxialExpansionChanger
 from armi.physics.neutronics import NeutronicsPlugin
 from armi.reactor.blocks import Block
 from armi.reactor.flags import Flags
@@ -42,13 +43,27 @@ def defineFlags():
         return {"SUPER_FLAG": utils.flags.auto()}
 
 
+class SillyAxialExpansionChanger(AxialExpansionChanger):
+    """Fake, test-specific axial expansion changer that a plugin will register."""
+
+
+class SillyAxialPlugin(plugins.ArmiPlugin):
+    """Trivial plugin that implements the axial expansion hook."""
+
+    @staticmethod
+    @plugins.HOOKIMPL
+    def getAxialExpansionChanger() -> type[SillyAxialExpansionChanger]:
+        return SillyAxialExpansionChanger
+
+
 class TestPluginRegistration(unittest.TestCase):
     def setUp(self):
         """
         Manipulate the standard App. We can't just configure our own, since the
         pytest environment bleeds between tests.
         """
-        self._backupApp = deepcopy(getApp())
+        self.app = getApp()
+        self._backupApp = deepcopy(self.app)
 
     def tearDown(self):
         """Restore the App to its original state."""
@@ -92,6 +107,18 @@ def test_defineFlags(self):
         # show the flag exists now
         self.assertEqual(type(Flags.SUPER_FLAG._value), int)
 
+    def test_axialExpansionHook(self):
+        """Test that plugins can override the axial expansion of assemblies via a hook."""
+        pm = self.app.pluginManager
+        first = pm.hook.getAxialExpansionChanger()
+        # By default, make sure we get the armi-shipped expansion class
+        self.assertIs(first, AxialExpansionChanger)
+        pm.register(SillyAxialPlugin)
+        second = pm.hook.getAxialExpansionChanger()
+        # Registering a plugin that implements the hook means we get
+        # that plugin's axial expander
+        self.assertIs(second, SillyAxialExpansionChanger)
+
 
 class TestPluginBasics(unittest.TestCase):
     def test_defineParameters(self):

doc/release/0.4.rst

@@ -11,8 +11,11 @@ New Features
 #. ARMI now supports Python 3.12. (`PR#1813 <https://github.com/terrapower/armi/pull/1813>`_)
 #. Removing the ``tabulate`` dependency by ingesting it to ``armi.utils.tabulate``. (`PR#1811 <https://github.com/terrapower/armi/pull/1811>`_)
 #. Adding ``--skip-inspection`` flag to ``CompareCases`` CLI. (`PR#1842 <https://github.com/terrapower/armi/pull/1842>`_)
+#. Allow merging a component with zero area into another component (`PR#1858 <https://github.com/terrapower/armi/pull/1858>`_)
 #. Provide utilities for determining location of a rotated object in a hexagonal lattice (``getIndexOfRotatedCell``). (`PR#1846 <https://github.com/terrapower/armi/1846`)
 #. Allow merging a component with zero area into another component. (`PR#1858 <https://github.com/terrapower/armi/pull/1858>`_)
+#. Plugins can provide the ``getAxialExpansionChanger`` hook to customize axial expansion.
+   (`PR#1870 <https://github.com/terrapower/armi/pull/1870`_)
 #. TBD
 
 API Changes