def setup_balance(dirname):
# set up grid and time stepping parameters
dt = 1.
tmax = 5.
deltax = 400
L = 2000.
H = 10000.
nlayers = int(H / deltax)
ncolumns = int(L / deltax)
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
output = OutputParameters(dirname=dirname + '/dry_balance',
dumpfreq=10,
dumplist=['u'])
parameters = CompressibleParameters()
state = State(mesh, dt=dt, output=output, parameters=parameters)
eqns = CompressibleEulerEquations(state, "CG", 1)
# Initial conditions
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# Isentropic background state
Tsurf = Constant(300.)
theta0.interpolate(Tsurf)
# Calculate hydrostatic exner
compressible_hydrostatic_balance(state, theta0, rho0, solve_for_rho=True)
state.set_reference_profiles([('rho', rho0), ('theta', theta0)])
# Set up transport schemes
transported_fields = [
ImplicitMidpoint(state, "u"),
SSPRK3(state, "rho"),
SSPRK3(state, "theta", options=EmbeddedDGOptions())
]
# Set up linear solver
linear_solver = CompressibleSolver(state, eqns)
# build time stepper
stepper = CrankNicolson(state,
eqns,
transported_fields,
linear_solver=linear_solver)
return stepper, tmax
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 mesh(geometry):
Lx = 100.
deltax = Lx / 5.
ncolumnsx = int(Lx/deltax)
if geometry == "periodic":
m = PeriodicIntervalMesh(ncolumnsx, Lx)
elif geometry == "non-periodic":
m = IntervalMesh(ncolumnsx, Lx)
return m
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 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 setup_fallout(dirname):
# declare grid shape, with length L and height H
L = 10.
H = 10.
nlayers = 10
ncolumns = 10
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
x = SpatialCoordinate(mesh)
dt = 0.1
output = OutputParameters(dirname=dirname + "/fallout",
dumpfreq=10,
dumplist=['rain_mixing_ratio'])
parameters = CompressibleParameters()
diagnostic_fields = [Precipitation()]
state = State(mesh,
dt=dt,
output=output,
parameters=parameters,
diagnostic_fields=diagnostic_fields)
Vrho = state.spaces("DG1_equispaced")
problem = [(AdvectionEquation(state, Vrho, "rho", ufamily="CG",
udegree=1), ForwardEuler(state))]
state.fields("rho").assign(1.)
physics_list = [Fallout(state)]
rain0 = state.fields("rain_mixing_ratio")
# set up rain
xc = L / 2
zc = H / 2
rc = H / 4
r = sqrt((x[0] - xc)**2 + (x[1] - zc)**2)
rain_expr = conditional(r > rc, 0., 1e-3 * (cos(pi * r / (rc * 2)))**2)
rain0.interpolate(rain_expr)
# build time stepper
stepper = PrescribedTransport(state, problem, physics_list=physics_list)
return stepper, 10.0
def mesh(geometry):
L = 100.
H = 100.
deltax = L / 5.
deltaz = H / 5.
nlayers = int(H / deltaz)
ncolumns = int(L / deltax)
if geometry == "periodic":
m = PeriodicIntervalMesh(ncolumns, L)
elif geometry == "non-periodic":
m = IntervalMesh(ncolumns, L)
extruded_mesh = ExtrudedMesh(m, layers=nlayers, layer_height=deltaz)
return extruded_mesh
def state(tmpdir, geometry):
"""
returns an instance of the State class, having set up either spherical
geometry or 2D vertical slice geometry
"""
output = OutputParameters(dirname=str(tmpdir), dumplist=["f"], dumpfreq=15)
if geometry == "sphere":
mesh = IcosahedralSphereMesh(radius=1, refinement_level=3, degree=1)
x = SpatialCoordinate(mesh)
mesh.init_cell_orientations(x)
family = "BDM"
vertical_degree = None
fieldlist = ["u", "D"]
dt = pi / 3. * 0.01
uexpr = as_vector([-x[1], x[0], 0.0])
if geometry == "slice":
m = PeriodicIntervalMesh(15, 1.)
mesh = ExtrudedMesh(m, layers=15, layer_height=1. / 15.)
family = "CG"
vertical_degree = 1
fieldlist = ["u", "rho", "theta"]
dt = 0.01
x = SpatialCoordinate(mesh)
uexpr = as_vector([1.0, 0.0])
timestepping = TimesteppingParameters(dt=dt)
state = State(mesh,
vertical_degree=vertical_degree,
horizontal_degree=1,
family=family,
timestepping=timestepping,
output=output,
fieldlist=fieldlist)
u0 = state.fields("u")
u0.project(uexpr)
return state
def setup_checkpointing(dirname):
nlayers = 5 # horizontal layers
columns = 15 # number of columns
L = 3.e5
m = PeriodicIntervalMesh(columns, L)
dt = 2.0
# build volume mesh
H = 1.0e4 # Height position of the model top
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
output = OutputParameters(dirname=dirname,
dumpfreq=1,
chkptfreq=2,
log_level='INFO')
parameters = CompressibleParameters()
state = State(mesh, dt=dt, output=output, parameters=parameters)
eqns = CompressibleEulerEquations(state, "CG", 1)
# Set up transport schemes
transported_fields = []
transported_fields.append(SSPRK3(state, "u"))
transported_fields.append(SSPRK3(state, "rho"))
transported_fields.append(SSPRK3(state, "theta"))
# Set up linear solver
linear_solver = CompressibleSolver(state, eqns)
# build time stepper
stepper = CrankNicolson(state,
eqns,
transported_fields,
linear_solver=linear_solver)
return state, stepper, dt
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 setup_tracer(dirname, hybridization):
# declare grid shape, with length L and height H
L = 1000.
H = 1000.
nlayers = int(H / 100.)
ncolumns = int(L / 100.)
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=10.0, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname+"/tracer",
dumpfreq=1,
dumplist=['u'],
perturbation_fields=['theta', 'rho'])
parameters = CompressibleParameters()
state = State(mesh, vertical_degree=1, horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist,
diagnostic_fields=[Difference('theta', 'tracer')])
# declare initial fields
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
# declare tracer field and a background field
tracer0 = state.fields("tracer", Vt)
# Isentropic background state
Tsurf = Constant(300.)
theta_b = Function(Vt).interpolate(Tsurf)
rho_b = Function(Vr)
# Calculate initial rho
compressible_hydrostatic_balance(state, theta_b, rho_b,
solve_for_rho=True)
# set up perturbation to theta
xc = 500.
zc = 350.
rc = 250.
x = SpatialCoordinate(mesh)
r = sqrt((x[0]-xc)**2 + (x[1]-zc)**2)
theta_pert = conditional(r > rc, 0., 0.25*(1. + cos((pi/rc)*r)))
theta0.interpolate(theta_b + theta_pert)
rho0.interpolate(rho_b)
tracer0.interpolate(theta0)
state.initialise([('u', u0),
('rho', rho0),
('theta', theta0),
('tracer', tracer0)])
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,
supg_params={"dg_direction": "horizontal"},
equation_form="advective")
# build advection dictionary
advected_fields = []
advected_fields.append(("u", ThetaMethod(state, u0, ueqn)))
advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn)))
advected_fields.append(("theta", SSPRK3(state, theta0, thetaeqn)))
advected_fields.append(("tracer", SSPRK3(state, tracer0, thetaeqn)))
# Set up linear solver
if hybridization:
linear_solver = HybridizedCompressibleSolver(state)
else:
linear_solver = CompressibleSolver(state)
compressible_forcing = CompressibleForcing(state)
# build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver,
compressible_forcing)
return stepper, 100.0
def setup_saturated(dirname):
# set up grid and time stepping parameters
dt = 1.
tmax = 3.
deltax = 400.
L = 2000.
H = 10000.
nlayers = int(H / deltax)
ncolumns = int(L / deltax)
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
# option to easily change between recovered and not if necessary
# default should be to use lowest order set of spaces
recovered = True
degree = 0 if recovered else 1
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname + '/saturated_balance',
dumpfreq=1,
dumplist=['u'],
perturbation_fields=['water_v'])
parameters = CompressibleParameters()
diagnostics = Diagnostics(*fieldlist)
diagnostic_fields = [Theta_e()]
state = State(mesh,
vertical_degree=degree,
horizontal_degree=degree,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist,
diagnostic_fields=diagnostic_fields)
# Initial conditions
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
water_v0 = state.fields("water_v", theta0.function_space())
water_c0 = state.fields("water_c", theta0.function_space())
moisture = ['water_v', 'water_c']
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
# Isentropic background state
Tsurf = Constant(300.)
total_water = Constant(0.02)
theta_e = Function(Vt).interpolate(Tsurf)
water_t = Function(Vt).interpolate(total_water)
# Calculate hydrostatic Pi
saturated_hydrostatic_balance(state, theta_e, water_t)
water_c0.assign(water_t - water_v0)
state.initialise([('u', u0), ('rho', rho0), ('theta', theta0),
('water_v', water_v0), ('water_c', water_c0)])
state.set_reference_profiles([('rho', rho0), ('theta', theta0),
('water_v', water_v0)])
# Set up advection schemes
if recovered:
VDG1 = FunctionSpace(mesh, "DG", 1)
VCG1 = FunctionSpace(mesh, "CG", 1)
Vt_brok = FunctionSpace(mesh, BrokenElement(Vt.ufl_element()))
Vu_DG1 = VectorFunctionSpace(mesh, "DG", 1)
Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1)
u_opts = RecoveredOptions(embedding_space=Vu_DG1,
recovered_space=Vu_CG1,
broken_space=Vu,
boundary_method=Boundary_Method.dynamics)
rho_opts = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=Vr,
boundary_method=Boundary_Method.dynamics)
theta_opts = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=Vt_brok)
ueqn = EmbeddedDGAdvection(state,
Vu,
equation_form="advective",
options=u_opts)
rhoeqn = EmbeddedDGAdvection(state,
Vr,
equation_form="continuity",
options=rho_opts)
thetaeqn = EmbeddedDGAdvection(state,
Vt,
equation_form="advective",
options=theta_opts)
else:
ueqn = EulerPoincare(state, Vu)
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
thetaeqn = EmbeddedDGAdvection(state,
Vt,
equation_form="advective",
options=EmbeddedDGOptions())
advected_fields = [('rho', SSPRK3(state, rho0, rhoeqn)),
('theta', SSPRK3(state, theta0, thetaeqn)),
('water_v', SSPRK3(state, water_v0, thetaeqn)),
('water_c', SSPRK3(state, water_c0, thetaeqn))]
if recovered:
advected_fields.append(('u', SSPRK3(state, u0, ueqn)))
else:
advected_fields.append(('u', ThetaMethod(state, u0, ueqn)))
linear_solver = CompressibleSolver(state, moisture=moisture)
# Set up forcing
if recovered:
compressible_forcing = CompressibleForcing(state,
moisture=moisture,
euler_poincare=False)
else:
compressible_forcing = CompressibleForcing(state, moisture=moisture)
# add physics
physics_list = [Condensation(state)]
# build time stepper
stepper = CrankNicolson(state,
advected_fields,
linear_solver,
compressible_forcing,
physics_list=physics_list)
return stepper, tmax
def run_cond_evap(dirname, process):
# declare grid shape, with length L and height H
L = 1000.
H = 1000.
nlayers = int(H / 10.)
ncolumns = int(L / 10.)
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
x, z = SpatialCoordinate(mesh)
dt = 2.0
tmax = dt
output = OutputParameters(dirname=dirname + "/cond_evap",
dumpfreq=1,
dumplist=['u'])
parameters = CompressibleParameters()
state = State(mesh,
dt=dt,
output=output,
parameters=parameters,
diagnostic_fields=[
Sum('vapour_mixing_ratio', 'cloud_liquid_mixing_ratio')
])
# spaces
Vt = state.spaces("theta", degree=1)
Vr = state.spaces("DG", "DG", degree=1)
# Set up equation -- use compressible to set up these spaces
# However the equation itself will be unused
_ = CompressibleEulerEquations(state, "CG", 1)
# Declare prognostic fields
rho0 = state.fields("rho")
theta0 = state.fields("theta")
water_v0 = state.fields("vapour_mixing_ratio", Vt)
water_c0 = state.fields("cloud_liquid_mixing_ratio", Vt)
# Set a background state with constant pressure and temperature
pressure = Function(Vr).interpolate(Constant(100000.))
temperature = Function(Vt).interpolate(Constant(300.))
theta_d = td.theta(parameters, temperature, pressure)
mv_sat = td.r_v(parameters, Constant(1.0), temperature, pressure)
Lv_over_cpT = td.Lv(parameters,
temperature) / (parameters.cp * temperature)
# Apply perturbation
xc = L / 2.
zc = H / 2.
rc = L / 4.
r = sqrt((x - xc)**2 + (z - zc)**2)
pert = conditional(r < rc, 1.0, 0.0)
if process == "evaporation":
water_v0.interpolate(0.96 * mv_sat)
water_c0.interpolate(0.005 * mv_sat * pert)
# Approximate answers
# Rate of change is roughly (m_sat - m_v) / 4 so should evaporate everything
mc_true = Function(Vt).interpolate(Constant(0.0))
theta_d_true = Function(Vt).interpolate(theta_d + 0.005 * mv_sat *
pert * Lv_over_cpT)
mv_true = Function(Vt).interpolate(mv_sat * (0.96 + 0.005 * pert))
elif process == "condensation":
water_v0.interpolate(mv_sat * (1.0 + 0.04 * pert))
# Approximate answers -- rate of change is roughly (m_v - m_sat) / 4
mc_true = Function(Vt).interpolate(0.01 * mv_sat * pert)
theta_d_true = Function(Vt).interpolate(theta_d - 0.01 * mv_sat *
pert * Lv_over_cpT)
mv_true = Function(Vt).interpolate(mv_sat * (1.0 + 0.03 * pert))
# Set prognostic variables
theta0.project(theta_d * (1 + water_v0 * parameters.R_v / parameters.R_d))
rho0.interpolate(pressure /
(temperature * parameters.R_d *
(1 + water_v0 * parameters.R_v / parameters.R_d)))
mc_init = Function(Vt).assign(water_c0)
# Have empty problem as only thing is condensation / evaporation
problem = []
physics_list = [Condensation(state)]
# build time stepper
stepper = PrescribedTransport(state, problem, physics_list=physics_list)
stepper.run(t=0, tmax=tmax)
return state, mv_true, mc_true, theta_d_true, mc_init
def setup_hori_limiters(dirname):
# declare grid shape
L = 400.
H = L
ncolumns = int(L / 10.)
nlayers = ncolumns
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
x, z = SpatialCoordinate(mesh)
fieldlist = ['u']
timestepping = TimesteppingParameters(dt=1.0, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname + "/limiting_hori",
dumpfreq=5,
dumplist=['u'],
perturbation_fields=['theta0', 'theta1'])
parameters = CompressibleParameters()
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist)
# make elements
# v is continuous in vertical, h is horizontal
cell = mesh._base_mesh.ufl_cell().cellname()
DG0_element = FiniteElement("DG", cell, 0)
CG1_element = FiniteElement("CG", cell, 1)
DG1_element = FiniteElement("DG", cell, 1)
CG2_element = FiniteElement("CG", cell, 2)
V0_element = TensorProductElement(DG0_element, CG1_element)
V1_element = TensorProductElement(DG1_element, CG2_element)
# spaces
Vpsi = FunctionSpace(mesh, "CG", 2)
VDG1 = FunctionSpace(mesh, "DG", 1)
VCG1 = FunctionSpace(mesh, "CG", 1)
V0 = FunctionSpace(mesh, V0_element)
V1 = FunctionSpace(mesh, V1_element)
V0_brok = FunctionSpace(mesh, BrokenElement(V0.ufl_element()))
V0_spaces = (VDG1, VCG1, V0_brok)
# declare initial fields
u0 = state.fields("u")
theta0 = state.fields("theta0", V0)
theta1 = state.fields("theta1", V1)
# make a gradperp
gradperp = lambda u: as_vector([-u.dx(1), u.dx(0)])
# Isentropic background state
Tsurf = 300.
thetab = Constant(Tsurf)
theta_b1 = Function(V1).interpolate(thetab)
theta_b0 = Function(V0).interpolate(thetab)
# set up bubble
xc = 200.
zc = 200.
rc = 100.
theta_expr = conditional(
sqrt((x - xc)**2.0) < rc,
conditional(sqrt((z - zc)**2.0) < rc, Constant(2.0), Constant(0.0)),
Constant(0.0))
theta_pert1 = Function(V1).interpolate(theta_expr)
theta_pert0 = Function(V0).interpolate(theta_expr)
# set up velocity field
u_max = Constant(10.0)
psi_expr = -u_max * z
psi0 = Function(Vpsi).interpolate(psi_expr)
u0.project(gradperp(psi0))
theta0.interpolate(theta_b0 + theta_pert0)
theta1.interpolate(theta_b1 + theta_pert1)
state.initialise([('u', u0), ('theta1', theta1), ('theta0', theta0)])
state.set_reference_profiles([('theta1', theta_b1), ('theta0', theta_b0)])
# set up advection schemes
thetaeqn1 = EmbeddedDGAdvection(state, V1, equation_form="advective")
thetaeqn0 = EmbeddedDGAdvection(state,
V0,
equation_form="advective",
recovered_spaces=V0_spaces)
# build advection dictionary
advected_fields = []
advected_fields.append(('u', NoAdvection(state, u0, None)))
advected_fields.append(('theta1',
SSPRK3(state,
theta1,
thetaeqn1,
limiter=ThetaLimiter(thetaeqn1))))
advected_fields.append(('theta0',
SSPRK3(state,
theta0,
thetaeqn0,
limiter=VertexBasedLimiter(VDG1))))
# build time stepper
stepper = AdvectionDiffusion(state, advected_fields)
return stepper, 40.0
def setup_recovered_space(dirname):
# declare grid shape, with length L and height H
L = 400.
H = 400.
nlayers = int(H / 20.)
ncolumns = int(L / 20.)
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
x, z = SpatialCoordinate(mesh)
fieldlist = ['u', 'rho']
timestepping = TimesteppingParameters(dt=1.0, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname + "/recovered_space_test",
dumpfreq=5,
dumplist=['u'])
parameters = CompressibleParameters()
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist)
# declare initial fields
u0 = state.fields("u")
# spaces
Vpsi = FunctionSpace(mesh, "CG", 2)
VDG0 = FunctionSpace(mesh, "DG", 0)
VDG1 = FunctionSpace(mesh, "DG", 1)
VCG1 = FunctionSpace(mesh, "CG", 1)
# set up tracer field
tracer0 = state.fields("tracer", VDG0)
# make a gradperp
gradperp = lambda u: as_vector([-u.dx(1), u.dx(0)])
# set up bubble
xc = 200.
zc = 200.
rc = 100.
tracer0.interpolate(
conditional(
sqrt((x - xc)**2.0) < rc,
conditional(
sqrt((z - zc)**2.0) < rc, Constant(0.2), Constant(0.0)),
Constant(0.0)))
# set up velocity field
u_max = Constant(10.0)
psi_expr = -u_max * z
psi0 = Function(Vpsi).interpolate(psi_expr)
u0.project(gradperp(psi0))
state.initialise([('u', u0), ('tracer', tracer0)])
# set up advection schemes
recovered_opts = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=VDG0,
boundary_method=Boundary_Method.dynamics)
tracereqn = EmbeddedDGAdvection(state,
VDG0,
equation_form="continuity",
options=recovered_opts)
# build advection dictionary
advected_fields = []
advected_fields.append(('tracer', SSPRK3(state, tracer0, tracereqn)))
# build time stepper
stepper = AdvectionDiffusion(state, advected_fields)
return stepper, 100.0
def setup_balance(dirname):
# set up grid and time stepping parameters
dt = 1.
tmax = 5.
deltax = 400
L = 2000.
H = 10000.
nlayers = int(H / deltax)
ncolumns = int(L / deltax)
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname + '/dry_balance',
dumpfreq=10,
dumplist=['u'])
parameters = CompressibleParameters()
diagnostics = Diagnostics(*fieldlist)
diagnostic_fields = []
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist,
diagnostic_fields=diagnostic_fields)
# Initial conditions
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
# Isentropic background state
Tsurf = Constant(300.)
theta0.interpolate(Tsurf)
# Calculate hydrostatic Pi
compressible_hydrostatic_balance(state, theta0, rho0, solve_for_rho=True)
state.initialise([('u', u0), ('rho', rho0), ('theta', theta0)])
state.set_reference_profiles([('rho', rho0), ('theta', theta0)])
# Set up advection schemes
ueqn = VectorInvariant(state, Vu)
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
thetaeqn = EmbeddedDGAdvection(state,
Vt,
equation_form="advective",
options=EmbeddedDGOptions())
advected_fields = [("u", ThetaMethod(state, u0, ueqn)),
("rho", SSPRK3(state, rho0, rhoeqn)),
("theta", SSPRK3(state, theta0, thetaeqn))]
# Set up linear solver
linear_solver = CompressibleSolver(state)
# Set up forcing
compressible_forcing = CompressibleForcing(state)
# build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver,
compressible_forcing)
return stepper, tmax
def do_simulation_loop(N,
variable_parameters,
simulation_parameters,
diagnostics=None,
fields_to_output=None,
expected_u=False):
"""
A recursive strategy for setting up a variable number of for loops
for the experiment.
:arg N: the level of loop.
:arg variable_parameters: an OrderedDict of the parameters to be varied.
:arg simulation_parameters: a dict storing the parameters for that simulation.
:arg diagnostics: a list of diagnostic values to be output.
:arg fields_to_output: a list of fields to be output.
"""
# we do the loop in reverse order to get the resolution loop on the outside
M = len(variable_parameters)
key = list(variable_parameters.items())[M - N][0]
have_setup = False
# we must turn the ordered dict into a list to iterate through it
for index, value in enumerate(
list(variable_parameters.items())[M - N][1][1]):
simulation_parameters[key] = (index, value)
# make mesh if loop is a resolution loop
if key == 'resolution':
mesh = PeriodicIntervalMesh(value, simulation_parameters['Ld'][-1])
simulation_parameters['mesh'] = (mesh, )
# do recursion if we aren't finished yet
if N > 1:
do_simulation_loop(N - 1,
variable_parameters,
simulation_parameters,
diagnostics=diagnostics,
fields_to_output=fields_to_output,
expected_u=expected_u)
# finally do simulation
elif N == 1:
if expected_u:
this_u = simulation(simulation_parameters,
diagnostic_values=diagnostics,
fields_to_output=fields_to_output,
expected_u=True)
if have_setup:
Eu.assign(counter * Eu + this_u)
counter.assign(counter + 1)
Eu.assign(Eu / counter)
else:
scheme = simulation_parameters['scheme'][-1]
mesh = simulation_parameters['mesh'][-1]
prognostic_variables = PrognosticVariables(scheme, mesh)
Eu = Function(prognostic_variables.Vu,
name='expected u').assign(this_u)
counter = Constant(1.0)
have_setup = True
else:
simulation(simulation_parameters,
diagnostic_values=diagnostics,
fields_to_output=fields_to_output)
if expected_u:
expected_u_file = File(simulation_parameters['dirname'][-1] +
'/expected_u.pvd')
expected_u_file.write(Eu)
def setup_saturated(dirname, recovered):
# set up grid and time stepping parameters
dt = 1.
tmax = 3.
deltax = 400.
L = 2000.
H = 10000.
nlayers = int(H / deltax)
ncolumns = int(L / deltax)
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
# option to easily change between recovered and not if necessary
# default should be to use lowest order set of spaces
degree = 0 if recovered else 1
output = OutputParameters(dirname=dirname + '/saturated_balance',
dumpfreq=1,
dumplist=['u'])
parameters = CompressibleParameters()
diagnostic_fields = [Theta_e()]
state = State(mesh,
dt=dt,
output=output,
parameters=parameters,
diagnostic_fields=diagnostic_fields)
tracers = [WaterVapour(), CloudWater()]
if recovered:
u_transport_option = "vector_advection_form"
else:
u_transport_option = "vector_invariant_form"
eqns = CompressibleEulerEquations(state,
"CG",
degree,
u_transport_option=u_transport_option,
active_tracers=tracers)
# Initial conditions
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
water_v0 = state.fields("vapour_mixing_ratio")
water_c0 = state.fields("cloud_liquid_mixing_ratio")
moisture = ['vapour_mixing_ratio', 'cloud_liquid_mixing_ratio']
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
# Isentropic background state
Tsurf = Constant(300.)
total_water = Constant(0.02)
theta_e = Function(Vt).interpolate(Tsurf)
water_t = Function(Vt).interpolate(total_water)
# Calculate hydrostatic exner
saturated_hydrostatic_balance(state, theta_e, water_t)
water_c0.assign(water_t - water_v0)
state.set_reference_profiles([('rho', rho0), ('theta', theta0)])
# Set up transport schemes
if recovered:
VDG1 = state.spaces("DG1_equispaced")
VCG1 = FunctionSpace(mesh, "CG", 1)
Vt_brok = FunctionSpace(mesh, BrokenElement(Vt.ufl_element()))
Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element())
Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1)
u_opts = RecoveredOptions(embedding_space=Vu_DG1,
recovered_space=Vu_CG1,
broken_space=Vu,
boundary_method=Boundary_Method.dynamics)
rho_opts = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=Vr,
boundary_method=Boundary_Method.dynamics)
theta_opts = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=Vt_brok)
wv_opts = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=Vt_brok)
wc_opts = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=Vt_brok)
else:
rho_opts = None
theta_opts = EmbeddedDGOptions()
wv_opts = EmbeddedDGOptions()
wc_opts = EmbeddedDGOptions()
transported_fields = [
SSPRK3(state, 'rho', options=rho_opts),
SSPRK3(state, 'theta', options=theta_opts),
SSPRK3(state, 'vapour_mixing_ratio', options=wv_opts),
SSPRK3(state, 'cloud_liquid_mixing_ratio', options=wc_opts)
]
if recovered:
transported_fields.append(SSPRK3(state, 'u', options=u_opts))
else:
transported_fields.append(ImplicitMidpoint(state, 'u'))
linear_solver = CompressibleSolver(state, eqns, moisture=moisture)
# add physics
physics_list = [Condensation(state)]
# build time stepper
stepper = CrankNicolson(state,
eqns,
transported_fields,
linear_solver=linear_solver,
physics_list=physics_list)
return stepper, tmax
def setup_limiters(dirname, direction, grid_params, ic_params):
# declare grid shape
L, H, ncolumns, nlayers = grid_params
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
x, z = SpatialCoordinate(mesh)
fieldlist = ['u']
timestepping = TimesteppingParameters(dt=1.0)
output = OutputParameters(dirname=dirname,
dumpfreq=5,
dumplist=['u'],
perturbation_fields=['theta0', 'theta1'])
parameters = CompressibleParameters()
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist)
# make elements
# v is continuous in vertical, h is horizontal
cell = mesh._base_mesh.ufl_cell().cellname()
DG0_element = FiniteElement("DG", cell, 0)
CG1_element = FiniteElement("CG", cell, 1)
DG1_element = FiniteElement("DG", cell, 1)
CG2_element = FiniteElement("CG", cell, 2)
V0_element = TensorProductElement(DG0_element, CG1_element)
V1_element = TensorProductElement(DG1_element, CG2_element)
# spaces
VDG1 = FunctionSpace(mesh, "DG", 1)
VCG1 = FunctionSpace(mesh, "CG", 1)
V0 = FunctionSpace(mesh, V0_element)
V1 = FunctionSpace(mesh, V1_element)
V0_brok = FunctionSpace(mesh, BrokenElement(V0.ufl_element()))
# declare initial fields
u0 = state.fields("u")
theta0 = state.fields("theta0", V0)
theta1 = state.fields("theta1", V1)
Tsurf, xc, zc, rc, u_max = ic_params
# Isentropic background state
thetab = Constant(Tsurf)
theta_b1 = Function(V1).interpolate(thetab)
theta_b0 = Function(V0).interpolate(thetab)
# set up bubble
theta_expr = conditional(
sqrt((x - xc)**2.0) < rc,
conditional(sqrt((z - zc)**2.0) < rc, Constant(2.0), Constant(0.0)),
Constant(0.0))
theta_pert1 = Function(V1).interpolate(theta_expr)
theta_pert0 = Function(V0).interpolate(theta_expr)
if direction == "horizontal":
Vpsi = FunctionSpace(mesh, "CG", 2)
psi_expr = -u_max * z
psi0 = Function(Vpsi).interpolate(psi_expr)
gradperp = lambda u: as_vector([-u.dx(1), u.dx(0)])
u0.project(gradperp(psi0))
elif direction == "vertical":
u0.project(as_vector([0, -u_max]))
theta0.interpolate(theta_b0 + theta_pert0)
theta1.interpolate(theta_b1 + theta_pert1)
state.initialise([('u', u0), ('theta1', theta1), ('theta0', theta0)])
state.set_reference_profiles([('theta1', theta_b1), ('theta0', theta_b0)])
# set up advection schemes
dg_opts = EmbeddedDGOptions()
recovered_opts = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=V0_brok)
thetaeqn1 = EmbeddedDGAdvection(state,
V1,
equation_form="advective",
options=dg_opts)
thetaeqn0 = EmbeddedDGAdvection(state,
V0,
equation_form="advective",
options=recovered_opts)
# build advection dictionary
advected_fields = []
advected_fields.append(
('theta1', SSPRK3(state, theta1, thetaeqn1, limiter=ThetaLimiter(V1))))
advected_fields.append(('theta0',
SSPRK3(state,
theta0,
thetaeqn0,
limiter=VertexBasedLimiter(VDG1))))
# build time stepper
stepper = AdvectionDiffusion(state, advected_fields)
return stepper, 40.0
def run_incompressible(tmpdir):
dt = 6.0
tmax = 2*dt
nlayers = 10 # horizontal layers
ncols = 10 # number of columns
Lx = 1000.0
Lz = 1000.0
m = PeriodicIntervalMesh(ncols, Lx)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=Lz/nlayers)
output = OutputParameters(dirname=tmpdir+"/incompressible",
dumpfreq=2, chkptfreq=2)
parameters = CompressibleParameters()
state = State(mesh,
dt=dt,
output=output,
parameters=parameters)
eqns = IncompressibleBoussinesqEquations(state, "CG", 1)
# Initial conditions
p0 = state.fields("p")
b0 = state.fields("b")
# z.grad(bref) = N**2
x, z = SpatialCoordinate(mesh)
N = parameters.N
bref = z*(N**2)
b_b = Function(b0.function_space()).interpolate(bref)
incompressible_hydrostatic_balance(state, b_b, p0)
state.initialise([('p', p0),
('b', b_b)])
# Add perturbation
r = sqrt((x-Lx/2)**2 + (z-Lz/2)**2)
b_pert = 0.1*exp(-(r/(Lx/5)**2))
b0.interpolate(b_b + b_pert)
# Set up transport schemes
b_opts = SUPGOptions()
transported_fields = [ImplicitMidpoint(state, "u"),
SSPRK3(state, "b", options=b_opts)]
# Set up linear solver for the timestepping scheme
linear_solver = IncompressibleSolver(state, eqns)
# build time stepper
stepper = CrankNicolson(state, eqns, transported_fields,
linear_solver=linear_solver)
# Run
stepper.run(t=0, tmax=tmax)
# State for checking checkpoints
checkpoint_name = 'incompressible_chkpt'
new_path = join(abspath(dirname(__file__)), '..', f'data/{checkpoint_name}')
check_output = OutputParameters(dirname=tmpdir+"/incompressible",
checkpoint_pickup_filename=new_path)
check_state = State(mesh, dt=dt, output=check_output)
check_eqn = IncompressibleBoussinesqEquations(check_state, "CG", 1)
check_stepper = CrankNicolson(check_state, check_eqn, [])
check_stepper.run(t=0, tmax=0, pickup=True)
return state, check_state
def setup_2d_recovery(dirname):
L = 100.
H = 100.
deltax = L / 5.
deltay = H / 5.
nlayers = int(H / deltay)
ncolumns = int(L / deltax)
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
x, y = SpatialCoordinate(mesh)
# horizontal base spaces
cell = mesh._base_mesh.ufl_cell().cellname()
u_hori = FiniteElement("CG", cell, 1)
w_hori = FiniteElement("DG", cell, 0)
# vertical base spaces
u_vert = FiniteElement("DG", interval, 0)
w_vert = FiniteElement("CG", interval, 1)
# build elements
u_element = HDiv(TensorProductElement(u_hori, u_vert))
w_element = HDiv(TensorProductElement(w_hori, w_vert))
theta_element = TensorProductElement(w_hori, w_vert)
v_element = u_element + w_element
# spaces
VDG0 = FunctionSpace(mesh, "DG", 0)
VCG1 = FunctionSpace(mesh, "CG", 1)
VDG1 = FunctionSpace(mesh, "DG", 1)
Vt = FunctionSpace(mesh, theta_element)
Vt_brok = FunctionSpace(mesh, BrokenElement(theta_element))
Vu = FunctionSpace(mesh, v_element)
VuCG1 = VectorFunctionSpace(mesh, "CG", 1)
VuDG1 = VectorFunctionSpace(mesh, "DG", 1)
# set up initial conditions
np.random.seed(0)
expr = np.random.randn() + np.random.randn() * y
# our actual theta and rho and v
rho_CG1_true = Function(VCG1).interpolate(expr)
theta_CG1_true = Function(VCG1).interpolate(expr)
v_CG1_true = Function(VuCG1).interpolate(as_vector([expr, expr]))
rho_Vt_true = Function(Vt).interpolate(expr)
# make the initial fields by projecting expressions into the lowest order spaces
rho_DG0 = Function(VDG0).interpolate(expr)
rho_CG1 = Function(VCG1)
theta_Vt = Function(Vt).interpolate(expr)
theta_CG1 = Function(VCG1)
v_Vu = Function(Vu).project(as_vector([expr, expr]))
v_CG1 = Function(VuCG1)
rho_Vt = Function(Vt)
# make the recoverers and do the recovery
rho_recoverer = Recoverer(rho_DG0,
rho_CG1,
VDG=VDG1,
boundary_method=Boundary_Method.dynamics)
theta_recoverer = Recoverer(theta_Vt,
theta_CG1,
VDG=VDG1,
boundary_method=Boundary_Method.dynamics)
v_recoverer = Recoverer(v_Vu,
v_CG1,
VDG=VuDG1,
boundary_method=Boundary_Method.dynamics)
rho_Vt_recoverer = Recoverer(rho_DG0,
rho_Vt,
VDG=Vt_brok,
boundary_method=Boundary_Method.physics)
rho_recoverer.project()
theta_recoverer.project()
v_recoverer.project()
rho_Vt_recoverer.project()
rho_diff = errornorm(rho_CG1, rho_CG1_true) / norm(rho_CG1_true)
theta_diff = errornorm(theta_CG1, theta_CG1_true) / norm(theta_CG1_true)
v_diff = errornorm(v_CG1, v_CG1_true) / norm(v_CG1_true)
rho_Vt_diff = errornorm(rho_Vt, rho_Vt_true) / norm(rho_Vt_true)
return (rho_diff, theta_diff, v_diff, rho_Vt_diff)
def run_dry_compressible(tmpdir):
dt = 6.0
tmax = 2 * dt
nlayers = 10 # horizontal layers
ncols = 10 # number of columns
Lx = 1000.0
Lz = 1000.0
m = PeriodicIntervalMesh(ncols, Lx)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=Lz / nlayers)
output = OutputParameters(dirname=tmpdir + "/dry_compressible",
dumpfreq=2,
chkptfreq=2)
parameters = CompressibleParameters()
R_d = parameters.R_d
g = parameters.g
state = State(mesh, dt=dt, output=output, parameters=parameters)
eqn = CompressibleEulerEquations(state, "CG", 1)
# Initial conditions
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# Approximate hydrostatic balance
x, z = SpatialCoordinate(mesh)
T = Constant(300.0)
zH = R_d * T / g
p = Constant(100000.0) * exp(-z / zH)
theta0.interpolate(tde.theta(parameters, T, p))
rho0.interpolate(p / (R_d * T))
state.set_reference_profiles([('rho', rho0), ('theta', theta0)])
# Add perturbation
r = sqrt((x - Lx / 2)**2 + (z - Lz / 2)**2)
theta_pert = 1.0 * exp(-(r / (Lx / 5))**2)
theta0.interpolate(theta0 + theta_pert)
# Set up transport schemes
transported_fields = [
ImplicitMidpoint(state, "u"),
SSPRK3(state, "rho"),
SSPRK3(state, "theta")
]
# Set up linear solver for the timestepping scheme
linear_solver = CompressibleSolver(state, eqn)
# build time stepper
stepper = CrankNicolson(state,
eqn,
transported_fields,
linear_solver=linear_solver)
# Run
stepper.run(t=0, tmax=tmax)
# State for checking checkpoints
checkpoint_name = 'dry_compressible_chkpt'
new_path = join(abspath(dirname(__file__)), '..',
f'data/{checkpoint_name}')
check_output = OutputParameters(dirname=tmpdir + "/dry_compressible",
checkpoint_pickup_filename=new_path)
check_state = State(mesh,
dt=dt,
output=check_output,
parameters=parameters)
check_eqn = CompressibleEulerEquations(check_state, "CG", 1)
check_stepper = CrankNicolson(check_state, check_eqn, [])
check_stepper.run(t=0, tmax=0, pickup=True)
return state, check_state
def setup_limiters(dirname):
dt = 0.01
Ld = 1.
tmax = 0.2
rotations = 0.1
m = PeriodicIntervalMesh(20, Ld)
mesh = ExtrudedMesh(m, layers=20, layer_height=(Ld / 20))
output = OutputParameters(
dirname=dirname,
dumpfreq=1,
dumplist=['u', 'chemical', 'moisture_higher', 'moisture_lower'])
parameters = CompressibleParameters()
timestepping = TimesteppingParameters(dt=dt, maxk=4, maxi=1)
fieldlist = [
'u', 'rho', 'theta', 'chemical', 'moisture_higher', 'moisture_lower'
]
diagnostic_fields = []
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist,
diagnostic_fields=diagnostic_fields)
x, z = SpatialCoordinate(mesh)
Vr = state.spaces("DG")
Vt = state.spaces("HDiv_v")
Vpsi = FunctionSpace(mesh, "CG", 2)
cell = mesh._base_mesh.ufl_cell().cellname()
DG0_element = FiniteElement("DG", cell, 0)
CG1_element = FiniteElement("CG", interval, 1)
Vt0_element = TensorProductElement(DG0_element, CG1_element)
Vt0 = FunctionSpace(mesh, Vt0_element)
Vt0_brok = FunctionSpace(mesh, BrokenElement(Vt0_element))
VCG1 = FunctionSpace(mesh, "CG", 1)
u = state.fields("u", dump=True)
chemical = state.fields("chemical", Vr, dump=True)
moisture_higher = state.fields("moisture_higher", Vt, dump=True)
moisture_lower = state.fields("moisture_lower", Vt0, dump=True)
x_lower = 2 * Ld / 5
x_upper = 3 * Ld / 5
z_lower = 6 * Ld / 10
z_upper = 8 * Ld / 10
bubble_expr_1 = conditional(
x > x_lower,
conditional(
x < x_upper,
conditional(z > z_lower, conditional(z < z_upper, 1.0, 0.0), 0.0),
0.0), 0.0)
bubble_expr_2 = conditional(
x > z_lower,
conditional(
x < z_upper,
conditional(z > x_lower, conditional(z < x_upper, 1.0, 0.0), 0.0),
0.0), 0.0)
chemical.assign(1.0)
moisture_higher.assign(280.)
chem_pert_1 = Function(Vr).interpolate(bubble_expr_1)
chem_pert_2 = Function(Vr).interpolate(bubble_expr_2)
moist_h_pert_1 = Function(Vt).interpolate(bubble_expr_1)
moist_h_pert_2 = Function(Vt).interpolate(bubble_expr_2)
moist_l_pert_1 = Function(Vt0).interpolate(bubble_expr_1)
moist_l_pert_2 = Function(Vt0).interpolate(bubble_expr_2)
chemical.assign(chemical + chem_pert_1 + chem_pert_2)
moisture_higher.assign(moisture_higher + moist_h_pert_1 + moist_h_pert_2)
moisture_lower.assign(moisture_lower + moist_l_pert_1 + moist_l_pert_2)
# set up solid body rotation for advection
# we do this slightly complicated stream function to make the velocity 0 at edges
# thus we avoid odd effects at boundaries
xc = Ld / 2
zc = Ld / 2
r = sqrt((x - xc)**2 + (z - zc)**2)
omega = rotations * 2 * pi / tmax
r_out = 9 * Ld / 20
r_in = 2 * Ld / 5
A = omega * r_in / (2 * (r_in - r_out))
B = -omega * r_in * r_out / (r_in - r_out)
C = omega * r_in**2 * r_out / (r_in - r_out) / 2
psi_expr = conditional(
r < r_in, omega * r**2 / 2,
conditional(r < r_out, A * r**2 + B * r + C,
A * r_out**2 + B * r_out + C))
psi = Function(Vpsi).interpolate(psi_expr)
gradperp = lambda v: as_vector([-v.dx(1), v.dx(0)])
u.project(gradperp(psi))
state.initialise([('u', u), ('chemical', chemical),
('moisture_higher', moisture_higher),
('moisture_lower', moisture_lower)])
# set up advection schemes
dg_opts = EmbeddedDGOptions()
recovered_opts = RecoveredOptions(embedding_space=Vr,
recovered_space=VCG1,
broken_space=Vt0_brok,
boundary_method=Boundary_Method.dynamics)
chemeqn = AdvectionEquation(state, Vr, equation_form="advective")
moisteqn_higher = EmbeddedDGAdvection(state,
Vt,
equation_form="advective",
options=dg_opts)
moisteqn_lower = EmbeddedDGAdvection(state,
Vt0,
equation_form="advective",
options=recovered_opts)
# build advection dictionary
advected_fields = []
advected_fields.append(('chemical',
SSPRK3(state,
chemical,
chemeqn,
limiter=VertexBasedLimiter(Vr))))
advected_fields.append(('moisture_higher',
SSPRK3(state,
moisture_higher,
moisteqn_higher,
limiter=ThetaLimiter(Vt))))
advected_fields.append(('moisture_lower',
SSPRK3(state,
moisture_lower,
moisteqn_lower,
limiter=VertexBasedLimiter(Vr))))
# build time stepper
stepper = AdvectionDiffusion(state, advected_fields)
return stepper, tmax
def setup_condens(dirname):
# declare grid shape, with length L and height H
L = 1000.
H = 1000.
nlayers = int(H / 100.)
ncolumns = int(L / 100.)
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
x = SpatialCoordinate(mesh)
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=1.0, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname + "/condens",
dumpfreq=1,
dumplist=['u'],
perturbation_fields=['theta', 'rho'])
parameters = CompressibleParameters()
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist,
diagnostic_fields=[Sum('water_v', 'water_c')])
# declare initial fields
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vpsi = FunctionSpace(mesh, "CG", 2)
Vt = theta0.function_space()
Vr = rho0.function_space()
# make a gradperp
gradperp = lambda u: as_vector([-u.dx(1), u.dx(0)])
# declare tracer field and a background field
water_v0 = state.fields("water_v", Vt)
water_c0 = state.fields("water_c", Vt)
# Isentropic background state
Tsurf = Constant(300.)
theta_b = Function(Vt).interpolate(Tsurf)
rho_b = Function(Vr)
# Calculate initial rho
compressible_hydrostatic_balance(state, theta_b, rho_b, solve_for_rho=True)
# set up water_v
xc = 500.
zc = 350.
rc = 250.
r = sqrt((x[0] - xc)**2 + (x[1] - zc)**2)
w_expr = conditional(r > rc, 0., 0.25 * (1. + cos((pi / rc) * r)))
# set up velocity field
u_max = 10.0
psi_expr = ((-u_max * L / pi) * sin(2 * pi * x[0] / L) *
sin(pi * x[1] / L))
psi0 = Function(Vpsi).interpolate(psi_expr)
u0.project(gradperp(psi0))
theta0.interpolate(theta_b)
rho0.interpolate(rho_b)
water_v0.interpolate(w_expr)
state.initialise([('u', u0), ('rho', rho0), ('theta', theta0),
('water_v', water_v0), ('water_c', water_c0)])
state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
# set up advection schemes
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
thetaeqn = SUPGAdvection(state,
Vt,
supg_params={"dg_direction": "horizontal"},
equation_form="advective")
# build advection dictionary
advected_fields = []
advected_fields.append(("u", NoAdvection(state, u0, None)))
advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn)))
advected_fields.append(("theta", SSPRK3(state, theta0, thetaeqn)))
advected_fields.append(("water_v", SSPRK3(state, water_v0, thetaeqn)))
advected_fields.append(("water_c", SSPRK3(state, water_c0, thetaeqn)))
physics_list = [Condensation(state)]
# build time stepper
stepper = AdvectionDiffusion(state,
advected_fields,
physics_list=physics_list)
return stepper, 5.0
def setup_sk(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)
# Set up points for output at the centre of the domain, edges and corners.
# The point at x=L*(13.0/30) is in the halo region for a two-way MPI decomposition
points_x = [0.0, L * (13.0 / 30), L / 2.0, L]
points_z = [0.0, H / 2.0, H]
points = np.array([p for p in itertools.product(points_x, points_z)])
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt)
output = OutputParameters(dirname=dirname + "/sk_nonlinear",
dumplist=['u'],
dumpfreq=5,
log_level=INFO,
point_data=[('rho', points), ('u', points)])
parameters = CompressibleParameters()
diagnostic_fields = [CourantNumber()]
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist,
diagnostic_fields=diagnostic_fields)
# Initial conditions
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
# 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)
x, z = SpatialCoordinate(mesh)
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)
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)
u0.project(as_vector([20.0, 0.0]))
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)
advected_fields = []
advected_fields.append(("u", ThetaMethod(state, u0, ueqn)))
advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn)))
advected_fields.append(("theta", SSPRK3(state, theta0, thetaeqn)))
# Set up linear solver
linear_solver = CompressibleSolver(state)
# Set up forcing
compressible_forcing = CompressibleForcing(state)
# build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver,
compressible_forcing)
return stepper, 2 * dt
def experiment(code,
Ld,
tmax,
resolutions=[],
dts=[],
sigmas=[],
seeds=[],
schemes=[],
timesteppings=[],
alphasq=[],
c0=[],
gamma=[],
ics=[],
num_Xis=[],
Xi_family=[],
diagnostics=None,
fields_to_output=None,
ndump=-1,
field_ndump=-1,
nXi_updates=1,
allow_fail=False,
t_kick=[],
sigma_kick=0.0,
smooth_t=None,
peakon_equations=False,
only_peakons=False,
true_peakon_data=None,
true_mean_peakon_data=None,
peakon_speed=None,
expected_u=False,
periodic=False,
peak_width=1.0,
peakon_method='ito_euler',
fixed_dW=None):
# set up dumping
dirname = 'results/' + code
if path.exists(dirname):
raise IOError("results directory '%s' already exists" % dirname)
else:
makedirs(dirname)
expmt_dict = {
'dt': (float, dts),
'resolution': (float, resolutions),
'seed': (int, seeds),
'sigma': (float, sigmas),
'scheme': ('S1', schemes),
'timestepping': ('S1', timesteppings),
'alphasq': (float, alphasq),
'c0': (float, c0),
'gamma': (float, gamma),
'ic': ('S1', ics),
'num_Xis': (int, num_Xis),
'Xi_family': ('S1', Xi_family),
}
# turns expmt_dict values into lists so they can be iterated over
for key, value in expmt_dict.items():
# need to not accidentally accept strings
if not isinstance(value, str):
try:
iter(value[1])
except TypeError:
expmt_dict[key] = (value[0], [value[1]])
else:
expmt_dict[key] = (value[0], [value[1]])
# we need a dictionary of parameters to form the loops over
# also a dictionary of fixed parameters for all simulations
# finally a dictionary of all parameters for individual simulations
# the first should be an ordered dict
variable_parameters = OrderedDict()
fixed_parameters = {}
simulation_parameters = {}
# need resolutions to be first in variable parameters dict (to loop over first)
if len(expmt_dict['resolution'][1]) > 1:
variable_parameters['resolution'] = None
# here we separate parameters into variable or fixed
# this is done by asking if they are iterable, and being careful to ignore strings
for key, value in expmt_dict.items():
if not isinstance(value[1], str) and len(value[1]) > 1:
variable_parameters[key] = value
simulation_parameters[key] = None
elif isinstance(value[1], str):
fixed_parameters[key] = value
simulation_parameters[key] = (value[1], )
elif len(value[1]) == 1:
fixed_parameters[key] = value
simulation_parameters[key] = (value[1][0], )
# Stuff where there can be only one value that is fixed for all experiments
for key, value in zip([
'dirname', 'Ld', 'tmax', 'ndump', 'field_ndump', 'nXi_updates',
'allow_fail', 'smooth_t', 'peakon_equations', 'only_peakons',
'true_peakon_data', 'true_mean_peakon_data', 'peakon_speed',
'periodic', 'fixed_dW', 'peak_width', 'peakon_method'
], [
dirname, Ld, tmax, ndump, field_ndump, nXi_updates, allow_fail,
smooth_t, peakon_equations, only_peakons, true_peakon_data,
true_mean_peakon_data, peakon_speed, periodic, fixed_dW,
peak_width, peakon_method
]):
simulation_parameters[key] = (value, )
output_arguments = ('time', ) + tuple(variable_parameters.keys())
list_of_peakon_diagnostics = [
'peakon_loc', 'peakon_min_du', 'peakon_max_du', 'peakon_min_du_loc',
'peakon_max_du_loc', 'peakon_max_u', 'peakon_mu', 'peakon_nu'
]
for i, dt in enumerate(expmt_dict['dt'][1]):
# make data file
if len(expmt_dict['dt'][1]) > 1:
file_name = dirname + '/data_dt' + str(i) + '.nc'
else:
file_name = dirname + '/data.nc'
simulation_parameters['file_name'] = (file_name, )
data_file = Dataset(file_name, 'w')
# create dimensions and variables
data_file.createDimension('time', None)
data_file.createVariable('time', float, ('time', ))
if 'resolution' in variable_parameters.keys():
data_file.createDimension(
'resolution', len(variable_parameters['resolution'][1]))
data_file.createVariable('deltax', float, ('resolution', ))
for j, res in enumerate(variable_parameters['resolution'][1]):
data_file['deltax'][
j:j + 1] = Ld / variable_parameters['resolution'][1][j]
# create a fixed x dimension for a and b fields
data_file.createDimension('x', 100)
data_file.createVariable('x', float, ('x', ))
data_file['x'][:] = np.linspace(0, Ld, num=100, endpoint=False)
simulation_parameters['store_coordinates'] = [False]
else:
# if the resolution is not varying, we can use the real resolution
num_points = simulation_parameters['resolution'][-1]
data_file.createDimension('x', num_points)
data_file.createVariable('x', float, ('x', ))
simulation_parameters['store_coordinates'] = [True]
for key, values in variable_parameters.items():
if key != 'resolution':
data_file.createDimension(key, len(values[1]))
data_file.createVariable(key, values[0], (key, ))
for j, value in enumerate(values[1]):
data_file[key][j:j + 1] = value
# create diagnostic variables, with dimensions dependent upon parameters
if diagnostics is not None:
for output in diagnostics:
if output == 'mu':
for i in range(4):
data_file.createVariable(output + '_' + str(i), float,
output_arguments)
data_file.createVariable(
'alt_' + output + '_' + str(i), float,
output_arguments)
elif output in ('a', 'b', 'u_field'):
data_file.createVariable(output, float,
output_arguments + ('x', ))
elif output == 'peakon_suite':
for peakon_diag in list_of_peakon_diagnostics:
data_file.createVariable(peakon_diag, float,
output_arguments)
else:
data_file.createVariable(output, float, output_arguments)
if peakon_equations:
data_file.createVariable('p', float, output_arguments)
data_file.createVariable('q', float, output_arguments)
# create diagnostics for wallclock time and failure time
data_file.createDimension('wallclock_times', 2)
data_file.createVariable('wallclock_time', float,
('wallclock_times', ) +
tuple(variable_parameters.keys()))
data_file.createVariable('failed_time', float,
tuple(variable_parameters.keys()))
data_file.close()
# we want to do a variable number of for loops so use recursive strategy
N = len(variable_parameters)
if N > 0:
# measure length of loops
# raise error if it's too long
num_sims = np.prod(
[len(value[1]) for value in variable_parameters.values()])
if num_sims > 1000:
raise ValueError(
'You are asking to do %d simulations! Please reduce your loops.'
% num_sims)
# do the simulation loop
if len(expmt_dict['resolution'][1]) == 1:
resolution = simulation_parameters['resolution'][-1]
mesh = PeriodicIntervalMesh(resolution, Ld)
simulation_parameters['mesh'] = (mesh, )
do_simulation_loop(N,
variable_parameters,
simulation_parameters,
diagnostics=diagnostics,
fields_to_output=fields_to_output,
expected_u=expected_u)
else:
resolution = simulation_parameters['resolution'][-1]
mesh = PeriodicIntervalMesh(resolution, Ld)
simulation_parameters['mesh'] = (mesh, )
simulation(simulation_parameters,
diagnostic_values=diagnostics,
fields_to_output=fields_to_output)
def setup_fallout(dirname):
# declare grid shape, with length L and height H
L = 10.
H = 10.
nlayers = 10
ncolumns = 10
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
x = SpatialCoordinate(mesh)
fieldlist = ['u', 'rho', 'theta', 'rain']
timestepping = TimesteppingParameters(dt=0.1, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname + "/fallout",
dumpfreq=10,
dumplist=['rain'])
parameters = CompressibleParameters()
diagnostic_fields = [Precipitation()]
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist,
diagnostic_fields=diagnostic_fields)
# declare initial fields
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vt = theta0.function_space()
# declare tracer field and a background field
rain0 = state.fields("rain", Vt)
# set up rain
xc = L / 2
zc = H / 2
rc = H / 4
r = sqrt((x[0] - xc)**2 + (x[1] - zc)**2)
rain_expr = conditional(r > rc, 0., 1e-3 * (cos(pi * r / (rc * 2)))**2)
rain0.interpolate(rain_expr)
rho0.assign(1.0)
state.initialise([('u', u0), ('rho', rho0), ('rain', rain0)])
# build advection dictionary
advected_fields = []
advected_fields.append(("u", NoAdvection(state, u0, None)))
advected_fields.append(("rho", NoAdvection(state, rho0, None)))
advected_fields.append(("rain", NoAdvection(state, rain0, None)))
physics_list = [Fallout(state)]
# build time stepper
stepper = AdvectionDiffusion(state,
advected_fields,
physics_list=physics_list)
return stepper, 10.0
dt = 10.
if '--running-tests' in sys.argv:
tmax = dt
else:
tmax = 3600.
if '--hybridization' in sys.argv:
hybridization = True
else:
hybridization = False
nlayers = 50 # horizontal layers
columns = 50 # number of columns
L = 3.0e5
m = PeriodicIntervalMesh(columns, L)
# build volume mesh
H = 1.0e4 # Height position of the model top
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt)
dirname = 'sk_linear'
if hybridization:
dirname += '_hybridization'
output = OutputParameters(dirname=dirname,
dumplist=['u'],
perturbation_fields=['theta', 'rho'])
# ---------------
# 1D mesh support ncells, Length (left coords), right coords (opt.)
# m= IntervalMesh(nx,-L,L)
m = IntervalMesh(nx, L)
xIM = m.coordinates.dat.data
# -> UnitIntervalMesh
# -------------------
m = UnitIntervalMesh(nx)
xUIM = m.coordinates.dat.data
# -> PeriodicInternalMesh
# -----------------------
m = PeriodicIntervalMesh(nx, L)
xPIM = m.coordinates.dat.data
y = np.ones([nx+1, ])
plt.plot(xIM, 0*y, color="black", marker="o", label="IntervalMesh")
plt.plot(xUIM, 1*y, color="red", marker="o", label="UnitIntervalMesh")
plt.plot(xPIM, 2*np.ones(np.shape(xPIM)), color="blue",
marker="o", label="PeriodicInternalMesh")
plt.legend()
plt.savefig("plots/oneDimensionMesh.png")
# TwoDimensional Meshes
# -> RectangleMesh
# ----------------
# support quadrilateral input (True or False)