Add type annotations around sumpy.fmm - #300
Conversation
| if self.previous_flag is not None or self.previous_no_cache_kernels is not None: | ||
| raise RuntimeError("cannot reuse the 'CacheMode' context manager") |
There was a problem hiding this comment.
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: ...
|
|
||
| @memoize_method | ||
| def multipole_expansion(self, order: int): | ||
| def multipole_expansion(self, order: int) -> MultipoleExpansionBase: |
There was a problem hiding this comment.
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).
| actx: ArrayContext, | ||
| fft_app: tuple[Any, FFTBackend], | ||
| input_vec: Any, | ||
| fft_app: tuple[lp.TranslationUnit | VkFFTApp, FFTBackend], |
There was a problem hiding this comment.
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?
|
Thanks for sifting through this! |
50db89b to
1eaa222
Compare
There was a problem hiding this comment.
🟡 Changes recommended
A few introduced issues can break runtime behavior (string DTypeLike handling in KernelComputation, PoCL/VkFFT skip condition drift, and use of assert for runtime type validation).
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR adds/updates type annotations across sumpy.fmm-related code paths and adjusts FFT backend selection plumbing (including tests and docs) to align with recent boxtree typing work.
Changes:
- Add richer type annotations and
@overridemarkers insumpy/fmm.pyand related utilities. - Refactor OpenCL FFT helpers to return a backend-specific “app” object directly (Loopy TU vs
VkFFTApp) and make backend selection configurable. - Update FMM tests to pass an
FFTBackendselection explicitly and modernize RNG usage.
File summaries
| File | Description |
|---|---|
| sumpy/toys.py | Updates get_opencl_fft_app calls to use the new keyword-only shape= argument. |
| sumpy/tools.py | Introduces FFT backend enum/API tweaks and adds type annotations across multiple helpers. |
| sumpy/test/test_fmm.py | Updates tests for typed FFT backend selection and RNG handling. |
| sumpy/fmm.py | Adds extensive typing/override annotations and threads FFT backend selection into tree-independent data. |
| sumpy/distributed.py | Minor whitespace cleanup. |
| sumpy/array_context.py | Import ordering adjustment. |
| sumpy/init.py | Adds type annotations and strengthens CacheMode context-manager behavior. |
| pyproject.toml | Adds pyvkfft to the type optional dependency group and tweaks lint config lists. |
| doc/conf.py | Adds pyvkfft intersphinx mapping and related nitpick_ignore entries. |
Review details
- Files reviewed: 9/10 changed files
- Comments generated: 3
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
xref: inducer/boxtree#125