Add SphereEllProduct and tests

Author: kburns

dedalus/core/operators.py

@@ -2654,6 +2654,36 @@ def operate(self, out):
                     comp_out += symbols * comp_in # TEMPORARY
 
 
+class SphereEllProduct(SeparableSphereOperator):
+
+    name = "SphereEllProduct"
+    complex_operator = False
+    subaxis_dependence = [False, True]
+
+    def __init__(self, operand, coordsys, ell_r_func, out=None):
+        super().__init__(operand, out=out)
+        self.ell_r_func = ell_r_func
+        self.coordsys = coordsys
+        self.operand = operand
+        self.input_basis = operand.domain.get_basis(coordsys)
+        self.output_basis = self.input_basis
+        self.first_axis = self.input_basis.first_axis
+        self.last_axis = self.input_basis.last_axis
+        # FutureField requirements
+        self.domain  = operand.domain
+        self.tensorsig = operand.tensorsig
+        self.dtype = operand.dtype
+
+    def symbol(self, spinindex_in, spinindex_out, local_ell, radius):
+        return self.ell_r_func(local_ell, radius)
+
+    def new_operand(self, operand, **kw):
+        return SphereEllProduct(operand, self.coordsys, self.ell_r_func, **kw)
+
+    def spinindex_out(self, spinindex_in):
+        return (spinindex_in,)
+
+
 class PolarMOperator(SpectralOperator):
 
     subaxis_dependence = [True, True]  # Depends on m and n

dedalus/tests/test_sphere2D_calculus.py

@@ -470,3 +470,36 @@ def test_divergence_cleaning(Nphi, Ntheta, dealias, dtype):
     h.change_scales(1)
     assert np.allclose(u['g'], h['g'])
 
+
+@pytest.mark.parametrize('Nphi', Nphi_range)
+@pytest.mark.parametrize('Ntheta', Ntheta_range)
+@pytest.mark.parametrize('dealias', dealias_range)
+@pytest.mark.parametrize('dtype', [np.float64, np.complex128])
+def test_sphere_ell_product_scalar(Nphi, Ntheta, dealias, dtype):
+    c, d, b, phi, theta = build_sphere(Nphi, Ntheta, dealias, dtype)
+    f = field.Field(dist=d, bases=(b,), dtype=dtype)
+    g = field.Field(dist=d, bases=(b,), dtype=dtype)
+    f.fill_random('g')
+    func = lambda ell, r: ell + 3
+    for ell, m_ind, ell_ind in b.ell_maps:
+        g['c'][m_ind, ell_ind] = func(ell, b.radius) * f['c'][m_ind, ell_ind]
+    h = operators.SphereEllProduct(f, c, func).evaluate()
+    assert np.allclose(g['c'], h['c'])
+
+
+@pytest.mark.parametrize('Nphi', Nphi_range)
+@pytest.mark.parametrize('Ntheta', Ntheta_range)
+@pytest.mark.parametrize('dealias', dealias_range)
+@pytest.mark.parametrize('dtype', [np.float64, np.complex128])
+def test_sphere_ell_product_vector(Nphi, Ntheta, dealias, dtype):
+    c, d, b, phi, theta = build_sphere(Nphi, Ntheta, dealias, dtype)
+    f = field.Field(dist=d, bases=(b,), dtype=dtype, tensorsig=(c,))
+    g = field.Field(dist=d, bases=(b,), dtype=dtype, tensorsig=(c,))
+    f.fill_random('g')
+    func = lambda ell, r: ell + 3
+    for ell, m_ind, ell_ind in b.ell_maps:
+        for i in range(c.dim):
+            g['c'][i, m_ind, ell_ind] = func(ell, b.radius) * f['c'][i, m_ind, ell_ind]
+    h = operators.SphereEllProduct(f, c, func).evaluate()
+    assert np.allclose(g['c'], h['c'])
+