Add some helpers for boundary bases. Reduce sparse and cast warnings.

dedalus/core/basis.py

@@ -333,12 +333,17 @@ def enum_indices(tensorsig):
                             if matrix.shape != (M,N):
                                 m, n = matrix.shape
                                 matrix = sparse.kron(sparse.eye(M//m, N//n), matrix)
-                        if G.imag != 0 and product.dtype == np.float64:
-                            raise NotImplementedError()
-                        block += G * matrix
+                        if product.dtype == np.float64:
+                            if G.imag != 0:
+                                raise NotImplementedError()
+                            else:
+                                block += G.real * matrix
+                        else:
+                            block += G * matrix
                 block_row.append(block)
             blocks.append(block_row)
-        return sparse.bmat(blocks, format='csr')
+        matrix = sparse.bmat(blocks, format='csr')
+        return matrix
         #return getattr(ncc_basis, self.ncc_method)(arg_basis, coeffs, ncc_ts, arg_ts, out_ts, subproblem, ncc_first, *gamma_args, cutoff=1e-6)
         # tshape = [cs.dim for cs in ncc.tensorsig]
         # self._ncc_matrices = [self._ncc_matrix_recursion(ncc.data[ind], ncc.domain.full_bases, operand.domain.full_bases, separability, **kw) for ind in np.ndindex(*tshape)]
@@ -2015,6 +2020,8 @@ def __init__(self, coordsystem, shape, dtype, radii=(1,2), k=0, alpha=(-0.5,-0.5
         self.rho = (radii[1] + radii[0])/self.dR
         self.alpha = tuple(alpha)
         self.grid_params = (coordsystem, self.radii, self.alpha, self.dealias)
+        self.inner_edge = self.S1_basis(radii[0])
+        self.outer_edge = self.S1_basis(radii[1])
 
     @CachedAttribute
     def radial_basis(self):
@@ -2266,6 +2273,7 @@ def __init__(self, coordsystem, shape, dtype, radius=1, k=0, alpha=0, dealias=(1
             logger.warning("You are using more azimuthal modes than can be resolved with your current radial resolution")
             #raise ValueError("shape[0] cannot be more than twice shape[1].")
         self.grid_params = (coordsystem, radius, alpha, self.dealias)
+        self.edge = self.S1_basis(radius)
 
     @CachedAttribute
     def radial_basis(self):
@@ -3151,11 +3159,11 @@ def _last_axis_component_ncc_matrix(cls, subproblem, ncc_basis, arg_basis, out_b
             # ell_min of data is based off |m|, but ell_min of matrix takes into account |s| and |m|
             lmin_out = max(abs(spintotal_out) - abs(m), 0)
             lmin_arg = max(abs(spintotal_arg) - abs(m), 0 )
-            matrix = sparse.csr_matrix((N, N), dtype=np.complex128)
+            matrix = sparse.lil_matrix((N, N), dtype=np.complex128)
             matrix[lmin_out:,lmin_arg:] = (prefactor @ clenshaw.matrix_clenshaw(coeffs_filter, A, B, f0, cutoff=cutoff))[:N-lmin_out,:N-lmin_arg]
             if m < 0:
                 matrix = matrix[::-1, ::-1] #for negative m, ell's go in opposite direction
-        return matrix
+        return matrix.tocsr()
 
 
 class ConvertConstantSphere(operators.ConvertConstant, operators.SeparableSphereOperator):
@@ -4023,9 +4031,9 @@ def _last_axis_component_ncc_matrix(cls, subproblem, ncc_basis, arg_basis, out_b
                 matrix = (prefactor @ clenshaw.matrix_clenshaw(coeffs_filter, A, B, f0, cutoff=cutoff))[:N, :N]
         else:
             if ncc_basis.dtype == np.float64:
-                matrix = sparse.csr_matrix((2*N, 2*N))
+                matrix = sparse.csr_matrix((2*N, 2*N), dtype=np.float64)
             elif ncc_basis.dtype == np.complex128:
-                matrix = sparse.csr_matrix((N, N))
+                matrix = sparse.csr_matrix((N, N), dtype=np.complex128)
         return matrix
 
 
@@ -4249,6 +4257,8 @@ def __init__(self, coordsystem, shape, dtype, radii=(1,2), alpha=(-0.5,-0.5), de
         self.backward_transforms = [self.backward_transform_azimuth,
                                     self.backward_transform_colatitude,
                                     self.backward_transform_radius]
+        self.inner_surface = self.S2_basis(radius=radii[0])
+        self.outer_surface = self.S2_basis(radius=radii[1])
 
     def __eq__(self, other):
         if isinstance(other, ShellBasis):
@@ -4434,6 +4444,7 @@ def __init__(self, coordsystem, shape, dtype, radius=1, k=0, alpha=0, dealias=(1
         self.backward_transforms = [self.backward_transform_azimuth,
                                     self.backward_transform_colatitude,
                                     self.backward_transform_radius]
+        self.surface = self.S2_basis(radius=radius)
 
     def __eq__(self, other):
         if isinstance(other, BallBasis):

dedalus/tools/array.py

@@ -345,9 +345,10 @@ def interleave_matrices(matrices):
     if N == 1:
         return matrices[0]
     sum = 0
+    P = sparse.lil_matrix((N, N))
     for i, matrix in enumerate(matrices):
-        P = sparse.csr_matrix((N, N))
         P[i, i] = 1
         sum += sparse.kron(matrix, P)
+        P[i, i] = 0
     return sum
 

examples/evp_1d_waves_on_a_string/waves_on_a_string.py

@@ -72,7 +72,8 @@
 x = dist.local_grid(xbasis)
 for n, idx in enumerate(np.argsort(solver.eigenvalues)[:5], start=1):
     solver.set_state(idx, solver.subsystems[0])
-    plt.plot(x, (u['g']/np.max(np.abs(u['g']))).real, label=f"n={n}")
+    ug = (u['g'] / u['g'][1]).real
+    plt.plot(x, ug/np.max(np.abs(ug)), label=f"n={n}")
 plt.xlim(0, 1)
 plt.legend(loc="lower right")
 plt.ylabel(r"mode structure")

examples/evp_shell_rotating_convection/rotating_convection.py

@@ -53,29 +53,29 @@
 # Bases
 coords = d3.SphericalCoordinates('phi', 'theta', 'r')
 dist = d3.Distributor(coords, dtype=dtype)
-basis = d3.ShellBasis(coords, shape=(Nphi, Ntheta, Nr), radii=(Ri, Ro), dtype=dtype)
-basis_S2 = basis.S2_basis()
-phi, theta, r = dist.local_grids(basis)
+shell = d3.ShellBasis(coords, shape=(Nphi, Ntheta, Nr), radii=(Ri, Ro), dtype=dtype)
+sphere = shell.outer_surface
+phi, theta, r = dist.local_grids(shell)
 
 # Fields
 om = dist.Field(name='om')
-u = dist.VectorField(coords, name='u', bases=basis)
-p = dist.Field(name='p', bases=basis)
-T = dist.Field(name='T', bases=basis)
-tau_u1 = dist.VectorField(coords, bases=basis_S2)
-tau_u2 = dist.VectorField(coords, bases=basis_S2)
-tau_T1 = dist.Field(bases=basis_S2)
-tau_T2 = dist.Field(bases=basis_S2)
+u = dist.VectorField(coords, name='u', bases=shell)
+p = dist.Field(name='p', bases=shell)
+T = dist.Field(name='T', bases=shell)
+tau_u1 = dist.VectorField(coords, bases=sphere)
+tau_u2 = dist.VectorField(coords, bases=sphere)
+tau_T1 = dist.Field(bases=sphere)
+tau_T2 = dist.Field(bases=sphere)
 tau_p = dist.Field()
 
 # Substitutions
 dt = lambda A: -1j*om*A
-rvec = dist.VectorField(coords, bases=basis.meridional_basis)
+rvec = dist.VectorField(coords, bases=shell.meridional_basis)
 rvec['g'][2] = r
-ez = dist.VectorField(coords, bases=basis.meridional_basis)
+ez = dist.VectorField(coords, bases=shell.meridional_basis)
 ez['g'][1] = -np.sin(theta)
 ez['g'][2] = np.cos(theta)
-lift_basis = basis.derivative_basis(1)
+lift_basis = shell.derivative_basis(1)
 lift = lambda A: d3.Lift(A, lift_basis, -1)
 grad_u = d3.grad(u) + rvec*lift(tau_u1) # First-order reduction
 grad_T = d3.grad(T) + rvec*lift(tau_T1) # First-order reduction

examples/ivp_ball_internally_heated_convection/internally_heated_convection.py

@@ -54,25 +54,25 @@
 # Bases
 coords = d3.SphericalCoordinates('phi', 'theta', 'r')
 dist = d3.Distributor(coords, dtype=dtype, mesh=mesh)
-basis = d3.BallBasis(coords, shape=(Nphi, Ntheta, Nr), radius=1, dealias=dealias, dtype=dtype)
-S2_basis = basis.S2_basis()
+ball = d3.BallBasis(coords, shape=(Nphi, Ntheta, Nr), radius=1, dealias=dealias, dtype=dtype)
+sphere = ball.surface
 
 # Fields
-u = dist.VectorField(coords, name='u',bases=basis)
-p = dist.Field(name='p', bases=basis)
-T = dist.Field(name='T', bases=basis)
+u = dist.VectorField(coords, name='u',bases=ball)
+p = dist.Field(name='p', bases=ball)
+T = dist.Field(name='T', bases=ball)
 tau_p = dist.Field(name='tau_p')
-tau_u = dist.VectorField(coords, name='tau u', bases=S2_basis)
-tau_T = dist.Field(name='tau T', bases=S2_basis)
+tau_u = dist.VectorField(coords, name='tau u', bases=sphere)
+tau_T = dist.Field(name='tau T', bases=sphere)
 
 # Substitutions
-phi, theta, r = dist.local_grids(basis)
-r_vec = dist.VectorField(coords, bases=basis.radial_basis)
+phi, theta, r = dist.local_grids(ball)
+r_vec = dist.VectorField(coords, bases=ball.radial_basis)
 r_vec['g'][2] = r
 T_source = 6
 kappa = (Rayleigh * Prandtl)**(-1/2)
 nu = (Rayleigh / Prandtl)**(-1/2)
-lift = lambda A: d3.Lift(A, basis, -1)
+lift = lambda A: d3.Lift(A, ball, -1)
 strain_rate = d3.grad(u) + d3.trans(d3.grad(u))
 shear_stress = d3.angular(d3.radial(strain_rate(r=1), index=1))
 

examples/ivp_disk_libration/libration.py

@@ -41,23 +41,23 @@
 # Bases
 coords = d3.PolarCoordinates('phi', 'r')
 dist = d3.Distributor(coords, dtype=dtype)
-basis = d3.DiskBasis(coords, shape=(Nphi, Nr), radius=1, dealias=dealias, dtype=dtype)
-S1_basis = basis.S1_basis(radius=1)
+disk = d3.DiskBasis(coords, shape=(Nphi, Nr), radius=1, dealias=dealias, dtype=dtype)
+edge = disk.edge
 
 # Fields
-u = dist.VectorField(coords, name='u', bases=basis)
-p = dist.Field(name='p', bases=basis)
-tau_u = dist.VectorField(coords, name='tau_u', bases=S1_basis)
+u = dist.VectorField(coords, name='u', bases=disk)
+p = dist.Field(name='p', bases=disk)
+tau_u = dist.VectorField(coords, name='tau_u', bases=edge)
 tau_p = dist.Field(name='tau_p')
 
 # Substitutions
-phi, r = dist.local_grids(basis)
+phi, r = dist.local_grids(disk)
 nu = Ekman
-lift = lambda A: d3.Lift(A, basis, -1)
+lift = lambda A: d3.Lift(A, disk, -1)
 
 # Background librating flow
-u0_real = dist.VectorField(coords, bases=basis)
-u0_imag = dist.VectorField(coords, bases=basis)
+u0_real = dist.VectorField(coords, bases=disk)
+u0_imag = dist.VectorField(coords, bases=disk)
 u0_real['g'][0] = Ro * np.real(jv(1, (1-1j)*r/np.sqrt(2*Ekman)) / jv(1, (1-1j)/np.sqrt(2*Ekman)))
 u0_imag['g'][0] = Ro * np.imag(jv(1, (1-1j)*r/np.sqrt(2*Ekman)) / jv(1, (1-1j)/np.sqrt(2*Ekman)))
 t = dist.Field()

examples/ivp_shell_convection/shell_convection.py

@@ -44,28 +44,28 @@
 # Bases
 coords = d3.SphericalCoordinates('phi', 'theta', 'r')
 dist = d3.Distributor(coords, dtype=dtype, mesh=mesh)
-basis = d3.ShellBasis(coords, shape=(Nphi, Ntheta, Nr), radii=(Ri, Ro), dealias=dealias, dtype=dtype)
-s2_basis = basis.S2_basis()
+shell = d3.ShellBasis(coords, shape=(Nphi, Ntheta, Nr), radii=(Ri, Ro), dealias=dealias, dtype=dtype)
+sphere = shell.outer_surface
 
 # Fields
-p = dist.Field(name='p', bases=basis)
-b = dist.Field(name='b', bases=basis)
-u = dist.VectorField(coords, name='u', bases=basis)
+p = dist.Field(name='p', bases=shell)
+b = dist.Field(name='b', bases=shell)
+u = dist.VectorField(coords, name='u', bases=shell)
 tau_p = dist.Field(name='tau_p')
-tau_b1 = dist.Field(name='tau_b1', bases=s2_basis)
-tau_b2 = dist.Field(name='tau_b2', bases=s2_basis)
-tau_u1 = dist.VectorField(coords, name='tau_u1', bases=s2_basis)
-tau_u2 = dist.VectorField(coords, name='tau_u2', bases=s2_basis)
+tau_b1 = dist.Field(name='tau_b1', bases=sphere)
+tau_b2 = dist.Field(name='tau_b2', bases=sphere)
+tau_u1 = dist.VectorField(coords, name='tau_u1', bases=sphere)
+tau_u2 = dist.VectorField(coords, name='tau_u2', bases=sphere)
 
 # Substitutions
 kappa = (Rayleigh * Prandtl)**(-1/2)
 nu = (Rayleigh / Prandtl)**(-1/2)
-phi, theta, r = dist.local_grids(basis)
-er = dist.VectorField(coords, bases=basis.radial_basis)
+phi, theta, r = dist.local_grids(shell)
+er = dist.VectorField(coords, bases=shell.radial_basis)
 er['g'][2] = 1
-rvec = dist.VectorField(coords, bases=basis.radial_basis)
+rvec = dist.VectorField(coords, bases=shell.radial_basis)
 rvec['g'][2] = r
-lift_basis = basis.derivative_basis(1)
+lift_basis = shell.derivative_basis(1)
 lift = lambda A: d3.Lift(A, lift_basis, -1)
 grad_u = d3.grad(u) + rvec*lift(tau_u1) # First-order reduction
 grad_b = d3.grad(b) + rvec*lift(tau_b1) # First-order reduction

examples/nlbvp_ball_lane_emden/lane_emden.py

@@ -42,8 +42,6 @@
 import logging
 logger = logging.getLogger(__name__)
 
-# TODO: print NCC bandwidths and optimize parameters
-
 
 # Parameters
 Nr = 64
@@ -56,22 +54,22 @@
 # Bases
 coords = d3.SphericalCoordinates('phi', 'theta', 'r')
 dist = d3.Distributor(coords, dtype=dtype)
-basis = d3.BallBasis(coords, (1, 1, Nr), radius=1, dtype=dtype, dealias=dealias)
+ball = d3.BallBasis(coords, (1, 1, Nr), radius=1, dtype=dtype, dealias=dealias)
 
 # Fields
-f = dist.Field(name='f', bases=basis)
-tau = dist.Field(name='tau', bases=basis.S2_basis(radius=1))
+f = dist.Field(name='f', bases=ball)
+tau = dist.Field(name='tau', bases=ball.surface)
 
 # Substitutions
-lift = lambda A: d3.Lift(A, basis, -1)
+lift = lambda A: d3.Lift(A, ball, -1)
 
 # Problem
 problem = d3.NLBVP([f, tau], namespace=locals())
 problem.add_equation("lap(f) + lift(tau) = - f**n")
 problem.add_equation("f(r=1) = 0")
 
 # Initial guess
-phi, theta, r = dist.local_grids(basis)
+phi, theta, r = dist.local_grids(ball)
 R0 = 5
 f['g'] = R0**(2/(n-1)) * (1 - r**2)**2
 
@@ -109,7 +107,7 @@
 
 # Plot solution
 plt.figure(figsize=(6, 4))
-_, _, r = dist.local_grids(basis, scales=(dealias,dealias,dealias))
+_, _, r = dist.local_grids(ball, scales=(dealias,dealias,dealias))
 alpha = np.linspace(0.2, 1, len(steps))
 color = ('C0',) * (len(steps)-1) + ('C1',)
 for i, step in enumerate(steps):