add types to new padua point routines

Author: alexfikl

modepy/nodes.py

@@ -388,15 +388,19 @@ def legendre_gauss_lobatto_tensor_product_nodes(dims: int, n: int) -> np.ndarray
 
 # {{{ Padua nodes
 
-def _make_padua_grid_nodes(alpha, beta, order):
+def _make_padua_grid_nodes(
+        alpha: float, beta: float, order: int
+        ) -> Tuple[np.ndarray, np.ndarray]:
     from modepy.quadrature.jacobi_gauss import jacobi_gauss_lobatto_nodes
     mu = jacobi_gauss_lobatto_nodes(alpha, beta, order)
     eta = jacobi_gauss_lobatto_nodes(alpha, beta, order + 1)
 
     return mu, eta
 
 
-def _make_padua_jacobi_nodes(mu, eta, odd_or_even):
+def _make_padua_jacobi_nodes(
+        mu: np.ndarray, eta: np.ndarray, odd_or_even: int
+        ) -> np.ndarray:
     nodes = np.stack(np.meshgrid(mu, eta, indexing="ij"))
     indices = np.sum(
             np.meshgrid(np.arange(mu.size), np.arange(eta.size), indexing="ij"),
@@ -405,30 +409,32 @@ def _make_padua_jacobi_nodes(mu, eta, odd_or_even):
     return nodes[:, indices % 2 == odd_or_even].reshape(2, -1)
 
 
-def _first_padua_jacobi_nodes(alpha, beta, order):
+def _first_padua_jacobi_nodes(alpha: float, beta: float, order: int) -> np.ndarray:
     mu, eta = _make_padua_grid_nodes(alpha, beta, order)
     return _make_padua_jacobi_nodes(mu, eta, 0)
 
 
-def _second_padua_jacobi_nodes(alpha, beta, order):
+def _second_padua_jacobi_nodes(alpha: float, beta: float, order: int) -> np.ndarray:
     # NOTE: these are just "rotated" by pi/2 from the first family
     mu, eta = _make_padua_grid_nodes(alpha, beta, order)
     return _make_padua_jacobi_nodes(eta, mu, 0)
 
 
-def _third_padua_jacobi_nodes(alpha, beta, order):
+def _third_padua_jacobi_nodes(alpha: float, beta: float, order: int) -> np.ndarray:
     # NOTE: these are just "rotated" by pi from the first family
     mu, eta = _make_padua_grid_nodes(alpha, beta, order)
     return _make_padua_jacobi_nodes(mu, eta, 1)
 
 
-def _fourth_padua_jacobi_nodes(alpha, beta, order):
+def _fourth_padua_jacobi_nodes(alpha: float, beta: float, order: int) -> np.ndarray:
     # NOTE: these are just "rotated" by 2 pi/3 from the first family
     mu, eta = _make_padua_grid_nodes(alpha, beta, order)
     return _make_padua_jacobi_nodes(eta, mu, 1)
 
 
-def padua_jacobi_nodes(alpha, beta, order, family="first"):
+def padua_jacobi_nodes(
+        alpha: float, beta: float, order: int,
+        family: str = "first") -> np.ndarray:
     r"""Generalized Padua-Jacobi nodes.
 
     The Padua-Jacobi nodes are constructed from an interlaced grid of
@@ -465,7 +471,7 @@ def padua_jacobi_nodes(alpha, beta, order, family="first"):
     return nodes
 
 
-def padua_nodes(order, family="first"):
+def padua_nodes(order: int, family: str = "first") -> np.ndarray:
     r"""Standard Padua nodes.
 
     Padua nodes are Padua-Jacobi nodes with :math:`\alpha = \beta = -0.5`,