def tracer_sphere(tmpdir, degree):
radius = 1
mesh = IcosahedralSphereMesh(radius=radius, refinement_level=3, degree=1)
x = SpatialCoordinate(mesh)
mesh.init_cell_orientations(x)
# Parameters chosen so that dt != 1
# Gaussian is translated from (lon=pi/2, lat=0) to (lon=0, lat=0)
# to demonstrate that transport is working correctly
dt = pi / 3. * 0.02
output = OutputParameters(dirname=str(tmpdir), dumpfreq=15)
state = State(mesh, dt=dt, output=output)
umax = 1.0
uexpr = as_vector([-umax * x[1] / radius, umax * x[0] / radius, 0.0])
tmax = pi / 2
f_init = exp(-x[2]**2 - x[0]**2)
f_end = exp(-x[2]**2 - x[1]**2)
tol = 0.05
return TracerSetup(state, tmax, f_init, f_end, "BDM", degree, uexpr, umax,
radius, tol)
def delta_expr(x0, x, y, z=None, sigma_x=2000.0):
"""Spatial function to apply source"""
sigma_x = Constant(sigma_x)
if z is None:
return exp(-sigma_x * ((x - x0[0]) ** 2 + (y - x0[1]) ** 2))
else:
return exp(-sigma_x * ((x - x0[0]) ** 2 + (y - x0[1]) ** 2 + (z - x0[2]) ** 2))
def update(self):
dp = 0
dq = (self.p/2)*(1+exp(-self.Ld/self.alpha))/(1-exp(-self.Ld/self.alpha)) * self.dt
if self.method == 'ito_euler':
for xi_field, xi_x_field, xi_xx_field, dW in zip(self.pure_xi_list, self.pure_xi_x_list, self.pure_xi_xx_list, self.dWs):
xi = xi_field.at(self.q, tolerance=1e-6)
xi_x = xi_x_field.at(self.q, tolerance=1e-6)
xi_xx = xi_xx_field.at(self.q, tolerance=1e-6)
dp += self.p/2*(xi_x*xi_x - xi*xi_xx)*self.dt - self.p*xi_x*dW
dq += 0.5*xi*xi_x*self.dt + xi*dW
elif self.method == 'milstein':
for xi_field, xi_x_field, xi_xx_field, dW in zip(self.pure_xi_list, self.pure_xi_x_list, self.pure_xi_xx_list, self.dWs):
xi = xi_field.at(self.q, tolerance=1e-6)
xi_x = xi_x_field.at(self.q, tolerance=1e-6)
xi_xx = xi_xx_field.at(self.q, tolerance=1e-6)
dp += self.p/2*(xi_x*xi_x - xi*xi_xx)*dW**2 - self.p*xi_x*dW
dq += 0.5*xi*xi_x*dW**2 + xi*dW
self.p += dp
self.q += dq
if self.q < 0:
self.q += self.Ld
if self.q > self.Ld:
self.q -= self.Ld
def update_constants(self):
"""
Update p and q constants from true peakon data.
"""
self.p.assign(self.true_peakon_file['p'][self.outputting.t_idx] * 0.5 *
(1 + exp(-self.Ld / sqrt(self.alphasq))) /
(1 - exp(-self.Ld / sqrt(self.alphasq))))
self.q.assign(self.true_peakon_file['q'][self.outputting.t_idx])
def setup_IPdiffusion(dirname, vector, DG):
dt = 0.01
L = 10.
m = PeriodicIntervalMesh(50, L)
mesh = ExtrudedMesh(m, layers=50, layer_height=0.2)
fieldlist = ['u', 'D']
timestepping = TimesteppingParameters(dt=dt)
parameters = CompressibleParameters()
output = OutputParameters(dirname=dirname)
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
parameters=parameters,
output=output,
fieldlist=fieldlist)
x = SpatialCoordinate(mesh)
if vector:
kappa = Constant([[0.05, 0.], [0., 0.05]])
if DG:
Space = VectorFunctionSpace(mesh, "DG", 1)
else:
Space = state.spaces("HDiv")
fexpr = as_vector([exp(-(L / 2. - x[0])**2 - (L / 2. - x[1])**2), 0.])
else:
kappa = 0.05
if DG:
Space = state.spaces("DG")
else:
Space = state.spaces("HDiv_v")
fexpr = exp(-(L / 2. - x[0])**2 - (L / 2. - x[1])**2)
f = state.fields("f", Space)
try:
f.interpolate(fexpr)
except NotImplementedError:
# finat elements raise NotImplementedError if they don't
# advertise a dual for interpolation.
f.project(fexpr)
mu = 5.
f_diffusion = InteriorPenalty(state,
f.function_space(),
kappa=kappa,
mu=mu)
diffused_fields = [("f", f_diffusion)]
stepper = AdvectionDiffusion(state, diffused_fields=diffused_fields)
return stepper
def run_advection_diffusion(tmpdir):
# Mesh, state and equation
L = 10
mesh = PeriodicIntervalMesh(20, L)
dt = 0.02
tmax = 1.0
diffusion_params = DiffusionParameters(kappa=0.75, mu=5)
output = OutputParameters(dirname=str(tmpdir), dumpfreq=25)
state = State(mesh, dt=dt, output=output)
V = state.spaces("DG", "DG", 1)
Vu = VectorFunctionSpace(mesh, "CG", 1)
equation = AdvectionDiffusionEquation(
state, V, "f", Vu=Vu, diffusion_parameters=diffusion_params)
problem = [(equation, ((SSPRK3(state), transport), (BackwardEuler(state),
diffusion)))]
# Initial conditions
x = SpatialCoordinate(mesh)
xc_init = 0.25 * L
xc_end = 0.75 * L
umax = 0.5 * L / tmax
# Get minimum distance on periodic interval to xc
x_init = conditional(
sqrt((x[0] - xc_init)**2) < 0.5 * L, x[0] - xc_init,
L + x[0] - xc_init)
x_end = conditional(
sqrt((x[0] - xc_end)**2) < 0.5 * L, x[0] - xc_end, L + x[0] - xc_end)
f_init = 5.0
f_end = f_init / 2.0
f_width_init = L / 10.0
f_width_end = f_width_init * 2.0
f_init_expr = f_init * exp(-(x_init / f_width_init)**2)
f_end_expr = f_end * exp(-(x_end / f_width_end)**2)
state.fields('f').interpolate(f_init_expr)
state.fields('u').interpolate(as_vector([Constant(umax)]))
f_end = state.fields('f_end', V).interpolate(f_end_expr)
# Time stepper
timestepper = PrescribedTransport(state, problem)
timestepper.run(0, tmax=tmax)
error = norm(state.fields('f') - f_end) / norm(f_end)
return error
def test_ilu():
"""Tests that ILU functionality gives correct solution."""
k = 10.0
num_cells = utils.h_to_num_cells(k**-1.5,2)
mesh = fd.UnitSquareMesh(num_cells,num_cells)
V = fd.FunctionSpace(mesh,"CG",1)
prob = hh.HelmholtzProblem(k,V)
angle = 2.0 * np.pi/7.0
d = [np.cos(angle),np.sin(angle)]
prob.f_g_plane_wave(d)
for fill_in in range(40):
prob.use_ilu_gmres(fill_in)
prob.solve()
x = fd.SpatialCoordinate(mesh)
# This error was found out by eye
assert np.abs(fd.norms.errornorm(fd.exp(1j * k * fd.dot(fd.as_vector(d),x)),uh=prob.u_h,norm_type='H1')) < 0.5
def interpolate(self, mesh, k0L):
V = fd.VectorFunctionSpace(mesh, "CG", 1)
x = fd.SpatialCoordinate(mesh)
pw = fd.interpolate(self.p * fd.exp(1j * k0L * fd.dot(self.s, x)), V)
return pw
def build(self):
"""
Build the gaussian function.
Raises:
AttributeError: If required properties are not defined.
Returns:
Function: firedrake Constant set to the given value.
"""
for k in self.properties:
if self._props[k] is None:
raise AttributeError('"{}" has not been defined.'.format(k))
mean = self._props['mean']
sd = self._props['sd']
scale = self._props['scale']
xs = SpatialCoordinate(self.mesh)
if not isinstance(mean, list):
mean = [mean] * len(xs)
if not isinstance(sd, list):
sd = [sd] * len(xs)
gaussian = Function(self.V)
components = [exp(-(x - m)**2 / 2 / s**2)
for x, m, s in zip(xs, mean, sd)]
product = components[0]
for c in components[1:]:
product *= c
gaussian.interpolate(scale * product)
return gaussian
def viscosityTheta(u, h, grounded, floating, A, theta):
"""Combine existing A on grounded ice with theta^2 from floating ice.
Parameters
----------
u : firedrake function
modelled velocity
h : firedrake function
model thickness
grounded : firedrake function
Mask with 1s for grounded 0 for floating.
floating : firedrake function
Mask with 1s for floating 0 for grounded.
A : firedrake function
A for grounded areas
theta : firedrake function
theta for A on floating ice as theta^2
Returns
-------
viscosityFunction function
Viscosity funciton for model
"""
# A = grounded * A + floating * theta**2
A0 = grounded * A + floating * firedrake.exp(theta)
# A 0 = firedrake.max_value(firedrake.min_value(A0, 1.), 100.)
return icepack.models.viscosity.viscosity_depth_averaged(u, h, A0)
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 viscosity(**kwargs):
u = kwargs['velocity']
h = kwargs['thickness']
θ = kwargs['log_fluidity']
A = A0 * firedrake.exp(θ)
return icepack.models.viscosity.viscosity_depth_averaged(
velocity=u, thickness=h, fluidity=A
)
def f(x):
"""Helper function for one_d_transition.
It didn't work when coded up directly....
"""
machine_eps = np.finfo(float).eps
return fd.exp(-new_max(x, machine_eps)**(-1.0)) * heaviside(x)
def viscosity(**kwargs):
u = kwargs["velocity"]
h = kwargs["thickness"]
q = kwargs["log_fluidity"]
A = A0 * firedrake.exp(-q / n)
return icepack.models.viscosity.viscosity_depth_averaged(velocity=u,
thickness=h,
fluidity=A)
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 __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 __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 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 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 taperedViscosity(velocity, thickness, fluidity, theta, groundedSmooth,
floatingSmooth):
''' This version uses feathered grouning to floating transition
'''
# Combine tapered A on grouned and theta for smooth
A = groundedSmooth * fluidity + floatingSmooth * firedrake.exp(theta)
viscosity = \
icepack.models.viscosity.viscosity_depth_averaged(velocity=velocity,
thickness=thickness,
fluidity=A)
return viscosity
def get_true_sol_expr(spatial_coord, kappa):
if mesh_dim == 3:
x, y, z = spatial_coord
norm = sqrt(x**2 + y**2 + z**2)
return Constant(1j / (4 * pi)) / norm * exp(1j * kappa * norm)
elif mesh_dim == 2:
x, y = spatial_coord
return Constant(1j / 4) * hankel_function(kappa * sqrt(x**2 + y**2),
n=hankel_cutoff)
raise ValueError("Only meshes of dimension 2, 3 supported")
def __init__(self, prognostic_variables, simulation_parameters, peakon_speed, method='ito_euler'):
self.dWs = prognostic_variables.dW_nums
self.pure_xi_list = prognostic_variables.pure_xi_list
self.pure_xi_x_list = prognostic_variables.pure_xi_x_list
self.pure_xi_xx_list = prognostic_variables.pure_xi_xx_list
self.u = prognostic_variables.u
mesh = simulation_parameters['mesh'][-1]
x, = SpatialCoordinate(mesh)
coords = Function(self.u.function_space()).project(x)
self.dt = simulation_parameters['dt'][-1]
self.Ld = simulation_parameters['Ld'][-1]
self.alpha = np.sqrt(simulation_parameters['alphasq'][-1].dat.data[0])
self.method = method
p0 = np.max(self.u.dat.data[:])*2*(1-exp(-self.Ld/self.alpha))/(1+exp(-self.Ld/self.alpha))
q0 = coords.dat.data[np.argmax(self.u.dat.data[:])]
self.p = peakon_speed if peakon_speed is not None else p0
self.q = q0
def test_HelmholtzProblem_solver_convergence():
"""Test that the solver is converging at the correct rate."""
k_range = [10.0,12.0,20.0,30.0,40.0]#[10.0,12.0]#[20.0,40.0]
num_levels = 2
tolerance = 0.05
log_err_L2 = np.empty((num_levels,len(k_range)))
log_err_H1 = np.empty((num_levels,len(k_range)))
for ii_k in range(len(k_range)):
k = k_range[ii_k]
print(k)
num_points = np.ceil(np.sqrt(2.0) * k**(1.5)) * 2.0**np.arange(float(num_levels))
log_h = np.log(np.sqrt(2.0) * 1.0 / num_points)
for ii_points in range(num_levels):
print(ii_points)
# Coarsest grid has h ~ k^{-1.5}, and then do uniform refinement
mesh = fd.UnitSquareMesh(num_points[ii_points],num_points[ii_points])
V = fd.FunctionSpace(mesh, "CG", 1)
prob = hh.HelmholtzProblem(k,V)
angle = 2.0*np.pi/7.0
d = [np.cos(angle),np.sin(angle)]
prob.f_g_plane_wave(d)
prob.solve()
x = fd.SpatialCoordinate(mesh)
exact_soln = fd.exp(1j * k * fd.dot(x,fd.as_vector(d)))
log_err_L2[ii_points,ii_k] = np.log(fd.norms.errornorm(exact_soln,prob.u_h,norm_type="L2"))
log_err_H1[ii_points,ii_k] = np.log(fd.norms.errornorm(exact_soln,prob.u_h,norm_type="H1"))
#plt.plot(log_h,log_err_L2[:,ii_k])
#plt.plot(log_h,log_err_H1[:,ii_k])
#plt.show()
fit_L2 = np.polyfit(log_h,log_err_L2[:,ii_k],deg=1)[0]
fit_H1 = np.polyfit(log_h,log_err_H1[:,ii_k],deg=1)[0]
print(fit_L2)
print(fit_H1)
assert abs(fit_L2 - 2.0) <= tolerance
assert abs(fit_H1 - 1.0) <= tolerance
def get_true_sol_expr(spatial_coord, wave_number):
mesh_dim = len(spatial_coord)
if mesh_dim == 3:
x, y, z = spatial_coord
norm = sqrt(x**2 + y**2 + z**2)
return Constant(1j / (4 * pi)) / norm * exp(1j * wave_number * norm)
elif mesh_dim == 2:
x, y = spatial_coord
return Constant(1j / 4) * hankel_function(
wave_number * sqrt(x**2 + y**2), n=80)
raise ValueError("Only meshes of dimension 2, 3 supported")
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 setup_gaussian(dirname):
n = 16
L = 1.
mesh = PeriodicSquareMesh(n, n, L)
fieldlist = ['u', 'D']
parameters = ShallowWaterParameters(H=1.0, g=1.0)
timestepping = TimesteppingParameters(dt=0.1)
output = OutputParameters(dirname=dirname + '/sw_plane_gaussian_subcycled')
diagnostic_fields = [CourantNumber()]
state = State(mesh,
horizontal_degree=1,
family="BDM",
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostic_fields=diagnostic_fields,
fieldlist=fieldlist)
u0 = state.fields("u")
D0 = state.fields("D")
x, y = SpatialCoordinate(mesh)
H = Constant(state.parameters.H)
D0.interpolate(H + exp(-50 * ((x - 0.5)**2 + (y - 0.5)**2)))
V = FunctionSpace(mesh, "CG", 1)
f = state.fields("coriolis", V)
f.interpolate(Constant(1.)) # Coriolis frequency (1/s)
state.initialise([("u", u0), ("D", D0)])
ueqn = EmbeddedDGAdvection(state,
u0.function_space(),
options=EmbeddedDGOptions())
Deqn = AdvectionEquation(state,
D0.function_space(),
equation_form="continuity")
advected_fields = []
advected_fields.append(("u", SSPRK3(state, u0, ueqn, subcycles=2)))
advected_fields.append(("D", SSPRK3(state, D0, Deqn, subcycles=2)))
linear_solver = ShallowWaterSolver(state)
# Set up forcing
sw_forcing = ShallowWaterForcing(state)
# build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver, sw_forcing)
return stepper
def tracer_blob_slice(tmpdir, degree):
dt = 0.01
L = 10.
m = PeriodicIntervalMesh(10, L)
mesh = ExtrudedMesh(m, layers=10, layer_height=1.)
output = OutputParameters(dirname=str(tmpdir), dumpfreq=25)
state = State(mesh, dt=dt, output=output)
tmax = 1.
x = SpatialCoordinate(mesh)
f_init = exp(-((x[0] - 0.5 * L)**2 + (x[1] - 0.5 * L)**2))
return TracerSetup(state, tmax, f_init, family="CG", degree=degree)
def tracer_slice(tmpdir, degree):
n = 30 if degree == 0 else 15
m = PeriodicIntervalMesh(n, 1.)
mesh = ExtrudedMesh(m, layers=n, layer_height=1. / n)
# Parameters chosen so that dt != 1 and u != 1
# Gaussian is translated by 1.5 times width of domain to demonstrate
# that transport is working correctly
dt = 0.01
tmax = 0.75
output = OutputParameters(dirname=str(tmpdir), dumpfreq=25)
state = State(mesh, dt=dt, output=output)
uexpr = as_vector([2.0, 0.0])
x = SpatialCoordinate(mesh)
width = 1. / 10.
f0 = 0.5
fmax = 2.0
xc_init = 0.25
xc_end = 0.75
r_init = sqrt((x[0] - xc_init)**2 + (x[1] - 0.5)**2)
r_end = sqrt((x[0] - xc_end)**2 + (x[1] - 0.5)**2)
f_init = f0 + (fmax - f0) * exp(-(r_init / width)**2)
f_end = f0 + (fmax - f0) * exp(-(r_end / width)**2)
tol = 0.12
return TracerSetup(state,
tmax,
f_init,
f_end,
"CG",
degree,
uexpr,
tol=tol)
def e_sat(parameters, T):
"""
Returns an expression for the saturated partial pressure
of water vapour as a function of T, based on Tetens' formula.
:arg parameters: a CompressibleParameters object.
:arg T: the temperature in K.
"""
w_sat2 = parameters.w_sat2
w_sat3 = parameters.w_sat3
w_sat4 = parameters.w_sat4
T_0 = parameters.T_0
return w_sat4 * exp(-w_sat2 * (T - T_0) / (T - w_sat3))
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)
# 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
p_0 = parameters.p_0
c_p = parameters.cp
R_d = parameters.R_d
kappa = parameters.kappa
# 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)
# Calculate hydrostatic Pi
params = {'pc_type': 'fieldsplit',
'pc_fieldsplit_type': 'schur',
'ksp_type': 'gmres',
'ksp_monitor_true_residual': True,
'ksp_max_it': 1000,
'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',
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
# Initial conditions u0, theta0, rho0 = Function(state.V[0]), Function(state.V[2]), Function(state.V[1]) # Initial conditions without u0 # Surface temperature G = g**2/(N**2*c_p) Ts = Function(W_CG1).assign(T_eq) # surface pressure psexp = p_eq*((Ts/T_eq)**(1.0/kappa)) ps = Function(W_CG1).interpolate(psexp) # Background pressure pexp = ps*(1 + G/Ts*(exp(-N**2*z/g)-1))**(1.0/kappa) p = Function(W_CG1).interpolate(pexp) # 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)