From 73cb4ed3657419a587ac0e587785c2890e6d1551 Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:54 -0800 Subject: [PATCH 01/27] Provide Composite.iterChildrenWithFlags Produces an iterable over all children that match the flag spec. Composite.getChildrenWithFlags throws these into a list so the user API is the same. Test added showing it does what we expect it to do. Replaced one usage of iteration over the list flavor with the iterable variant. --- armi/reactor/blocks.py | 2 +- armi/reactor/composites.py | 15 ++++++++++----- armi/reactor/tests/test_composites.py | 8 ++++++++ 3 files changed, 19 insertions(+), 6 deletions(-) diff --git a/armi/reactor/blocks.py b/armi/reactor/blocks.py index 023e86361..4c3edd8ab 100644 --- a/armi/reactor/blocks.py +++ b/armi/reactor/blocks.py @@ -1532,7 +1532,7 @@ def getPinLocations(self) -> list[grids.IndexLocation]: :meth:`getPinCoordinates` - companion for this method. """ items = [] - for clad in self.getChildrenWithFlags(Flags.CLAD): + for clad in self.iterChildrenWithFlags(Flags.CLAD): if isinstance(clad.spatialLocator, grids.MultiIndexLocation): items.extend(clad.spatialLocator) else: diff --git a/armi/reactor/composites.py b/armi/reactor/composites.py index 62a0e329d..196cb90ed 100644 --- a/armi/reactor/composites.py +++ b/armi/reactor/composites.py @@ -525,6 +525,10 @@ def getChildren(self, deep=False, generationNum=1, includeMaterials=False): """Return the children of this object.""" raise NotImplementedError + def iterChildrenWithFlags(self, typeSpec: TypeSpec, exactMatch=True): + """Produce an iterator of children that have given flags.""" + raise NotImplementedError() + def getChildrenWithFlags(self, typeSpec: TypeSpec, exactMatch=True): """Get all children that have given flags.""" raise NotImplementedError @@ -2701,13 +2705,14 @@ def getChildren( return children + def iterChildrenWithFlags(self, typeSpec: TypeSpec, exactMatch=False): + """Produce an iterator over all children of a specific type.""" + checker = operator.methodcaller("hasFlags", typeSpec, exactMatch) + return filter(checker, self) + def getChildrenWithFlags(self, typeSpec: TypeSpec, exactMatch=False): """Get all children of a specific type.""" - children = [] - for child in self: - if child.hasFlags(typeSpec, exact=exactMatch): - children.append(child) - return children + return list(self.iterChildrenWithFlags(typeSpec, exactMatch)) def getChildrenOfType(self, typeName): """Get children that have a specific input type name.""" diff --git a/armi/reactor/tests/test_composites.py b/armi/reactor/tests/test_composites.py index 5a60cc5a1..4232034bc 100644 --- a/armi/reactor/tests/test_composites.py +++ b/armi/reactor/tests/test_composites.py @@ -417,6 +417,14 @@ def test_syncParameters(self): numSynced = self.container._syncParameters(data, {}) self.assertEqual(numSynced, 2) + def test_iterChildrenWithFlags(self): + expectedChildren = {c for c in self.container if c.hasFlags(Flags.DUCT)} + found = set() + for c in self.container.iterChildrenWithFlags(Flags.DUCT): + self.assertIn(c, expectedChildren) + found.add(c) + self.assertSetEqual(found, expectedChildren) + class TestCompositeTree(unittest.TestCase): blueprintYaml = """ From a9d3c7598b609d8a51a6562a81c759c1fff8f84b Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:54 -0800 Subject: [PATCH 02/27] Provide Composite.iterChildrenOfType --- armi/reactor/composites.py | 12 ++++++------ armi/reactor/tests/test_composites.py | 8 ++++++++ 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/armi/reactor/composites.py b/armi/reactor/composites.py index 196cb90ed..a5c8b412b 100644 --- a/armi/reactor/composites.py +++ b/armi/reactor/composites.py @@ -2714,13 +2714,13 @@ def getChildrenWithFlags(self, typeSpec: TypeSpec, exactMatch=False): """Get all children of a specific type.""" return list(self.iterChildrenWithFlags(typeSpec, exactMatch)) - def getChildrenOfType(self, typeName): + def iterChildrenOfType(self, typeName: str): + """Produce an iterator over all children with a specific input type name""" + return filter(lambda c: c.getType() == typeName, self) + + def getChildrenOfType(self, typeName: str): """Get children that have a specific input type name.""" - children = [] - for child in self: - if child.getType() == typeName: - children.append(child) - return children + return list(self.iterChildrenOfType(typeName)) def getComponents(self, typeSpec: TypeSpec = None, exact=False): return list(self.iterComponents(typeSpec, exact)) diff --git a/armi/reactor/tests/test_composites.py b/armi/reactor/tests/test_composites.py index 4232034bc..1a1fe5c65 100644 --- a/armi/reactor/tests/test_composites.py +++ b/armi/reactor/tests/test_composites.py @@ -106,6 +106,7 @@ def setUp(self): container.add(leaf) nested = DummyComposite("clad", 98) nested.setType("clad") + self.cladChild = nested self.secondGen = DummyComposite("liner", 97) self.thirdGen = DummyLeaf("pin 77", 33) self.secondGen.add(self.thirdGen) @@ -425,6 +426,13 @@ def test_iterChildrenWithFlags(self): found.add(c) self.assertSetEqual(found, expectedChildren) + def test_iterChildrenOfType(self): + clads = self.container.iterChildrenOfType("clad") + first = next(clads) + self.assertIs(first, self.cladChild) + with self.assertRaises(StopIteration): + next(clads) + class TestCompositeTree(unittest.TestCase): blueprintYaml = """ From 68adeb1d86c335c004b423da5f3554380d41f1bc Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:54 -0800 Subject: [PATCH 03/27] Use generator inside Composite.getMass not list --- armi/reactor/composites.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/armi/reactor/composites.py b/armi/reactor/composites.py index a5c8b412b..072e3ca33 100644 --- a/armi/reactor/composites.py +++ b/armi/reactor/composites.py @@ -898,7 +898,7 @@ def getMaxArea(self): """ raise NotImplementedError() - def getMass(self, nuclideNames=None): + def getMass(self, nuclideNames=None) -> float: """ Determine the mass in grams of nuclide(s) and/or elements in this object. @@ -921,7 +921,7 @@ def getMass(self, nuclideNames=None): mass : float The mass in grams. """ - return sum([c.getMass(nuclideNames=nuclideNames) for c in self]) + return sum(c.getMass(nuclideNames=nuclideNames) for c in self) def getMassFrac(self, nucName): """ From 63f852a868a215430c32ce88a029e4bdfcb7ce93 Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:54 -0800 Subject: [PATCH 04/27] Provide and use Assembly.iterBlocks in select places --- armi/reactor/assemblies.py | 38 +++++++++++++++++++++++++++++--------- 1 file changed, 29 insertions(+), 9 deletions(-) diff --git a/armi/reactor/assemblies.py b/armi/reactor/assemblies.py index da55c7694..465a91ee6 100644 --- a/armi/reactor/assemblies.py +++ b/armi/reactor/assemblies.py @@ -326,10 +326,8 @@ def getPinPlenumVolumeInCubicMeters(self): ------- This is a bit design-specific for pinned assemblies """ - plenumBlocks = self.getBlocks(Flags.PLENUM) - plenumVolume = 0.0 - for b in plenumBlocks: + for b in self.iterBlocks(Flags.PLENUM): cladId = b.getComponent(Flags.CLAD).getDimension("id") length = b.getHeight() plenumVolume += ( @@ -339,7 +337,7 @@ def getPinPlenumVolumeInCubicMeters(self): def getAveragePlenumTemperature(self): """Return the average of the plenum block outlet temperatures.""" - plenumBlocks = self.getBlocks(Flags.PLENUM) + plenumBlocks = self.iterBlocks(Flags.PLENUM) plenumTemps = [b.p.THcoolantOutletT for b in plenumBlocks] # no plenum blocks, use the top block of the assembly for plenum temperature @@ -814,6 +812,31 @@ def dump(self, fName=None): with open(fName, "w") as pkl: pickle.dump(self, pkl) + def iterBlocks(self, typeSpec=None, exact=False): + """Produce an iterator over all blocks in this assembly from bottom to top. + + Parameters + ---------- + typeSpec : Flags or list of Flags, optional + Restrict returned blocks to have these flags. + exact : bool, optional + If true, only produce blocks that have those exact flags. + + Returns + ------- + iterable of Block + + See also + -------- + * :meth:`__iter__` - if no type spec provided, assemblies can be + naturally iterated upon. + * :meth:`iterChildrenWithFlags` - alternative if you know you have + a type spec that isn't ``None``. + """ + if typeSpec is None: + return iter(self) + return self.iterChildrenWithFlags(typeSpec, exact) + def getBlocks(self, typeSpec=None, exact=False): """ Get blocks in an assembly from bottom to top. @@ -830,10 +853,7 @@ def getBlocks(self, typeSpec=None, exact=False): blocks : list List of blocks. """ - if typeSpec is None: - return self.getChildren() - else: - return self.getChildrenWithFlags(typeSpec, exactMatch=exact) + return list(self.iterBlocks(typeSpec, exact)) def getBlocksAndZ(self, typeSpec=None, returnBottomZ=False, returnTopZ=False): """ @@ -1191,7 +1211,7 @@ def countBlocksWithFlags(self, blockTypeSpec=None): blockCounter : int number of blocks of this type """ - return len(self.getBlocks(blockTypeSpec)) + return sum(1 for _ in self.iterBlocks(blockTypeSpec)) def getDim(self, typeSpec, dimName): """ From 78d260c17ecc31004077af233326cec402cdc42c Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:54 -0800 Subject: [PATCH 05/27] Naturally iterate over assemlby in hasContinuousCoolantChannel --- armi/reactor/assemblies.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/armi/reactor/assemblies.py b/armi/reactor/assemblies.py index 465a91ee6..d9cc59041 100644 --- a/armi/reactor/assemblies.py +++ b/armi/reactor/assemblies.py @@ -901,9 +901,7 @@ def getBlocksAndZ(self, typeSpec=None, returnBottomZ=False, returnTopZ=False): return zip(blocks, zCoords) def hasContinuousCoolantChannel(self): - return all( - b.containsAtLeastOneChildWithFlags(Flags.COOLANT) for b in self.getBlocks() - ) + return all(b.containsAtLeastOneChildWithFlags(Flags.COOLANT) for b in self) def getFirstBlock(self, typeSpec=None, exact=False): bs = self.getBlocks(typeSpec, exact=exact) From 00cf8b320a2fdf786b2a57769c22a4b155f0ba90 Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:55 -0800 Subject: [PATCH 06/27] Use Assembly.iterBlocks in getFirstBlock Rather than generate a potentially large list just to use the first element, try to advance the iterator. If it raises a StopIteration error, we don't have any blocks that match the spec. --- armi/reactor/assemblies.py | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/armi/reactor/assemblies.py b/armi/reactor/assemblies.py index d9cc59041..badf33eb4 100644 --- a/armi/reactor/assemblies.py +++ b/armi/reactor/assemblies.py @@ -904,10 +904,25 @@ def hasContinuousCoolantChannel(self): return all(b.containsAtLeastOneChildWithFlags(Flags.COOLANT) for b in self) def getFirstBlock(self, typeSpec=None, exact=False): - bs = self.getBlocks(typeSpec, exact=exact) - if bs: - return bs[0] - else: + """Find the first block that matches the spec. + + Parameters + ---------- + typeSpec : flag or list of flags, optional + Specification to require on the returned block. + exact : bool, optional + Require block to exactly match ``typeSpec`` + + Returns + ------- + Block or None + First block that matches if such a block could be found. + """ + try: + # Create an iterator and attempt to advance it to the first value. + return next(self.iterBlocks(typeSpec, exact)) + except StopIteration: + # No items found in the iteration -> no blocks match the request return None def getFirstBlockByType(self, typeName): From 5878363f23364cad3b8609d87422e0fd93ad10e8 Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:55 -0800 Subject: [PATCH 07/27] Use iterator in Assembly.getFirstBlockByType I cannot find any uses of this method outside of the one method where it is tested --- armi/reactor/assemblies.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/armi/reactor/assemblies.py b/armi/reactor/assemblies.py index badf33eb4..af34b9bac 100644 --- a/armi/reactor/assemblies.py +++ b/armi/reactor/assemblies.py @@ -926,14 +926,11 @@ def getFirstBlock(self, typeSpec=None, exact=False): return None def getFirstBlockByType(self, typeName): - bs = [ - b - for b in self.getChildren(deep=False) - if isinstance(b, blocks.Block) and b.getType() == typeName - ] - if bs: - return bs[0] - return None + blocks = filter(lambda b: b.getType() == typeName, self) + try: + return next(blocks) + except StopIteration: + return None def getBlockAtElevation(self, elevation): """ From a28d3795357c05e262aa56c4e7145f62ed6ccc4b Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:55 -0800 Subject: [PATCH 08/27] Fix a bug in fuel handler test Calling SFP.getChildren with a flags argument should really be SFP.getChildrenWithFlags --- armi/physics/fuelCycle/tests/test_fuelHandlers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/armi/physics/fuelCycle/tests/test_fuelHandlers.py b/armi/physics/fuelCycle/tests/test_fuelHandlers.py index b6bbc14a0..3967e3bf5 100644 --- a/armi/physics/fuelCycle/tests/test_fuelHandlers.py +++ b/armi/physics/fuelCycle/tests/test_fuelHandlers.py @@ -811,7 +811,7 @@ def test_dischargeSwap(self): # grab an arbitrary fuel assembly from the core and from the SFP a1 = self.r.core.getAssemblies(Flags.FUEL)[0] - a2 = self.r.excore["sfp"].getChildren(Flags.FUEL)[0] + a2 = self.r.excore["sfp"].getChildrenWithFlags(Flags.FUEL)[0] # grab the stationary blocks pre swap a1PreSwapStationaryBlocks = [ From 5dbd04700553806c3b6c16ea744e69ea2004776c Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:55 -0800 Subject: [PATCH 09/27] Make test_getChildren more declarative and precise Rather than checking just the length of things, compare we get the objects we expect --- armi/reactor/tests/test_composites.py | 26 +++++++++++++++++++++----- 1 file changed, 21 insertions(+), 5 deletions(-) diff --git a/armi/reactor/tests/test_composites.py b/armi/reactor/tests/test_composites.py index 1a1fe5c65..a20750b5d 100644 --- a/armi/reactor/tests/test_composites.py +++ b/armi/reactor/tests/test_composites.py @@ -15,6 +15,7 @@ """Tests for the composite pattern.""" from copy import deepcopy import logging +import itertools import unittest from armi import nuclearDataIO @@ -113,6 +114,14 @@ def setUp(self): nested.add(self.secondGen) container.add(nested) self.container = container + # Composite tree structure in list of lists for testing + # tree[i] contains the children at "generation" or "depth" i + self.tree: list[list[composites.Composite]] = [ + [self.container], + list(self.container), + [self.secondGen], + [self.thirdGen], + ] def test_composite(self): """Test basic Composite things. @@ -140,22 +149,29 @@ def test_getChildren(self): :id: T_ARMI_CMP1 :tests: R_ARMI_CMP """ - # There are 5 leaves and 1 composite in container. The composite has one leaf. firstGen = self.container.getChildren() - self.assertEqual(len(firstGen), 6) + self.assertEqual(firstGen, self.tree[1]) + secondGen = self.container.getChildren(generationNum=2) - self.assertEqual(len(secondGen), 1) + self.assertEqual(secondGen, self.tree[2]) + self.assertIs(secondGen[0], self.secondGen) third = self.container.getChildren(generationNum=3) - self.assertEqual(len(third), 1) + self.assertEqual(third, self.tree[3]) self.assertIs(third[0], self.thirdGen) + allC = self.container.getChildren(deep=True) - self.assertEqual(len(allC), 8) + expected = self.tree[1] + self.tree[2] + self.tree[3] + self.assertTrue( + all(a is e for a, e in itertools.zip_longest(allC, expected)), + msg=f"Deep traversal differs: {allC=} != {expected=}", + ) onlyLiner = self.container.getChildren( deep=True, predicate=lambda o: o.p.type == "liner" ) self.assertEqual(len(onlyLiner), 1) + self.assertIs(onlyLiner[0], self.secondGen) def test_getName(self): """Test the getName method. From d783368e8dc7e805b892e6cab4c24db8e35ca7fd Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:55 -0800 Subject: [PATCH 10/27] Add a test for Composite.getChildren(includeMaterials=True) --- armi/reactor/tests/test_composites.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/armi/reactor/tests/test_composites.py b/armi/reactor/tests/test_composites.py index a20750b5d..b6a764019 100644 --- a/armi/reactor/tests/test_composites.py +++ b/armi/reactor/tests/test_composites.py @@ -77,6 +77,8 @@ def __init__(self, name, i=0): composites.Composite.__init__(self, name) self.p.type = name self.spatialLocator = grids.IndexLocation(i, i, i, _testGrid) + # Some special material attribute for testing getChildren(includeMaterials=True) + self.material = ("hello", "world") def getChildren( self, deep=False, generationNum=1, includeMaterials=False, predicate=None @@ -173,6 +175,22 @@ def test_getChildren(self): self.assertEqual(len(onlyLiner), 1) self.assertIs(onlyLiner[0], self.secondGen) + def test_getChildrenWithMaterials(self): + """Test the ability for getChildren to place the material after the object.""" + withMaterials = self.container.getChildren(deep=True, includeMaterials=True) + # Grab the iterable so we can control the progression + items = iter(withMaterials) + for item in items: + expectedMat = getattr(item, "material", None) + if expectedMat is None: + continue + # Material should be the next item in the list + actualMat = next(items) + self.assertIs(actualMat, expectedMat) + break + else: + raise RuntimeError("No materials found with includeMaterials=True") + def test_getName(self): """Test the getName method. From f764750965e16fb47aa735f274a1e6e0a2e20b17 Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:55 -0800 Subject: [PATCH 11/27] Add tests for Composite.removeAll Came up in some debugging that this isn't totally tested. Turns out a couple places override this and that was causing my problem. But a test is still good to add! --- armi/reactor/tests/test_composites.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/armi/reactor/tests/test_composites.py b/armi/reactor/tests/test_composites.py index b6a764019..aafa5cd59 100644 --- a/armi/reactor/tests/test_composites.py +++ b/armi/reactor/tests/test_composites.py @@ -467,6 +467,17 @@ def test_iterChildrenOfType(self): with self.assertRaises(StopIteration): next(clads) + def test_removeAll(self): + """Test the ability to remove all children of a composite.""" + self.container.removeAll() + self.assertEqual(len(self.container), 0) + # Nothing to iterate over + items = iter(self.container) + with self.assertRaises(StopIteration): + next(items) + for child in self.tree[1]: + self.assertIsNone(child.parent) + class TestCompositeTree(unittest.TestCase): blueprintYaml = """ From fef2e5d7a22c7d1d9bae2b2bd94021394ac4a04a Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:56 -0800 Subject: [PATCH 12/27] Add test for Composite.setChildren This should still pass on main. But I added it during some debugging because we didn't have a test just on this method. Tests are still good though --- armi/reactor/tests/test_composites.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/armi/reactor/tests/test_composites.py b/armi/reactor/tests/test_composites.py index aafa5cd59..9ef2c7c3f 100644 --- a/armi/reactor/tests/test_composites.py +++ b/armi/reactor/tests/test_composites.py @@ -478,6 +478,17 @@ def test_removeAll(self): for child in self.tree[1]: self.assertIsNone(child.parent) + def test_setChildren(self): + """Test the ability to override children on a composite.""" + newChildren = self.tree[2] + self.tree[3] + oldChildren = list(self.container) + self.container.setChildren(newChildren) + self.assertEqual(len(self.container), len(newChildren)) + for old in oldChildren: + self.assertIsNone(old.parent) + for actualNew, expectedNew in zip(newChildren, self.container): + self.assertIs(actualNew, expectedNew) + class TestCompositeTree(unittest.TestCase): blueprintYaml = """ From 9e7ef533cd7134101a21c29742fbf9ece5496e5d Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:56 -0800 Subject: [PATCH 13/27] Provide Composite.iterChildren Produces an iterator of the children rather than a list of all children. Does not support the `includeMaterials` option --- armi/reactor/composites.py | 70 ++++++++++++++++++++++++++- armi/reactor/tests/test_composites.py | 42 ++++++++++++++++ 2 files changed, 111 insertions(+), 1 deletion(-) diff --git a/armi/reactor/composites.py b/armi/reactor/composites.py index 072e3ca33..196168033 100644 --- a/armi/reactor/composites.py +++ b/armi/reactor/composites.py @@ -36,7 +36,7 @@ import itertools import operator import timeit -from typing import Dict, List, Optional, Tuple, Type, Union +from typing import Dict, List, Optional, Tuple, Type, Union, Iterator, Callable import numpy as np import six @@ -2520,6 +2520,8 @@ class is a child-class of the ARMIObject class and provides a structure """ + _children: list["Composite"] + def __init__(self, name): ArmiObject.__init__(self, name) self.childrenByLocator = {} @@ -2617,6 +2619,72 @@ def setChildren(self, items): for c in items: self.add(c) + def iterChildren( + self, + deep=False, + generationNum=1, + predicate: Optional[Callable[["Composite"], bool]] = None, + ) -> Iterator["Composite"]: + """Iterate over children objects of this composite. + + Parameters + ---------- + deep : bool, optional + generationNum: int, optional + predicate: f(Composite) -> bool, optional, + + Returns + ------- + iterator of Composite + + See Also + -------- + :meth:`getChildren` produces a list for situations where you need to perform + multiple iterations or do list operations (append, indexing, sorting, containment, etc.) + + Composites are naturally iterable. The following are identical:: + + >>> for child in c.getChildren(): + ... pass + >>> for child in c.iterChildren(): + ... pass + >>> for child in c: + ... pass + + If you do not need any depth-traversal, natural iteration should be sufficient. + + The :func:`filter` command may be sufficient if you do not wish to pass a predicate. The following + are identical:: + >>> checker = lambda c: len(c.name) % 3 + >>> for child in c.getChildren(predicate=checker): + ... pass + >>> for child in c.iterChildren(predicate=checker): + ... pass + >>> for child in filter(checker, c): + ... pass + + If you're going to be doing traversal beyond the first generation, this method will help you. + + """ + if deep and generationNum > 1: + raise RuntimeError( + "Cannot get children with a generation number set and the deep flag set" + ) + if predicate is None: + checker = lambda _: True + else: + checker = predicate + yield from self._iterChildren(deep, generationNum, checker) + + def _iterChildren( + self, deep: bool, generationNum: int, checker: Callable[["Composite"], bool] + ) -> Iterator["Composite"]: + if deep or generationNum == 1: + yield from filter(checker, self) + if deep or generationNum > 1: + for c in self: + yield from c._iterChildren(deep, generationNum - 1, checker) + def getChildren( self, deep=False, generationNum=1, includeMaterials=False, predicate=None ): diff --git a/armi/reactor/tests/test_composites.py b/armi/reactor/tests/test_composites.py index 9ef2c7c3f..048b6931c 100644 --- a/armi/reactor/tests/test_composites.py +++ b/armi/reactor/tests/test_composites.py @@ -191,6 +191,48 @@ def test_getChildrenWithMaterials(self): else: raise RuntimeError("No materials found with includeMaterials=True") + def test_iterChildren(self): + """Detailed testing on Composite.iterChildren.""" + + def compareIterables(actual, expected: list[composites.Composite]): + for e in expected: + a = next(actual) + self.assertIs(a, e) + # Ensure we've consumed the actual iterator and there's nothing left + with self.assertRaises(StopIteration): + next(actual) + + compareIterables(self.container.iterChildren(), self.tree[1]) + compareIterables(self.container.iterChildren(generationNum=2), self.tree[2]) + compareIterables(self.container.iterChildren(generationNum=3), self.tree[3]) + compareIterables( + self.container.iterChildren(deep=True), + self.tree[1] + self.tree[2] + self.tree[3], + ) + + def test_iterAndGetChildren(self): + """Compare that iter children and get children are consistent.""" + self._compareIterGetChildren() + self._compareIterGetChildren(deep=True) + self._compareIterGetChildren(generationNum=2) + # Some wacky predicate just to check we can use that too + self._compareIterGetChildren(deep=True, predicate=lambda c: len(c.name) % 3) + + def _compareIterGetChildren(self, **kwargs): + fromIter = self.container.iterChildren(**kwargs) + fromGetter = self.container.getChildren(**kwargs) + msg = repr(kwargs) + # Use zip longest just in case one iterator comes up short + for count, (it, gt) in enumerate(itertools.zip_longest(fromIter, fromGetter)): + self.assertIs(it, gt, msg=f"{count=} :: {msg}") + + def test_simpleIterChildren(self): + """Test that C.iterChildren() is identical to iter(C).""" + for count, (fromNative, fromIterChildren) in enumerate( + itertools.zip_longest(self.container, self.container.iterChildren()) + ): + self.assertIs(fromIterChildren, fromNative, msg=count) + def test_getName(self): """Test the getName method. From 0d788c400a4a6e8911fad230f68cc5670de2bca5 Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:56 -0800 Subject: [PATCH 14/27] Composite.getChildren relies on iterChildren method Special handling for `includeMaterials=True` --- armi/reactor/composites.py | 56 ++++++++++++++++++-------------------- 1 file changed, 27 insertions(+), 29 deletions(-) diff --git a/armi/reactor/composites.py b/armi/reactor/composites.py index 196168033..fa63da5a6 100644 --- a/armi/reactor/composites.py +++ b/armi/reactor/composites.py @@ -2686,8 +2686,12 @@ def _iterChildren( yield from c._iterChildren(deep, generationNum - 1, checker) def getChildren( - self, deep=False, generationNum=1, includeMaterials=False, predicate=None - ): + self, + deep=False, + generationNum=1, + includeMaterials=False, + predicate: Optional[Callable[["Composite"], bool]] = None, + ) -> list["Composite"]: """ Return the children objects of this composite. @@ -2729,6 +2733,11 @@ def getChildren( to meet the predicate only affects the object in question; children will still be considered. + See Also + -------- + :meth:`iterChildren` if you do not need to produce a full list, e.g., just iterating + over objects. + Examples -------- >>> obj.getChildren() @@ -2745,33 +2754,22 @@ def getChildren( [grandchild1, grandchild3] """ - _pred = predicate or (lambda x: True) - if deep and generationNum > 1: - raise RuntimeError( - "Cannot get children with a generation number set and the deep flag set" - ) - - children = [] - for child in self._children: - if generationNum == 1 or deep: - if _pred(child): - children.append(child) - - if generationNum > 1 or deep: - children.extend( - child.getChildren( - deep=deep, - generationNum=generationNum - 1, - includeMaterials=includeMaterials, - predicate=predicate, - ) - ) - if includeMaterials: - material = getattr(self, "material", None) - if material: - children.append(material) - - return children + children = self.iterChildren( + deep=deep, generationNum=generationNum, predicate=predicate + ) + if not includeMaterials: + return list(children) + # Each entry is either (c, ) or (c, c.material) if the child has a material attribute + stitched = map( + lambda c: ( + (c, ) if getattr(c, "material", None) is None else (c, c.material) + ), + children, + ) + # Iterator that iterates over each "sub" iterator. If we have ((c0, ), (c1, m1)), this produces a single + # iterator of (c0, c1, m1) + flattened = itertools.chain.from_iterable(stitched) + return list(flattened) def iterChildrenWithFlags(self, typeSpec: TypeSpec, exactMatch=False): """Produce an iterator over all children of a specific type.""" From 2bf2d8ca2da70f0981eee21fcd43318665fe481c Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:56 -0800 Subject: [PATCH 15/27] The fake _Parent in test_components has a fake iter Don't need `getChildren` either because `Composite.getChildren` just needs `Composite.__iter__` to be implemented for the traversal. --- armi/reactor/tests/test_components.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/armi/reactor/tests/test_components.py b/armi/reactor/tests/test_components.py index 7ce082a54..6c4d70a38 100644 --- a/armi/reactor/tests/test_components.py +++ b/armi/reactor/tests/test_components.py @@ -165,8 +165,9 @@ def getHeight(self): def clearCache(self): pass - def getChildren(self): - return [] + def __iter__(self): + """Act like an iterator but don't actually iterate.""" + return iter(()) derivedMustUpdate = False From 263316c0aa911718cd011c2bfa56b1624d30941d Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:56 -0800 Subject: [PATCH 16/27] Use Composite.__iter__ over Composite.getChildren A lot of places have the pattern ```python for child in self.getChildren() ``` that isn't necessary. If you want to iterate over the children of a composite, you can use `Composite.__iter__` for natural iteration. --- armi/bookkeeping/report/reportInterface.py | 4 ++-- .../fuelCycle/tests/test_fuelHandlers.py | 6 +++--- .../globalFlux/globalFluxInterface.py | 2 +- .../tests/test_globalFluxInterface.py | 2 +- .../isotopicDepletion/crossSectionTable.py | 2 +- armi/reactor/blocks.py | 4 ++-- armi/reactor/blueprints/__init__.py | 2 +- armi/reactor/components/__init__.py | 2 +- armi/reactor/components/component.py | 2 +- armi/reactor/composites.py | 12 ++++++------ armi/reactor/spentFuelPool.py | 4 ++-- armi/reactor/tests/test_composites.py | 18 +++++++++--------- armi/reactor/tests/test_reactors.py | 2 +- 13 files changed, 31 insertions(+), 31 deletions(-) diff --git a/armi/bookkeeping/report/reportInterface.py b/armi/bookkeeping/report/reportInterface.py index 0105f70bb..96cac18f8 100644 --- a/armi/bookkeeping/report/reportInterface.py +++ b/armi/bookkeeping/report/reportInterface.py @@ -184,7 +184,7 @@ def reportSFP(sfp): runLog.important(title) runLog.important("-" * len(title)) totFis = 0.0 - for a in sfp.getChildren(): + for a in sfp: runLog.important( "{assembly:15s} discharged at t={dTime:10f} after {residence:10f} yrs. It entered at cycle: {cycle}. " "It has {fiss:10f} kg (x {mult}) fissile and peak BU={bu:.2f} %.".format( @@ -217,7 +217,7 @@ def countAssembliesSFP(sfp): a = sfp.getChildren()[0] lastTime = a.getAge() / units.DAYS_PER_YEAR + a.p.chargeTime - for a in sfp.getChildren(): + for a in sfp: thisTime = a.getAge() / units.DAYS_PER_YEAR + a.p.chargeTime if thisTime != lastTime: diff --git a/armi/physics/fuelCycle/tests/test_fuelHandlers.py b/armi/physics/fuelCycle/tests/test_fuelHandlers.py index 3967e3bf5..305a7ab67 100644 --- a/armi/physics/fuelCycle/tests/test_fuelHandlers.py +++ b/armi/physics/fuelCycle/tests/test_fuelHandlers.py @@ -468,7 +468,7 @@ def runShuffling(self, fh): self.r.p.cycle = cycle fh.cycle = cycle fh.manageFuel(cycle) - for a in self.r.excore["sfp"].getChildren(): + for a in self.r.excore["sfp"]: self.assertEqual(a.getLocation(), "SFP") for b in self.r.core.getBlocks(Flags.FUEL): self.assertGreater(b.p.kgHM, 0.0, "b.p.kgHM not populated!") @@ -498,7 +498,7 @@ def test_repeatShuffles(self): ensure repeatability. """ # check labels before shuffling: - for a in self.r.excore["sfp"].getChildren(): + for a in self.r.excore["sfp"]: self.assertEqual(a.getLocation(), "SFP") # do some shuffles @@ -532,7 +532,7 @@ def test_repeatShuffles(self): # make sure the shuffle was repeated perfectly. for a in self.r.core.getAssemblies(): self.assertEqual(a.getName(), firstPassResults[a.getLocation()]) - for a in self.r.excore["sfp"].getChildren(): + for a in self.r.excore["sfp"]: self.assertEqual(a.getLocation(), "SFP") # Do some cleanup, since the fuelHandler Interface has code that gets diff --git a/armi/physics/neutronics/globalFlux/globalFluxInterface.py b/armi/physics/neutronics/globalFlux/globalFluxInterface.py index 9e51adab5..bb5cf8342 100644 --- a/armi/physics/neutronics/globalFlux/globalFluxInterface.py +++ b/armi/physics/neutronics/globalFlux/globalFluxInterface.py @@ -277,7 +277,7 @@ def getTightCouplingValue(self): return self.r.core.p.keff if self.coupler.parameter == "power": scaledCorePowerDistribution = [] - for a in self.r.core.getChildren(): + for a in self.r.core: scaledPower = [] assemPower = sum(b.p.power for b in a) for b in a: diff --git a/armi/physics/neutronics/globalFlux/tests/test_globalFluxInterface.py b/armi/physics/neutronics/globalFlux/tests/test_globalFluxInterface.py index 09adacaf8..ae9c0e96c 100644 --- a/armi/physics/neutronics/globalFlux/tests/test_globalFluxInterface.py +++ b/armi/physics/neutronics/globalFlux/tests/test_globalFluxInterface.py @@ -286,7 +286,7 @@ def test_getTightCouplingValue(self): self._setTightCouplingTrue() self.assertEqual(self.gfi.getTightCouplingValue(), 1.0) # set in setUp self.gfi.coupler.parameter = "power" - for a in self.r.core.getChildren(): + for a in self.r.core: for b in a: b.p.power = 10.0 self.assertEqual( diff --git a/armi/physics/neutronics/isotopicDepletion/crossSectionTable.py b/armi/physics/neutronics/isotopicDepletion/crossSectionTable.py index a7921e3bc..3103f7f31 100644 --- a/armi/physics/neutronics/isotopicDepletion/crossSectionTable.py +++ b/armi/physics/neutronics/isotopicDepletion/crossSectionTable.py @@ -223,7 +223,7 @@ def makeReactionRateTable(obj, nuclides: List = None): for nucName in nuclides } - for armiObject in obj.getChildren(): + for armiObject in obj: for nucName in nuclides: rxnRates = armiObject.getReactionRates(nucName, nDensity=1.0) for rxName, rxRate in rxnRates.items(): diff --git a/armi/reactor/blocks.py b/armi/reactor/blocks.py index 4c3edd8ab..eadf14b2f 100644 --- a/armi/reactor/blocks.py +++ b/armi/reactor/blocks.py @@ -627,7 +627,7 @@ def getArea(self, cold=False): return area a = 0.0 - for c in self.getChildren(): + for c in self: myArea = c.getArea(cold=cold) a += myArea fullArea = a @@ -1283,7 +1283,7 @@ def getPinPitch(self, cold=False): def getDimensions(self, dimension): """Return dimensional values of the specified dimension.""" dimVals = set() - for c in self.getChildren(): + for c in self: try: dimVal = c.getDimension(dimension) except parameters.ParameterError: diff --git a/armi/reactor/blueprints/__init__.py b/armi/reactor/blueprints/__init__.py index 48532ed58..3bee854e4 100644 --- a/armi/reactor/blueprints/__init__.py +++ b/armi/reactor/blueprints/__init__.py @@ -519,7 +519,7 @@ def _checkAssemblyAreaConsistency(self, cs): runLog.error("CURRENT COMPARISON BLOCK:") b.printContents(includeNuclides=False) - for c in b.getChildren(): + for c in b: runLog.error( "{0} area {1} effective area {2}" "".format(c, c.getArea(), c.getVolume() / b.getHeight()) diff --git a/armi/reactor/components/__init__.py b/armi/reactor/components/__init__.py index 516cdae81..f4fe2c6d8 100644 --- a/armi/reactor/components/__init__.py +++ b/armi/reactor/components/__init__.py @@ -383,7 +383,7 @@ def _deriveVolumeAndArea(self): # Determine the volume/areas of the non-derived shape components within the parent. siblingVolume = 0.0 siblingArea = 0.0 - for sibling in self.parent.getChildren(): + for sibling in self.parent: if sibling is self: continue elif not self and isinstance(sibling, DerivedShape): diff --git a/armi/reactor/components/component.py b/armi/reactor/components/component.py index 5a4da2f3b..be2ce62db 100644 --- a/armi/reactor/components/component.py +++ b/armi/reactor/components/component.py @@ -964,7 +964,7 @@ def clearLinkedCache(self): def getLinkedComponents(self): """Find other components that are linked to this component.""" dependents = [] - for child in self.parent.getChildren(): + for child in self.parent: for dimName in child.DIMENSION_NAMES: isLinked = child.dimensionIsLinked(dimName) if isLinked and child.p[dimName].getLinkedComponent() is self: diff --git a/armi/reactor/composites.py b/armi/reactor/composites.py index fa63da5a6..e9e7d1c21 100644 --- a/armi/reactor/composites.py +++ b/armi/reactor/composites.py @@ -471,7 +471,7 @@ def duplicate(self): def clearCache(self): """Clear the cache so all new values are recomputed.""" self.cached = {} - for child in self.getChildren(): + for child in self: child.clearCache() def _getCached(self, name): # TODO: stop the "returns None" nonsense? @@ -617,7 +617,7 @@ def copyParamsToChildren(self, paramNames): """ for paramName in paramNames: myVal = self.p[paramName] - for c in self.getChildren(): + for c in self: c.p[paramName] = myVal @classmethod @@ -1631,7 +1631,7 @@ def setLumpedFissionProducts(self, lfpCollection): self._lumpedFissionProducts = lfpCollection def setChildrenLumpedFissionProducts(self, lfpCollection): - for c in self.getChildren(): + for c in self: c.setLumpedFissionProducts(lfpCollection) def getFissileMassEnrich(self): @@ -1932,7 +1932,7 @@ def getNuclides(self): List of nuclide names that exist in this """ nucs = set() - for child in self.getChildren(): + for child in self: nucs.update(child.getNuclides()) return nucs @@ -3027,7 +3027,7 @@ def getLumpedFissionProductCollection(self): """ lfps = ArmiObject.getLumpedFissionProductCollection(self) if lfps is None: - for c in self.getChildren(): + for c in self: lfps = c.getLumpedFissionProductCollection() if lfps is not None: break @@ -3165,7 +3165,7 @@ def getReactionRates(self, nucName, nDensity=None): def printContents(self, includeNuclides=True): """Display information about all the comprising children in this object.""" runLog.important(self) - for c in self.getChildren(): + for c in self: c.printContents(includeNuclides=includeNuclides) def _genChildByLocationLookupTable(self): diff --git a/armi/reactor/spentFuelPool.py b/armi/reactor/spentFuelPool.py index 666c646b4..a3e374de6 100644 --- a/armi/reactor/spentFuelPool.py +++ b/armi/reactor/spentFuelPool.py @@ -77,7 +77,7 @@ def add(self, assem, loc=None): def getAssembly(self, name): """Get a specific assembly by name.""" - for a in self.getChildren(): + for a in self: if a.getName() == name: return a @@ -123,7 +123,7 @@ def normalizeNames(self, startIndex=0): The new max Assembly number. """ ind = startIndex - for a in self.getChildren(): + for a in self: oldName = a.getName() newName = a.makeNameFromAssemNum(ind) if oldName == newName: diff --git a/armi/reactor/tests/test_composites.py b/armi/reactor/tests/test_composites.py index 048b6931c..318acd34c 100644 --- a/armi/reactor/tests/test_composites.py +++ b/armi/reactor/tests/test_composites.py @@ -247,26 +247,26 @@ def test_getName(self): def test_sort(self): # in this case, the children should start sorted - c0 = [c.name for c in self.container.getChildren()] + c0 = [c.name for c in self.container] self.container.sort() - c1 = [c.name for c in self.container.getChildren()] + c1 = [c.name for c in self.container] self.assertNotEqual(c0, c1) # verify repeated sortings behave for _ in range(3): self.container.sort() - ci = [c.name for c in self.container.getChildren()] + ci = [c.name for c in self.container] self.assertEqual(c1, ci) # break the order children = self.container.getChildren() self.container._children = children[2:] + children[:2] - c2 = [c.name for c in self.container.getChildren()] + c2 = [c.name for c in self.container] self.assertNotEqual(c1, c2) # verify the sort order self.container.sort() - c3 = [c.name for c in self.container.getChildren()] + c3 = [c.name for c in self.container] self.assertEqual(c1, c3) def test_areChildernOfType(self): @@ -397,12 +397,12 @@ def test_setChildrenLumpedFissionProducts(self): # validate that the LFP collection is None self.container.setChildrenLumpedFissionProducts(None) - for c in self.container.getChildren(): + for c in self.container: self.assertIsNone(c._lumpedFissionProducts) # validate that the LFP collection is not None self.container.setChildrenLumpedFissionProducts(lfps) - for c in self.container.getChildren(): + for c in self.container: self.assertIsNotNone(c._lumpedFissionProducts) def test_requiresLumpedFissionProducts(self): @@ -900,7 +900,7 @@ def test_getNumberDensities(self): # sum nuc densities from children components totalVolume = self.obj.getVolume() childDensities = {} - for o in self.obj.getChildren(): + for o in self.obj: m = o.getVolume() d = o.getNumberDensities() for nuc, val in d.items(): @@ -938,7 +938,7 @@ def test_getNumberDensitiesWithExpandedFissionProducts(self): # sum nuc densities from children components totalVolume = self.obj.getVolume() childDensities = {} - for o in self.obj.getChildren(): + for o in self.obj: # get the number densities with and without fission products d0 = o.getNumberDensities(expandFissionProducts=False) d = o.getNumberDensities(expandFissionProducts=True) diff --git a/armi/reactor/tests/test_reactors.py b/armi/reactor/tests/test_reactors.py index 6187155e6..69790e6c7 100644 --- a/armi/reactor/tests/test_reactors.py +++ b/armi/reactor/tests/test_reactors.py @@ -360,7 +360,7 @@ def test_growToFullCore(self): self.assertFalse(self.r.core.isFullCore) self.r.core.growToFullCore(self.o.cs) aNums = [] - for a in self.r.core.getChildren(): + for a in self.r.core: self.assertNotIn(a.getNum(), aNums) aNums.append(a.getNum()) From 267430875742f62e9378590b872c8be9a854d4f5 Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:57 -0800 Subject: [PATCH 17/27] Use Composite.__len__ and Composite.__getitem__ for some reporting --- armi/bookkeeping/report/reportInterface.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/armi/bookkeeping/report/reportInterface.py b/armi/bookkeeping/report/reportInterface.py index 96cac18f8..c970db6e4 100644 --- a/armi/bookkeeping/report/reportInterface.py +++ b/armi/bookkeeping/report/reportInterface.py @@ -208,13 +208,13 @@ def reportSFP(sfp): @staticmethod def countAssembliesSFP(sfp): """Report on the count of assemblies in the SFP at each timestep.""" - if not sfp.getChildren(): + if not len(sfp): return runLog.important("Count:") totCount = 0 thisTimeCount = 0 - a = sfp.getChildren()[0] + a = sfp[0] lastTime = a.getAge() / units.DAYS_PER_YEAR + a.p.chargeTime for a in sfp: From 5b85cf280843f0bf00a6a0286eaa6dae2277ff04 Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:57 -0800 Subject: [PATCH 18/27] Use Composite._iter__ over Composite.getChildren in fuelHandlers We just need a list of assemblies. `list(SFP)` will do that for us --- armi/physics/fuelCycle/fuelHandlers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/armi/physics/fuelCycle/fuelHandlers.py b/armi/physics/fuelCycle/fuelHandlers.py index 9f99d5a8f..5996f3d15 100644 --- a/armi/physics/fuelCycle/fuelHandlers.py +++ b/armi/physics/fuelCycle/fuelHandlers.py @@ -654,7 +654,7 @@ def _getAssembliesInRings( f"or otherwise instantiate a SpentFuelPool object as r.excore['sfp']" ) else: - sfpAssems = self.r.excore["sfp"].getChildren() + sfpAssems = list(self.r.excore["sfp"]) assemblyList = [[] for _i in range(len(ringList))] # empty lists for each ring if exclusions is None: From 824c843be759b3f35b7a2cac9ab78bd73877c7fc Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:57 -0800 Subject: [PATCH 19/27] Use Reactor.iterChildren with a predicate during a read-only db load --- armi/bookkeeping/db/database.py | 8 ++++---- armi/reactor/reactorParameters.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/armi/bookkeeping/db/database.py b/armi/bookkeeping/db/database.py index 5a718cda8..85d555d0e 100644 --- a/armi/bookkeeping/db/database.py +++ b/armi/bookkeeping/db/database.py @@ -829,11 +829,11 @@ def loadReadOnly(self, cycle, node, statePointName=None): return r @staticmethod - def _setParamsBeforeFreezing(r): + def _setParamsBeforeFreezing(r: Reactor): """Set some special case parameters before they are made read-only.""" - for child in r.getChildren(deep=True): - if not isinstance(child, Component): - continue + for child in r.iterChildren( + deep=True, predicate=lambda c: isinstance(c, Component) + ): # calling Component.getVolume() sets the volume parameter child.getVolume() diff --git a/armi/reactor/reactorParameters.py b/armi/reactor/reactorParameters.py index 90eeb185b..3f7577d27 100644 --- a/armi/reactor/reactorParameters.py +++ b/armi/reactor/reactorParameters.py @@ -789,5 +789,5 @@ def makeParametersReadOnly(r): Once you make one Reactor read-only, you cannot make it writeable again. """ r.p.readOnly = True - for child in r.getChildren(deep=True): + for child in r.iterChildren(deep=True): child.p.readOnly = True From 8803d03e6fd42f8cc10ff309b7bd88f1aa404dc3 Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:57 -0800 Subject: [PATCH 20/27] Block.removeAll uses a list of self over getChildren --- armi/reactor/blocks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/armi/reactor/blocks.py b/armi/reactor/blocks.py index eadf14b2f..3e0e1fda0 100644 --- a/armi/reactor/blocks.py +++ b/armi/reactor/blocks.py @@ -930,7 +930,7 @@ def add(self, c): self._updatePitchComponent(c) def removeAll(self, recomputeAreaFractions=True): - for c in self.getChildren(): + for c in list(self): self.remove(c, recomputeAreaFractions=False) if recomputeAreaFractions: # only do this once self.getVolumeFractions() From 329fb09da3ba18ec3aac6fcba3d4fec8389f6a30 Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:57 -0800 Subject: [PATCH 21/27] Try to reduce list comprehension in ExpansionData targets --- .../converters/axialExpansionChanger/expansionData.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/armi/reactor/converters/axialExpansionChanger/expansionData.py b/armi/reactor/converters/axialExpansionChanger/expansionData.py index 39d4dacaa..985de6f13 100644 --- a/armi/reactor/converters/axialExpansionChanger/expansionData.py +++ b/armi/reactor/converters/axialExpansionChanger/expansionData.py @@ -288,19 +288,17 @@ def determineTargetComponent( if flagOfInterest is None: # Follow expansion of most neutronically important component, fuel then control/poison for targetFlag in TARGET_FLAGS_IN_PREFERRED_ORDER: - candidates = [c for c in b.getChildren() if c.hasFlags(targetFlag)] + candidates = b.getChildrenWithFlags(targetFlag) if candidates: break # some blocks/components are not included in the above list but should still be found if not candidates: candidates = [c for c in b.getChildren() if c.p.flags in b.p.flags] else: - candidates = [c for c in b.getChildren() if c.hasFlags(flagOfInterest)] + candidates = b.getChildrenWithFlags(flagOfInterest) if len(candidates) == 0: # if only 1 solid, be smart enought to snag it - solidMaterials = list( - c for c in b if not isinstance(c.material, material.Fluid) - ) + solidMaterials = getSolidComponents(b) if len(solidMaterials) == 1: candidates = solidMaterials if len(candidates) == 0: From c631c90077895c6bb0eafae63774a5b048c1d7f5 Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:57 -0800 Subject: [PATCH 22/27] fixup docs for Composite.iterChildren --- armi/reactor/composites.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/armi/reactor/composites.py b/armi/reactor/composites.py index e9e7d1c21..c444f04c5 100644 --- a/armi/reactor/composites.py +++ b/armi/reactor/composites.py @@ -2626,22 +2626,27 @@ def iterChildren( predicate: Optional[Callable[["Composite"], bool]] = None, ) -> Iterator["Composite"]: """Iterate over children objects of this composite. - + Parameters ---------- deep : bool, optional + If true, traverse the entire composite tree. Otherwise, go as far as ``generationNum``. generationNum: int, optional - predicate: f(Composite) -> bool, optional, - + Produce composites at this depth. A depth of ``1`` includes children of ``self``, ``2`` + is children of children, and so on. + predicate: f(Composite) -> bool, optional + Function to check on a composite before producing it. All items in the iteration + will pass this check. + Returns ------- iterator of Composite - + See Also -------- :meth:`getChildren` produces a list for situations where you need to perform multiple iterations or do list operations (append, indexing, sorting, containment, etc.) - + Composites are naturally iterable. The following are identical:: >>> for child in c.getChildren(): @@ -2652,7 +2657,7 @@ def iterChildren( ... pass If you do not need any depth-traversal, natural iteration should be sufficient. - + The :func:`filter` command may be sufficient if you do not wish to pass a predicate. The following are identical:: >>> checker = lambda c: len(c.name) % 3 @@ -2664,7 +2669,7 @@ def iterChildren( ... pass If you're going to be doing traversal beyond the first generation, this method will help you. - + """ if deep and generationNum > 1: raise RuntimeError( From e1b01917c6abc337d606450855f4273dca71c620 Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Tue, 10 Dec 2024 14:56:58 -0800 Subject: [PATCH 23/27] Provide and use Composite.iterChildrenWithMaterials --- armi/reactor/composites.py | 68 ++++++++++++++++++++++++++------------ 1 file changed, 46 insertions(+), 22 deletions(-) diff --git a/armi/reactor/composites.py b/armi/reactor/composites.py index c444f04c5..ff35d56e4 100644 --- a/armi/reactor/composites.py +++ b/armi/reactor/composites.py @@ -2690,6 +2690,28 @@ def _iterChildren( for c in self: yield from c._iterChildren(deep, generationNum - 1, checker) + def iterChildrenWithMaterials(self, *args, **kwargs) -> Iterator: + """Produce an iterator that also includes any materials found on descendants. + + Arguments are forwarded to :meth:`iterChildren` and control the depth of traversal + and filtering of objects. + + This is useful for sending state across MPI tasks where you need a more full + representation of the composite tree. Which includes the materials attached + to components. + """ + children = self.iterChildren(*args, **kwargs) + # Each entry is either (c, ) or (c, c.material) if the child has a material attribute + stitched = map( + lambda c: ( + (c,) if getattr(c, "material", None) is None else (c, c.material) + ), + children, + ) + # Iterator that iterates over each "sub" iterator. If we have ((c0, ), (c1, m1)), this produces a single + # iterator of (c0, c1, m1) + return itertools.chain.from_iterable(stitched) + def getChildren( self, deep=False, @@ -2759,22 +2781,15 @@ def getChildren( [grandchild1, grandchild3] """ - children = self.iterChildren( - deep=deep, generationNum=generationNum, predicate=predicate - ) if not includeMaterials: - return list(children) - # Each entry is either (c, ) or (c, c.material) if the child has a material attribute - stitched = map( - lambda c: ( - (c, ) if getattr(c, "material", None) is None else (c, c.material) - ), - children, - ) - # Iterator that iterates over each "sub" iterator. If we have ((c0, ), (c1, m1)), this produces a single - # iterator of (c0, c1, m1) - flattened = itertools.chain.from_iterable(stitched) - return list(flattened) + items = self.iterChildren( + deep=deep, generationNum=generationNum, predicate=predicate + ) + else: + items = self.iterChildrenWithMaterials( + deep=deep, generationNum=generationNum, predicate=predicate + ) + return list(items) def iterChildrenWithFlags(self, typeSpec: TypeSpec, exactMatch=False): """Produce an iterator over all children of a specific type.""" @@ -2845,8 +2860,11 @@ def syncMpiState(self): startTime = timeit.default_timer() # sync parameters... - allComps = [self] + self.getChildren(deep=True, includeMaterials=True) - allComps = [c for c in allComps if hasattr(c, "p")] + genItems = itertools.chain( + [self], + self.iterChildrenWithMaterials(deep=True), + ) + allComps = [c for c in genItems if hasattr(c, "p")] sendBuf = [c.p.getSyncData() for c in allComps] runLog.debug("syncMpiState has {} comps".format(len(allComps))) @@ -2951,7 +2969,11 @@ def _markSynchronized(self): SINCE_LAST_DISTRIBUTE_STATE. """ paramDefs = set() - for child in [self] + self.getChildren(deep=True, includeMaterials=True): + items = itertools.chain( + [self], + self.iterChildrenWithMaterials(deep=True), + ) + for child in items: # Materials don't have a "p" / Parameter attribute to sync if hasattr(child, "p"): # below reads as: assigned & everything_but(SINCE_LAST_DISTRIBUTE_STATE) @@ -3216,7 +3238,7 @@ class StateRetainer: """ - def __init__(self, composite, paramsToApply=None): + def __init__(self, composite: Composite, paramsToApply=None): """ Create an instance of a StateRetainer. @@ -3244,9 +3266,11 @@ def _enterExitHelper(self, func): ``backUp()`` or ``restoreBackup()``. """ paramDefs = set() - for child in [self.composite] + self.composite.getChildren( - deep=True, includeMaterials=True - ): + items = itertools.chain( + (self.composite,), + self.composite.iterChildrenWithMaterials(deep=True), + ) + for child in items: if hasattr(child, "p"): # materials don't have Parameters paramDefs.update(child.p.paramDefs) From bc1aa10afb70036735a254c5efcf392d625c0b73 Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Wed, 11 Dec 2024 08:45:54 -0800 Subject: [PATCH 24/27] Lint cleanup in docstrings --- armi/reactor/assemblies.py | 2 +- armi/reactor/composites.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/armi/reactor/assemblies.py b/armi/reactor/assemblies.py index af34b9bac..8a16634c8 100644 --- a/armi/reactor/assemblies.py +++ b/armi/reactor/assemblies.py @@ -826,7 +826,7 @@ def iterBlocks(self, typeSpec=None, exact=False): ------- iterable of Block - See also + See Also -------- * :meth:`__iter__` - if no type spec provided, assemblies can be naturally iterated upon. diff --git a/armi/reactor/composites.py b/armi/reactor/composites.py index ff35d56e4..19fd6c22a 100644 --- a/armi/reactor/composites.py +++ b/armi/reactor/composites.py @@ -2801,7 +2801,7 @@ def getChildrenWithFlags(self, typeSpec: TypeSpec, exactMatch=False): return list(self.iterChildrenWithFlags(typeSpec, exactMatch)) def iterChildrenOfType(self, typeName: str): - """Produce an iterator over all children with a specific input type name""" + """Produce an iterator over all children with a specific input type name.""" return filter(lambda c: c.getType() == typeName, self) def getChildrenOfType(self, typeName: str): From 7252468afe84bd6b99c80f3f3d9543b88f405b50 Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Fri, 13 Dec 2024 14:50:32 -0800 Subject: [PATCH 25/27] Add Composite.iter* release notes --- doc/release/0.5.rst | 3 +++ 1 file changed, 3 insertions(+) diff --git a/doc/release/0.5.rst b/doc/release/0.5.rst index f730979d6..1605e949b 100644 --- a/doc/release/0.5.rst +++ b/doc/release/0.5.rst @@ -9,6 +9,9 @@ Release Date: TBD New Features ------------ #. Move instead of copy files from TemporaryDirectoryChanger. (`PR#2022 `_) +#. Provide ``Composite.iterChildren``, ``Composite.iterChildrenWithFlags``, ``Composite.iterChildrenOfType``, + ``Composite.iterChildrenWithMaterials``, and ``Assembly.iterBlocks`` for efficient composite tree traversal. + (`PR#2031 `_) API Changes ----------- From 355e89f484da623d7b21ac125c8185da0a35ce8d Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Wed, 18 Dec 2024 10:45:58 -0800 Subject: [PATCH 26/27] Remove Assembly.iterBlocks --- armi/reactor/assemblies.py | 49 ++++++++++++++------------------------ doc/release/0.5.rst | 2 +- 2 files changed, 19 insertions(+), 32 deletions(-) diff --git a/armi/reactor/assemblies.py b/armi/reactor/assemblies.py index 8a16634c8..c46a73aa6 100644 --- a/armi/reactor/assemblies.py +++ b/armi/reactor/assemblies.py @@ -327,7 +327,7 @@ def getPinPlenumVolumeInCubicMeters(self): This is a bit design-specific for pinned assemblies """ plenumVolume = 0.0 - for b in self.iterBlocks(Flags.PLENUM): + for b in self.iterChildrenWithFlags(Flags.PLENUM): cladId = b.getComponent(Flags.CLAD).getDimension("id") length = b.getHeight() plenumVolume += ( @@ -337,7 +337,7 @@ def getPinPlenumVolumeInCubicMeters(self): def getAveragePlenumTemperature(self): """Return the average of the plenum block outlet temperatures.""" - plenumBlocks = self.iterBlocks(Flags.PLENUM) + plenumBlocks = self.iterChildrenWithFlags(Flags.PLENUM) plenumTemps = [b.p.THcoolantOutletT for b in plenumBlocks] # no plenum blocks, use the top block of the assembly for plenum temperature @@ -812,32 +812,7 @@ def dump(self, fName=None): with open(fName, "w") as pkl: pickle.dump(self, pkl) - def iterBlocks(self, typeSpec=None, exact=False): - """Produce an iterator over all blocks in this assembly from bottom to top. - - Parameters - ---------- - typeSpec : Flags or list of Flags, optional - Restrict returned blocks to have these flags. - exact : bool, optional - If true, only produce blocks that have those exact flags. - - Returns - ------- - iterable of Block - - See Also - -------- - * :meth:`__iter__` - if no type spec provided, assemblies can be - naturally iterated upon. - * :meth:`iterChildrenWithFlags` - alternative if you know you have - a type spec that isn't ``None``. - """ - if typeSpec is None: - return iter(self) - return self.iterChildrenWithFlags(typeSpec, exact) - - def getBlocks(self, typeSpec=None, exact=False): + def getBlocks(self, typeSpec=None, exact=False) -> list[blocks.Block]: """ Get blocks in an assembly from bottom to top. @@ -853,7 +828,11 @@ def getBlocks(self, typeSpec=None, exact=False): blocks : list List of blocks. """ - return list(self.iterBlocks(typeSpec, exact)) + if typeSpec is None: + items = iter(self) + else: + items = self.iterChildrenWithFlags(typeSpec, exact) + return list(items) def getBlocksAndZ(self, typeSpec=None, returnBottomZ=False, returnTopZ=False): """ @@ -918,9 +897,13 @@ def getFirstBlock(self, typeSpec=None, exact=False): Block or None First block that matches if such a block could be found. """ + if typeSpec is None: + items = iter(self) + else: + items = self.iterChildrenWithFlags(typeSpec, exact) try: # Create an iterator and attempt to advance it to the first value. - return next(self.iterBlocks(typeSpec, exact)) + return next(items) except StopIteration: # No items found in the iteration -> no blocks match the request return None @@ -1221,7 +1204,11 @@ def countBlocksWithFlags(self, blockTypeSpec=None): blockCounter : int number of blocks of this type """ - return sum(1 for _ in self.iterBlocks(blockTypeSpec)) + if blockTypeSpec is None: + items = iter(self) + else: + items = self.iterChildrenWithFlags(blockTypeSpec) + return sum(1 for _ in items) def getDim(self, typeSpec, dimName): """ diff --git a/doc/release/0.5.rst b/doc/release/0.5.rst index 1605e949b..2eab84399 100644 --- a/doc/release/0.5.rst +++ b/doc/release/0.5.rst @@ -10,7 +10,7 @@ New Features ------------ #. Move instead of copy files from TemporaryDirectoryChanger. (`PR#2022 `_) #. Provide ``Composite.iterChildren``, ``Composite.iterChildrenWithFlags``, ``Composite.iterChildrenOfType``, - ``Composite.iterChildrenWithMaterials``, and ``Assembly.iterBlocks`` for efficient composite tree traversal. + ``Composite.iterChildrenWithMaterials``, for efficient composite tree traversal. (`PR#2031 `_) API Changes From fc86fd156d018f3ff4e189113b3e074b6e0a1646 Mon Sep 17 00:00:00 2001 From: Andrew Johnson Date: Thu, 19 Dec 2024 09:25:01 -0800 Subject: [PATCH 27/27] Promote more Composite.iter* and Composite.get* to ArmiObject Once we have `ArmiObject.iterChildren` specified on `ArmiObject`, we can use that to build more base functionality: - `ArmiObject.iterChildrenWithFlags` - `ArmiObject.getChildrenWithFlags` - `ArmiObject.iterChildrenOfType` - `ArmiObject.getChildrenOfType` Additionally, I noted that `ArmiObject.getChildrenWithFlags` had a different default `exactMatch=True` than `Composite.getChildrenWithFlags`. Both now default to `exactMatch=False` for consistency. --- armi/reactor/composites.py | 52 +++++++++++++++++++++----------------- 1 file changed, 29 insertions(+), 23 deletions(-) diff --git a/armi/reactor/composites.py b/armi/reactor/composites.py index 19fd6c22a..25e024c1e 100644 --- a/armi/reactor/composites.py +++ b/armi/reactor/composites.py @@ -521,17 +521,40 @@ def updateParamsFrom(self, new): for paramName, val in new.p.items(): self.p[paramName] = val - def getChildren(self, deep=False, generationNum=1, includeMaterials=False): + def iterChildren( + self, + deep=False, + generationNum=1, + predicate: Optional[Callable[["ArmiObject"], bool]] = None, + ) -> Iterator["ArmiObject"]: + """Iterate over children of this object.""" + raise NotImplementedError() + + def getChildren( + self, deep=False, generationNum=1, includeMaterials=False + ) -> list["ArmiObject"]: """Return the children of this object.""" - raise NotImplementedError + raise NotImplementedError() - def iterChildrenWithFlags(self, typeSpec: TypeSpec, exactMatch=True): + def iterChildrenWithFlags( + self, typeSpec: TypeSpec, exactMatch=False + ) -> Iterator["ArmiObject"]: """Produce an iterator of children that have given flags.""" - raise NotImplementedError() + return self.iterChildren(predicate=lambda o: o.hasFlags(typeSpec, exactMatch)) - def getChildrenWithFlags(self, typeSpec: TypeSpec, exactMatch=True): + def getChildrenWithFlags( + self, typeSpec: TypeSpec, exactMatch=False + ) -> list["ArmiObject"]: """Get all children that have given flags.""" - raise NotImplementedError + return list(self.iterChildrenWithFlags(typeSpec, exactMatch)) + + def iterChildrenOfType(self, typeName: str) -> Iterator["ArmiObject"]: + """Iterate over children that have a specific input type name.""" + return self.iterChildren(predicate=lambda o: o.getType() == typeName) + + def getChildrenOfType(self, typeName: str) -> list["ArmiObject"]: + """Produce a list of children that have a specific input type name.""" + return list(self.iterChildrenOfType(typeName)) def getComponents(self, typeSpec: TypeSpec = None, exact=False): """ @@ -2791,23 +2814,6 @@ def getChildren( ) return list(items) - def iterChildrenWithFlags(self, typeSpec: TypeSpec, exactMatch=False): - """Produce an iterator over all children of a specific type.""" - checker = operator.methodcaller("hasFlags", typeSpec, exactMatch) - return filter(checker, self) - - def getChildrenWithFlags(self, typeSpec: TypeSpec, exactMatch=False): - """Get all children of a specific type.""" - return list(self.iterChildrenWithFlags(typeSpec, exactMatch)) - - def iterChildrenOfType(self, typeName: str): - """Produce an iterator over all children with a specific input type name.""" - return filter(lambda c: c.getType() == typeName, self) - - def getChildrenOfType(self, typeName: str): - """Get children that have a specific input type name.""" - return list(self.iterChildrenOfType(typeName)) - def getComponents(self, typeSpec: TypeSpec = None, exact=False): return list(self.iterComponents(typeSpec, exact))