def __init__(self, state, linear=False):
self.state = state
g = state.parameters.g
f = state.f
Vu = state.V[0]
W = state.W
self.x0 = Function(W) # copy x to here
u0, D0 = split(self.x0)
n = FacetNormal(state.mesh)
un = 0.5 * (dot(u0, n) + abs(dot(u0, n)))
F = TrialFunction(Vu)
w = TestFunction(Vu)
self.uF = Function(Vu)
outward_normals = CellNormal(state.mesh)
perp = lambda u: cross(outward_normals, u)
a = inner(w, F) * dx
L = (-f * inner(w, perp(u0)) + g * div(w) * D0) * dx - g * inner(
jump(w, n), un("+") * D0("+") - un("-") * D0("-")
) * dS
if not linear:
L -= 0.5 * div(w) * inner(u0, u0) * dx
u_forcing_problem = LinearVariationalProblem(a, L, self.uF)
self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem)
def _setup_solver(self):
state = self.state
H = state.parameters.H
g = state.parameters.g
beta = state.timestepping.dt*state.timestepping.alpha
# Split up the rhs vector (symbolically)
u_in, D_in = split(state.xrhs)
W = state.W
w, phi = TestFunctions(W)
u, D = TrialFunctions(W)
eqn = (
inner(w, u) - beta*g*div(w)*D
- inner(w, u_in)
+ phi*D + beta*H*phi*div(u)
- phi*D_in
)*dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Place to put result of u rho solver
self.uD = Function(W)
# Solver for u, D
uD_problem = LinearVariationalProblem(
aeqn, Leqn, self.state.dy)
self.uD_solver = LinearVariationalSolver(uD_problem,
solver_parameters=self.params,
options_prefix='SWimplicit')
def _build_forcing_solver(self, linear):
"""
Only put forcing terms into the u equation.
"""
state = self.state
self.scaling = Constant(1.0)
Vu = state.V[0]
W = state.W
self.x0 = Function(W) # copy x to here
u0, p0, b0 = split(self.x0)
F = TrialFunction(Vu)
w = TestFunction(Vu)
self.uF = Function(Vu)
Omega = state.Omega
mu = state.mu
a = inner(w, F) * dx
L = (
self.scaling * div(w) * p0 * dx # pressure gradient
+ self.scaling * b0 * inner(w, state.k) * dx # gravity term
)
if not linear:
L -= self.scaling * 0.5 * div(w) * inner(u0, u0) * dx
if Omega is not None:
L -= self.scaling * inner(w, cross(2 * Omega, u0)) * dx # Coriolis term
if mu is not None:
self.mu_scaling = Constant(1.0)
L -= self.mu_scaling * mu * inner(w, state.k) * inner(u0, state.k) * dx
bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")]
u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs)
self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem)
Vp = state.V[1]
p = TrialFunction(Vp)
q = TestFunction(Vp)
self.divu = Function(Vp)
a = p * q * dx
L = q * div(u0) * dx
divergence_problem = LinearVariationalProblem(a, L, self.divu)
self.divergence_solver = LinearVariationalSolver(divergence_problem)
def _build_forcing_solver(self, linear):
"""
Only put forcing terms into the u equation.
"""
state = self.state
self.scaling = Constant(1.0)
Vu = state.V[0]
W = state.W
self.x0 = Function(W) # copy x to here
u0, rho0, theta0 = split(self.x0)
F = TrialFunction(Vu)
w = TestFunction(Vu)
self.uF = Function(Vu)
Omega = state.Omega
cp = state.parameters.cp
mu = state.mu
n = FacetNormal(state.mesh)
pi = exner(theta0, rho0, state)
a = inner(w, F) * dx
L = self.scaling * (
+cp * div(theta0 * w) * pi * dx # pressure gradient [volume]
- cp * jump(w * theta0, n) * avg(pi) * dS_v # pressure gradient [surface]
)
if state.parameters.geopotential:
Phi = state.Phi
L += self.scaling * div(w) * Phi * dx # gravity term
else:
g = state.parameters.g
L -= self.scaling * g * inner(w, state.k) * dx # gravity term
if not linear:
L -= self.scaling * 0.5 * div(w) * inner(u0, u0) * dx
if Omega is not None:
L -= self.scaling * inner(w, cross(2 * Omega, u0)) * dx # Coriolis term
if mu is not None:
self.mu_scaling = Constant(1.0)
L -= self.mu_scaling * mu * inner(w, state.k) * inner(u0, state.k) * dx
bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")]
u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs)
self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem)
def _build_forcing_solvers(self):
super(CompressibleForcing, self)._build_forcing_solvers()
# build forcing for theta equation
if self.moisture is not None:
_, _, theta0 = split(self.x0)
Vt = self.state.spaces("HDiv_v")
p = TrialFunction(Vt)
q = TestFunction(Vt)
self.thetaF = Function(Vt)
a = p * q * dx
L = self.theta_forcing()
L = q * L * dx
theta_problem = LinearVariationalProblem(a, L, self.thetaF)
solver_parameters = {}
if self.state.output.log_level == DEBUG:
solver_parameters["ksp_monitor_true_residual"] = True
self.theta_solver = LinearVariationalSolver(
theta_problem,
solver_parameters=solver_parameters,
option_prefix="ThetaForcingSolver")
def moist_hydrostatic_balance(state, theta_e, water_t, pi_boundary=Constant(1.0)):
"""
Given a wet equivalent potential temperature, theta_e, and the total moisture
content, water_t, compute a hydrostatically balance virtual potential temperature,
dry density and water vapour profile.
:arg state: The :class:`State` object.
:arg theta_e: The initial wet equivalent potential temperature profile.
:arg water_t: The total water pseudo-mixing ratio profile.
:arg pi_boundary: the value of pi on the lower boundary of the domain.
"""
theta0 = state.fields('theta')
rho0 = state.fields('rho')
water_v0 = state.fields('water_v')
# Calculate hydrostatic Pi
Vt = theta0.function_space()
Vr = rho0.function_space()
Vv = state.fields('u').function_space()
n = FacetNormal(state.mesh)
g = state.parameters.g
cp = state.parameters.cp
R_d = state.parameters.R_d
p_0 = state.parameters.p_0
VDG = state.spaces("DG")
if any(deg > 2 for deg in VDG.ufl_element().degree()):
state.logger.warning("default quadrature degree most likely not sufficient for this degree element")
quadrature_degree = (5, 5)
params = {'ksp_type': 'preonly',
'ksp_monitor_true_residual': True,
'ksp_converged_reason': True,
'snes_converged_reason': True,
'ksp_max_it': 100,
'mat_type': 'aij',
'pc_type': 'lu',
'pc_factor_mat_solver_type': 'mumps'}
theta0.interpolate(theta_e)
water_v0.interpolate(water_t)
Pi = Function(Vr)
epsilon = 0.9 # relaxation constant
# set up mixed space
Z = MixedFunctionSpace((Vt, Vt))
z = Function(Z)
gamma, phi = TestFunctions(Z)
theta_v, w_v = z.split()
# give first guesses for trial functions
theta_v.assign(theta0)
w_v.assign(water_v0)
theta_v, w_v = split(z)
# define variables
T = thermodynamics.T(state.parameters, theta_v, Pi, r_v=w_v)
p = thermodynamics.p(state.parameters, Pi)
w_sat = thermodynamics.r_sat(state.parameters, T, p)
dxp = dx(degree=(quadrature_degree))
# set up weak form of theta_e and w_sat equations
F = (-gamma * theta_e * dxp
+ gamma * thermodynamics.theta_e(state.parameters, T, p, w_v, water_t) * dxp
- phi * w_v * dxp
+ phi * w_sat * dxp)
problem = NonlinearVariationalProblem(F, z)
solver = NonlinearVariationalSolver(problem, solver_parameters=params)
theta_v, w_v = z.split()
Pi_h = Function(Vr).interpolate((p / p_0) ** (R_d / cp))
# solve for Pi with theta_v and w_v constant
# then solve for theta_v and w_v with Pi constant
for i in range(5):
compressible_hydrostatic_balance(state, theta0, rho0, pi0=Pi_h, water_t=water_t)
Pi.assign(Pi * (1 - epsilon) + epsilon * Pi_h)
solver.solve()
theta0.assign(theta0 * (1 - epsilon) + epsilon * theta_v)
water_v0.assign(water_v0 * (1 - epsilon) + epsilon * w_v)
# now begin on Newton solver, setup up new mixed space
Z = MixedFunctionSpace((Vt, Vt, Vr, Vv))
z = Function(Z)
gamma, phi, psi, w = TestFunctions(Z)
theta_v, w_v, pi, v = z.split()
# use previous values as first guesses for newton solver
theta_v.assign(theta0)
w_v.assign(water_v0)
pi.assign(Pi)
theta_v, w_v, pi, v = split(z)
# define variables
T = thermodynamics.T(state.parameters, theta_v, pi, r_v=w_v)
p = thermodynamics.p(state.parameters, pi)
w_sat = thermodynamics.r_sat(state.parameters, T, p)
F = (-gamma * theta_e * dxp
+ gamma * thermodynamics.theta_e(state.parameters, T, p, w_v, water_t) * dxp
- phi * w_v * dxp
+ phi * w_sat * dxp
+ cp * inner(v, w) * dxp
- cp * div(w * theta_v / (1.0 + water_t)) * pi * dxp
+ psi * div(theta_v * v / (1.0 + water_t)) * dxp
+ cp * inner(w, n) * pi_boundary * theta_v / (1.0 + water_t) * ds_b
+ g * inner(w, state.k) * dxp)
bcs = [DirichletBC(Z.sub(3), 0.0, "top")]
problem = NonlinearVariationalProblem(F, z, bcs=bcs)
solver = NonlinearVariationalSolver(problem, solver_parameters=params)
solver.solve()
theta_v, w_v, pi, v = z.split()
# assign final values
theta0.assign(theta_v)
water_v0.assign(w_v)
# find rho
compressible_hydrostatic_balance(state, theta0, rho0, water_t=water_t, solve_for_rho=True)
def __init__(self, prognostic_variables, simulation_parameters):
mesh = simulation_parameters['mesh'][-1]
self.scheme = simulation_parameters['scheme'][-1]
self.timestepping = simulation_parameters['timestepping'][-1]
alphasq = simulation_parameters['alphasq'][-1]
c0 = simulation_parameters['c0'][-1]
gamma = simulation_parameters['gamma'][-1]
Dt = Constant(simulation_parameters['dt'][-1])
self.solvers = []
if alphasq.values()[0] > 0.0 and gamma.values()[0] == 0.0:
self.setup = 'ch'
if self.scheme == 'upwind' and self.timestepping == 'ssprk3':
Vm = prognostic_variables.Vm
Vu = prognostic_variables.Vu
self.m = prognostic_variables.m
self.u = prognostic_variables.u
self.Xi = prognostic_variables.dXi
self.m0 = Function(Vm).assign(self.m)
# now make problem for the actual problem
psi = TestFunction(Vm)
self.m_trial = Function(Vm)
self.dm = Function(
Vm
) # introduce this as the advection operator for a single step
us = Dt * self.u + self.Xi
nhat = FacetNormal(mesh)
un = 0.5 * (dot(us, nhat) + abs(dot(us, nhat)))
ones = Function(Vu).project(as_vector([Constant(1.)]))
Lm = (psi * self.dm * dx -
psi.dx(0) * self.m_trial * dot(ones, us) * dx +
psi * self.m_trial * dot(ones, us.dx(0)) * dx +
jump(psi) * (un('+') * self.m_trial('+') -
un('-') * self.m_trial('-')) * dS)
mprob = NonlinearVariationalProblem(Lm, self.dm)
self.msolver = NonlinearVariationalSolver(mprob,
solver_parameters={
'ksp_type':
'preonly',
'pc_type':
'bjacobi',
'sub_pc_type':
'ilu'
})
phi = TestFunction(Vu)
Lu = (dot(phi, ones) * self.m * dx - dot(phi, self.u) * dx -
alphasq * dot(self.u.dx(0), phi.dx(0)) * dx)
uprob = NonlinearVariationalProblem(Lu, self.u)
self.usolver = NonlinearVariationalSolver(uprob,
solver_parameters={
'ksp_type':
'preonly',
'pc_type': 'lu'
})
elif self.scheme == 'hydrodynamic' and self.timestepping == 'midpoint':
Vu = prognostic_variables.Vu
self.u = prognostic_variables.u
W = MixedFunctionSpace((Vu, ) * 3)
psi, phi, zeta = TestFunctions(W)
w1 = Function(W)
self.u1, dFh, dGh = split(w1)
uh = (self.u1 + self.u) / 2
dXi = prognostic_variables.dXi
dXi_x = prognostic_variables.dXi_x
dXi_xx = prognostic_variables.dXi_xx
dvh = Dt * uh + dXi
Lu = (psi * (self.u1 - self.u) * dx +
psi * uh.dx(0) * dvh * dx - psi.dx(0) * dFh * dx +
psi * dGh * dx + phi * dFh * dx +
alphasq * phi.dx(0) * dFh.dx(0) * dx -
phi * uh * uh * Dt * dx -
0.5 * alphasq * phi * uh.dx(0) * uh.dx(0) * Dt * dx +
zeta * dGh * dx + alphasq * zeta.dx(0) * dGh.dx(0) * dx -
2 * zeta * uh * dXi_x * dx -
alphasq * zeta * uh.dx(0) * dXi_xx * dx)
self.u1, dFh, dGh = w1.split()
uprob = NonlinearVariationalProblem(Lu, w1)
self.usolver = NonlinearVariationalSolver(uprob,
solver_parameters={
'mat_type':
'aij',
'ksp_type':
'preonly',
'pc_type': 'lu'
})
elif self.scheme == 'no_gradient' and self.timestepping == 'midpoint':
# a version of the hydrodynamic form but without exploiting the gradient
Vu = prognostic_variables.Vu
self.u = prognostic_variables.u
W = MixedFunctionSpace((Vu, ) * 3)
psi, phi, zeta = TestFunctions(W)
w1 = Function(W)
self.u1, dFh, dGh = split(w1)
uh = (self.u1 + self.u) / 2
dXi = prognostic_variables.dXi
dXi_x = prognostic_variables.dXi_x
dXi_xx = prognostic_variables.dXi_xx
dvh = Dt * uh + dXi
Lu = (psi * (self.u1 - self.u) * dx +
psi * uh.dx(0) * dvh * dx + psi * dFh.dx(0) * dx +
psi * dGh * dx + phi * dFh * dx +
alphasq * phi.dx(0) * dFh.dx(0) * dx -
phi * uh * uh * Dt * dx -
0.5 * alphasq * phi * uh.dx(0) * uh.dx(0) * Dt * dx +
zeta * dGh * dx + alphasq * zeta.dx(0) * dGh.dx(0) * dx -
2 * zeta * uh * dXi_x * dx -
alphasq * zeta * uh.dx(0) * dXi_xx * dx)
self.u1, dFh, dGh = w1.split()
uprob = NonlinearVariationalProblem(Lu, w1)
self.usolver = NonlinearVariationalSolver(uprob,
solver_parameters={
'mat_type':
'aij',
'ksp_type':
'preonly',
'pc_type': 'lu'
})
elif self.scheme == 'test' and self.timestepping == 'midpoint':
self.u = prognostic_variables.u
Vu = prognostic_variables.Vu
psi = TestFunction(Vu)
self.u1 = Function(Vu)
uh = (self.u1 + self.u) / 2
dvh = Dt * uh + prognostic_variables.dXi
eqn = (psi * (self.u1 - self.u) * dx -
psi * uh * dvh.dx(0) * dx)
prob = NonlinearVariationalProblem(eqn, self.u1)
self.usolver = NonlinearVariationalSolver(prob,
solver_parameters={
'mat_type':
'aij',
'ksp_type':
'preonly',
'pc_type': 'lu'
})
else:
raise ValueError(
'Scheme %s and timestepping %s either not compatible or not recognised.'
% (self.scheme, self.timestepping))
elif alphasq.values()[0] == 0.0 and gamma.values()[0] > 0.0:
self.setup = 'kdv'
if self.scheme == 'upwind' and self.timestepping == 'ssprk3':
raise NotImplementedError(
'Scheme %s and timestepping %s not yet implemented.' %
(self.scheme, self.timestepping))
elif self.scheme == 'upwind' and self.timestepping == 'midpoint':
raise NotImplementedError(
'Scheme %s and timestepping %s not yet implemented.' %
(self.scheme, self.timestepping))
elif self.scheme == 'hydrodynamic' and self.timestepping == 'midpoint':
raise NotImplementedError(
'Scheme %s and timestepping %s not yet implemented.' %
(self.scheme, self.timestepping))
else:
raise ValueError(
'Scheme %s and timestepping %s either not compatible or not recognised.'
% (self.scheme, self.timestepping))
else:
raise NotImplementedError(
'Schemes for your values of alpha squared %.3f and gamma %.3f are not yet implemented.'
% (alphasq, gamma))
def __init__(self,
solution: fe.Function,
time: float = 0.,
time_stencil_size: int = 2,
timestep_size: float = 1.,
quadrature_degree: int = None,
solver_parameters: dict = {
"snes_type": "newtonls",
"snes_monitor": None,
"ksp_type": "preonly",
"pc_type": "lu",
"mat_type": "aij",
"pc_factor_mat_solver_type": "mumps"},
output_directory_path: str = "output/",
fieldnames: typing.Iterable[str] = None):
"""
Instantiating this class requires enough information to fully
specify the FE spatial discretization and weak form residual.
boundary conditions, and initial values. All of these required
arguments are Firedrake objects used according to Firedrake
conventions.
Backward Difference Formula time discretizations are
automatically implemented. To use a different time
discretization, inherit this class and redefine
`time_discrete_terms`.
Args:
solution: Solution for a single time step.
As a `fe.Function`, this also defines the
mesh, element, and solution function space.
time: The initial time.
time_stencil_size: The number of solutions at
discrete times used for approximating time derivatives.
This also determines the number of stored solutions.
Must be greater than zero.
Defaults to 2. Set to 1 for steady state problems.
Increase for higher-order time accuracy.
timestep_size: The size of discrete time steps.
Defaults to 1.
Higher order time discretizations are assumed to use
a constant time step size.
quadrature_degree: The quadrature degree used for
numerical integration.
Defaults to `None`, in which case Firedrake will
automatically choose a suitable quadrature degree.
solver_parameters: The solver parameters dictionary
which Firedrake uses to configure PETSc.
output_directory_path: String that will be converted
to a Path where output files will be written.
Defaults to "output/".
fieldnames: A list of names for the components of `solution`.
Defaults to `None`.
These names can be used when indexing solutions that are split
either by `firedrake.split` or `firedrake.Function.split`.
If not `None`, then the `dict` `self.solution_fields` will be created.
The `dict` will have two items for each field,
containing the results of either splitting method.
The results of `firedrake.split` will be suffixed with "_ufl".
"""
assert(time_stencil_size > 0)
self.fieldcount = len(solution.split())
if fieldnames is None:
fieldnames = ["w_{}" for i in range(self.fieldcount)]
assert(len(fieldnames) == self.fieldcount)
self.fieldnames = fieldnames
self.solution = solution
self.time = fe.Constant(time)
self.solution_space = self.solution.function_space()
self.mesh = self.solution_space.mesh()
self.unit_vectors = unit_vectors(self.mesh)
self.element = self.solution_space.ufl_element()
self.timestep_size = fe.Constant(timestep_size)
self.quadrature_degree = quadrature_degree
self.dx = fe.dx(degree = self.quadrature_degree)
self.solver_parameters = solver_parameters
initial_values = self.initial_values()
if initial_values is not None:
self.solution = self.solution.assign(initial_values)
# States for time dependent simulation and checkpointing
self.solutions = [self.solution,]
self.times = [self.time,]
self.state = {
"solution": self.solution,
"time": self.time,
"index": 0}
self.states = [self.state,]
for i in range(1, time_stencil_size):
self.solutions.append(fe.Function(self.solution))
self.times.append(fe.Constant(self.time - i*timestep_size))
self.states.append({
"solution": self.solutions[i],
"time": self.times[i],
"index": -i})
# Continuation helpers
self.backup_solution = fe.Function(self.solution)
# Mixed solution indexing helpers
self.solution_fields = {}
self.solution_subfunctions = {}
self.test_functions = {}
self.time_discrete_terms = {}
self.solution_subspaces = {}
for name, field, field_pp, testfun, timeterm in zip(
fieldnames,
fe.split(self.solution),
self.solution.split(),
fe.TestFunctions(self.solution_space),
time_discrete_terms(
solutions = self.solutions,
timestep_size = self.timestep_size)):
self.solution_fields[name] = field
self.solution_subfunctions[name] = field_pp
self.test_functions[name] = testfun
self.time_discrete_terms[name] = timeterm
self.solution_subspaces[name] = self.solution_space.sub(
fieldnames.index(name))
# Output controls
self.output_directory_path = pathlib.Path(output_directory_path)
self.output_directory_path.mkdir(parents = True, exist_ok = True)
self.vtk_solution_file = None
self.plotvars = None
self.snes_iteration_count = 0
def heat_exchanger_optimization(mu=0.03, n_iters=1000):
output_dir = "2D/"
path = os.path.abspath(__file__)
dir_path = os.path.dirname(path)
mesh = fd.Mesh(f"{dir_path}/2D_mesh.msh")
# Perturb the mesh coordinates. Necessary to calculate shape derivatives
S = fd.VectorFunctionSpace(mesh, "CG", 1)
s = fd.Function(S, name="deform")
mesh.coordinates.assign(mesh.coordinates + s)
# Initial level set function
x, y = fd.SpatialCoordinate(mesh)
PHI = fd.FunctionSpace(mesh, "CG", 1)
phi_expr = sin(y * pi / 0.2) * cos(x * pi / 0.2) - fd.Constant(0.8)
# Avoid recording the operation interpolate into the tape.
# Otherwise, the shape derivatives will not be correct
with fda.stop_annotating():
phi = fd.interpolate(phi_expr, PHI)
phi.rename("LevelSet")
fd.File(output_dir + "phi_initial.pvd").write(phi)
# Physics
mu = fd.Constant(mu) # viscosity
alphamin = 1e-12
alphamax = 2.5 / (2e-4)
parameters = {
"mat_type": "aij",
"ksp_type": "preonly",
"ksp_converged_reason": None,
"pc_type": "lu",
"pc_factor_mat_solver_type": "mumps",
}
stokes_parameters = parameters
temperature_parameters = parameters
u_inflow = 2e-3
tin1 = fd.Constant(10.0)
tin2 = fd.Constant(100.0)
P2 = fd.VectorElement("CG", mesh.ufl_cell(), 2)
P1 = fd.FiniteElement("CG", mesh.ufl_cell(), 1)
TH = P2 * P1
W = fd.FunctionSpace(mesh, TH)
U = fd.TrialFunction(W)
u, p = fd.split(U)
V = fd.TestFunction(W)
v, q = fd.split(V)
epsilon = fd.Constant(10000.0)
def hs(phi, epsilon):
return fd.Constant(alphamax) * fd.Constant(1.0) / (
fd.Constant(1.0) + exp(-epsilon * phi)) + fd.Constant(alphamin)
def stokes(phi, BLOCK_INLET_MOUTH, BLOCK_OUTLET_MOUTH):
a_fluid = mu * inner(grad(u), grad(v)) - div(v) * p - q * div(u)
darcy_term = inner(u, v)
return (a_fluid * dx + hs(phi, epsilon) * darcy_term * dx(0) +
alphamax * darcy_term *
(dx(BLOCK_INLET_MOUTH) + dx(BLOCK_OUTLET_MOUTH)))
# Dirichlet boundary conditions
inflow1 = fd.as_vector([
u_inflow * sin(
((y - (line_sep -
(dist_center + inlet_width))) * pi) / inlet_width),
0.0,
])
inflow2 = fd.as_vector([
u_inflow * sin(((y - (line_sep + dist_center)) * pi) / inlet_width),
0.0,
])
noslip = fd.Constant((0.0, 0.0))
# Stokes 1
bcs1_1 = fd.DirichletBC(W.sub(0), noslip, WALLS)
bcs1_2 = fd.DirichletBC(W.sub(0), inflow1, INLET1)
bcs1_3 = fd.DirichletBC(W.sub(1), fd.Constant(0.0), OUTLET1)
bcs1_4 = fd.DirichletBC(W.sub(0), noslip, INLET2)
bcs1_5 = fd.DirichletBC(W.sub(0), noslip, OUTLET2)
bcs1 = [bcs1_1, bcs1_2, bcs1_3, bcs1_4, bcs1_5]
# Stokes 2
bcs2_1 = fd.DirichletBC(W.sub(0), noslip, WALLS)
bcs2_2 = fd.DirichletBC(W.sub(0), inflow2, INLET2)
bcs2_3 = fd.DirichletBC(W.sub(1), fd.Constant(0.0), OUTLET2)
bcs2_4 = fd.DirichletBC(W.sub(0), noslip, INLET1)
bcs2_5 = fd.DirichletBC(W.sub(0), noslip, OUTLET1)
bcs2 = [bcs2_1, bcs2_2, bcs2_3, bcs2_4, bcs2_5]
# Forward problems
U1, U2 = fd.Function(W), fd.Function(W)
L = inner(fd.Constant((0.0, 0.0, 0.0)), V) * dx
problem = fd.LinearVariationalProblem(stokes(-phi, INMOUTH2, OUTMOUTH2),
L,
U1,
bcs=bcs1)
solver_stokes1 = fd.LinearVariationalSolver(
problem,
solver_parameters=stokes_parameters,
options_prefix="stokes_1")
solver_stokes1.solve()
problem = fd.LinearVariationalProblem(stokes(phi, INMOUTH1, OUTMOUTH1),
L,
U2,
bcs=bcs2)
solver_stokes2 = fd.LinearVariationalSolver(
problem,
solver_parameters=stokes_parameters,
options_prefix="stokes_2")
solver_stokes2.solve()
# Convection difussion equation
ks = fd.Constant(1e0)
cp_value = 5.0e5
cp = fd.Constant(cp_value)
T = fd.FunctionSpace(mesh, "DG", 1)
t = fd.Function(T, name="Temperature")
w = fd.TestFunction(T)
# Mesh-related functions
n = fd.FacetNormal(mesh)
h = fd.CellDiameter(mesh)
u1, p1 = fd.split(U1)
u2, p2 = fd.split(U2)
def upwind(u):
return (dot(u, n) + abs(dot(u, n))) / 2.0
u1n = upwind(u1)
u2n = upwind(u2)
# Penalty term
alpha = fd.Constant(500.0)
# Bilinear form
a_int = dot(grad(w), ks * grad(t) - cp * (u1 + u2) * t) * dx
a_fac = (fd.Constant(-1.0) * ks * dot(avg(grad(w)), jump(t, n)) * dS +
fd.Constant(-1.0) * ks * dot(jump(w, n), avg(grad(t))) * dS +
ks("+") *
(alpha("+") / avg(h)) * dot(jump(w, n), jump(t, n)) * dS)
a_vel = (dot(
jump(w),
cp * (u1n("+") + u2n("+")) * t("+") - cp *
(u1n("-") + u2n("-")) * t("-"),
) * dS + dot(w,
cp * (u1n + u2n) * t) * ds)
a_bnd = (dot(w,
cp * dot(u1 + u2, n) * t) * (ds(INLET1) + ds(INLET2)) +
w * t * (ds(INLET1) + ds(INLET2)) - w * tin1 * ds(INLET1) -
w * tin2 * ds(INLET2) + alpha / h * ks * w * t *
(ds(INLET1) + ds(INLET2)) - ks * dot(grad(w), t * n) *
(ds(INLET1) + ds(INLET2)) - ks * dot(grad(t), w * n) *
(ds(INLET1) + ds(INLET2)))
aT = a_int + a_fac + a_vel + a_bnd
LT_bnd = (alpha / h * ks * tin1 * w * ds(INLET1) +
alpha / h * ks * tin2 * w * ds(INLET2) -
tin1 * ks * dot(grad(w), n) * ds(INLET1) -
tin2 * ks * dot(grad(w), n) * ds(INLET2))
problem = fd.LinearVariationalProblem(derivative(aT, t), LT_bnd, t)
solver_temp = fd.LinearVariationalSolver(
problem,
solver_parameters=temperature_parameters,
options_prefix="temperature",
)
solver_temp.solve()
# fd.solve(eT == 0, t, solver_parameters=temperature_parameters)
# Cost function: Flux at the cold outlet
scale_factor = 4e-4
Jform = fd.assemble(
fd.Constant(-scale_factor * cp_value) * inner(t * u1, n) * ds(OUTLET1))
# Constraints: Pressure drop on each fluid
power_drop = 1e-2
Power1 = fd.assemble(p1 / power_drop * ds(INLET1))
Power2 = fd.assemble(p2 / power_drop * ds(INLET2))
phi_pvd = fd.File("phi_evolution.pvd")
def deriv_cb(phi):
with stop_annotating():
phi_pvd.write(phi[0])
c = fda.Control(s)
# Reduced Functionals
Jhat = LevelSetFunctional(Jform, c, phi, derivative_cb_pre=deriv_cb)
P1hat = LevelSetFunctional(Power1, c, phi)
P1control = fda.Control(Power1)
P2hat = LevelSetFunctional(Power2, c, phi)
P2control = fda.Control(Power2)
Jhat_v = Jhat(phi)
print("Initial cost function value {:.5f}".format(Jhat_v), flush=True)
print("Power drop 1 {:.5f}".format(Power1), flush=True)
print("Power drop 2 {:.5f}".format(Power2), flush=True)
beta_param = 0.08
# Regularize the shape derivatives only in the domain marked with 0
reg_solver = RegularizationSolver(S,
mesh,
beta=beta_param,
gamma=1e5,
dx=dx,
design_domain=0)
tol = 1e-5
dt = 0.05
params = {
"alphaC": 1.0,
"debug": 5,
"alphaJ": 1.0,
"dt": dt,
"K": 1e-3,
"maxit": n_iters,
"maxtrials": 5,
"itnormalisation": 10,
"tol_merit":
5e-3, # new merit can be within 0.5% of the previous merit
# "normalize_tol" : -1,
"tol": tol,
}
solver_parameters = {
"reinit_solver": {
"h_factor": 2.0,
}
}
# Optimization problem
problem = InfDimProblem(
Jhat,
reg_solver,
ineqconstraints=[
Constraint(P1hat, 1.0, P1control),
Constraint(P2hat, 1.0, P2control),
],
solver_parameters=solver_parameters,
)
results = nlspace_solve(problem, params)
return results
DG0 = fd.FunctionSpace(mesh, "DQ", 0)
vDG = fd.VectorFunctionSpace(mesh, "DQ", 1)
else:
DG1 = fd.FunctionSpace(mesh, "DG", order)
vDG1 = fd.VectorFunctionSpace(mesh, "DG", order)
Mh = fd.FunctionSpace(mesh, "HDiv Trace", order)
DG0 = fd.FunctionSpace(mesh, "DG", 0)
vDG = fd.VectorFunctionSpace(mesh, "DG", 1)
W = DG1 * vDG1 * Mh
# 3.1) Define trial and test functions
w = fd.Function(W)
w.assign(0.0)
ch, qh, lmbd_h = fd.split(w)
wh, vh, mu_h = fd.TestFunctions(W)
# 3.2) Set initial conditions
# ---- previous solution
# concentrations
c0 = fd.Function(DG1, name="c0")
c0.assign(0.0)
# ----------------------
# 3.4) Variational Form
# ----------------------
# coefficients
dt = np.sqrt(tol)
dtc = fd.Constant(dt)
n = fd.FacetNormal(mesh)
def initialize(self, pc):
from firedrake import TrialFunction, TestFunction, dx, \
assemble, inner, grad, split, Constant, parameters
from firedrake.assemble import allocate_matrix, create_assembly_callable
if pc.getType() != "python":
raise ValueError("Expecting PC type python")
prefix = pc.getOptionsPrefix() + "pcd_"
# we assume P has things stuffed inside of it
_, P = pc.getOperators()
context = P.getPythonContext()
test, trial = context.a.arguments()
if test.function_space() != trial.function_space():
raise ValueError("Pressure space test and trial space differ")
Q = test.function_space()
p = TrialFunction(Q)
q = TestFunction(Q)
mass = p * q * dx
# Regularisation to avoid having to think about nullspaces.
stiffness = inner(grad(p), grad(q)) * dx + Constant(1e-6) * p * q * dx
opts = PETSc.Options()
# we're inverting Mp and Kp, so default them to assembled.
# Fp only needs its action, so default it to mat-free.
# These can of course be overridden.
# only Fp is referred to in update, so that's the only
# one we stash.
default = parameters["default_matrix_type"]
Mp_mat_type = opts.getString(prefix + "Mp_mat_type", default)
Kp_mat_type = opts.getString(prefix + "Kp_mat_type", default)
self.Fp_mat_type = opts.getString(prefix + "Fp_mat_type", "matfree")
Mp = assemble(mass,
form_compiler_parameters=context.fc_params,
mat_type=Mp_mat_type,
options_prefix=prefix + "Mp_")
Kp = assemble(stiffness,
form_compiler_parameters=context.fc_params,
mat_type=Kp_mat_type,
options_prefix=prefix + "Kp_")
Mp.force_evaluation()
Kp.force_evaluation()
# FIXME: Should we transfer nullspaces over. I think not.
Mksp = PETSc.KSP().create(comm=pc.comm)
Mksp.incrementTabLevel(1, parent=pc)
Mksp.setOptionsPrefix(prefix + "Mp_")
Mksp.setOperators(Mp.petscmat)
Mksp.setUp()
Mksp.setFromOptions()
self.Mksp = Mksp
Kksp = PETSc.KSP().create(comm=pc.comm)
Kksp.incrementTabLevel(1, parent=pc)
Kksp.setOptionsPrefix(prefix + "Kp_")
Kksp.setOperators(Kp.petscmat)
Kksp.setUp()
Kksp.setFromOptions()
self.Kksp = Kksp
state = context.appctx["state"]
Re = context.appctx.get("Re", 1.0)
velid = context.appctx["velocity_space"]
u0 = split(state)[velid]
fp = 1.0 / Re * inner(grad(p), grad(q)) * dx + inner(u0,
grad(p)) * q * dx
self.Re = Re
self.Fp = allocate_matrix(fp,
form_compiler_parameters=context.fc_params,
mat_type=self.Fp_mat_type,
options_prefix=prefix + "Fp_")
self._assemble_Fp = create_assembly_callable(
fp,
tensor=self.Fp,
form_compiler_parameters=context.fc_params,
mat_type=self.Fp_mat_type)
self._assemble_Fp()
self.Fp.force_evaluation()
Fpmat = self.Fp.petscmat
self.workspace = [Fpmat.createVecLeft() for i in (0, 1)]
def value_form(self):
"""Evaluate misfit functional."""
u = fd.split(self.pde_solver.solution)[0]
nu = self.pde_solver.nu
return 0.5 * nu * fd.inner(fd.grad(u), fd.grad(u)) * fd.dx
def lmbdainv(s):
# Total mobility
return mu_rel * mu_w / (mu_rel * s**2 + (1.0 - s)**2)
def Fw(s):
# Fractional flow function
return mu_rel * s**2 / (mu_rel * s**2 + (1.0 - s)**2)
# ---
# 3.3) initial cond.
U0 = fd.Function(W)
u0, p0, s0 = fd.split(U0)
U = fd.Function(W)
u, p, s = fd.split(U)
# ----
# set boundary conditions
# The strongly enforced boundary conditions on the BDM space on the top and
# bottom of the domain are declared as: ::
bc0 = fd.DirichletBC(W.sub(0), fd.Constant(qbar), inlet)
bc1 = fd.DirichletBC(W.sub(0), fd.Constant(q0bar), noflow)
bc3 = fd.DirichletBC(W.sub(2), fd.Constant(sbar), inlet, method='geometric')
# -------
# 3.4) Variational Form
# Time step
def _setup_solver(self):
state = self.state # just cutting down line length a bit
Dt = state.timestepping.dt
beta_ = Dt*state.timestepping.alpha
mu = state.mu
Vu = state.spaces("HDiv")
Vb = state.spaces("HDiv_v")
Vp = state.spaces("DG")
# Store time-stepping coefficients as UFL Constants
dt = Constant(Dt)
beta = Constant(beta_)
# Split up the rhs vector (symbolically)
u_in, p_in, b_in = split(state.xrhs)
# Build the reduced function space for u,p
M = MixedFunctionSpace((Vu, Vp))
w, phi = TestFunctions(M)
u, p = TrialFunctions(M)
# Get background fields
bbar = state.fields("bbar")
# Analytical (approximate) elimination of theta
k = state.k # Upward pointing unit vector
b = -dot(k, u)*dot(k, grad(bbar))*beta + b_in
# vertical projection
def V(u):
return k*inner(u, k)
eqn = (
inner(w, (u - u_in))*dx
- beta*div(w)*p*dx
- beta*inner(w, k)*b*dx
+ phi*div(u)*dx
)
if mu is not None:
eqn += dt*mu*inner(w, k)*inner(u, k)*dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Place to put result of u p solver
self.up = Function(M)
# Boundary conditions (assumes extruded mesh)
bcs = None if len(self.state.bcs) == 0 else self.state.bcs
# Solver for u, p
up_problem = LinearVariationalProblem(aeqn, Leqn, self.up, bcs=bcs)
# Provide callback for the nullspace of the trace system
def trace_nullsp(T):
return VectorSpaceBasis(constant=True)
appctx = {"trace_nullspace": trace_nullsp}
self.up_solver = LinearVariationalSolver(up_problem,
solver_parameters=self.solver_parameters,
appctx=appctx)
# Reconstruction of b
b = TrialFunction(Vb)
gamma = TestFunction(Vb)
u, p = self.up.split()
self.b = Function(Vb)
b_eqn = gamma*(b - b_in
+ dot(k, u)*dot(k, grad(bbar))*beta)*dx
b_problem = LinearVariationalProblem(lhs(b_eqn),
rhs(b_eqn),
self.b)
self.b_solver = LinearVariationalSolver(b_problem)
def compressible_hydrostatic_balance(state,
theta0,
rho0,
exner0=None,
top=False,
exner_boundary=Constant(1.0),
mr_t=None,
solve_for_rho=False,
params=None):
"""
Compute a hydrostatically balanced density given a potential temperature
profile. By default, this uses a vertically-oriented hybridization
procedure for solving the resulting discrete systems.
:arg state: The :class:`State` object.
:arg theta0: :class:`.Function`containing the potential temperature.
:arg rho0: :class:`.Function` to write the initial density into.
:arg top: If True, set a boundary condition at the top. Otherwise, set
it at the bottom.
:arg exner_boundary: a field or expression to use as boundary data for exner
on the top or bottom as specified.
:arg mr_t: the initial total water mixing ratio field.
"""
# Calculate hydrostatic Pi
VDG = state.spaces("DG")
Vu = state.spaces("HDiv")
Vv = FunctionSpace(state.mesh, Vu.ufl_element()._elements[-1])
W = MixedFunctionSpace((Vv, VDG))
v, exner = TrialFunctions(W)
dv, dexner = TestFunctions(W)
n = FacetNormal(state.mesh)
cp = state.parameters.cp
# add effect of density of water upon theta
theta = theta0
if mr_t is not None:
theta = theta0 / (1 + mr_t)
alhs = ((cp * inner(v, dv) - cp * div(dv * theta) * exner) * dx +
dexner * div(theta * v) * dx)
if top:
bmeasure = ds_t
bstring = "bottom"
else:
bmeasure = ds_b
bstring = "top"
arhs = -cp * inner(dv, n) * theta * exner_boundary * bmeasure
# Possibly make g vary with spatial coordinates?
g = state.parameters.g
arhs -= g * inner(dv, state.k) * dx
bcs = [DirichletBC(W.sub(0), zero(), bstring)]
w = Function(W)
exner_problem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs)
if params is None:
params = {
'ksp_type': 'preonly',
'pc_type': 'python',
'mat_type': 'matfree',
'pc_python_type': 'gusto.VerticalHybridizationPC',
# Vertical trace system is only coupled vertically in columns
# block ILU is a direct solver!
'vert_hybridization': {
'ksp_type': 'preonly',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
}
}
exner_solver = LinearVariationalSolver(exner_problem,
solver_parameters=params,
options_prefix="exner_solver")
exner_solver.solve()
v, exner = w.split()
if exner0 is not None:
exner0.assign(exner)
if solve_for_rho:
w1 = Function(W)
v, rho = w1.split()
rho.interpolate(thermodynamics.rho(state.parameters, theta0, exner))
v, rho = split(w1)
dv, dexner = TestFunctions(W)
exner = thermodynamics.exner_pressure(state.parameters, rho, theta0)
F = ((cp * inner(v, dv) - cp * div(dv * theta) * exner) * dx +
dexner * div(theta0 * v) * dx +
cp * inner(dv, n) * theta * exner_boundary * bmeasure)
F += g * inner(dv, state.k) * dx
rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs)
rhosolver = NonlinearVariationalSolver(rhoproblem,
solver_parameters=params,
options_prefix="rhosolver")
rhosolver.solve()
v, rho_ = w1.split()
rho0.assign(rho_)
else:
rho0.interpolate(thermodynamics.rho(state.parameters, theta0, exner))
def initialize(self, pc):
from firedrake import TrialFunction, TestFunction, dx, \
assemble, inner, grad, split, Constant, parameters
from firedrake.assemble import allocate_matrix, create_assembly_callable
prefix = pc.getOptionsPrefix() + "pcd_"
# we assume P has things stuffed inside of it
_, P = pc.getOperators()
context = P.getPythonContext()
test, trial = context.a.arguments()
if test.function_space() != trial.function_space():
raise ValueError("Pressure space test and trial space differ")
Q = test.function_space()
p = TrialFunction(Q)
q = TestFunction(Q)
mass = p*q*dx
# Regularisation to avoid having to think about nullspaces.
stiffness = inner(grad(p), grad(q))*dx + Constant(1e-6)*p*q*dx
opts = PETSc.Options()
# we're inverting Mp and Kp, so default them to assembled.
# Fp only needs its action, so default it to mat-free.
# These can of course be overridden.
# only Fp is referred to in update, so that's the only
# one we stash.
default = parameters["default_matrix_type"]
Mp_mat_type = opts.getString(prefix+"Mp_mat_type", default)
Kp_mat_type = opts.getString(prefix+"Kp_mat_type", default)
self.Fp_mat_type = opts.getString(prefix+"Fp_mat_type", "matfree")
Mp = assemble(mass, form_compiler_parameters=context.fc_params,
mat_type=Mp_mat_type)
Kp = assemble(stiffness, form_compiler_parameters=context.fc_params,
mat_type=Kp_mat_type)
Mp.force_evaluation()
Kp.force_evaluation()
# FIXME: Should we transfer nullspaces over. I think not.
Mksp = PETSc.KSP().create()
Mksp.setOptionsPrefix(prefix + "Mp_")
Mksp.setOperators(Mp.petscmat)
Mksp.setUp()
Mksp.setFromOptions()
self.Mksp = Mksp
Kksp = PETSc.KSP().create()
Kksp.setOptionsPrefix(prefix + "Kp_")
Kksp.setOperators(Kp.petscmat)
Kksp.setUp()
Kksp.setFromOptions()
self.Kksp = Kksp
state = context.appctx["state"]
Re = context.appctx.get("Re", 1.0)
velid = context.appctx["velocity_space"]
u0 = split(state)[velid]
fp = 1.0/Re * inner(grad(p), grad(q))*dx + inner(u0, grad(p))*q*dx
self.Re = Re
self.Fp = allocate_matrix(fp, form_compiler_parameters=context.fc_params,
mat_type=self.Fp_mat_type)
self._assemble_Fp = create_assembly_callable(fp, tensor=self.Fp,
form_compiler_parameters=context.fc_params,
mat_type=self.Fp_mat_type)
self._assemble_Fp()
self.Fp.force_evaluation()
Fpmat = self.Fp.petscmat
self.workspace = [Fpmat.createVecLeft() for i in (0, 1)]
else:
innerp = fs.UflInnerProductFromForm(extension,
Q,
fixed_bids=fixed_bids,
direct_solve=True)
# import IPython; IPython.embed()
res = [0, 1, 10, 50, 100, 150, 200, 250, 300, 400, 499, 625, 750, 875, 999]
res = [r for r in res if r <= optre - 1]
if res[-1] != optre - 1:
res.append(optre - 1)
results = run_solver(solver, res, args)
# import sys; sys.exit()
u, _ = fd.split(solver.z)
nu = solver.nu
objective_form = nu * fd.inner(fd.sym(fd.grad(u)), fd.sym(fd.grad(u))) * fd.dx
solver.setup_adjoint(objective_form)
solver.solver_adjoint.solve()
# import sys; sys.exit()
class Constraint(fs.PdeConstraint):
def solve(self):
super().solve()
solver.solve(optre)
class Objective(fs.ShapeObjective):
def getForm(F,
butch,
t,
dt,
u0,
bcs=None,
bc_type="DAE",
splitting=AI,
nullspace=None):
"""Given a time-dependent variational form and a
:class:`ButcherTableau`, produce UFL for the s-stage RK method.
:arg F: UFL form for the semidiscrete ODE/DAE
:arg butch: the :class:`ButcherTableau` for the RK method being used to
advance in time.
:arg t: a :class:`Constant` referring to the current time level.
Any explicit time-dependence in F is included
:arg dt: a :class:`Constant` referring to the size of the current
time step.
:arg splitting: a callable that maps the (floating point) Butcher matrix
a to a pair of matrices `A1, A2` such that `butch.A = A1 A2`. This is used
to vary between the classical RK formulation and Butcher's reformulation
that leads to a denser mass matrix with block-diagonal stiffness.
Some choices of function will assume that `butch.A` is invertible.
:arg u0: a :class:`Function` referring to the state of
the PDE system at time `t`
:arg bcs: optionally, a :class:`DirichletBC` object (or iterable thereof)
containing (possible time-dependent) boundary conditions imposed
on the system.
:arg bc_type: How to manipulate the strongly-enforced boundary
conditions to derive the stage boundary conditions. Should
be a string, either "DAE", which implements BCs as
constraints in the style of a differential-algebraic
equation, or "ODE", which takes the time derivative of the
boundary data and evaluates this for the stage values
:arg nullspace: A list of tuples of the form (index, VSB) where
index is an index into the function space associated with `u`
and VSB is a :class: `firedrake.VectorSpaceBasis` instance to
be passed to a `firedrake.MixedVectorSpaceBasis` over the
larger space associated with the Runge-Kutta method
On output, we return a tuple consisting of four parts:
- Fnew, the :class:`Form`
- k, the :class:`firedrake.Function` holding all the stages.
It lives in a :class:`firedrake.FunctionSpace` corresponding to the
s-way tensor product of the space on which the semidiscrete
form lives.
- `bcnew`, a list of :class:`firedrake.DirichletBC` objects to be posed
on the stages,
- 'nspnew', the :class:`firedrake.MixedVectorSpaceBasis` object
that represents the nullspace of the coupled system
- `gblah`, a list of tuples of the form (f, expr, method),
where f is a :class:`firedrake.Function` and expr is a
:class:`ufl.Expr`. At each time step, each expr needs to be
re-interpolated/projected onto the corresponding f in order
for Firedrake to pick up that time-dependent boundary
conditions need to be re-applied. The
interpolation/projection is encoded in method, which is
either `f.interpolate(expr-c*u0)` or `f.project(expr-c*u0)`, depending
on whether the function space for f supports interpolation or
not.
"""
v = F.arguments()[0]
V = v.function_space()
assert V == u0.function_space()
c = numpy.array([Constant(ci) for ci in butch.c], dtype=object)
bA1, bA2 = splitting(butch.A)
try:
bA1inv = numpy.linalg.inv(bA1)
except numpy.linalg.LinAlgError:
bA1inv = None
try:
bA2inv = numpy.linalg.inv(bA2)
A2inv = numpy.array([[ConstantOrZero(aa) for aa in arow]
for arow in bA2inv],
dtype=object)
except numpy.linalg.LinAlgError:
raise NotImplementedError("We require A = A1 A2 with A2 invertible")
A1 = numpy.array([[ConstantOrZero(aa) for aa in arow] for arow in bA1],
dtype=object)
if bA1inv is not None:
A1inv = numpy.array([[ConstantOrZero(aa) for aa in arow]
for arow in bA1inv],
dtype=object)
else:
A1inv = None
num_stages = butch.num_stages
num_fields = len(V)
Vbig = reduce(mul, (V for _ in range(num_stages)))
vnew = TestFunction(Vbig)
w = Function(Vbig)
if len(V) == 1:
u0bits = [u0]
vbits = [v]
if num_stages == 1:
vbigbits = [vnew]
wbits = [w]
else:
vbigbits = split(vnew)
wbits = split(w)
else:
u0bits = split(u0)
vbits = split(v)
vbigbits = split(vnew)
wbits = split(w)
wbits_np = numpy.zeros((num_stages, num_fields), dtype=object)
for i in range(num_stages):
for j in range(num_fields):
wbits_np[i, j] = wbits[i * num_fields + j]
A1w = A1 @ wbits_np
A2invw = A2inv @ wbits_np
Fnew = Zero()
for i in range(num_stages):
repl = {t: t + c[i] * dt}
for j, (ubit, vbit) in enumerate(zip(u0bits, vbits)):
repl[ubit] = ubit + dt * A1w[i, j]
repl[vbit] = vbigbits[num_fields * i + j]
repl[TimeDerivative(ubit)] = A2invw[i, j]
if (len(ubit.ufl_shape) == 1):
for kk in range(len(A1w[i, j])):
repl[TimeDerivative(ubit[kk])] = A2invw[i, j][kk]
repl[ubit[kk]] = repl[ubit][kk]
repl[vbit[kk]] = repl[vbit][kk]
Fnew += replace(F, repl)
bcnew = []
gblah = []
if bcs is None:
bcs = []
if bc_type == "ODE":
assert splitting == AI, "ODE-type BC aren't implemented for this splitting strategy"
u0_mult_np = numpy.divide(1.0,
butch.c,
out=numpy.zeros_like(butch.c),
where=butch.c != 0)
u0_mult = numpy.array([ConstantOrZero(mi) / dt for mi in u0_mult_np],
dtype=object)
def bc2gcur(bc, i):
gorig = bc._original_arg
gfoo = expand_derivatives(diff(gorig, t))
return replace(gfoo, {t: t + c[i] * dt}) + u0_mult[i] * gorig
elif bc_type == "DAE":
if bA1inv is None:
raise NotImplementedError(
"Cannot have DAE BCs for this Butcher Tableau/splitting")
u0_mult_np = A1inv @ numpy.ones_like(butch.c)
u0_mult = numpy.array([ConstantOrZero(mi) / dt for mi in u0_mult_np],
dtype=object)
def bc2gcur(bc, i):
gorig = as_ufl(bc._original_arg)
gcur = 0
for j in range(num_stages):
gcur += ConstantOrZero(bA1inv[i, j]) / dt * replace(
gorig, {t: t + c[j] * dt})
return gcur
else:
raise ValueError("Unrecognised bc_type: %s", bc_type)
# This logic uses information set up in the previous section to
# set up the new BCs for either method
for bc in bcs:
if num_fields == 1: # not mixed space
comp = bc.function_space().component
if comp is not None: # check for sub-piece of vector-valued
Vsp = V.sub(comp)
Vbigi = lambda i: Vbig[i].sub(comp)
else:
Vsp = V
Vbigi = lambda i: Vbig[i]
else: # mixed space
sub = bc.function_space_index()
comp = bc.function_space().component
if comp is not None: # check for sub-piece of vector-valued
Vsp = V.sub(sub).sub(comp)
Vbigi = lambda i: Vbig[sub + num_fields * i].sub(comp)
else:
Vsp = V.sub(sub)
Vbigi = lambda i: Vbig[sub + num_fields * i]
for i in range(num_stages):
gcur = bc2gcur(bc, i)
blah = BCStageData(Vsp, gcur, u0, u0_mult, i, t, dt)
gdat, gcr, gmethod = blah.gstuff
gblah.append((gdat, gcr, gmethod))
bcnew.append(DirichletBC(Vbigi(i), gdat, bc.sub_domain))
nspnew = getNullspace(V, Vbig, butch, nullspace)
return Fnew, w, bcnew, nspnew, gblah
def _setup_solver(self):
state = self.state # just cutting down line length a bit
Dt = state.timestepping.dt
beta_ = Dt*state.timestepping.alpha
cp = state.parameters.cp
mu = state.mu
Vu = state.spaces("HDiv")
Vtheta = state.spaces("HDiv_v")
Vrho = state.spaces("DG")
# Store time-stepping coefficients as UFL Constants
dt = Constant(Dt)
beta = Constant(beta_)
beta_cp = Constant(beta_ * cp)
# Split up the rhs vector (symbolically)
u_in, rho_in, theta_in = split(state.xrhs)
# Build the reduced function space for u,rho
M = MixedFunctionSpace((Vu, Vrho))
w, phi = TestFunctions(M)
u, rho = TrialFunctions(M)
n = FacetNormal(state.mesh)
# Get background fields
thetabar = state.fields("thetabar")
rhobar = state.fields("rhobar")
pibar = thermodynamics.pi(state.parameters, rhobar, thetabar)
pibar_rho = thermodynamics.pi_rho(state.parameters, rhobar, thetabar)
pibar_theta = thermodynamics.pi_theta(state.parameters, rhobar, thetabar)
# Analytical (approximate) elimination of theta
k = state.k # Upward pointing unit vector
theta = -dot(k, u)*dot(k, grad(thetabar))*beta + theta_in
# Only include theta' (rather than pi') in the vertical
# component of the gradient
# the pi prime term (here, bars are for mean and no bars are
# for linear perturbations)
pi = pibar_theta*theta + pibar_rho*rho
# vertical projection
def V(u):
return k*inner(u, k)
# specify degree for some terms as estimated degree is too large
dxp = dx(degree=(self.quadrature_degree))
dS_vp = dS_v(degree=(self.quadrature_degree))
# add effect of density of water upon theta
if self.moisture is not None:
water_t = Function(Vtheta).assign(0.0)
for water in self.moisture:
water_t += self.state.fields(water)
theta_w = theta / (1 + water_t)
thetabar_w = thetabar / (1 + water_t)
else:
theta_w = theta
thetabar_w = thetabar
eqn = (
inner(w, (state.h_project(u) - u_in))*dx
- beta_cp*div(theta_w*V(w))*pibar*dxp
# following does nothing but is preserved in the comments
# to remind us why (because V(w) is purely vertical).
# + beta_cp*jump(theta*V(w), n)*avg(pibar)*dS_v
- beta_cp*div(thetabar_w*w)*pi*dxp
+ beta_cp*jump(thetabar_w*w, n)*avg(pi)*dS_vp
+ (phi*(rho - rho_in) - beta*inner(grad(phi), u)*rhobar)*dx
+ beta*jump(phi*u, n)*avg(rhobar)*(dS_v + dS_h)
)
if mu is not None:
eqn += dt*mu*inner(w, k)*inner(u, k)*dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Place to put result of u rho solver
self.urho = Function(M)
# Boundary conditions (assumes extruded mesh)
bcs = [DirichletBC(M.sub(0), 0.0, "bottom"),
DirichletBC(M.sub(0), 0.0, "top")]
# Solver for u, rho
urho_problem = LinearVariationalProblem(
aeqn, Leqn, self.urho, bcs=bcs)
self.urho_solver = LinearVariationalSolver(urho_problem,
solver_parameters=self.solver_parameters,
options_prefix='ImplicitSolver')
# Reconstruction of theta
theta = TrialFunction(Vtheta)
gamma = TestFunction(Vtheta)
u, rho = self.urho.split()
self.theta = Function(Vtheta)
theta_eqn = gamma*(theta - theta_in
+ dot(k, u)*dot(k, grad(thetabar))*beta)*dx
theta_problem = LinearVariationalProblem(lhs(theta_eqn),
rhs(theta_eqn),
self.theta)
self.theta_solver = LinearVariationalSolver(theta_problem,
options_prefix='thetabacksubstitution')
def compressible_hydrostatic_balance(state, theta0, rho0, pi0=None,
top=False, pi_boundary=Constant(1.0),
water_t=None,
solve_for_rho=False,
params=None):
"""
Compute a hydrostatically balanced density given a potential temperature
profile.
:arg state: The :class:`State` object.
:arg theta0: :class:`.Function`containing the potential temperature.
:arg rho0: :class:`.Function` to write the initial density into.
:arg top: If True, set a boundary condition at the top. Otherwise, set
it at the bottom.
:arg pi_boundary: a field or expression to use as boundary data for pi on
the top or bottom as specified.
:arg water_t: the initial total water mixing ratio field.
"""
# Calculate hydrostatic Pi
VDG = state.spaces("DG")
Vv = state.spaces("Vv")
W = MixedFunctionSpace((Vv, VDG))
v, pi = TrialFunctions(W)
dv, dpi = TestFunctions(W)
n = FacetNormal(state.mesh)
cp = state.parameters.cp
# add effect of density of water upon theta
theta = theta0
if water_t is not None:
theta = theta0 / (1 + water_t)
alhs = (
(cp*inner(v, dv) - cp*div(dv*theta)*pi)*dx
+ dpi*div(theta*v)*dx
)
if top:
bmeasure = ds_t
bstring = "bottom"
else:
bmeasure = ds_b
bstring = "top"
arhs = -cp*inner(dv, n)*theta*pi_boundary*bmeasure
g = state.parameters.g
arhs -= g*inner(dv, state.k)*dx
bcs = [DirichletBC(W.sub(0), 0.0, bstring)]
w = Function(W)
PiProblem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs)
if params is None:
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': 'gamg',
'fieldsplit_1_pc_gamg_sym_graph': True,
'fieldsplit_1_mg_levels_ksp_type': 'chebyshev',
'fieldsplit_1_mg_levels_ksp_chebyshev_esteig': True,
'fieldsplit_1_mg_levels_ksp_max_it': 5,
'fieldsplit_1_mg_levels_pc_type': 'bjacobi',
'fieldsplit_1_mg_levels_sub_pc_type': 'ilu'}
PiSolver = LinearVariationalSolver(PiProblem,
solver_parameters=params)
PiSolver.solve()
v, Pi = w.split()
if pi0 is not None:
pi0.assign(Pi)
if solve_for_rho:
w1 = Function(W)
v, rho = w1.split()
rho.interpolate(thermodynamics.rho(state.parameters, theta0, Pi))
v, rho = split(w1)
dv, dpi = TestFunctions(W)
pi = thermodynamics.pi(state.parameters, rho, theta0)
F = (
(cp*inner(v, dv) - cp*div(dv*theta)*pi)*dx
+ dpi*div(theta0*v)*dx
+ cp*inner(dv, n)*theta*pi_boundary*bmeasure
)
F += g*inner(dv, state.k)*dx
rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs)
rhosolver = NonlinearVariationalSolver(rhoproblem, solver_parameters=params)
rhosolver.solve()
v, rho_ = w1.split()
rho0.assign(rho_)
else:
rho0.interpolate(thermodynamics.rho(state.parameters, theta0, Pi))
def split(self, fields):
from firedrake import replace, as_vector, split
from firedrake_ts.ts_solver import DAEProblem
from firedrake.bcs import DirichletBC, EquationBC
fields = tuple(tuple(f) for f in fields)
splits = self._splits.get(tuple(fields))
if splits is not None:
return splits
splits = []
problem = self._problem
splitter = ExtractSubBlock()
for field in fields:
F = splitter.split(problem.F, argument_indices=(field, ))
J = splitter.split(problem.J, argument_indices=(field, field))
us = problem.u.split()
V = F.arguments()[0].function_space()
# Exposition:
# We are going to make a new solution Function on the sub
# mixed space defined by the relevant fields.
# But the form may refer to the rest of the solution
# anyway.
# So we pull it apart and will make a new function on the
# subspace that shares data.
pieces = [us[i].dat for i in field]
if len(pieces) == 1:
(val, ) = pieces
subu = function.Function(V, val=val)
subsplit = (subu, )
else:
val = op2.MixedDat(pieces)
subu = function.Function(V, val=val)
# Split it apart to shove in the form.
subsplit = split(subu)
# Permutation from field indexing to indexing of pieces
field_renumbering = dict([f, i] for i, f in enumerate(field))
vec = []
for i, u in enumerate(us):
if i in field:
# If this is a field we're keeping, get it from
# the new function. Otherwise just point to the
# old data.
u = subsplit[field_renumbering[i]]
if u.ufl_shape == ():
vec.append(u)
else:
for idx in numpy.ndindex(u.ufl_shape):
vec.append(u[idx])
# So now we have a new representation for the solution
# vector in the old problem. For the fields we're going
# to solve for, it points to a new Function (which wraps
# the original pieces). For the rest, it points to the
# pieces from the original Function.
# IOW, we've reinterpreted our original mixed solution
# function as being made up of some spaces we're still
# solving for, and some spaces that have just become
# coefficients in the new form.
u = as_vector(vec)
F = replace(F, {problem.u: u})
J = replace(J, {problem.u: u})
if problem.Jp is not None:
Jp = splitter.split(problem.Jp,
argument_indices=(field, field))
Jp = replace(Jp, {problem.u: u})
else:
Jp = None
bcs = []
for bc in problem.bcs:
if isinstance(bc, DirichletBC):
bc_temp = bc.reconstruct(
field=field,
V=V,
g=bc.function_arg,
sub_domain=bc.sub_domain,
method=bc.method,
)
elif isinstance(bc, EquationBC):
bc_temp = bc.reconstruct(field, V, subu, u)
if bc_temp is not None:
bcs.append(bc_temp)
new_problem = DAEProblem(
F,
subu,
problem.udot,
problem.tspan,
bcs=bcs,
J=J,
Jp=Jp,
form_compiler_parameters=problem.form_compiler_parameters,
)
new_problem._constant_jacobian = problem._constant_jacobian
splits.append(
type(self)(
new_problem,
mat_type=self.mat_type,
pmat_type=self.pmat_type,
appctx=self.appctx,
transfer_manager=self.transfer_manager,
))
return self._splits.setdefault(tuple(fields), splits)
def gravity_term(self):
_, _, b0 = split(self.x0)
L = b0 * inner(self.test, self.state.k) * dx
return L
def repl(t):
"""
Function returned by replace_subject to return a new :class:`Term` with
the subject replaced by the variable `new`. It is built around the ufl
replace routine.
Returns a new :class:`Term`.
:arg t: the original :class:`Term`.
"""
subj = t.get(subject)
# Build a dictionary to pass to the ufl.replace routine
# The dictionary matches variables in the old term with those in the new
replace_dict = {}
# Consider cases that subj is normal Function or MixedFunction
# vs cases of new being Function vs MixedFunction vs tuple
# Ideally catch all cases or fail gracefully
if type(subj.ufl_element()) is MixedElement:
if type(new) == tuple:
assert len(new) == len(subj.function_space())
for k, v in zip(split(subj), new):
replace_dict[k] = v
elif type(new) == ufl.algebra.Sum:
replace_dict[subj] = new
# Otherwise fail if new is not a function
elif not isinstance(new, Function):
raise ValueError(
f'new must be a tuple or Function, not type {type(new)}')
# Now handle MixedElements separately as these need indexing
elif type(new.ufl_element()) is MixedElement:
assert len(new.function_space()) == len(subj.function_space())
# If idx specified, replace only that component
if idx is not None:
replace_dict[split(subj)[idx]] = split(new)[idx]
# Otherwise replace all components
else:
for k, v in zip(split(subj), split(new)):
replace_dict[k] = v
# Otherwise 'new' is a normal Function
else:
replace_dict[split(subj)[idx]] = new
# subj is a normal Function
else:
if type(new) is tuple:
if idx is None:
raise ValueError(
'idx must be specified to replace_subject' +
' when new is a tuple')
replace_dict[subj] = new[idx]
elif not isinstance(new, Function):
raise ValueError(
f'new must be a Function, not type {type(new)}')
elif type(new.ufl_element()) == MixedElement:
if idx is None:
raise ValueError(
'idx must be specified to replace_subject' +
' when new is a tuple')
replace_dict[subj] = split(new)[idx]
else:
replace_dict[subj] = new
new_form = ufl.replace(t.form, replace_dict)
return Term(new_form, t.labels)
def coriolis_term(self):
u0 = split(self.x0)[0]
return -inner(self.test, cross(2 * self.state.Omega, u0)) * dx
def split(self, fields):
from firedrake import replace, as_vector, split
from firedrake import NonlinearVariationalProblem as NLVP
fields = tuple(tuple(f) for f in fields)
splits = self._splits.get(tuple(fields))
if splits is not None:
return splits
splits = []
problem = self._problem
splitter = ExtractSubBlock()
for field in fields:
F = splitter.split(problem.F, argument_indices=(field, ))
J = splitter.split(problem.J, argument_indices=(field, field))
us = problem.u.split()
V = F.arguments()[0].function_space()
# Exposition:
# We are going to make a new solution Function on the sub
# mixed space defined by the relevant fields.
# But the form may refer to the rest of the solution
# anyway.
# So we pull it apart and will make a new function on the
# subspace that shares data.
pieces = [us[i].dat for i in field]
if len(pieces) == 1:
val, = pieces
subu = function.Function(V, val=val)
subsplit = (subu, )
else:
val = op2.MixedDat(pieces)
subu = function.Function(V, val=val)
# Split it apart to shove in the form.
subsplit = split(subu)
# Permutation from field indexing to indexing of pieces
field_renumbering = dict([f, i] for i, f in enumerate(field))
vec = []
for i, u in enumerate(us):
if i in field:
# If this is a field we're keeping, get it from
# the new function. Otherwise just point to the
# old data.
u = subsplit[field_renumbering[i]]
if u.ufl_shape == ():
vec.append(u)
else:
for idx in numpy.ndindex(u.ufl_shape):
vec.append(u[idx])
# So now we have a new representation for the solution
# vector in the old problem. For the fields we're going
# to solve for, it points to a new Function (which wraps
# the original pieces). For the rest, it points to the
# pieces from the original Function.
# IOW, we've reinterpreted our original mixed solution
# function as being made up of some spaces we're still
# solving for, and some spaces that have just become
# coefficients in the new form.
u = as_vector(vec)
F = replace(F, {problem.u: u})
J = replace(J, {problem.u: u})
if problem.Jp is not None:
Jp = splitter.split(problem.Jp, argument_indices=(field, field))
Jp = replace(Jp, {problem.u: u})
else:
Jp = None
bcs = []
for bc in problem.bcs:
Vbc = bc.function_space()
if Vbc.parent is not None and isinstance(Vbc.parent.ufl_element(), VectorElement):
index = Vbc.parent.index
else:
index = Vbc.index
cmpt = Vbc.component
# TODO: need to test this logic
if index in field:
if len(field) == 1:
W = V
else:
W = V.sub(field_renumbering[index])
if cmpt is not None:
W = W.sub(cmpt)
bcs.append(type(bc)(W,
bc.function_arg,
bc.sub_domain,
method=bc.method))
new_problem = NLVP(F, subu, bcs=bcs, J=J, Jp=Jp,
form_compiler_parameters=problem.form_compiler_parameters)
new_problem._constant_jacobian = problem._constant_jacobian
splits.append(type(self)(new_problem, mat_type=self.mat_type, pmat_type=self.pmat_type,
appctx=self.appctx))
return self._splits.setdefault(tuple(fields), splits)
def euler_poincare_term(self):
u0 = split(self.x0)[0]
return -0.5 * div(self.test) * inner(self.state.h_project(u0), u0) * dx
def pressure_gradient_term(self):
_, p0, _ = split(self.x0)
L = div(self.test) * p0 * dx
return L
def initialize(self, pc):
from firedrake import TrialFunction, TestFunction, Function, DirichletBC, dx, \
Mesh, inner, grad, split, Constant, parameters
from firedrake.assemble import allocate_matrix, create_assembly_callable
prefix = pc.getOptionsPrefix() + "pcd_"
_, P = pc.getOperators()
context = P.getPythonContext()
test, trial = context.a.arguments()
Q = test.function_space()
self.Q = Q
p = TrialFunction(Q)
q = TestFunction(Q)
nu = context.appctx.get("nu", 1.0)
self.nu = nu
mass = Constant(1.0 / self.nu) * p * q * dx
stiffness = inner(grad(p), grad(q)) * dx
state = context.appctx["state"]
velid = context.appctx["velocity_space"]
opts = PETSc.Options()
default = parameters["default_matrix_type"]
Mp_mat_type = opts.getString(prefix + "Mp_mat_type", default)
Kp_mat_type = opts.getString(prefix + "Kp_mat_type", default)
self.Fp_mat_type = opts.getString(prefix + "Fp_mat_type", "matfree")
Mp = assemble(mass,
form_compiler_parameters=context.fc_params,
mat_type=Mp_mat_type,
options_prefix=prefix + "Mp_")
Kp = assemble(stiffness,
form_compiler_parameters=context.fc_params,
mat_type=Kp_mat_type,
options_prefix=prefix + "Kp_")
Mksp = PETSc.KSP().create(comm=pc.comm)
Mksp.incrementTabLevel(1, parent=pc)
Mksp.setOptionsPrefix(prefix + "Mp_")
Mksp.setOperators(Mp.petscmat)
Mksp.setUp()
Mksp.setFromOptions()
self.Mksp = Mksp
Kksp = PETSc.KSP().create(comm=pc.comm)
Kksp.incrementTabLevel(1, parent=pc)
Kksp.setOptionsPrefix(prefix + "Kp_")
Kksp.setOperators(Kp.petscmat)
Kksp.setUp()
Kksp.setFromOptions()
self.Kksp = Kksp
u0 = split(state)[velid]
fp = Constant(self.nu) * inner(grad(p), grad(q)) * dx + inner(
u0, grad(p)) * q * dx
self.Fp = allocate_matrix(fp,
form_compiler_parameters=context.fc_params,
mat_type=self.Fp_mat_type,
options_prefix=prefix + "Fp_")
self._assemble_Fp = create_assembly_callable(
fp,
tensor=self.Fp,
form_compiler_parameters=context.fc_params,
mat_type=self.Fp_mat_type)
self._assemble_Fp()
Fpmat = self.Fp.petscmat
self.workspace = [Fpmat.createVecLeft() for i in (0, 1)]
self.tmp = Function(self.Q)
def coriolis_term(self):
f = self.state.fields("coriolis")
u0, _ = split(self.x0)
L = -f * inner(self.test, self.state.perp(u0)) * dx
return L
def getFormStage(F,
butch,
u0,
t,
dt,
bcs=None,
splitting=None,
nullspace=None):
"""Given a time-dependent variational form and a
:class:`ButcherTableau`, produce UFL for the s-stage RK method.
:arg F: UFL form for the semidiscrete ODE/DAE
:arg butch: the :class:`ButcherTableau` for the RK method being used to
advance in time.
:arg u0: a :class:`Function` referring to the state of
the PDE system at time `t`
:arg t: a :class:`Constant` referring to the current time level.
Any explicit time-dependence in F is included
:arg dt: a :class:`Constant` referring to the size of the current
time step.
:arg splitting: a callable that maps the (floating point) Butcher matrix
a to a pair of matrices `A1, A2` such that `butch.A = A1 A2`. This is used
to vary between the classical RK formulation and Butcher's reformulation
that leads to a denser mass matrix with block-diagonal stiffness.
Only `AI` and `IA` are currently supported.
:arg bcs: optionally, a :class:`DirichletBC` object (or iterable thereof)
containing (possible time-dependent) boundary conditions imposed
on the system.
:arg nullspace: A list of tuples of the form (index, VSB) where
index is an index into the function space associated with `u`
and VSB is a :class: `firedrake.VectorSpaceBasis` instance to
be passed to a `firedrake.MixedVectorSpaceBasis` over the
larger space associated with the Runge-Kutta method
On output, we return a tuple consisting of several parts:
- Fnew, the :class:`Form`
- possibly a 4-tuple containing information needed to solve a mass matrix to update
the solution (this is empty for RadauIIA methods for which there is a trivial
update function.
- UU, the :class:`firedrake.Function` holding all the stage time values.
It lives in a :class:`firedrake.FunctionSpace` corresponding to the
s-way tensor product of the space on which the semidiscrete
form lives.
- `bcnew`, a list of :class:`firedrake.DirichletBC` objects to be posed
on the stages,
- 'nspnew', the :class:`firedrake.MixedVectorSpaceBasis` object
that represents the nullspace of the coupled system
- `gblah`, a list of tuples of the form (f, expr, method),
where f is a :class:`firedrake.Function` and expr is a
:class:`ufl.Expr`. At each time step, each expr needs to be
re-interpolated/projected onto the corresponding f in order
for Firedrake to pick up that time-dependent boundary
conditions need to be re-applied. The
interpolation/projection is encoded in method, which is
either `f.interpolate(expr-c*u0)` or `f.project(expr-c*u0)`, depending
on whether the function space for f supports interpolation or
not.
"""
v = F.arguments()[0]
V = v.function_space()
assert V == u0.function_space()
num_stages = butch.num_stages
num_fields = len(V)
# s-way product space for the stage variables
Vbig = reduce(mul, (V for _ in range(num_stages)))
VV = TestFunction(Vbig)
UU = Function(Vbig)
vecconst = np.vectorize(Constant)
C = vecconst(butch.c)
A = vecconst(butch.A)
# set up the pieces we need to work with to do our substitutions
nsxnf = (num_stages, num_fields)
if num_fields == 1:
u0bits = np.array([u0], dtype="O")
vbits = np.array([v], dtype="O")
if num_stages == 1: # single-stage method
VVbits = np.array([[VV]], dtype="O")
UUbits = np.array([[UU]], dtype="O")
else: # multi-stage methods
VVbits = np.reshape(split(VV), nsxnf)
UUbits = np.reshape(split(UU), nsxnf)
else:
u0bits = np.array(list(split(u0)), dtype="O")
vbits = np.array(list(split(v)), dtype="O")
VVbits = np.reshape(split(VV), nsxnf)
UUbits = np.reshape(split(UU), nsxnf)
split_form = extract_terms(F)
Fnew = Zero()
# first, process terms with a time derivative. I'm
# assuming we have something of the form inner(Dt(g(u0)), v)*dx
# For each stage i, this gets replaced with
# inner((g(stages[i]) - g(u0))/dt, v)*dx
# but we have to carefully handle the cases where g indexes into
# pieces of u
dtless = strip_dt_form(split_form.time)
if splitting is None or splitting == AI:
for i in range(num_stages):
repl = {t: t + C[i] * dt}
for j in range(num_fields):
repl[u0bits[j]] = UUbits[i][j] - u0bits[j]
repl[vbits[j]] = VVbits[i][j]
# Also get replacements right for indexing.
for j in range(num_fields):
for ii in np.ndindex(u0bits[j].ufl_shape):
repl[u0bits[j][ii]] = UUbits[i][j][ii] - u0bits[j][ii]
repl[vbits[j][ii]] = VVbits[i][j][ii]
Fnew += replace(dtless, repl)
# Now for the non-time derivative parts
for i in range(num_stages):
# replace test function
repl = {}
for k in range(num_fields):
repl[vbits[k]] = VVbits[i][k]
for ii in np.ndindex(vbits[k].ufl_shape):
repl[vbits[k][ii]] = VVbits[i][k][ii]
Ftmp = replace(split_form.remainder, repl)
# replace the solution with stage values
for j in range(num_stages):
repl = {t: t + C[j] * dt}
for k in range(num_fields):
repl[u0bits[k]] = UUbits[j][k]
for ii in np.ndindex(u0bits[k].ufl_shape):
repl[u0bits[k][ii]] = UUbits[j][k][ii]
# and sum the contribution
Fnew += A[i, j] * dt * replace(Ftmp, repl)
elif splitting == IA:
Ainv = np.vectorize(Constant)(np.linalg.inv(butch.A))
# time derivative part gets inverse of Butcher matrix.
for i in range(num_stages):
repl = {}
for k in range(num_fields):
repl[vbits[k]] = VVbits[i][k]
for ii in np.ndindex(vbits[k].ufl_shape):
repl[vbits[k][ii]] = VVbits[i][k][ii]
Ftmp = replace(dtless, repl)
for j in range(num_stages):
repl = {t: t + C[j] * dt}
for k in range(num_fields):
repl[u0bits[k]] = (UUbits[j][k] - u0bits[k])
for ii in np.ndindex(u0bits[k].ufl_shape):
repl[u0bits[k][ii]] = UUbits[j][k][ii] - u0bits[k][ii]
Fnew += Ainv[i, j] * replace(Ftmp, repl)
# rest of the operator: just diagonal!
for i in range(num_stages):
repl = {t: t + C[i] * dt}
for j in range(num_fields):
repl[u0bits[j]] = UUbits[i][j]
repl[vbits[j]] = VVbits[i][j]
# Also get replacements right for indexing.
for j in range(num_fields):
for ii in np.ndindex(u0bits[j].ufl_shape):
repl[u0bits[j][ii]] = UUbits[i][j][ii]
repl[vbits[j][ii]] = VVbits[i][j][ii]
Fnew += dt * replace(split_form.remainder, repl)
else:
raise NotImplementedError("Can't handle that splitting type")
if bcs is None:
bcs = []
bcsnew = []
gblah = []
# For each BC, we need a new BC for each stage
# so we need to figure out how the function is indexed (mixed + vec)
# and then set it to have the value of the original argument at
# time t+C[i]*dt.
for bc in bcs:
if num_fields == 1: # not mixed space
comp = bc.function_space().component
if comp is not None: # check for sub-piece of vector-valued
Vsp = V.sub(comp)
Vbigi = lambda i: Vbig[i].sub(comp)
else:
Vsp = V
Vbigi = lambda i: Vbig[i]
else: # mixed space
sub = bc.function_space_index()
comp = bc.function_space().component
if comp is not None: # check for sub-piece of vector-valued
Vsp = V.sub(sub).sub(comp)
Vbigi = lambda i: Vbig[sub + num_fields * i].sub(comp)
else:
Vsp = V.sub(sub)
Vbigi = lambda i: Vbig[sub + num_fields * i]
bcarg = bc._original_arg
for i in range(num_stages):
try:
gdat = interpolate(bcarg, Vsp)
gmethod = lambda gd, gc: gd.interpolate(gc)
except: # noqa: E722
gdat = project(bcarg, Vsp)
gmethod = lambda gd, gc: gd.project(gc)
gcur = replace(bcarg, {t: t + C[i] * dt})
bcsnew.append(DirichletBC(Vbigi(i), gdat, bc.sub_domain))
gblah.append((gdat, gcur, gmethod))
nspacenew = getNullspace(V, Vbig, butch, nullspace)
# For RIIA, we have an optimized update rule and don't need to
# build the variational form for doing updates.
# But something's broken with null spaces, so that's a TO-DO.
unew = Function(V)
Fupdate = inner(unew - u0, v) * dx
B = vectorize(Constant)(butch.b)
C = vectorize(Constant)(butch.c)
for i in range(num_stages):
repl = {t: t + C[i] * dt}
for k in range(num_fields):
repl[u0bits[k]] = UUbits[i][k]
for ii in np.ndindex(u0bits[k].ufl_shape):
repl[u0bits[k][ii]] = UUbits[i][k][ii]
eFFi = replace(split_form.remainder, repl)
Fupdate += dt * B[i] * eFFi
# And the BC's for the update -- just the original BC at t+dt
update_bcs = []
update_bcs_gblah = []
for bc in bcs:
if num_fields == 1: # not mixed space
comp = bc.function_space().component
if comp is not None: # check for sub-piece of vector-valued
Vsp = V.sub(comp)
else:
Vsp = V
else: # mixed space
sub = bc.function_space_index()
comp = bc.function_space().component
if comp is not None: # check for sub-piece of vector-valued
Vsp = V.sub(sub).sub(comp)
else:
Vsp = V.sub(sub)
bcarg = bc._original_arg
try:
gdat = interpolate(bcarg, Vsp)
gmethod = lambda gd, gc: gd.interpolate(gc)
except: # noqa: E722
gdat = project(bcarg, Vsp)
gmethod = lambda gd, gc: gd.project(gc)
gcur = replace(bcarg, {t: t + dt})
update_bcs.append(DirichletBC(Vsp, gdat, bc.sub_domain))
update_bcs_gblah.append((gdat, gcur, gmethod))
return (Fnew, (unew, Fupdate, update_bcs, update_bcs_gblah), UU, bcsnew,
gblah, nspacenew)
def sponge_term(self):
u0 = split(self.x0)[0]
return self.state.mu * inner(self.test, self.state.k) * inner(
u0, self.state.k) * dx
def sdhm_form(self, problem, mesh, bcs_p, bcs_u):
rho = problem.rho
mu = problem.mu
k = problem.k
f = problem.f
q, p, lambda_h = fire.split(self.solution)
w, v, mu_h = fire.TestFunctions(self._W)
n = fire.FacetNormal(mesh)
h = fire.CellDiameter(mesh)
# Stabilizing parameters
has_mesh_characteristic_length = True
beta_0 = fire.Constant(1e-15)
delta_0 = fire.Constant(1)
delta_1 = fire.Constant(-1 / 2)
delta_2 = fire.Constant(1 / 2)
delta_3 = fire.Constant(1 / 2)
# h_avg = (h('+') + h('-')) / 2.
beta = beta_0 / h
beta_avg = beta_0 / h("+")
if has_mesh_characteristic_length:
delta_2 = delta_2 * h * h
delta_3 = delta_3 * h * h
kappa = rho * k / mu
inv_kappa = 1.0 / kappa
# Classical mixed terms
a = (dot(inv_kappa * q, w) - div(w) * p - delta_0 * v * div(q)) * dx
L = -delta_0 * f * v * dx
# Hybridization terms
a += lambda_h("+") * dot(w, n)("+") * dS + mu_h("+") * dot(q,
n)("+") * dS
a += beta_avg * kappa("+") * (lambda_h("+") - p("+")) * (mu_h("+") -
v("+")) * dS
# Add the contributions of the pressure boundary conditions to L
primal_bc_markers = list(mesh.exterior_facets.unique_markers)
for pboundary, iboundary in bcs_p:
primal_bc_markers.remove(iboundary)
a += (pboundary * dot(w, n) + mu_h * dot(q, n)) * ds(iboundary)
a += beta * kappa * (lambda_h - pboundary) * mu_h * ds(iboundary)
unprescribed_primal_bc = primal_bc_markers
for bc_marker in unprescribed_primal_bc:
a += (lambda_h * dot(w, n) + mu_h * dot(q, n)) * ds(bc_marker)
a += beta * kappa * lambda_h * mu_h * ds(bc_marker)
# Add the (weak) contributions of the velocity boundary conditions to L
for uboundary, iboundary, component in bcs_u:
if component is not None:
dim = mesh.geometric_dimension()
bc_array = []
for _ in range(dim):
bc_array.append(0.0)
bc_array[component] = uboundary
bc_as_vector = fire.Constant(bc_array)
L += mu_h * dot(bc_as_vector, n) * ds(iboundary)
else:
L += mu_h * dot(uboundary, n) * ds(iboundary)
# Stabilizing terms
a += (delta_1 * inner(kappa * (inv_kappa * q + grad(p)),
delta_0 * inv_kappa * w + grad(v)) * dx)
a += delta_2 * inv_kappa * div(q) * div(w) * dx
a += delta_3 * inner(kappa * curl(inv_kappa * q), curl(
inv_kappa * w)) * dx
L += delta_2 * inv_kappa * f * div(w) * dx
return a, L
def hydrostatic_term(self):
u0 = split(self.x0)[0]
return inner(u0, self.state.k) * inner(self.test, self.state.k) * dx
def _setup_solver(self):
import numpy as np
state = self.state
dt = state.dt
beta_ = dt * self.alpha
cp = state.parameters.cp
Vu = state.spaces("HDiv")
Vu_broken = FunctionSpace(state.mesh, BrokenElement(Vu.ufl_element()))
Vtheta = state.spaces("theta")
Vrho = state.spaces("DG")
# Store time-stepping coefficients as UFL Constants
beta = Constant(beta_)
beta_cp = Constant(beta_ * cp)
h_deg = Vrho.ufl_element().degree()[0]
v_deg = Vrho.ufl_element().degree()[1]
Vtrace = FunctionSpace(state.mesh, "HDiv Trace", degree=(h_deg, v_deg))
# Split up the rhs vector (symbolically)
self.xrhs = Function(self.equations.function_space)
u_in, rho_in, theta_in = split(self.xrhs)[0:3]
# Build the function space for "broken" u, rho, and pressure trace
M = MixedFunctionSpace((Vu_broken, Vrho, Vtrace))
w, phi, dl = TestFunctions(M)
u, rho, l0 = TrialFunctions(M)
n = FacetNormal(state.mesh)
# Get background fields
thetabar = state.fields("thetabar")
rhobar = state.fields("rhobar")
exnerbar = thermodynamics.exner_pressure(state.parameters, rhobar,
thetabar)
exnerbar_rho = thermodynamics.dexner_drho(state.parameters, rhobar,
thetabar)
exnerbar_theta = thermodynamics.dexner_dtheta(state.parameters, rhobar,
thetabar)
# Analytical (approximate) elimination of theta
k = state.k # Upward pointing unit vector
theta = -dot(k, u) * dot(k, grad(thetabar)) * beta + theta_in
# Only include theta' (rather than exner') in the vertical
# component of the gradient
# The exner prime term (here, bars are for mean and no bars are
# for linear perturbations)
exner = exnerbar_theta * theta + exnerbar_rho * rho
# Vertical projection
def V(u):
return k * inner(u, k)
# hydrostatic projection
h_project = lambda u: u - k * inner(u, k)
# Specify degree for some terms as estimated degree is too large
dxp = dx(degree=(self.quadrature_degree))
dS_vp = dS_v(degree=(self.quadrature_degree))
dS_hp = dS_h(degree=(self.quadrature_degree))
ds_vp = ds_v(degree=(self.quadrature_degree))
ds_tbp = (ds_t(degree=(self.quadrature_degree)) +
ds_b(degree=(self.quadrature_degree)))
# Add effect of density of water upon theta
if self.moisture is not None:
water_t = Function(Vtheta).assign(0.0)
for water in self.moisture:
water_t += self.state.fields(water)
theta_w = theta / (1 + water_t)
thetabar_w = thetabar / (1 + water_t)
else:
theta_w = theta
thetabar_w = thetabar
_l0 = TrialFunction(Vtrace)
_dl = TestFunction(Vtrace)
a_tr = _dl('+') * _l0('+') * (
dS_vp + dS_hp) + _dl * _l0 * ds_vp + _dl * _l0 * ds_tbp
def L_tr(f):
return _dl('+') * avg(f) * (
dS_vp + dS_hp) + _dl * f * ds_vp + _dl * f * ds_tbp
cg_ilu_parameters = {
'ksp_type': 'cg',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
}
# Project field averages into functions on the trace space
rhobar_avg = Function(Vtrace)
exnerbar_avg = Function(Vtrace)
rho_avg_prb = LinearVariationalProblem(a_tr, L_tr(rhobar), rhobar_avg)
exner_avg_prb = LinearVariationalProblem(a_tr, L_tr(exnerbar),
exnerbar_avg)
rho_avg_solver = LinearVariationalSolver(
rho_avg_prb,
solver_parameters=cg_ilu_parameters,
options_prefix='rhobar_avg_solver')
exner_avg_solver = LinearVariationalSolver(
exner_avg_prb,
solver_parameters=cg_ilu_parameters,
options_prefix='exnerbar_avg_solver')
with timed_region("Gusto:HybridProjectRhobar"):
rho_avg_solver.solve()
with timed_region("Gusto:HybridProjectExnerbar"):
exner_avg_solver.solve()
# "broken" u, rho, and trace system
# NOTE: no ds_v integrals since equations are defined on
# a periodic (or sphere) base mesh.
if any([t.has_label(hydrostatic) for t in self.equations.residual]):
u_mass = inner(w, (h_project(u) - u_in)) * dx
else:
u_mass = inner(w, (u - u_in)) * dx
eqn = (
# momentum equation
u_mass - beta_cp * div(theta_w * V(w)) * exnerbar * dxp
# following does nothing but is preserved in the comments
# to remind us why (because V(w) is purely vertical).
# + beta_cp*jump(theta_w*V(w), n=n)*exnerbar_avg('+')*dS_vp
+ beta_cp * jump(theta_w * V(w), n=n) * exnerbar_avg('+') * dS_hp +
beta_cp * dot(theta_w * V(w), n) * exnerbar_avg * ds_tbp -
beta_cp * div(thetabar_w * w) * exner * dxp
# trace terms appearing after integrating momentum equation
+ beta_cp * jump(thetabar_w * w, n=n) * l0('+') * (dS_vp + dS_hp) +
beta_cp * dot(thetabar_w * w, n) * l0 * (ds_tbp + ds_vp)
# mass continuity equation
+ (phi *
(rho - rho_in) - beta * inner(grad(phi), u) * rhobar) * dx +
beta * jump(phi * u, n=n) * rhobar_avg('+') * (dS_v + dS_h)
# term added because u.n=0 is enforced weakly via the traces
+ beta * phi * dot(u, n) * rhobar_avg * (ds_tb + ds_v)
# constraint equation to enforce continuity of the velocity
# through the interior facets and weakly impose the no-slip
# condition
+ dl('+') * jump(u, n=n) * (dS_vp + dS_hp) + dl * dot(u, n) *
(ds_tbp + ds_vp))
# contribution of the sponge term
if hasattr(self.equations, "mu"):
eqn += dt * self.equations.mu * inner(w, k) * inner(u, k) * dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Function for the hybridized solutions
self.urhol0 = Function(M)
hybridized_prb = LinearVariationalProblem(aeqn, Leqn, self.urhol0)
hybridized_solver = LinearVariationalSolver(
hybridized_prb,
solver_parameters=self.solver_parameters,
options_prefix='ImplicitSolver')
self.hybridized_solver = hybridized_solver
# Project broken u into the HDiv space using facet averaging.
# Weight function counting the dofs of the HDiv element:
shapes = {
"i": Vu.finat_element.space_dimension(),
"j": np.prod(Vu.shape, dtype=int)
}
weight_kernel = """
for (int i=0; i<{i}; ++i)
for (int j=0; j<{j}; ++j)
w[i*{j} + j] += 1.0;
""".format(**shapes)
self._weight = Function(Vu)
par_loop(weight_kernel, dx, {"w": (self._weight, INC)})
# Averaging kernel
self._average_kernel = """
for (int i=0; i<{i}; ++i)
for (int j=0; j<{j}; ++j)
vec_out[i*{j} + j] += vec_in[i*{j} + j]/w[i*{j} + j];
""".format(**shapes)
# HDiv-conforming velocity
self.u_hdiv = Function(Vu)
# Reconstruction of theta
theta = TrialFunction(Vtheta)
gamma = TestFunction(Vtheta)
self.theta = Function(Vtheta)
theta_eqn = gamma * (theta - theta_in + dot(k, self.u_hdiv) *
dot(k, grad(thetabar)) * beta) * dx
theta_problem = LinearVariationalProblem(lhs(theta_eqn),
rhs(theta_eqn), self.theta)
self.theta_solver = LinearVariationalSolver(
theta_problem,
solver_parameters=cg_ilu_parameters,
options_prefix='thetabacksubstitution')
# Store boundary conditions for the div-conforming velocity to apply
# post-solve
self.bcs = self.equations.bcs['u']
def _setup_solver(self):
state = self.state # just cutting down line length a bit
dt = state.timestepping.dt
beta = dt*state.timestepping.alpha
mu = state.mu
# Split up the rhs vector (symbolically)
u_in, p_in, b_in = split(state.xrhs)
# Build the reduced function space for u,p
M = MixedFunctionSpace((state.V[0], state.V[1]))
w, phi = TestFunctions(M)
u, p = TrialFunctions(M)
# Get background fields
bbar = state.bbar
# Analytical (approximate) elimination of theta
k = state.k # Upward pointing unit vector
b = -dot(k,u)*dot(k,grad(bbar))*beta + b_in
# vertical projection
def V(u):
return k*inner(u,k)
eqn = (
inner(w, (u - u_in))*dx
- beta*div(w)*p*dx
- beta*inner(w,k)*b*dx
+ phi*div(u)*dx
)
if mu is not None:
eqn += dt*mu*inner(w,k)*inner(u,k)*dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Place to put result of u p solver
self.up = Function(M)
# Boundary conditions (assumes extruded mesh)
dim = M.sub(0).ufl_element().value_shape()[0]
bc = ("0.0",)*dim
bcs = [DirichletBC(M.sub(0), Expression(bc), "bottom"),
DirichletBC(M.sub(0), Expression(bc), "top")]
# preconditioner equation
L = self.L
Ap = (
inner(w,u) + L*L*div(w)*div(u) +
phi*p/L/L
)*dx
# Solver for u, p
up_problem = LinearVariationalProblem(
aeqn, Leqn, self.up, bcs=bcs, aP=Ap)
nullspace = MixedVectorSpaceBasis(M,
[M.sub(0),
VectorSpaceBasis(constant=True)])
self.up_solver = LinearVariationalSolver(up_problem,
solver_parameters=self.params,
nullspace=nullspace)
# Reconstruction of b
b = TrialFunction(state.V[2])
gamma = TestFunction(state.V[2])
u, p = self.up.split()
self.b = Function(state.V[2])
b_eqn = gamma*(b - b_in +
dot(k,u)*dot(k,grad(bbar))*beta)*dx
b_problem = LinearVariationalProblem(lhs(b_eqn),
rhs(b_eqn),
self.b)
self.b_solver = LinearVariationalSolver(b_problem)
def _setup_solver(self):
state = self.state # just cutting down line length a bit
dt = state.dt
beta_ = dt * self.alpha
Vu = state.spaces("HDiv")
Vb = state.spaces("theta")
Vp = state.spaces("DG")
# Store time-stepping coefficients as UFL Constants
beta = Constant(beta_)
# Split up the rhs vector (symbolically)
self.xrhs = Function(self.equations.function_space)
u_in, p_in, b_in = split(self.xrhs)
# Build the reduced function space for u,p
M = MixedFunctionSpace((Vu, Vp))
w, phi = TestFunctions(M)
u, p = TrialFunctions(M)
# Get background fields
bbar = state.fields("bbar")
# Analytical (approximate) elimination of theta
k = state.k # Upward pointing unit vector
b = -dot(k, u) * dot(k, grad(bbar)) * beta + b_in
# vertical projection
def V(u):
return k * inner(u, k)
eqn = (inner(w, (u - u_in)) * dx - beta * div(w) * p * dx -
beta * inner(w, k) * b * dx + phi * div(u) * dx)
if hasattr(self.equations, "mu"):
eqn += dt * self.equations.mu * inner(w, k) * inner(u, k) * dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Place to put result of u p solver
self.up = Function(M)
# Boundary conditions (assumes extruded mesh)
# BCs are declared for the plain velocity space. As we need them in
# a mixed problem, we replicate the BCs but for subspace of M
bcs = [
DirichletBC(M.sub(0), bc.function_arg, bc.sub_domain)
for bc in self.equations.bcs['u']
]
# Solver for u, p
up_problem = LinearVariationalProblem(aeqn, Leqn, self.up, bcs=bcs)
# Provide callback for the nullspace of the trace system
def trace_nullsp(T):
return VectorSpaceBasis(constant=True)
appctx = {"trace_nullspace": trace_nullsp}
self.up_solver = LinearVariationalSolver(
up_problem,
solver_parameters=self.solver_parameters,
appctx=appctx)
# Reconstruction of b
b = TrialFunction(Vb)
gamma = TestFunction(Vb)
u, p = self.up.split()
self.b = Function(Vb)
b_eqn = gamma * (b - b_in + dot(k, u) * dot(k, grad(bbar)) * beta) * dx
b_problem = LinearVariationalProblem(lhs(b_eqn), rhs(b_eqn), self.b)
self.b_solver = LinearVariationalSolver(b_problem)
def _setup_solver(self):
state = self.state # just cutting down line length a bit
dt = state.timestepping.dt
beta = dt*state.timestepping.alpha
cp = state.parameters.cp
mu = state.mu
# Split up the rhs vector (symbolically)
u_in, rho_in, theta_in = split(state.xrhs)
# Build the reduced function space for u,rho
M = MixedFunctionSpace((state.V[0], state.V[1]))
w, phi = TestFunctions(M)
u, rho = TrialFunctions(M)
n = FacetNormal(state.mesh)
# Get background fields
thetabar = state.thetabar
rhobar = state.rhobar
pibar = exner(thetabar, rhobar, state)
pibar_rho = exner_rho(thetabar, rhobar, state)
pibar_theta = exner_theta(thetabar, rhobar, state)
# Analytical (approximate) elimination of theta
k = state.k # Upward pointing unit vector
theta = -dot(k,u)*dot(k,grad(thetabar))*beta + theta_in
# Only include theta' (rather than pi') in the vertical
# component of the gradient
# the pi prime term (here, bars are for mean and no bars are
# for linear perturbations)
pi = pibar_theta*theta + pibar_rho*rho
# vertical projection
def V(u):
return k*inner(u,k)
eqn = (
inner(w, (u - u_in))*dx
- beta*cp*div(theta*V(w))*pibar*dx
# following does nothing but is preserved in the comments
# to remind us why (because V(w) is purely vertical.
# + beta*cp*jump(theta*V(w),n)*avg(pibar)*dS_v
- beta*cp*div(thetabar*w)*pi*dx
+ beta*cp*jump(thetabar*w,n)*avg(pi)*dS_v
+ (phi*(rho - rho_in) - beta*inner(grad(phi), u)*rhobar)*dx
+ beta*jump(phi*u, n)*avg(rhobar)*(dS_v + dS_h)
)
if mu is not None:
eqn += dt*mu*inner(w,k)*inner(u,k)*dx
aeqn = lhs(eqn)
Leqn = rhs(eqn)
# Place to put result of u rho solver
self.urho = Function(M)
# Boundary conditions (assumes extruded mesh)
dim = M.sub(0).ufl_element().value_shape()[0]
bc = ("0.0",)*dim
bcs = [DirichletBC(M.sub(0), Expression(bc), "bottom"),
DirichletBC(M.sub(0), Expression(bc), "top")]
# Solver for u, rho
urho_problem = LinearVariationalProblem(
aeqn, Leqn, self.urho, bcs=bcs)
self.urho_solver = LinearVariationalSolver(urho_problem,
solver_parameters=self.params,
options_prefix='ImplicitSolver')
# Reconstruction of theta
theta = TrialFunction(state.V[2])
gamma = TestFunction(state.V[2])
u, rho = self.urho.split()
self.theta = Function(state.V[2])
theta_eqn = gamma*(theta - theta_in +
dot(k,u)*dot(k,grad(thetabar))*beta)*dx
theta_problem = LinearVariationalProblem(lhs(theta_eqn),
rhs(theta_eqn),
self.theta)
self.theta_solver = LinearVariationalSolver(theta_problem,
options_prefix='thetabacksubstitution')
def compressible_hydrostatic_balance(state, theta0, rho0, pi0=None,
top=False, pi_boundary=Constant(1.0),
solve_for_rho=False,
params=None):
"""
Compute a hydrostatically balanced density given a potential temperature
profile.
:arg state: The :class:`State` object.
:arg theta0: :class:`.Function`containing the potential temperature.
:arg rho0: :class:`.Function` to write the initial density into.
:arg top: If True, set a boundary condition at the top. Otherwise, set
it at the bottom.
:arg pi_boundary: a field or expression to use as boundary data for pi on
the top or bottom as specified.
"""
# Calculate hydrostatic Pi
W = MixedFunctionSpace((state.Vv,state.V[1]))
v, pi = TrialFunctions(W)
dv, dpi = TestFunctions(W)
n = FacetNormal(state.mesh)
cp = state.parameters.cp
alhs = (
(cp*inner(v,dv) - cp*div(dv*theta0)*pi)*dx
+ dpi*div(theta0*v)*dx
)
if top:
bmeasure = ds_t
bstring = "bottom"
else:
bmeasure = ds_b
bstring = "top"
arhs = -cp*inner(dv,n)*theta0*pi_boundary*bmeasure
if state.parameters.geopotential:
Phi = state.Phi
arhs += div(dv)*Phi*dx - inner(dv,n)*Phi*bmeasure
else:
g = state.parameters.g
arhs -= g*inner(dv,state.k)*dx
if(state.mesh.geometric_dimension() == 2):
bcs = [DirichletBC(W.sub(0), Expression(("0.", "0.")), bstring)]
elif(state.mesh.geometric_dimension() == 3):
bcs = [DirichletBC(W.sub(0), Expression(("0.", "0.", "0.")), bstring)]
w = Function(W)
PiProblem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs)
if(params is None):
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': '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'}
PiSolver = LinearVariationalSolver(PiProblem,
solver_parameters=params)
PiSolver.solve()
v, Pi = w.split()
if pi0 is not None:
pi0.assign(Pi)
kappa = state.parameters.kappa
R_d = state.parameters.R_d
p_0 = state.parameters.p_0
if solve_for_rho:
w1 = Function(W)
v, rho = w1.split()
rho.interpolate(p_0*(Pi**((1-kappa)/kappa))/R_d/theta0)
v, rho = split(w1)
dv, dpi = TestFunctions(W)
pi = ((R_d/p_0)*rho*theta0)**(kappa/(1.-kappa))
F = (
(cp*inner(v,dv) - cp*div(dv*theta0)*pi)*dx
+ dpi*div(theta0*v)*dx
+ cp*inner(dv,n)*theta0*pi_boundary*bmeasure
)
if state.parameters.geopotential:
F += - div(dv)*Phi*dx + inner(dv,n)*Phi*bmeasure
else:
F += g*inner(dv,state.k)*dx
rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs)
rhosolver = NonlinearVariationalSolver(rhoproblem, solver_parameters=params)
rhosolver.solve()
v, rho_ = w1.split()
rho0.assign(rho_)
else:
rho0.interpolate(p_0*(Pi**((1-kappa)/kappa))/R_d/theta0)
