Remove intermediate tools.jacobi interface

Author: kburns

dedalus/core/basis.py

@@ -19,13 +19,11 @@
 from ..tools.cache import CachedAttribute
 from ..tools.cache import CachedMethod
 from ..tools.cache import CachedClass
-from ..tools import jacobi
 from ..tools import clenshaw
 from ..tools.array import reshape_vector, axindex, axslice, interleave_matrices
 from ..tools.dispatch import MultiClass, SkipDispatchException
 from ..tools.general import unify, DeferredTuple
 
-from .spaces import ParityInterval, Disk
 from .coords import Coordinate, CartesianCoordinates, S2Coordinates, SphericalCoordinates, PolarCoordinates, AzimuthalCoordinate
 from .domain import Domain
 from .field  import Operand, LockedField
@@ -478,23 +476,19 @@ def __init__(self, coord, size, bounds, a, b, a0=None, b0=None, dealias=1, libra
         self.b0 = float(b0)
         self.library = library
         self.grid_params = (coord, bounds, a0, b0)
-        self.constant_mode_value = 1 / np.sqrt(jacobi.mass(self.a, self.b))
+        self.constant_mode_value = (1 / np.sqrt(dedalus_sphere.jacobi.mass(self.a, self.b))).astype(np.float64)
 
     def _native_grid(self, scale):
         """Native flat global grid."""
         N, = self.grid_shape((scale,))
-        return jacobi.build_grid(N, a=self.a0, b=self.b0)
+        grid, weights = dedalus_sphere.jacobi.quadrature(N, self.a0, self.b0)
+        return grid.astype(np.float64)
 
     @CachedMethod
     def transform_plan(self, grid_size):
         """Build transform plan."""
         return self.transforms[self.library](grid_size, self.size, self.a, self.b, self.a0, self.b0)
 
-    # def weights(self, scales):
-    #     """Gauss-Jacobi weights."""
-    #     N = self.grid_shape(scales)[0]
-    #     return jacobi.build_weights(N, a=self.a, b=self.b)
-
     # def __str__(self):
     #     space = self.space
     #     cls = self.__class__
@@ -556,14 +550,30 @@ def Jacobi_matrix(self, size):
             size = self.size
         return dedalus_sphere.jacobi.operator('Z')(size, self.a, self.b).square
 
+    @staticmethod
+    def conversion_matrix(N, a0, b0, a1, b1):
+        if not float(a1-a0).is_integer():
+            raise ValueError("a0 and a1 must be integer-separated")
+        if not float(b1-b0).is_integer():
+            raise ValueError("b0 and b1 must be integer-separated")
+        if a0 > a1:
+            raise ValueError("a0 must be less than or equal to a1")
+        if b0 > b1:
+            raise ValueError("b0 must be less than or equal to b1")
+        A = dedalus_sphere.jacobi.operator('A')(+1)
+        B = dedalus_sphere.jacobi.operator('B')(+1)
+        da, db = int(a1-a0), int(b1-b0)
+        conv = A**da @ B**db
+        return conv(N, a0, b0).astype(np.float64)
+
     def ncc_matrix(self, arg_basis, coeffs, cutoff=1e-6):
         """Build NCC matrix via Clenshaw algorithm."""
         if arg_basis is None:
             return super().ncc_matrix(arg_basis, coeffs)
         # Kronecker Clenshaw on argument Jacobi matrix
         elif isinstance(arg_basis, Jacobi):
             N = self.size
-            J = jacobi.jacobi_matrix(N, arg_basis.a, arg_basis.b)
+            J = dedalus_sphere.jacobi.operator('Z')(N, arg_basis.a, arg_basis.b).square.astype(np.float64)
             A, B = clenshaw.jacobi_recursion(N, self.a, self.b, J)
             f0 = self.const * sparse.identity(N)
             total = clenshaw.kronecker_clenshaw(coeffs, A, B, f0, cutoff=cutoff)
@@ -595,7 +605,7 @@ def _last_axis_component_ncc_matrix(cls, subproblem, ncc_basis, arg_basis, out_b
         A, B = clenshaw.jacobi_recursion(Nmat, a_ncc, b_ncc, J)
         f0 = dedalus_sphere.jacobi.polynomials(1, a_ncc, b_ncc, 1)[0].astype(np.float64) * sparse.identity(Nmat)
         matrix = clenshaw.matrix_clenshaw(coeffs.ravel(), A, B, f0, cutoff=cutoff)
-        convert = jacobi.conversion_matrix(Nmat, arg_basis.a, arg_basis.b, out_basis.a, out_basis.b)
+        convert = Jacobi.conversion_matrix(Nmat, arg_basis.a, arg_basis.b, out_basis.a, out_basis.b)
         matrix = convert @ matrix
         return matrix[:N, :N]
 
@@ -647,7 +657,7 @@ def _full_matrix(input_basis, output_basis):
         N = input_basis.size
         a0, b0 = input_basis.a, input_basis.b
         a1, b1 = output_basis.a, output_basis.b
-        matrix = jacobi.conversion_matrix(N, a0, b0, a1, b1)
+        matrix = Jacobi.conversion_matrix(N, a0, b0, a1, b1)
         return matrix.tocsr()
 
 
@@ -686,8 +696,9 @@ def _output_basis(input_basis):
     def _full_matrix(input_basis, output_basis):
         N = input_basis.size
         a, b = input_basis.a, input_basis.b
-        matrix = jacobi.differentiation_matrix(N, a, b) / input_basis.COV.stretch
-        return matrix.tocsr()
+        native_matrix = dedalus_sphere.jacobi.operator('D')(+1)(N, a, b).square.astype(np.float64)
+        problem_matrix = native_matrix / input_basis.COV.stretch
+        return problem_matrix.tocsr()
 
 
 class InterpolateJacobi(operators.Interpolate, operators.SpectralOperator1D):
@@ -709,9 +720,9 @@ def _full_matrix(input_basis, output_basis, position):
         N = input_basis.size
         a, b = input_basis.a, input_basis.b
         x = input_basis.COV.native_coord(position)
-        interp_vector = jacobi.build_polynomials(N, a, b, x)
-        # Return with shape (1, N)
-        return interp_vector[None, :]
+        x = np.array([x])
+        matrix = dedalus_sphere.jacobi.polynomials(N, a, b, x).T
+        return matrix.astype(np.float64)
 
 
 class IntegrateJacobi(operators.Integrate, operators.SpectralOperator1D):
@@ -731,7 +742,7 @@ def _full_matrix(input_basis, output_basis):
         # Build native integration vector
         N = input_basis.size
         a, b = input_basis.a, input_basis.b
-        integ_vector = jacobi.integration_vector(N, a, b)
+        integ_vector = dedalus_sphere.jacobi.polynomial_integrals(N, a, b).astype(np.float64)
         # Rescale and return with shape (1, N)
         return integ_vector[None, :] * input_basis.COV.stretch
 
@@ -753,7 +764,7 @@ def _full_matrix(input_basis, output_basis):
         # Build native integration vector
         N = input_basis.size
         a, b = input_basis.a, input_basis.b
-        integ_vector = jacobi.integration_vector(N, a, b)
+        integ_vector = dedalus_sphere.jacobi.polynomial_integrals(N, a, b).astype(np.float64)
         ave_vector = integ_vector / 2
         # Rescale and return with shape (1, N)
         return ave_vector[None, :]

dedalus/core/spaces.py

@@ -4,7 +4,6 @@
 
 import numpy as np
 
-from ..tools import jacobi
 from ..tools.array import reshape_vector
 from ..tools.cache import CachedMethod, CachedAttribute
 

dedalus/core/transforms.py

@@ -10,7 +10,6 @@
 
 from . import basis
 from ..libraries.fftw import fftw_wrappers as fftw
-from ..tools import jacobi
 from ..tools.array import apply_matrix, apply_dense, axslice, splu_inverse, apply_sparse, prod
 from ..tools.cache import CachedAttribute
 from ..tools.cache import CachedMethod
@@ -118,10 +117,9 @@ def forward_matrix(self):
         a, a0 = self.a, self.a0
         b, b0 = self.b, self.b0
         # Gauss quadrature with base (a0, b0) polynomials
-        base_grid = jacobi.build_grid(N, a=a0, b=b0)
-        base_polynomials = jacobi.build_polynomials(max(M, N), a0, b0, base_grid)
-        base_weights = jacobi.build_weights(N, a=a0, b=b0)
-        base_transform = (base_polynomials * base_weights)
+        base_grid, base_weights = dedalus_sphere.jacobi.quadrature(N, a0, b0)
+        base_polynomials = dedalus_sphere.jacobi.polynomials(max(M, N), a0, b0, base_grid)
+        base_transform = (base_polynomials * base_weights).astype(np.float64)
         # Zero higher coefficients for transforms with grid_size < coeff_size
         base_transform[N:, :] = 0
         if DEALIAS_BEFORE_CONVERTING():
@@ -131,7 +129,7 @@ def forward_matrix(self):
         if (a == a0) and (b == b0):
             forward_matrix = base_transform
         else:
-            conversion = jacobi.conversion_matrix(base_transform.shape[0], a0, b0, a, b)
+            conversion = basis.Jacobi.conversion_matrix(base_transform.shape[0], a0, b0, a, b)
             forward_matrix = conversion @ base_transform
         if not DEALIAS_BEFORE_CONVERTING():
             # Truncate to specified coeff_size
@@ -146,8 +144,8 @@ def backward_matrix(self):
         a, a0 = self.a, self.a0
         b, b0 = self.b, self.b0
         # Construct polynomials on the base grid
-        base_grid = jacobi.build_grid(N, a=a0, b=b0)
-        polynomials = jacobi.build_polynomials(M, a, b, base_grid)
+        base_grid, base_weights = dedalus_sphere.jacobi.quadrature(N, a0, b0)
+        polynomials = dedalus_sphere.jacobi.polynomials(M, a, b, base_grid).astype(np.float64)
         # Zero higher polynomials for transforms with grid_size < coeff_size
         polynomials[N:, :] = 0
         # Transpose and ensure C ordering for fast dot products
@@ -817,12 +815,12 @@ def __init__(self, grid_size, coeff_size, a, b, a0, b0):
         else:
             # Conversion matrices
             if DEALIAS_BEFORE_CONVERTING() and (self.M_orig < self.N): # truncate prior to conversion matrix
-                self.forward_conversion = jacobi.conversion_matrix(self.M_orig, a0, b0, a, b)
+                self.forward_conversion = basis.Jacobi.conversion_matrix(self.M_orig, a0, b0, a, b)
             else: # input to conversion matrix not truncated
-                self.forward_conversion = jacobi.conversion_matrix(self.N, a0, b0, a, b)
+                self.forward_conversion = basis.Jacobi.conversion_matrix(self.N, a0, b0, a, b)
                 self.forward_conversion.resize(self.M_orig, self.N)
                 self.forward_conversion = self.forward_conversion.tocsr()
-            self.backward_conversion = jacobi.conversion_matrix(self.M_orig, a0, b0, a, b)
+            self.backward_conversion = basis.Jacobi.conversion_matrix(self.M_orig, a0, b0, a, b)
             self.backward_conversion = splu_inverse(self.backward_conversion)
             self.resize_rescale_forward = self._resize_rescale_forward_convert
             self.resize_rescale_backward = self._resize_rescale_backward_convert

dedalus/libraries/dedalus_sphere/jacobi.py

@@ -113,6 +113,42 @@ def polynomials(n, a, b, z, init=None, Newton=False, normalized=True, dtype=None
     return P[:n]
 
 
+def polynomial_integrals(n, a, b, dtype=None, **kw):
+    """
+    Definite integrals of the Jacobi polynomials, evaluated with Gauss-Legendre quadrature.
+
+    Parameters
+    ----------
+    n : int
+        Number of polynomials to compute (max degree + 1).
+    a, b : float
+        Jacobi parameters.
+    dtype : dtype, optional
+        Data type. Default: module-level DEFAULT_GRID_DTYPE.
+    **kw : dict, optional
+        Other keywords passed to jacobi.polynomials.
+
+    Returns
+    -------
+    integrals : array
+        Vector of polynomial integrals.
+    """
+    if dtype is None:
+        dtype = DEFAULT_GRID_DTYPE
+    # Build Legendre quadrature
+    grid, weights = quadrature(n, a=0, b=0, dtype=dtype)
+    # Evaluate polynomials on Legendre grid
+    polys = polynomials(n, a, b, grid, dtype=dtype, **kw)
+    # Compute integrals using Legendre quadrature
+    integrals = weights @ polys.T
+    # Eliminate known zeros
+    if a == b == 0:
+        integrals[1:] = 0
+    elif a == b:
+        integrals[1::2] = 0
+    return integrals
+
+
 def quadrature(n, a, b, iterations=3, probability=False, dtype=None):
     """
     Gauss-Jacobi quadrature grid z and weights w.

dedalus/tests/test_clenshaw.py

@@ -4,7 +4,6 @@
 from scipy import sparse
 from dedalus.core import coords, distributor, basis, field, operators
 from dedalus.tools.array import apply_matrix
-from dedalus.tools import jacobi
 from dedalus.tools import clenshaw
 from ..libraries import dedalus_sphere
 

dedalus/tools/clenshaw.py

@@ -3,7 +3,7 @@
 import numpy as np
 from scipy import sparse
 
-from . import jacobi
+from ..libraries import dedalus_sphere
 from .general import DeferredTuple
 
 
@@ -78,7 +78,7 @@ def jacobi_recursion(N, a, b, X):
         B[n] = - J[n,n-1]/J[n,n+1]
     """
     # Jacobi matrix
-    J = jacobi.jacobi_matrix(N, a, b)
+    J = dedalus_sphere.jacobi.operator('Z')(N, a, b).square.astype(np.float64)
     JA = J.A
     # Identity element
     if np.isscalar(X):