Refactor backend classes to use generic type parameters and improve type hints across various backends#807
Conversation
…ype hints across various backends
There was a problem hiding this comment.
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: ignoresuppressions.
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.
| 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 |
There was a problem hiding this comment.
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).
| from tqdm.auto import tqdm as progress_bar_class # type: ignore | |
| from tqdm.auto import tqdm as progress_bar_class |
| arr = np_rfftn(arr) | ||
|
|
||
| # scale according to frequency | ||
| arr *= scaling | ||
|
|
||
| arr *= scaling # type: ignore |
There was a problem hiding this comment.
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.
- 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
…ve mps device support in tests
…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.
…d update solver compatibility in tests
No description provided.