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):