Provide more Composite.iter* methods
ARMI does a lot of iteration over the composite tree. Understandably so, we have lots of methods that support iterating over children of an object, with some filtering. These usually follow the pattern of parent.getChildren* like Composite.getChildren, Composite.getChildrenWithFlags, and Composite.getChildrenOfType, further expanded to things like Assembly.getBlocks.
A lot of instances of the usage of these get* methods are for iteration
https://github.com/terrapower/armi/blob/96a5c88d6167cee989378599b6d341c351c291bd/armi/reactor/components/component.py#L965
https://github.com/terrapower/armi/blob/96a5c88d6167cee989378599b6d341c351c291bd/armi/reactor/converters/axialExpansionChanger/expansionData.py#L253
https://github.com/terrapower/armi/blob/96a5c88d6167cee989378599b6d341c351c291bd/armi/reactor/blocks.py#L633
The tradeoff is these get* methods create and populate lists, and then return those lists back to the user. This is great if you need to do multiple iterations or check sizing of the returned object, but it's inefficient for iteration.
I propose we provide iter* methods with similar signatures as their get* counterparts that do not return lists of things. Instead, they could yield back components, produce a filter object, etc. something that allows one-pass iteration over things. The get* methods could then be distilled to
def getChildrenWithFlags(self, ...):
return list(self.iterChildrenWithFlags(...))
for backwards compatibility.
Candidate methods
To keep the scope clear, I propose we add, test, and use (where appropriate)
-
Composite.iterChildren -
Composite.iterChildrenWithFlags -
Composite.iterChildrenOfType
There are other composite subclasses that I believe could benefit from more iteration methods. But for start, I'll keep the scope limited to Composite methods
Future work
Things that could be looked at later
-
Assembly.iterBlocks -
Assembly.iterBlocksBetweenHeights -
Block.iterComponentsThatAreLinkedTo -
Core.iterAssemblies -
Core.iterAssembliesInRing -
Core.iterBlocks
Non-exhaustive
Things that could be looked at later
Assembly.iterBlocks
Assembly.iterBlocksBetweenHeights
Block.iterComponentsThatAreLinkedTo
Core.iterAssemblies
Core.iterAssembliesInRing
Core.iterBlocks
If I were to be choosy, I would do Core.iterAssemblies and Core.iterAssembliesInRing last. In an ideal world I would wait until we separate out Core and HexCore or something.
Just thinking aloud.
If I were to be choosy, I would do Core.iterAssemblies and Core.iterAssembliesInRing last. In an ideal world I would wait until we separate out Core and HexCore or something.
Agree. I don't intend on tackling things outside Composite.iter* in the first go
Some additional digging.
Composite.getChildren has an option that also adds the materials to the list of returned things found in the recursive dig. That feels weird.
That mode, includeMaterials=True is only used twice in armi and internal projects I can access. Both instances are for MPI related distribution and syncing
https://github.com/terrapower/armi/blob/f1f3dbea05c2564404a03d8e053f72c2af15daa9/armi/reactor/composites.py#L2778
https://github.com/terrapower/armi/blob/f1f3dbea05c2564404a03d8e053f72c2af15daa9/armi/reactor/composites.py#L2884
But Composite.getChildren is used heavily! And a lot of iteration, e.g., for c in thing.getChildren. But in a lot of cases, where no arguments are provided, it seems to be a way to just iterate over the children of a composite. Which is naturally supported with __iter__
https://github.com/terrapower/armi/blob/f1f3dbea05c2564404a03d8e053f72c2af15daa9/armi/reactor/composites.py#L2535-L2536
There are few recursive usage of getChildren, like for VTK visualization
https://github.com/terrapower/armi/blob/f1f3dbea05c2564404a03d8e053f72c2af15daa9/armi/bookkeeping/visualization/vtk.py#L100-L103
So maybe some wide-spread improvements could be
- Replace single-level, material-less usage of
for child in parent.getChildrenwith plain iterationfor child in parent. - Recursive iterator that digs through the composite tree, completely or at a specific
generationNum - Update the recursive
getChildren, with materials, to use the recursive iterator