def __init__(self, L_max, S_max):
self.L_max = L_max
self.S_max = S_max
# Domain
phi_basis = de.Fourier('phi', 2 * (L_max + 1), interval=(0, 2 * np.pi))
theta_basis = de.Fourier('theta', L_max + 1, interval=(0, np.pi))
self.domain = de.Domain([phi_basis, theta_basis],
grid_dtype=np.float64)
# Local m
layout0 = self.domain.distributor.layouts[0]
self.m_start = layout0.start(1)[0]
self.m_len = layout0.local_shape(1)[0]
self.m_end = self.m_start + self.m_len - 1
self.local_m = np.arange(self.m_start, self.m_end + 1)
# Sphere wrapper
self.sphere_wrapper = sphere_wrapper.Sphere(L_max,
S_max,
m_min=self.m_start,
m_max=self.m_end)
# Grids
self.phi_grid = self.domain.grid(0)
self.global_theta_grid = self.sphere_wrapper.grid[None, :]
theta_slice = self.domain.distributor.layouts[-1].slices(1)[1]
self.local_theta_grid = self.global_theta_grid[:, theta_slice]
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)
domain = de.Domain([phi_basis, theta_basis], grid_dtype=np.float64)
# set up sphere
m_start = domain.distributor.coeff_layout.start(1)[0]
m_len = domain.distributor.coeff_layout.local_shape(1)[0]
m_end = m_start + m_len - 1
N_theta = int((L_max + 1) * L_dealias)
S = sph.Sphere(L_max, S_max, N_theta=N_theta, m_min=m_start, m_max=m_end)
phi = domain.grids(L_dealias)[0]
theta_slice = domain.distributor.grid_layout.slices(domain.dealias)[1]
theta_len = domain.local_grid_shape(domain.dealias)[1]
theta_global = S.grid
theta = S.grid[theta_slice].reshape((1, theta_len))
u = sph.TensorField(1, S, domain)
h = sph.TensorField(0, S, domain)
c = sph.TensorField(0, S, domain)
Du = sph.TensorField(2, S, domain)
uh = sph.TensorField(1, S, domain)
Dc = sph.TensorField(1, S, domain)
divuh = sph.TensorField(0, S, domain)
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
self.ell_min, self.ell_max = ell_min, ell_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.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)
# grid and weights for radial integral using volume measure
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)
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)
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)
# downcast to double precision
self.radius = np.sqrt((z_projection + 1) / 2).astype(np.float64)
self.dV = self.dV.astype(np.float64)
self.LU_grad_initialized = []
self.LU_grad = []
self.LU_curl_initialized = []
self.LU_curl = []
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)
if timing:
self.radial_transform_time = 0.
self.angular_transform_time = 0.
self.transpose_time = 0
# work arrays v_thth = domain.new_field() v_thph = domain.new_field() v_phth = domain.new_field() v_phph = domain.new_field() F_th = domain.new_field() F_ph = domain.new_field() # set up sphere p.require_coeff_space() m_start = p.layout.start(1)[0] m_len = p.layout.local_shape(1)[0] m_end = m_start + m_len - 1 S = sph.Sphere(L_max, S_max, m_min=m_start, m_max=m_end) # Calculate theta grid theta_slice = domain.distributor.layouts[-1].slices(1)[1] theta_len = domain.local_grid_shape(1)[1] theta = S.grid[theta_slice].reshape([1, theta_len]) # Grid initial conditions # Right now we specify initial conditions in terms of spectral coefficients # If you want to specify initial conditions in terms of theta & phi # do that here. # Move initial conditions into coefficient space # move data into (m,theta):
