Skip to content

Refactor backend classes to use generic type parameters and improve type hints across various backends#807

Merged
david-zwicker merged 16 commits into
masterfrom
generics
Apr 18, 2026
Merged

Refactor backend classes to use generic type parameters and improve type hints across various backends#807
david-zwicker merged 16 commits into
masterfrom
generics

Conversation

@david-zwicker
Copy link
Copy Markdown
Member

No description provided.

Copy link
Copy Markdown

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

This PR refactors the backend layer to use a generic native-array type parameter, updating concrete backends (NumPy/JAX/Torch) accordingly, and applies small typing improvements across PDE implementations and utility helpers.

Changes:

  • Introduce BackendBase[TNativeArray] and update backend subclasses to specify their native array type (e.g., jax.Array, torch.Tensor, NumericArray).
  • Improve type hints in several PDEs by explicitly typing operator-construction kwargs dictionaries.
  • Adjust a few numerical/util functions with updated annotations and additional type: ignore suppressions.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 5 comments.

Show a summary per file
File Description
pde/backends/base.py Make backend base class generic and update operator return types to use TNativeArray.
pde/backends/numpy/backend.py Parameterize NumpyBackend with NumericArray and adjust imports/signatures for typing.
pde/backends/jax/backend.py Parameterize JaxBackend with jax.Array and remove some signature-level ignores.
pde/backends/torch/backend.py Parameterize TorchBackend with torch.Tensor and remove some signature-level ignores.
pde/backends/numba/backend.py Remove a # type: ignore on make_pde_rhs signature (typing alignment cleanup).
pde/solvers/base.py Add a suppression on Euler step to accommodate new generic backend typing.
pde/tools/spectral.py Minor numeric annotation tweaks and add suppression on FFT-domain scaling.
pde/tools/output.py Add suppression on tqdm.auto import alias.
pde/pdes/swift_hohenberg.py Type the operator kwargs dict as dict[str, Any].
pde/pdes/kuramoto_sivashinsky.py Type the operator kwargs dict as dict[str, Any].
pde/pdes/kpz_interface.py Type the operator kwargs dict as dict[str, Any].
pde/pdes/cahn_hilliard.py Type the operator kwargs dict as dict[str, Any].

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread pde/tools/output.py
else:
# use the fancier version of the progress bar in jupyter
from tqdm.auto import tqdm as progress_bar_class
from tqdm.auto import tqdm as progress_bar_class # type: ignore
Copy link

Copilot AI Apr 1, 2026

Choose a reason for hiding this comment

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

Avoid using a blanket # type: ignore on the tqdm.auto import. The repo’s mypy config already sets ignore_missing_imports = true for tqdm.* and also enables warn_unused_ignores, so this can become an “unused ignore” warning (and mask real issues). Prefer removing the ignore; if a specific type-checker error remains, suppress it narrowly (e.g., by adjusting typing under TYPE_CHECKING or using a local cast).

Suggested change
from tqdm.auto import tqdm as progress_bar_class # type: ignore
from tqdm.auto import tqdm as progress_bar_class

Copilot uses AI. Check for mistakes.
Comment thread pde/tools/spectral.py
Comment on lines 346 to +348
arr = np_rfftn(arr)

# scale according to frequency
arr *= scaling

arr *= scaling # type: ignore
Copy link

Copilot AI Apr 1, 2026

Choose a reason for hiding this comment

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

Please avoid introducing a blanket # type: ignore for the in-place multiply. With warn_unused_ignores = true, this can easily turn into an “unused ignore” once types are adjusted, and it also hides the underlying typing mismatch. Prefer making the types explicit here (e.g., assign the FFT result to a separately-typed variable / use a cast for the complex FFT result) so the operation type-checks without suppression.

Copilot uses AI. Check for mistakes.
Comment thread pde/solvers/base.py
Comment thread pde/backends/base.py
Comment thread pde/backends/base.py Outdated
- Introduced a unified `_apply_operator` method across backends to streamline operator application.
- Updated backend implementations (NumPy, Numba, JAX, Torch) to utilize the new operator application method.
- Removed the `native` parameter from operator creation methods to simplify backend interface.
- Enhanced data handling to ensure compatibility with various backends, ensuring that data is processed in the native format of the backend.
- Updated documentation to reflect changes in data layout and backend operations.
- Modified tests to utilize the new operator application methods, ensuring consistency across different backends.
- Added a new Torch backend with support for various solvers including Euler and Crank-Nicolson.
- Introduced Torch-specific types and utilities for handling tensor operations.
- Refactored existing solver interfaces to accommodate the new backend.
- Updated the PDE base class to support noise realization in stochastic equations using Torch.
- Enhanced the integration tests to include scenarios for the new Torch backend.
- Temporarily skipped tests for the Numba backend during integration testing.
…cross various backends and PDE implementations
…prove logging, and update tests for adaptive and stochastic solvers
- Updated backend registration methods to use `register_class` instead of instantiating backends directly in the `__init__.py` files for Numba, NumPy, SciPy, and Torch backends.
- Enhanced `BackendRegistry` to manage backend packages, classes, and instantiated backends more effectively, including adding methods for registering classes and backends.
- Modified backend configuration retrieval to ensure consistent access and linking of configurations across different backends.
- Refactored operator registration in SciPy and Torch backends to utilize the backend class directly instead of instantiated backend objects.
- Updated various backend-related tests to reflect changes in backend registration and configuration handling.
- Adjusted conversion methods in backend classes to improve clarity and consistency in data handling between NumPy and native backend formats.
- Updated operator functions in various files (cartesian.py, cylindrical_sym.py, polar_sym.py, spherical_sym.py, common.py) to accept a backend parameter of type NumbaBackend.
- Modified the operator registration to use NumbaBackend directly instead of the previous numba_backend reference.
- Ensured that configuration details are accessed through the backend parameter for consistency and improved flexibility.
@david-zwicker david-zwicker merged commit 5e7c3e0 into master Apr 18, 2026
16 of 17 checks passed
@david-zwicker david-zwicker deleted the generics branch April 18, 2026 16:50
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