def test_convert_constant_tensor(Nphi, Ntheta, Nr, k, dealias, basis, dtype): c, d, b, phi, theta, r, x, y, z = basis(Nphi, Ntheta, Nr, k, dealias, dtype) f = field.Field(dist=d, dtype=dtype, tensorsig=(c, c)) f['g'][0, 0] = f['g'][1, 1] = f['g'][2, 2] = 1 g = operators.Convert(f, b).evaluate() assert np.allclose(f['g'], g['g'])
def test_implicit_transpose_tensor(Nphi, Ntheta, Nr, k, dealias, dtype, basis): c, d, b, phi, theta, r, x, y, z = basis(Nphi, Ntheta, Nr, k, dealias, dtype) ct, st, cp, sp = np.cos(theta), np.sin(theta), np.cos(phi), np.sin(phi) u = field.Field(dist=d, bases=(b, ), tensorsig=(c, ), dtype=dtype) u.preset_scales(b.domain.dealias) u['g'][2] = r**2 * st * (2 * ct**2 * cp - r * ct**3 * sp + r**3 * cp**3 * st**5 * sp**3 + r * ct * st**2 * (cp**3 + sp**3)) u['g'][1] = r**2 * (2 * ct**3 * cp - r * cp**3 * st**4 + r**3 * ct * cp**3 * st**5 * sp**3 - 1 / 16 * r * np.sin(2 * theta)**2 * (-7 * sp + np.sin(3 * phi))) u['g'][0] = r**2 * sp * (-2 * ct**2 + r * ct * cp * st**2 * sp - r**3 * cp**2 * st**5 * sp**3) T = operators.Gradient(u, c).evaluate() Ttg = np.transpose(np.copy(T['g']), (1, 0, 2, 3, 4)) Tt = field.Field(dist=d, bases=(b, ), tensorsig=( c, c, ), dtype=dtype) trans = lambda A: operators.TransposeComponents(A) problem = problems.LBVP([Tt]) problem.add_equation((trans(Tt), T)) # Solver solver = solvers.LinearBoundaryValueSolver(problem) solver.solve() assert np.allclose(Tt['g'], Ttg)
def test_convert_constant_scalar(Nphi, Ntheta, Nr, k, dealias, basis, dtype): c, d, b, phi, theta, r, x, y, z = basis(Nphi, Ntheta, Nr, k, dealias, dtype) f = field.Field(dist=d, dtype=dtype) f['g'] = 1 g = operators.Convert(f, b).evaluate() assert np.allclose(f['g'], g['g'])
def test_angular_component_tensor(Nphi, Ntheta, Nr, k, dealias, dtype, basis, radius): c, d, b, phi, theta, r, x, y, z = basis(Nphi, Ntheta, Nr, k, dealias, dtype) T = field.Field(dist=d, bases=(b, ), tensorsig=(c, c), dtype=dtype) T.preset_scales(b.domain.dealias) T['g'][2, 2] = (6 * x**2 + 4 * y * z) / r**2 T['g'][2, 1] = T['g'][1, 2] = -2 * (y**3 + x**2 * (y - 3 * z) - y * z**2) / (r**3 * np.sin(theta)) T['g'][2, 0] = T['g'][0, 2] = 2 * x * (z - 3 * y) / (r**2 * np.sin(theta)) T['g'][1, 1] = 6 * x**2 / (r**2 * np.sin(theta)**2) - (6 * x**2 + 4 * y * z) / r**2 T['g'][1, 0] = T['g'][0, 1] = -2 * x * (x**2 + y**2 + 3 * y * z) / (r**3 * np.sin(theta)**2) T['g'][0, 0] = 6 * y**2 / (x**2 + y**2) A = operators.AngularComponent(operators.interpolate(T, r=radius), index=1).evaluate() Ag = 0 * A['g'] Ag[2, 1] = 6 * np.cos(theta) * np.cos(phi)**2 * np.sin(theta) + 2 * np.cos( 2 * theta) * np.sin(phi) Ag[2, 0] = 2 * np.cos(phi) * (np.cos(theta) - 3 * np.sin(theta) * np.sin(phi)) Ag[1, 1] = 2 * np.cos(theta) * (3 * np.cos(theta) * np.cos(phi)**2 - 2 * np.sin(theta) * np.sin(phi)) Ag[1, 0] = Ag[0, 1] = -2 * np.cos(phi) * (np.sin(theta) + 3 * np.cos(theta) * np.sin(phi)) Ag[0, 0] = 6 * np.sin(phi)**2 assert np.allclose(A['g'], Ag)
def test_tensor_prod_scalar(Nphi, Nr, basis, ncc_first, dealias, dtype): c, d, b, phi, r, x, y = basis(Nphi, Nr, dealias=dealias, dtype=dtype) f = field.Field(dist=d, bases=(b.radial_basis, ), dtype=dtype) g = field.Field(dist=d, bases=(b, ), dtype=dtype) f['g'] = r**4 g['g'] = 3 * x**2 + 2 * y T = operators.Gradient(operators.Gradient(f, c), c).evaluate() vars = [g] if ncc_first: U0 = T * g else: U0 = g * T U1 = U0.reinitialize(ncc=True, ncc_vars=vars) problem = problems.LBVP(vars) problem.add_equation((f * g, 0)) solver = solvers.LinearBoundaryValueSolver( problem, matsolver='SuperluNaturalSpsolve', matrix_coupling=[False, True]) U1.store_ncc_matrices(vars, solver.subproblems) U0 = U0.evaluate() U0.change_scales(1) U1 = U1.evaluate_as_ncc() assert np.allclose(U0['g'], U1['g'])
def test_tensor_dot_tensor(Nphi, Ntheta, Nr, basis, alpha, k_ncc, k_arg, dealias, dtype): c, d, b, phi, theta, r, x, y, z = basis(Nphi, Ntheta, Nr, alpha, dealias=dealias, dtype=dtype) b_ncc = b.radial_basis.clone_with(k=k_ncc) b_arg = b.clone_with(k=k_arg) T = field.Field(dist=d, bases=(b_ncc, ), tensorsig=(c, c), dtype=dtype) U = field.Field(dist=d, bases=(b_arg, ), tensorsig=(c, c), dtype=dtype) T['g'] = np.random.randn(*T['g'].shape) U['g'] = np.random.randn(*U['g'].shape) if isinstance(b, BallBasis): k_out = k_out_ball(k_ncc, k_arg, T, U) else: k_out = k_out_shell(k_ncc, k_arg) if k_out > 0 and dealias < 2: T['c'][:, :, :, :, -k_out:] = 0 U['c'][:, :, :, :, -k_out:] = 0 vars = [U] w0 = dot(T, U) w1 = w0.reinitialize(ncc=True, ncc_vars=vars) problem = problems.LBVP(vars) problem.add_equation((operators.Laplacian(U, c), 0)) solver = solvers.LinearBoundaryValueSolver(problem, ) w1.store_ncc_matrices(vars, solver.subproblems) w0 = w0.evaluate() w1 = w1.evaluate_as_ncc() assert np.allclose(w0['c'], w1['c'])
def test_scalar_prod_scalar(Nphi, Ntheta, Nr, basis, alpha, k_ncc, k_arg, dealias, dtype): c, d, b, phi, theta, r, x, y, z = basis(Nphi, Ntheta, Nr, alpha, dealias=dealias, dtype=dtype) b_ncc = b.radial_basis.clone_with(k=k_ncc) b_arg = b.clone_with(k=k_arg) f = field.Field(dist=d, bases=(b_ncc, ), dtype=dtype) g = field.Field(dist=d, bases=(b_arg, ), dtype=dtype) f['g'] = np.random.randn(*f['g'].shape) g['g'] = np.random.randn(*g['g'].shape) if isinstance(b, BallBasis): k_out = k_out_ball(k_ncc, k_arg, f, g) else: k_out = k_out_shell(k_ncc, k_arg) if k_out > 0 and dealias < 2: f['c'][:, :, -k_out:] = 0 g['c'][:, :, -k_out:] = 0 vars = [g] w0 = f * g w1 = w0.reinitialize(ncc=True, ncc_vars=vars) problem = problems.LBVP(vars) problem.add_equation((w1, 0)) solver = solvers.LinearBoundaryValueSolver(problem) w1.store_ncc_matrices(vars, solver.subproblems) w0 = w0.evaluate() w1 = w1.evaluate_as_ncc() assert np.allclose(w0['c'], w1['c'])
def test_scalar_prod_vector(Nphi, Nr, basis, ncc_first, dealias, dtype): c, d, b, phi, r, x, y = basis(Nphi, Nr, dealias=dealias, dtype=dtype) f = field.Field(dist=d, bases=(b.radial_basis, ), dtype=dtype) g = field.Field(dist=d, bases=(b, ), dtype=dtype) f['g'] = r**2 g['g'] = 3 * x**2 + 2 * y u = operators.Gradient(g, c).evaluate() vars = [u] if ncc_first: w0 = f * u else: w0 = u * f w1 = w0.reinitialize(ncc=True, ncc_vars=vars) problem = problems.LBVP(vars) problem.add_equation((w1, 0)) solver = solvers.LinearBoundaryValueSolver( problem, matsolver='SuperluNaturalSpsolve', matrix_coupling=[False, True]) w1.store_ncc_matrices(vars, solver.subproblems) w0 = w0.evaluate() w0.change_scales(1) w1 = w1.evaluate_as_ncc() assert np.allclose(w0['g'], w1['g'])
def test_angular_component_vector(Nphi, Ntheta, Nr, k, dealias, dtype, basis, radius): c, d, b, phi, theta, r, x, y, z = basis(Nphi, Ntheta, Nr, k, dealias, dtype) ct, st, cp, sp = np.cos(theta), np.sin(theta), np.cos(phi), np.sin(phi) u = field.Field(dist=d, bases=(b, ), tensorsig=(c, ), dtype=dtype) u.preset_scales(b.domain.dealias) u['g'][2] = r**2 * st * (2 * ct**2 * cp - r * ct**3 * sp + r**3 * cp**3 * st**5 * sp**3 + r * ct * st**2 * (cp**3 + sp**3)) u['g'][1] = r**2 * (2 * ct**3 * cp - r * cp**3 * st**4 + r**3 * ct * cp**3 * st**5 * sp**3 - 1 / 16 * r * np.sin(2 * theta)**2 * (-7 * sp + np.sin(3 * phi))) u['g'][0] = r**2 * sp * (-2 * ct**2 + r * ct * cp * st**2 * sp - r**3 * cp**2 * st**5 * sp**3) v = operators.AngularComponent(operators.interpolate(u, r=radius)).evaluate() vg = 0 * v['g'] vg[0] = radius**2 * sp * (-2 * ct**2 + radius * ct * cp * st**2 * sp - radius**3 * cp**2 * st**5 * sp**3) vg[1] = radius**2 * (2 * ct**3 * cp - radius * cp**3 * st**4 + radius**3 * ct * cp**3 * st**5 * sp**3 - 1 / 16 * radius * np.sin(2 * theta)**2 * (-7 * sp + np.sin(3 * phi))) assert np.allclose(v['g'], vg)
def test_interpolate_radius_tensor(Nphi, Ntheta, Nr, k, dealias, dtype, basis, r_interp): c, d, b, phi, theta, r, x, y, z = basis(Nphi, Ntheta, Nr, k, dealias, dtype) T = field.Field(dist=d, bases=(b, ), tensorsig=(c, c), dtype=dtype) T.preset_scales(b.domain.dealias) T['g'][2, 2] = (6 * x**2 + 4 * y * z) / r**2 T['g'][2, 1] = T['g'][1, 2] = -2 * (y**3 + x**2 * (y - 3 * z) - y * z**2) / (r**3 * np.sin(theta)) T['g'][2, 0] = T['g'][0, 2] = 2 * x * (z - 3 * y) / (r**2 * np.sin(theta)) T['g'][1, 1] = 6 * x**2 / (r**2 * np.sin(theta)**2) - (6 * x**2 + 4 * y * z) / r**2 T['g'][1, 0] = T['g'][0, 1] = -2 * x * (x**2 + y**2 + 3 * y * z) / (r**3 * np.sin(theta)**2) T['g'][0, 0] = 6 * y**2 / (x**2 + y**2) A = operators.interpolate(T, r=r_interp).evaluate() Ag = 0 * A['g'] r = np.array([[[r_interp]]]) x, y, z = c.cartesian(phi, theta, r) Ag[2, 2] = (6 * x**2 + 4 * y * z) / r**2 Ag[2, 1] = Ag[1, 2] = -2 * (y**3 + x**2 * (y - 3 * z) - y * z**2) / (r**3 * np.sin(theta)) Ag[2, 0] = Ag[0, 2] = 2 * x * (z - 3 * y) / (r**2 * np.sin(theta)) Ag[1, 1] = 6 * x**2 / (r**2 * np.sin(theta)**2) - (6 * x**2 + 4 * y * z) / r**2 Ag[1, 0] = Ag[0, 1] = -2 * x * (x**2 + y**2 + 3 * y * z) / (r**3 * np.sin(theta)**2) Ag[0, 0] = 6 * y**2 / (x**2 + y**2) assert np.allclose(A['g'], Ag)
def test_interpolate_colatitude_vector(Nphi, Ntheta, Nr, k, dealias, dtype, basis, theta_interp): c, d, b, phi, theta, r, x, y, z = basis(Nphi, Ntheta, Nr, k, dealias, dtype) ct, st, cp, sp = np.cos(theta), np.sin(theta), np.cos(phi), np.sin(phi) u = field.Field(dist=d, bases=(b, ), tensorsig=(c, ), dtype=dtype) u.preset_scales(b.domain.dealias) u['g'][0] = r**2 * sp * (-2 * ct**2 + r * ct * cp * st**2 * sp - r**3 * cp**2 * st**5 * sp**3) u['g'][1] = r**2 * (2 * ct**3 * cp - r * cp**3 * st**4 + r**3 * ct * cp**3 * st**5 * sp**3 - 1 / 16 * r * np.sin(2 * theta)**2 * (-7 * sp + np.sin(3 * phi))) u['g'][2] = r**2 * st * (2 * ct**2 * cp - r * ct**3 * sp + r**3 * cp**3 * st**5 * sp**3 + r * ct * st**2 * (cp**3 + sp**3)) v = operators.interpolate(u, theta=theta_interp).evaluate() vg = 0 * v['g'] theta = np.array([[[theta_interp]]]) ct, st = np.cos(theta), np.sin(theta) vg[0] = r**2 * sp * (-2 * ct**2 + r * ct * cp * st**2 * sp - r**3 * cp**2 * st**5 * sp**3) vg[1] = r**2 * (2 * ct**3 * cp - r * cp**3 * st**4 + r**3 * ct * cp**3 * st**5 * sp**3 - 1 / 16 * r * np.sin(2 * theta)**2 * (-7 * sp + np.sin(3 * phi))) vg[2] = r**2 * st * (2 * ct**2 * cp - r * ct**3 * sp + r**3 * cp**3 * st**5 * sp**3 + r * ct * st**2 * (cp**3 + sp**3)) assert np.allclose(v['g'], vg)
def test_azimuthal_average_scalar(Nphi, Nr, k, dealias, dtype, basis): c, d, b, phi, r, x, y = basis(Nphi, Nr, k, dealias, dtype) f = field.Field(dist=d, bases=(b, ), dtype=dtype) f.preset_scales(b.domain.dealias) f['g'] = r**2 + x h = operators.Average(f, c.coords[0]).evaluate() hg = r**2 assert np.allclose(h['g'], hg)
def test_sphere_complex_scalar_backward(Nphi, Ntheta, radius, basis, dealias): c, d, b, phi, theta = basis(Nphi, Ntheta, radius, dealias, np.complex128) f = field.Field(dist=d, bases=(b, ), dtype=np.complex128) f.preset_scales(dealias) m, ell, *_ = d.coeff_layout.local_group_arrays(b.domain, scales=1) f['c'][(m == -2) * (ell == 2)] = 1 fg = np.sqrt(15) / 4 * np.sin(theta)**2 * np.exp(-2j * phi) assert np.allclose(fg, f['g'])
def test_laplacian_scalar(Nphi, Nr, dealias, basis, dtype): c, d, b, phi, r, x, y = basis(Nphi, Nr, dealias, dtype) f = field.Field(dist=d, bases=(b, ), dtype=dtype) f.preset_scales(b.domain.dealias) f['g'] = x**4 + 2 * y**4 h = operators.Laplacian(f, c).evaluate() hg = 12 * x**2 + 24 * y**2 assert np.allclose(h['g'], hg)
def test_gradient_radial_scalar(Nr, dealias, basis, dtype): Nphi = 1 c, d, b, phi, r, x, y = basis(Nphi, Nr, dealias, dtype) f = field.Field(dist=d, bases=(b, ), dtype=dtype) f.preset_scales(b.domain.dealias) f['g'] = r**4 u = operators.Gradient(f, c).evaluate() ug = [0 * r * phi, 4 * r**3 + 0 * phi] assert np.allclose(u['g'], ug)
def test_divergence_radial_vector(Nr, dealias, basis, dtype): c, d, b, phi, theta, r, x, y, z = basis(1, 1, Nr, 1, dtype=dtype) f = field.Field(dist=d, bases=(b,), dtype=dtype) f.preset_scales(b.domain.dealias) f['g'] = r**4/3 u = operators.Gradient(f, c) h = operators.Divergence(u).evaluate() hg = 20/3*r**2 assert np.allclose(h['g'], hg)
def test_laplacian_radial_scalar(Nr, dealias, basis, dtype): Nphi = 1 c, d, b, phi, r, x, y = basis(Nphi, Nr, dealias, dtype=dtype) f = field.Field(dist=d, bases=(b, ), dtype=dtype) f.preset_scales(b.domain.dealias) f['g'] = r**2 h = operators.Laplacian(f, c).evaluate() hg = 4 assert np.allclose(h['g'], hg)
def test_spherical_average_scalar(Nphi, Ntheta, Nr, k, dealias, dtype, basis): c, d, b, phi, theta, r, x, y, z = basis(Nphi, Ntheta, Nr, k, dealias, dtype) f = field.Field(dist=d, bases=(b, ), dtype=dtype) f.preset_scales(b.domain.dealias) f['g'] = r**2 + x + z h = operators.Average(f, c.S2coordsys).evaluate() hg = r**2 assert np.allclose(h['g'], hg)
def test_divergence_vector(Nphi, Ntheta, Nr, dealias, basis, dtype): c, d, b, phi, theta, r, x, y, z = basis(Nphi, Ntheta, Nr, dealias, dtype) f = field.Field(dist=d, bases=(b,), dtype=dtype) f.preset_scales(b.domain.dealias) f['g'] = x**3 + 2*y**3 + 3*z**3 u = operators.Gradient(f, c) h = operators.Divergence(u).evaluate() hg = 6*x + 12*y + 18*z assert np.allclose(h['g'], hg)
def test_multiply_vector_vector(Nphi, Ntheta, Nr, dealias, basis): c, d, b, phi, theta, r, x, y, z = basis(Nphi, Ntheta, Nr, dealias, np.complex128) f = field.Field(dist=d, bases=(b, ), dtype=np.complex128) f['g'] = x**3 + 2 * y**3 + 3 * z**3 u = operators.Gradient(f, c).evaluate() T = (u * u).evaluate() Tg = u['g'][None, ...] * u['g'][:, None, ...] assert np.allclose(T['g'], Tg)
def test_gradient_scalar(Nphi, Nr, dealias, basis, dtype): c, d, b, phi, r, x, y = basis(Nphi, Nr, dealias, dtype) f = field.Field(dist=d, bases=(b, ), dtype=dtype) f.preset_scales(b.domain.dealias) f['g'] = 3 * x**2 + 2 * y u = operators.Gradient(f, c).evaluate() ex = np.array([-np.sin(phi) + 0. * r, np.cos(phi) + 0. * r]) ey = np.array([np.cos(phi) + 0. * r, np.sin(phi) + 0. * r]) ug = 6 * x * ex + 2 * ey assert np.allclose(u['g'], ug)
def test_curl_vector(Nphi, Nr, dealias, basis, dtype): c, d, b, phi, r, x, y = basis(Nphi, Nr, dealias, dtype) v = field.Field(dist=d, bases=(b, ), tensorsig=(c, ), dtype=dtype) v.preset_scales(b.domain.dealias) ex = np.array([-np.sin(phi) + 0. * r, np.cos(phi) + 0. * r]) ey = np.array([np.cos(phi) + 0. * r, np.sin(phi) + 0. * r]) v['g'] = 4 * x**3 * ey + 3 * y**2 * ey u = operators.Curl(v).evaluate() ug = 12 * x**2 assert np.allclose(u['g'], ug)
def test_sphere_real_scalar_forward(Nphi, Ntheta, radius, basis, dealias): c, d, b, phi, theta = basis(Nphi, Ntheta, radius, dealias, np.float64) f = field.Field(dist=d, bases=(b, ), dtype=np.float64) f.preset_scales(dealias) m, ell, *_ = d.coeff_layout.local_group_arrays(b.domain, scales=1) f['g'] = np.sqrt(15) / 4 * np.sin(theta)**2 * (np.cos(2 * phi) - np.sin(2 * phi)) fc = np.zeros_like(f['c']) fc[(m == 2) * (ell == 2)] = 1 assert np.allclose(fc, f['c'])
def test_laplacian_vector(Nphi, Nr, dealias, basis, dtype): c, d, b, phi, r, x, y = basis(Nphi, Nr, dealias, dtype) v = field.Field(dist=d, tensorsig=(c, ), bases=(b, ), dtype=dtype) v.preset_scales(b.domain.dealias) ex = np.array([-np.sin(phi) + 0. * r, np.cos(phi) + 0. * r]) ey = np.array([np.cos(phi) + 0. * r, np.sin(phi) + 0. * r]) v['g'] = 4 * x**3 * ey + 3 * y**2 * ey U = operators.Laplacian(v, c).evaluate() Ug = (24 * x + 6) * ey assert np.allclose(U['g'], Ug)
def test_interpolate_radius_scalar(Nphi, Nr, k, dealias, basis, dtype, r_interp): c, d, b, phi, r, x, y = basis(Nphi, Nr, k, dealias, dtype) f = field.Field(dist=d, bases=(b, ), dtype=dtype) f.preset_scales(b.domain.dealias) f['g'] = x**4 + 2 * y**4 h = operators.interpolate(f, r=r_interp).evaluate() x, y = c.cartesian(phi, np.array([[r_interp]])) hg = x**4 + 2 * y**4 assert np.allclose(h['g'], hg)
def test_laplacian_radial_vector(Nr, dealias, basis, dtype): Nphi = 1 c, d, b, phi, r, x, y = basis(Nphi, Nr, dealias, dtype=dtype) u = field.Field(dist=d, bases=(b, ), tensorsig=(c, ), dtype=dtype) u.preset_scales(b.domain.dealias) u['g'][1] = 4 * r**3 v = operators.Laplacian(u, c).evaluate() vg = 0 * v['g'] vg[1] = 32 * r assert np.allclose(v['g'], vg)
def test_transpose_explicit(basis, Nphi, Nr, k, dealias, dtype, layout): c, d, b, phi, r, x, y = basis(Nphi, Nr, k, dealias, dtype) # Random tensor field f = d.TensorField((c, c), bases=b) f.fill_random(layout='g') f.low_pass_filter(scales=0.75) # Evaluate transpose f.change_layout(layout) g = operators.transpose(f).evaluate() assert np.allclose(g['g'], np.transpose(f['g'], (1, 0, 2, 3)))
def test_convert_scalar(Nphi, Nr, k, dealias, basis, dtype, layout): c, d, b, phi, r, x, y = basis(Nphi, Nr, k, dealias, dtype) f = field.Field(dist=d, bases=(b, ), dtype=dtype) f.preset_scales(b.domain.dealias) f['g'] = 3 * x**2 + 2 * y g = operators.Laplacian(f, c).evaluate() f.change_layout(layout) g.change_layout(layout) h = (f + g).evaluate() assert np.allclose(h['g'], f['g'] + g['g'])
def test_multiply_tensor_tensor(Nphi, Ntheta, Nr, dealias, basis): c, d, b, phi, theta, r, x, y, z = basis(Nphi, Ntheta, Nr, dealias, np.complex128) f = field.Field(dist=d, bases=(b, ), dtype=np.complex128) f['g'] = x**3 + 2 * y**3 + 3 * z**3 u = operators.Gradient(f, c).evaluate() T = operators.Gradient(u, c).evaluate() Q = (T * T).evaluate() Qg = T['g'][:, :, None, None, ...] * T['g'][None, None, ...] assert np.allclose(Q['g'], Qg)
def test_multiply_scalar_scalar(Nphi, Ntheta, Nr, dealias, basis): c, d, b, phi, theta, r, x, y, z = basis(Nphi, Ntheta, Nr, dealias, np.complex128) f = field.Field(dist=d, bases=(b, ), dtype=np.complex128) f['g'] = x**3 + 2 * y**3 + 3 * z**3 h = (f * f).evaluate() phi, theta, r = b.local_grids(b.domain.dealias) x, y, z = c.cartesian(phi, theta, r) hg = (x**3 + 2 * y**3 + 3 * z**3)**2 assert np.allclose(h['g'], hg)