diff --git a/armi/reactor/parameters/parameterCollections.py b/armi/reactor/parameters/parameterCollections.py index 82d7d1a39..c9621e541 100644 --- a/armi/reactor/parameters/parameterCollections.py +++ b/armi/reactor/parameters/parameterCollections.py @@ -14,7 +14,7 @@ import copy import pickle -from typing import Any, Optional, List, Set +from typing import Any, Optional, List, Set, Iterator, Callable import sys import numpy as np @@ -367,7 +367,7 @@ def __contains__(self, name): else: return name in self._hist - def __eq__(self, other): + def __eq__(self, other: "ParameterCollection"): if not isinstance(other, self.__class__): return False @@ -382,7 +382,8 @@ def __eq__(self, other): return True - def __iter__(self): + def __iter__(self) -> Iterator[str]: + """Iterate over names of assigned parameters define on this collection.""" return ( pd.name for pd in self.paramDefs @@ -501,6 +502,32 @@ def restoreBackup(self, paramsToApply): pd.assigned = SINCE_ANYTHING self.assigned = SINCE_ANYTHING + def where( + self, f: Callable[[parameterDefinitions.Parameter], bool] + ) -> Iterator[parameterDefinitions.Parameter]: + """Produce an iterator over parameters that meet some criteria. + + Parameters + ---------- + f : callable function f(parameter) -> bool + Function to check if a parameter should be fetched during the iteration. + + Returns + ------- + iterator of :class:`armi.reactor.parameters.Parameter` + Iterator, **not** list or tuple, that produces each parameter that + meets ``f(parameter) == True``. + + Examples + -------- + >>> block = r.core[0][0] + >>> pdef = block.p.paramDefs + >>> for param in pdef.where(lambda pd: pd.atLocation(ParamLocation.EDGES)): + ... print(param.name, block.p[param.name]) + + """ + return filter(f, self.paramDefs) + def collectPluginParameters(pm): """Apply parameters from plugins to their respective object classes.""" diff --git a/armi/reactor/parameters/parameterDefinitions.py b/armi/reactor/parameters/parameterDefinitions.py index 472c7367a..7ad5ccec0 100644 --- a/armi/reactor/parameters/parameterDefinitions.py +++ b/armi/reactor/parameters/parameterDefinitions.py @@ -421,6 +421,10 @@ def atLocation(self, loc): """True if parameter is defined at location.""" return self.location and self.location & loc + def hasCategory(self, category: str) -> bool: + """True if a parameter has a specific category.""" + return category in self.categories + class ParameterDefinitionCollection: """ diff --git a/armi/reactor/tests/test_parameters.py b/armi/reactor/tests/test_parameters.py index e2068e8ae..38f47141b 100644 --- a/armi/reactor/tests/test_parameters.py +++ b/armi/reactor/tests/test_parameters.py @@ -456,10 +456,18 @@ class MockPC(parameters.ParameterCollection): self.assertEqual(p2.categories, set(["awesome", "stuff", "bacon"])) self.assertEqual(p3.categories, set(["bacon"])) + for p in [p1, p2, p3]: + self._testCategoryConsistency(p) + self.assertEqual(set(pc.paramDefs.inCategory("awesome")), set([p1, p2])) self.assertEqual(set(pc.paramDefs.inCategory("stuff")), set([p1, p2])) self.assertEqual(set(pc.paramDefs.inCategory("bacon")), set([p2, p3])) + def _testCategoryConsistency(self, p: parameters.Parameter): + for category in p.categories: + self.assertTrue(p.hasCategory(category)) + self.assertFalse(p.hasCategory("this_shouldnot_exist")) + def test_parameterCollectionsHave__slots__(self): """Tests we prevent accidental creation of attributes.""" self.assertEqual( @@ -502,3 +510,80 @@ class MockPCChild(MockPC): pcc = MockPCChild() with self.assertRaises(AssertionError): pcc.whatever = 33 + + +class ParamCollectionWhere(unittest.TestCase): + """Tests for ParameterCollection.where.""" + + class ScopeParamCollection(parameters.ParameterCollection): + pDefs = parameters.ParameterDefinitionCollection() + with pDefs.createBuilder() as pb: + pb.defParam( + name="empty", + description="Bare", + location=None, + categories=None, + units="", + ) + pb.defParam( + name="keff", + description="keff", + location=parameters.ParamLocation.VOLUME_INTEGRATED, + categories=[parameters.Category.neutronics], + units="", + ) + pb.defParam( + name="cornerFlux", + description="corner flux", + location=parameters.ParamLocation.CORNERS, + categories=[ + parameters.Category.neutronics, + ], + units="", + ) + pb.defParam( + name="edgeTemperature", + description="edge temperature", + location=parameters.ParamLocation.EDGES, + categories=[parameters.Category.thermalHydraulics], + units="", + ) + + @classmethod + def setUpClass(cls) -> None: + """Define a couple useful parameters with categories, locations, etc.""" + cls.pc = cls.ScopeParamCollection() + + def test_onCategory(self): + """Test the use of Parameter.hasCategory on filtering.""" + names = {"keff", "cornerFlux"} + for p in self.pc.where( + lambda pd: pd.hasCategory(parameters.Category.neutronics) + ): + self.assertTrue(p.hasCategory(parameters.Category.neutronics), msg=p) + names.remove(p.name) + self.assertFalse(names, msg=f"{names=} should be empty!") + + def test_onLocation(self): + """Test the use of Parameter.atLocation in filtering.""" + names = {"edgeTemperature"} + for p in self.pc.where( + lambda pd: pd.atLocation(parameters.ParamLocation.EDGES) + ): + self.assertTrue(p.atLocation(parameters.ParamLocation.EDGES), msg=p) + names.remove(p.name) + self.assertFalse(names, msg=f"{names=} should be empty!") + + def test_complicated(self): + """Test a multi-condition filter.""" + names = {"cornerFlux"} + + def check(p: parameters.Parameter) -> bool: + return p.atLocation(parameters.ParamLocation.CORNERS) and p.hasCategory( + parameters.Category.neutronics + ) + + for p in self.pc.where(check): + self.assertTrue(check(p), msg=p) + names.remove(p.name) + self.assertFalse(names, msg=f"{names=} should be empty") diff --git a/doc/release/0.4.rst b/doc/release/0.4.rst index 24bc49730..45f7dba34 100644 --- a/doc/release/0.4.rst +++ b/doc/release/0.4.rst @@ -14,6 +14,10 @@ New Features #. Allow merging a component with zero area into another component (`PR#1858 `_) #. Provide utilities for determining location of a rotated object in a hexagonal lattice (``getIndexOfRotatedCell``). (`PR#1846 `_) +#. Provide ``Parameter.hasCategory`` for quickly checking if a parameter is defined with a given category. + (`PR#1899 `_) +#. Provide ``ParameterCollection.where`` for efficient iteration over parameters who's definition. + matches a given condition. (`PR#1899 `_) #. Plugins can provide the ``getAxialExpansionChanger`` hook to customize axial expansion. (`PR#1870