Skip to content

Add type annotations around sumpy.fmm - #300

Open
alexfikl wants to merge 6 commits into
inducer:mainfrom
alexfikl:type-fmm
Open

Add type annotations around sumpy.fmm#300
alexfikl wants to merge 6 commits into
inducer:mainfrom
alexfikl:type-fmm

Conversation

@alexfikl

Copy link
Copy Markdown
Collaborator

Comment thread sumpy/__init__.py
Comment thread sumpy/__init__.py
Comment on lines +125 to +126
if self.previous_flag is not None or self.previous_no_cache_kernels is not None:
raise RuntimeError("cannot reuse the 'CacheMode' context manager")

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

pyright was complaining about setting attributes outside of __init__, so I reworked this a bit. It now yells if it's nested like with ... as cm: with cm: ...

Comment thread sumpy/tools.py
import numpy
from numpy.typing import DTypeLike
from optype.numpy import Array2D
from pyvkfft.opencl import VkFFTApp

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

Is it ok to make this a dependency for typing? It's mostly just used in some return types.

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.

Oh sure! Just make sure it's in whatever uv uses for dev dependencies, so that it's not pulled in at runtime.

Comment thread sumpy/tools.py
Comment thread sumpy/tools.py
Comment thread sumpy/fmm.py

@memoize_method
def multipole_expansion(self, order: int):
def multipole_expansion(self, order: int) -> MultipoleExpansionBase:

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

I'm not sure about returning the base classes here. It makes sense, but they're currently not implemented very nicely for something like this (e.g. pyright complains MultipoleExpansionBase is not callable because it has no __call__ stub).

Comment thread sumpy/tools.py
actx: ArrayContext,
fft_app: tuple[Any, FFTBackend],
input_vec: Any,
fft_app: tuple[lp.TranslationUnit | VkFFTApp, FFTBackend],

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

Not a big fan of this being a tuple (the first and second element could get out of sync?). Any reason to not just do an isinstance?

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.

I agree, that would be nicer.

Comment thread sumpy/fmm.py
Comment on lines +597 to +599
def reorder_potentials(
self, potentials: Array | obj_array.ObjectArray1D[Array]
) -> Array | obj_array.ObjectArray1D[Array]:

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

This should also be added to boxtree. That currently just takes Array.

Comment thread sumpy/fmm.py
Comment thread sumpy/fmm.py

def finalize_potentials(self, actx: ArrayContext, potentials):
@override
def finalize_potentials(self, actx: ArrayContext, potentials: Array) -> Array:

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

This should probably also be a Array | ObjectArray, but needs a little update in boxtree as well.

@alexfikl
alexfikl marked this pull request as ready for review August 19, 2026 18:27
@inducer

inducer commented Aug 20, 2026

Copy link
Copy Markdown
Owner

Thanks for sifting through this!

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