terrapower / armi

An open-source nuclear reactor analysis automation framework that helps design teams increase efficiency and quality
https://terrapower.github.io/armi/
Apache License 2.0
231 stars 89 forks source link

Provide more Composite.iter* methods #1915

Open drewj-tp opened 1 month ago

drewj-tp commented 1 month ago

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)

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 commented 1 month ago

Future work

Things that could be looked at later

Non-exhaustive

john-science commented 1 month ago
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.

drewj-tp commented 1 month ago

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 commented 1 month ago

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