# Make domain
phi_basis = de.Fourier('phi',
2 * (L_max + 1),
interval=(0, 2 * np.pi),
dealias=L_dealias)
theta_basis = de.Fourier('theta',
L_max + 1,
interval=(0, np.pi),
dealias=L_dealias)
r_basis = de.Fourier('r', N_max + 1, interval=(0, 1), dealias=N_dealias)
domain = de.Domain([phi_basis, theta_basis, r_basis],
grid_dtype=np.float64,
mesh=mesh)
domain.global_coeff_shape = np.array([L_max + 1, L_max + 1, N_max + 1])
domain.distributor = Distributor(domain, comm, mesh)
mesh = domain.distributor.mesh
if len(mesh) == 0:
phi_layout = domain.distributor.layouts[3]
th_m_layout = domain.distributor.layouts[2]
ell_r_layout = domain.distributor.layouts[1]
r_ell_layout = domain.distributor.layouts[1]
elif len(mesh) == 1:
phi_layout = domain.distributor.layouts[4]
th_m_layout = domain.distributor.layouts[2]
ell_r_layout = domain.distributor.layouts[1]
r_ell_layout = domain.distributor.layouts[1]
elif len(mesh) == 2:
phi_layout = domain.distributor.layouts[5]
th_m_layout = domain.distributor.layouts[3]
phi_basis_3D = de.Fourier('phi',
2 * (L_max + 1),
interval=(0, 2 * np.pi),
dealias=L_dealias)
theta_basis_3D = de.Fourier('theta',
L_max + 1,
interval=(0, np.pi),
dealias=L_dealias)
r_basis_3D = de.Fourier('r', N_max + 1, interval=(0, 1), dealias=N_dealias)
domain_3D = de.Domain([phi_basis_3D, theta_basis_3D, r_basis_3D],
grid_dtype=np.float64,
comm=comm,
mesh=mesh)
domain_3D.global_coeff_shape = np.array([L_max + 1, L_max + 1, N_max + 1])
domain_3D.distributor = Distributor(domain_3D, comm, mesh)
phi_layout = domain_3D.distributor.layouts[3]
grid_slices = phi_layout.slices(domain_3D.dealias)
phi = domain_3D.grid(0, scales=domain_3D.dealias)
theta = B_3D.grid(1, dimensions=3)[:, grid_slices[1], :] # local
r = B_3D.grid(2, dimensions=3)[:, :, grid_slices[2]] # local
H = ball.TensorField_3D(1, B_3D, domain_3D)
H2 = ball.TensorField_3D(1, B_3D, domain_3D)
A = ball.TensorField_3D(1, B_3D, domain_3D)
H['g'][1] = -3. / 2. * r * (-1 + 4 * r**2 - 6 * r**4 +
3 * r**6) * (np.cos(phi) + np.sin(phi))
H['g'][2] = -3./4.*r*(-1+r**2)*np.cos(theta)* \
( 3*r*(2-5*r**2+4*r**4)*np.sin(theta)