Skip to content

Improve types around potentials in ExpansionWranglerInterface - #126

Merged
inducer merged 3 commits into
inducer:mainfrom
alexfikl:type-fmm
Sep 2, 2026
Merged

Improve types around potentials in ExpansionWranglerInterface#126
inducer merged 3 commits into
inducer:mainfrom
alexfikl:type-fmm

Conversation

@alexfikl

Copy link
Copy Markdown
Contributor

This adds a little PotentialArray = Array | ObjectArray[Array] alias and uses it in all the functions working with potentials: reorder_potentials, eval_multipoles, etc.

I tried doing some sort of PotentialArrayT = TypeVar(..., bound=CanAdd), but e.g. the pyfmmlib wrangler decides at runtime if the potential is an array or an object array, so that doesn't work.

xref: inducer/sumpy#300

Comment thread boxtree/fmm.py
Comment on lines -758 to +763
def reorder_potentials(self, potentials: Array) -> Array:
return potentials[self.tree.sorted_target_ids]
def reorder_potentials(self, potentials: PotentialArray) -> PotentialArray:
def reorder(x: Array) -> Array:
return x[self.tree.sorted_target_ids]

return obj_array.vectorize(reorder, potentials)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this tested anywhere? It seems like it would break when the result is an object array.

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ostensibly not.

def real(x: Array) -> Array:
return x.real

potentials = obj_array.vectorize(real, potentials)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

obj_array_real isn't typed 😢

Comment thread boxtree/tools.py
@alexfikl

Copy link
Copy Markdown
Contributor Author

For what it's worth, this passes on Sphinx 9.1.0: https://github.com/inducer/boxtree/actions/runs/32374671129/job/96443075259?pr=126.

@alexfikl
alexfikl marked this pull request as ready for review August 20, 2026 13:47
@inducer

inducer commented Sep 2, 2026

Copy link
Copy Markdown
Owner

Thx!

@inducer
inducer merged commit 82a86d8 into inducer:main Sep 2, 2026
14 of 16 checks passed
@alexfikl
alexfikl deleted the type-fmm branch September 4, 2026 07:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants