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_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_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_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_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 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_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 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 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)
def test_fourier_integrate(N, bounds, dtype):
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)
x = b.local_grid(1)
f = field.Field(dist=d, bases=(b, ), dtype=dtype)
k = 4 * np.pi / (bounds[1] - bounds[0])
f['g'] = 1 + np.sin(k * x + 0.1)
fi = operators.Integrate(f, c).evaluate()
assert np.allclose(fi['g'], bounds[1] - bounds[0])
def test_RF_scalar_roundtrip(N, dealias):
if comm.size == 1:
c = coords.Coordinate('x')
d = distributor.Distributor([c])
xb = basis.RealFourier(c, size=N, bounds=(0, 1), dealias=dealias)
x = xb.local_grid(dealias)
# Scalar transforms
u = field.Field(dist=d, bases=(xb, ), dtype=np.float64)
u.preset_scales(dealias)
u['g'] = ug = np.cos(2 * np.pi * x + np.pi / 4)
u['c']
assert np.allclose(u['g'], ug)
else:
pytest.skip("Can only test 1D transform in serial")
def test_jacobi_convert_implicit(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()
g = field.Field(dist=d, bases=(b,), dtype=dtype)
problem = problems.LBVP([g])
problem.add_equation((g, fx))
solver = solvers.LinearBoundaryValueSolver(problem)
solver.solve()
assert np.allclose(g['g'], fx['g'])
def test_J_scalar_roundtrip(a, b, N, dealias, dtype):
if comm.size == 1:
c = coords.Coordinate('x')
d = distributor.Distributor([c])
xb = basis.Jacobi(c, a=a, b=b, size=N, bounds=(0, 1), dealias=dealias)
x = xb.local_grid(dealias)
# Scalar transforms
u = field.Field(dist=d, bases=(xb, ), dtype=dtype)
u.preset_scales(dealias)
u['g'] = ug = 2 * x**2 - 1
u['c']
assert np.allclose(u['g'], ug)
else:
pytest.skip("Can only test 1D transform in serial")
def test_jacobi_ncc_eval(N, a0, b0, k_ncc, k_arg, dealias, dtype):
"""
This tests for aliasing errors when evaluating the product
f(x) * g(x)
as both an NCC operator and with the pseudospectral method.
With 3/2 dealiasing, the 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)
f['g'] = np.random.randn(*f['g'].shape)
g['g'] = np.random.randn(*g['g'].shape)
x = b.local_grid(1)
f['g'] = np.cos(6*x)
g['g'] = np.cos(6*x)
kmax = max(k_ncc, k_arg)
if kmax > 0 and dealias < 2:
f['c'][-kmax:] = 0
g['c'][-kmax:] = 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()
w2 = field.Field(dist=d, bases=(b,), dtype=dtype)
w2.preset_scales(w2.domain.dealias)
w2['g'] = w0['g']
w2['c']
w0['c']
print(np.max(np.abs(w2['g'] - w0['g'])))
w2 = operators.Convert(w2, w0.domain.bases[0]).evaluate()
w1 = w1.evaluate_as_ncc()
assert np.allclose(w0['c'], w2['c'])
def test_fourier_interpolate(N, bounds, dtype):
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)
x = b.local_grid(1)
f = field.Field(dist=d, bases=(b, ), dtype=dtype)
k = 4 * np.pi / (bounds[1] - bounds[0])
f['g'] = 1 + np.sin(k * x + 0.1)
results = []
for p in [
bounds[0], bounds[1],
bounds[0] + (bounds[1] - bounds[0]) * np.random.rand()
]:
fp = operators.Interpolate(f, c, p).evaluate()
results.append(np.allclose(fp['g'], 1 + np.sin(k * p + 0.1)))
assert all(results)
def test_heat_1d_periodic(x_basis_class, Nx, dtype):
# Bases
c = coords.Coordinate('x')
d = distributor.Distributor((c, ))
xb = x_basis_class(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)
a = field.Field(name='a', dist=d, dtype=dtype)
# Problem
dx = lambda A: operators.Differentiate(A, c)
problem = problems.EVP([u], a)
problem.add_equation((a * u + dx(dx(u)), 0))
# Solver
solver = solvers.EigenvalueSolver(problem, matrix_coupling=[True])
solver.solve_dense(solver.subproblems[0])
# Check solution
k = xb.wavenumbers
if x_basis_class is basis.RealFourier:
k = k[1:] # Drop one k=0 for msin
assert np.allclose(solver.eigenvalues, k**2)
def test_real_fourier_libraries_forward(N, dealias, dtype, library):
"""Tests that fast real Fourier transforms match matrix transforms."""
c = coords.Coordinate('x')
d = distributor.Distributor([c])
# Matrix
b_mat = basis.RealFourier(c,
size=N,
bounds=(0, 2 * np.pi),
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.RealFourier(c,
size=N,
bounds=(0, 2 * np.pi),
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'])
dtype = np.float64
# Parameters
radius = 1
Mmax = 31
Nmax = 31
nz = 32
Ampl = 1e-6
Re = 3162.
nu = 1./Re
zmin = 0
zmax = 16*np.pi*2*radius
# Bases
zc = coords.Coordinate('z')
c = coords.PolarCoordinates('phi', 'r')
d = distributor.Distributor((zc,c))
b = basis.DiskBasis(c, (Mmax+1, Nmax+1), radius=radius, dtype=dtype)
zb = basis.RealFourier(zc, nz, bounds=(zmin, zmax))
cb = b.S1_basis(radius=radius)
phi, r = b.local_grids((1, 1))
z = zb.local_grids((1,))
# Fields
u = field.Field(dist=d, bases=(zb, b), tensorsig=(c,),dtype=dtype)
w = field.Field(dist=d, bases=(zb, b), dtype=dtype)
p = field.Field(dist=d, bases=(zb, b), dtype=dtype)
tau_u = field.Field(dist=d, bases=(cb,), tensorsig=(c,), dtype=dtype)
tau_w = field.Field(dist=d, bases=(cb,), dtype=dtype)
def main(N, alpha, method, tau):
# Bases
c = coords.Coordinate('x')
d = distributor.Distributor((c,))
xb = basis.ChebyshevT(c, size=N, bounds=(-1, 1))
x = xb.local_grid(1)
# Fields
u = field.Field(name='u', dist=d, bases=(xb,), dtype=np.float64)
ux = field.Field(name='ux', dist=d, bases=(xb,), dtype=np.float64)
f = field.Field(name='f', dist=d, bases=(xb,), dtype=np.float64)
t1 = field.Field(name='t1', dist=d, dtype=np.float64)
t2 = field.Field(name='t2', dist=d, dtype=np.float64)
pi, sin, cos = np.pi, np.sin, np.cos
f['g'] = 64*pi**3*(4*pi*sin(4*pi*x)**2*sin(4*pi*cos(4*pi*x)) + cos(4*pi*x)*cos(4*pi*cos(4*pi*x))) + alpha*sin(4*pi*cos(4*pi*x))
# Tau polynomials
xb1 = xb._new_a_b(xb.a+1, xb.b+1)
xb2 = xb._new_a_b(xb.a+2, xb.b+2)
if tau == 0:
# First-order classical Chebyshev tau: T[-1], dx(T[-1])
p1 = field.Field(name='p1', dist=d, bases=(xb,), dtype=np.float64)
p2 = field.Field(name='p2', dist=d, bases=(xb,), dtype=np.float64)
p1['c'][-1] = 1
p2['c'][-1] = 1
elif tau == 1:
# First-order ultraspherical tau: U[-1], dx(U[-1])
p1 = field.Field(name='p1', dist=d, bases=(xb1,), dtype=np.float64)
p2 = field.Field(name='p2', dist=d, bases=(xb1,), dtype=np.float64)
p1['c'][-1] = 1
p2['c'][-1] = 1
# Problem
dx = lambda A: operators.Differentiate(A, c)
problem = problems.LBVP([u, ux, t1, t2])
problem.add_equation((alpha*u - dx(ux) + t1*p1, f))
problem.add_equation((ux - dx(u) + t2*p2, 0))
problem.add_equation((u(x=-1), 0))
problem.add_equation((u(x=+1), 0))
solver = solvers.LinearBoundaryValueSolver(problem)
# Methods
if method == 0:
# Condition number
L_exp = solver.subproblems[0].L_exp
result = np.linalg.cond(L_exp.A)
elif method == 1:
# Roundtrip roundoff with uniform u
A = solver.subproblems[0].L_exp
v = np.random.rand(A.shape[1])
f = A * v
u = solver.subproblem_matsolvers[solver.subproblems[0]].solve(f)
result = np.max(np.abs(u-v))
elif method == 2:
# Manufactured solution
from mpi4py_fft import fftw as mpi4py_fftw
solver.solve()
ue = np.sin(4*np.pi*np.cos(4*np.pi*x))
d = mpi4py_fftw.aligned(N, fill=0)
k = 2*(1 + np.arange((N-1)//2))
d[::2] = (2./N)/np.hstack((1., 1.-k*k))
w = mpi4py_fftw.aligned_like(d)
dct = mpi4py_fftw.dctn(w, axes=(0,), type=3)
weights = dct(d, w)
result = np.sqrt(np.sum(weights*(u['g']-ue)**2))
elif method == 3:
# Roundtrip roundoff with uniform f
A = solver.subproblems[0].L_exp
g = np.random.rand(A.shape[0])
v = solver.subproblem_matsolvers[solver.subproblems[0]].solve(g)
f = A * v
u = solver.subproblem_matsolvers[solver.subproblems[0]].solve(f)
result = np.max(np.abs(u-v))
return result
