Type point_calculus

Author: inducer

pyproject.toml

@@ -40,10 +40,14 @@ dependencies = [
 
 [dependency-groups]
 dev = [
+    {include-group = "type"},
     {include-group = "doc"},
     {include-group = "test"},
     {include-group = "lint"},
 ]
+type = [
+    "optype>=0.14"
+]
 lint = [
     "pylint",
     # https://github.com/astral-sh/ruff/issues/16943

sumpy/point_calculus.py

@@ -22,6 +22,7 @@
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE.
 """
+from typing import TYPE_CHECKING, Literal, TypeAlias
 
 import numpy as np
 import numpy.linalg as la
@@ -30,13 +31,22 @@
 from pytools import memoize_method
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable, Sequence
+
+    from optype.numpy import Array1D, Array2D
+
+
 __doc__ = """
 .. autoclass:: CalculusPatch
 
 .. autofunction:: frequency_domain_maxwell
 """
 
 
+NodesKind: TypeAlias = Literal["chebyshev", "equispaced", "legendre"]
+
+
 class CalculusPatch:
     """Sets up a grid of points on which derivatives can be calculated. Useful
     to verify that an evaluated potential actually solves a PDE.
@@ -64,7 +74,16 @@ class CalculusPatch:
     .. automethod:: plot_nodes
     .. automethod:: plot
     """
-    def __init__(self, center, h=1e-1, order=4, nodes="chebyshev"):
+    dim: int
+    center: Array1D[np.floating]
+    points: Array2D[np.floating]
+    npoints: int
+
+    def __init__(self,
+            center: Array1D[np.floating],
+            h: float = 1e-1,
+            order: int = 4,
+            nodes: NodesKind = "chebyshev"):
         self.center = center
 
         npoints = order + 1
@@ -119,7 +138,7 @@ def _zero_eval_vec_1d(self):
         # The zeroth coefficient--all others involve x=0.
         return self._vandermonde_1d()[0]
 
-    def basis(self):
+    def basis(self) -> Sequence[Callable[[Array2D[np.floating]], Array1D[np.floating]]]:
         """
         :returns: a :class:`list` containing functions that realize
             a high-order interpolation basis on the :py:attr:`points`.
@@ -129,7 +148,7 @@ def basis(self):
 
         from pytools import indices_in_shape
 
-        def eval_basis(ind, x):
+        def eval_basis(ind: tuple[int, ...], x: Array2D[np.floating]):
             result = 1
             for i in range(self.dim):
                 coord = (x[i] - self.center[i])/(self.h/2)
@@ -172,7 +191,11 @@ def _diff_mat_1d(self, nderivs):
         deriv_coeffs_mat = la.solve(vandermonde.T, n_diff_mat.T).T
         return vandermonde.dot(deriv_coeffs_mat)
 
-    def diff(self, axis, f_values, nderivs=1):
+    def diff(self,
+                axis: int,
+                f_values: Array1D[np.inexact],
+                nderivs: int = 1
+            ) -> Array1D[np.inexact] | Literal[0]:
         """Return the derivative along *axis* of *f_values*.
 
         :arg f_values: an array of shape ``(npoints_total,)``
@@ -197,16 +220,16 @@ def diff(self, axis, f_values, nderivs=1):
                 self._diff_mat_1d(nderivs),
                 f_values.reshape(*self._pshape)).reshape(-1)
 
-    def dx(self, f_values):
+    def dx(self, f_values: Array1D[np.inexact]):
         return self.diff(0, f_values)
 
-    def dy(self, f_values):
+    def dy(self, f_values: Array1D[np.inexact]):
         return self.diff(1, f_values)
 
-    def dz(self, f_values):
+    def dz(self, f_values: Array1D[np.inexact]):
         return self.diff(2, f_values)
 
-    def laplace(self, f_values):
+    def laplace(self, f_values: Array1D[np.inexact]):
         """Return the Laplacian of *f_values*.
 
         :arg f_values: an array of shape ``(npoints_total,)``
@@ -215,18 +238,22 @@ def laplace(self, f_values):
 
         return sum(self.diff(iaxis, f_values, 2) for iaxis in range(self.dim))
 
-    def div(self, arg):
+    def div(self,
+                arg: obj_array.ObjectArray1D[Array1D[np.inexact]]
+            ) -> Array1D[np.inexact] | int:
         r"""
         :arg arg: an object array containing
             :class:`numpy.ndarray`\ s with shape ``(npoints_total,)``.
         """
-        result = 0
+        result: Array1D[np.inexact] | int = 0
         for i, arg_i in enumerate(arg):
             result = result + self.diff(i, arg_i)
 
         return result
 
-    def curl(self, arg):
+    def curl(self,
+                arg: obj_array.ObjectArray1D[Array1D[np.inexact]]
+            ) -> obj_array.ObjectArray1D[Array1D[np.inexact]]:
         r"""Take the curl of the vector quantity *arg*.
 
         :arg arg: an object array of shape ``(3,)`` containing
@@ -254,15 +281,15 @@ def eval_at_center(self, f_values):
         return f_values
 
     @property
-    def x(self):
+    def x(self) -> Array1D[np.floating]:
         return self.points[0]
 
     @property
-    def y(self):
+    def y(self) -> Array1D[np.floating]:
         return self.points[1]
 
     @property
-    def z(self):
+    def z(self) -> Array1D[np.floating]:
         return self.points[2]
 
     def norm(self, arg, p):
@@ -284,15 +311,20 @@ def plot_nodes(self):
             self._points_shaped[1].reshape(-1),
             "o")
 
-    def plot(self, f):
+    def plot(self, f: Array1D[np.floating]):
         f = f.reshape(*self._pshape)
 
         import matplotlib.pyplot as plt
         plt.gca().set_aspect("equal")
         plt.contourf(self._points_1d, self._points_1d, f)
 
 
-def frequency_domain_maxwell(cpatch, e, h, k):
+def frequency_domain_maxwell(
+            cpatch: CalculusPatch,
+            e: obj_array.ObjectArray1D[Array1D[np.complexfloating]],
+            h: obj_array.ObjectArray1D[Array1D[np.complexfloating]],
+            k: complex
+        ):
     mu = 1
     epsilon = 1
     c = 1/np.sqrt(mu*epsilon)