From 647d1481e5d70738d6a54ca7106431e490e9eba3 Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Sat, 2 Oct 2021 19:30:28 -0500 Subject: [PATCH] add types to new padua point routines --- modepy/nodes.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/modepy/nodes.py b/modepy/nodes.py index a7cdc13..a03159e 100644 --- a/modepy/nodes.py +++ b/modepy/nodes.py @@ -409,7 +409,9 @@ 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) @@ -417,7 +419,9 @@ def _make_padua_grid_nodes(alpha, beta, order): 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"), @@ -426,30 +430,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 @@ -486,7 +492,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`,