armi icon indicating copy to clipboard operation
armi copied to clipboard

Provide more Composite.iter* methods

Open drewj-tp opened this issue 1 year ago • 4 comments

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

drewj-tp avatar Sep 30 '24 16:09 drewj-tp

Future work

Things that could be looked at later

  • Assembly.iterBlocks
  • Assembly.iterBlocksBetweenHeights
  • Block.iterComponentsThatAreLinkedTo
  • Core.iterAssemblies
  • Core.iterAssembliesInRing
  • Core.iterBlocks

Non-exhaustive

drewj-tp avatar Sep 30 '24 16:09 drewj-tp

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.

john-science avatar Oct 01 '24 01:10 john-science

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

drewj-tp avatar Oct 01 '24 16:10 drewj-tp

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

  1. Replace single-level, material-less usage of for child in parent.getChildren with plain iteration for child in parent.
  2. Recursive iterator that digs through the composite tree, completely or at a specific generationNum
  3. Update the recursive getChildren, with materials, to use the recursive iterator

drewj-tp avatar Oct 02 '24 17:10 drewj-tp