def u_evaluation(t):
psi_expr = ((-u_max * L / pi) * sin(2 * pi * x[0] / L) *
sin(pi * x[1] / L)) * sin(2 * pi * t / tmax)
psi0 = Function(Vpsi).interpolate(psi_expr)
return gradperp(psi0)
def sphere_to_cartesian(mesh, u_zonal, u_merid):
theta, lamda = latlon_coords(mesh)
cartesian_u_expr = -u_zonal*sin(lamda) - u_merid*sin(theta)*cos(lamda)
cartesian_v_expr = u_zonal*cos(lamda) - u_merid*sin(theta)*sin(lamda)
cartesian_w_expr = u_merid*cos(theta)
return as_vector((cartesian_u_expr, cartesian_v_expr, cartesian_w_expr))
def solve_something(mesh):
V = fd.FunctionSpace(mesh, "CG", 1)
u = fd.Function(V)
v = fd.TestFunction(V)
x, y = fd.SpatialCoordinate(mesh)
# f = fd.sin(x) * fd.sin(y) + x**2 + y**2
# uex = x**4 * y**4
uex = fd.sin(x) * fd.sin(y) #*(x*y)**3
# def source(xs, ys):
# return 1/((x-xs)**2+(y-ys)**2 + 0.1)
# uex = source(0, 0)
uex = uex - fd.assemble(uex * fd.dx) / fd.assemble(1 * fd.dx(domain=mesh))
# f = fd.conditional(fd.ge(abs(x)-abs(y), 0), 1, 0)
from firedrake import inner, grad, dx, ds, div, sym
eps = fd.Constant(0.0)
f = uex - div(grad(uex)) + eps * div(grad(div(grad(uex))))
n = fd.FacetNormal(mesh)
g = inner(grad(uex), n)
g1 = inner(grad(div(grad(uex))), n)
g2 = div(grad(uex))
# F = 0.1 * inner(u, v) * dx + inner(grad(u), grad(v)) * dx + inner(grad(grad(u)), grad(grad(v))) * dx - f * v * dx - g * v * ds
F = inner(u, v) * dx + inner(grad(u),
grad(v)) * dx - f * v * dx - g * v * ds
F += eps * inner(div(grad(u)), div(grad(v))) * dx
F += eps * g1 * v * ds
F -= eps * g2 * inner(grad(v), n) * ds
# f = -div(grad(uex))
# F = inner(grad(u), grad(v)) * dx - f * v * dx
# bc = fd.DirichletBC(V, uex, "on_boundary")
bc = None
fd.solve(F == 0,
u,
bcs=bc,
solver_parameters={
"ksp_type": "cg",
"ksp_atol": 1e-13,
"ksp_rtol": 1e-13,
"ksp_dtol": 1e-13,
"ksp_stol": 1e-13,
"pc_type": "jacobi",
"pc_factor_mat_solver_type": "mumps",
"snes_type": "ksponly",
"ksp_converged_reason": None
})
print("||u-uex|| =", fd.norm(u - uex))
print("||grad(u-uex)|| =", fd.norm(grad(u - uex)))
print("||grad(grad(u-uex))|| =", fd.norm(grad(grad(u - uex))))
err = fd.Function(
fd.TensorFunctionSpace(mesh, "DG",
V.ufl_element().degree() - 2)).interpolate(
grad(grad(u - uex)))
# err = fd.Function(fd.FunctionSpace(mesh, "DG", V.ufl_element().degree())).interpolate(u-uex)
fd.File(outdir + "sln.pvd").write(u)
fd.File(outdir + "err.pvd").write(err)
def f_end(geometry, state):
"""
returns an expression for the expected final state
"""
x = SpatialCoordinate(state.mesh)
if geometry == "sphere":
fexpr = exp(-x[2]**2 - x[0]**2)
if geometry == "slice":
fexpr = sin(2 * pi * (x[0] - 0.5)) * sin(2 * pi * x[1])
return fexpr
def __init__(self):
super().__init__()
X = self.X
# volume and boundary integrals
self.name = "LevelSet Example 2"
self.J = (sin(X[0]) * cos(X[1]) * dx +
pow(1.3 + X[0], 4.2) * pow(1.4 + X[1], 3.3) * dx +
exp(sin(X[0]) + cos(X[1])) * dx +
ln(5 + sin(X[0]) + cos(X[1])) * ds)
self.set_quadrature(8)
def sphere_to_cartesian(mesh, u_zonal, u_merid):
"""Reformulate vector function in theta, lamda directions
as function in x,y,z directions"""
theta, lamda = latlon_coords(mesh)
cartesian_u_expr = -u_zonal*sin(lamda) - u_merid*sin(theta)*cos(lamda)
cartesian_v_expr = u_zonal*cos(lamda) - u_merid*sin(theta)*sin(lamda)
cartesian_w_expr = u_merid*cos(theta)
return as_vector((cartesian_u_expr, cartesian_v_expr, cartesian_w_expr))
def f_init(geometry, state):
"""
returns an expression for the initial condition
"""
x = SpatialCoordinate(state.mesh)
if geometry == "sphere":
fexpr = exp(-x[2]**2 - x[1]**2)
if geometry == "slice":
fexpr = sin(2 * pi * x[0]) * sin(2 * pi * x[1])
return fexpr
def __init__(self):
super().__init__()
X = self.X
# volume and boundary integrals
self.name = "LevelSet Example 3"
V = FunctionSpace(self.mesh, "CG", 1)
u = interpolate(sin(X[0]) * cos(X[1])**2, V)
n = FacetNormal(self.mesh)
self.J = (sin(X[0]) * cos(X[1]) * dx +
pow(1.3 + X[0], 4.2) * pow(1.4 + X[1], 3.3) * dx +
exp(sin(X[0]) + cos(X[1])) * dx +
ln(5 + sin(X[0]) + cos(X[1])) * ds + inner(grad(u), n) * ds)
self.set_quadrature(8)
def test_sharp_cutoff_pre_ufl():
"""Tests that the sharp cutoff function does what it should when the
preconditioning coefficient is given by ufl."""
k = 10.0
mesh = fd.UnitSquareMesh(10,10)
V = fd.FunctionSpace(mesh,"CG",1)
x = fd.SpatialCoordinate(mesh)
n_pre = 1.0 + fd.sin(30*x[0])
prob = hh.HelmholtzProblem(k,V,n_pre=n_pre,A_pre = fd.as_matrix([[1.0,0.0],[0.0,1.0]]))
prob.sharp_cutoff(np.array([0.5,0.5]),0.5,True)
V_DG = fd.FunctionSpace(mesh,"DG",0)
n_fn = fd.Function(V_DG)
n_fn.interpolate(prob._n_pre)
# As above
assert n_fn.dat.data_ro[97] == 1.0
def test_n_min_pre_ufl():
"""Tests that the sharp cutoff function does what it should when n_pre is given by a UFL expression."""
k = 10.0
mesh = fd.UnitSquareMesh(10,10)
V = fd.FunctionSpace(mesh,"CG",1)
x = fd.SpatialCoordinate(mesh)
n_pre = 1.0 + fd.sin(30*x[0])
prob = hh.HelmholtzProblem(k,V,n_pre=n_pre,A_pre = fd.as_matrix([[1.0,0.0],[0.0,1.0]]))
n_min_val = 2.0
prob.n_min(n_min_val,True)
V_DG = fd.FunctionSpace(mesh,"DG",0)
n_fn = fd.Function(V_DG)
n_fn.interpolate(prob._n_pre)
assert (n_fn.dat.data_ro >= n_min_val).all()
def eval(self, value, X):
if X[2] > d:
value[:] = [
A * fd.sin((X[2] - d) / l * fd.pi / 2.)**2, 0., 0.
]
else:
value[:] = [0., 0., 0.]
def test_n_min_ufl():
"""Tests that the sharp cutoff function does what it should when n is given by a ufl expression."""
k = 10.0
mesh = fd.UnitSquareMesh(10,10)
V = fd.FunctionSpace(mesh,"CG",1)
x = fd.SpatialCoordinate(mesh)
n = 1.0 + fd.sin(30*x[0])
prob = hh.HelmholtzProblem(k,V,n=n)
n_min_val = 2.0
prob.n_min(n_min_val)
V_DG = fd.FunctionSpace(mesh,"DG",0)
n_fn = fd.Function(V_DG)
n_fn.interpolate(prob._n)
assert (n_fn.dat.data_ro >= n_min_val).all()
def test_sharp_cutoff_ufl():
"""Tests that the sharp cutoff function does what it should when the
coefficient is given by a ufl expression."""
k = 10.0
mesh = fd.UnitSquareMesh(10,10)
V = fd.FunctionSpace(mesh,"CG",1)
x = fd.SpatialCoordinate(mesh)
n = 1.0 + fd.sin(30*x[0])
prob = hh.HelmholtzProblem(k,V,n=n)
prob.sharp_cutoff(np.array([0.5,0.5]),0.5)
V_DG = fd.FunctionSpace(mesh,"DG",0)
n_fn = fd.Function(V_DG)
n_fn.interpolate(prob._n)
# This is a rudimentary test that it's 1 on the boundary
# Yes, I kind of made this pass by changing the value to check until it did.
# But I've confirmed that it's doing (roughly) the right thing visually, so I'm content
assert n_fn.dat.data_ro[97] == 1.0
def setup_SUPGadvection(direction):
nlayers = 25 # horizontal layers
columns = 25 # number of columns
L = 1.0
m = PeriodicIntervalMesh(columns, L)
H = 1.0 # Height position of the model top
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers)
# Space for initialising velocity
W_VectorCG1 = VectorFunctionSpace(mesh, "CG", 1)
W_CG1 = FunctionSpace(mesh, "CG", 1)
# vertical coordinate and normal
z = Function(W_CG1).interpolate(Expression("x[1]"))
k = Function(W_VectorCG1).interpolate(Expression(("0.","1.")))
fieldlist = ['u','rho', 'theta']
timestepping = TimesteppingParameters(dt=0.01)
parameters = CompressibleParameters()
state = CompressibleState(mesh, vertical_degree=1, horizontal_degree=1,
family="CG",
z=z, k=k,
timestepping=timestepping,
parameters=parameters,
fieldlist=fieldlist)
# interpolate initial conditions
u0 = Function(state.V[0], name="velocity")
uexpr = as_vector([1.0, 0.0])
u0.project(uexpr)
if len(direction) == 0:
space = W_CG1
else:
space = state.V[2]
f = Function(space, name='f')
x = SpatialCoordinate(mesh)
f_expr = sin(2*pi*x[0])*sin(2*pi*x[1])
f.interpolate(f_expr)
f_end = Function(space)
f_end_expr = sin(2*pi*(x[0]-0.5))*sin(2*pi*x[1])
f_end.interpolate(f_end_expr)
return state, u0, f, f_end
def D_integrand(th):
"""Initial D field is calculated by integrating D_integrand w.r.t.
theta. Assumes the input is between theta0 and theta1.
Note that this function operates on vectorized input."""
from scipy import sin, exp, tan
f = 2.0 * Omega * sin(th)
u_zon = (80.0 / e_n) * exp(1.0 / ((th - theta_0) * (th - theta_1)))
return u_zon * (f + tan(th) * u_zon / R)
def b(mesh, fs):
v = TestFunction(fs)
f = Function(fs)
x = SpatialCoordinate(mesh)
f.interpolate((4.*pi*pi)*sin(x[0]*pi*2))
L = f * v * dx
b = assemble(L)
return b
def test_solver_no_flow_region():
mesh = fd.Mesh("./2D_mesh.msh")
no_flow = [2]
no_flow_markers = [1]
mesh = mark_no_flow_regions(mesh, no_flow, no_flow_markers)
P2 = fd.VectorElement("CG", mesh.ufl_cell(), 1)
P1 = fd.FiniteElement("CG", mesh.ufl_cell(), 1)
TH = P2 * P1
W = fd.FunctionSpace(mesh, TH)
(v, q) = fd.TestFunctions(W)
# Stokes 1
w_sol1 = fd.Function(W)
nu = fd.Constant(0.05)
F = NavierStokesBrinkmannForm(W, w_sol1, nu, beta_gls=2.0)
x, y = fd.SpatialCoordinate(mesh)
u_mms = fd.as_vector(
[sin(2.0 * pi * x) * sin(pi * y),
sin(pi * x) * sin(2.0 * pi * y)])
p_mms = -0.5 * (u_mms[0]**2 + u_mms[1]**2)
f_mms_u = (grad(u_mms) * u_mms + grad(p_mms) -
2.0 * nu * div(sym(grad(u_mms))))
f_mms_p = div(u_mms)
F += -inner(f_mms_u, v) * dx - f_mms_p * q * dx
bc1 = fd.DirichletBC(W.sub(0), u_mms, "on_boundary")
bc2 = fd.DirichletBC(W.sub(1), p_mms, "on_boundary")
bc_no_flow = InteriorBC(W.sub(0), fd.Constant((0.0, 0.0)), no_flow_markers)
solver_parameters = {"ksp_max_it": 500, "ksp_monitor": None}
problem1 = fd.NonlinearVariationalProblem(F,
w_sol1,
bcs=[bc1, bc2, bc_no_flow])
solver1 = NavierStokesBrinkmannSolver(
problem1,
options_prefix="navier_stokes",
solver_parameters=solver_parameters,
)
solver1.solve()
u_sol, _ = w_sol1.split()
u_mms_func = fd.interpolate(u_mms, W.sub(0))
error = fd.errornorm(u_sol, u_mms_func)
assert error < 0.07
def __init__(self):
super().__init__()
u, v, X = self.u, self.v, self.X
self.name = "nonlinear PDE constraint"
f = sin(X[1]) * cos(X[0])
g = exp(f)
self.F = (u * v + (1 + u**2) * inner(grad(u), grad(v)) -
f * v) * dx + g * v * ds
self.J = u * dx + pow(1 + u * u, 2.5) * ds
self.set_quadrature(10)
def __init__(self):
super().__init__()
u, v, X = self.u, self.v, self.X
# nonhomogeneous Neumann bc and nonlinear functional on bdry
self.name = "PDE constrained Example 2"
f = sin(X[1]) * cos(X[0])
g = exp(f)
self.F = (u * v + inner(grad(u), grad(v)) - f * v) * dx + g * v * ds
self.J = u * dx + pow(1 + u * u, 2.5) * ds
self.set_quadrature(10)
def run_solver(r):
mesh = fd.UnitSquareMesh(2**r, 2**r)
P2 = fd.VectorElement("CG", mesh.ufl_cell(), 1)
P1 = fd.FiniteElement("CG", mesh.ufl_cell(), 1)
TH = P2 * P1
W = fd.FunctionSpace(mesh, TH)
(v, q) = fd.TestFunctions(W)
# Stokes 1
w_sol1 = fd.Function(W)
nu = fd.Constant(0.05)
F = NavierStokesBrinkmannForm(W, w_sol1, nu, beta_gls=2.0)
from firedrake import sin, grad, pi, sym, div, inner
x, y = fd.SpatialCoordinate(mesh)
u_mms = fd.as_vector(
[sin(2.0 * pi * x) * sin(pi * y),
sin(pi * x) * sin(2.0 * pi * y)])
p_mms = -0.5 * (u_mms[0]**2 + u_mms[1]**2)
f_mms_u = (grad(u_mms) * u_mms + grad(p_mms) -
2.0 * nu * div(sym(grad(u_mms))))
f_mms_p = div(u_mms)
F += -inner(f_mms_u, v) * dx - f_mms_p * q * dx
bc1 = fd.DirichletBC(W.sub(0), u_mms, "on_boundary")
bc2 = fd.DirichletBC(W.sub(1), p_mms, "on_boundary")
solver_parameters = {"ksp_max_it": 200}
problem1 = fd.NonlinearVariationalProblem(F, w_sol1, bcs=[bc1, bc2])
solver1 = NavierStokesBrinkmannSolver(
problem1,
options_prefix="navier_stokes",
solver_parameters=solver_parameters,
)
solver1.solve()
u_sol, _ = w_sol1.split()
fd.File("test_u_sol.pvd").write(u_sol)
u_mms_func = fd.interpolate(u_mms, W.sub(0))
error = fd.errornorm(u_sol, u_mms_func)
print(f"Error: {error}")
return error
def surface_force(self):
A = 1.
d = 0.9 * self.Lz
l = 0.1 * self.Lz
T = fd.Function(self.V)
T.interpolate(
fd.conditional(
self.X[2] > d,
fd.as_vector(
[A * fd.sin((self.X[2] - d) / l * fd.pi / 2.)**2, 0., 0.]),
fd.as_vector([0., 0., 0.]))) # surface force / area
return T
def __init__(self):
super().__init__()
self.name = "PDE-constraint with DirBC"
u, v, X = self.u, self.v, self.X
f = sin(X[1]) * cos(X[0])
g = f / 2.
self.g = g
self.with_DirBC = 1 # special case for DirichletBC in TaylorTest
self.bc = DirichletBC(self.V, 0., "on_boundary")
self.F = (inner(grad(u + g), grad(v)) + (u + g) * v - f * v) * dx
self.J = (u + g) * (u + g) * dx + pow(1 + u * u, 2.5) * ds
self.set_quadrature(10)
def set_initial_value(mesh, T, extent):
"""
Setup T to hold the initial value for the problem.
mesh: Mesh, The mesh to define the function for
T: Function, The function to define the initial value on
extent: list<float>, The length of each side of the mesh (assumes rect)
"""
x = SpatialCoordinate(mesh)
val = 1
for pos, length in zip(x, extent):
val *= sin(2*pos/length*pi)
T.interpolate(100 + 1000**val)
def create_S(mesh, V, extent):
"""
This is the source term.
mesh: Mesh, The mesh to define the function for.
V: FunctionSpace, The function space that the function should be in
extent: list<float>, The length of each side of the mesh (assumes rect)
"""
x = SpatialCoordinate(mesh)
val = 1
for pos, length in zip(x, extent):
val *= sin(pos / length * pi)**2
return Function(V).interpolate(1e9 * val)
def heat(n, deg, time_stages, stage_type="deriv", splitting=IA):
N = 2**n
msh = UnitIntervalMesh(N)
params = {
"snes_type": "ksponly",
"ksp_type": "preonly",
"mat_type": "aij",
"pc_type": "lu"
}
V = FunctionSpace(msh, "CG", deg)
x, = SpatialCoordinate(msh)
t = Constant(0.0)
dt = Constant(2.0 / N)
uexact = exp(-t) * sin(pi * x)
rhs = expand_derivatives(diff(uexact, t)) - div(grad(uexact))
butcher_tableau = GaussLegendre(time_stages)
u = project(uexact, V)
v = TestFunction(V)
F = (inner(Dt(u), v) * dx + inner(grad(u), grad(v)) * dx -
inner(rhs, v) * dx)
bc = DirichletBC(V, Constant(0), "on_boundary")
stepper = TimeStepper(F,
butcher_tableau,
t,
dt,
u,
bcs=bc,
solver_parameters=params,
stage_type=stage_type,
splitting=splitting)
while (float(t) < 1.0):
if (float(t) + float(dt) > 1.0):
dt.assign(1.0 - float(t))
stepper.advance()
t.assign(float(t) + float(dt))
return errornorm(uexact, u) / norm(uexact)
def wave(n, deg, butcher_tableau, splitting=AI):
N = 2**n
msh = UnitIntervalMesh(N)
params = {
"snes_type": "ksponly",
"ksp_type": "preonly",
"mat_type": "aij",
"pc_type": "lu"
}
V = FunctionSpace(msh, "CG", deg)
W = FunctionSpace(msh, "DG", deg - 1)
Z = V * W
x, = SpatialCoordinate(msh)
t = Constant(0.0)
dt = Constant(2.0 / N)
up = project(as_vector([0, sin(pi * x)]), Z)
u, p = split(up)
v, w = TestFunctions(Z)
F = (inner(Dt(u), v) * dx + inner(u.dx(0), w) * dx + inner(Dt(p), w) * dx -
inner(p, v.dx(0)) * dx)
E = 0.5 * (inner(u, u) * dx + inner(p, p) * dx)
stepper = TimeStepper(F,
butcher_tableau,
t,
dt,
up,
solver_parameters=params,
splitting=splitting)
energies = []
while (float(t) < 1.0):
if (float(t) + float(dt) > 1.0):
dt.assign(1.0 - float(t))
stepper.advance()
t.assign(float(t) + float(dt))
energies.append(assemble(E))
return np.array(energies)
def test_time_dependant_result_sine(self):
"""
Test that the solve creates a correct result for a time
dependant problem with a sine wave initial value.
T(x,0) = 10*sin(3*pi*x)
T(0,t) = T(1,t) = 0
Analytic solution: T(x, t) = 10*sin(3*pi*x)*exp(-pi*pi*9*t)
"""
m = UnitIntervalMesh(500)
V = FunctionSpace(m, 'CG', 2)
prob = SimpleTimeDep(mesh=m, V=V)
file_path = os.path.join(self.out_dir.name, 'out.pvd')
prob.set_function('C', Constant(1))
prob.set_function('K', Constant(1))
prob.set_function('S', Constant(0))
prob.set_timescale(steps=100, dt=0.00001)
prob.add_boundary('dirichlet', g=0, surface='all')
x = SpatialCoordinate(m)
prob.T_.interpolate(10 * sin(x[0] * pi * 3))
prob.T.assign(prob.T_)
prob.set_method('CrankNicolson')
solver = Solver(prob)
def analytical(x, t):
return 10 * np.sin(3 * np.pi * x) * np.exp(-np.pi * np.pi * 9 * t)
coords = np.array([i / 10 for i in range(11)])
t = 0
for i in range(10):
t = prob.dt * (i + 1) * prob.steps
solver.solve(file_path=file_path)
value = solver.u.at(coords)
expected = analytical(coords, t)
print(value)
print(expected)
self.assertTrue(np.isclose(value, expected).all())
def initialise_fields(state):
L = 1.e5
H = 1.0e4 # Height position of the model top
# Initial conditions
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vt = theta0.function_space()
Vr = rho0.function_space()
# Thermodynamic constants required for setting initial conditions
# and reference profiles
g = state.parameters.g
N = state.parameters.N
# N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz)
x, z = SpatialCoordinate(state.mesh)
Tsurf = 300.
thetab = Tsurf * exp(N**2 * z / g)
theta_b = Function(Vt).interpolate(thetab)
rho_b = Function(Vr)
# Calculate hydrostatic exner
compressible_hydrostatic_balance(state, theta_b, rho_b)
a = 5.0e3
deltaTheta = 1.0e-2
theta_pert = deltaTheta * sin(pi * z / H) / (1 + (x - L / 2)**2 / a**2)
theta0.interpolate(theta_b + theta_pert)
rho0.assign(rho_b)
u0.project(as_vector([20.0, 0.0]))
state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
x, z = SpatialCoordinate(mesh)
# N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz)
Tsurf = 300.
thetab = Tsurf * exp(N**2 * z / g)
theta_b = Function(Vt).interpolate(thetab)
rho_b = Function(Vr)
# Calculate hydrostatic Pi
compressible_hydrostatic_balance(state, theta_b, rho_b)
W_DG1 = FunctionSpace(mesh, "DG", 1)
a = 5.0e3
deltaTheta = 1.0e-2
theta_pert = deltaTheta * sin(np.pi * z / H) / (1 + (x - L / 2)**2 / a**2)
theta0.interpolate(theta_b + theta_pert)
rho0.assign(rho_b)
state.initialise([('u', u0), ('rho', rho0), ('theta', theta0)])
state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
# Set up advection schemes
rhoeqn = LinearAdvection(state,
Vr,
qbar=rho_b,
ibp="once",
equation_form="continuity")
thetaeqn = LinearAdvection(state, Vt, qbar=theta_b)
advected_fields = []
advected_fields.append(("u", NoAdvection(state, u0, None)))
def run_profliler(hybridization,
model_degree,
model_family,
mesh_degree,
cfl,
refinements,
layers,
debug,
rtol,
flexsolver=True,
stronger_smoother=False,
suppress_data_output=False):
nlayers = layers # Number of vertical layers
refinements = refinements # Number of horiz. cells = 20*(4^refinements)
hybrid = bool(hybridization)
# Set up problem parameters
parameters = CompressibleParameters()
a_ref = 6.37122e6 # Radius of the Earth (m)
X = 1.0 # Reduced-size Earth reduction factor
a = a_ref / X # Scaled radius of planet (m)
g = parameters.g # Acceleration due to gravity (m/s^2)
N = parameters.N # Brunt-Vaisala frequency (1/s)
p_0 = parameters.p_0 # Reference pressure (Pa, not hPa)
c_p = parameters.cp # SHC of dry air at const. pressure (J/kg/K)
R_d = parameters.R_d # Gas constant for dry air (J/kg/K)
kappa = parameters.kappa # R_d/c_p
T_eq = 300.0 # Isothermal atmospheric temperature (K)
p_eq = 1000.0 * 100.0 # Reference surface pressure at the equator
u_0 = 20.0 # Maximum amplitude of the zonal wind (m/s)
d = 5000.0 # Width parameter for Theta'
lamda_c = 2.0 * np.pi / 3.0 # Longitudinal centerpoint of Theta'
phi_c = 0.0 # Lat. centerpoint of Theta' (equator)
deltaTheta = 1.0 # Maximum amplitude of Theta' (K)
L_z = 20000.0 # Vert. wave length of the Theta' perturb.
gamma = (1 - kappa) / kappa
cs = sqrt(c_p * T_eq / gamma) # Speed of sound in an air parcel
if model_family == "RTCF":
# Cubed-sphere mesh
m = CubedSphereMesh(radius=a,
refinement_level=refinements,
degree=mesh_degree)
elif model_family == "RT" or model_family == "BDFM":
m = IcosahedralSphereMesh(radius=a,
refinement_level=refinements,
degree=mesh_degree)
else:
raise ValueError("Unknown family: %s" % model_family)
cell_vs = interpolate(CellVolume(m), FunctionSpace(m, "DG", 0))
a_max = fmax(cell_vs)
dx_max = sqrt(a_max)
u_max = u_0
PETSc.Sys.Print("\nDetermining Dt from specified horizontal CFL: %s" % cfl)
dt = int(cfl * (dx_max / cs))
# Height position of the model top (m)
z_top = 1.0e4
deltaz = z_top / nlayers
vertical_cfl = dt * (cs / deltaz)
PETSc.Sys.Print("""
Problem parameters:\n
Profiling linear solver for the compressible Euler equations.\n
Speed of sound in compressible atmosphere: %s,\n
Hybridized compressible solver: %s,\n
Model degree: %s,\n
Model discretization: %s,\n
Mesh degree: %s,\n
Horizontal refinements: %s,\n
Vertical layers: %s,\n
Dx (max, m): %s,\n
Dz (m): %s,\n
Dt (s): %s,\n
horizontal CFL: %s,\n
vertical CFL: %s.
""" % (cs, hybrid, model_degree, model_family, mesh_degree, refinements,
nlayers, dx_max, deltaz, dt, cfl, vertical_cfl))
# Build volume mesh
mesh = ExtrudedMesh(m,
layers=nlayers,
layer_height=deltaz,
extrusion_type="radial")
x = SpatialCoordinate(mesh)
# Create polar coordinates:
# Since we use a CG1 field, this is constant on layers
W_Q1 = FunctionSpace(mesh, "CG", 1)
z_expr = sqrt(x[0] * x[0] + x[1] * x[1] + x[2] * x[2]) - a
z = Function(W_Q1).interpolate(z_expr)
lat_expr = asin(x[2] / sqrt(x[0] * x[0] + x[1] * x[1] + x[2] * x[2]))
lat = Function(W_Q1).interpolate(lat_expr)
lon = Function(W_Q1).interpolate(atan_2(x[1], x[0]))
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt, maxk=1, maxi=1)
dirname = 'meanflow_ref'
if hybrid:
dirname += '_hybridization'
# No output
output = OutputParameters(dumpfreq=3600,
dirname=dirname,
perturbation_fields=['theta', 'rho'],
dump_vtus=False,
dump_diagnostics=False,
checkpoint=False,
log_level='INFO')
diagnostics = Diagnostics(*fieldlist)
state = State(mesh,
vertical_degree=model_degree,
horizontal_degree=model_degree,
family=model_family,
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist)
# Initial conditions
u0 = state.fields.u
theta0 = state.fields.theta
rho0 = state.fields.rho
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
x = SpatialCoordinate(mesh)
# Random velocity field
CG2 = VectorFunctionSpace(mesh, "CG", 2)
urand = Function(CG2)
urand.dat.data[:] += np.random.randn(*urand.dat.data.shape)
u0.project(urand)
# Surface temperature
G = g**2 / (N**2 * c_p)
Ts_expr = G + (T_eq - G) * exp(-(u_max * N**2 / (4 * g * g)) * u_max *
(cos(2.0 * lat) - 1.0))
Ts = Function(W_Q1).interpolate(Ts_expr)
# Surface pressure
ps_expr = p_eq * exp((u_max / (4.0 * G * R_d)) * u_max *
(cos(2.0 * lat) - 1.0)) * (Ts / T_eq)**(1.0 / kappa)
ps = Function(W_Q1).interpolate(ps_expr)
# Background pressure
p_expr = ps * (1 + G / Ts * (exp(-N**2 * z / g) - 1))**(1.0 / kappa)
p = Function(W_Q1).interpolate(p_expr)
# Background temperature
Tb_expr = G * (1 - exp(N**2 * z / g)) + Ts * exp(N**2 * z / g)
Tb = Function(W_Q1).interpolate(Tb_expr)
# Background potential temperature
thetab_expr = Tb * (p_0 / p)**kappa
thetab = Function(W_Q1).interpolate(thetab_expr)
theta_b = Function(theta0.function_space()).interpolate(thetab)
rho_b = Function(rho0.function_space())
sin_tmp = sin(lat) * sin(phi_c)
cos_tmp = cos(lat) * cos(phi_c)
r = a * acos(sin_tmp + cos_tmp * cos(lon - lamda_c))
s = (d**2) / (d**2 + r**2)
theta_pert = deltaTheta * s * sin(2 * np.pi * z / L_z)
theta0.interpolate(theta_pert)
# Compute the balanced density
PETSc.Sys.Print("Computing balanced density field...\n")
# Use vert. hybridization preconditioner for initialization
pi_params = {
'ksp_type': 'preonly',
'pc_type': 'python',
'mat_type': 'matfree',
'pc_python_type': 'gusto.VerticalHybridizationPC',
'vert_hybridization': {
'ksp_type': 'gmres',
'pc_type': 'gamg',
'pc_gamg_sym_graph': True,
'ksp_rtol': 1e-12,
'ksp_atol': 1e-12,
'mg_levels': {
'ksp_type': 'richardson',
'ksp_max_it': 3,
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
}
}
}
if debug:
pi_params['vert_hybridization']['ksp_monitor_true_residual'] = None
compressible_hydrostatic_balance(state,
theta_b,
rho_b,
top=False,
pi_boundary=(p / p_0)**kappa,
solve_for_rho=False,
params=pi_params)
# Random potential temperature perturbation
theta0.assign(0.0)
theta0.dat.data[:] += np.random.randn(len(theta0.dat.data))
# Random density field
rho0.assign(0.0)
rho0.dat.data[:] += np.random.randn(len(rho0.dat.data))
state.initialise([('u', u0), ('rho', rho0), ('theta', theta0)])
state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
# Set up advection schemes
ueqn = EulerPoincare(state, Vu)
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
thetaeqn = SUPGAdvection(state, Vt, equation_form="advective")
advected_fields = []
advected_fields.append(("u", ThetaMethod(state, u0, ueqn)))
advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn, subcycles=2)))
advected_fields.append(
("theta", SSPRK3(state, theta0, thetaeqn, subcycles=2)))
# Set up linear solver
if hybrid:
outer_solver_type = "Hybrid_SCPC"
PETSc.Sys.Print("""
Setting up hybridized solver on the traces.""")
if flexsolver:
inner_solver_type = "fgmres_gamg_gmres_smoother"
inner_parameters = {
'ksp_type': 'fgmres',
'ksp_rtol': rtol,
'ksp_max_it': 500,
'ksp_gmres_restart': 30,
'pc_type': 'gamg',
'pc_gamg_sym_graph': None,
'mg_levels': {
'ksp_type': 'gmres',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu',
'ksp_max_it': 3
}
}
else:
inner_solver_type = "fgmres_ml_richardson"
inner_parameters = {
'ksp_type': 'fgmres',
'ksp_rtol': rtol,
'ksp_max_it': 500,
'ksp_gmres_restart': 30,
'pc_type': 'ml',
'pc_mg_cycles': 1,
'pc_ml_maxNlevels': 25,
'mg_levels': {
'ksp_type': 'richardson',
'ksp_richardson_scale': 0.8,
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu',
'ksp_max_it': 3
}
}
if stronger_smoother:
inner_parameters['mg_levels']['ksp_max_it'] = 5
inner_solver_type += "_stronger"
if debug:
PETSc.Sys.Print("""Debugging on.""")
inner_parameters['ksp_monitor_true_residual'] = None
PETSc.Sys.Print("Inner solver: %s" % inner_solver_type)
# Use Firedrake static condensation interface
solver_parameters = {
'mat_type': 'matfree',
'pmat_type': 'matfree',
'ksp_type': 'preonly',
'pc_type': 'python',
'pc_python_type': 'firedrake.SCPC',
'pc_sc_eliminate_fields': '0, 1',
'condensed_field': inner_parameters
}
linear_solver = HybridizedCompressibleSolver(
state,
solver_parameters=solver_parameters,
overwrite_solver_parameters=True)
else:
outer_solver_type = "gmres_SchurPC"
PETSc.Sys.Print("""
Setting up GCR fieldsplit solver with Schur complement PC.""")
solver_parameters = {
'pc_type': 'fieldsplit',
'pc_fieldsplit_type': 'schur',
'ksp_type': 'fgmres',
'ksp_max_it': 100,
'ksp_rtol': rtol,
'pc_fieldsplit_schur_fact_type': 'FULL',
'pc_fieldsplit_schur_precondition': 'selfp',
'fieldsplit_0': {
'ksp_type': 'preonly',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
},
'fieldsplit_1': {
'ksp_type': 'preonly',
'ksp_max_it': 30,
'ksp_monitor_true_residual': None,
'pc_type': 'hypre',
'pc_hypre_type': 'boomeramg',
'pc_hypre_boomeramg_max_iter': 1,
'pc_hypre_boomeramg_agg_nl': 0,
'pc_hypre_boomeramg_coarsen_type': 'Falgout',
'pc_hypre_boomeramg_smooth_type': 'Euclid',
'pc_hypre_boomeramg_eu_bj': 1,
'pc_hypre_boomeramg_interptype': 'classical',
'pc_hypre_boomeramg_P_max': 0,
'pc_hypre_boomeramg_agg_nl': 0,
'pc_hypre_boomeramg_strong_threshold': 0.25,
'pc_hypre_boomeramg_max_levels': 25,
'pc_hypre_boomeramg_no_CF': False
}
}
inner_solver_type = "hypre"
if debug:
solver_parameters['ksp_monitor_true_residual'] = None
linear_solver = CompressibleSolver(state,
solver_parameters=solver_parameters,
overwrite_solver_parameters=True)
# Set up forcing
compressible_forcing = CompressibleForcing(state)
param_info = ParameterInfo(dt=dt,
deltax=dx_max,
deltaz=deltaz,
horizontal_courant=cfl,
vertical_courant=vertical_cfl,
family=model_family,
model_degree=model_degree,
mesh_degree=mesh_degree,
solver_type=outer_solver_type,
inner_solver_type=inner_solver_type)
# Build profiler
profiler = Profiler(parameterinfo=param_info,
state=state,
advected_fields=advected_fields,
linear_solver=linear_solver,
forcing=compressible_forcing,
suppress_data_output=suppress_data_output)
PETSc.Sys.Print("Starting profiler...\n")
profiler.run(t=0, tmax=dt)
if smooth_z:
dirname += '_smootherz'
zh = 5000.
xexpr = as_vector([x, conditional(z < zh, z + cos(0.5*pi*z/zh)**6*zs, z)])
else:
xexpr = as_vector([x, z + ((H-z)/H)*zs])
new_coords = Function(Vc).interpolate(xexpr)
mesh = Mesh(new_coords)
# sponge function
W_DG = FunctionSpace(mesh, "DG", 2)
x, z = SpatialCoordinate(mesh)
zc = H-10000.
mubar = 0.15/dt
mu_top = conditional(z <= zc, 0.0, mubar*sin((pi/2.)*(z-zc)/(H-zc))**2)
mu = Function(W_DG).interpolate(mu_top)
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt)
output = OutputParameters(dirname=dirname,
dumpfreq=18,
dumplist=['u'],
perturbation_fields=['theta', 'rho'],
log_level='INFO')
parameters = CompressibleParameters(g=9.80665, cp=1004.)
diagnostics = Diagnostics(*fieldlist)
diagnostic_fields = [CourantNumber(), VelocityZ()]
state = State(mesh, vertical_degree=1, horizontal_degree=1,
# first setup the background buoyancy profile # z.grad(bref) = N**2 # the following is symbolic algebra, using the default buoyancy frequency # from the parameters class. x[1]=z and comes from x=SpatialCoordinate(mesh) N = parameters.N bref = x[1]*(N**2) # interpolate the expression to the function b_b = Function(state.V[2]).interpolate(bref) # setup constants a = Constant(5.0e3) deltab = Constant(1.0e-2) H = Constant(H) L = Constant(L) b_pert = deltab*sin(np.pi*x[1]/H)/(1 + (x[0] - L/2)**2/a**2) # interpolate the expression to the function b0.interpolate(b_b + b_pert) # interpolate velocity to vector valued function space uinit = Function(W_VectorCG1).interpolate(as_vector([20.0,0.0])) # project to the function space we actually want to use # this step is purely because it is not yet possible to interpolate to the # vector function spaces we require for the compatible finite element # methods that we use u0.project(uinit) # pass these initial conditions to the state.initialise method state.initialise([u0, p0, b0]) # set the background buoyancy state.set_reference_profiles(b_b)
def setup_sk(dirname):
nlayers = 10 # horizontal layers
columns = 30 # number of columns
L = 1.e5
m = PeriodicIntervalMesh(columns, L)
dt = 15.0
# build volume mesh
H = 1.0e4 # Height position of the model top
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers)
# Space for initialising velocity
W_VectorCG1 = VectorFunctionSpace(mesh, "CG", 1)
W_CG1 = FunctionSpace(mesh, "CG", 1)
# vertical coordinate and normal
z = Function(W_CG1).interpolate(Expression("x[1]"))
k = Function(W_VectorCG1).interpolate(Expression(("0.","1.")))
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt)
output = OutputParameters(dirname=dirname+"/sk_linear", dumplist=['u'], dumpfreq=150)
diagnostics = Diagnostics(*fieldlist)
parameters = CompressibleParameters()
diagnostic_fields = [CourantNumber()]
state = CompressibleState(mesh, vertical_degree=1, horizontal_degree=1,
family="CG",
z=z, k=k,
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist,
diagnostic_fields=diagnostic_fields,
on_sphere=False)
# Initial conditions
u0, rho0, theta0 = Function(state.V[0]), Function(state.V[1]), Function(state.V[2])
# Thermodynamic constants required for setting initial conditions
# and reference profiles
g = parameters.g
N = parameters.N
# N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz)
Tsurf = 300.
thetab = Tsurf*exp(N**2*z/g)
theta_b = Function(state.V[2]).interpolate(thetab)
rho_b = Function(state.V[1])
compressible_hydrostatic_balance(state, theta_b, rho_b)
W_DG1 = FunctionSpace(mesh, "DG", 1)
x = Function(W_DG1).interpolate(Expression("x[0]"))
a = 5.0e3
deltaTheta = 1.0e-2
theta_pert = deltaTheta*sin(np.pi*z/H)/(1 + (x - L/2)**2/a**2)
theta0.interpolate(theta_b + theta_pert)
rho0.assign(rho_b)
state.initialise([u0, rho0, theta0])
state.set_reference_profiles(rho_b, theta_b)
state.output.meanfields = {'rho':state.rhobar, 'theta':state.thetabar}
# Set up advection schemes
advection_dict = {}
advection_dict["u"] = NoAdvection(state)
advection_dict["rho"] = LinearAdvection_V3(state, state.V[1], rho_b)
advection_dict["theta"] = LinearAdvection_Vt(state, state.V[2], theta_b)
# Set up linear solver
schur_params = {'pc_type': 'fieldsplit',
'pc_fieldsplit_type': 'schur',
'ksp_type': 'gmres',
'ksp_monitor_true_residual': True,
'ksp_max_it': 100,
'ksp_gmres_restart': 50,
'pc_fieldsplit_schur_fact_type': 'FULL',
'pc_fieldsplit_schur_precondition': 'selfp',
'fieldsplit_0_ksp_type': 'richardson',
'fieldsplit_0_ksp_max_it': 5,
'fieldsplit_0_pc_type': 'bjacobi',
'fieldsplit_0_sub_pc_type': 'ilu',
'fieldsplit_1_ksp_type': 'richardson',
'fieldsplit_1_ksp_max_it': 5,
"fieldsplit_1_ksp_monitor_true_residual": True,
'fieldsplit_1_pc_type': 'gamg',
'fieldsplit_1_pc_gamg_sym_graph': True,
'fieldsplit_1_mg_levels_ksp_type': 'chebyshev',
'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues': True,
'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues_random': True,
'fieldsplit_1_mg_levels_ksp_max_it': 5,
'fieldsplit_1_mg_levels_pc_type': 'bjacobi',
'fieldsplit_1_mg_levels_sub_pc_type': 'ilu'}
linear_solver = CompressibleSolver(state, params=schur_params)
# Set up forcing
compressible_forcing = CompressibleForcing(state, linear=True)
# build time stepper
stepper = Timestepper(state, advection_dict, linear_solver,
compressible_forcing)
return stepper, dt
# Background temperature
# Tbexp = Ts*(p/ps)**kappa/(Ts/G*((p/ps)**kappa - 1) + 1)
Tbexp = G*(1 - exp(N**2*z/g)) + Ts*exp(N**2*z/g)
Tb = Function(W_CG1).interpolate(Tbexp)
# Background potential temperature
thetabexp = Tb*(p_0/p)**kappa
thetab = Function(W_CG1).interpolate(thetabexp)
theta_b = Function(state.V[2]).interpolate(thetab)
rho_b = Function(state.V[1])
sin_tmp = sin(lat) * sin(phi_c)
cos_tmp = cos(lat) * cos(phi_c)
r = a*acos(sin_tmp + cos_tmp*cos(lon-lamda_c))
s = (d**2)/(d**2 + r**2)
theta_pert = deltaTheta*s*sin(2*np.pi*z/L_z)
theta0.interpolate(theta_b)
# Compute the balanced density
compressible_hydrostatic_balance(state, theta_b, rho_b, top=False,
pi_boundary=(p/p_0)**kappa)
theta0.interpolate(theta_pert)
theta0 += theta_b
rho0.assign(rho_b)
def setup_gw(dirname):
nlayers = 10 # horizontal layers
columns = 30 # number of columns
L = 1.e5
m = PeriodicIntervalMesh(columns, L)
dt = 6.0
# build volume mesh
H = 1.0e4 # Height position of the model top
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers)
# Space for initialising velocity
W_VectorCG1 = VectorFunctionSpace(mesh, "CG", 1)
W_CG1 = FunctionSpace(mesh, "CG", 1)
# vertical coordinate and normal
z = Function(W_CG1).interpolate(Expression("x[1]"))
k = Function(W_VectorCG1).interpolate(Expression(("0.","1.")))
fieldlist = ['u', 'p', 'b']
timestepping = TimesteppingParameters(dt=dt)
output = OutputParameters(dirname=dirname+"/gw_incompressible", dumplist=['u'], dumpfreq=5)
diagnostics = Diagnostics(*fieldlist)
parameters = CompressibleParameters(geopotential=False)
diagnostic_fields = [CourantNumber()]
state = IncompressibleState(mesh, vertical_degree=1, horizontal_degree=1,
family="CG",
z=z, k=k,
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist,
diagnostic_fields=diagnostic_fields,
on_sphere=False)
# Initial conditions
u0, p0, b0 = Function(state.V[0]), Function(state.V[1]), Function(state.V[2])
# Thermodynamic constants required for setting initial conditions
# and reference profiles
N = parameters.N
# z.grad(bref) = N**2
N = parameters.N
bref = z*(N**2)
b_b = Function(state.V[2]).interpolate(bref)
W_DG1 = FunctionSpace(mesh, "DG", 1)
x = Function(W_DG1).interpolate(Expression("x[0]"))
a = 5.0e3
deltaTheta = 1.0e-2
theta_pert = deltaTheta*sin(np.pi*z/H)/(1 + (x - L/2)**2/a**2)
b0.interpolate(b_b + theta_pert)
u0.project(as_vector([20.0,0.0]))
state.initialise([u0, p0, b0])
state.set_reference_profiles(b_b)
state.output.meanfields = {'b':state.bbar}
# Set up advection schemes
Vtdg = FunctionSpace(mesh, "DG", 1)
advection_dict = {}
advection_dict["u"] = EulerPoincareForm(state, state.V[0])
advection_dict["b"] = EmbeddedDGAdvection(state, state.V[2],
Vdg=Vtdg, continuity=False)
# Set up linear solver
params = {'ksp_type':'gmres',
'pc_type':'fieldsplit',
'pc_fieldsplit_type':'additive',
'fieldsplit_0_pc_type':'lu',
'fieldsplit_1_pc_type':'lu',
'fieldsplit_0_ksp_type':'preonly',
'fieldsplit_1_ksp_type':'preonly'}
linear_solver = IncompressibleSolver(state, L, params=params)
# Set up forcing
forcing = IncompressibleForcing(state)
# build time stepper
stepper = Timestepper(state, advection_dict, linear_solver,
forcing)
return stepper, 10*dt
# We declare a function over our function space and give it the
# value of our right hand side function::
# Permeability tensor
# Homogeneous
one = fd.Constant(1.0)
zero = fd.Constant(0.0)
Khom = fd.interpolate(one, V)
Khom.rename('K', 'Permeability')
# Heterogeneous
x, y, z = fd.SpatialCoordinate(W.mesh())
Kx = fd.Function(V).interpolate(1.0+y)
Ky = fd.Function(V).interpolate(1.0+0.25*fd.sin(fd.pi/Ly*y/2))
Kz = fd.Function(V).interpolate(z)
Kx.rename('Kx', 'Permeability in x direction.')
Ky.rename('Ky', 'Permeability in y direction.')
Kz.rename('Kz', 'Permeability in z direction.')
#Kx = fd.Constant(1.0)
#Ky = fd.Constant(1.0)
#Kz = fd.Constant(1.0)
# Permeability field harmonic interpolation to facets
Kx_facet = fd.conditional(fd.gt(fd.avg(Kx), 0.0), Kx('+')*Kx('-') / fd.avg(Kx), 0.0)
Ky_facet = fd.conditional(fd.gt(fd.avg(Ky), 0.0), Ky('+')*Ky('-') / fd.avg(Ky), 0.0)
Kz_facet = fd.conditional(fd.gt(fd.avg(Kz), 0.0), Kz('+')*Kz('-') / fd.avg(Kz), 0.0)
# N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz)
Tsurf = 300.
thetab = Tsurf*exp(N**2*z/g)
theta_b = Function(state.V[2]).interpolate(thetab)
rho_b = Function(state.V[1])
# Calculate hydrostatic Pi
compressible_hydrostatic_balance(state, theta_b, rho_b)
W_DG1 = FunctionSpace(mesh, "DG", 1)
x = Function(W_DG1).interpolate(Expression("x[0]"))
a = 5.0e3
deltaTheta = 1.0e-2
theta_pert = deltaTheta*sin(np.pi*z/H)/(1 + (x - L/2)**2/a**2)
theta0.interpolate(theta_b + theta_pert)
rho0.assign(rho_b)
state.initialise([u0, rho0, theta0])
state.set_reference_profiles(rho_b, theta_b)
state.output.meanfields = {'rho':state.rhobar, 'theta':state.thetabar}
# Set up advection schemes
advection_dict = {}
advection_dict["u"] = NoAdvection(state)
advection_dict["rho"] = LinearAdvection_V3(state, state.V[1], rho_b)
advection_dict["theta"] = LinearAdvection_Vt(state, state.V[2], theta_b)
# Set up linear solver
schur_params = {'pc_type': 'fieldsplit',
