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
tmax = 15. * 60.
ndumps = 4
L = 51200.
# build volume mesh
H = 6400. # Height position of the model top
for delta, dt in res_dt.items():
dirname = "straka_dx%s_dt%s" % (delta, dt)
nlayers = int(H / delta) # horizontal layers
columns = int(L / delta) # number of columns
m = PeriodicIntervalMesh(columns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
dumpfreq = int(tmax / (ndumps * dt))
output = OutputParameters(dirname=dirname,
dumpfreq=dumpfreq,
dumplist=['u'],
perturbation_fields=['theta', 'rho'],
log_level='INFO')
parameters = CompressibleParameters()
diagnostic_fields = [CourantNumber()]
state = State(mesh,
dt=dt,
output=output,
parameters=parameters,
def setup_unsaturated(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)
recovered = True
degree = 0 if recovered else 1
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname + '/unsaturated_balance',
dumpfreq=1,
dumplist=['u', 'rho', 'theta'],
perturbation_fields=['water_v'])
parameters = CompressibleParameters()
diagnostics = Diagnostics(*fieldlist)
diagnostic_fields = [Theta_d(), RelativeHumidity()]
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.)
humidity = Constant(0.5)
theta_d = Function(Vt).interpolate(Tsurf)
RH = Function(Vt).interpolate(humidity)
# Calculate hydrostatic Pi
unsaturated_hydrostatic_balance(state, theta_d, RH)
water_c0.assign(0.0)
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 = state.spaces("DG1")
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)
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)
# Set up 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 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)
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt)
output = OutputParameters(dirname=dirname + "/sk_nonlinear",
dumplist=['u'],
dumpfreq=5,
log_level=INFO)
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,
supg_params={"dg_direction": "horizontal"})
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
form, parameters=parameters)[0]
indices = IndexDict(
{idx: sympy.symbols(name)
for idx, name in index_names})
expr = expression(impero_kernel.tree,
impero_kernel.temporaries,
indices,
top=True)
p1 = sympy.symbols("p") + 1
'''Currently assume p+1 quad points in each direction.'''
return expr.subs([(i, p1) for i in indices.values()]).expand()
m = ExtrudedMesh(UnitSquareMesh(2, 2, quadrilateral=True), 2)
mass = form.mass(m, 6)
poisson = form.poisson(m, 6)
hyperelasticity = form.hyperelasticity(m, 6)
curl_curl = form.curl_curl(m, 6)
parameters = firedrake.parameters['form_compiler'].copy()
parameters['return_impero'] = True
parameters['mode'] = 'spectral'
for mode, action in (("assembly", False), ("action", True)):
print(mode)
print(" mass: ", complexity(mass, parameters, action))
print(" laplacian: ", complexity(poisson, parameters, action))
print(" hyperelasticity:", complexity(hyperelasticity, parameters,
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.)
tmax = 10.0
# 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 = 20.0
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)
u0.project(u_evaluation(0))
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, 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)))
prescribed_fields = [('u', u_evaluation)]
physics_list = [Condensation(state)]
# build time stepper
stepper = AdvectionDiffusion(state,
advected_fields,
physics_list=physics_list,
prescribed_fields=prescribed_fields)
return stepper, tmax
