diff --git a/dedalus/core/basis.py b/dedalus/core/basis.py index f669f4cd..0f2ddd63 100644 --- a/dedalus/core/basis.py +++ b/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,6 +550,22 @@ 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: @@ -563,7 +573,7 @@ def ncc_matrix(self, arg_basis, coeffs, cutoff=1e-6): # 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) @@ -593,9 +603,9 @@ def _last_axis_component_ncc_matrix(cls, subproblem, ncc_basis, arg_basis, out_b Nmat = 3*((N+1)//2) + min((N+1)//2, (da+db+1)//2) J = arg_basis.Jacobi_matrix(size=Nmat) A, B = clenshaw.jacobi_recursion(Nmat, a_ncc, b_ncc, J) - f0 = dedalus_sphere.jacobi.polynomials(1, a_ncc, b_ncc, 1)[0] * sparse.identity(Nmat) + 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, :] @@ -2082,7 +2093,7 @@ def _radius_weights(self, scale): Q0 = dedalus_sphere.jacobi.polynomials(N, self.alpha[0], self.alpha[1], z0) Q_proj = dedalus_sphere.jacobi.polynomials(N, self.alpha[0], self.alpha[1], z_proj) normalization = self.dR/2 - return normalization * ( (Q0 @ weights0).T ) @ (weights_proj*Q_proj) + return (normalization * ( (Q0 @ weights0).T ) @ (weights_proj*Q_proj)).astype(np.float64) def global_radius_weights(self, scale=None): if scale == None: scale = 1 @@ -2101,7 +2112,7 @@ def local_radius_weights(self, scale=None): def constant_mode_value(self): # Note the zeroth mode is constant only for k=0 Q0 = dedalus_sphere.jacobi.polynomials(1, self.alpha[0], self.alpha[1], np.array([0.0])) - return Q0[0,0] + return np.float64(Q0[0,0]) def _new_k(self, k): return AnnulusBasis(self.coordsystem, self.shape, radii = self.radii, k=k, alpha=self.alpha, dealias=self.dealias, dtype=self.dtype, @@ -2179,7 +2190,7 @@ def _interpolation(self, position): a = self.alpha[0] + self.k b = self.alpha[1] + self.k radial_factor = (self.dR/position)**(self.k) - return radial_factor*dedalus_sphere.jacobi.polynomials(self.n_size(0), a, b, native_position) + return radial_factor*dedalus_sphere.jacobi.polynomials(self.n_size(0), a, b, native_position).astype(np.float64) @CachedMethod def operator_matrix(self,op,m,spintotal, size=None): @@ -2231,7 +2242,7 @@ def _last_axis_component_ncc_matrix(cls, subproblem, ncc_basis, arg_basis, out_b Nmat = 3*((N0+1)//2) + ncc_basis.k J = arg_basis.operator_matrix('Z', m, spintotal_arg, size=Nmat) A, B = clenshaw.jacobi_recursion(Nmat, a_ncc, b_ncc, J) - f0 = dedalus_sphere.jacobi.polynomials(1, a_ncc, b_ncc, 1)[0] * sparse.identity(Nmat) + f0 = dedalus_sphere.jacobi.polynomials(1, a_ncc, b_ncc, 1)[0].astype(np.float64) * sparse.identity(Nmat) # Conversions to account for radial prefactors prefactor = arg_basis.jacobi_conversion(m, dk=ncc_basis.k, size=Nmat) if ncc_basis.dtype == np.float64: @@ -2342,7 +2353,7 @@ def local_radius_weights(self, scale=None): @CachedAttribute def constant_mode_value(self): Qk = dedalus_sphere.zernike.polynomials(2, 1, self.alpha+self.k, 0, np.array([0])) - return Qk[0] + return np.float64(Qk[0]) def _new_k(self, k): return DiskBasis(self.coordsystem, self.shape, radius = self.radius, k=k, alpha=self.alpha, dealias=self.dealias, dtype=self.dtype, @@ -2442,7 +2453,7 @@ def operator_matrix(self, op, m, spin, size=None): def interpolation(self, m, spintotal, position): native_position = self.radial_COV.native_coord(position) native_z = 2*native_position**2 - 1 - return dedalus_sphere.zernike.polynomials(2, self.n_size(m), self.alpha + self.k, np.abs(m + spintotal), native_z) + return dedalus_sphere.zernike.polynomials(2, self.n_size(m), self.alpha + self.k, np.abs(m + spintotal), native_z).astype(np.float64) @CachedMethod def radius_multiplication_matrix(self, m, spintotal, order, d): @@ -2480,7 +2491,7 @@ def _last_axis_component_ncc_matrix(cls, subproblem, ncc_basis, arg_basis, out_b J = arg_basis.operator_matrix('Z', m, spintotal_arg) A, B = clenshaw.jacobi_recursion(N, a_ncc, b_ncc, J) # assuming that we're doing ball for now... - f0 = dedalus_sphere.zernike.polynomials(2, 1, a_ncc, b_ncc, 1)[0] * sparse.identity(N) + f0 = dedalus_sphere.zernike.polynomials(2, 1, a_ncc, b_ncc, 1)[0].astype(np.float64) * sparse.identity(N) prefactor = arg_basis.radius_multiplication_matrix(m, spintotal_arg, diff_regtotal, d) if ncc_basis.dtype == np.float64: coeffs_filter = coeffs.ravel()[:2*N] @@ -2926,7 +2937,7 @@ def local_grid_colatitude(self, scale): def _native_colatitude_grid(self, scale): N = int(np.ceil(scale * self.shape[1])) cos_theta, weights = dedalus_sphere.sphere.quadrature(Lmax=N-1) - theta = np.arccos(cos_theta).astype(np.float64) + theta = np.arccos(cos_theta.astype(np.float64)) # TODO: xprec doesn't yet support arccos return theta def global_colatitude_weights(self, scale=None): @@ -2945,7 +2956,7 @@ def local_colatitude_weights(self, scale=None): @CachedMethod def transform_plan(self, Ntheta, s): """Build transform plan.""" - return self.transforms[self.colatitude_library](Ntheta, self.Lmax, self.m_maps, s) + return self.transforms[self.colatitude_library](Ntheta, self.Lmax, self.m_maps, s, self.dtype) def forward_transform_azimuth_Mmax0(self, field, axis, gdata, cdata): slice_axis = axis + len(field.tensorsig) @@ -3670,13 +3681,13 @@ def _radius_weights(self, scale): Q0 = dedalus_sphere.jacobi.polynomials(N, self.alpha[0], self.alpha[1], z0) Q_proj = dedalus_sphere.jacobi.polynomials(N, self.alpha[0], self.alpha[1], z_proj) normalization = self.dR/2 - return normalization * ( (Q0 @ weights0).T ) @ (weights_proj*Q_proj) + return (normalization * ( (Q0 @ weights0).T ) @ (weights_proj*Q_proj)).astype(np.float64) @CachedAttribute def constant_mode_value(self): # Note the zeroth mode is constant only for k=0 Q0 = dedalus_sphere.jacobi.polynomials(1, self.alpha[0], self.alpha[1], np.array([0.0])) - return Q0[0,0] + return np.float64(Q0[0,0]) @CachedMethod def radial_transform_factor(self, scale, data_axis, dk): @@ -3689,7 +3700,7 @@ def interpolation(self, position): a = self.alpha[0] + self.k b = self.alpha[1] + self.k radial_factor = (self.dR/position)**(self.k) - return radial_factor*dedalus_sphere.jacobi.polynomials(self.n_size(0), a, b, native_position) + return radial_factor*dedalus_sphere.jacobi.polynomials(self.n_size(0), a, b, native_position).astype(np.float64) @CachedMethod def transform_plan(self, grid_size, k): @@ -3778,7 +3789,7 @@ def _last_axis_component_ncc_matrix(cls, subproblem, ncc_basis, arg_basis, out_b Nmat = 3*((N0+1)//2) + ncc_basis.k J = arg_radial_basis.operator_matrix('Z', ell, regtotal_arg, size=Nmat) A, B = clenshaw.jacobi_recursion(Nmat, a_ncc, b_ncc, J) - f0 = dedalus_sphere.jacobi.polynomials(1, a_ncc, b_ncc, 1)[0] * sparse.identity(Nmat) + f0 = dedalus_sphere.jacobi.polynomials(1, a_ncc, b_ncc, 1)[0].astype(np.float64) * sparse.identity(Nmat) # Conversions to account for radial prefactors prefactor = arg_radial_basis.jacobi_conversion(ell, dk=ncc_basis.k, size=Nmat) if ncc_basis.dtype == np.float64: @@ -3892,18 +3903,18 @@ def _native_radius_grid(self, scale): def _radius_weights(self, scale): N = int(np.ceil(scale * self.shape[2])) z, weights = dedalus_sphere.zernike.quadrature(3, N, k=self.alpha) - return weights + return weights.astype(np.float64) @CachedAttribute def constant_mode_value(self): Qk = dedalus_sphere.zernike.polynomials(3, 1, self.alpha+self.k, 0, np.array([0])) - return Qk[0] + return np.float64(Qk[0]) @CachedMethod def interpolation(self, ell, regtotal, position): native_position = self.radial_COV.native_coord(position) native_z = 2*native_position**2 - 1 - return dedalus_sphere.zernike.polynomials(3, self.n_size(ell), self.alpha + self.k, ell + regtotal, native_z) + return dedalus_sphere.zernike.polynomials(3, self.n_size(ell), self.alpha + self.k, ell + regtotal, native_z).astype(np.float64) @CachedMethod def transform_plan(self, grid_shape, regindex, axis, regtotal, k, alpha): @@ -4008,7 +4019,7 @@ def _last_axis_component_ncc_matrix(cls, subproblem, ncc_basis, arg_basis, out_b if (d >= 0) and (d % 2 == 0): J = arg_radial_basis.operator_matrix('Z', ell, regtotal_arg, size=Nmat) A, B = clenshaw.jacobi_recursion(N0, a_ncc, b_ncc, J) - f0 = dedalus_sphere.zernike.polynomials(3, 1, a_ncc, regtotal_ncc, 1)[0] * sparse.identity(Nmat) + f0 = dedalus_sphere.zernike.polynomials(3, 1, a_ncc, regtotal_ncc, 1)[0].astype(np.float64) * sparse.identity(Nmat) radial_factor = arg_radial_basis.radius_multiplication_matrix(ell, regtotal_arg, diff_regtotal, d, size=Nmat) conversion = arg_radial_basis.conversion_matrix(ell, regtotal_out, dk, size=Nmat) prefactor = conversion @ radial_factor @@ -5143,7 +5154,7 @@ def _radial_matrix(basis, m): N = basis.shape[1] z0, w0 = dedalus_sphere.zernike.quadrature(2, N, k=0) Qk = dedalus_sphere.zernike.polynomials(2, n_size, basis.alpha+basis.k, abs(m), z0) - matrix = (w0[None, :] @ Qk.T).astype(basis.dtype) + matrix = (w0[None, :] @ Qk.T).astype(np.float64) matrix *= basis.radius**2 matrix *= 2 * np.pi # Fourier contribution else: @@ -5165,8 +5176,8 @@ def _radial_matrix(basis, m): z0, w0 = dedalus_sphere.jacobi.quadrature(N, a=0, b=0) r0 = basis.dR / 2 * (z0 + basis.rho) Qk = dedalus_sphere.jacobi.polynomials(n_size, basis.alpha[0]+basis.k, basis.alpha[1]+basis.k, z0) - w0_geom = r0 * w0 * (r0 / basis.dR)**(-basis.k) - matrix = (w0_geom[None, :] @ Qk.T).astype(basis.dtype) + w0_geom = r0 * w0 * (r0.astype(np.float64) / basis.dR)**(-basis.k) # TODO: xprec does not yet support power + matrix = (w0_geom[None, :] @ Qk.T).astype(np.float64) matrix *= basis.dR / 2 matrix *= 2 * np.pi # Fourier contribution else: @@ -5232,7 +5243,7 @@ def _radial_matrix(basis, ell): N = basis.shape[2] z0, w0 = dedalus_sphere.zernike.quadrature(3, N, k=0) Qk = dedalus_sphere.zernike.polynomials(3, n_size, basis.alpha+basis.k, ell, z0) - matrix = (w0[None, :] @ Qk.T).astype(basis.dtype) + matrix = (w0[None, :] @ Qk.T).astype(np.float64) matrix *= basis.radius**3 matrix *= 4 * np.pi / np.sqrt(2) # SWSH contribution else: @@ -5254,8 +5265,8 @@ def _radial_matrix(basis, ell): z0, w0 = dedalus_sphere.jacobi.quadrature(N, a=0, b=0) r0 = basis.dR / 2 * (z0 + basis.rho) Qk = dedalus_sphere.jacobi.polynomials(n_size, basis.alpha[0]+basis.k, basis.alpha[1]+basis.k, z0) - w0_geom = r0**2 * w0 * (r0 / basis.dR)**(-basis.k) - matrix = (w0_geom[None, :] @ Qk.T).astype(basis.dtype) + w0_geom = r0**2 * w0 * (r0.astype(np.float64) / basis.dR)**(-basis.k) # TODO: xprec does not yet support power + matrix = (w0_geom[None, :] @ Qk.T).astype(np.float64) matrix *= basis.dR / 2 matrix *= 4 * np.pi / np.sqrt(2) # SWSH contribution else: @@ -5387,7 +5398,7 @@ def _interpolation_vectors(sphere_basis, Ntheta, s, theta): Lmin = max(abs(m), abs(s)) interp_m = dedalus_sphere.sphere.harmonics(sphere_basis.Lmax, m, s, z)[None, :] forward_m = forward._forward_SWSH_matrices[m][Lmin-abs(m):] - interp_vectors[m] = interp_m @ forward_m + interp_vectors[m] = interp_m.astype(np.float64) @ forward_m else: interp_vectors[m] = np.zeros((1, Ntheta)) return interp_vectors diff --git a/dedalus/core/operators.py b/dedalus/core/operators.py index 43b8edd2..fdfd1be4 100644 --- a/dedalus/core/operators.py +++ b/dedalus/core/operators.py @@ -2858,14 +2858,14 @@ def radial_matrix(self, spinindex_in, spinindex_out, m): radial_basis = self.input_basis spintotal = radial_basis.spintotal(spinindex_in) if spinindex_out in self.spinindex_out(spinindex_in): - return self._radial_matrix(radial_basis.Lmax, spintotal, m, self.dtype) + return self._radial_matrix(radial_basis.Lmax, spintotal, m) else: raise ValueError("This should never happen") @staticmethod @CachedMethod - def _radial_matrix(Lmax, spintotal, m, dtype): - matrix = dedalus_sphere.sphere.operator('Cos', dtype)(Lmax, m, spintotal).square + def _radial_matrix(Lmax, spintotal, m): + matrix = dedalus_sphere.sphere.operator('Cos')(Lmax, m, spintotal).square # Pad to include invalid ells trunc = abs(spintotal) - abs(m) if trunc > 0: diff --git a/dedalus/core/spaces.py b/dedalus/core/spaces.py index 45794066..a34b27b8 100644 --- a/dedalus/core/spaces.py +++ b/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 diff --git a/dedalus/core/transforms.py b/dedalus/core/transforms.py index a41d5b8c..af1c3363 100644 --- a/dedalus/core/transforms.py +++ b/dedalus/core/transforms.py @@ -6,12 +6,12 @@ import scipy import scipy.fft import scipy.fftpack -from ..libraries import dedalus_sphere from . import basis +from ..libraries import dedalus_sphere 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.general import is_real_dtype, is_complex_dtype from ..tools.cache import CachedAttribute from ..tools.cache import CachedMethod @@ -118,10 +118,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 +130,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 +145,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 +816,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 @@ -1234,44 +1233,62 @@ def backward(self, cdata, gdata, axis): self.backward_reduced(cdata, gdata) -@register_transform(basis.SphereBasis, 'matrix') -class SWSHColatitudeTransform(NonSeparableTransform): +class AssociatedLegendreTransform(NonSeparableTransform): + """ + Spin generalizations of the associated Legendre transform for transforming + spin-weighted spherical harmonics in colatitude. + """ - def __init__(self, Ntheta, Lmax, m_maps, s): + def __init__(self, Ntheta, Lmax, m_maps, s, dtype): self.Ntheta = Ntheta self.Lmax = Lmax self.m_maps = m_maps self.s = s + self.dtype = dtype def forward_reduced(self, gdata, cdata): # local_m = self.local_m # if gdata.shape[1] != len(local_m): # do we want to do this check??? # raise ValueError("gdata.shape[1]: %i, len(local_m): %i" %(gdata.shape[1], len(local_m))) - m_matrices = self._forward_SWSH_matrices - Lmax = self.Lmax for m, mg_slice, mc_slice, ell_slice in self.m_maps: # Skip transforms when |m| > Lmax - if abs(m) <= Lmax: + if abs(m) <= self.Lmax: # Use rectangular transform matrix, padded with zeros when Lmin > abs(m) - grm = gdata[:, mg_slice, :, :] - crm = cdata[:, mc_slice, ell_slice, :] - apply_matrix(m_matrices[m], grm, axis=2, out=crm) + gdata_reduced_m = gdata[:, mg_slice, :, :] + cdata_reduced_m = cdata[:, mc_slice, ell_slice, :] + self.forward_reduced_m(m, gdata_reduced_m, cdata_reduced_m) def backward_reduced(self, cdata, gdata): # local_m = self.local_m # if gdata.shape[1] != len(local_m): # do we want to do this check??? # raise ValueError("gdata.shape[1]: %i, len(local_m): %i" %(gdata.shape[1], len(local_m))) - m_matrices = self._backward_SWSH_matrices - Lmax = self.Lmax for m, mg_slice, mc_slice, ell_slice in self.m_maps: - if abs(m) > Lmax: + if abs(m) > self.Lmax: # Write zeros because they'll be used by the inverse azimuthal transform gdata[:, mg_slice, :, :] = 0 else: # Use rectangular transform matrix, padded with zeros when Lmin > abs(m) - grm = gdata[:, mg_slice, :, :] - crm = cdata[:, mc_slice, ell_slice, :] - apply_matrix(m_matrices[m], crm, axis=2, out=grm) + gdata_reduced_m = gdata[:, mg_slice, :, :] + cdata_reduced_m = cdata[:, mc_slice, ell_slice, :] + self.backward_reduced_m(m, cdata_reduced_m, gdata_reduced_m) + + def forward_reduced_m(self, m, gdata_reduced_m, cdata_reduced_m): + raise NotImplementedError("Subclasses must implement.") + + def backward_reduced_m(self, m, cdata_reduced_m, gdata_reduced_m): + raise NotImplementedError("Subclasses must implement.") + + +@register_transform(basis.SphereBasis, 'matrix') +class AssociatedLegendreMMT(AssociatedLegendreTransform): + + def forward_reduced_m(self, m, gdata_reduced_m, cdata_reduced_m): + m_matrix = self._forward_SWSH_matrices[m] + apply_matrix(m_matrix, gdata_reduced_m, axis=2, out=cdata_reduced_m) + + def backward_reduced_m(self, m, cdata_reduced_m, gdata_reduced_m): + m_matrix = self._backward_SWSH_matrices[m] + apply_matrix(m_matrix, cdata_reduced_m, axis=2, out=gdata_reduced_m) @CachedAttribute def _quadrature(self): @@ -1325,6 +1342,189 @@ def _backward_SWSH_matrices(self): m_matrices[m] = np.asarray(Yfull, order='C') return m_matrices + +@register_transform(basis.SphereBasis, 'shtns') +class AssociatedLegendreSHTNS(AssociatedLegendreTransform): + """SHTNS-based associated Legendre transform.""" + + def __init__(self, *args, **kw): + super().__init__(*args, **kw) + if self.s != 0: + raise ValueError("SHTNS does not support spin-weighted transforms.") + # Setup transform based on Ntheta + import shtns + self.shtns_obj = shtns.sht(lmax=self.Lmax, nthreads=1) + self.shtns_obj.set_grid(nlat=self.Ntheta) + self.is_complex = is_complex_dtype(self.dtype) + + def forward_reduced_m(self, m, grm, crm): + shtns_obj = self.shtns_obj + # Make contiguous along theta, m + gdata = np.ascontiguousarray(grm.transpose((0, 3, 2, 1))) + # Complexify along m axis if real + gdata.dtype = np.complex128 + # Drop singleton m axis and combine leading axes + gdata.shape = (-1, grm.shape[2]) + # Setup contiguous cdata + cdata = np.zeros((gdata.shape[0], crm.shape[2]), dtype=np.complex128) + # Call transforms on contiguous vectors + if self.is_complex and m == 0: + # m = 0 transforms only do real part apparently + gtemp = np.empty_like(gdata[0]) + ctemp = np.empty_like(cdata[0]) + for i in range(gdata.shape[0]): + gtemp[:] = gdata[i].real + shtns_obj.spat_to_SH_m(gtemp, ctemp, m) + cdata[i].real = ctemp.real + gtemp[:] = gdata[i].imag + shtns_obj.spat_to_SH_m(gtemp, ctemp, m) + cdata[i].imag = ctemp.real + else: + for i in range(gdata.shape[0]): + shtns_obj.spat_to_SH_m(gdata[i], cdata[i], m) + # Rescale cdata + mod = 1 / np.sqrt(2 * np.pi) + if m < 0: + mod *= (-1) ** (m%2) + cdata *= mod + cdata[:, 1::2] *= -1 + # Unpack cdata + cdata.shape = (crm.shape[0], crm.shape[3], crm.shape[2], 1) + cdata.dtype = crm.dtype + crm[:] = cdata.transpose((0, 3, 2, 1)) + + def backward_reduced_m(self, m, crm, grm): + shtns_obj = self.shtns_obj + # Make contiguous along theta, m + cdata = np.ascontiguousarray(crm.transpose((0, 3, 2, 1))) + # Complexify along m axis if real + cdata.dtype = np.complex128 + # Drop singleton m axis and combine leading axes + cdata.shape = (-1, crm.shape[2]) + # Rescale cdata + mod = np.sqrt(2 * np.pi) + if m < 0: + mod *= (-1) ** (m%2) + cdata *= mod + cdata[:, 1::2] *= -1 + # Setup contiguous gdata + gdata = np.zeros((cdata.shape[0], grm.shape[2]), dtype=np.complex128) + # Call transforms on contiguous vectors + if self.is_complex and m == 0: + # m = 0 transforms only do real part apparently + gtemp = np.empty_like(gdata[0]) + ctemp = np.empty_like(cdata[0]) + for i in range(gdata.shape[0]): + ctemp[:] = cdata[i].real + shtns_obj.SH_to_spat_m(ctemp, gtemp, m) + gdata[i].real = gtemp.real + ctemp[:] = cdata[i].imag + shtns_obj.SH_to_spat_m(ctemp, gtemp, m) + gdata[i].imag = gtemp.real + else: + for i in range(gdata.shape[0]): + shtns_obj.SH_to_spat_m(cdata[i], gdata[i], m) + # Unpack gdata + gdata.shape = (grm.shape[0], grm.shape[3], grm.shape[2], 1) + gdata.dtype = grm.dtype + grm[:] = gdata.transpose((0, 3, 2, 1)) + + +@register_transform(basis.SphereBasis, 'ducc') +class AssociatedLegendreDUCC(AssociatedLegendreTransform): + """DUCC-based associated Legendre transform.""" + + def __init__(self, *args, **kw): + super().__init__(*args, **kw) + # Setup grid + import ducc0 + self.theta = ducc0.sht.experimental.GL_thetas(self.Ntheta) + self.ducc_forward = ducc0.sht.experimental.leg2alm + self.ducc_backward = ducc0.sht.experimental.alm2leg + + def forward_reduced_m(self, m, grm, crm): + # Make contiguous along theta, m + gdata = np.ascontiguousarray(grm.transpose((0, 3, 2, 1))) + # Complexify along m axis if real + gdata.dtype = np.complex128 + # Drop singleton m and combine other axes + gdata.shape = (-1, grm.shape[2]) + # Repack according to DUCC layout + ncomp = 1 + (self.s != 0) + if self.s == 0: + gdata.shape = (-1, 1, grm.shape[2], 1) + else: + gdata_padded = np.zeros((gdata.shape[0], ncomp, grm.shape[2], 1), dtype=np.complex128) + if self.s > 0: + gdata_padded[:, 0, :, 0] = gdata + else: + gdata_padded[:, 1, :, 0] = gdata + gdata = gdata_padded + # Setup contiguous cdata with extra padding + cdata = np.zeros((gdata.shape[0], ncomp, self.Lmax + 1), dtype=np.complex128) + # Call transforms on contiguous vectors + mval = np.array([m]) + mstart = np.array([0]) + for i in range(gdata.shape[0]): + self.ducc_forward(leg=gdata[i], lmax=self.Lmax, theta=self.theta, spin=abs(self.s), mval=mval, mstart=mstart, alm=cdata[i]) + # Drop extra padding + if self.s >= 0: + cdata = cdata[:, 0, abs(m):] + else: + cdata = cdata[:, 1, abs(m):] + # Rescale cdata + mod = 1 / np.sqrt(2 * np.pi) + if m < 0: + mod *= (-1) ** (m%2) + cdata *= mod + cdata[:, 1::2] *= -1 + # Unpack cdata + cdata.shape = (crm.shape[0], crm.shape[3], crm.shape[2], 1) + cdata.dtype = crm.dtype + crm[:] = cdata.transpose((0, 3, 2, 1)) + + def backward_reduced_m(self, m, crm, grm): + # Make contiguous along theta, m + cdata = np.ascontiguousarray(crm.transpose((0, 3, 2, 1))) + # Complexify along m axis if real + cdata.dtype = np.complex128 + # Drop singleton m axis and combine leading axes + cdata.shape = (-1, crm.shape[2]) + # Rescale cdata + mod = np.sqrt(2 * np.pi) + if m < 0: + mod *= (-1) ** (m%2) + cdata *= mod + cdata[:, 1::2] *= -1 + # Repack according to DUCC layout + ncomp = 1 + (self.s != 0) + if self.s == 0 and m == 0: + cdata.shape = (-1, 1, crm.shape[2]) + else: + cdata_padded = np.zeros((cdata.shape[0], ncomp, self.Lmax+1), dtype=np.complex128) + if self.s >= 0: + cdata_padded[:, 0, abs(m):] = cdata + else: + cdata_padded[:, 1, abs(m):] = cdata + cdata = cdata_padded + # Setup contiguous gdata with extra padding + gdata = np.zeros((cdata.shape[0], ncomp, grm.shape[2], 1), dtype=np.complex128) + # Call DUCC transforms + mval = np.array([m]) + mstart = np.array([0]) + for i in range(gdata.shape[0]): + self.ducc_backward(alm=cdata[i], lmax=self.Lmax, theta=self.theta, spin=abs(self.s), mval=mval, mstart=mstart, leg=gdata[i]) + # Drop extra padding + if self.s >= 0: + gdata = gdata[:, 0, :, 0] + else: + gdata = gdata[:, 1, :, 0] + # Unpack gdata + gdata.shape = (grm.shape[0], grm.shape[3], grm.shape[2], 1) + gdata.dtype = grm.dtype + grm[:] = gdata.transpose((0, 3, 2, 1)) + + @register_transform(basis.DiskBasis, 'matrix') class DiskRadialTransform(NonSeparableTransform): """ @@ -1401,7 +1601,7 @@ def _forward_matrices(self): # Spectral conversion if self.k > 0: conversion = dedalus_sphere.zernike.operator(2, 'E')(+1)**self.k - W = conversion(W.shape[0], self.alpha, abs(m + self.s)) @ W + W = conversion(W.shape[0], self.alpha, abs(m + self.s)) @ W.astype(np.float64) if not DEALIAS_BEFORE_CONVERTING(): # Truncate to specified coeff_size W = W[:max(self.N2c-Nmin,0)] @@ -1510,7 +1710,7 @@ def _forward_GSZP_matrix(self): # Spectral conversion if self.k > 0: conversion = dedalus_sphere.zernike.operator(3, 'E')(+1)**self.k - W = conversion(W.shape[0], self.alpha, ell+self.regtotal) @ W + W = conversion(W.shape[0], self.alpha, ell+self.regtotal) @ W.astype(np.float64) if not DEALIAS_BEFORE_CONVERTING(): # Truncate to specified coeff_size W = W[:max(self.N3c-Nmin,0)] diff --git a/dedalus/libraries/dedalus_sphere/annulus.py b/dedalus/libraries/dedalus_sphere/annulus.py index e5956fc6..e43d1972 100644 --- a/dedalus/libraries/dedalus_sphere/annulus.py +++ b/dedalus/libraries/dedalus_sphere/annulus.py @@ -1,51 +1,62 @@ -import numpy as np +import numpy as np from . import jacobi # The defalut configurations for the base Jacobi parameters. -alpha = (-0.5,-0.5) +alpha = (-0.5, -0.5) -def quadrature(Nmax,alpha=alpha,**kw): - return jacobi.quadrature(Nmax,alpha[0],alpha[1],**kw) -def trial_functions(Nmax,z,alpha=alpha): +def quadrature(Nmax, alpha=alpha, **kw): + return jacobi.quadrature(Nmax, alpha[0], alpha[1], **kw) - init = 1/np.sqrt(jacobi.mass(alpha[0],alpha[1])) + 0*z - return jacobi.recursion(Nmax,alpha[0],alpha[1],z,init) -def operator(dimension,op,Nmax,k,ell,radii,pad=0,alpha=alpha): +def trial_functions(Nmax, z, alpha=alpha): + + init = 1 / np.sqrt(jacobi.mass(alpha[0], alpha[1])) + 0 * z + return jacobi.recursion(Nmax, alpha[0], alpha[1], z, init) + + +def operator(dimension, op, Nmax, k, ell, radii, pad=0, alpha=alpha): # Pad the matrices by a safe amount before outputting the correct size. - - if radii[1] <= radii[0]: raise ValueError('Inner radius must be greater than outer radius.') - - gapwidth = radii[1] - radii[0] - aspectratio = (radii[1] + radii[0])/gapwidth - - a, b, N = k+alpha[0], k+alpha[1], Nmax + pad - + + if radii[1] <= radii[0]: + raise ValueError('Inner radius must be greater than outer radius.') + + gapwidth = radii[1] - radii[0] + aspectratio = (radii[1] + radii[0]) / gapwidth + + a, b, N = k + alpha[0], k + alpha[1], Nmax + pad + # zeros - if (op == '0'): return jacobi.operator('0',N,a,b) - + if op == '0': + return jacobi.operator('0', N, a, b) + # identity - if (op == 'I'): return jacobi.operator('I',N,a,b) - - Z = aspectratio*jacobi.operator('I',N+2,a,b) - Z += jacobi.operator('J',N+2,a,b) - + if op == 'I': + return jacobi.operator('I', N, a, b) + + Z = aspectratio * jacobi.operator('I', N + 2, a, b) + Z += jacobi.operator('J', N + 2, a, b) + # r multiplication - if op == 'R': return (gapwidth/2)*Z[:N+1,:N+1] - - E = jacobi.operator('A+',N+2,a,b+1) @ jacobi.operator('B+',N+2,a,b) - + if op == 'R': + return (gapwidth / 2) * Z[: N + 1, : N + 1] + + E = jacobi.operator('A+', N + 2, a, b + 1) @ jacobi.operator('B+', N + 2, a, b) + # conversion - if op == 'E': return 0.5*(E @ Z)[:N+1,:N+1] - - D = jacobi.operator('D+',N+2,a,b) @ Z - + if op == 'E': + return 0.5 * (E @ Z)[: N + 1, : N + 1] + + D = jacobi.operator('D+', N + 2, a, b) @ Z + # derivatives - if op == 'D+': return (D - (ell+k+1 )*E)[:N+1,:N+1]/gapwidth - if op == 'D-': return (D + (ell-k+dimension-3)*E)[:N+1,:N+1]/gapwidth - - # restriction - if op == 'r=Ri': return jacobi.operator('z=-1',N,a,b) - if op == 'r=Ro': return jacobi.operator('z=+1',N,a,b) + if op == 'D+': + return (D - (ell + k + 1) * E)[: N + 1, : N + 1] / gapwidth + if op == 'D-': + return (D + (ell - k + dimension - 3) * E)[: N + 1, : N + 1] / gapwidth + # restriction + if op == 'r=Ri': + return jacobi.operator('z=-1', N, a, b) + if op == 'r=Ro': + return jacobi.operator('z=+1', N, a, b) diff --git a/dedalus/libraries/dedalus_sphere/ball_wrapper.py b/dedalus/libraries/dedalus_sphere/ball_wrapper.py index 1d866d5b..58f78e88 100644 --- a/dedalus/libraries/dedalus_sphere/ball_wrapper.py +++ b/dedalus/libraries/dedalus_sphere/ball_wrapper.py @@ -1,5 +1,4 @@ - -import numpy as np +import numpy as np from scipy.sparse import linalg as spla from . import sphere_wrapper as sph from . import ball128 as ball @@ -10,63 +9,72 @@ import time from dedalus.tools.config import config + STORE_LU_TRANSFORM = config['transforms'].getboolean('STORE_LU_TRANSFORM') import logging + logger = logging.getLogger(__name__) barrier = False timing = False + class Ball: - def __init__(self,N_max,L_max,R_max=0,a=0,N_r=None,N_theta=None,ell_min=None,ell_max=None,m_min=None,m_max=None): - self.N_max, self.L_max, self.R_max = N_max, L_max, R_max - if N_r == None: self.N_r = self.N_max + 1 - else: self.N_r = N_r + def __init__(self, N_max, L_max, R_max=0, a=0, N_r=None, N_theta=None, ell_min=None, ell_max=None, m_min=None, m_max=None): + self.N_max, self.L_max, self.R_max = N_max, L_max, R_max + if N_r == None: + self.N_r = self.N_max + 1 + else: + self.N_r = N_r self.a = a - if ell_min == None: ell_min = 0 - if ell_max == None: ell_max = L_max - if m_min == None: m_min = 0 - if m_max == None: m_max = L_max + if ell_min == None: + ell_min = 0 + if ell_max == None: + ell_max = L_max + if m_min == None: + m_min = 0 + if m_max == None: + m_max = L_max self.ell_min, self.ell_max = ell_min, ell_max - self.m_min, self.m_max = m_min, m_max + self.m_min, self.m_max = m_min, m_max # Spherical Harmonic Transforms - self.S = sph.Sphere(self.L_max,S_max=self.R_max,N_theta=N_theta,m_min=m_min,m_max=m_max) + self.S = sph.Sphere(self.L_max, S_max=self.R_max, N_theta=N_theta, m_min=m_min, m_max=m_max) - self.theta = self.S.grid + self.theta = self.S.grid self.cos_theta = self.S.cos_grid self.sin_theta = self.S.sin_grid # grid and weights for the radial transforms - z_projection, weights_projection = ball.quadrature(self.N_r-1,niter=3,a=a,report_error=False) + z_projection, weights_projection = ball.quadrature(self.N_r - 1, niter=3, a=a, report_error=False) # grid and weights for radial integral using volume measure - z0, weights0 = ball.quadrature(self.N_r-1,a=0.0) + z0, weights0 = ball.quadrature(self.N_r - 1, a=0.0) - Q0 = ball.polynomial(self.N_r-1,0,0,z0,a=a) - Q_projection = ball.polynomial(self.N_r-1,0,0,z_projection,a=a) + Q0 = ball.polynomial(self.N_r - 1, 0, 0, z0, a=a) + Q_projection = ball.polynomial(self.N_r - 1, 0, 0, z_projection, a=a) - self.dV = ((Q0.dot(weights0)).T).dot(weights_projection*Q_projection) + self.dV = ((Q0.dot(weights0)).T).dot(weights_projection * Q_projection) self.pushW, self.pullW = {}, {} self.Q = {} - for ell in range( max(ell_min-R_max,0) , ell_max+R_max+1): - W = ball.polynomial(self.N_max+self.R_max-self.N_min(ell),0,ell,z_projection,a=a) - self.pushW[(ell)] = (weights_projection*W).astype(np.float64) + for ell in range(max(ell_min - R_max, 0), ell_max + R_max + 1): + W = ball.polynomial(self.N_max + self.R_max - self.N_min(ell), 0, ell, z_projection, a=a) + self.pushW[(ell)] = (weights_projection * W).astype(np.float64) self.pullW[(ell)] = (W.T).astype(np.float64) - for ell in range( ell_min, ell_max+1): - self.Q[(ell,0)] = np.array([[1]]) - for deg in range(1,R_max+1): - self.Q[(ell,deg)] = ball.recurseQ(self.Q[(ell,deg-1)],ell,deg) + for ell in range(ell_min, ell_max + 1): + self.Q[(ell, 0)] = np.array([[1]]) + for deg in range(1, R_max + 1): + self.Q[(ell, deg)] = ball.recurseQ(self.Q[(ell, deg - 1)], ell, deg) # downcast to double precision - self.radius = np.sqrt( (z_projection+1)/2 ).astype(np.float64) + self.radius = np.sqrt((z_projection + 1) / 2).astype(np.float64) self.dV = self.dV.astype(np.float64) self.LU_grad_initialized = [] @@ -74,358 +82,376 @@ def __init__(self,N_max,L_max,R_max=0,a=0,N_r=None,N_theta=None,ell_min=None,ell self.LU_curl_initialized = [] self.LU_curl = [] - for ell in range( ell_min, ell_max+1): + for ell in range(ell_min, ell_max + 1): # this hard codes an assumption that there are two ranks (e.g., 0 and 1) - self.LU_grad_initialized.append([False]*2) - self.LU_grad.append([None]*2) - self.LU_curl_initialized.append([False]*2) - self.LU_curl.append([None]*2) + self.LU_grad_initialized.append([False] * 2) + self.LU_grad.append([None] * 2) + self.LU_curl_initialized.append([False] * 2) + self.LU_curl.append([None] * 2) if timing: - self.radial_transform_time = 0. - self.angular_transform_time = 0. + self.radial_transform_time = 0.0 + self.angular_transform_time = 0.0 self.transpose_time = 0 @CachedMethod def grid(self, axis, dimensions=2): - if axis == 0 and dimensions == 2: grid = self.theta - if axis == 1 and dimensions == 2: grid = self.radius - if axis == 1 and dimensions == 3: grid = self.theta - if axis == 2 and dimensions == 3: grid = self.radius + if axis == 0 and dimensions == 2: + grid = self.theta + if axis == 1 and dimensions == 2: + grid = self.radius + if axis == 1 and dimensions == 3: + grid = self.theta + if axis == 2 and dimensions == 3: + grid = self.radius return reshape_vector(grid, dimensions, axis) @CachedMethod - def weight(self,axis,dimensions=2): - if axis == 0 and dimensions == 2: weight = self.S.weights - if axis == 1 and dimensions == 2: weight = self.dV - if axis == 1 and dimensions == 3: weight = self.S.weights - if axis == 2 and dimensions == 3: weight = self.dV + def weight(self, axis, dimensions=2): + if axis == 0 and dimensions == 2: + weight = self.S.weights + if axis == 1 and dimensions == 2: + weight = self.dV + if axis == 1 and dimensions == 3: + weight = self.S.weights + if axis == 2 and dimensions == 3: + weight = self.dV return reshape_vector(weight, dimensions, axis) @CachedMethod - def op(self,op_name,N,k,ell,dtype=np.float64,a=None): - if a == None: a = self.a - return ball.operator(op_name,N,k,ell,a=a).astype(dtype) + def op(self, op_name, N, k, ell, dtype=np.float64, a=None): + if a == None: + a = self.a + return ball.operator(op_name, N, k, ell, a=a).astype(dtype) - def xi(self,mu,ell): + def xi(self, mu, ell): # returns xi for ell > 0 or ell = 0 and mu = +1 # otherwise returns 0. if (ell > 0) or (ell == 0 and mu == 1): - return ball.xi(mu,ell) - return 0. + return ball.xi(mu, ell) + return 0.0 @CachedMethod - def unitary3D(self,rank=1,adjoint=False): - return ball.unitary3D(rank=rank,adjoint=adjoint) + def unitary3D(self, rank=1, adjoint=False): + return ball.unitary3D(rank=rank, adjoint=adjoint) @CachedMethod - def spins(self,rank): + def spins(self, rank): return ball.spins(rank) - def forward_angle(self,m,rank,data_in,data_out): + def forward_angle(self, m, rank, data_in, data_out): - if timing: start_time = time.time() + if timing: + start_time = time.time() if rank == 0: - data_out[0,int(self.S.L_min(m,0)):] = self.S.forward_spin(m,0,data_in[0]) + data_out[0, int(self.S.L_min(m, 0)) :] = self.S.forward_spin(m, 0, data_in[0]) return - spins, unitary = self.spins(rank), self.unitary3D(rank=rank,adjoint=True) + spins, unitary = self.spins(rank), self.unitary3D(rank=rank, adjoint=True) - data_in = np.einsum("ij,j...->i...",unitary,data_in) + data_in = np.einsum("ij,j...->i...", unitary, data_in) # This may benefit from some cython. Maybe, maybe not. Keaton? for i in range(3**rank): - data_out[i,int(self.S.L_min(m,spins[i])):] = self.S.forward_spin(m,spins[i],data_in[i]) + data_out[i, int(self.S.L_min(m, spins[i])) :] = self.S.forward_spin(m, spins[i], data_in[i]) if timing: end_time = time.time() - self.angular_transform_time += (end_time - start_time) + self.angular_transform_time += end_time - start_time - def backward_angle(self,m,rank,data_in,data_out): + def backward_angle(self, m, rank, data_in, data_out): - if timing: start_time = time.time() + if timing: + start_time = time.time() if rank == 0: - data_out[0] = self.S.backward_spin(m,0,data_in[0,self.S.L_min(m,0):]) + data_out[0] = self.S.backward_spin(m, 0, data_in[0, self.S.L_min(m, 0) :]) return spins = self.spins(rank) # This may benefit from some cython. Maybe, maybe not. Keaton? for i in range(3**rank): - data_out[i] = self.S.backward_spin(m,spins[i],data_in[i,int(self.S.L_min(m,spins[i])):]) + data_out[i] = self.S.backward_spin(m, spins[i], data_in[i, int(self.S.L_min(m, spins[i])) :]) if timing: end_time = time.time() - self.angular_transform_time += (end_time - start_time) + self.angular_transform_time += end_time - start_time @CachedMethod - def N_min(self,ell): + def N_min(self, ell): return ball.N_min(ell) # data is coeff representation of the ncc def ncc_matrix(self, N, k, ell, deg_in, deg_out, data, cutoff=1e-6, name=""): q_in = self.a - m_in = deg_in + 1/2 + m_in = deg_in + 1 / 2 q_out = k + self.a - m_out= ell + deg_out + 1/2 + m_out = ell + deg_out + 1 / 2 n_terms, max_term, matrix = clenshaw.ncc_matrix(N, q_in, m_in, q_out, m_out, data, cutoff=cutoff) - matrix /= (0.5)**(3/4) + matrix /= (0.5) ** (3 / 4) logger.debug("Expanded NCC {:s} to mode {:d} with {:d} terms.".format(name, max_term, n_terms)) return matrix - def forward_component(self,ell,deg,data): + def forward_component(self, ell, deg, data): # grid --> coefficients - N = self.N_max - self.N_min(ell-self.R_max) + 1 - if ell+deg >= 0: return (self.pushW[(ell+deg)][:N,:]).dot(data) + N = self.N_max - self.N_min(ell - self.R_max) + 1 + if ell + deg >= 0: + return (self.pushW[(ell + deg)][:N, :]).dot(data) else: shape = np.array(data.shape) shape[0] = N return np.zeros(shape) - def backward_component(self,ell,deg,data): + def backward_component(self, ell, deg, data): # coefficients --> grid - N = self.N_max - self.N_min(ell-self.R_max) + 1 - if ell+deg >= 0: - return self.pullW[(ell+deg)][:,:N].dot(data) + N = self.N_max - self.N_min(ell - self.R_max) + 1 + if ell + deg >= 0: + return self.pullW[(ell + deg)][:, :N].dot(data) else: shape = np.array(data.shape) shape[0] = self.N_r return np.zeros(shape) - def radial_forward(self,ell,rank,data_in,data_out): + def radial_forward(self, ell, rank, data_in, data_out): - if timing: start_time = time.time() + if timing: + start_time = time.time() if rank == 0: - np.copyto(data_out,self.forward_component(ell,0,data_in[0])) + np.copyto(data_out, self.forward_component(ell, 0, data_in[0])) return degs = self.spins(rank) - N = self.N_max - self.N_min(ell-self.R_max) + 1 + N = self.N_max - self.N_min(ell - self.R_max) + 1 - data_in = np.einsum("ij,j...->i...",self.Q[(ell,rank)].T,data_in) # note transpose + data_in = np.einsum("ij,j...->i...", self.Q[(ell, rank)].T, data_in) # note transpose - for i in range(3**rank): # Geoff thinks python can make this faster (??) "VECTORS MAN!" -- Geoff, 2017 - data_out[i*N:(i+1)*N] = self.forward_component(ell,degs[i],data_in[i]) + for i in range(3**rank): # Geoff thinks python can make this faster (??) "VECTORS MAN!" -- Geoff, 2017 + data_out[i * N : (i + 1) * N] = self.forward_component(ell, degs[i], data_in[i]) if timing: end_time = time.time() - self.radial_transform_time += (end_time - start_time) + self.radial_transform_time += end_time - start_time - def radial_backward(self,ell,rank,data_in,data_out): + def radial_backward(self, ell, rank, data_in, data_out): - if timing: start_time = time.time() + if timing: + start_time = time.time() if rank == 0: - data_out[0] = self.backward_component(ell,0,data_in) + data_out[0] = self.backward_component(ell, 0, data_in) return degs = self.spins(rank) - N =self.N_max - self.N_min(ell-self.R_max) + 1 + N = self.N_max - self.N_min(ell - self.R_max) + 1 for i in range(3**rank): - data_out[i] = self.backward_component(ell,degs[i],data_in[i*N:(i+1)*N]) + data_out[i] = self.backward_component(ell, degs[i], data_in[i * N : (i + 1) * N]) if timing: end_time = time.time() - self.radial_transform_time += (end_time - start_time) + self.radial_transform_time += end_time - start_time - def unpack(self,ell,rank,data_in): - N = self.N_max + 1 - self.N_min(ell-self.R_max) + def unpack(self, ell, rank, data_in): + N = self.N_max + 1 - self.N_min(ell - self.R_max) data_out = [] for i in range(3**rank): - data_out.append(data_in[i*N:(i+1)*N]) + data_out.append(data_in[i * N : (i + 1) * N]) return data_out - def pack(self,ell,rank,data_in,data_out): - N = self.N_max + 1 - self.N_min(ell-self.R_max) + def pack(self, ell, rank, data_in, data_out): + N = self.N_max + 1 - self.N_min(ell - self.R_max) for i in range(3**rank): - data_out[i*N:(i+1)*N] = data_in[i] + data_out[i * N : (i + 1) * N] = data_in[i] - def rank(self,length): + def rank(self, length): if length == 1: return 0 else: - return 1 + self.rank(length//3) + return 1 + self.rank(length // 3) - def grad(self,ell,rank,data_in): + def grad(self, ell, rank, data_in): shape = np.array(data_in.shape) shape[0] *= 3 data_dtype = data_in.dtype - data_out = np.zeros(shape,dtype=data_in.dtype) + data_out = np.zeros(shape, dtype=data_in.dtype) if STORE_LU_TRANSFORM: - i_LU = ell-self.ell_min + i_LU = ell - self.ell_min if not self.LU_grad_initialized[i_LU][rank]: logger.debug("LU_grad not initialized l={},rank={}".format(ell, rank)) - self.LU_grad[i_LU][rank] = [None]*(4*(3**rank)) + self.LU_grad[i_LU][rank] = [None] * (4 * (3**rank)) for i in range(3**rank): - tau_bar = ball.bar(i,rank) - N = self.N_max - self.N_min(ell-self.R_max) + tau_bar = ball.bar(i, rank) + N = self.N_max - self.N_min(ell - self.R_max) - if ell+tau_bar >= 1: - Cm = self.op('E',N,0,ell+tau_bar-1,data_dtype) - Dm = self.op('D-',N,0,ell+tau_bar,data_dtype) - xim = self.xi(-1,ell+tau_bar) + if ell + tau_bar >= 1: + Cm = self.op('E', N, 0, ell + tau_bar - 1, data_dtype) + Dm = self.op('D-', N, 0, ell + tau_bar, data_dtype) + xim = self.xi(-1, ell + tau_bar) index = i if STORE_LU_TRANSFORM: if not self.LU_grad_initialized[i_LU][rank]: self.LU_grad[i_LU][rank][index] = spla.splu(Cm) - data_out[i*(N+1):(i+1)*(N+1)] = self.LU_grad[i_LU][rank][index].solve(xim*Dm.dot(data_in[i*(N+1):(i+1)*(N+1)])) + data_out[i * (N + 1) : (i + 1) * (N + 1)] = self.LU_grad[i_LU][rank][index].solve( + xim * Dm.dot(data_in[i * (N + 1) : (i + 1) * (N + 1)]) + ) else: - data_out[i*(N+1):(i+1)*(N+1)] = spla.spsolve(Cm,xim*Dm.dot(data_in[i*(N+1):(i+1)*(N+1)])) + data_out[i * (N + 1) : (i + 1) * (N + 1)] = spla.spsolve(Cm, xim * Dm.dot(data_in[i * (N + 1) : (i + 1) * (N + 1)])) - if ell+tau_bar >= 0: - Cp = self.op('E',N,0,ell+tau_bar+1,data_dtype) - Dp = self.op('D+',N,0,ell+tau_bar,data_dtype) - xip = self.xi(+1,ell+tau_bar) - index = i+2*(3**rank) + if ell + tau_bar >= 0: + Cp = self.op('E', N, 0, ell + tau_bar + 1, data_dtype) + Dp = self.op('D+', N, 0, ell + tau_bar, data_dtype) + xip = self.xi(+1, ell + tau_bar) + index = i + 2 * (3**rank) if STORE_LU_TRANSFORM: if not self.LU_grad_initialized[i_LU][rank]: self.LU_grad[i_LU][rank][index] = spla.splu(Cp) - data_out[index*(N+1):(index+1)*(N+1)] = self.LU_grad[i_LU][rank][index].solve(xip*Dp.dot(data_in[i*(N+1):(i+1)*(N+1)])) + data_out[index * (N + 1) : (index + 1) * (N + 1)] = self.LU_grad[i_LU][rank][index].solve( + xip * Dp.dot(data_in[i * (N + 1) : (i + 1) * (N + 1)]) + ) else: - data_out[index*(N+1):(index+1)*(N+1)] = spla.spsolve(Cp,xip*Dp.dot(data_in[i*(N+1):(i+1)*(N+1)])) + data_out[index * (N + 1) : (index + 1) * (N + 1)] = spla.spsolve(Cp, xip * Dp.dot(data_in[i * (N + 1) : (i + 1) * (N + 1)])) if STORE_LU_TRANSFORM and not self.LU_grad_initialized[i_LU][rank]: self.LU_grad_initialized[i_LU][rank] = True return data_out - def curl(self,ell,rank,data_in,data_out): + def curl(self, ell, rank, data_in, data_out): data_dtype = data_in.dtype if STORE_LU_TRANSFORM: - i_LU = ell-self.ell_min + i_LU = ell - self.ell_min if not self.LU_curl_initialized[i_LU][rank]: logger.debug("LU_curl not initialized l={},rank={}".format(ell, rank)) - self.LU_curl[i_LU][rank] = [None]*(3) + self.LU_curl[i_LU][rank] = [None] * (3) - N = self.N_max - self.N_min(ell-self.R_max) - xim = self.xi(-1,ell) - xip = self.xi(+1,ell) + N = self.N_max - self.N_min(ell - self.R_max) + xim = self.xi(-1, ell) + xip = self.xi(+1, ell) if ell >= 1: - Cm = self.op('E',N,0,ell-1,data_dtype) - Dm = self.op('D-',N,0,ell,data_dtype) + Cm = self.op('E', N, 0, ell - 1, data_dtype) + Dm = self.op('D-', N, 0, ell, data_dtype) if STORE_LU_TRANSFORM: index = 0 if not self.LU_curl_initialized[ell][rank]: self.LU_curl[i_LU][rank][index] = spla.splu(Cm) - data_out[:N+1] = self.LU_curl[i_LU][rank][index].solve(-1j*xip*Dm.dot(data_in[(N+1):2*(N+1)])) + data_out[: N + 1] = self.LU_curl[i_LU][rank][index].solve(-1j * xip * Dm.dot(data_in[(N + 1) : 2 * (N + 1)])) else: - data_out[:N+1] = spla.spsolve(Cm,-1j*xip*Dm.dot(data_in[(N+1):2*(N+1)])) + data_out[: N + 1] = spla.spsolve(Cm, -1j * xip * Dm.dot(data_in[(N + 1) : 2 * (N + 1)])) else: - data_out[:N+1] = 0. + data_out[: N + 1] = 0.0 - - C0 = self.op('E',N,0,ell,data_dtype) - Dm = self.op('D-',N,0,ell+1,data_dtype) + C0 = self.op('E', N, 0, ell, data_dtype) + Dm = self.op('D-', N, 0, ell + 1, data_dtype) if STORE_LU_TRANSFORM: index = 1 if not self.LU_curl_initialized[ell][rank]: self.LU_curl[i_LU][rank][index] = spla.splu(C0) if ell >= 1: - Dp = self.op('D+',N,0,ell-1,data_dtype) + Dp = self.op('D+', N, 0, ell - 1, data_dtype) if STORE_LU_TRANSFORM: - data_out[(N+1):2*(N+1)] = self.LU_curl[i_LU][rank][index].solve( 1j*xim*Dm.dot(data_in[2*(N+1):]) - -1j*xip*Dp.dot(data_in[:(N+1)])) + data_out[(N + 1) : 2 * (N + 1)] = self.LU_curl[i_LU][rank][index].solve( + 1j * xim * Dm.dot(data_in[2 * (N + 1) :]) - 1j * xip * Dp.dot(data_in[: (N + 1)]) + ) else: - data_out[(N+1):2*(N+1)] = spla.spsolve(C0, 1j*xim*Dm.dot(data_in[2*(N+1):]) - -1j*xip*Dp.dot(data_in[:(N+1)])) + data_out[(N + 1) : 2 * (N + 1)] = spla.spsolve(C0, 1j * xim * Dm.dot(data_in[2 * (N + 1) :]) - 1j * xip * Dp.dot(data_in[: (N + 1)])) else: if STORE_LU_TRANSFORM: - data_out[(N+1):2*(N+1)] = self.LU_curl[i_LU][rank][index].solve(1j*xim*Dm.dot(data_in[2*(N+1):])) + data_out[(N + 1) : 2 * (N + 1)] = self.LU_curl[i_LU][rank][index].solve(1j * xim * Dm.dot(data_in[2 * (N + 1) :])) else: - data_out[(N+1):2*(N+1)] = spla.spsolve(C0,1j*xim*Dm.dot(data_in[2*(N+1):])) + data_out[(N + 1) : 2 * (N + 1)] = spla.spsolve(C0, 1j * xim * Dm.dot(data_in[2 * (N + 1) :])) - Cp = self.op('E',N,0,ell+1,data_dtype) - Dp = self.op('D+',N,0,ell,data_dtype) + Cp = self.op('E', N, 0, ell + 1, data_dtype) + Dp = self.op('D+', N, 0, ell, data_dtype) if STORE_LU_TRANSFORM: index = 2 if not self.LU_curl_initialized[ell][rank]: self.LU_curl[i_LU][rank][index] = spla.splu(Cp) if STORE_LU_TRANSFORM: - data_out[2*(N+1):] = self.LU_curl[i_LU][rank][index].solve(1j*xim*Dp.dot(data_in[(N+1):2*(N+1)])) + data_out[2 * (N + 1) :] = self.LU_curl[i_LU][rank][index].solve(1j * xim * Dp.dot(data_in[(N + 1) : 2 * (N + 1)])) else: - data_out[2*(N+1):] = spla.spsolve(Cp,1j*xim*Dp.dot(data_in[(N+1):2*(N+1)])) + data_out[2 * (N + 1) :] = spla.spsolve(Cp, 1j * xim * Dp.dot(data_in[(N + 1) : 2 * (N + 1)])) if STORE_LU_TRANSFORM and not self.LU_curl_initialized[i_LU][rank]: self.LU_curl_initialized[i_LU][rank] = True - def div(self,data_in): + def div(self, data_in): rank = self.rank(len(data_in)) data_dtype = data_in[0].dtype - data_out = [None]*(3**(rank-1)) + data_out = [None] * (3 ** (rank - 1)) - for i in range(3**(rank-1)): - tau_bar = ball.bar(i,rank-1) + for i in range(3 ** (rank - 1)): + tau_bar = ball.bar(i, rank - 1) m_tau_bar = -1 + tau_bar - p_tau_bar = 1 + tau_bar + p_tau_bar = 1 + tau_bar # initialise arrays data_out[i] = [] - for ell in range(self.L_max+1): + for ell in range(self.L_max + 1): - N = self.N_max - self.N_min(ell-self.R_max) + N = self.N_max - self.N_min(ell - self.R_max) - if ell+tau_bar == 0: - C = self.op('E',N,0,ell+tau_bar,data_dtype) - Dm = self.op('D-',N,0,ell+p_tau_bar,data_dtype) - xip = self.xi(+1,ell+tau_bar) + if ell + tau_bar == 0: + C = self.op('E', N, 0, ell + tau_bar, data_dtype) + Dm = self.op('D-', N, 0, ell + p_tau_bar, data_dtype) + xip = self.xi(+1, ell + tau_bar) - data_out[i].append(spla.spsolve(C,xip*Dm.dot(data_in[i+2*(3**(rank-1))][ell]))) + data_out[i].append(spla.spsolve(C, xip * Dm.dot(data_in[i + 2 * (3 ** (rank - 1))][ell]))) - elif ell+tau_bar > 0: - C = self.op('E',N,0,ell+tau_bar,data_dtype) - Dm = self.op('D-',N,0,ell+p_tau_bar,data_dtype) - Dp = self.op('D+',N,0,ell+m_tau_bar,data_dtype) - xim, xip = self.xi([-1,+1],ell+tau_bar) + elif ell + tau_bar > 0: + C = self.op('E', N, 0, ell + tau_bar, data_dtype) + Dm = self.op('D-', N, 0, ell + p_tau_bar, data_dtype) + Dp = self.op('D+', N, 0, ell + m_tau_bar, data_dtype) + xim, xip = self.xi([-1, +1], ell + tau_bar) - data_out[i].append(spla.spsolve(C,xip*Dm.dot(data_in[i+2*(3**(rank-1))][ell])+xim*Dp.dot(data_in[i][ell]))) + data_out[i].append(spla.spsolve(C, xip * Dm.dot(data_in[i + 2 * (3 ** (rank - 1))][ell]) + xim * Dp.dot(data_in[i][ell]))) else: - data_out[i].append(0*data_in[i][ell]) + data_out[i].append(0 * data_in[i][ell]) return data_out - def div_grad(self,data_in,ell_start=0,ell_end=None): - if ell_end == None: ell_end = self.L_max + def div_grad(self, data_in, ell_start=0, ell_end=None): + if ell_end == None: + ell_end = self.L_max rank = self.rank(len(data_in)) - data_out = [None]*(3**rank) + data_out = [None] * (3**rank) for i in range(3**rank): - tau_bar = ball.bar(i,rank) + tau_bar = ball.bar(i, rank) # initialise arrays data_out[i] = [] - for ell in range(ell_start,ell_end+1): + for ell in range(ell_start, ell_end + 1): ell_local = ell - ell_start - N = self.N_max - self.N_min(ell-self.R_max) + N = self.N_max - self.N_min(ell - self.R_max) - if ell+tau_bar >= 0: - CC = self.op('E',N,1,ell+tau_bar).dot(self.op('E',N,0,ell+tau_bar)) - DD = self.op('D-',N,1,ell+tau_bar+1).dot(self.op('D+',N,0,ell+tau_bar)) - Lap = spla.spsolve(CC,DD.dot(data_in[i][ell_local])) + if ell + tau_bar >= 0: + CC = self.op('E', N, 1, ell + tau_bar).dot(self.op('E', N, 0, ell + tau_bar)) + DD = self.op('D-', N, 1, ell + tau_bar + 1).dot(self.op('D+', N, 0, ell + tau_bar)) + Lap = spla.spsolve(CC, DD.dot(data_in[i][ell_local])) data_out[i].append(Lap) else: - data_out[i].append(0*data_in[i][ell_local]) + data_out[i].append(0 * data_in[i][ell_local]) return data_out - def cross_grid(self,a,b): - return np.array([a[1]*b[2]-a[2]*b[1],a[2]*b[0]-a[0]*b[2],a[0]*b[1]-a[1]*b[0]]) + def cross_grid(self, a, b): + return np.array([a[1] * b[2] - a[2] * b[1], a[2] * b[0] - a[0] * b[2], a[0] * b[1] - a[1] * b[0]]) - def dot_grid(self,a,b): - return a[0]*b[0] + a[1]*b[1] + a[2]*b[2] + def dot_grid(self, a, b): + return a[0] * b[0] + a[1] * b[1] + a[2] * b[2] -class TensorField: - def __init__(self,rank,B,domain): +class TensorField: + def __init__(self, rank, B, domain): self.domain = domain self.B = B self.rank = rank @@ -463,15 +489,15 @@ def require_layout(self, layout): elif layout == 'c' and self._layout == 'g': self.require_coeff_space() -class TensorField_2D(TensorField): - def __init__(self,rank,m,B,domain): +class TensorField_2D(TensorField): + def __init__(self, rank, m, B, domain): - TensorField.__init__(self,rank,B,domain) + TensorField.__init__(self, rank, B, domain) self.m = m mesh = self.domain.distributor.mesh - if len(mesh) == 0: #serial + if len(mesh) == 0: # serial self.ell_r_layout = domain.distributor.layouts[1] self.r_ell_layout = domain.distributor.layouts[1] else: @@ -479,23 +505,23 @@ def __init__(self,rank,m,B,domain): self.r_ell_layout = domain.distributor.layouts[1] local_grid_shape = self.ell_r_layout.local_shape(scales=1) - local_grid_shape = (int(domain.dealias[0]*local_grid_shape[0]), - int(domain.dealias[1]*local_grid_shape[1])) + local_grid_shape = (int(domain.dealias[0] * local_grid_shape[0]), int(domain.dealias[1] * local_grid_shape[1])) local_ellr_shape = self.ell_r_layout.local_shape(scales=domain.dealias) local_rell_shape = self.r_ell_layout.local_shape(scales=domain.dealias) - grid_shape = np.append(3**rank,np.array(local_grid_shape)) - ellr_shape = np.append(3**rank,np.array(local_ellr_shape)) - rell_shape = np.append(3**rank,np.array(local_rell_shape)) + grid_shape = np.append(3**rank, np.array(local_grid_shape)) + ellr_shape = np.append(3**rank, np.array(local_ellr_shape)) + rell_shape = np.append(3**rank, np.array(local_rell_shape)) - self.grid_data = np.zeros(grid_shape,dtype=np.complex128) - self.ellr_data = np.zeros(ellr_shape,dtype=np.complex128) - self.rell_data = np.zeros(rell_shape,dtype=np.complex128) + self.grid_data = np.zeros(grid_shape, dtype=np.complex128) + self.ellr_data = np.zeros(ellr_shape, dtype=np.complex128) + self.rell_data = np.zeros(rell_shape, dtype=np.complex128) self.fields = domain.new_fields(3**rank) - for field in self.fields: field.preset_scales(domain.dealias) + for field in self.fields: + field.preset_scales(domain.dealias) self.coeff_data = [] - for ell in range(self.ell_min,self.ell_max+1): - N = B.N_max - B.N_min(ell-B.R_max) + 1 - self.coeff_data.append(np.zeros(N*(3**rank),dtype=np.complex128)) + for ell in range(self.ell_min, self.ell_max + 1): + N = B.N_max - B.N_min(ell - B.R_max) + 1 + self.coeff_data.append(np.zeros(N * (3**rank), dtype=np.complex128)) self._layout = 'g' self.data = self.grid_data @@ -504,7 +530,7 @@ def require_coeff_space(self): """Transform from grid space to coeff space""" rank = self.rank - self.B.forward_angle(self.m,rank,self.grid_data,self.ellr_data) + self.B.forward_angle(self.m, rank, self.grid_data, self.ellr_data) for i, field in enumerate(self.fields): field.layout = self.ell_r_layout @@ -513,9 +539,7 @@ def require_coeff_space(self): for ell in range(self.ell_min, self.ell_max + 1): ell_local = ell - self.ell_min - self.B.radial_forward(ell,rank, - [self.fields[i].data[ell_local] for i in range(3**rank)], - self.coeff_data[ell_local]) + self.B.radial_forward(ell, rank, [self.fields[i].data[ell_local] for i in range(3**rank)], self.coeff_data[ell_local]) self.data = self.coeff_data self._layout = 'c' @@ -526,75 +550,75 @@ def require_grid_space(self): rank = self.rank for ell in range(self.ell_min, self.ell_max + 1): ell_local = ell - self.ell_min - self.B.radial_backward(ell,rank,self.coeff_data[ell_local],self.rell_data[:,ell_local,:]) - self.rell_data[:,ell_local,:] = np.einsum("ij,j...->i...",self.B.Q[(ell,rank)],self.rell_data[:,ell_local,:]) + self.B.radial_backward(ell, rank, self.coeff_data[ell_local], self.rell_data[:, ell_local, :]) + self.rell_data[:, ell_local, :] = np.einsum("ij,j...->i...", self.B.Q[(ell, rank)], self.rell_data[:, ell_local, :]) for i, field in enumerate(self.fields): field.layout = self.r_ell_layout field.data = self.rell_data[i] field.require_layout(self.ell_r_layout) - self.B.backward_angle(self.m,rank,np.array([self.fields[i].data for i in range(3**rank)]),self.grid_data) + self.B.backward_angle(self.m, rank, np.array([self.fields[i].data for i in range(3**rank)]), self.grid_data) if rank > 0: - self.grid_data = np.einsum("ij,j...->i...",self.B.unitary3D(rank=rank,adjoint=False),self.grid_data) + self.grid_data = np.einsum("ij,j...->i...", self.B.unitary3D(rank=rank, adjoint=False), self.grid_data) self.data = self.grid_data self._layout = 'g' class TensorField_3D(TensorField): + def __init__(self, rank, B, domain): - def __init__(self,rank,B,domain): - - TensorField.__init__(self,rank,B,domain) + TensorField.__init__(self, rank, B, domain) mesh = self.domain.distributor.mesh - if len(mesh) == 0: # serial - self.phi_layout = domain.distributor.layouts[3] - self.th_m_layout = domain.distributor.layouts[2] + if len(mesh) == 0: # serial + self.phi_layout = domain.distributor.layouts[3] + self.th_m_layout = domain.distributor.layouts[2] self.ell_r_layout = domain.distributor.layouts[1] self.r_ell_layout = domain.distributor.layouts[1] - elif len(mesh) == 1: # 1D domain decomposition - self.phi_layout = domain.distributor.layouts[4] - self.th_m_layout = domain.distributor.layouts[2] + elif len(mesh) == 1: # 1D domain decomposition + self.phi_layout = domain.distributor.layouts[4] + self.th_m_layout = domain.distributor.layouts[2] self.ell_r_layout = domain.distributor.layouts[1] self.r_ell_layout = domain.distributor.layouts[1] - elif len(mesh) == 2: # 2D domain decomposition - self.phi_layout = domain.distributor.layouts[5] - self.th_m_layout = domain.distributor.layouts[3] + elif len(mesh) == 2: # 2D domain decomposition + self.phi_layout = domain.distributor.layouts[5] + self.th_m_layout = domain.distributor.layouts[3] self.ell_r_layout = domain.distributor.layouts[2] self.r_ell_layout = domain.distributor.layouts[1] # allocating arrays local_grid_shape = self.phi_layout.local_shape(scales=domain.dealias) - grid_shape = np.append(3**rank,np.array(local_grid_shape)) - self.grid_data = np.zeros(grid_shape,dtype=np.float64) + grid_shape = np.append(3**rank, np.array(local_grid_shape)) + self.grid_data = np.zeros(grid_shape, dtype=np.float64) - scales = (1,1,domain.dealias[2]) + scales = (1, 1, domain.dealias[2]) local_ellr_shape = self.ell_r_layout.local_shape(scales=scales) - ellr_shape = np.append(3**rank,np.array(local_ellr_shape)) - self.mlr_ell_data = np.zeros(ellr_shape,dtype=np.complex128) + ellr_shape = np.append(3**rank, np.array(local_ellr_shape)) + self.mlr_ell_data = np.zeros(ellr_shape, dtype=np.complex128) local_rell_shape = self.r_ell_layout.local_shape(scales=scales) - rell_shape = np.append(3**rank,np.array(local_rell_shape)) - self.mlr_r_data = np.zeros(rell_shape,dtype=np.complex128) - rlm_shape = np.append(3**rank,np.array(local_rell_shape)[::-1]) - self.rlm_data = np.zeros(rlm_shape,dtype=np.complex128) + rell_shape = np.append(3**rank, np.array(local_rell_shape)) + self.mlr_r_data = np.zeros(rell_shape, dtype=np.complex128) + rlm_shape = np.append(3**rank, np.array(local_rell_shape)[::-1]) + self.rlm_data = np.zeros(rlm_shape, dtype=np.complex128) - scales = (1,domain.dealias[1],self.domain.dealias[2]) + scales = (1, domain.dealias[1], self.domain.dealias[2]) local_mthr_shape = self.th_m_layout.local_shape(scales=scales) - mthr_shape = np.append(3**rank,np.array(local_mthr_shape)) - self.mthr_data = np.zeros(mthr_shape,dtype=np.complex128) + mthr_shape = np.append(3**rank, np.array(local_mthr_shape)) + self.mthr_data = np.zeros(mthr_shape, dtype=np.complex128) self.fields = domain.new_fields(3**rank) - for field in self.fields: field.preset_scales(domain.dealias) + for field in self.fields: + field.preset_scales(domain.dealias) m_size = B.m_max - B.m_min + 1 self.coeff_data = [] - for ell in range(self.ell_min,self.ell_max+1): - N = B.N_max - B.N_min(ell-B.R_max) + 1 - self.coeff_data.append(np.zeros( (N*(3**rank),m_size) ,dtype=np.complex128)) + for ell in range(self.ell_min, self.ell_max + 1): + N = B.N_max - B.N_min(ell - B.R_max) + 1 + self.coeff_data.append(np.zeros((N * (3**rank), m_size), dtype=np.complex128)) self._layout = 'g' self.data = self.grid_data @@ -611,50 +635,60 @@ def decrement_layout(self): B = self.B if self._layout == 'g': - if barrier: self.domain.dist.comm_cart.Barrier() - if timing: start_time = time.time() + if barrier: + self.domain.dist.comm_cart.Barrier() + if timing: + start_time = time.time() for i, field in enumerate(self.fields): field.layout = self.phi_layout - np.copyto(field.data,self.data[i]) + np.copyto(field.data, self.data[i]) field.require_layout(self.th_m_layout) - np.copyto(self.mthr_data[i],field.data) + np.copyto(self.mthr_data[i], field.data) self._layout = 3 if timing: end_time = time.time() - self.B.transpose_time += end_time-start_time - if barrier: self.domain.dist.comm_cart.Barrier() + self.B.transpose_time += end_time - start_time + if barrier: + self.domain.dist.comm_cart.Barrier() elif self._layout == 3: - if barrier: self.domain.dist.comm_cart.Barrier() - for m in range(B.m_min,B.m_max+1): + if barrier: + self.domain.dist.comm_cart.Barrier() + for m in range(B.m_min, B.m_max + 1): m_local = m - B.m_min - B.forward_angle(m,rank,np.array(self.mthr_data[:,m_local]),self.mlr_ell_data[:,m_local,:,:]) + B.forward_angle(m, rank, np.array(self.mthr_data[:, m_local]), self.mlr_ell_data[:, m_local, :, :]) self._layout = 2 - if barrier: self.domain.dist.comm_cart.Barrier() + if barrier: + self.domain.dist.comm_cart.Barrier() elif self._layout == 2: - if barrier: self.domain.dist.comm_cart.Barrier() - if timing: start_time = time.time() + if barrier: + self.domain.dist.comm_cart.Barrier() + if timing: + start_time = time.time() if self.ell_r_layout != self.r_ell_layout: for i, field in enumerate(self.fields): field.layout = self.ell_r_layout - np.copyto(field.data,self.mlr_ell_data[i]) + np.copyto(field.data, self.mlr_ell_data[i]) self.domain.distributor.paths[1].decrement([field]) - np.copyto(self.rlm_data[i],field.data.T) + np.copyto(self.rlm_data[i], field.data.T) else: for i, field in enumerate(self.fields): - np.copyto(self.rlm_data[i],self.mlr_ell_data[i].T) + np.copyto(self.rlm_data[i], self.mlr_ell_data[i].T) self._layout = 1 if timing: end_time = time.time() - self.B.transpose_time += end_time-start_time - if barrier: self.domain.dist.comm_cart.Barrier() + self.B.transpose_time += end_time - start_time + if barrier: + self.domain.dist.comm_cart.Barrier() elif self._layout == 1: - if barrier: self.domain.dist.comm_cart.Barrier() + if barrier: + self.domain.dist.comm_cart.Barrier() for ell in range(B.ell_min, B.ell_max + 1): ell_local = ell - B.ell_min - self.B.radial_forward(ell,rank,self.rlm_data[:,:,ell_local,:],self.coeff_data[ell_local]) + self.B.radial_forward(ell, rank, self.rlm_data[:, :, ell_local, :], self.coeff_data[ell_local]) self.data = self.coeff_data self._layout = 'c' - if barrier: self.domain.dist.comm_cart.Barrier() + if barrier: + self.domain.dist.comm_cart.Barrier() def require_grid_space(self): """Transform from coeff space to grid space""" @@ -668,51 +702,61 @@ def increment_layout(self): B = self.B if self._layout == 'c': - if barrier: self.domain.dist.comm_cart.Barrier() - for ell in range(B.ell_min, B.ell_max+1): + if barrier: + self.domain.dist.comm_cart.Barrier() + for ell in range(B.ell_min, B.ell_max + 1): ell_local = ell - B.ell_min - B.radial_backward(ell,rank,self.data[ell_local],self.rlm_data[:,:,ell_local,:]) - self.rlm_data[:,:,ell_local,:] = np.einsum("ij,j...->i...",self.B.Q[(ell,rank)],self.rlm_data[:,:,ell_local,:]) - np.copyto(self.mlr_r_data,self.rlm_data.transpose(0,3,2,1),casting='no') + B.radial_backward(ell, rank, self.data[ell_local], self.rlm_data[:, :, ell_local, :]) + self.rlm_data[:, :, ell_local, :] = np.einsum("ij,j...->i...", self.B.Q[(ell, rank)], self.rlm_data[:, :, ell_local, :]) + np.copyto(self.mlr_r_data, self.rlm_data.transpose(0, 3, 2, 1), casting='no') self._layout = 1 - if barrier: self.domain.dist.comm_cart.Barrier() + if barrier: + self.domain.dist.comm_cart.Barrier() elif self._layout == 1: - if barrier: self.domain.dist.comm_cart.Barrier() - if timing: start_time = time.time() + if barrier: + self.domain.dist.comm_cart.Barrier() + if timing: + start_time = time.time() if self.ell_r_layout != self.r_ell_layout: for i, field in enumerate(self.fields): field.layout = self.r_ell_layout - np.copyto(field.data,self.mlr_r_data[i]) + np.copyto(field.data, self.mlr_r_data[i]) self.domain.distributor.paths[1].increment([field]) - np.copyto(self.mlr_ell_data[i],field.data) + np.copyto(self.mlr_ell_data[i], field.data) else: for i, field in enumerate(self.fields): - np.copyto(self.mlr_ell_data[i],self.mlr_r_data[i]) + np.copyto(self.mlr_ell_data[i], self.mlr_r_data[i]) self._layout = 2 if timing: end_time = time.time() - self.B.transpose_time += end_time-start_time - if barrier: self.domain.dist.comm_cart.Barrier() + self.B.transpose_time += end_time - start_time + if barrier: + self.domain.dist.comm_cart.Barrier() elif self._layout == 2: - if barrier: self.domain.dist.comm_cart.Barrier() + if barrier: + self.domain.dist.comm_cart.Barrier() for m in range(B.m_min, B.m_max + 1): m_local = m - B.m_min - B.backward_angle(m,rank,self.mlr_ell_data[:,m_local,:,:],self.mthr_data[:,m_local,:,:]) + B.backward_angle(m, rank, self.mlr_ell_data[:, m_local, :, :], self.mthr_data[:, m_local, :, :]) if rank > 0: - self.mthr_data = np.einsum("ij,j...->i...",self.B.unitary3D(rank=rank,adjoint=False),self.mthr_data) + self.mthr_data = np.einsum("ij,j...->i...", self.B.unitary3D(rank=rank, adjoint=False), self.mthr_data) self._layout = 3 - if barrier: self.domain.dist.comm_cart.Barrier() + if barrier: + self.domain.dist.comm_cart.Barrier() elif self._layout == 3: - if barrier: self.domain.dist.comm_cart.Barrier() - if timing: start_time = time.time() + if barrier: + self.domain.dist.comm_cart.Barrier() + if timing: + start_time = time.time() for i, field in enumerate(self.fields): field.layout = self.th_m_layout - np.copyto(field.data,self.mthr_data[i]) + np.copyto(field.data, self.mthr_data[i]) field.require_layout(self.phi_layout) - np.copyto(self.grid_data[i],field.data) + np.copyto(self.grid_data[i], field.data) self.data = self.grid_data if timing: end_time = time.time() - self.B.transpose_time += end_time-start_time + self.B.transpose_time += end_time - start_time self._layout = 'g' - if barrier: self.domain.dist.comm_cart.Barrier() + if barrier: + self.domain.dist.comm_cart.Barrier() diff --git a/dedalus/libraries/dedalus_sphere/clenshaw.py b/dedalus/libraries/dedalus_sphere/clenshaw.py index 01f9efaa..47e5afb8 100644 --- a/dedalus/libraries/dedalus_sphere/clenshaw.py +++ b/dedalus/libraries/dedalus_sphere/clenshaw.py @@ -1,4 +1,3 @@ - from scipy import sparse import numpy as np from . import jacobi @@ -8,15 +7,17 @@ # arg_basis is the basis of the thing we're multiplying by # i.e., if we are doing u.grad X, then arg_basis is the basis of u + def ncc_matrix(N, a_ncc, b_ncc, a_arg, b_arg, coeffs, cutoff=1e-6): """Build NCC matrix via Clenshaw algorithm.""" # Kronecker Clenshaw on argument Jacobi matrix J = jacobi_matrix(N, a_arg, b_arg) A, B = jacobi_recursion(N, a_ncc, b_ncc, J) - f0 = 1/np.sqrt(jacobi.mass(a_ncc,b_ncc)) * sparse.identity(N) + f0 = 1 / np.sqrt(jacobi.mass(a_ncc, b_ncc)) * sparse.identity(N) total = matrix_clenshaw(coeffs, A, B, f0, cutoff=cutoff) return total + def jacobi_recursion(N, a, b, X): """ Build Clenshaw recurrence coefficients for Jacobi polynomials. @@ -40,17 +41,19 @@ def jacobi_recursion(N, a, b, X): I = sparse.identity(X.shape[0]) # Clenshaw coefficients def compute_A(n): - if 0 <= n < (N-1): - return (X - JA[n,n]*I) / JA[n,n+1] + if 0 <= n < (N - 1): + return (X - JA[n, n] * I) / JA[n, n + 1] else: - return 0*I + return 0 * I + def compute_B(n): - if 0 < n < (N-1): - return (-J[n,n-1] / J[n,n+1]) * I + if 0 < n < (N - 1): + return (-J[n, n - 1] / J[n, n + 1]) * I else: - return 0*I - A = DeferredTuple(compute_A, N+1) - B = DeferredTuple(compute_B, N+1) + return 0 * I + + A = DeferredTuple(compute_A, N + 1) + B = DeferredTuple(compute_B, N + 1) return A, B @@ -62,27 +65,27 @@ def matrix_clenshaw(c, A, B, f0, cutoff): N = len(c) I = sparse.identity(f0.shape[0]) # Clenshaw - b0, b1 = 0*I, 0*I + b0, b1 = 0 * I, 0 * I n_terms = max_term = 0 for n in reversed(range(N)): b1, b2 = b0, b1 if abs(c[n]) > cutoff: - b0 = (c[n] * I) + (A[n] @ b1) + (B[n+1] @ b2) + b0 = (c[n] * I) + (A[n] @ b1) + (B[n + 1] @ b2) n_terms += 1 - if max_term == 0 : + if max_term == 0: # reversed range, so first term is max_term max_term = n else: - b0 = (A[n] @ b1) + (B[n+1] @ b2) + b0 = (A[n] @ b1) + (B[n + 1] @ b2) return n_terms, max_term, (b0 @ f0) + def jacobi_matrix(N, a, b): - J = jacobi.operator('J',N-1,a,b) + J = jacobi.operator('J', N - 1, a, b) return J.tocsr().astype(np.float64) class DeferredTuple: - def __init__(self, entry_function, size): self.entry_function = entry_function self.size = size @@ -92,7 +95,7 @@ def __getitem__(self, key): if key < 0: key += len(self) if key >= len(self): - raise IndexError("The index (%d) is out of range." %key) + raise IndexError("The index (%d) is out of range." % key) return self.entry_function(key) else: raise TypeError("Invalid argument type.") diff --git a/dedalus/libraries/dedalus_sphere/jacobi.py b/dedalus/libraries/dedalus_sphere/jacobi.py index 635fd816..eaa215ea 100644 --- a/dedalus/libraries/dedalus_sphere/jacobi.py +++ b/dedalus/libraries/dedalus_sphere/jacobi.py @@ -1,212 +1,401 @@ -import numpy as np -from scipy.sparse import dia_matrix as banded +"""Jacobi polynomial functions.""" +import numpy as np +#import xprec +from scipy.sparse import dia_matrix +from scipy.special import beta, betaln +from scipy.linalg import eigvalsh_tridiagonal from .operators import Operator, Codomain, infinite_csr -dtype='float64' -def coefficient_connection(N,ab,cd,init_ab=1,init_cd=1): - """The connection matrix between any bases coefficients: - Pab(z) = Pcd(z) @ Cab2cd --> Acd = Cab2cd @ Aab - The output is always a dense matrix format. - Parameters - ---------- - N : int - ab, cd : tuples of input and output Jacobi parameters. - """ - - a,b = ab - c,d = cd - - zcd, wcd = quadrature(N,c,d) - - wcd /= np.sum(wcd) - - Pab = polynomials(N,a,b,zcd,init_ab + 0*zcd) - Pcd = polynomials(N,c,d,zcd,init_cd + 0*zcd) +# Default dtypes +DEFAULT_GRID_DTYPE = np.longdouble +DEFAULT_OPERATOR_DTYPE = np.float64 - return Pcd @ (wcd*Pab).T -def polynomials(n,a,b,z,init=None,Newton=False,normalised=True,dtype=dtype,internal='longdouble'): +def coefficient_connection(n, ab, cd, init_ab=1, init_cd=1): """ - Jacobi polynomials, P(n,a,b,z), of type (a,b) up to degree n-1. - - Jacobi.polynomials(n,a,b,z)[k] gives P(k,a,b,z). + The connection matrix between any bases coefficients: + Pab(z) = Pcd(z) @ Cab2cd --> Acd = Cab2cd @ Aab - Newton = True: cubic-converging update of P(n-1,a,b,z) = 0. + The output is always a dense matrix format. Parameters ---------- - a,b > -1 - z: float, np.ndarray. + n : int + Number of polynomials (max degree + 1). + ab, cd : tuple of floats + Tuples of input and output Jacobi parameters. + init_ab, init_cd : floats or arrays, optional. + Initial envelopes for input and output polynomials. + + Returns + ------- + conn : array + Coupling coefficient matrix from (a,b) to (c,d) polynomials. + """ + a, b = ab + c, d = cd + zcd, wcd = quadrature(n, c, d) + wcd /= np.sum(wcd) + Pab = polynomials(n, a, b, zcd, init_ab) + Pcd = polynomials(n, c, d, zcd, init_cd) + return Pcd @ (wcd * Pab).T - init: float, np.ndarray or None -> 1+0*z, or (1+0*z)/sqrt(mass) if normalised. - normalised: classical or unit-integral normalisation. - dtype: 'float64','longdouble' output dtype. - internal: internal dtype. +def generate_polynomials(n, a, b, z, init=None, normalized=True, dtype=None): """ + Jacobi polynomials of type (a,b) up to degree n-1. + Parameters + ---------- + n : int + Number of polynomials to compute (max degree + 1). + a, b : float + Jacobi parameters. + z : array + Grid array. + init : float or array, optional + Initial envelope for polynomial recursion. Default determined by 'normalized' option. + normalized : bool, optional + True to initialize with 1/sqrt(mass) if init not provided. Default: True. + dtype : dtype, optional + Data type. Default: module-level DEFAULT_GRID_DTYPE. + + Yields + ------ + polynomials : array + Polynomials evalauted at specified z points. + Yield order follows polynomial degree. + """ + if dtype is None: + dtype = DEFAULT_GRID_DTYPE + # Return empty array if n < 1 if n < 1: return np.zeros((0, z.size), dtype=dtype) - + # Cast z to working dtype + if np.isscalar(z): + z = np.dtype(dtype).type(z) + else: + z = z.astype(dtype) + # Setup initial envelope if init is None: - init = 1 + 0*z - if normalised: - init /= np.sqrt(mass(a,b),dtype=internal) - - Z = operator('Z',normalised=normalised,dtype=internal) - Z = banded(Z(n+1,a,b).T).data - - shape = n+1 - if type(z) == np.ndarray: - shape = (shape,len(z)) - - P = np.zeros(shape,dtype=internal) - P[0] = init - + init = 1 + 0 * z + if normalized: + init /= np.sqrt(mass(a, b, dtype=dtype)) + else: + init = init + 0 * z # Cast to working dtype + # Get Jacobi operator + Z = operator('Z', normalized=normalized) + Z = dia_matrix(Z(n + 1, a, b).T).data.astype(dtype) + # Generate polynomials via recursion in degree + P0 = init + yield P0 + if n == 1: + return if len(Z) == 2: - P[1] = z*P[0]/Z[1,1] - for k in range(2,n+1): - P[k] = (z*P[k-1] - Z[0,k-2]*P[k-2])/Z[1,k] + P1 = z * P0 / Z[1, 1] + yield P1 + Pkm2 = P0 + Pkm1 = P1 + for k in range(2, n): + Pk = (z * Pkm1 - Z[0, k - 2] * Pkm2) / Z[1, k] + yield Pk + Pkm2 = Pkm1 + Pkm1 = Pk else: - P[1] = (z-Z[1,0])*P[0]/Z[2,1] - for k in range(2,n+1): - P[k] = ((z-Z[1,k-1])*P[k-1] - Z[0,k-2]*P[k-2])/Z[2,k] + P1 = (z - Z[1, 0]) * P0 / Z[2, 1] + yield P1 + Pkm2 = P0 + Pkm1 = P1 + for k in range(2, n): + Pk = ((z - Z[1, k - 1]) * Pkm1 - Z[0, k - 2] * Pkm2) / Z[2, k] + yield Pk + Pkm2 = Pkm1 + Pkm1 = Pk + + +def polynomials(n, a, b, z, init=None, Newton=False, normalized=True, dtype=None): + """ + Jacobi polynomials of type (a,b) up to degree n-1. + Parameters + ---------- + n : int + Number of polynomials to compute (max degree + 1). + a, b : float + Jacobi parameters. + z : array + Grid array. + init : float or array, optional + Initial envelope for polynomial recursion. Default determined by 'normalized' option. + Newton : bool, optional + Return cubicly converging update to grid as roots of P[n]. Default: False. + normalized : bool, optional + True to initialize with 1/sqrt(mass) if init not provided. Default: True. + dtype : dtype, optional + Data type. Default: module-level DEFAULT_GRID_DTYPE. + + Returns + ------- + grid : array, optional + Refined grid as roots of P[n]. Only returned if Newton is True. + polynomials : array + Array of polynomials evaluated at specified z points. + First axis is polynomial degree. + """ + if dtype is None: + dtype = DEFAULT_GRID_DTYPE + # Return empty array if n < 1 + if n < 1: + return np.zeros((0, z.size), dtype=dtype) + # Cast z to working dtype + if np.isscalar(z): + z = np.dtype(dtype).type(z) + z_shape = tuple() + else: + z = z.astype(dtype) + z_shape = z.shape + # Setup initial envelope + if init is None: + init = 1 + 0 * z + if normalized: + init /= np.sqrt(mass(a, b, dtype=dtype)) + else: + init = init + 0 * z # Cast to working dtype + # Get Jacobi operator + Z = operator('Z', normalized=normalized) + Z = dia_matrix(Z(n + 1, a, b).T).data.astype(dtype) + # Build polynomials via recursion in degree + shape = (n + 1,) + z_shape + P = np.zeros(shape, dtype=dtype) + P[0] = init + if len(Z) == 2: + P[1] = z * P[0] / Z[1, 1] + for k in range(2, n + 1): + P[k] = (z * P[k - 1] - Z[0, k - 2] * P[k - 2]) / Z[1, k] + else: + P[1] = (z - Z[1, 0]) * P[0] / Z[2, 1] + for k in range(2, n + 1): + P[k] = ((z - Z[1, k - 1]) * P[k - 1] - Z[0, k - 2] * P[k - 2]) / Z[2, k] + # Newton update to grid if Newton: - L = n + (a+b)/2 - return z + (1-z**2)*P[n-1]/(L*Z[-1,n]*P[n]-(L-1)*Z[0,n-2]*P[n-2]), P[:n-1] + L = n + (a + b) / 2 + return z + (1 - z * z) * P[n - 1] / (L * Z[-1, n] * P[n] - (L - 1) * Z[0, n - 2] * P[n - 2]), P[: n - 1] + return P[:n] - return P[:n].astype(dtype) -def quadrature(n,a,b,days=3,probability=False,dtype=dtype,internal='longdouble'): +def polynomial_integrals(n, a, b, dtype=None, **kw): """ - Jacobi 'roots' grid and weights; solutions to - - P(n,a,b,z) = 0; len(z) = n. - - sum(weights*polynomial) = integrate_(-1,+1)(polynomial), - exactly up to degree(polynomial) = 2n - 1. + Definite integrals of the Jacobi polynomials, evaluated with Gauss-Legendre quadrature. Parameters ---------- - n: int > 0. - a,b: float > -1. - days: number of Newton updates. - probability: sum(weights) = 1 or sum(weights) = mass(a,b) - dtype: 'float64','longdouble' output dtype. - internal: internal dtype (Newton uses by default). - + 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. + The grid points are the roots of P[n]: P(n, a, b, z[i]) = 0, len(z) = n. + For all polynomials p up to degree 2*n - 1, the weights satisfy: sum(w[i]*p(z[i])) = integ(p). - z = grid_guess(n,a,b,dtype=internal) - + Parameters + ---------- + n : int + Number of quadrature nodes. + a, b : float + Jacobi parameters. + iterations : int, optional + Number of Newton updates. Default: 3. + probability : bool, optional. + True for sum(weights) = 1 or False for sum(weights) = mass(a,b). Default: True. + dtype : dtype, optional + Data type. Default: module-level DEFAULT_GRID_DTYPE. + + Returns + ------- + z : array + Quadrature grid. + w : array + Quadrature weights. + """ + if dtype is None: + dtype = DEFAULT_GRID_DTYPE + # Initial guesses + z = grid_guess(n, a, b, dtype=dtype) if probability: - w = 1 + w = 1 / n + 0 * z else: - w = mass(a,b) - - if (a == b == -1/2): - return z.astype(dtype), (w/n + 0*z).astype(dtype) - elif (a == b == +1/2): - P = polynomials(n+1,a,b,z,dtype=internal)[:n] + m = mass(a, b, dtype=dtype) + w = mass(a, b, dtype=dtype) / int(n) + 0 * z + # Chebyshev T: grid guess exact, uniform weights + if a == b == -1 / 2: + return z, w + # Chebyshev U: grid guess exact + elif a == b == 1 / 2: + P = polynomials(n, a, b, z, dtype=dtype) + # Newton iterations for other Jacobi else: - for _ in range(days): - z, P = polynomials(n+1,a,b,z,Newton=True) + for _ in range(iterations): + z, P = polynomials(n + 1, a, b, z, dtype=dtype, Newton=True) + # Compute weights + w *= n * P[0] * P[0] / np.sum(P * P, axis=0, initial=0.0) + return z, w - P[0] /= np.sqrt(np.sum(P**2,axis=0)) - w *= P[0]**2 - return z.astype(dtype), w.astype(dtype) - -def grid_guess(n,a,b,dtype='longdouble',quick=False): +def grid_guess(n, a, b, quick=False, dtype=None): """ - Approximate solution to - - P(n,a,b,z) = 0 + Approximate guesses for roots of P^(a,b)_n(z) using Golub-Welsch. + Parameters + ---------- + n : int + Degree / number of roots in grid. + a, b : ints or floats + Jacobi parameters. + quick : bool, optional + Use a quick guess. Default: False. + dtype : dtype + Data type. Default: module-level DEFAULT_GRID_DTYPE. + + Returns + ------- + grid : array + Grid guess. """ - - if a == b == -1/2 : - quick = True - + if dtype is None: + dtype = DEFAULT_GRID_DTYPE + # Use constant for n = 1 if n == 1: - return operator('Z')(n,a,b).A[0] - + return operator('Z')(n, a, b).A[0].astype(dtype) + # Quick guess is exact for Chebyshev + if a == b == -1 / 2: + quick = True + # Compute guess if quick: - return np.cos(np.pi*(np.arange(4*n-1,2,-4,dtype=dtype)+2*a)/(4*n+2*(a+b+1))) - - from scipy.linalg import eigvalsh_tridiagonal as eigs - - Z = banded(operator('Z')(n,a,b)) - - return eigs(Z.diagonal(0),Z.diagonal(1)).astype(dtype) + return np.cos(np.pi * (np.arange(4 * n - 1, 2, -4).astype(dtype) + 2 * a) / (4 * int(n) + 2 * (a + b + 1))) + else: + # Use Golub-Welsch + Z = dia_matrix(operator('Z')(n, a, b)) + return eigvalsh_tridiagonal(Z.diagonal(0), Z.diagonal(1)).astype(dtype) -def measure(a,b,z,probability=True,log=False): +def measure(a, b, z, probability=True, log=False, dtype=None): """ - - mu(a,b,z) = (1-z)**a (1+z)**b - - if normalised: ((1-z)/2)**a ((1+z)/2)**b / (2*Beta(a+1,b+1)) + Jacobi measure evaluated on a grid. Classical/unnormalized measure: + mu(a,b,z) = (1-z)^a (1+z)^b Parameters ---------- - a,b > -1 - + a, b : ints or floats + Jacobi parameters. + z : float or array + Points for evaluating measure. + probability : bool, optional + False for classical measure, True to renormalize to have unit integral. Default: True. + log : bool, optional + Return log of measure. Default: False. + dtype : dtype, optional. + Data type. Default: module-level DEFAULT_GRID_DTYPE. + + Returns + ------- + measure : float or array + Measure evaluated at a z. """ - - if not log: - - w = 1 - - if a != 0: - w *= (1-z)**a - - if b != 0: - w *= (1+z)**b - + if dtype is None: + dtype = DEFAULT_GRID_DTYPE + # Cast z to working dtype + if np.isscalar(z): + z = np.dtype(dtype).type(z) + else: + z = z.astype(dtype) + # Compute measure + if log: + log_w = 0 * z + if a: + log_w += a * np.log(1 - z) + if b: + log_w += b * np.log(1 + z) if probability: - w /= mass(a,b) - + log_w -= mass(a, b, log=True, dtype=dtype) + if np.isscalar(z): + if np.isnan(log_w): + log_w = np.dtype(dtype).type(-np.inf) + else: + log_w[np.isnan(log_w)] = -np.inf # log(0) gives nan instead of -inf + return log_w + else: + w = 1 + 0 * z + if a: + w *= (1 - z) ** a + if b: + w *= (1 + z) ** b + if probability: + w /= mass(a, b, dtype=dtype) return w - S = 0 - - if a != 0: - S += a*np.log(1-z) - if b != 0: - S += b*np.log(1+z) - - if probability: - S -= mass(a,b,log=True) - - return S - -def mass(a,b,log=False): +def mass(a, b, log=False, dtype=None): """ - - 2**(a+b+1)*Beta(a+1,b+1) = integrate_(-1,+1)( (1-z)**a (1+z)**b ) + Integral of classical Jacobi measure: + \integ_{-1,+1} (1-z)^a (1+z)^b dz = 2^(a+b+1) beta(a+1,b+1) Parameters __________ - a,b > -1 - log: optional - + a, b : ints or floats + Jacobi parameters. + log: bool, optional + Return natural logarithm of mass. Default: False. + dtype : dtype, optional + Data type. Default: module-level DEFAULT_GRID_DTYPE. + + Returns + ------- + mass : float + Integral of Jacobi measure. """ + if dtype is None: + dtype = DEFAULT_GRID_DTYPE + # beta and betaln do not natively support xprec, so outputs must be casted + type = np.dtype(dtype).type + if log: + return (a + b + 1) * type(np.log(2)) + type(betaln(a + 1, b + 1)) + else: + return 2 ** (a + b + 1) * type(beta(a + 1, b + 1)) - if not log: - from scipy.special import beta - return 2**(a+b+1)*beta(a+1,b+1) - - from scipy.special import betaln - return (a+b+1)*np.log(2) + betaln(a+1,b+1) -def norm_ratio(dn,da,db,n,a,b,squared=False): +def norm_ratio(dn, da, db, n, a, b, squared=False): """ - Ratio of classical Jacobi normalisation. + Ratio of classical Jacobi normalization. N(n,a,b) = integrate_(-1,+1)( (1-z)**a (1+z)**b P(n,a,b,z)**2 ) @@ -221,69 +410,89 @@ def norm_ratio(dn,da,db,n,a,b,squared=False): Parameters ---------- - dn,da,db: int - n: np.ndarray, int > 0 - a,b: float > -1 - squared: return N(n,a,b) or sqrt(N(n,a,b)) (defalut) - + dn, da, db : int + Increments in n, a, b + n : int or array + Starting polynomial degrees + a, b : float + Starting Jacobi parameters. + squared : bool, optinoal + True for direct norm ratio, False for square root of norm ratio. Default: False. + + Returns + ------- + ratios : float or array + Norm ratios for specified n. """ - - if not all(type(d) == int for d in (dn,da,db)): + if not all(type(d) == int for d in (dn, da, db)): raise TypeError('can only increment by integers.') - def tricky(n,a,b): + def tricky(n, a, b): # 0/0 = 1 - if a+b != -1: - return (2*n+a+b+1)/(n+a+b+1) - return 2 - (n==0) - - def n_ratio(d,n,a,b): - if d < 0: return 1/n_ratio(-d,n+d,a,b) - if d == 0: return 1 + 0*n + if a + b != -1: + return (2 * n + a + b + 1) / (n + a + b + 1) + return 2 - (n == 0) + + def n_ratio(d, n, a, b): + if d < 0: + return 1 / n_ratio(-d, n + d, a, b) + if d == 0: + return 1 + 0 * n if d == 1: - return ((n+a+1)*(n+b+1)/((n+1)*(2*n+a+b+3))) * tricky(n,a,b) - return n_ratio(1,n+d-1,a,b)*n_ratio(d-1,n,a,b) - - def ab_ratio(d,n,a,b): - if d < 0: return 1/ab_ratio(-d,n,a+d,b) - if d == 0: return 1 + 0*n + return ((n + a + 1) * (n + b + 1) / ((n + 1) * (2 * n + a + b + 3))) * tricky(n, a, b) + return n_ratio(1, n + d - 1, a, b) * n_ratio(d - 1, n, a, b) + + def ab_ratio(d, n, a, b): + if d < 0: + return 1 / ab_ratio(-d, n, a + d, b) + if d == 0: + return 1 + 0 * n if d == 1: - return (2*(n+a+1)/(2*n+a+b+2)) * tricky(n,a,b) - return ab_ratio(1,n,a+d-1,b)*ab_ratio(d-1,n,a,b) + return (2 * (n + a + 1) / (2 * n + a + b + 2)) * tricky(n, a, b) + return ab_ratio(1, n, a + d - 1, b) * ab_ratio(d - 1, n, a, b) - ratio = n_ratio(dn,n,a+da,b+db)*ab_ratio(da,n,a,b+db)*ab_ratio(db,n,b,a) - - if not squared: + ratio = n_ratio(dn, n, a + da, b + db) * ab_ratio(da, n, a, b + db) * ab_ratio(db, n, b, a) + if squared: + return ratio + else: return np.sqrt(ratio) - return ratio -def operator(name,normalised=True,dtype=dtype): +def operator(name, **kw): """ - Interface to base JacobiOperator class. + Build Jacobi operators by name. Parameters ---------- - name: A, B, C, D, Id, Pi, N, Z (Jacobi matrix) - normalised: True --> unit-integral, False --> classical. - dtype: output dtype - + name : str + Jacobi operator name: 'A', 'B', 'C', 'D', 'Id', 'Pi', 'N', or 'Z' + normalized : bool, optional + True for unit-integral normalization, False for classical normalization. Default: True. + **kw : dict, optional + Other keywords passed to JacobiOperator. + + Returns + ------- + operator : JacobiOperator + Specified operator. """ if name == 'Id': - return JacobiOperator.identity(dtype=dtype) - if name == 'Pi': - return JacobiOperator.parity(dtype=dtype) - if name == 'N': - return JacobiOperator.number(dtype=dtype) - if name == 'Z': - A = JacobiOperator('A',normalised=normalised,dtype=dtype) - B = JacobiOperator('B',normalised=normalised,dtype=dtype) - return (B(-1) @ B(+1) - A(-1) @ A(+1))/2 - return JacobiOperator(name,normalised=normalised,dtype=dtype) - -class JacobiOperator(): + return JacobiOperator.identity(**kw) + elif name == 'Pi': + return JacobiOperator.parity(**kw) + elif name == 'N': + return JacobiOperator.number(**kw) + elif name == 'Z': + A = JacobiOperator('A', **kw) + B = JacobiOperator('B', **kw) + return (B(-1) @ B(+1) - A(-1) @ A(+1)) / 2 + else: + return JacobiOperator(name, **kw) + + +class JacobiOperator: """ - The base class for primary operators acting on finite row vectors of Jacobi polynomials. + Operators acting on finite row vectors of Jacobi polynomials. n+{self[0]},a->a+{self[1]},b->b+{self[2]})' - if self[3]: s = s.replace('a->a','a->b').replace('b->b','b->a') - return s.replace('+0','').replace('+-','-') + if self[3]: + s = s.replace('a->a', 'a->b').replace('b->b', 'b->a') + return s.replace('+0', '').replace('+-', '-') - def __add__(self,other): - return self.Output(*self(*other[:3],evaluate=False),self[3]^other[3]) + def __add__(self, other): + return self.Output(*self(*other[:3], evaluate=False), self[3] ^ other[3]) - def __call__(self,*args,evaluate=True): - n,a,b = args[:3] - if self[3]: a,b = b,a + def __call__(self, *args, evaluate=True): + n, a, b = args[:3] + if self[3]: + a, b = b, a n, a, b = self[0] + n, self[1] + a, self[2] + b if evaluate and (a <= -1 or b <= -1): raise ValueError('invalid Jacobi parameter.') - return n,a,b + return n, a, b def __neg__(self): - a,b = -self[1],-self[2] - if self[3]: a,b = b,a - return self.Output(-self[0],a,b,self[3]) + a, b = -self[1], -self[2] + if self[3]: + a, b = b, a + return self.Output(-self[0], a, b, self[3]) - def __eq__(self,other): + def __eq__(self, other): return self[1:] == other[1:] - def __or__(self,other): + def __or__(self, other): if self != other: raise TypeError('operators have incompatible codomains.') if self[0] >= other[0]: diff --git a/dedalus/libraries/dedalus_sphere/old_intertwiner.py b/dedalus/libraries/dedalus_sphere/old_intertwiner.py index 4f8c0684..5fc80968 100644 --- a/dedalus/libraries/dedalus_sphere/old_intertwiner.py +++ b/dedalus/libraries/dedalus_sphere/old_intertwiner.py @@ -1,54 +1,65 @@ import numpy as np -def forbidden_spin(ell,spin): - if type(spin) == int: spin = [spin] + +def forbidden_spin(ell, spin): + if type(spin) == int: + spin = [spin] return ell < abs(sum(spin)) -def forbidden_regularity(ell,regularity): - if type(regularity) == int: regularity = [regularity] + +def forbidden_regularity(ell, regularity): + if type(regularity) == int: + regularity = [regularity] walk = [ell] for r in regularity[::-1]: walk += [walk[-1] + r] - if walk[-1] < 0 or ((walk[-1] == 0) and (walk[-2] == 0)): return True + if walk[-1] < 0 or ((walk[-1] == 0) and (walk[-2] == 0)): + return True return False -def _replace(t,i,nu): - return tuple(nu if i==j else t[j] for j in range(len(t))) -def regularity2spinMap(ell,spin,regularity): +def _replace(t, i, nu): + return tuple(nu if i == j else t[j] for j in range(len(t))) + - if spin == (): return 1 +def regularity2spinMap(ell, spin, regularity): + + if spin == (): + return 1 + + if forbidden_spin(ell, spin) or forbidden_regularity(ell, regularity): + return 0 - if forbidden_spin(ell,spin) or forbidden_regularity(ell,regularity): return 0 - if type(spin) == int: rank = 1 sigma, a = spin, regularity - tau, b = (), () + tau, b = (), () else: rank = len(spin) - sigma, a = spin[0], regularity[0] - tau, b = spin[1:], regularity[1:] + sigma, a = spin[0], regularity[0] + tau, b = spin[1:], regularity[1:] R = 0 - for i in range(rank-1): + for i in range(rank - 1): if tau[i] == -sigma: - R -= regularity2spinMap(ell,_replace(tau,i,0),b) + R -= regularity2spinMap(ell, _replace(tau, i, 0), b) if tau[i] == 0: - R += regularity2spinMap(ell,_replace(tau,i,sigma),b) - - Qold = regularity2spinMap(ell,tau,b) - - degree = ell+sum(b) - kangle = -sigma*np.sqrt((ell-sigma*sum(tau))*(ell+sigma*sum(tau)+1)/2) + R += regularity2spinMap(ell, _replace(tau, i, sigma), b) - R -= kangle*Qold - if sigma != 0: Qold = 0 + Qold = regularity2spinMap(ell, tau, b) - if a == -1: return (Qold*degree - R)/np.sqrt(degree*(2*degree+1)) - if a == 0: return sigma*R/np.sqrt(degree*(degree+1)) - if a == +1: return (Qold*(degree+1) + R)/np.sqrt((degree+1)*(2*degree+1)) + degree = ell + sum(b) + kangle = -sigma * np.sqrt((ell - sigma * sum(tau)) * (ell + sigma * sum(tau) + 1) / 2) + R -= kangle * Qold + if sigma != 0: + Qold = 0 + if a == -1: + return (Qold * degree - R) / np.sqrt(degree * (2 * degree + 1)) + if a == 0: + return sigma * R / np.sqrt(degree * (degree + 1)) + if a == +1: + return (Qold * (degree + 1) + R) / np.sqrt((degree + 1) * (2 * degree + 1)) diff --git a/dedalus/libraries/dedalus_sphere/operators.py b/dedalus/libraries/dedalus_sphere/operators.py index 8698da07..fbee762b 100644 --- a/dedalus/libraries/dedalus_sphere/operators.py +++ b/dedalus/libraries/dedalus_sphere/operators.py @@ -3,7 +3,8 @@ from scipy.sparse import coo_matrix from scipy.sparse import identity as id_matrix -class Operator(): + +class Operator: """ Class for deffered (lazy) evaluation of matrix-valued functions between parameterised vector spaces. @@ -74,12 +75,13 @@ class Operator(): """ - def __init__(self,function,codomain,Output=None): - if Output == None: Output = Operator + def __init__(self, function, codomain, Output=None): + if Output == None: + Output = Operator self.__function = function self.__codomain = codomain - self.__Output = Output + self.__Output = Output @property def function(self): @@ -93,40 +95,45 @@ def codomain(self): def Output(self): return self.__Output - def __call__(self,*args): + def __call__(self, *args): return self.__function(*args) - def __matmul__(self,other): + def __matmul__(self, other): def function(*args): return self(*other.codomain(*args)) @ other(*args) + return self.Output(function, self.codomain + other.codomain) @property def T(self): codomain = -self.codomain + def function(*args): return self(*codomain(*args)).T - return self.Output(function,codomain) + + return self.Output(function, codomain) @property def identity(self): def function(*args): return self(*args).identity - return self.Output(function,0*self.codomain) - def __pow__(self,exponent): + return self.Output(function, 0 * self.codomain) + + def __pow__(self, exponent): if exponent < 0: raise TypeError('exponent must be a non-negative integer.') if exponent == 0: return self.identity - return self @ self**(exponent-1) + return self @ self ** (exponent - 1) - def __add__(self,other): + def __add__(self, other): - if other == 0: return self + if other == 0: + return self - if not isinstance(other,Operator): - other = other*self.identity + if not isinstance(other, Operator): + other = other * self.identity codomain = self.codomain | other.codomain @@ -135,45 +142,44 @@ def function(*args): return self.Output(function, codomain) - def __mul__(self,other): - if isinstance(other,Operator): + def __mul__(self, other): + if isinstance(other, Operator): return self @ other - other @ self def function(*args): - return other*self(*args) - return self.Output(function,self.codomain) - - # def function(*args): - # a,b = args[:len(self.codomain)], args[len(self.codomain):] - # return Kronecker(self(*a),other(*b)) - # codomain = Codomain(self.codomain,other.codomain) - # return Operator(function,codomain) + return other * self(*args) + return self.Output(function, self.codomain) + # def function(*args): + # a,b = args[:len(self.codomain)], args[len(self.codomain):] + # return Kronecker(self(*a),other(*b)) + # codomain = Codomain(self.codomain,other.codomain) + # return Operator(function,codomain) - def __radd__(self,other): + def __radd__(self, other): return self + other - def __rmul__(self,other): - return self*other + def __rmul__(self, other): + return self * other - def __truediv__(self,other): - return self * (1/other) + def __truediv__(self, other): + return self * (1 / other) def __pos__(self): return self def __neg__(self): - return (-1)*self + return (-1) * self - def __sub__(self,other): + def __sub__(self, other): return self + (-other) - def __rsub__(self,other): + def __rsub__(self, other): return -self + other -class Codomain(): +class Codomain: """Base class for Codomain objects. @@ -187,10 +193,11 @@ class Codomain(): """ - def __init__(self,*arrow,Output=None): - if Output == None: Output = Codomain + def __init__(self, *arrow, Output=None): + if Output == None: + Output = Codomain - self.__arrow = arrow + self.__arrow = arrow self.__Output = Output @property @@ -201,7 +208,7 @@ def arrow(self): def Output(self): return self.__Output - def __getitem__(self,item): + def __getitem__(self, item): return self.__arrow[(item)] def __len__(self): @@ -213,41 +220,40 @@ def __str__(self): def __repr__(self): return str(self) - def __add__(self,other): - return self.Output(*tuple(a+b for a,b in zip(self[:],other[:]))) + def __add__(self, other): + return self.Output(*tuple(a + b for a, b in zip(self[:], other[:]))) - def __call__(self,*args): - return tuple(a+b for a,b in zip(self.arrow,args)) + def __call__(self, *args): + return tuple(a + b for a, b in zip(self.arrow, args)) - def __eq__(self,other): + def __eq__(self, other): return self[:] == other[:] - def __or__(self,other): + def __or__(self, other): if self != other: raise TypeError('operators have incompatible codomains.') - return self.Output(*tuple(a|b for a,b in zip(self[:],other[:]))) + return self.Output(*tuple(a | b for a, b in zip(self[:], other[:]))) def __neg__(self): return self.Output(*tuple(-a for a in self.arrow)) - def __mul__(self,other): + def __mul__(self, other): if type(other) != int: raise TypeError('only integer multiplication defined.') if other == 0: - return self.Output(*(len(self.arrow)*(0,))) + return self.Output(*(len(self.arrow) * (0,))) if other < 0: - return -self + (other+1)*self + return -self + (other + 1) * self - def __rmul__(self,other): - return self*other + def __rmul__(self, other): + return self * other - def __sub__(self,other): + def __sub__(self, other): return self + (-other) - class infinite_csr(csr_matrix): """ Base class for extendable addition with csr_matrix types. @@ -275,15 +281,15 @@ class infinite_csr(csr_matrix): """ - def __init__(self,*args,**kwargs): - csr_matrix.__init__(self,*args,**kwargs) + def __init__(self, *args, **kwargs): + csr_matrix.__init__(self, *args, **kwargs) def __repr__(self): s = csr_matrix(self).__repr__() i = s.find('sparse matrix') j = s.find('with') k = s.find('stored elements') - return s[:i] + 'Infinite Compressed Sparse Row matrix; ' + s[j:k+15] + '>' + return s[:i] + 'Infinite Compressed Sparse Row matrix; ' + s[j : k + 15] + '>' @property def T(self): @@ -295,9 +301,9 @@ def identity(self): @property def square(self): - return csr_matrix(self[:self.shape[1]]) + return csr_matrix(self[: self.shape[1]]) - def __getitem__(self,item): + def __getitem__(self, item): if type(item) == tuple: i, j = item[0], item[1:] @@ -305,7 +311,7 @@ def __getitem__(self,item): i, j = item, () if type(i) != slice: - i = slice(i, i+1, 1) + i = slice(i, i + 1, 1) s = self.shape @@ -313,12 +319,11 @@ def __getitem__(self,item): return infinite_csr(csr_matrix(self)[item]) coo_self = coo_matrix(self) - output = coo_matrix((coo_self.data, (coo_self.row, coo_self.col)), shape=(i.stop+1,s[1]), dtype=self.dtype) + output = coo_matrix((coo_self.data, (coo_self.row, coo_self.col)), shape=(i.stop + 1, s[1]), dtype=self.dtype) output = lil_matrix(output) return infinite_csr(output[(i,) + j]) - - def __add__(self,other): + def __add__(self, other): ns, no = self.shape[0], other.shape[0] @@ -333,8 +338,7 @@ def __add__(self,other): sum_ = lil_matrix(other) sum_[:ns] += self - return infinite_csr(sum_) - def __radd__(self,other): + def __radd__(self, other): return self + other diff --git a/dedalus/libraries/dedalus_sphere/shell.py b/dedalus/libraries/dedalus_sphere/shell.py index 78fd378e..92d4a8af 100644 --- a/dedalus/libraries/dedalus_sphere/shell.py +++ b/dedalus/libraries/dedalus_sphere/shell.py @@ -1,12 +1,13 @@ import numpy as np -from . import jacobi as Jacobi -from .operators import Operator, Codomain, infinite_csr +from . import jacobi as Jacobi +from .operators import Operator, Codomain, infinite_csr # The defalut configuration for the base Jacobi parameter. -alpha = (-1/2,-1/2) +alpha = (-1 / 2, -1 / 2) -def operator(dimension,radii,name,alpha=alpha): + +def operator(dimension, radii, name, alpha=alpha): """ Shell.operator function. @@ -15,71 +16,79 @@ def operator(dimension,radii,name,alpha=alpha): """ - width = radii[1] - radii[0] + width = radii[1] - radii[0] if name == 'Z': + def Z(n, k): - return Jacobi.operator('Z')(n, k+alpha[0], k+alpha[1]) + return Jacobi.operator('Z')(n, k + alpha[0], k + alpha[1]) + return Operator(Z, ShellCodomain(0, 0)) - Z = (radii[1] + radii[0])/width + Jacobi.operator('Z') + Z = (radii[1] + radii[0]) / width + Jacobi.operator('Z') if name == 'Id': - def I(n,k): - return Jacobi.operator('Id')(n,k+alpha[0],k+alpha[1]) - return Operator(I,ShellCodomain(0,0)) + + def I(n, k): + return Jacobi.operator('Id')(n, k + alpha[0], k + alpha[1]) + + return Operator(I, ShellCodomain(0, 0)) if name == 'R': - def R(n,k): - return (0.5*width)*Z(n,k+alpha[0],k+alpha[1]) - return Operator(R,ShellCodomain(1,0)) + + def R(n, k): + return (0.5 * width) * Z(n, k + alpha[0], k + alpha[1]) + + return Operator(R, ShellCodomain(1, 0)) AB = Jacobi.operator('A')(+1) @ Jacobi.operator('B')(+1) if name == 'AB': + def AB_(n, k): - return AB(n, k+alpha[0], k+alpha[1]) + return AB(n, k + alpha[0], k + alpha[1]) + return Operator(AB_, ShellCodomain(0, 1)) if name == 'E': - def E(n,k): - return 0.5 * (AB @ Z)(n,k+alpha[0],k+alpha[1]) - return Operator(E,ShellCodomain(1,1)) - if name == 'D': + def E(n, k): + return 0.5 * (AB @ Z)(n, k + alpha[0], k + alpha[1]) - def D(dl,l): + return Operator(E, ShellCodomain(1, 1)) - def D(n,k): + if name == 'D': + + def D(dl, l): + def D(n, k): D = Jacobi.operator('D')(+1) @ Z K = Jacobi.operator('A')(0) - alpha[0] - K += dl*l + (dl == -1)*(2-dimension) + K += dl * l + (dl == -1) * (2 - dimension) - D = ( D - K @ AB )/width + D = (D - K @ AB) / width - return D(n,k+alpha[0],k+alpha[1]) + return D(n, k + alpha[0], k + alpha[1]) - return Operator(D,ShellCodomain(0,1)) + return Operator(D, ShellCodomain(0, 1)) return D class ShellCodomain(Codomain): - - def __init__(self,dn=0,dk=0): - Codomain.__init__(self,dn,dk,Output=ShellCodomain) + def __init__(self, dn=0, dk=0): + Codomain.__init__(self, dn, dk, Output=ShellCodomain) def __str__(self): s = f'(n->n+{self[0]},k->k+{self[1]})' - return s.replace('+0','').replace('+-','-') + return s.replace('+0', '').replace('+-', '-') - def __eq__(self,other): + def __eq__(self, other): return self[1] == other[1] - def __or__(self,other): + def __or__(self, other): if self != other: raise TypeError('operators have incompatible codomains.') if self[0] >= other[0]: diff --git a/dedalus/libraries/dedalus_sphere/sphere.py b/dedalus/libraries/dedalus_sphere/sphere.py index f653813f..74e473f9 100644 --- a/dedalus/libraries/dedalus_sphere/sphere.py +++ b/dedalus/libraries/dedalus_sphere/sphere.py @@ -1,113 +1,165 @@ -import numpy as np -from . import jacobi as Jacobi -from scipy.sparse import dia_matrix as banded -from .operators import Operator, infinite_csr +"""Spin-weighted spherical harmonic functions.""" -dtype = 'longdouble' +import numpy as np +from scipy.sparse import dia_matrix +from . import jacobi +from .operators import Operator, infinite_csr -def quadrature(Lmax,dtype=dtype): - """Generates the Gauss quadrature grid and weights for spherical harmonics transform. - Returns cos_theta, weights - Will integrate polynomials on (-1,+1) exactly up to degree = 2*Lmax+1. +def quadrature(Lmax, **kw): + """ + Gauss quadrature grid and weights for SWSH transform. + Will exactly integrate polynomials up to degree 2*Lmax+1. Parameters ---------- - Lmax: int >=0; spherical-harmonic degree. - + Lmax : int + Maximum spherical harmonic degree. + **kw : dict, optional + Other keywords passed to jacobi.quadrature. + + Returns + ------- + cos_theta : array + Quadrature grid in z = cos(theta). + w : array + Quadrature weights. """ - - return Jacobi.quadrature(Lmax+1,0,0,dtype=dtype) - - -def spin2Jacobi(Lmax,m,s,ds=None,dm=None): - - n = Lmax + 1 - max(abs(m),abs(s)) - a, b = abs(m+s), abs(m-s) - - if ds == dm == None: - return n,a,b - - if ds == None: ds = 0 - if dm == None: dm = 0 - + return jacobi.quadrature(Lmax + 1, a=0, b=0, **kw) + + +def spin2Jacobi(Lmax, m, s, ds=None, dm=None): + """Compute Jacobi parameters from SWSH parameters.""" + n = Lmax + 1 - max(abs(m), abs(s)) + a, b = abs(m + s), abs(m - s) + if ds is None and dm is None: + return n, a, b + if ds is None: + ds = 0 + if dm is None: + dm = 0 m += dm s += ds + dn = Lmax + 1 - max(abs(m), abs(s)) - n + da, db = abs(m + s) - a, abs(m - s) - b + return n, a, b, dn, da, db - dn = Lmax + 1 - max(abs(m),abs(s)) - n - da,db = abs(m+s) - a, abs(m-s) - b - - return n,a,b,dn,da,db - -def harmonics(Lmax,m,s,cos_theta,**kwargs): +def harmonics(Lmax, m, s, cos_theta, **kw): """ - Gives spin-wieghted spherical harmonic functions on the Gauss quaduature grid. - Returns an array with - shape = ( Lmax - Lmin(m,s) + 1, len(z) ) - or (Lmax - Lmin(m,s) + 1,) if z is a single point. - - Parameters - ---------- - Lmax: int >=0; spherical-harmonic degree. - m,s : int - spherical harmonic parameters. - cos_theta: np.ndarray or float. - dtype: output dtype. internal dtype = 'longdouble'. - """ - - n,a,b = spin2Jacobi(Lmax,m,s) + SWSH of type (m,s) up to degree Lmax. - init = np.exp(0.5*Jacobi.measure(a,b,cos_theta,log=True)) - init *= ((-1.)**max(m,-s)) - - return Jacobi.polynomials(n,a,b,cos_theta,init,**kwargs) - - -def operator(name,dtype=dtype): + Parameters + ---------- + Lmax : int + Maximum spherical-harmonic degree. + m, s : int + SWSH parameters. + cos_theta : float or array + Grid array in z = cos(theta). + **kw : dict, optional + Other keywords passed to jacobi.polynomials. + + Returns + ------- + polynomials : array + Array of polynomials evaluated at specified z points. + First axis is polynomial degree. + """ + # Compute Jacobi parameters + n, a, b = spin2Jacobi(Lmax, m, s) + # Build envelope + if np.isscalar(cos_theta): + # Call non-log version to avoid issues with |z| = 1 + init = np.sqrt(jacobi.measure(a, b, cos_theta, dtype=np.float64)) + else: + init = np.exp(0.5 * jacobi.measure(a, b, cos_theta, log=True)) + init *= (-1.0) ** max(m, -s) + # Build polynomials + return jacobi.polynomials(n, a, b, cos_theta, init, **kw) + + +def generate_harmonics(Lmax, m, s, cos_theta, **kw): """ - Interface to base ShereOperator class. + SWSH of type (m,s) up to degree Lmax. Parameters ---------- - + Lmax : int + Maximum spherical-harmonic degree. + m, s : int + SWSH parameters. + cos_theta : float or array + Grid array in z = cos(theta). + **kw : dict, optional + Other keywords passed to jacobi.polynomials. + + Returns + ------- + polynomials : array + Array of polynomials evaluated at specified z points. + First axis is polynomial degree. + """ + # Compute Jacobi parameters + n, a, b = spin2Jacobi(Lmax, m, s) + # Build envelope + if np.isscalar(cos_theta): + # Call non-log version to avoid issues with |z| = 1 + init = np.sqrt(jacobi.measure(a, b, cos_theta, dtype=np.float64)) + else: + init = np.exp(0.5 * jacobi.measure(a, b, cos_theta, log=True)) + init *= (-1.0) ** max(m, -s) + # Generate polynomials + yield from jacobi.generate_polynomials(n, a, b, cos_theta, init, **kw) + + +def operator(name, **kw): """ + Build SWSH operators by name. + Parameters + ---------- + name : str + Sphere operator name: 'D', 'Sin', 'Cos', 'Id', 'Pi', 'L', 'M', or 'S' + **kw : dict, optional + Other keywords passed to SphereOperator or JacobiOperator. + + Returns + ------- + operator : SphereOperator or JacobiOperator + Specified operator. + """ if name == 'Id': - return SphereOperator.identity(dtype=dtype) - + return SphereOperator.identity(**kw) if name == 'Pi': - return SphereOperator.parity(dtype=dtype) - + return SphereOperator.parity(**kw) if name == 'L': - return SphereOperator.L(dtype=dtype) - + return SphereOperator.L(**kw) if name == 'M': - return SphereOperator.M(dtype=dtype) - + return SphereOperator.M(**kw) if name == 'S': - return SphereOperator.S(dtype=dtype) - + return SphereOperator.S(**kw) if name == 'Cos': - def Cos(Lmax,m,s): - return Jacobi.operator('Z',dtype=dtype)(*spin2Jacobi(Lmax,m,s)) - #return Jacobi.operator('Z',dtype=dtype)(Lmax+1, abs(m+s), abs(m-s)) - return Operator(Cos,SphereCodomain(1,0,0,0)) - - return SphereOperator(name,dtype=dtype) + def Cos(Lmax, m, s): + return jacobi.operator('Z', **kw)(*spin2Jacobi(Lmax, m, s)) + #return Jacobi.operator('Z',dtype=dtype)(Lmax+1, abs(m+s), abs(m-s)) -class SphereOperator(): + return Operator(Cos, SphereCodomain(1, 0, 0, 0)) + return SphereOperator(name, **kw) - def __init__(self,name,radius=1,dtype=dtype): - self.__function = getattr(self,f'_SphereOperator__{name}') +class SphereOperator: + """Operators acting on finite row vectors of SWSH.""" + def __init__(self, name, radius=1, dtype=None): + self.__function = getattr(self, f'_SphereOperator__{name}') self.__radius = radius - + if dtype is None: + dtype = jacobi.DEFAULT_OPERATOR_DTYPE self.__dtype = dtype - def __call__(self,ds): + def __call__(self, ds): return Operator(*self.__function(ds)) @property @@ -118,71 +170,78 @@ def radius(self): def dtype(self): return self.__dtype - def __D(self,ds): - - def D(Lmax,m,s): - - n,a,b,dn,da,db = spin2Jacobi(Lmax,m,s,ds=ds) - - D = Jacobi.operator('C' if da+db == 0 else 'D',dtype=self.dtype)(da) - - return (-ds*np.sqrt(0.5)/self.radius)*D(n,a,b) - - return D, SphereCodomain(0,0,ds,0) + def __D(self, ds): + def D(Lmax, m, s): + n, a, b, dn, da, db = spin2Jacobi(Lmax, m, s, ds=ds) + D = jacobi.operator('C' if da + db == 0 else 'D', dtype=self.dtype)(da) + return (-ds * np.sqrt(0.5) / self.radius) * D(n, a, b) - def __Sin(self,ds): + return D, SphereCodomain(0, 0, ds, 0) - def Sin(Lmax,m,s): + def __Sin(self, ds): + def Sin(Lmax, m, s): + n, a, b, dn, da, db = spin2Jacobi(Lmax, m, s, ds=ds) + S = jacobi.operator('A', dtype=self.dtype)(da) + S = S @ jacobi.operator('B', dtype=self.dtype)(db) + return (da * ds) * S(n, a, b) - n,a,b,dn,da,db = spin2Jacobi(Lmax,m,s,ds=ds) - - S = Jacobi.operator('A',dtype=self.dtype)(da) - S = S @ Jacobi.operator('B',dtype=self.dtype)(db) - - return (da*ds) * S(n,a,b) - - return Sin, SphereCodomain(1,0,ds,0) + return Sin, SphereCodomain(1, 0, ds, 0) @staticmethod - def identity(dtype=dtype): + def identity(dtype=None): + if dtype is None: + dtype = jacobi.DEFAULT_OPERATOR_DTYPE - def I(Lmax,m,s): - n = spin2Jacobi(Lmax,m,s)[0] - N = np.ones(n,dtype=dtype) - return infinite_csr(banded((N,[0]),(max(n,0),max(n,0)))) + def I(Lmax, m, s): + n = spin2Jacobi(Lmax, m, s)[0] + N = np.ones(n, dtype=dtype) + return infinite_csr(dia_matrix((N, [0]), (max(n, 0), max(n, 0)))) - return Operator(I,SphereCodomain(0,0,0,0)) + return Operator(I, SphereCodomain(0, 0, 0, 0)) @staticmethod - def parity(dtype=dtype): + def parity(dtype=None): + if dtype is None: + dtype = jacobi.DEFAULT_OPERATOR_DTYPE - def Pi(Lmax,m,s): - return Jacobi.operator('Pi',dtype=dtype)(*spin2Jacobi(Lmax,m,s)) + def Pi(Lmax, m, s): + return jacobi.operator('Pi', dtype=dtype)(*spin2Jacobi(Lmax, m, s)) - return Operator(Pi,SphereCodomain(0,0,0,1)) + return Operator(Pi, SphereCodomain(0, 0, 0, 1)) @staticmethod - def L(dtype=dtype): + def L(dtype=None): + if dtype is None: + dtype = jacobi.DEFAULT_OPERATOR_DTYPE - def L(Lmax,m,s): - n = spin2Jacobi(Lmax,m,s)[0] - N = np.arange(Lmax+1-n,Lmax+1,dtype=dtype) - return infinite_csr(banded((N,[0]),(max(n,0),max(n,0)))) + def L(Lmax, m, s): + n = spin2Jacobi(Lmax, m, s)[0] + N = np.arange(Lmax + 1 - n, Lmax + 1, dtype=dtype) + return infinite_csr(dia_matrix((N, [0]), (max(n, 0), max(n, 0)))) - return Operator(L,SphereCodomain(0,0,0,0)) + return Operator(L, SphereCodomain(0, 0, 0, 0)) @staticmethod - def M(dtype=dtype): + def M(dtype=None): + if dtype is None: + dtype = jacobi.DEFAULT_OPERATOR_DTYPE - def M(Lmax,m,s): - n = spin2Jacobi(Lmax,m,s)[0] - N = m*np.ones(n,dtype=dtype) - return infinite_csr(banded((N,[0]),(max(n,0),max(n,0)))) + def M(Lmax, m, s): + n = spin2Jacobi(Lmax, m, s)[0] + N = m * np.ones(n, dtype=dtype) + return infinite_csr(dia_matrix((N, [0]), (max(n, 0), max(n, 0)))) - return Operator(M,SphereCodomain(0,0,0,0)) + return Operator(M, SphereCodomain(0, 0, 0, 0)) @staticmethod - def S(dtype=dtype): + def S(dtype=None): + if dtype is None: + dtype = jacobi.DEFAULT_OPERATOR_DTYPE + + def S(Lmax, m, s): + n = spin2Jacobi(Lmax, m, s)[0] + N = s * np.ones(n, dtype=dtype) + return infinite_csr(dia_matrix((N, [0]), (max(n, 0), max(n, 0)))) def S(Lmax,m,s): n = spin2Jacobi(Lmax,m,s)[0] @@ -192,23 +251,26 @@ def S(Lmax,m,s): return Operator(S,SphereCodomain(0,0,0,0)) -class SphereCodomain(Jacobi.JacobiCodomain): +class SphereCodomain(jacobi.JacobiCodomain): + """Sphere codomain.""" - def __init__(self,dL=0,dm=0,ds=0,pi=0): - Jacobi.JacobiCodomain.__init__(self,*(dL,dm,ds,pi),Output=SphereCodomain) + def __init__(self, dL=0, dm=0, ds=0, pi=0): + jacobi.JacobiCodomain.__init__(self, *(dL, dm, ds, pi), Output=SphereCodomain) def __str__(self): s = f'(L->L+{self[0]},m->m+{self[1]},s->s+{self[2]})' - if self[3]: s = s.replace('s->s','s->-s') - return s.replace('+0','').replace('+-','-') - - def __call__(self,*args,evaluate=True): - L,m,s = args[:3] - if self[3]: s *= -1 + if self[3]: + s = s.replace('s->s', 's->-s') + return s.replace('+0', '').replace('+-', '-') + + def __call__(self, *args, evaluate=True): + L, m, s = args[:3] + if self[3]: + s *= -1 return self[0] + L, self[1] + m, self[2] + s def __neg__(self): - m,s = -self[1],-self[2] - if self[3]: s *= -1 - return SphereCodomain(-self[0],m,s,self[3]) - + m, s = -self[1], -self[2] + if self[3]: + s *= -1 + return SphereCodomain(-self[0], m, s, self[3]) diff --git a/dedalus/libraries/dedalus_sphere/sphere_wrapper.py b/dedalus/libraries/dedalus_sphere/sphere_wrapper.py index f96b7a75..03618366 100644 --- a/dedalus/libraries/dedalus_sphere/sphere_wrapper.py +++ b/dedalus/libraries/dedalus_sphere/sphere_wrapper.py @@ -1,133 +1,134 @@ - -import numpy as np +import numpy as np from . import sphere128 as sph from dedalus.tools.cache import CachedMethod + class Sphere: - def __init__(self,L_max,S_max=0,N_theta=None,m_min=None,m_max=None): - self.L_max, self.S_max = L_max, S_max - if N_theta == None: N_theta = L_max+1 + def __init__(self, L_max, S_max=0, N_theta=None, m_min=None, m_max=None): + self.L_max, self.S_max = L_max, S_max + if N_theta == None: + N_theta = L_max + 1 self.N_theta = N_theta - if m_min == None: m_min = -L_max - if m_max == None: m_max = L_max + if m_min == None: + m_min = -L_max + if m_max == None: + m_max = L_max # grid and weights for the all transforms - self.cos_grid,self.weights = sph.quadrature(self.N_theta-1,niter=3,report_error=False) + self.cos_grid, self.weights = sph.quadrature(self.N_theta - 1, niter=3, report_error=False) self.grid = np.arccos(self.cos_grid) - self.sin_grid = np.sqrt(1-self.cos_grid**2) + self.sin_grid = np.sqrt(1 - self.cos_grid**2) self.pushY, self.pullY = {}, {} - for s in range(-S_max,S_max+1): - for m in range(m_min,m_max+1): - Y = sph.Y(self.L_max,m,s,self.cos_grid) - self.pushY[(m,s)] = (self.weights*Y).astype(np.float64) - self.pullY[(m,s)] = (Y.T).astype(np.float64) + for s in range(-S_max, S_max + 1): + for m in range(m_min, m_max + 1): + Y = sph.Y(self.L_max, m, s, self.cos_grid) + self.pushY[(m, s)] = (self.weights * Y).astype(np.float64) + self.pullY[(m, s)] = (Y.T).astype(np.float64) # downcast to double precision - self.grid = self.grid.astype(np.float64) - self.weights = self.weights.astype(np.float64) + self.grid = self.grid.astype(np.float64) + self.weights = self.weights.astype(np.float64) self.sin_grid = self.sin_grid.astype(np.float64) self.cos_grid = self.cos_grid.astype(np.float64) @CachedMethod - def op(self,op_name,m,s): - return sph.operator(op_name,self.L_max,m,s).astype(np.float64) + def op(self, op_name, m, s): + return sph.operator(op_name, self.L_max, m, s).astype(np.float64) @CachedMethod - def L_min(self,m,s): - return sph.L_min(m,s) + def L_min(self, m, s): + return sph.L_min(m, s) - def zeros(self,m,s_out,s_in): - return sph.zeros(self.L_max,m,s_out,s_in) + def zeros(self, m, s_out, s_in): + return sph.zeros(self.L_max, m, s_out, s_in) - def forward_spin(self,m,s,data): + def forward_spin(self, m, s, data): # grid --> coefficients - return self.pushY[(m,s)].dot(data) + return self.pushY[(m, s)].dot(data) - def backward_spin(self,m,s,data): + def backward_spin(self, m, s, data): # coefficients --> grid - return self.pullY[(m,s)].dot(data) + return self.pullY[(m, s)].dot(data) @CachedMethod - def tensor_index(self,m,rank): + def tensor_index(self, m, rank): num = np.arange(2**rank) - spin = (-1)**num - for k in range(2,rank+1): - spin += ((-1)**(num//2**(k-1))).astype(np.int64) + spin = (-1) ** num + for k in range(2, rank + 1): + spin += ((-1) ** (num // 2 ** (k - 1))).astype(np.int64) - if rank == 0: spin = [0] + if rank == 0: + spin = [0] start_index = [0] end_index = [] for k in range(2**rank): - end_index.append(start_index[k]+self.L_max-sph.L_min(m,spin[k])+1) - if k < 2**rank-1: + end_index.append(start_index[k] + self.L_max - sph.L_min(m, spin[k]) + 1) + if k < 2**rank - 1: start_index.append(end_index[k]) - return (start_index,end_index,spin) + return (start_index, end_index, spin) @CachedMethod - def unitary(self,rank=1,adjoint=False): - return sph.unitary(rank=rank,adjoint=adjoint) + def unitary(self, rank=1, adjoint=False): + return sph.unitary(rank=rank, adjoint=adjoint) - def forward(self,m,rank,data,unitary=None): + def forward(self, m, rank, data, unitary=None): if rank == 0: - return self.forward_spin(m,0,data) + return self.forward_spin(m, 0, data) - (start_index,end_index,spin) = self.tensor_index(m,rank) + (start_index, end_index, spin) = self.tensor_index(m, rank) if not unitary: - unitary = self.unitary(rank=rank,adjoint=True) + unitary = self.unitary(rank=rank, adjoint=True) - data = np.einsum("ij,j...->i...",unitary,data) + data = np.einsum("ij,j...->i...", unitary, data) shape = np.array(np.array(data).shape[1:]) shape[0] = end_index[-1] - data_c = np.zeros(shape,dtype=np.complex128) + data_c = np.zeros(shape, dtype=np.complex128) for i in range(2**rank): - data_c[start_index[i]:end_index[i]] = self.forward_spin(m,spin[i],data[i]) + data_c[start_index[i] : end_index[i]] = self.forward_spin(m, spin[i], data[i]) return data_c - def backward(self,m,rank,data,unitary=None): + def backward(self, m, rank, data, unitary=None): if rank == 0: - return self.backward_spin(m,0,data) + return self.backward_spin(m, 0, data) - (start_index,end_index,spin) = self.tensor_index(m,rank) + (start_index, end_index, spin) = self.tensor_index(m, rank) if not unitary: - unitary = self.unitary(rank=rank,adjoint=False) + unitary = self.unitary(rank=rank, adjoint=False) shape = np.array(np.array(data).shape) - shape = np.concatenate(([2**rank],shape)) + shape = np.concatenate(([2**rank], shape)) shape[1] = self.N_theta - data_g = np.zeros(shape,dtype=np.complex128) + data_g = np.zeros(shape, dtype=np.complex128) for i in range(2**rank): - data_g[i] = self.backward_spin(m,spin[i],data[start_index[i]:end_index[i]]) - return np.einsum("ij,j...->i...",unitary,data_g) + data_g[i] = self.backward_spin(m, spin[i], data[start_index[i] : end_index[i]]) + return np.einsum("ij,j...->i...", unitary, data_g) - def grad(self,m,rank_in,data,data_out): + def grad(self, m, rank_in, data, data_out): # data and data_out are in coefficient space - (start_index_in,end_index_in,spin_in) = self.tensor_index(m,rank_in) - rank_out = rank_in+1 - (start_index_out,end_index_out,spin_out) = self.tensor_index(m,rank_out) + (start_index_in, end_index_in, spin_in) = self.tensor_index(m, rank_in) + rank_out = rank_in + 1 + (start_index_out, end_index_out, spin_out) = self.tensor_index(m, rank_out) - half = 2**(rank_out-1) - for i in range(2**(rank_out)): - if i//half == 0: - operator = self.op('k+',m,spin_in[i%half]) + half = 2 ** (rank_out - 1) + for i in range(2 ** (rank_out)): + if i // half == 0: + operator = self.op('k+', m, spin_in[i % half]) else: - operator = self.op('k-',m,spin_in[i%half]) - - np.copyto( data_out[start_index_out[i]:end_index_out[i]], - operator.dot(data[start_index_in[i%half]:end_index_in[i%half]]) ) - + operator = self.op('k-', m, spin_in[i % half]) + np.copyto(data_out[start_index_out[i] : end_index_out[i]], operator.dot(data[start_index_in[i % half] : end_index_in[i % half]])) diff --git a/dedalus/libraries/dedalus_sphere/spin_operators.py b/dedalus/libraries/dedalus_sphere/spin_operators.py index a5733e40..83ccfffc 100644 --- a/dedalus/libraries/dedalus_sphere/spin_operators.py +++ b/dedalus/libraries/dedalus_sphere/spin_operators.py @@ -3,9 +3,10 @@ from .tuple_tools import * from .operators import Operator, Codomain -indexing = (-1,0,1) +indexing = (-1, 0, 1) threshold = 1e-12 + class TensorOperator(Operator): """ Class for lazy evaluation of spin/regularity tensor operations. @@ -34,9 +35,9 @@ class TensorOperator(Operator): """ - def __init__(self,function,codomain,indexing=indexing,threshold=threshold): - Operator.__init__(self,function,codomain,Output=TensorOperator) - self.__indexing = indexing + def __init__(self, function, codomain, indexing=indexing, threshold=threshold): + Operator.__init__(self, function, codomain, Output=TensorOperator) + self.__indexing = indexing self.__threshold = threshold @property @@ -51,28 +52,29 @@ def threshold(self): def dimension(self): return len(self.indexing) - def __call__(self,*args): + def __call__(self, *args): output = self.function(*args) - np.where(np.abs(output) < self.threshold,0,output) + np.where(np.abs(output) < self.threshold, 0, output) return output @int2tuple - def __getitem__(self,i): - sigma,tau = i[0],i[1] - i = tuple2index(sigma,self.indexing) - j = tuple2index(tau,self.indexing) - return self(len(tau))[i,j] + def __getitem__(self, i): + sigma, tau = i[0], i[1] + i = tuple2index(sigma, self.indexing) + j = tuple2index(tau, self.indexing) + return self(len(tau))[i, j] - def range(self,rank): - return product(*(rank*(self.indexing,))) + def range(self, rank): + return product(*(rank * (self.indexing,))) - def array(self,ranks): + def array(self, ranks): T = np.zeros(tuple(self.dimension**r for r in ranks)) for i, sigma in enumerate(self.range(ranks[0])): for j, tau in enumerate(self.range(ranks[1])): - T[i,j] = self[sigma,tau] + T[i, j] = self[sigma, tau] return T + class Identity(TensorOperator): """ Spin/regularity space identity transformation of arbitrary rank. @@ -83,13 +85,13 @@ class Identity(TensorOperator): """ - def __init__(self,**kwargs): + def __init__(self, **kwargs): - identity = lambda rank: self.array((rank,rank)) - TensorOperator.__init__(self,identity,TensorCodomain(0),**kwargs) + identity = lambda rank: self.array((rank, rank)) + TensorOperator.__init__(self, identity, TensorCodomain(0), **kwargs) @int2tuple - def __getitem__(self,i): + def __getitem__(self, i): return int(i[0] == i[1]) @@ -105,13 +107,13 @@ class Metric(TensorOperator): """ - def __init__(self,**kwargs): + def __init__(self, **kwargs): - metric = lambda rank: self.array((rank,rank)) - TensorOperator.__init__(self,metric,TensorCodomain(0),**kwargs) + metric = lambda rank: self.array((rank, rank)) + TensorOperator.__init__(self, metric, TensorCodomain(0), **kwargs) @int2tuple - def __getitem__(self,i): + def __getitem__(self, i): return int(i[0] == dual(i[1])) @@ -135,10 +137,10 @@ class Transpose(TensorOperator): """ - def __init__(self,permutation=(1,0),**kwargs): + def __init__(self, permutation=(1, 0), **kwargs): - transpose = lambda rank: self.array((rank,rank)) - TensorOperator.__init__(self,transpose,TensorCodomain(0),**kwargs) + transpose = lambda rank: self.array((rank, rank)) + TensorOperator.__init__(self, transpose, TensorCodomain(0), **kwargs) self.__permutation = permutation @property @@ -146,9 +148,10 @@ def permutation(self): return self.__permutation @int2tuple - def __getitem__(self,i): + def __getitem__(self, i): return int(i[0] == apply(self.permutation)(i[1])) + class Trace(TensorOperator): """ Class for contracting arbitrary indices down to a scalar in those indices: @@ -198,23 +201,25 @@ class Trace(TensorOperator): """ - def __init__(self,indices,**kwargs): - if type(indices) == int: indices = (indices,) + def __init__(self, indices, **kwargs): + if type(indices) == int: + indices = (indices,) - trace = lambda rank: self.array((rank-len(indices),rank)) - TensorOperator.__init__(self,trace,TensorCodomain(-len(indices)),**kwargs) + trace = lambda rank: self.array((rank - len(indices), rank)) + TensorOperator.__init__(self, trace, TensorCodomain(-len(indices)), **kwargs) - self.__indices = indices + self.__indices = indices @property def indices(self): return self.__indices @int2tuple - def __getitem__(self,i): + def __getitem__(self, i): return int(i[0] == remove(self.indices)(i[1]) and sum_(self.indices)(i[1]) == 0) + class TensorProduct(TensorOperator): """ Action of multiplication by single spin-tensor basis element: @@ -233,13 +238,14 @@ class TensorProduct(TensorOperator): """ - def __init__(self,element,action='left',**kwargs): - if type(element) == int: element = (element,) + def __init__(self, element, action='left', **kwargs): + if type(element) == int: + element = (element,) - product = lambda rank: self.array((rank+len(element),rank)) - TensorOperator.__init__(self,product,TensorCodomain(len(element)),**kwargs) + product = lambda rank: self.array((rank + len(element), rank)) + TensorOperator.__init__(self, product, TensorCodomain(len(element)), **kwargs) self.__element = element - self.__action = action + self.__action = action @property def element(self): @@ -250,27 +256,27 @@ def action(self): return self.__action @int2tuple - def __getitem__(self,i): + def __getitem__(self, i): if self.action == 'left': return int(i[0] == self.element + i[1]) if self.action == 'right': return int(i[0] == i[1] + self.element) -def xi(mu,ell): +def xi(mu, ell): """ - Normalised derivative scale factors. xi(-1,ell)**2 + xi(+1,ell)**2 = 1. + Normalised derivative scale factors. xi(-1,ell)**2 + xi(+1,ell)**2 = 1. - Parameters - ---------- - mu : int - regularity; -1,+1,0. xi(0,ell) = 0 by definition. - ell : int - spherical-harmonic degree. + Parameters + ---------- + mu : int + regularity; -1,+1,0. xi(0,ell) = 0 by definition. + ell : int + spherical-harmonic degree. - """ + """ - return np.abs(mu)*np.sqrt((1 + mu/(2*ell+1))/2) + return np.abs(mu) * np.sqrt((1 + mu / (2 * ell + 1)) / 2) class Intertwiner(TensorOperator): @@ -297,25 +303,25 @@ class Intertwiner(TensorOperator): """ - def __init__(self,L,**kwargs): + def __init__(self, L, **kwargs): - intertwiner = lambda rank: self.array((rank,rank)) - TensorOperator.__init__(self,intertwiner,TensorCodomain(0),**kwargs) + intertwiner = lambda rank: self.array((rank, rank)) + TensorOperator.__init__(self, intertwiner, TensorCodomain(0), **kwargs) self.__ell = L @property def L(self): return self.__ell - def k(self,mu,s): - return -mu*np.sqrt((self.L-s*mu)*(self.L+s*mu+1)/2) + def k(self, mu, s): + return -mu * np.sqrt((self.L - s * mu) * (self.L + s * mu + 1) / 2) @int2tuple - def forbidden_spin(self,spin): + def forbidden_spin(self, spin): return self.L < abs(sum(spin)) @int2tuple - def forbidden_regularity(self,regularity): + def forbidden_regularity(self, regularity): # Fast return for clearly allowed cases if self.L >= len(regularity): return False @@ -323,12 +329,12 @@ def forbidden_regularity(self,regularity): walk = (self.L,) for r in regularity[::-1]: walk += (walk[-1] + r,) - if walk[-1] < 0 or walk[-2:] == (0,0): + if walk[-1] < 0 or walk[-2:] == (0, 0): return True return False @int2tuple - def __getitem__(self,i): + def __getitem__(self, i): spin, regularity = i[0], i[1] @@ -338,24 +344,29 @@ def __getitem__(self,i): if self.forbidden_spin(spin) or self.forbidden_regularity(regularity): return 0 - sigma, a = spin[0], regularity[0] - tau, b = spin[1:], regularity[1:] + sigma, a = spin[0], regularity[0] + tau, b = spin[1:], regularity[1:] R = 0 - for i,t in enumerate(tau): - if t+sigma == 0: R -= self[replace(i,0)(tau),b] - if t == 0: R += self[replace(i,sigma)(tau),b] + for i, t in enumerate(tau): + if t + sigma == 0: + R -= self[replace(i, 0)(tau), b] + if t == 0: + R += self[replace(i, sigma)(tau), b] - Q = self[tau,b] - R -= self.k(sigma,sum(tau))*Q - J = self.L + sum(b) + Q = self[tau, b] + R -= self.k(sigma, sum(tau)) * Q + J = self.L + sum(b) if sigma != 0: Q = 0 - if a == -1: return (Q * J - R)/np.sqrt(J*(2*J+1)) - if a == 0: return sigma*R /np.sqrt(J*(J+1)) - if a == +1: return (Q*(J+1) + R)/np.sqrt((J+1)*(2*J+1)) + if a == -1: + return (Q * J - R) / np.sqrt(J * (2 * J + 1)) + if a == 0: + return sigma * R / np.sqrt(J * (J + 1)) + if a == +1: + return (Q * (J + 1) + R) / np.sqrt((J + 1) * (2 * J + 1)) class TensorCodomain(Codomain): @@ -371,9 +382,9 @@ class TensorCodomain(Codomain): """ - def __init__(self,rank_change): - Codomain.__init__(self,rank_change,Output=TensorCodomain) + def __init__(self, rank_change): + Codomain.__init__(self, rank_change, Output=TensorCodomain) def __str__(self): s = f'(rank->rank+{self[0]})' - return s.replace('+0','').replace('+-','-') + return s.replace('+0', '').replace('+-', '-') diff --git a/dedalus/libraries/dedalus_sphere/timesteppers.py b/dedalus/libraries/dedalus_sphere/timesteppers.py index 7079d85e..75aa5280 100644 --- a/dedalus/libraries/dedalus_sphere/timesteppers.py +++ b/dedalus/libraries/dedalus_sphere/timesteppers.py @@ -1,14 +1,15 @@ - from collections import deque import numpy as np from scipy.sparse import linalg from dedalus.tools.config import config -#STORE_LU = config['linear algebra'].getboolean('store_LU') + +# STORE_LU = config['linear algebra'].getboolean('store_LU') STORE_LU = True PERMC_SPEC = config['linear algebra']['permc_spec'] USE_UMFPACK = config['linear algebra'].getboolean('use_umfpack') + class MultistepIMEX: """ Base class for implicit-explicit multistep methods. @@ -50,7 +51,7 @@ def __init__(self, CoeffSystem, *args): # Create deque for storing recent timesteps N = max(self.amax, self.bmax, self.cmax) - self.dt = deque([0.]*N) + self.dt = deque([0.0] * N) # Create coefficient systems for multistep history self.MX = MX = deque() @@ -96,25 +97,25 @@ def step(self, dt, state_vector, basis, L, M, P, NL, LU): b0 = b[0] if STORE_LU: - update_LHS = ((a0, b0) != self._LHS_params) + update_LHS = (a0, b0) != self._LHS_params self._LHS_params = (a0, b0) for dl, l in enumerate(basis.local_l): - P[dl] = a0*M[dl] + b0*L[dl] + P[dl] = a0 * M[dl] + b0 * L[dl] for dm, m in enumerate(basis.local_m): if m <= l: MX0.data[dl][dm] = M[dl].dot(state_vector.data[dl][dm]) LX0.data[dl][dm] = L[dl].dot(state_vector.data[dl][dm]) - np.copyto(F0.data[dl][dm],NL.data[dl][dm]) + np.copyto(F0.data[dl][dm], NL.data[dl][dm]) # Build RHS - RHS.data[dl][dm] *= 0. + RHS.data[dl][dm] *= 0.0 for j in range(1, len(c)): - RHS.data[dl][dm] += c[j] * F[j-1].data[dl][dm] + RHS.data[dl][dm] += c[j] * F[j - 1].data[dl][dm] for j in range(1, len(a)): - RHS.data[dl][dm] -= a[j] * MX[j-1].data[dl][dm] + RHS.data[dl][dm] -= a[j] * MX[j - 1].data[dl][dm] for j in range(1, len(b)): - RHS.data[dl][dm] -= b[j] * LX[j-1].data[dl][dm] + RHS.data[dl][dm] -= b[j] * LX[j - 1].data[dl][dm] # Solve if STORE_LU: @@ -123,7 +124,8 @@ def step(self, dt, state_vector, basis, L, M, P, NL, LU): pLHS = LU[dl] state_vector.data[dl][dm] = pLHS.solve(RHS.data[dl][dm]) else: - state_vector.data[dl][dm] = linalg.spsolve(P[dl],RHS.data[dl][dm], use_umfpack=USE_UMFPACK, permc_spec=PERMC_SPEC) + state_vector.data[dl][dm] = linalg.spsolve(P[dl], RHS.data[dl][dm], use_umfpack=USE_UMFPACK, permc_spec=PERMC_SPEC) + class SBDF1(MultistepIMEX): """ @@ -141,9 +143,9 @@ class SBDF1(MultistepIMEX): @classmethod def compute_coefficients(self, timesteps, iteration): - a = np.zeros(self.amax+1) - b = np.zeros(self.bmax+1) - c = np.zeros(self.cmax+1) + a = np.zeros(self.amax + 1) + b = np.zeros(self.bmax + 1) + c = np.zeros(self.cmax + 1) k0, *rest = timesteps @@ -174,14 +176,14 @@ def compute_coefficients(self, timesteps, iteration): if iteration < 1: return SBDF1.compute_coefficients(timesteps, iteration) - a = np.zeros(self.amax+1) - b = np.zeros(self.bmax+1) - c = np.zeros(self.cmax+1) + a = np.zeros(self.amax + 1) + b = np.zeros(self.bmax + 1) + c = np.zeros(self.cmax + 1) k1, k0, *rest = timesteps w1 = k1 / k0 - a[0] = (1 + 2*w1) / (1 + w1) / k1 + a[0] = (1 + 2 * w1) / (1 + w1) / k1 a[1] = -(1 + w1) / k1 a[2] = w1**2 / (1 + w1) / k1 b[0] = 1 @@ -210,22 +212,22 @@ def compute_coefficients(self, timesteps, iteration): if iteration < 2: return SBDF2.compute_coefficients(timesteps, iteration) - a = np.zeros(self.amax+1) - b = np.zeros(self.bmax+1) - c = np.zeros(self.cmax+1) + a = np.zeros(self.amax + 1) + b = np.zeros(self.bmax + 1) + c = np.zeros(self.cmax + 1) k2, k1, k0, *rest = timesteps w2 = k2 / k1 w1 = k1 / k0 - a[0] = (1 + w2/(1 + w2) + w1*w2/(1 + w1*(1 + w2))) / k2 - a[1] = (-1 - w2 - w1*w2*(1 + w2)/(1 + w1)) / k2 - a[2] = w2**2 * (w1 + 1/(1 + w2)) / k2 - a[3] = -w1**3 * w2**2 * (1 + w2) / (1 + w1) / (1 + w1 + w1*w2) / k2 + a[0] = (1 + w2 / (1 + w2) + w1 * w2 / (1 + w1 * (1 + w2))) / k2 + a[1] = (-1 - w2 - w1 * w2 * (1 + w2) / (1 + w1)) / k2 + a[2] = w2**2 * (w1 + 1 / (1 + w2)) / k2 + a[3] = -(w1**3) * w2**2 * (1 + w2) / (1 + w1) / (1 + w1 + w1 * w2) / k2 b[0] = 1 - c[1] = (1 + w2)*(1 + w1*(1 + w2)) / (1 + w1) - c[2] = -w2*(1 + w1*(1 + w2)) - c[3] = w1*w1*w2*(1 + w2) / (1 + w1) + c[1] = (1 + w2) * (1 + w1 * (1 + w2)) / (1 + w1) + c[2] = -w2 * (1 + w1 * (1 + w2)) + c[3] = w1 * w1 * w2 * (1 + w2) / (1 + w1) return a, b, c @@ -249,29 +251,29 @@ def compute_coefficients(self, timesteps, iteration): if iteration < 3: return SBDF3.compute_coefficients(timesteps, iteration) - a = np.zeros(self.amax+1) - b = np.zeros(self.bmax+1) - c = np.zeros(self.cmax+1) + a = np.zeros(self.amax + 1) + b = np.zeros(self.bmax + 1) + c = np.zeros(self.cmax + 1) k3, k2, k1, k0, *rest = timesteps w3 = k3 / k2 w2 = k2 / k1 w1 = k1 / k0 - A1 = 1 + w1*(1 + w2) - A2 = 1 + w2*(1 + w3) - A3 = 1 + w1*A2 + A1 = 1 + w1 * (1 + w2) + A2 = 1 + w2 * (1 + w3) + A3 = 1 + w1 * A2 - a[0] = (1 + w3/(1 + w3) + w2*w3/A2 + w1*w2*w3/A3) / k3 - a[1] = (-1 - w3*(1 + w2*(1 + w3)/(1 + w2)*(1 + w1*A2/A1))) / k3 - a[2] = w3 * (w3/(1 + w3) + w2*w3*(A3 + w1)/(1 + w1)) / k3 - a[3] = -w2**3 * w3**2 * (1 + w3) / (1 + w2) * A3 / A2 / k3 + a[0] = (1 + w3 / (1 + w3) + w2 * w3 / A2 + w1 * w2 * w3 / A3) / k3 + a[1] = (-1 - w3 * (1 + w2 * (1 + w3) / (1 + w2) * (1 + w1 * A2 / A1))) / k3 + a[2] = w3 * (w3 / (1 + w3) + w2 * w3 * (A3 + w1) / (1 + w1)) / k3 + a[3] = -(w2**3) * w3**2 * (1 + w3) / (1 + w2) * A3 / A2 / k3 a[4] = (1 + w3) / (1 + w1) * A2 / A1 * w1**4 * w2**3 * w3**2 / A3 / k3 b[0] = 1 - c[1] = w2 * (1 + w3) / (1 + w2) * ((1 + w3)*(A3 + w1) + (1 + w1)/w2) / A1 + c[1] = w2 * (1 + w3) / (1 + w2) * ((1 + w3) * (A3 + w1) + (1 + w1) / w2) / A1 c[2] = -A2 * A3 * w3 / (1 + w1) c[3] = w2**2 * w3 * (1 + w3) / (1 + w2) * A3 - c[4] = -w1**3 * w2**2 * w3 * (1 + w3) / (1 + w1) * A2 / A1 + c[4] = -(w1**3) * w2**2 * w3 * (1 + w3) / (1 + w1) * A2 / A1 return a, b, c @@ -292,9 +294,9 @@ class CNAB1(MultistepIMEX): @classmethod def compute_coefficients(self, timesteps, iteration): - a = np.zeros(self.amax+1) - b = np.zeros(self.bmax+1) - c = np.zeros(self.cmax+1) + a = np.zeros(self.amax + 1) + b = np.zeros(self.bmax + 1) + c = np.zeros(self.cmax + 1) k0, *rest = timesteps @@ -324,12 +326,12 @@ class CNAB2(MultistepIMEX): def compute_coefficients(self, timesteps, iteration): if iteration < 1: -# return CNAB1.compute_coefficients(timesteps, iteration) + # return CNAB1.compute_coefficients(timesteps, iteration) return SBDF1.compute_coefficients(timesteps, iteration) - a = np.zeros(self.amax+1) - b = np.zeros(self.bmax+1) - c = np.zeros(self.cmax+1) + a = np.zeros(self.amax + 1) + b = np.zeros(self.bmax + 1) + c = np.zeros(self.cmax + 1) k1, k0, *rest = timesteps w1 = k1 / k0 @@ -338,11 +340,12 @@ def compute_coefficients(self, timesteps, iteration): a[1] = -1 / k1 b[0] = 1 / 2 b[1] = 1 / 2 - c[1] = 1 + w1/2 + c[1] = 1 + w1 / 2 c[2] = -w1 / 2 return a, b, c + class RungeKuttaIMEX: """ Base class for implicit-explicit multistep methods. @@ -386,7 +389,6 @@ def __init__(self, CoeffSystem, *args): self.RHS = CoeffSystem(*args) - self.MX0 = CoeffSystem(*args) self.LX = LX = [CoeffSystem(*args) for i in range(self.stages)] self.NL = NL = [CoeffSystem(*args) for i in range(self.stages)] @@ -408,7 +410,7 @@ def step(self, dt, state_vector, B, L, M, P, nonlinear, LU): k = dt if STORE_LU: - update_LHS = (k != self._LHS_params) + update_LHS = k != self._LHS_params self._LHS_params = k ell_start = B.ell_min @@ -417,45 +419,45 @@ def step(self, dt, state_vector, B, L, M, P, nonlinear, LU): m_end = B.m_max m_size = m_end - m_start + 1 - for ell in range(ell_start,ell_end+1): - ell_local = ell-ell_start - for m in range(m_start,m_end+1): - m_local = m-m_start - index = ell_local*m_size+m_local + for ell in range(ell_start, ell_end + 1): + ell_local = ell - ell_start + for m in range(m_start, m_end + 1): + m_local = m - m_start + index = ell_local * m_size + m_local MX0.data[index] = M[ell_local].dot(state_vector.data[index]) if STORE_LU and update_LHS: - LU[ell_local] = [None] * (self.stages+1) + LU[ell_local] = [None] * (self.stages + 1) # Compute stages # (M + k Hii L).X(n,i) = M.X(n,0) + k Aij F(n,j) - k Hij L.X(n,j) - for i in range(1, self.stages+1): + for i in range(1, self.stages + 1): # Compute F(n,i-1) - nonlinear(state_vector,NL[i-1],0) + nonlinear(state_vector, NL[i - 1], 0) # Compute L.X(n,i-1) - for ell in range(ell_start,ell_end+1): - ell_local = ell-ell_start - P[ell_local] = M[ell_local] + (k*H[i,i])*L[ell_local] + for ell in range(ell_start, ell_end + 1): + ell_local = ell - ell_start + P[ell_local] = M[ell_local] + (k * H[i, i]) * L[ell_local] if STORE_LU and update_LHS: LU[ell_local][i] = linalg.splu(P[ell_local].tocsc(), permc_spec=PERMC_SPEC) - for m in range(m_start,m_end+1): - m_local = m-m_start - index = ell_local*m_size+m_local + for m in range(m_start, m_end + 1): + m_local = m - m_start + index = ell_local * m_size + m_local - LX[i-1].data[index] = L[ell_local].dot(state_vector.data[index]) + LX[i - 1].data[index] = L[ell_local].dot(state_vector.data[index]) - # Construct RHS(n,i) - np.copyto(RHS.data[index],MX0.data[index]) - for j in range(0,i): - RHS.data[index] += k * A[i,j] * NL[j].data[index] - RHS.data[index] -= k * H[i,j] * LX[j].data[index] + # Construct RHS(n,i) + np.copyto(RHS.data[index], MX0.data[index]) + for j in range(0, i): + RHS.data[index] += k * A[i, j] * NL[j].data[index] + RHS.data[index] -= k * H[i, j] * LX[j].data[index] - # Solve + # Solve if STORE_LU: state_vector.data[index] = LU[ell_local][i].solve(RHS.data[index]) else: - state_vector.data[index] = linalg.spsolve(P[ell_local],RHS.data[index], use_umfpack=USE_UMFPACK, permc_spec=PERMC_SPEC) + state_vector.data[index] = linalg.spsolve(P[ell_local], RHS.data[index], use_umfpack=USE_UMFPACK, permc_spec=PERMC_SPEC) class RK111(RungeKuttaIMEX): @@ -465,11 +467,9 @@ class RK111(RungeKuttaIMEX): c = np.array([0, 1]) - A = np.array([[0, 0], - [1, 0]]) + A = np.array([[0, 0], [1, 0]]) - H = np.array([[0, 0], - [0, 1]]) + H = np.array([[0, 0], [0, 1]]) class RK222(RungeKuttaIMEX): @@ -482,32 +482,22 @@ class RK222(RungeKuttaIMEX): c = np.array([0, γ, 1]) - A = np.array([[0, 0 , 0], - [γ, 0 , 0], - [δ, 1-δ, 0]]) + A = np.array([[0, 0, 0], [γ, 0, 0], [δ, 1 - δ, 0]]) + + H = np.array([[0, 0, 0], [0, γ, 0], [0, 1 - γ, γ]]) - H = np.array([[0, 0 , 0], - [0, γ , 0], - [0, 1-γ, γ]]) class RK443(RungeKuttaIMEX): """3rd-order 4-stage DIRK+ERK scheme [Ascher 1997 sec 2.8]""" stages = 4 - c = np.array([0, 1/2, 2/3, 1/2, 1]) + c = np.array([0, 1 / 2, 2 / 3, 1 / 2, 1]) + + A = np.array([[0, 0, 0, 0, 0], [1 / 2, 0, 0, 0, 0], [11 / 18, 1 / 18, 0, 0, 0], [5 / 6, -5 / 6, 1 / 2, 0, 0], [1 / 4, 7 / 4, 3 / 4, -7 / 4, 0]]) - A = np.array([[ 0 , 0 , 0 , 0 , 0], - [ 1/2 , 0 , 0 , 0 , 0], - [11/18, 1/18, 0 , 0 , 0], - [ 5/6 , -5/6 , 1/2, 0 , 0], - [ 1/4 , 7/4 , 3/4, -7/4, 0]]) + H = np.array([[0, 0, 0, 0, 0], [0, 1 / 2, 0, 0, 0], [0, 1 / 6, 1 / 2, 0, 0], [0, -1 / 2, 1 / 2, 1 / 2, 0], [0, 3 / 2, -3 / 2, 1 / 2, 1 / 2]]) - H = np.array([[0, 0 , 0 , 0 , 0 ], - [0, 1/2, 0 , 0 , 0 ], - [0, 1/6, 1/2, 0 , 0 ], - [0, -1/2, 1/2, 1/2, 0 ], - [0, 3/2, -3/2, 1/2, 1/2]]) class RKGFY(RungeKuttaIMEX): """2nd-order 2-stage scheme from Hollerbach and Marti""" @@ -516,10 +506,6 @@ class RKGFY(RungeKuttaIMEX): c = np.array([0, 1, 1]) - A = np.array([[ 0, 0 , 0], - [ 1, 0 , 0], - [0.5, 0.5, 0]]) + A = np.array([[0, 0, 0], [1, 0, 0], [0.5, 0.5, 0]]) - H = np.array([[0 , 0 , 0], - [0.5 , 0.5, 0], - [0.5 , 0 , 0.5]]) + H = np.array([[0, 0, 0], [0.5, 0.5, 0], [0.5, 0, 0.5]]) diff --git a/dedalus/libraries/dedalus_sphere/tuple_tools.py b/dedalus/libraries/dedalus_sphere/tuple_tools.py index 44460f23..9a109636 100644 --- a/dedalus/libraries/dedalus_sphere/tuple_tools.py +++ b/dedalus/libraries/dedalus_sphere/tuple_tools.py @@ -1,27 +1,32 @@ - # tuple helper functions -dual = lambda t: tuple(-e for e in t) -apply = lambda p: lambda t: tuple(t[int(i)] for i in p) -sum_ = lambda k: lambda t: sum(t[int(i)] for i in k if 0 <= i < len(t)) -remove = lambda k: lambda t: tuple(s for i,s in enumerate(t) if not i in k) -replace = lambda j,n: lambda t: tuple(s if i!=j else n for i,s in enumerate(t)) +dual = lambda t: tuple(-e for e in t) +apply = lambda p: lambda t: tuple(t[int(i)] for i in p) +sum_ = lambda k: lambda t: sum(t[int(i)] for i in k if 0 <= i < len(t)) +remove = lambda k: lambda t: tuple(s for i, s in enumerate(t) if not i in k) +replace = lambda j, n: lambda t: tuple(s if i != j else n for i, s in enumerate(t)) # converts integer args to tuple singleton i --> (i,). def int2tuple(func): if func.__name__ == '__getitem__': + def wrapper(*args): self, args = args[0], args[1] - if type(args) == slice: return func(self,(args,)) - if type(args) == int: return func(self,(args,)) - args = (self,tuple((s,) if type(s)==int else s for s in args)) + if type(args) == slice: + return func(self, (args,)) + if type(args) == int: + return func(self, (args,)) + args = (self, tuple((s,) if type(s) == int else s for s in args)) return func(*args) + return wrapper - return lambda *args: func(*tuple((s,) if type(s)==int else s for s in args)) - + return lambda *args: func(*tuple((s,) if type(s) == int else s for s in args)) + + # forward and inverse mapping between spin/regularity elements and integers. -def tuple2index(tup,indexing): - return int('0'+''.join(str(indexing.index(s)) for s in tup),len(indexing)) +def tuple2index(tup, indexing): + return int('0' + ''.join(str(indexing.index(s)) for s in tup), len(indexing)) + -def index2tuple(index,rank,indexing): - s = np.base_repr(index,len(indexing),rank) - return apply(s[(rank==0)-rank:])(indexing) +def index2tuple(index, rank, indexing): + s = np.base_repr(index, len(indexing), rank) + return apply(s[(rank == 0) - rank :])(indexing) diff --git a/dedalus/libraries/dedalus_sphere/zernike.py b/dedalus/libraries/dedalus_sphere/zernike.py index 1050bce2..52507e89 100644 --- a/dedalus/libraries/dedalus_sphere/zernike.py +++ b/dedalus/libraries/dedalus_sphere/zernike.py @@ -1,45 +1,49 @@ import numpy as np -from . import jacobi as Jacobi -from .operators import Operator, infinite_csr +from . import jacobi as Jacobi +from .operators import Operator, infinite_csr # The defalut configuration for the base Jacobi parameter. alpha = 0 -def mass(dimension,k=alpha): - return Jacobi.mass(k,dimension/2 - 1)/2**( k + dimension/2 + 1 ) -def quadrature(dimension,n,k=alpha): +def mass(dimension, k=alpha): + return Jacobi.mass(k, dimension / 2 - 1) / 2 ** (k + dimension / 2 + 1) + + +def quadrature(dimension, n, k=alpha): """ Weights associated with dV = (1-r*r)**k * r**(dimension-1) dr, where 0 <= r <= 1. """ - z, w = Jacobi.quadrature(n,k,dimension/2 - 1) + z, w = Jacobi.quadrature(n, k, dimension / 2 - 1) - w /= 2**( k + dimension/2 + 1 ) + w /= 2 ** (k + dimension / 2 + 1) return z, w + def min_degree(l): - return max(l//2,0) + return max(l // 2, 0) + -def polynomials(dimension,n,k,l,z): +def polynomials(dimension, n, k, l, z): """ - Unit normalised: + Unit normalised: - integral(Q**2 dV) = 1 + integral(Q**2 dV) = 1 """ - b = l + dimension/2 - 1 + b = l + dimension / 2 - 1 - init = Jacobi.measure(0,l,z,log=True,probability=False) - init -= Jacobi.mass(k,b,log=True) - np.log(2)*(k + dimension/2 + 1) - init = np.exp(0.5*init) + init = Jacobi.measure(0, l, z, log=True, probability=False) + init -= Jacobi.mass(k, b, log=True) - np.log(2) * (k + dimension / 2 + 1) + init = np.exp(0.5 * init) - return Jacobi.polynomials(n,k,b,z,init) + return Jacobi.polynomials(n, k, b, z, init) def operator(dimension, name, radius=1): @@ -52,26 +56,30 @@ def operator(dimension, name, radius=1): """ if name == 'Id': - def I(n,k,l): - return Jacobi.operator('Id')(n,k,l+dimension/2 - 1) - return Operator(I,ZernikeCodomain(0,0,0)) + + def I(n, k, l): + return Jacobi.operator('Id')(n, k, l + dimension / 2 - 1) + + return Operator(I, ZernikeCodomain(0, 0, 0)) if name == 'Z': - def Z(n,k,l): - return Jacobi.operator('Z')(n,k,l+dimension/2 - 1) - return Operator(Z,ZernikeCodomain(1,0,0)) + + def Z(n, k, l): + return Jacobi.operator('Z')(n, k, l + dimension / 2 - 1) + + return Operator(Z, ZernikeCodomain(1, 0, 0)) return ZernikeOperator(dimension, name, radius=radius) -class ZernikeOperator(): - def __init__(self,dimension,name,radius=1): +class ZernikeOperator: + def __init__(self, dimension, name, radius=1): - self.__function = getattr(self,f'_ZernikeOperator__{name}') - self.__dimension = dimension - self.__radius = radius + self.__function = getattr(self, f'_ZernikeOperator__{name}') + self.__dimension = dimension + self.__radius = radius - def __call__(self,p): + def __call__(self, p): return Operator(*self.__function(p)) @property @@ -82,40 +90,35 @@ def dimension(self): def radius(self): return self.__radius - def b(self,l): - return l + self.dimension/2 - 1 - - def __D(self,dl): + def b(self, l): + return l + self.dimension / 2 - 1 - def D(n,k,l): + def __D(self, dl): + def D(n, k, l): D = Jacobi.operator('D' if dl > 0 else 'C')(+1) - return (2/self.radius)*D(n,k,self.b(l)) - - return D, ZernikeCodomain(-(1+dl)//2,1,dl) + return (2 / self.radius) * D(n, k, self.b(l)) - def __E(self,dk): + return D, ZernikeCodomain(-(1 + dl) // 2, 1, dl) - def E(n,k,l): + def __E(self, dk): + def E(n, k, l): E = Jacobi.operator('A')(dk) - return np.sqrt(0.5)*E(n,k,self.b(l)) + return np.sqrt(0.5) * E(n, k, self.b(l)) - return E, ZernikeCodomain((1-dk)//2,dk,0) + return E, ZernikeCodomain((1 - dk) // 2, dk, 0) - def __R(self,dl): - - def R(n,k,l): + def __R(self, dl): + def R(n, k, l): R = Jacobi.operator('B')(dl) - return (np.sqrt(0.5)*self.radius)*R(n,k,self.b(l)) + return (np.sqrt(0.5) * self.radius) * R(n, k, self.b(l)) - return R, ZernikeCodomain((1-dl)//2,0,dl) + return R, ZernikeCodomain((1 - dl) // 2, 0, dl) class ZernikeCodomain(Jacobi.JacobiCodomain): - - def __init__(self,dn=0,dk=0,dl=0,pi=0): - Jacobi.JacobiCodomain.__init__(self,dn,dk,dl,0,Output=ZernikeCodomain) + def __init__(self, dn=0, dk=0, dl=0, pi=0): + Jacobi.JacobiCodomain.__init__(self, dn, dk, dl, 0, Output=ZernikeCodomain) def __str__(self): s = f'(n->n+{self[0]},k->k+{self[1]},l->l+{self[2]})' - return s.replace('+0','').replace('+-','-') - + return s.replace('+0', '').replace('+-', '-') diff --git a/dedalus/tests/test_clenshaw.py b/dedalus/tests/test_clenshaw.py index 56f8a2e6..ed759974 100644 --- a/dedalus/tests/test_clenshaw.py +++ b/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 @@ -40,7 +39,7 @@ def test_clenshaw(N, regtotal_in, k1, k2, ell): coeffs_filter = ncc['c'][0,0,:n_size] J = basis_in.operator_matrix('Z', ell, regtotal_in) A, B = clenshaw.jacobi_recursion(n_size, a_ncc, b_ncc, J) - f0 = dedalus_sphere.zernike.polynomials(3, 1, a_ncc, regtotal_ncc, 1)[0] * sparse.identity(n_size) + f0 = dedalus_sphere.zernike.polynomials(3, 1, a_ncc, regtotal_ncc, 1)[0].astype(np.float64) * sparse.identity(n_size) matrix = clenshaw.matrix_clenshaw(coeffs_filter, A, B, f0, cutoff=1e-6) assert np.allclose(J.todense(), matrix.todense()) @@ -74,7 +73,7 @@ def test_clenshaw_vector(N, regtotal_in, k1, k2, ell): J = basis_in.operator_matrix('Z', ell, regtotal_in) I = basis_in.operator_matrix('Id', ell, regtotal_in) A, B = clenshaw.jacobi_recursion(n_size, a_ncc, b_ncc, J) - f0 = dedalus_sphere.zernike.polynomials(3, 1, a_ncc, regtotal_ncc, 1)[0] * sparse.identity(n_size) + f0 = dedalus_sphere.zernike.polynomials(3, 1, a_ncc, regtotal_ncc, 1)[0].astype(np.float64) * sparse.identity(n_size) matrix = clenshaw.matrix_clenshaw(coeffs_filter, A, B, f0, cutoff=1e-6) assert np.allclose(J.todense(), matrix.todense()) diff --git a/dedalus/tests/test_transforms.py b/dedalus/tests/test_transforms.py index 01930bb7..8b42abb6 100644 --- a/dedalus/tests/test_transforms.py +++ b/dedalus/tests/test_transforms.py @@ -367,6 +367,54 @@ def test_sphere_roundtrip_noise(Nphi, Ntheta, radius, basis, dealias, dtype, lay assert np.allclose(f_layout, f[layout]) +@pytest.mark.parametrize('Ntheta', [16]) +#@pytest.mark.parametrize('dealias', [0.5, 1, 1.5]) +#@pytest.mark.parametrize('dtype', [np.float64, np.complex128]) +@pytest.mark.parametrize('dealias', [1, 1.5]) +@pytest.mark.parametrize('dtype', [np.float64]) +@pytest.mark.parametrize('library', ['shtns', 'ducc']) +def test_colatitude_libraries_backward(Ntheta, dealias, dtype, library): + """Tests that ALT library results match matrix transforms.""" + shape = (2 * Ntheta, Ntheta) + c = coords.S2Coordinates('phi', 'theta') + d = distributor.Distributor(c, dtype=dtype) + # Matrix + b_mat = basis.SphereBasis(c, shape=shape, dealias=dealias, dtype=dtype, colatitude_library='matrix') + u_mat = d.Field(bases=b_mat) + u_mat.preset_scales(dealias) + u_mat.fill_random(layout='c') + # Library + b_lib = basis.SphereBasis(c, shape=shape, dealias=dealias, dtype=dtype, colatitude_library=library) + u_lib = d.Field(bases=b_lib) + u_lib.preset_scales(dealias) + u_lib['c'] = u_mat['c'] + assert np.allclose(u_mat['g'], u_lib['g']) + + +@pytest.mark.parametrize('Ntheta', [16]) +#@pytest.mark.parametrize('dealias', [0.5, 1, 1.5]) +#@pytest.mark.parametrize('dtype', [np.float64, np.complex128]) +@pytest.mark.parametrize('dealias', [1, 1.5]) +@pytest.mark.parametrize('dtype', [np.float64]) +@pytest.mark.parametrize('library', ['shtns', 'ducc']) +def test_colatitude_libraries_forward(Ntheta, dealias, dtype, library): + """Tests that ALT library results match matrix transforms.""" + shape = (2 * Ntheta, Ntheta) + c = coords.S2Coordinates('phi', 'theta') + d = distributor.Distributor(c, dtype=dtype) + # Matrix + b_mat = basis.SphereBasis(c, shape=shape, dealias=dealias, dtype=dtype, colatitude_library='matrix') + u_mat = d.Field(bases=b_mat) + u_mat.preset_scales(dealias) + u_mat.fill_random(layout='g') + # Library + b_lib = basis.SphereBasis(c, shape=shape, dealias=dealias, dtype=dtype, colatitude_library=library) + u_lib = d.Field(bases=b_lib) + u_lib.preset_scales(dealias) + u_lib['g'] = u_mat['g'] + print(u_lib['c']/u_mat['c']) + assert np.allclose(u_mat['c'], u_lib['c']) + ## Polar bases Nphi_range = [16] diff --git a/dedalus/tools/clenshaw.py b/dedalus/tools/clenshaw.py index d1a40d70..7dc9fd1a 100644 --- a/dedalus/tools/clenshaw.py +++ b/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): diff --git a/dedalus/tools/jacobi.py b/dedalus/tools/jacobi.py deleted file mode 100644 index 6c09d87f..00000000 --- a/dedalus/tools/jacobi.py +++ /dev/null @@ -1,264 +0,0 @@ -# import numpy as np -# import scipy as sp -# import scipy.sparse as sparse -# import scipy.linalg as linear -import scipy.special as fun - - - -# dtype = np.float128 - -# def sparse_symm_to_banded(matrix): -# """Convert sparse symmetric to upper-banded form.""" -# diag = matrix.todia() -# B, I = max(abs(diag.offsets)), diag.data.shape[1] -# banded = np.zeros((B+1, I), dtype=diag.dtype) -# for i, b in enumerate(diag.offsets): -# if b >= 0: -# banded[B-b] = diag.data[i] -# return banded - -# def grid_guess(Jacobi_matrix,symmetric=True): -# if symmetric: -# J_banded = sparse_symm_to_banded(Jacobi_matrix) -# return (np.sort(linear.eigvals_banded(J_banded).real)).astype(dtype) -# else: -# J_dense = Jacobi_matrix.todense() -# return (np.sort(linear.eigvals(J_dense).real)).astype(dtype) - -# def remainders(Jacobi_matrix,grid): - -# J, z, N = sparse.dia_matrix(Jacobi_matrix).data, grid, len(grid) - -# P = np.zeros((N,N),dtype=dtype) -# D = np.zeros((N,N),dtype=dtype) - -# P[0] = np.ones(N,dtype=dtype) -# D[1] = P[0]/J[-1][1] - -# if np.shape(J)[0] == 3: -# P[1] = (z-J[1][0])*P[0]/J[-1][1] -# for n in range(2,N): -# P[n] = ((z-J[1][n-1])*P[n-1] - J[0][n-2]*P[n-2])/J[-1][n] -# D[n] = ((z-J[1][n-1])*D[n-1] - J[0][n-2]*D[n-2])/J[-1][n] + P[n-1]/J[-1][n] -# return ((z-J[1][N-1])*P[N-1] - J[0][N-2]*P[N-2]), ((z-J[1][N-1])*D[N-1] - J[0][N-2]*D[N-2]+P[N-1]) - -# P[1] = z*P[0]/J[-1][1] -# for n in range(2,N): -# P[n] = (z*P[n-1] - J[0][n-2]*P[n-2])/J[-1][n] -# D[n] = (z*D[n-1] - J[0][n-2]*D[n-2])/J[-1][n] + P[n-1]/J[-1][n] - -# return (z*P[N-1] - J[0][N-2]*P[N-2]), (z*D[N-1] - J[0][N-2]*D[N-2] + P[N-1]) - -# def gauss_quadrature(Jacobi_matrix,mass=1,niter=3,guess=None,report_error=False): - -# if guess == None: -# z = grid_guess(Jacobi_matrix) -# else: -# z = guess - -# for i in range(niter): -# P, D = remainders(Jacobi_matrix,z) -# if report_error: print(np.max(np.abs(P/D))) -# z = z - P/D - -# w = 1/((1-z**2)*D**2) -# w = (mass/np.sum(w))*w - -# return z, w - -# def three_term_recursion(Jacobi_matrix,grid,max_degree,init): - -# J, z, N = sparse.dia_matrix(Jacobi_matrix).data, grid, max_degree+1 - -# P = np.zeros((N,len(grid)),dtype=dtype) -# P[0] = init - -# if np.shape(J)[0] == 3: -# P[1] = (z-J[1][0])*P[0]/J[-1][1] -# for n in range(2,N): -# P[n] = ((z-J[1][n-1])*P[n-1] - J[0][n-2]*P[n-2])/J[-1][n] -# return P - -# P[1] = z*P[0]/J[-1][1] -# for n in range(2,N): -# P[n] = (z*P[n-1] - J[0][n-2]*P[n-2])/J[-1][n] -# return P - -# def polish_norm(functions,weights): -# for k in range(len(functions)): -# functions[k] = functions[k]/np.sqrt(np.sum(weights*functions[k]**2)) - -# def Jacobi_quadrature(max_degree,a,b,niter=3,report_error=False): - -# mu = mass(a,b) -# J = Jacobi_operator('J',a,b,max_degree) -# return gauss_quadrature(J,mass=mu,niter=niter,report_error=report_error) - -# def Jacobi_envelope(a,b,a0,b0,z): - -# mu = mass(a,b) -# return np.exp( ((a-a0)/2)*np.log(1-z) + ((b-b0)/2)*np.log(1+z) )/np.sqrt(mu) - -# def Jacobi_recursion(max_degree,a,b,grid,init): - -# J = Jacobi_operator('J',a,b,max_degree) - -# return three_term_recursion(J,grid,max_degree,init) - -# def push(operator,data): -# return (operator).dot(data) - -# def pull(operator,data): -# return (operator.transpose()).dot(data) - -def mass(a,b): - return np.exp( (a+b+1)*np.log(2) + fun.gammaln(a+1) + fun.gammaln(b+1) - fun.gammaln(a+b+2) ) - -# def Jacobi_operator(op,a,b,max_degree,format='csr',rescale=None): - -# def diag(bands,locs): -# return sparse.dia_matrix((bands,locs),shape=(len(bands[0]),len(bands[0]))) - -# N = max_degree+1 -# n = np.arange(0,N,dtype=dtype) -# na, nb, nab, nnab = n+a, n+b, n+a+b, 2*n+a+b - -# # (1-z) 0): -# if a+b==0: -# middle = na/(2*n+1) -# lower = (nb+1)/(2*n+1) -# middle[0] = 2*a -# else: -# middle = 2*na*nab/(nnab*(nnab+1)) -# lower = 2*(n+1)*(nb+1)/((nnab+1)*(nnab+2)) -# operator = diag([-np.sqrt(lower),np.sqrt(middle)],[-1,0]) - -# # 0): -# if a+b == 0: -# middle = nb/(2*n+1) -# lower = (na+1)/(2*n+1) -# middle[0] = 2*b -# else: -# middle = 2*nb*nab/(nnab*(nnab+1)) -# lower = 2*(n+1)*(na+1)/((nnab+1)*(nnab+2)) -# operator = diag([np.sqrt(lower),np.sqrt(middle)],[-1,0]) - -# # 0): -# operator = diag([np.sqrt(na*(nb+1))],[0]) - -# # ( b + (1+z)*d/dz ) 0): -# operator = diag([np.sqrt((na+1)*nb)],[0]) - -# # ( a(1+z) - b(1-z) - (1-z^2)*d/dz ) 0) and (b > 0): -# operator = diag([np.sqrt((n+1)*nab)],[-1]) - -# # d/dz 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 = jacobi.operator('A')(+1) - B = jacobi.operator('B')(+1) - - da, db = int(a1-a0), int(b1-b0) - conv = A**da @ B**db - - return conv(N, a0, b0).astype(output_dtype) - -def differentiation_matrix(N, a, b): - return jacobi.operator('D')(+1)(N, a, b).square.astype(output_dtype) - -def jacobi_matrix(N, a, b): - return jacobi.operator('Z')(N, a, b).square.astype(output_dtype) - -def integration_vector(N, a, b): - # Build Legendre quadrature - leg_grid, leg_weights = jacobi.quadrature(N, 0, 0) - # Evaluate polynomials on Legendre grid - interp = build_polynomials(N, a, b, leg_grid).T - # Compute integrals using Legendre quadrature - integ = leg_weights @ interp * (2 / mass(0, 0)) - # Zero out entries below output eps - cutoff = np.finfo(output_dtype).resolution - integ[np.abs(integ) <= cutoff] = 0. - return integ.astype(output_dtype) -