def test_jacobi_ncc_solve(N, a0, b0, k_ncc, k_arg, dealias, dtype):
"""
This tests for aliasing errors when solving the equation
f(x)*u(x) = f(x)*g(x).
With 3/2 dealiasing, the RHS product will generally contain aliasing
errors in the last 2*max(k_ncc, k_arg) modes. We can eliminate these
by zeroing the corresponding number of modes of f(x) and/or g(x).
"""
c = coords.Coordinate('x')
d = distributor.Distributor((c,))
b = basis.Jacobi(c, size=N, a=a0, b=b0, bounds=(0, 1), dealias=dealias)
b_ncc = b.clone_with(a=a0+k_ncc, b=b0+k_ncc)
b_arg = b.clone_with(a=a0+k_arg, b=b0+k_arg)
f = field.Field(dist=d, bases=(b_ncc,), dtype=dtype)
g = field.Field(dist=d, bases=(b_arg,), dtype=dtype)
u = 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)
kmax = max(k_ncc, k_arg)
if kmax > 0 and dealias < 2:
f['c'][-kmax:] = 0
g['c'][-kmax:] = 0
problem = problems.LBVP([u])
problem.add_equation((f*u, f*g))
solver = solvers.LinearBoundaryValueSolver(problem)
solver.solve()
assert np.allclose(u['c'], g['c'])
def test_clenshaw_vector(N, regtotal_in, k1, k2, ell):
dtype = np.complex128
c = coords.SphericalCoordinates('phi', 'theta', 'r')
d = distributor.Distributor((c,))
b = basis.BallBasis(c, (N, N, N), dtype=dtype, k=k1, radius=1)
basis_in = basis.BallBasis(c, (N, N, N), dtype=dtype, k=k2, radius=1)
phi, theta, r = b.local_grids((1, 1, 1))
x = r * np.sin(theta) * np.cos(phi)
y = r * np.sin(theta) * np.sin(phi)
z = r * np.cos(theta)
ncc = field.Field(dist=d, bases=(b.radial_basis,), tensorsig=(c,), dtype=dtype)
ncc['g'][2] = r*(2*r**2-1)
ncc_basis = ncc.domain.bases[0]
a_ncc = ncc_basis.alpha + ncc_basis.k
regtotal_ncc = 1
b_ncc = 1/2 + regtotal_ncc
n_size = b.n_size(ell)
coeffs_filter = ncc['c'][1,0,0,:n_size]
J = basis_in.operator_matrix('Z', ell, regtotal_in)
I = basis_in.operator_matrix('Id', ell, regtotal_in)
A, B = clenshaw.jacobi_recursion(n_size, a_ncc, b_ncc, J)
f0 = dedalus_sphere.zernike.polynomials(3, 1, a_ncc, regtotal_ncc, 1)[0] * sparse.identity(n_size)
matrix = clenshaw.matrix_clenshaw(coeffs_filter, A, B, f0, cutoff=1e-6)
assert np.allclose(J.todense(), matrix.todense())
def test_heat_ball_nlbvp(Nr, dtype, dealias):
radius = 2
ncc_cutoff = 1e-10
tolerance = 1e-10
# Bases
c = coords.SphericalCoordinates('phi', 'theta', 'r')
d = distributor.Distributor((c,))
b = basis.BallBasis(c, (1, 1, Nr), radius=radius, dtype=dtype, dealias=dealias)
bs = b.S2_basis(radius=radius)
phi, theta, r = b.local_grids((1, 1, 1))
# Fields
u = field.Field(name='u', dist=d, bases=(b,), dtype=dtype)
τ = field.Field(name='τ', dist=d, bases=(bs,), dtype=dtype)
F = field.Field(name='F', dist=d, bases=(b,), dtype=dtype) # todo: make this constant
F['g'] = 6
# Problem
Lap = lambda A: operators.Laplacian(A, c)
Lift = lambda A: operators.Lift(A, b, -1)
problem = problems.NLBVP([u, τ])
problem.add_equation((Lap(u) + Lift(τ), F))
problem.add_equation((u(r=radius), 0))
# Solver
solver = solvers.NonlinearBoundaryValueSolver(problem, ncc_cutoff=ncc_cutoff)
# Initial guess
u['g'] = 1
# Iterations
def error(perts):
return np.sum([np.sum(np.abs(pert['c'])) for pert in perts])
err = np.inf
while err > tolerance:
solver.newton_iteration()
err = error(solver.perturbations)
u_true = r**2 - radius**2
u.change_scales(1)
assert np.allclose(u['g'], u_true)
def test_waves_1d_first_order(x_basis_class, Nx, dtype):
# Bases
c = coords.Coordinate('x')
d = distributor.Distributor((c, ))
xb = x_basis_class(c, size=Nx, bounds=(0, np.pi))
x = xb.local_grid(1)
# Fields
u = field.Field(name='u', dist=d, bases=(xb, ), dtype=dtype)
ux = field.Field(name='ux', dist=d, bases=(xb, ), dtype=dtype)
a = field.Field(name='a', dist=d, dtype=dtype)
τ1 = field.Field(name='τ1', dist=d, dtype=dtype)
τ2 = field.Field(name='τ2', dist=d, dtype=dtype)
# Problem
dx = lambda A: operators.Differentiate(A, c)
xb1 = dx(u).domain.bases[0]
P1 = field.Field(name='P1', dist=d, bases=(xb1, ), dtype=dtype)
P2 = field.Field(name='P2', dist=d, bases=(xb1, ), dtype=dtype)
P1['c'][-1] = 1
P2['c'][-1] = 1
problem = problems.EVP([u, ux, τ1, τ2], a)
problem.add_equation((a * u + dx(ux) + P1 * τ1, 0))
problem.add_equation((ux - dx(u) + P2 * τ2, 0))
problem.add_equation((u(x=0), 0))
problem.add_equation((u(x=np.pi), 0))
# Solver
solver = solvers.EigenvalueSolver(problem, matrix_coupling=[True])
solver.solve_dense(solver.subproblems[0])
i_sort = np.argsort(solver.eigenvalues)
solver.eigenvalues = solver.eigenvalues[i_sort]
solver.eigenvectors = solver.eigenvectors[:, i_sort]
# Check solution
solver.set_state(0, solver.subproblems[0].subsystems[0])
eigenfunction = u['g'] / np.max(u['g'])
sol = np.sin(x) / np.max(np.sin(x))
assert np.allclose(eigenfunction, sol)
def test_box_AdvectiveCFL(x_basis_class, Nx, Nz, timestepper, dtype,
z_velocity_mag, dealias):
# Bases
Lx = 2
Lz = 1
c = coords.CartesianCoordinates('x', 'z')
d = distributor.Distributor((c, ))
xb = x_basis_class(c.coords[0], size=Nx, bounds=(0, Lx), dealias=dealias)
x = xb.local_grid(1)
zb = basis.ChebyshevT(c.coords[1],
size=Nz,
bounds=(0, Lz),
dealias=dealias)
z = zb.local_grid(1)
b = (xb, zb)
# Fields
u = field.Field(name='u', dist=d, tensorsig=(c, ), bases=b, dtype=dtype)
print(u.domain.get_basis(c))
# Test Fourier CFL
fourier_velocity = lambda x: np.sin(4 * np.pi * x / Lx)
chebyshev_velocity = lambda z: -z_velocity_mag * z
u['g'][0] = fourier_velocity(x)
u['g'][1] = chebyshev_velocity(z)
# AdvectiveCFL initialization
cfl = operators.AdvectiveCFL(u, c)
cfl_freq = cfl.evaluate()['g']
comparison_freq = np.abs(u['g'][0]) / cfl.cfl_spacing(u)[0]
comparison_freq += np.abs(u['g'][1]) / cfl.cfl_spacing(u)[1]
assert np.allclose(cfl_freq, comparison_freq)
def test_heat_1d_periodic(x_basis_class, Nx, timestepper, dtype):
# Bases
c = coords.Coordinate('x')
d = distributor.Distributor((c, ))
xb = basis.ComplexFourier(c, size=Nx, bounds=(0, 2 * np.pi))
x = xb.local_grid(1)
# Fields
u = field.Field(name='u', dist=d, bases=(xb, ), dtype=dtype)
F = field.Field(name='F', dist=d, bases=(xb, ), dtype=dtype)
F['g'] = -np.sin(x)
# Problem
dx = lambda A: operators.Differentiate(A, c)
dt = operators.TimeDerivative
problem = problems.IVP([u])
problem.add_equation((-dt(u) + dx(dx(u)), F))
# Solver
solver = solvers.InitialValueSolver(problem, timestepper)
dt = 1e-5
iter = 10
for i in range(iter):
solver.step(dt)
# Check solution
amp = 1 - np.exp(-solver.sim_time)
u_true = amp * np.sin(x)
assert np.allclose(u['g'], u_true)
def test_sin_nlbvp(Nx, dtype, dealias, basis_class):
ncc_cutoff = 1e-10
tolerance = 1e-10
# Bases
c = coords.Coordinate('x')
d = distributor.Distributor((c,))
xb = basis_class(c, size=Nx, bounds=(0, 1), dealias=dealias)
x = xb.local_grid(1)
# Fields
u = field.Field(name='u', dist=d, bases=(xb,), dtype=dtype)
τ = field.Field(name='τ', dist=d, dtype=dtype)
xb1 = xb.clone_with(a=xb.a+1, b=xb.b+1)
P = field.Field(name='P', dist=d, bases=(xb1,), dtype=dtype)
P['c'][-1] = 1
# Problem
dx = lambda A: operators.Differentiate(A, c)
problem = problems.NLBVP([u, τ])
problem.add_equation((dx(u) + τ*P, np.sqrt(1-u*u)))
problem.add_equation((u(x=0), 0))
# Solver
solver = solvers.NonlinearBoundaryValueSolver(problem, ncc_cutoff=ncc_cutoff)
# Initial guess
u['g'] = x
# Iterations
def error(perts):
return np.sum([np.sum(np.abs(pert['c'])) for pert in perts])
err = np.inf
while err > tolerance:
solver.newton_iteration()
err = error(solver.perturbations)
# Check solution
u_true = np.sin(x)
u.change_scales(1)
assert np.allclose(u['g'], u_true)
def build_shell(Nphi, Ntheta, Nr, dealias, dtype):
c = coords.SphericalCoordinates('phi', 'theta', 'r')
d = distributor.Distributor((c,))
b = basis.ShellBasis(c, (Nphi, Ntheta, Nr), radii=radii_shell, dealias=(dealias, dealias, dealias), dtype=dtype)
phi, theta, r = b.local_grids(b.domain.dealias)
x, y, z = c.cartesian(phi, theta, r)
return c, d, b, phi, theta, r, x, y, z
def test_chebyshev_libraries_forward(N, alpha, dealias, dtype, library):
"""Tests that fast Chebyshev transforms match matrix transforms."""
c = coords.Coordinate('x')
d = distributor.Distributor([c])
# Matrix
b_mat = basis.Ultraspherical(c,
size=N,
alpha0=0,
alpha=alpha,
bounds=(-1, 1),
dealias=dealias,
library='matrix')
u_mat = field.Field(dist=d, bases=(b_mat, ), dtype=dtype)
u_mat.preset_scales(dealias)
u_mat['g'] = np.random.randn(int(np.ceil(dealias * N)))
# Library
b_lib = basis.Ultraspherical(c,
size=N,
alpha0=0,
alpha=alpha,
bounds=(-1, 1),
dealias=dealias,
library=library)
u_lib = field.Field(dist=d, bases=(b_lib, ), dtype=dtype)
u_lib.preset_scales(dealias)
u_lib['g'] = u_mat['g']
assert np.allclose(u_mat['c'], u_lib['c'])
def test_disk_bessel_zeros(Nphi, Nr, m, radius, dtype):
# Bases
c = coords.PolarCoordinates('phi', 'r')
d = distributor.Distributor((c, ))
b = basis.DiskBasis(c, (Nphi, Nr), radius=radius, dtype=dtype)
b_S1 = b.S1_basis()
phi, r = b.local_grids((1, 1))
# Fields
f = field.Field(dist=d, bases=(b, ), dtype=dtype)
τ_f = field.Field(dist=d, bases=(b_S1, ), dtype=dtype)
k2 = field.Field(name='k2', dist=d, dtype=dtype)
# Parameters and operators
lap = lambda A: operators.Laplacian(A, c)
Lift = lambda A: operators.Lift(A, b, -1)
# Bessel equation: k^2*f + lap(f) = 0
problem = problems.EVP([f, τ_f], k2)
problem.add_equation((k2 * f + lap(f) + Lift(τ_f), 0))
problem.add_equation((f(r=radius), 0))
# Solver
solver = solvers.EigenvalueSolver(problem)
print(solver.subproblems[0].group)
for sp in solver.subproblems:
if sp.group[0] == m:
break
else:
raise ValueError("Could not find subproblem with m = %i" % m)
solver.solve_dense(sp)
# Compare eigenvalues
n_compare = 5
selected_eigenvalues = np.sort(solver.eigenvalues)[:n_compare]
analytic_eigenvalues = (spec.jn_zeros(m, n_compare) / radius)**2
assert np.allclose(selected_eigenvalues, analytic_eigenvalues)
def test_waves_1d(x_basis_class, Nx, dtype):
# Bases
c = coords.Coordinate('x')
d = distributor.Distributor((c, ))
xb = x_basis_class(c, size=Nx, bounds=(0, np.pi))
x = xb.local_grid(1)
# Fields
u = field.Field(name='u', dist=d, bases=(xb, ), dtype=dtype)
a = field.Field(name='a', dist=d, dtype=dtype)
τ1 = field.Field(name='τ1', dist=d, dtype=dtype)
τ2 = field.Field(name='τ2', dist=d, dtype=dtype)
# Problem
dx = lambda A: operators.Differentiate(A, c)
xb2 = dx(dx(u)).domain.bases[0]
P1 = field.Field(name='P1', dist=d, bases=(xb2, ), dtype=dtype)
P2 = field.Field(name='P2', dist=d, bases=(xb2, ), dtype=dtype)
P1['c'][-1] = 1
P2['c'][-2] = 1
problem = problems.EVP([u, τ1, τ2], a)
problem.add_equation((a * u + dx(dx(u)) + P1 * τ1 + P2 * τ2, 0))
problem.add_equation((u(x=0), 0))
problem.add_equation((u(x=np.pi), 0))
# Solver
solver = solvers.EigenvalueSolver(problem, matrix_coupling=[True])
solver.solve_dense(solver.subproblems[0])
i_sort = np.argsort(solver.eigenvalues)
sorted_eigenvalues = solver.eigenvalues[i_sort]
# Check solution
Nmodes = Nx // 4
k = np.arange(Nmodes) + 1
assert np.allclose(sorted_eigenvalues[:Nmodes], k**2)
def test_cartesian_output(dtype, dealias, output_scales, output_layout):
Nx = Ny = Nz = 16
Lx = Ly = Lz = 2 * np.pi
# Bases
c = coords.CartesianCoordinates('x', 'y', 'z')
d = distributor.Distributor((c, ))
Fourier = {
np.float64: basis.RealFourier,
np.complex128: basis.ComplexFourier
}[dtype]
xb = Fourier(c.coords[0], size=Nx, bounds=(0, Lx), dealias=dealias)
yb = Fourier(c.coords[1], size=Ny, bounds=(0, Ly), dealias=dealias)
zb = basis.ChebyshevT(c.coords[2],
size=Nz,
bounds=(0, Lz),
dealias=dealias)
x = xb.local_grid(1)
y = yb.local_grid(1)
z = zb.local_grid(1)
# Fields
u = field.Field(name='u', dist=d, bases=(xb, yb, zb), dtype=dtype)
u['g'] = np.sin(x) * np.sin(y) * np.sin(z)
# Problem
dt = operators.TimeDerivative
problem = problems.IVP([u])
problem.add_equation((dt(u) + u, 0))
# Solver
solver = solvers.InitialValueSolver(problem, timesteppers.RK222)
# Output
tasks = [
u,
u(x=0),
u(y=0),
u(z=0),
u(x=0, y=0),
u(x=0, z=0),
u(y=0, z=0),
u(x=0, y=0, z=0)
]
with tempfile.TemporaryDirectory(dir='.') as tempdir:
tempdir = pathlib.Path(tempdir).stem
output = solver.evaluator.add_file_handler(tempdir, iter=1)
for task in tasks:
output.add_task(task,
layout=output_layout,
name=str(task),
scales=output_scales)
solver.evaluator.evaluate_handlers([output])
output.process_virtual_file()
# Check solution
#post.merge_process_files('test_output')
errors = []
with h5py.File(f'{tempdir}/{tempdir}_s1.h5', mode='r') as file:
for task in tasks:
task_saved = file['tasks'][str(task)][-1]
task = task.evaluate()
task.change_scales(output_scales)
errors.append(np.max(np.abs(task[output_layout] - task_saved)))
assert np.allclose(errors, 0)
def test_jacobi_differentiate(N, a, b, k, dtype):
c = coords.Coordinate('x')
d = distributor.Distributor((c,))
b = basis.Jacobi(c, size=N, a0=a, b0=b, a=a+k, b=b+k, bounds=(0, 1))
x = b.local_grid(1)
f = field.Field(dist=d, bases=(b,), dtype=dtype)
f['g'] = x**5
fx = operators.Differentiate(f, c).evaluate()
assert np.allclose(fx['g'], 5*x**4)
def test_jacobi_convert_constant(N, a, b, k, dtype, layout):
c = coords.Coordinate('x')
d = distributor.Distributor((c,))
b = basis.Jacobi(c, size=N, a0=a, b0=b, a=a+k, b=b+k, bounds=(0, 1))
fc = field.Field(dist=d, dtype=dtype)
fc['g'] = 1
fc[layout]
f = operators.convert(fc, (b,)).evaluate()
assert np.allclose(fc['g'], f['g'])
def build_sphere(Nphi, Ntheta, dealias, dtype):
c = coords.S2Coordinates('phi', 'theta')
d = distributor.Distributor(c, dtype=dtype)
b = basis.SphereBasis(c, (Nphi, Ntheta),
radius=radius,
dealias=(dealias, dealias),
dtype=dtype)
phi, theta = b.local_grids(b.domain.dealias)
return c, d, b, phi, theta
def build_sphere_3d(Nphi, Ntheta, radius, dealias, dtype):
c = coords.SphericalCoordinates('phi', 'theta', 'r')
d = distributor.Distributor((c, ))
b = basis.SphereBasis(c, (Nphi, Ntheta),
radius=radius,
dealias=(dealias, dealias),
dtype=dtype)
phi, theta = b.local_grids((dealias, dealias, dealias))
return c, d, b, phi, theta
def test_ball_diffusion(Lmax, Nmax, Leig, radius, bc, dtype):
# Bases
c = coords.SphericalCoordinates('phi', 'theta', 'r')
d = distributor.Distributor((c, ))
b = basis.BallBasis(c, (2 * (Lmax + 1), Lmax + 1, Nmax + 1),
radius=radius,
dtype=dtype)
b_S2 = b.S2_basis()
phi, theta, r = b.local_grids((1, 1, 1))
# Fields
A = field.Field(dist=d, bases=(b, ), tensorsig=(c, ), dtype=dtype)
φ = field.Field(dist=d, bases=(b, ), dtype=dtype)
τ_A = field.Field(dist=d, bases=(b_S2, ), tensorsig=(c, ), dtype=dtype)
λ = field.Field(name='λ', dist=d, dtype=dtype)
# Parameters and operators
div = lambda A: operators.Divergence(A)
grad = lambda A: operators.Gradient(A, c)
curl = lambda A: operators.Curl(A)
lap = lambda A: operators.Laplacian(A, c)
trans = lambda A: operators.TransposeComponents(A)
radial = lambda A, index: operators.RadialComponent(A, index=index)
angular = lambda A, index: operators.AngularComponent(A, index=index)
Lift = lambda A: operators.Lift(A, b, -1)
# Problem
problem = problems.EVP([φ, A, τ_A], λ)
problem.add_equation((div(A), 0))
problem.add_equation((-λ * A + grad(φ) - lap(A) + Lift(τ_A), 0))
if bc == 'no-slip':
problem.add_equation((A(r=radius), 0))
elif bc == 'stress-free':
E = 1 / 2 * (grad(A) + trans(grad(A)))
problem.add_equation((radial(A(r=radius), 0), 0))
problem.add_equation((radial(angular(E(r=radius), 0), 1), 0))
elif bc == 'potential':
ell_func = lambda ell: ell + 1
ell_1 = lambda A: operators.SphericalEllProduct(A, c, ell_func)
problem.add_equation(
(radial(grad(A)(r=radius), 0) + ell_1(A)(r=radius) / radius, 0))
elif bc == 'conducting':
problem.add_equation((φ(r=radius), 0))
problem.add_equation((angular(A(r=radius), 0), 0))
elif bc == 'pseudo':
problem.add_equation((radial(A(r=radius), 0), 0))
problem.add_equation((angular(curl(A)(r=radius), 0), 0))
# Solver
solver = solvers.EigenvalueSolver(problem)
if not solver.subproblems[Leig].group[1] == Leig:
raise ValueError("subproblems indexed in a strange way")
solver.solve_dense(solver.subproblems[Leig])
i_sort = np.argsort(solver.eigenvalues)
solver.eigenvalues = solver.eigenvalues[i_sort]
λ_analytic = analytic_eigenvalues(Leig, Nmax + 1, bc, r0=radius)
if (bc == 'stress-free' and Leig == 1):
# add null space solution
λ_analytic = np.append(0, λ_analytic)
assert np.allclose(solver.eigenvalues[:Nmax // 4], λ_analytic[:Nmax // 4])
def build_disk(Nphi, Nr, dealias, dtype):
c = coords.PolarCoordinates('phi', 'r')
d = distributor.Distributor((c, ))
b = basis.DiskBasis(c, (Nphi, Nr),
radius=radius_disk,
dealias=(dealias, dealias),
dtype=dtype)
phi, r = b.local_grids()
x, y = c.cartesian(phi, r)
return c, d, b, phi, r, x, y
def build_ball(Nphi, Ntheta, Nr, dtype, dealias, radius=1):
c = coords.SphericalCoordinates('phi', 'theta', 'r')
d = distributor.Distributor((c, ))
b = basis.BallBasis(c, (Nphi, Ntheta, Nr),
radius=radius,
dtype=dtype,
dealias=(dealias, dealias, dealias))
phi, theta, r = b.local_grids()
x, y, z = c.cartesian(phi, theta, r)
return c, d, b, phi, theta, r, x, y, z
def build_disk(Nphi, Nr, radius, alpha, k, dealias, dtype):
c = coords.PolarCoordinates('phi', 'r')
d = distributor.Distributor((c, ))
b = basis.DiskBasis(c, (Nphi, Nr),
dtype=dtype,
radius=radius,
alpha=alpha,
k=k,
dealias=(dealias, dealias))
return c, d, b
def build_annulus(Nphi, Nr, dealias, dtype):
c = coords.PolarCoordinates('phi', 'r')
d = distributor.Distributor((c, ))
b = basis.AnnulusBasis(c, (Nphi, Nr),
radii=radii_annulus,
dealias=(dealias, dealias),
dtype=dtype)
phi, r = b.local_grids(b.domain.dealias)
x, y = c.cartesian(phi, r)
return c, d, b, phi, r, x, y
def build_annulus(Nphi, Nr, radius, alpha, k, dealias, dtype):
c = coords.PolarCoordinates('phi', 'r')
d = distributor.Distributor((c, ))
alpha = (alpha - 1 / 2, alpha - 1 / 2)
b = basis.AnnulusBasis(c, (Nphi, Nr),
radii=(radius, radius + 1.3),
alpha=alpha,
k=k,
dealias=(dealias, dealias),
dtype=dtype)
return c, d, b
def build_S2(Nphi, Ntheta, dealias, dtype=np.complex128, grid_scale=1):
c = coords.S2Coordinates('phi', 'theta')
d = distributor.Distributor((c, ))
dealias_tuple = (dealias, dealias)
sb = basis.SphereBasis(c, (Nphi, Ntheta),
radius=1,
dealias=dealias_tuple,
dtype=dtype)
grid_scale_tuple = (grid_scale, grid_scale)
phi, theta = sb.local_grids(grid_scale_tuple)
return c, d, sb, phi, theta
def test_jacobi_convert_explicit(N, a, b, k, dtype, layout):
c = coords.Coordinate('x')
d = distributor.Distributor((c,))
b = basis.Jacobi(c, size=N, a0=a, b0=b, a=a+k, b=b+k, bounds=(0, 1))
x = b.local_grid(1)
f = field.Field(dist=d, bases=(b,), dtype=dtype)
f['g'] = x**5
fx = operators.Differentiate(f, c)
f[layout]
g = operators.convert(f, fx.domain.bases).evaluate()
assert np.allclose(g['g'], f['g'])
def test_jacobi_interpolate(N, a, b, k, dtype):
c = coords.Coordinate('x')
d = distributor.Distributor((c,))
b = basis.Jacobi(c, size=N, a0=a, b0=b, a=a+k, b=b+k, bounds=(0, 1))
x = b.local_grid(1)
f = field.Field(dist=d, bases=(b,), dtype=dtype)
f['g'] = x**5
results = []
for p in [0, 1, np.random.rand()]:
fp = operators.Interpolate(f, c, p).evaluate()
results.append(np.allclose(fp['g'], p**5))
assert all(results)
def test_fourier_convert_constant(N, bounds, dtype, layout):
c = coords.Coordinate('x')
d = distributor.Distributor((c, ))
if dtype == np.float64:
b = basis.RealFourier(c, size=N, bounds=bounds)
elif dtype == np.complex128:
b = basis.ComplexFourier(c, size=N, bounds=bounds)
fc = field.Field(dist=d, dtype=dtype)
fc['g'] = 1
fc[layout]
f = operators.convert(fc, (b, )).evaluate()
assert np.allclose(fc['g'], f['g'])
def build_disk(Nphi, Nr, radius, dealias, dtype=np.float64, grid_scale=1):
c = coords.PolarCoordinates('phi', 'r')
d = distributor.Distributor((c, ))
dealias_tuple = (dealias, dealias)
b = basis.DiskBasis(c, (Nphi, Nr),
radius=radius,
dealias=dealias_tuple,
dtype=dtype)
grid_scale_tuple = (grid_scale, grid_scale)
phi, r = b.local_grids(grid_scale_tuple)
x, y = c.cartesian(phi, r)
return c, d, b, phi, r, x, y
def build_shell(Nphi, Ntheta, Nr, radii_shell, dealias, dtype, grid_scale=1):
c = coords.SphericalCoordinates('phi', 'theta', 'r')
d = distributor.Distributor((c, ))
dealias_tuple = (dealias, dealias, dealias)
b = basis.ShellBasis(c, (Nphi, Ntheta, Nr),
radii=radii_shell,
dealias=dealias_tuple,
dtype=dtype)
grid_scale_tuple = (grid_scale, grid_scale, grid_scale)
phi, theta, r = b.local_grids(grid_scale_tuple)
x, y, z = c.cartesian(phi, theta, r)
return c, d, b, phi, theta, r, x, y, z
def build_2d_box(Nx, Nz, dealias, dtype, k=0):
c = coords.CartesianCoordinates('x', 'z')
d = distributor.Distributor((c,))
if dtype == np.complex128:
xb = basis.ComplexFourier(c.coords[0], size=Nx, bounds=(0, Lx))
elif dtype == np.float64:
xb = basis.RealFourier(c.coords[0], size=Nx, bounds=(0, Lx))
zb = basis.ChebyshevT(c.coords[1], size=Nz, bounds=(0, Lz))
b = (xb, zb)
x = xb.local_grid(1)
z = zb.local_grid(1)
return c, d, b, x, z
def test_jacobi_ufunc(N, a, b, dtype, dealias, func):
c = coords.Coordinate('x')
d = distributor.Distributor((c, ))
b = basis.Jacobi(c, size=N, a=a, b=b, bounds=(0, 1), dealias=dealias)
x = b.local_grid(1)
f = field.Field(dist=d, bases=(b, ), dtype=dtype)
if func is np.arccosh:
f['g'] = 1 + x**2
else:
f['g'] = x**2
g0 = func(f['g'])
g = func(f).evaluate()
assert np.allclose(g['g'], g0)
