def test_curl_implicit_FFF(basis, N, dealias, dtype):
c, d, b, r = basis(N, dealias, dtype)
Lz = b[2].bounds[1]
integ = lambda A: d3.Integrate(
d3.Integrate(d3.Integrate(A, c.coords[0]), c.coords[1]), c.coords[2])
tau1 = d.VectorField(c, name='tau1')
tau2 = d.Field(name='tau2')
# ABC vector field
k = 2 * np.pi * np.array([1 / Lx, 1 / Ly, 1 / Lz])
f = d.VectorField(c, bases=b)
f.preset_scales(dealias)
f['g'][0] = np.sin(k[2] * r[2]) + np.cos(k[1] * r[1])
f['g'][1] = np.sin(k[0] * r[0]) + np.cos(k[2] * r[2])
f['g'][2] = np.sin(k[1] * r[1]) + np.cos(k[0] * r[0])
g = d.VectorField(c, bases=b)
g.preset_scales(dealias)
g['g'][0] = k[2] * np.sin(k[2] * r[2]) + k[1] * np.cos(k[1] * r[1])
g['g'][1] = k[0] * np.sin(k[0] * r[0]) + k[2] * np.cos(k[2] * r[2])
g['g'][2] = k[1] * np.sin(k[1] * r[1]) + k[0] * np.cos(k[0] * r[0])
# Skew LBVP
u = d.VectorField(c, name='u', bases=b)
phi = d.Field(name='phi', bases=b)
problem = d3.LBVP([u, phi, tau1, tau2], namespace=locals())
problem.add_equation("curl(u) + grad(phi) + tau1 = g")
problem.add_equation("div(u) + tau2 = 0")
problem.add_equation("integ(phi) = 0")
problem.add_equation("integ(comp(u,index=0,coord=c['x'])) = 0")
problem.add_equation("integ(comp(u,index=0,coord=c['y'])) = 0")
problem.add_equation("integ(comp(u,index=0,coord=c['z'])) = 0")
solver = problem.build_solver()
solver.print_subproblem_ranks()
solver.solve()
assert np.allclose(u['c'], f['c'])
def test_poisson_1d_fourier(Nx, dtype, matrix_coupling):
# Bases
coord = d3.Coordinate('x')
dist = d3.Distributor(coord, dtype=dtype)
if dtype == np.complex128:
basis = d3.ComplexFourier(coord, size=Nx, bounds=(0, 2 * np.pi))
elif dtype == np.float64:
basis = d3.RealFourier(coord, size=Nx, bounds=(0, 2 * np.pi))
x = basis.local_grid(1)
# Fields
u = dist.Field(name='u', bases=basis)
g = dist.Field(name='c')
# Substitutions
dx = lambda A: d3.Differentiate(A, coord)
integ = lambda A: d3.Integrate(A, coord)
F = dist.Field(bases=basis)
F['g'] = -np.sin(x)
# Problem
problem = d3.LBVP([u, g], namespace=locals())
problem.add_equation("dx(dx(u)) + g = F")
problem.add_equation("integ(u) = 0")
# Solver
solver = problem.build_solver(matrix_coupling=[matrix_coupling])
solver.solve()
# Check solution
u_true = np.sin(x)
assert np.allclose(u['g'], u_true)
def mcwilliams_preconditioner(domain,
h=0.1,
lx=1 / 4,
ly=1 / 8,
f=1,
US=1 / 16):
# Poisson equation
problem = de.LBVP(domain, variables=['ψ', 'ψz'])
problem.parameters['h'] = h
problem.parameters['lx'] = lx
problem.parameters['ly'] = ly
problem.parameters['f'] = f
problem.parameters['N'] = h / (f * lx)
problem.parameters['US'] = f * lx
problem.substitutions[
'uSy'] = "- y * US / ly**2 * exp(z/h - x**2 / (2 * lx**2) - y**2 / (2 * ly**2))"
problem.substitutions[
'uSz'] = "US / h * exp(z/h - x**2 / (2 * lx**2) - y**2 / (2 * ly**2))"
problem.add_equation("dx(dx(ψ)) + dy(dy(ψ)) + f**2 / N**2 * dz(ψz) = -uSy",
condition="(nx != 0) or (ny != 0)")
problem.add_equation("ψ = 0", condition="(nx == 0) and (ny == 0)")
problem.add_equation("ψz = 0", condition="(nx == 0) and (ny == 0)")
problem.add_equation("dz(ψ) - ψz = 0", condition="(nx != 0) or (ny != 0)")
problem.add_bc("left(ψz) = 0", condition="(nx != 0) or (ny != 0)")
problem.add_bc("right(ψz) = 0", condition="(nx != 0) or (ny != 0)")
solver = problem.build_solver()
return solver
def test_curl_implicit_FFC(basis, N, dealias, dtype):
c, d, b, r = basis(N, dealias, dtype)
Lz = b[2].bounds[1]
lift_basis = b[2].clone_with(a=1 / 2, b=1 / 2) # First derivative basis
lift = lambda A, n: d3.Lift(A, lift_basis, n)
integ = lambda A: d3.Integrate(
d3.Integrate(d3.Integrate(A, c.coords[0]), c.coords[1]), c.coords[2])
tau1 = d.VectorField(c, name='tau1', bases=b[0:2])
tau2 = d.Field(name='tau2', bases=b[0:2])
# ABC vector field
k = 2 * np.pi * np.array([1 / Lx, 1 / Ly, 1 / Lz])
f = d.VectorField(c, bases=b)
f.preset_scales(dealias)
f['g'][0] = np.sin(k[2] * r[2]) + np.cos(k[1] * r[1])
f['g'][1] = np.sin(k[0] * r[0]) + np.cos(k[2] * r[2])
f['g'][2] = np.sin(k[1] * r[1]) + np.cos(k[0] * r[0])
g = d.VectorField(c, bases=b)
g.preset_scales(dealias)
g['g'][0] = k[2] * np.sin(k[2] * r[2]) + k[1] * np.cos(k[1] * r[1])
g['g'][1] = k[0] * np.sin(k[0] * r[0]) + k[2] * np.cos(k[2] * r[2])
g['g'][2] = k[1] * np.sin(k[1] * r[1]) + k[0] * np.cos(k[0] * r[0])
# Skew LBVP
u = d.VectorField(c, name='u', bases=b)
phi = d.Field(name='phi', bases=b)
problem = d3.LBVP([u, phi, tau1, tau2], namespace=locals())
problem.add_equation("curl(u) + grad(phi) + lift(tau1,-1) = g")
problem.add_equation("div(u) + lift(tau2,-1) = 0")
problem.add_equation("u(z=0) = f(z=0)")
problem.add_equation("phi(z=0) = 0")
solver = problem.build_solver()
solver.solve()
assert np.allclose(u['c'], f['c'])
def test_poisson_1d_jacobi(Nx, a0, b0, da, db, dtype):
# Bases
coord = d3.Coordinate('x')
dist = d3.Distributor(coord, dtype=dtype)
basis = d3.Jacobi(coord,
size=Nx,
bounds=(0, 2 * np.pi),
a=a0 + da,
b=b0 + db,
a0=a0,
b0=b0)
x = basis.local_grid(1)
# Fields
u = dist.Field(name='u', bases=basis)
tau1 = dist.Field(name='tau1')
tau2 = dist.Field(name='tau2')
# Substitutions
dx = lambda A: d3.Differentiate(A, coord)
lift_basis = basis.clone_with(a=a0 + da + 2, b=b0 + db + 2)
lift = lambda A, n: d3.Lift(A, lift_basis, n)
F = dist.Field(bases=basis)
F['g'] = -np.sin(x)
# Problem
problem = d3.LBVP([u, tau1, tau2], namespace=locals())
problem.add_equation("dx(dx(u)) + lift(tau1,-1) + lift(tau2,-2) = F")
problem.add_equation("u(x='left') = 0")
problem.add_equation("u(x='right') = 0")
# Solver
solver = problem.build_solver()
solver.solve()
# Check solution
u_true = np.sin(x)
assert np.allclose(u['g'], u_true)
def test_poisson_2d_periodic_firstorder(benchmark, x_basis_class,
y_basis_class, Nx, Ny, dtype):
# Bases and domain
x_basis = x_basis_class('x', Nx, interval=(0, 2 * np.pi))
y_basis = y_basis_class('y', Ny, interval=(0, 2 * np.pi))
domain = de.Domain([x_basis, y_basis], grid_dtype=dtype)
# Forcing
F = domain.new_field(name='F')
F.meta['x']['parity'] = -1
F.meta['y']['parity'] = -1
x, y = domain.grids()
F['g'] = -2 * np.sin(x) * np.sin(y)
# Problem
problem = de.LBVP(domain, variables=['u', 'ux', 'uy'])
problem.meta['u']['x']['parity'] = -1
problem.meta['u']['y']['parity'] = -1
problem.meta['ux']['x']['parity'] = 1
problem.meta['ux']['y']['parity'] = -1
problem.meta['uy']['x']['parity'] = -1
problem.meta['uy']['y']['parity'] = 1
problem.parameters['F'] = F
problem.add_equation("ux - dx(u) = 0")
problem.add_equation("uy - dy(u) = 0")
problem.add_equation("dx(ux) + dy(uy) = F",
condition="(nx != 0) or (ny != 0)")
problem.add_equation("u = 0", condition="(nx == 0) and (ny == 0)")
# Solver
solver = problem.build_solver()
solver.solve()
# Check solution
u_true = np.sin(x) * np.sin(y)
u = solver.state['u']
assert np.allclose(u['g'], u_true)
def linear_tide_2d(param, atmos=None, dtype=np.float64, comm=MPI.COMM_WORLD):
"""Solve linear tide in 2D."""
if atmos is None:
atmos = Atmosphere(param, dim=2, dtype=dtype, comm=comm)
# Adiabatic viscous fully-compressible hydrodynamics
problem = de.LBVP(atmos.domain,
variables=['a1', 'p1', 'u', 'w', 'uz', 'wz'],
ncc_cutoff=param.ivp_cutoff,
entry_cutoff=param.matrix_cutoff)
problem.meta[:]['z']['dirichlet'] = True
problem.parameters['a0'] = atmos.a0
problem.parameters['p0'] = atmos.p0
problem.parameters['a0z'] = atmos.a0z
problem.parameters['p0z'] = atmos.p0z
problem.substitutions['a0x'] = '0'
problem.substitutions['p0x'] = '0'
problem.parameters['U'] = param.U
problem.parameters['μ'] = param.μ
problem.parameters['γ'] = param.γ
problem.parameters['k'] = param.k_tide
problem.parameters['ω'] = param.ω_tide
problem.parameters['σ'] = param.σ_tide
problem.parameters['A'] = param.A_tide
problem.parameters['Lx'] = param.Lx
problem.parameters['Lz'] = param.Lz
problem.substitutions['dt(Q)'] = "0*Q"
problem.substitutions['ux'] = "dx(u)"
problem.substitutions['wx'] = "dx(w)"
problem.substitutions['div_u'] = "ux + wz"
problem.substitutions['txx'] = "μ*(2*ux - 2/3*div_u)"
problem.substitutions['txz'] = "μ*(wx + uz)"
problem.substitutions['tzz'] = "μ*(2*wz - 2/3*div_u)"
problem.substitutions['φ'] = "A*cos(k*x)*exp(k*(z - Lz))"
problem.substitutions['cs20'] = "γ*p0*a0"
problem.add_equation(
"dt(u) + U*ux + a0*dx(p1) + a1*p0x - a0*(dx(txx) + dz(txz)) = - dx(φ)",
condition="nx != 0")
problem.add_equation(
"dt(w) + U*wx + a0*dz(p1) + a1*p0z - a0*(dx(txz) + dz(tzz)) = - dz(φ)",
condition="nx != 0")
problem.add_equation("dt(a1) + U*dx(a1) + u*a0x + w*a0z - a0*div_u = 0",
condition="nx != 0")
problem.add_equation("dt(p1) + U*dx(p1) + u*p0x + w*p0z + γ*p0*div_u = 0",
condition="nx != 0")
problem.add_equation("uz - dz(u) = 0", condition="nx != 0")
problem.add_equation("wz - dz(w) = 0", condition="nx != 0")
problem.add_bc("left(txz/μ) = 0", condition="nx != 0")
problem.add_bc("right(txz/μ) = 0", condition="nx != 0")
problem.add_bc("left(w) = 0", condition="nx != 0")
problem.add_bc("right(w) = 0", condition="nx != 0")
# kx = 0 equations
problem.add_equation("u = 0", condition="nx == 0")
problem.add_equation("w = 0", condition="nx == 0")
problem.add_equation("a1 = 0", condition="nx == 0")
problem.add_equation("p1 = 0", condition="nx == 0")
problem.add_equation("uz = 0", condition="nx == 0")
problem.add_equation("wz = 0", condition="nx == 0")
return atmos, problem
def test_transpose_implicit(basis, N, dealias, dtype):
c, d, b, r = basis(N, dealias, dtype)
# Random tensor field
f = d.TensorField((c, c), bases=b)
f.fill_random(layout='g')
# Transpose LBVP
u = d.TensorField((c, c), bases=b)
problem = d3.LBVP([u], namespace=locals())
problem.add_equation("transpose(u) = transpose(f)")
solver = problem.build_solver()
solver.solve()
assert np.allclose(u['c'], f['c'])
def test_skew_implicit(basis, N, dealias, dtype):
c, d, b, r = basis(N, dealias, dtype)
# Random vector field
f = d.VectorField(c, bases=b)
f.fill_random(layout='g')
# Skew LBVP
u = d.VectorField(c, bases=b)
problem = d3.LBVP([u], namespace=locals())
problem.add_equation("skew(u) = skew(f)")
solver = problem.build_solver()
solver.solve()
assert np.allclose(u['c'][:, 0], f['c'][:, 0])
def linear_tide_1d(param, comm=MPI.COMM_SELF):
"""Solve linear tide with k=k_tide in 1D."""
# Background BVP
dtype = np.complex128
domain, p_full, a_full = background.solve_hydrostatic_pressure(param,
dtype,
comm=comm)
p_full, p_trunc, a_full, a_trunc, heq, N2 = background.truncate_background(
param, p_full, a_full)
# Adiabatic viscous fully-compressible hydrodynamics
problem = de.LBVP(domain,
variables=['a1', 'p1', 'u', 'w', 'uz', 'wz'],
ncc_cutoff=param.ivp_cutoff,
entry_cutoff=param.matrix_cutoff)
problem.meta[:]['z']['dirichlet'] = True
problem.parameters['a0'] = a_trunc
problem.parameters['p0'] = p_trunc
problem.parameters['a0z'] = a_trunc.differentiate('z')
problem.parameters['p0z'] = p_trunc.differentiate('z')
problem.parameters['U'] = param.U
problem.parameters['μ'] = param.μ
problem.parameters['γ'] = param.γ
problem.parameters['k'] = param.k_tide
problem.parameters['ω'] = param.ω_tide
problem.parameters['σ'] = param.σ_tide
problem.parameters['A'] = param.A_tide
problem.parameters['Lx'] = param.Lx
problem.parameters['Lz'] = param.Lz
problem.substitutions['a0x'] = '0'
problem.substitutions['p0x'] = '0'
problem.substitutions['dt(Q)'] = "0*Q"
problem.substitutions['dx(Q)'] = "1j*k*Q"
problem.substitutions['ux'] = "dx(u)"
problem.substitutions['wx'] = "dx(w)"
problem.substitutions['div_u'] = "ux + wz"
problem.substitutions['txx'] = "μ*(2*ux - 2/3*div_u)"
problem.substitutions['txz'] = "μ*(wx + uz)"
problem.substitutions['tzz'] = "μ*(2*wz - 2/3*div_u)"
problem.substitutions['φ'] = "A/2*exp(k*(z - Lz))"
problem.substitutions['cs20'] = "γ*p0*a0"
problem.add_equation(
"dt(u) + U*ux + a0*dx(p1) + a1*p0x - a0*(dx(txx) + dz(txz)) = - dx(φ)")
problem.add_equation(
"dt(w) + U*wx + a0*dz(p1) + a1*p0z - a0*(dx(txz) + dz(tzz)) = - dz(φ)")
problem.add_equation("dt(a1) + U*dx(a1) + u*a0x + w*a0z - a0*div_u = 0")
problem.add_equation("dt(p1) + U*dx(p1) + u*p0x + w*p0z + γ*p0*div_u = 0")
problem.add_equation("uz - dz(u) = 0")
problem.add_equation("wz - dz(w) = 0")
problem.add_bc("left(txz/μ) = 0")
problem.add_bc("right(txz/μ) = 0")
problem.add_bc("left(w) = 0")
problem.add_bc("right(w) = 0")
return domain, problem
def test_trace_implicit(basis, N, dealias, dtype):
c, d, b, r = basis(N, dealias, dtype)
# Random scalar field
f = d.Field(bases=b)
f.fill_random(layout='g')
# Trace LBVP
u = d.Field(bases=b)
I = d.TensorField((c, c))
dim = len(r)
for i in range(dim):
I['g'][i, i] = 1
problem = d3.LBVP([u], namespace=locals())
problem.add_equation("trace(I*u) = dim*f")
solver = problem.build_solver()
solver.solve()
assert np.allclose(u['c'], f['c'])
def test_exponential_free(benchmark, x_basis_class, Nx, dtype):
# Bases and domain
x_basis = x_basis_class('x', Nx, edge=0)
domain = de.Domain([x_basis], grid_dtype=dtype)
# Problem
problem = de.LBVP(domain, variables=['u'])
problem.add_equation("dx(u) + u = 0")
problem.add_bc("left(u) = 1")
# Solver
solver = problem.build_solver()
solver.solve()
# Check solution
x = domain.grid(0)
u_true = np.exp(-x)
u = solver.state['u']
assert np.allclose(u['g'], u_true)
def solve_mcwilliams_preconditioner_problem(
buoyancy_bc="right(ψz) = 0",
h=0.1,
lx=1 / 2,
ly=1 / 4,
f=1,
nx=32,
nz=32,
):
# Create bases and domain
x_basis = de.Fourier('x', nx, interval=(-np.pi, np.pi))
y_basis = de.Fourier('y', nx, interval=(-np.pi, np.pi))
z_basis = de.Chebyshev('z', nz, interval=(-1, 0))
domain = de.Domain([x_basis, y_basis, z_basis], grid_dtype=np.float64)
# Poisson equation
problem = de.LBVP(domain, variables=['ψ', 'ψz'])
problem.parameters['h'] = h
problem.parameters['lx'] = lx
problem.parameters['ly'] = ly
problem.parameters['f'] = f
problem.parameters['N'] = h / (f * lx)
problem.parameters['US'] = f * lx
problem.substitutions[
'uSy'] = "- y * US / ly**2 * exp(z/h - x**2 / (2 * lx**2) - y**2 / (2 * ly**2))"
problem.substitutions[
'uSz'] = "US / h * exp(z/h - x**2 / (2 * lx**2) - y**2 / (2 * ly**2))"
problem.add_equation("dx(dx(ψ)) + dy(dy(ψ)) + f**2 / N**2 * dz(ψz) = -uSy",
condition="(nx != 0) or (ny != 0)")
problem.add_equation("ψ = 0", condition="(nx == 0) and (ny == 0)")
problem.add_equation("ψz = 0", condition="(nx == 0) and (ny == 0)")
problem.add_equation("dz(ψ) - ψz = 0", condition="(nx != 0) or (ny != 0)")
problem.add_bc("left(ψz) = 0", condition="(nx != 0) or (ny != 0)")
problem.add_bc(buoyancy_bc, condition="(nx != 0) or (ny != 0)")
# Build solver
solver = problem.build_solver()
solver.solve()
return solver
def diagonal_solver(Nx, Ny, dtype):
# Bases and domain
x_basis = de.Fourier('x', Nx, interval=(0, 2 * np.pi))
y_basis = de.Fourier('y', Ny, interval=(0, 2 * np.pi))
domain = de.Domain([x_basis, y_basis], grid_dtype=dtype)
# Forcing
F = domain.new_field(name='F')
x, y = domain.grids()
F['g'] = -2 * np.sin(x) * np.sin(y)
# Problem
problem = de.LBVP(domain, variables=['u'])
problem.parameters['F'] = F
problem.add_equation("dx(dx(u)) + dy(dy(u)) = F",
condition="(nx != 0) or (ny != 0)")
problem.add_equation("u = 0", condition="(nx == 0) and (ny == 0)")
# Solver
solver = problem.build_solver()
return solver
def test_gaussian(benchmark, x_basis_class, Nx, dtype):
# Bases and domain
x_basis = x_basis_class('x', Nx, interval=(-1, 1))
domain = de.Domain([x_basis], grid_dtype=dtype)
# Forcing
F = domain.new_field(name='F')
x = domain.grid(0)
F['g'] = -2*x*np.exp(-x**2)
# Problem
problem = de.LBVP(domain, variables=['u'])
problem.parameters['F'] = F
problem.add_equation("dx(u) = F", tau=False)
# Solver
solver = problem.build_solver()
solver.solve()
# Check solution
u_true = np.exp(-x**2)
u = solver.state['u']
assert np.allclose(u['g'], u_true)
def test_double_exponential_forced(benchmark, x_basis_class, Nx, dtype):
# Bases and domain
x_basis = x_basis_class('x', Nx, center=0)
domain = de.Domain([x_basis], grid_dtype=dtype)
# Forcing
F = domain.new_field(name='F')
x = domain.grid(0)
F['g'] = -np.sign(x) * np.exp(-np.sign(x) * x)
# Problem
problem = de.LBVP(domain, variables=['u'])
problem.parameters['F'] = F
problem.add_equation("dx(u) = F", tau=False)
# Solver
solver = problem.build_solver()
solver.solve()
# Check solution
u_true = np.exp(-np.sign(x) * x)
u = solver.state['u']
assert np.allclose(u['g'], u_true)
def test_double_exponential_free(benchmark, x_basis_class, Nx, dtype):
# Bases and domain
x_basis = x_basis_class('x', Nx, center=0)
domain = de.Domain([x_basis], grid_dtype=dtype)
# NCC
x = domain.grid(0)
sign_x = domain.new_field()
sign_x.meta['x']['envelope'] = False
sign_x['g'] = np.sign(x)
# Problem
problem = de.LBVP(domain, variables=['u'])
problem.parameters['sign_x'] = sign_x
problem.add_equation("dx(u) + sign_x*u = 0", tau=True)
problem.add_equation("integ(u) = 2")
# Solver
solver = problem.build_solver()
solver.solve()
# Check solution
u_true = np.exp(-np.sign(x) * x)
u = solver.state['u']
assert np.allclose(u['g'], u_true)
def coupled_solver(Nx, Ny, dtype):
# Bases and domain
x_basis = de.Fourier('x', Nx, interval=(0, 2 * np.pi))
y_basis = de.Chebyshev('y', Ny, interval=(0, 2 * np.pi))
domain = de.Domain([x_basis, y_basis], grid_dtype=dtype)
# Forcing
F = domain.new_field(name='F')
x, y = domain.grids()
F['g'] = -2 * np.sin(x) * np.sin(y)
# Problem
problem = de.LBVP(domain, variables=['u', 'ux', 'uy', 'Lu'])
problem.parameters['F'] = F
problem.add_equation("ux - dx(u) = 0")
problem.add_equation("uy - dy(u) = 0")
problem.add_equation("Lu - dx(ux) - dy(uy) = 0")
problem.add_equation("Lu = F")
problem.add_bc("left(u) - right(u) = 0")
problem.add_bc("left(uy) - right(uy) = 0", condition="nx != 0")
problem.add_bc("left(u) = 0", condition="nx == 0")
# Solver
solver = problem.build_solver()
return solver
def test_gaussian_free(benchmark, x_basis_class, Nx, dtype):
# Stretch Laguerres
if x_basis_class is de.Hermite:
stretch = 1.0
else:
stretch = 0.1
# Bases and domain
x_basis = x_basis_class('x', Nx, center=0, stretch=stretch)
domain = de.Domain([x_basis], grid_dtype=dtype)
# Problem
problem = de.LBVP(domain, variables=['u'])
problem.parameters['pi'] = np.pi
problem.add_equation("dx(u) + 2*x*u = 0", tau=True)
problem.add_bc("integ(u) = sqrt(pi)")
# Solver
solver = problem.build_solver()
solver.solve()
# Check solution
x = domain.grid(0)
u_true = np.exp(-x**2)
u = solver.state['u']
assert np.allclose(u['g'], u_true)
def test_gaussian_forced(benchmark, x_basis_class, Nx, dtype):
# Stretch Laguerres
if x_basis_class is de.Hermite:
stretch = 1.0
else:
stretch = 0.1
# Bases and domain
x_basis = x_basis_class('x', Nx, center=0, stretch=stretch)
domain = de.Domain([x_basis], grid_dtype=dtype)
# Forcing
F = domain.new_field(name='F')
x = domain.grid(0)
F['g'] = -2 * x * np.exp(-x**2)
# Problem
problem = de.LBVP(domain, variables=['u'])
problem.parameters['F'] = F
problem.add_equation("dx(u) = F", tau=False)
# Solver
solver = problem.build_solver()
solver.solve()
# Check solution
u_true = np.exp(-x**2)
u = solver.state['u']
assert np.allclose(u['g'], u_true)
def make_solver(domain, h=1, lx=2, ly=1, f=1, N=20, US=1):
problem = de.LBVP(domain, variables=['ψ', 'ψz'])
problem.parameters["h"] = h
problem.parameters["lx"] = lx
problem.parameters["ly"] = ly
problem.parameters["f"] = f
problem.parameters["N"] = N
problem.parameters["US"] = US
problem.substitutions["uSy"] = "- y / ly**2 * US * exp(z/h - x**2 / (2 * lx**2) - y**2 / (2 * ly**2))"
problem.add_equation("dx(dx(ψ)) + dy(dy(ψ)) + f**2 / N**2 * dz(ψz) = - uSy",
condition="(nx != 0) or (ny != 0)")
problem.add_equation("ψ = 0", condition="(nx == 0) and (ny == 0)")
problem.add_equation("ψz = 0", condition="(nx == 0) and (ny == 0)")
problem.add_equation("dz(ψ) - ψz = 0", condition="(nx != 0) or (ny != 0)")
problem.add_bc("left(ψz) = 0", condition="(nx != 0) or (ny != 0)")
problem.add_bc("right(ψz) = 0", condition="(nx != 0) or (ny != 0)")
# Build solver
return problem.build_solver()
# file name that results are saved in
directory = '/home/jacob/dedalus/UpdatedMixingStability/'
# build domain
z_basis = de.Chebyshev('z', nz, interval=(0, H))
domain = de.Domain([z_basis], np.complex128, comm=MPI.COMM_SELF)
# non-constant coefficients
kap = domain.new_field(name='kap')
z = domain.grid(0)
kap['g'] = kap_0 + kap_1*np.exp(-z/h)
# STEADY STATE
# setup problem
problem = de.LBVP(domain, variables=['U', 'V', 'B', 'Uz', 'Vz', 'Bz'])
problem.parameters['N'] = N
problem.parameters['f'] = f
problem.parameters['tht'] = tht
problem.parameters['kap'] = kap
problem.parameters['Pr'] = Pr
problem.add_equation(('-f*V*cos(tht) - B*sin(tht) - Pr*(dz(kap)*Uz'
'+ kap*dz(Uz)) = 0'))
problem.add_equation('f*U*cos(tht) - Pr*(dz(kap)*Vz + kap*dz(Vz)) = 0')
problem.add_equation(('U*N**2*sin(tht) - dz(kap)*Bz - kap*dz(Bz)'
'= dz(kap)*N**2*cos(tht)'))
problem.add_equation('Uz - dz(U) = 0')
problem.add_equation('Vz - dz(V) = 0')
problem.add_equation('Bz - dz(B) = 0')
problem.add_bc('left(U) = 0')
problem.add_bc('left(V) = 0')
cu_ref.meta['x']['constant'] = True
cu_ref.meta['y0']['parity'] = 1
cu_ref.meta['y1']['parity'] = 1
x, y0, y1 = domain.grids()
# Build as 1D function of y0
k_cu = param.cu_k
k_cu_parity = (param.cu_k + 1) # opposite parity wrt centerline
cu_ref['g'] = param.cu_ampl * (
param.cu_lambda * np.cos(k_cu_parity * y0 * np.pi / param.Ly) +
(1 - param.cu_lambda) * np.cos(k_cu * y0 * np.pi / param.Ly))
# Diagonalize
cu_ref = Diag(cu_ref, 'y0', 'y1').evaluate()
ic_problem = de.LBVP(domain, variables=['cs'])
ic_problem.meta['cs']['x']['constant'] = True
ic_problem.meta['cs']['y0']['parity'] = 1
ic_problem.meta['cs']['y1']['parity'] = -1
ic_problem.parameters['cu_ref'] = cu_ref
ic_problem.add_equation("cs = 0", condition="(nx != 0) or (ny0 != 0)")
ic_problem.add_equation(
"cs = 0", condition="(nx == 0) and (ny1 == 0) and (ny0 == 0)")
ic_problem.add_equation(
"-dy1(cs) = cu_ref",
condition="(nx == 0) and (ny0 == 0) and (ny1 != 0)")
ic_solver = ic_problem.build_solver()
ic_solver.solve()
cs['c'] = ic_solver.state['cs']['c']
h = 0.1
lx = 1 / 2
ly = 1 / 4
f = 1
nx = 32
nz = 32
# Create bases and domain
x_basis = de.Fourier('x', nx, interval=(-np.pi, np.pi))
y_basis = de.Fourier('y', nx, interval=(-np.pi, np.pi))
z_basis = de.Chebyshev('z', nz, interval=(-1, 0))
domain = de.Domain([x_basis, y_basis, z_basis], grid_dtype=np.float64)
# Poisson equation
problem = de.LBVP(domain, variables=['ψ', 'ψz'])
problem.parameters["h"] = h
problem.parameters["lx"] = lx
problem.parameters["ly"] = ly
problem.parameters["f"] = f
problem.parameters["N"] = h / (f * lx)
problem.parameters["US"] = f * lx
problem.substitutions[
"uS"] = "US * exp(z/h - x**2 / (2 * lx**2) - y**2 / (2 * ly**2))"
#problem.substitutions["uSz"] = "dz(uS)"
problem.substitutions[
"uSy"] = "- y * US / ly**2 * exp(z/h - x**2 / (2 * lx**2) - y**2 / (2 * ly**2))"
problem.add_equation("dx(dx(ψ)) + dy(dy(ψ)) + f**2 / N**2 * dz(ψz) = - uSy",
max_iter = 10
newton_tolerance = 1e-16
# Bases and domain
x_basis = de.Fourier('x', Nx, interval=(0, Lx), dealias=3 / 2)
y_basis = de.Fourier('y', Ny, interval=(0, Ly), dealias=3 / 2)
z_basis = de.Chebyshev('z', Nz, interval=(-Lz / 2, Lz / 2), dealias=3 / 2)
domain = de.Domain([x_basis, y_basis, z_basis], grid_dtype=np.float64)
# Fields
field_names = ['p', 'u', 'v', 'w', 'uz', 'vz', 'wz']
X0 = dev.system.FieldSystem(field_names, domain)
X1 = dev.system.FieldSystem(field_names, domain)
# Define Y1 problem
problem = de.LBVP(domain, variables=field_names)
problem.parameters['Re'] = Re
for name in field_names:
problem.parameters[name + '0'] = X0[name]
problem.parameters[name + '1'] = X1[name]
problem.substitutions['U'] = "z"
problem.substitutions['L(a,az)'] = "dx(dx(a)) + dy(dy(a)) + dz(az)"
problem.add_equation("dx(u) + dy(v) + wz = 0")
problem.add_equation(
"(1/Re)*L(u,uz) - dx(p) - U*dx(u) - w*dz(U) = -(dx(u1*u0) + dx(u1*u0) + dy(v1*u0) + dy(u1*v0) + dz(w1*u0) + dz(u1*w0))"
)
problem.add_equation(
"(1/Re)*L(v,vz) - dy(p) - U*dx(v) = -(dx(u1*v0) + dx(v1*u0) + dy(v1*v0) + dy(v1*v0) + dz(w1*v0) + dz(v1*w0))"
)
problem.add_equation(
"(1/Re)*L(w,wz) - dz(p) - U*dx(w) = -(dx(u1*w0) + dx(w1*u0) + dy(v1*w0) + dy(w1*v0) + dz(w1*w0) + dz(w1*w0))"
problem.add_equation("dt(zeta) + beta*v - HD(zeta, 8) = J(zeta, psi) ",
condition="(nx != 0) or (ny != 0)")
problem.add_equation("psi = 0", condition="(nx == 0) and (ny == 0)")
#problem.add_equation("zeta = L(psi)", condition="(nx != 0) or (ny != 0)")
solver = problem.build_solver(de.timesteppers.SBDF3)
solver.stop_sim_time = .1
solver.stop_wall_time = np.inf
solver.stop_iteration = np.inf
# vorticity & velocity are no longer states of the system. They are true diagnostic variables.
# But you still might want to set initial condiitons based on vorticity (for example).
# To do this you'll have to solve for the streamfunction.
# This will solve for an inital psi, given a vorticity field.
init = de.LBVP(domain, variables=['init_psi'])
gshape = domain.dist.grid_layout.global_shape(scales=1)
slices = domain.dist.grid_layout.slices(scales=1)
rand = np.random.RandomState(seed=42)
noise = rand.standard_normal(gshape)[slices]
init_vorticity = domain.new_field()
init_vorticity.set_scales(1)
k = domain.bases[0].wavenumbers[:, np.newaxis]
l = domain.bases[1].wavenumbers[np.newaxis, :]
ksq = k**2 + l**2
ck = np.zeros_like(ksq)
ck = np.sqrt(ksq + (1.0 + (ksq / 36.0)**2))**-1
x, y = dist.local_grids(xbasis, ybasis)
f = dist.Field(bases=(xbasis, ybasis))
g = dist.Field(bases=xbasis)
h = dist.Field(bases=xbasis)
f.fill_random('g', seed=40)
f.low_pass_filter(shape=(64, 32))
g['g'] = np.sin(8 * x) * 0.025
h['g'] = 0
# Substitutions
dy = lambda A: d3.Differentiate(A, coords['y'])
lift_basis = ybasis.derivative_basis(2)
lift = lambda A, n: d3.Lift(A, lift_basis, n)
# Problem
problem = d3.LBVP([u, tau_1, tau_2], namespace=locals())
problem.add_equation("lap(u) + lift(tau_1,-1) + lift(tau_2,-2) = f")
problem.add_equation("u(y=0) = g")
problem.add_equation("dy(u)(y=Ly) = h")
# Solver
solver = problem.build_solver()
solver.solve()
# Gather global data
x = xbasis.global_grid()
y = ybasis.global_grid()
ug = u.allgather_data('g')
# Plot
if dist.comm.rank == 0:
def spheromak_A(domain, center=(0, 0, 0), B0=1, R=1, L=1):
"""Solve
Laplacian(A) = - J0
J0 = S(r) l_sph [ -pi J1(a r) cos(pi z) rhat + l_sph*J1(a r)*sin(pi z)
"""
j1_zero1 = jn_zeros(1, 1)[0]
kr = j1_zero1 / R
kz = np.pi / L
lam = np.sqrt(kr**2 + kz**2)
J0 = B0 / lam
problem = de.LBVP(domain, variables=['Ax', 'Ay', 'Az'])
problem.meta['Ax']['y', 'z']['parity'] = -1
problem.meta['Ax']['x']['parity'] = 1
problem.meta['Ay']['x', 'z']['parity'] = -1
problem.meta['Ay']['y']['parity'] = 1
problem.meta['Az']['x', 'y']['parity'] = -1
problem.meta['Az']['z']['parity'] = 1
J0_x = domain.new_field()
J0_y = domain.new_field()
J0_z = domain.new_field()
xx, yy, zz = domain.grids()
r = np.sqrt((xx - center[0])**2 + (yy - center[1])**2)
theta = np.arctan2(yy, xx)
z = zz - center[2]
S = getS(r, z, L, R, center[2])
J_r = S * lam * (-np.pi * j1(kr * r) * np.cos(z * kz))
J_t = S * lam * (lam * j1(kr * r) * np.sin(z * kz))
J0_x['g'] = J0 * (J_r * np.cos(theta) - J_t * np.sin(theta))
J0_y['g'] = J0 * (J_r * np.sin(theta) + J_t * np.cos(theta))
J0_z['g'] = J0 * S * lam * (kr * j0(kr * r) * np.sin(z * kz))
J0_x.meta['y', 'z']['parity'] = -1
J0_x.meta['x']['parity'] = 1
J0_y.meta['x', 'z']['parity'] = -1
J0_y.meta['y']['parity'] = 1
J0_z.meta['x', 'y']['parity'] = -1
J0_z.meta['z']['parity'] = 1
problem.parameters['J0_x'] = J0_x
problem.parameters['J0_y'] = J0_y
problem.parameters['J0_z'] = J0_z
problem.add_equation("dx(dx(Ax)) + dy(dy(Ax)) + dz(dz(Ax)) = J0_x",
condition="(nx != 0) or (ny != 0) or (nz != 0)")
problem.add_equation("Ax = 0",
condition="(nx == 0) and (ny == 0) and (nz == 0)")
problem.add_equation("dx(dx(Ay)) + dy(dy(Ay)) + dz(dz(Ay)) = J0_y",
condition="(nx != 0) or (ny != 0) or (nz != 0)")
problem.add_equation("Ay = 0",
condition="(nx == 0) and (ny == 0) and (nz == 0)")
problem.add_equation("dx(dx(Az)) + dy(dy(Az)) + dz(dz(Az)) = J0_z",
condition="(nx != 0) or (ny != 0) or (nz != 0)")
problem.add_equation("Az = 0",
condition="(nx == 0) and (ny == 0) and (nz == 0)")
# Build solver
solver = problem.build_solver()
solver.solve()
return solver.state['Ax']['g'], solver.state['Ay']['g'], solver.state[
'Az']['g']
zcross = lambda A: d3.MulCosine(d3.skew(A))
# Initial conditions: zonal jet
phi, theta = dist.local_grids(basis)
lat = np.pi / 2 - theta + 0 * phi
umax = 80 * meter / second
lat0 = np.pi / 7
lat1 = np.pi / 2 - lat0
en = np.exp(-4 / (lat1 - lat0)**2)
jet = (lat0 <= lat) * (lat <= lat1)
u_jet = umax / en * np.exp(1 / (lat[jet] - lat0) / (lat[jet] - lat1))
u['g'][0][jet] = u_jet
# Initial conditions: balanced height
c = dist.Field(name='c')
problem = d3.LBVP([h, c], namespace=locals())
problem.add_equation("g*lap(h) + c = - div(u@grad(u) + 2*Omega*zcross(u))")
problem.add_equation("ave(h) = 0")
solver = problem.build_solver()
solver.solve()
# Initial conditions: perturbation
lat2 = np.pi / 4
hpert = 120 * meter
alpha = 1 / 3
beta = 1 / 15
h['g'] += hpert * np.cos(lat) * np.exp(-(phi / alpha)**2) * np.exp(-(
(lat2 - lat) / beta)**2)
# Problem
problem = d3.IVP([u, h], namespace=locals())
P_init, B_init, Q_init = domain.new_fields(3)
for f in [P_init, B_init, Q_init]:
f.set_scales(domain.dealias)
# Filtered noise.
rand = np.random.RandomState(seed=42)
gshape = domain.dist.grid_layout.global_shape(scales=domain.dealias)[:2]
gslices = domain.dist.grid_layout.slices(scales=domain.dealias)[:2]
noise1 = rand.standard_normal(gshape)[gslices][:, :, None]
noise2 = rand.standard_normal(gshape)[gslices][:, :, None]
P_init['g'] = 30 * (noise1 * np.cos(kz * z) + noise2)
P_init.set_scales(1 / 16)
P_init.require_grid_space()
# Need to solve for W_init using dt(P) --> Pt_init as a slack variable.
init_problem = de.LBVP(domain, variables=['Pt', 'W'])
init_problem.meta[:]['z']['dirichlet'] = True
init_problem.substitutions = problem.substitutions
init_problem.parameters = problem.parameters
init_problem.parameters['P'] = P_init
init_problem.parameters['Q'] = Q_init
init_problem.add_equation(
" L(Pt) - dz(W) = -U*dx(zeta) - beta*v - nu*HD(zeta) - D(zeta)")
init_problem.add_equation(
"dz(Pt) + Gamma*W = -U*dx(B) + dz(U)*v - kappa*HD(B) - D(B)")
init_problem.add_bc(" left(W) = left(r*zeta)")
init_problem.add_bc("right(W) = 0", condition="(nx != 0) or (ny != 0)")