Improved some types

pde/fields/base.py

@@ -1684,7 +1684,7 @@ def _apply_operator(
         )
 
     def make_dot_operator(
-        self, backend: str = "numba", *, conjugate: bool = True
+        self, backend: Literal["numpy", "numba"] = "numba", *, conjugate: bool = True
     ) -> Callable[[np.ndarray, np.ndarray, Optional[np.ndarray]], np.ndarray]:
         """return operator calculating the dot product between two fields
 

pde/fields/vectorial.py

@@ -6,7 +6,17 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Sequence,
+    Union,
+)
 
 import numba as nb
 import numpy as np
@@ -247,7 +257,7 @@ def outer_product(
         return out
 
     def make_outer_prod_operator(
-        self, backend: str = "numba"
+        self, backend: Literal["numpy", "numba"] = "numba"
     ) -> Callable[[np.ndarray, np.ndarray, Optional[np.ndarray]], np.ndarray]:
         """return operator calculating the outer product of two vector fields
 

pde/solvers/base.py

@@ -24,6 +24,7 @@
 from ..tools.math import OnlineStatistics
 from ..tools.misc import classproperty
 from ..tools.numba import is_jitted, jit
+from ..tools.typing import BackendType
 
 
 class SolverBase(metaclass=ABCMeta):
@@ -38,7 +39,7 @@ class SolverBase(metaclass=ABCMeta):
     _subclasses: Dict[str, Type[SolverBase]] = {}
     """dict: dictionary of all inheriting classes"""
 
-    def __init__(self, pde: PDEBase, *, backend: str = "auto"):
+    def __init__(self, pde: PDEBase, *, backend: BackendType = "auto"):
         """
         Args:
             pde (:class:`~pde.pdes.base.PDEBase`):
@@ -140,7 +141,7 @@ def modify_after_step(state_data: np.ndarray) -> float:
         return modify_after_step  # type: ignore
 
     def _make_pde_rhs(
-        self, state: FieldBase, backend: str = "auto"
+        self, state: FieldBase, backend: BackendType = "auto"
     ) -> Callable[[np.ndarray, float], np.ndarray]:
         """obtain a function for evaluating the right hand side
 
@@ -328,7 +329,7 @@ def __init__(
         self,
         pde: PDEBase,
         *,
-        backend: str = "auto",
+        backend: BackendType = "auto",
         adaptive: bool = True,
         tolerance: float = 1e-4,
     ):

pde/solvers/explicit.py

@@ -13,6 +13,7 @@
 from ..pdes.base import PDEBase
 from ..tools.math import OnlineStatistics
 from ..tools.numba import jit
+from ..tools.typing import BackendType
 from .base import AdaptiveSolverBase
 
 
@@ -26,7 +27,7 @@ def __init__(
         pde: PDEBase,
         scheme: Literal["euler", "runge-kutta", "rk", "rk45"] = "euler",
         *,
-        backend: str = "auto",
+        backend: BackendType = "auto",
         adaptive: bool = False,
         tolerance: float = 1e-4,
     ):

pde/solvers/explicit_mpi.py

@@ -14,6 +14,7 @@
 from ..pdes.base import PDEBase
 from ..tools import mpi
 from ..tools.math import OnlineStatistics
+from ..tools.typing import BackendType
 from .explicit import ExplicitSolver
 
 
@@ -80,7 +81,7 @@ def __init__(
         scheme: Literal["euler", "runge-kutta", "rk", "rk45"] = "euler",
         decomposition: Union[int, List[int]] = -1,
         *,
-        backend: str = "auto",
+        backend: BackendType = "auto",
         adaptive: bool = False,
         tolerance: float = 1e-4,
     ):

pde/solvers/implicit.py

@@ -11,6 +11,7 @@
 
 from ..fields.base import FieldBase
 from ..pdes.base import PDEBase
+from ..tools.typing import BackendType
 from .base import SolverBase
 
 
@@ -28,7 +29,7 @@ def __init__(
         pde: PDEBase,
         maxiter: int = 100,
         maxerror: float = 1e-4,
-        backend: str = "auto",
+        backend: BackendType = "auto",
     ):
         """
         Args:

pde/solvers/scipy.py

@@ -10,6 +10,7 @@
 
 from ..fields.base import FieldBase
 from ..pdes.base import PDEBase
+from ..tools.typing import BackendType
 from .base import SolverBase
 
 
@@ -22,7 +23,7 @@ class ScipySolver(SolverBase):
 
     name = "scipy"
 
-    def __init__(self, pde: PDEBase, backend: str = "auto", **kwargs):
+    def __init__(self, pde: PDEBase, backend: BackendType = "auto", **kwargs):
         r"""
         Args:
             pde (:class:`~pde.pdes.base.PDEBase`):

pde/storage/base.py

@@ -54,7 +54,7 @@ class StorageBase(metaclass=ABCMeta):
 
     times: Sequence[float]  # :class:`~numpy.ndarray`): stored time points
     data: Any  # actual data for all the stored times
-    write_mode: str  # mode determining how the storage behaves
+    write_mode: WriteModeType  # mode determining how the storage behaves
 
     def __init__(
         self,

pde/tools/spectral.py

@@ -89,8 +89,6 @@ def noise_normal():
     scaling = 2 * np.pi * scale * k2s ** (exponent / 4)
     scaling.flat[0] = 0
 
-    # TODO: accelerate the FFT using the pyfftw package
-
     def noise_colored() -> np.ndarray:
         """return array of colored noise"""
         # random field

pde/tools/typing.py

@@ -4,7 +4,7 @@
 .. codeauthor:: David Zwicker <[email protected]>
 """
 
-from typing import TYPE_CHECKING, Protocol, Tuple, Union
+from typing import TYPE_CHECKING, Literal, Protocol, Tuple, Union
 
 import numpy as np
 from numpy.typing import ArrayLike  # @UnusedImport
@@ -16,6 +16,7 @@
 Number = Union[Real, complex]
 NumberOrArray = Union[Number, np.ndarray]
 FloatNumerical = Union[float, np.ndarray]
+BackendType = Literal["auto", "numpy", "numba"]
 
 
 class OperatorType(Protocol):