Some typing improvements in examples

Author: inducer

examples/curve-pot.py

@@ -1,5 +1,6 @@
 import numpy as np
 import numpy.linalg as la
+from numpy.typing import NDArray
 
 import pyopencl as cl
 
@@ -120,7 +121,7 @@ def draw_pot_figure(aspect_ratio,
     a = 1
     b = 1/aspect_ratio
 
-    def map_to_curve(t):
+    def map_to_curve(t: NDArray[np.floating]):
         t = t*(2*np.pi)
 
         x = a*np.cos(t)
@@ -185,8 +186,8 @@ def map_to_curve(t):
 
         from sumpy.tools import build_matrix
 
-        def apply_lpot(x):
-            xovsmp = np.dot(fim, x)
+        def apply_lpot(x: NDArray[np.inexact]) -> NDArray[np.inexact]:
+            xovsmp = fim @ x
             _evt, (y,) = lpot(actx.queue,
                     sources,
                     ovsmp_sources,
@@ -197,7 +198,7 @@ def apply_lpot(x):
 
             return actx.to_numpy(y)
 
-        op = LinearOperator((nsrc, nsrc), apply_lpot)
+        op = LinearOperator((nsrc, nsrc), np.dtype(np.complex128), apply_lpot)
         mat = build_matrix(op, dtype=np.complex128)
         w, _v = la.eig(mat)
         plt.plot(w.real, "o-")

examples/curve.py

@@ -1,8 +1,11 @@
+from typing import final
+
 import numpy as np
 import scipy as sp
 import scipy.fftpack
 
 
+@final
 class CurveGrid:
     def __init__(self, x, y):
         self.pos = np.vstack([x, y]).copy()

examples/fourier.py

@@ -1,13 +1,14 @@
 import numpy as np
+from numpy.typing import NDArray
 
 
-def make_fourier_vdm(n, inverse):
+def make_fourier_vdm(n: int, inverse: bool) -> NDArray[np.complex128]:
     i = np.arange(n, dtype=np.float64)
     imat = i[:, np.newaxis]*i/n
     result = np.exp((2j*np.pi)*imat)
 
     if inverse:
-        result = result.T.conj()/n
+        result = np.conj(result.T)/n
     return result
 
 
@@ -30,9 +31,7 @@ def make_fourier_mode_extender(m, n, dtype):
     return result
 
 
-def make_fourier_interp_matrix(m, n):
-    return np.dot(
-            np.dot(
-                make_fourier_vdm(m, inverse=False),
-                make_fourier_mode_extender(m, n, np.float64)),
-            make_fourier_vdm(n, inverse=True))
+def make_fourier_interp_matrix(m: int, n: int):
+    return (make_fourier_vdm(m, inverse=False)
+        @ make_fourier_mode_extender(m, n, np.float64)
+        @  make_fourier_vdm(n, inverse=True))