def continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False):
if outflow and ibp == IntegrateByParts.NEVER:
raise ValueError(
"outflow is True and ibp is None are incompatible options")
Vu = state.spaces("HDiv")
dS_ = (dS_v + dS_h) if Vu.extruded else dS
ubar = Function(Vu)
if ibp == IntegrateByParts.ONCE:
L = -inner(grad(test), outer(q, ubar)) * dx
else:
L = inner(test, div(outer(q, ubar))) * dx
if ibp != IntegrateByParts.NEVER:
n = FacetNormal(state.mesh)
un = 0.5 * (dot(ubar, n) + abs(dot(ubar, n)))
L += dot(jump(test), (un('+') * q('+') - un('-') * q('-'))) * dS_
if ibp == IntegrateByParts.TWICE:
L -= (inner(test('+'),
dot(ubar('+'), n('+')) * q('+')) +
inner(test('-'),
dot(ubar('-'), n('-')) * q('-'))) * dS_
if outflow:
n = FacetNormal(state.mesh)
un = 0.5 * (dot(ubar, n) + abs(dot(ubar, n)))
L += test * un * q * (ds_v + ds_t + ds_b)
form = transporting_velocity(L, ubar)
return ibp_label(transport(form, TransportEquationType.conservative), ibp)
def interior_penalty_diffusion_form(state, test, q, parameters):
"""
Interior penalty diffusion form
:arg state: :class:`.State` object.
:arg V: Function space of diffused field
:arg direction: list containing directions in which function space
:arg: mu: the penalty weighting function, which is
:recommended to be proportional to 1/dx
:arg: kappa: strength of diffusion
"""
dS_ = (dS_v + dS_h) if state.mesh.extruded else dS
kappa = parameters.kappa
mu = parameters.mu
n = FacetNormal(state.mesh)
form = inner(grad(test), grad(q)*kappa)*dx
def get_flux_form(dS, M):
fluxes = (
-inner(2*avg(outer(q, n)), avg(grad(test)*M))
- inner(avg(grad(q)*M), 2*avg(outer(test, n)))
+ mu*inner(2*avg(outer(q, n)), 2*avg(outer(test, n)*kappa))
)*dS
return fluxes
form += get_flux_form(dS_, kappa)
return diffusion(form)
def add_dirichlet(self, g, surface):
"""
Adds dirichlet boundary conditions to the problem.
Args:
g (Function, int, or float):
The function to apply on the boundary.
surface (int or list of int):
The index of the boundary to apply the condition to.
"""
# Explicitly check against False as None should not be caught here.
if self._has_boundary is False:
raise AttributeError('Cannot add boundary after declaring that '
'there are no boundaries')
norm = FacetNormal(self.mesh)
if isinstance(g, (float, int)):
g = Constant(g)
integrand = -1 * self.v * dot(self.q, norm)
if surface == 'all':
dbc = DirichletBC(V=self.V, g=g, sub_domain="on_boundary")
self.a += integrand * ds
else:
dbc = DirichletBC(V=self.V, g=g, sub_domain=surface)
try:
self.a += sum(integrand * ds(s) for s in surface)
except TypeError:
self.a += integrand * ds(surface)
self.bcs.append(dbc)
self._has_boundary = True
def __init__(self, state, V, kappa, mu, bcs=None):
super(InteriorPenalty, self).__init__(state)
dt = state.timestepping.dt
gamma = TestFunction(V)
phi = TrialFunction(V)
self.phi1 = Function(V)
n = FacetNormal(state.mesh)
a = inner(gamma, phi) * dx + dt * inner(grad(gamma),
grad(phi) * kappa) * dx
def get_flux_form(dS, M):
fluxes = (-inner(2 * avg(outer(phi, n)), avg(grad(gamma) * M)) -
inner(avg(grad(phi) * M), 2 * avg(outer(gamma, n))) +
mu * inner(2 * avg(outer(phi, n)),
2 * avg(outer(gamma, n) * kappa))) * dS
return fluxes
a += dt * get_flux_form(dS_v, kappa)
a += dt * get_flux_form(dS_h, kappa)
L = inner(gamma, phi) * dx
problem = LinearVariationalProblem(a,
action(L, self.phi1),
self.phi1,
bcs=bcs)
self.solver = LinearVariationalSolver(problem)
def __init__(self, A, pyt_grad_op, pyt_op, actx, kappa):
"""
:arg kappa: The wave number
"""
self.actx = actx
self.k = kappa
self.pyt_op = pyt_op
self.pyt_grad_op = pyt_grad_op
self.A = A
self.meshmode_src_connection = meshmode_src_connection
# {{{ Create some functions needed for multing
self.x_fntn = Function(fspace)
# CG
self.potential_int = Function(fspace)
self.potential_int.dat.data[:] = 0.0
self.grad_potential_int = Function(vfspace)
self.grad_potential_int.dat.data[:] = 0.0
self.pyt_result = Function(fspace)
self.n = FacetNormal(mesh)
self.v = TestFunction(fspace)
# some meshmode ones
self.x_mm_fntn = self.meshmode_src_connection.discr.empty(
self.actx, dtype='c')
def __init__(self, mesh, conditions, timestepping, params, output, solver_params):
self.timestepping = timestepping
self.timestep = timestepping.timestep
self.timescale = timestepping.timescale
self.params = params
if output is None:
raise RuntimeError("You must provide a directory name for dumping results")
else:
self.output = output
self.outfile = File(output.dirname)
self.dump_count = 0
self.dump_freq = output.dumpfreq
self.solver_params = solver_params
self.mesh = mesh
self.conditions = conditions
if conditions.steady_state == True:
self.ind = 1
else:
self.ind = 1
family = conditions.family
self.x, self.y = SpatialCoordinate(mesh)
self.n = FacetNormal(mesh)
self.V = VectorFunctionSpace(mesh, family, conditions.order + 1)
self.U = FunctionSpace(mesh, family, conditions.order + 1)
self.U1 = FunctionSpace(mesh, 'DG', conditions.order)
self.S = TensorFunctionSpace(mesh, 'DG', conditions.order)
self.D = FunctionSpace(mesh, 'DG', 0)
self.W1 = MixedFunctionSpace([self.V, self.S])
self.W2 = MixedFunctionSpace([self.V, self.U1, self.U1])
self.W3 = MixedFunctionSpace([self.V, self.S, self.U1, self.U1])
def eady_initial_v(state, p0, v):
f = state.parameters.f
x, y, z = SpatialCoordinate(state.mesh)
# get pressure gradient
Vu = state.spaces("HDiv")
g = TrialFunction(Vu)
wg = TestFunction(Vu)
n = FacetNormal(state.mesh)
a = inner(wg, g)*dx
L = -div(wg)*p0*dx + inner(wg, n)*p0*ds_tb
pgrad = Function(Vu)
solve(a == L, pgrad)
# get initial v
Vp = p0.function_space()
phi = TestFunction(Vp)
m = TrialFunction(Vp)
a = f*phi*m*dx
L = phi*pgrad[0]*dx
solve(a == L, v)
return v
def compressible_eady_initial_v(state, theta0, rho0, v):
f = state.parameters.f
cp = state.parameters.cp
# exner function
Vr = rho0.function_space()
Pi = Function(Vr).interpolate(thermodynamics.pi(state.parameters, rho0, theta0))
# get Pi gradient
Vu = state.spaces("HDiv")
g = TrialFunction(Vu)
wg = TestFunction(Vu)
n = FacetNormal(state.mesh)
a = inner(wg, g)*dx
L = -div(wg)*Pi*dx + inner(wg, n)*Pi*ds_tb
pgrad = Function(Vu)
solve(a == L, pgrad)
# get initial v
m = TrialFunction(Vr)
phi = TestFunction(Vr)
a = phi*f*m*dx
L = phi*cp*theta0*pgrad[0]*dx
solve(a == L, v)
return v
def setup(self, state):
if not self._initialised:
space = state.spaces("Vv")
super().setup(state, space=space)
rho = state.fields("rho")
rhobar = state.fields("rhobar")
theta = state.fields("theta")
thetabar = state.fields("thetabar")
pi = thermodynamics.pi(state.parameters, rho, theta)
pibar = thermodynamics.pi(state.parameters, rhobar, thetabar)
cp = Constant(state.parameters.cp)
n = FacetNormal(state.mesh)
F = TrialFunction(space)
w = TestFunction(space)
a = inner(w, F)*dx
L = (- cp*div((theta-thetabar)*w)*pibar*dx
+ cp*jump((theta-thetabar)*w, n)*avg(pibar)*dS_v
- cp*div(thetabar*w)*(pi-pibar)*dx
+ cp*jump(thetabar*w, n)*avg(pi-pibar)*dS_v)
bcs = [DirichletBC(space, 0.0, "bottom"),
DirichletBC(space, 0.0, "top")]
imbalanceproblem = LinearVariationalProblem(a, L, self.field, bcs=bcs)
self.imbalance_solver = LinearVariationalSolver(imbalanceproblem)
def setUp(self):
"""
Prepare for tests.
"""
self.mesh = UnitCubeMesh(10, 10, 10)
self.V = FunctionSpace(self.mesh, 'CG', 1)
self.problem = MockProblem(self.mesh, self.V)
self.norm = FacetNormal(self.mesh)
def advection_term(self, q):
if self.continuity:
n = FacetNormal(self.state.mesh)
L = (-dot(grad(self.test), self.ubar) * self.qbar * dx +
jump(self.ubar * self.test, n) * avg(self.qbar) * self.dS)
else:
L = self.test * dot(self.ubar, self.state.k) * dot(
self.state.k, grad(self.qbar)) * dx
return L
def topography_term(self):
g = self.state.parameters.g
u0, _ = split(self.x0)
b = self.state.fields("topography")
n = FacetNormal(self.state.mesh)
un = 0.5 * (dot(u0, n) + abs(dot(u0, n)))
L = g * div(self.test) * b * dx - g * inner(
jump(self.test, n),
un('+') * b('+') - un('-') * b('-')) * dS
return L
def pressure_gradient_term(self):
g = self.state.parameters.g
u0, D0 = split(self.x0)
n = FacetNormal(self.state.mesh)
un = 0.5 * (dot(u0, n) + abs(dot(u0, n)))
L = g * (div(self.test) * D0 * dx -
inner(jump(self.test, n),
un('+') * D0('+') - un('-') * D0('-')) * dS)
return L
def advection_term(self, q):
if self.continuity:
if self.ibp == IntegrateByParts.ONCE:
L = -inner(grad(self.test), outer(q, self.ubar)) * dx
else:
L = inner(self.test, div(outer(q, self.ubar))) * dx
else:
if self.ibp == IntegrateByParts.ONCE:
L = -inner(div(outer(self.test, self.ubar)), q) * dx
else:
L = inner(outer(self.test, self.ubar), grad(q)) * dx
if self.dS is not None and self.ibp != IntegrateByParts.NEVER:
n = FacetNormal(self.state.mesh)
un = 0.5 * (dot(self.ubar, n) + abs(dot(self.ubar, n)))
L += dot(jump(self.test),
(un('+') * q('+') - un('-') * q('-'))) * self.dS
if self.ibp == IntegrateByParts.TWICE:
L -= (inner(self.test('+'),
dot(self.ubar('+'), n('+')) * q('+')) +
inner(self.test('-'),
dot(self.ubar('-'), n('-')) * q('-'))) * self.dS
if self.outflow:
n = FacetNormal(self.state.mesh)
un = 0.5 * (dot(self.ubar, n) + abs(dot(self.ubar, n)))
L += self.test * un * q * (ds_v + ds_t + ds_b)
if self.vector_manifold:
n = FacetNormal(self.state.mesh)
w = self.test
dS = self.dS
u = q
L += un('+') * inner(w('-'),
n('+') + n('-')) * inner(u('+'), n('+')) * dS
L += un('-') * inner(w('+'),
n('+') + n('-')) * inner(u('-'), n('-')) * dS
return L
def __init__(self):
super().__init__()
X = self.X
# volume and boundary integrals
self.name = "LevelSet Example 3"
V = FunctionSpace(self.mesh, "CG", 1)
u = interpolate(sin(X[0]) * cos(X[1])**2, V)
n = FacetNormal(self.mesh)
self.J = (sin(X[0]) * cos(X[1]) * dx +
pow(1.3 + X[0], 4.2) * pow(1.4 + X[1], 3.3) * dx +
exp(sin(X[0]) + cos(X[1])) * dx +
ln(5 + sin(X[0]) + cos(X[1])) * ds + inner(grad(u), n) * ds)
self.set_quadrature(8)
def vector_invariant_form(state, test, q, ibp=IntegrateByParts.ONCE):
Vu = state.spaces("HDiv")
dS_ = (dS_v + dS_h) if Vu.extruded else dS
ubar = Function(Vu)
n = FacetNormal(state.mesh)
Upwind = 0.5 * (sign(dot(ubar, n)) + 1)
if state.mesh.topological_dimension() == 3:
if ibp != IntegrateByParts.ONCE:
raise NotImplementedError
# <w,curl(u) cross ubar + grad( u.ubar)>
# =<curl(u),ubar cross w> - <div(w), u.ubar>
# =<u,curl(ubar cross w)> -
# <<u_upwind, [[n cross(ubar cross w)cross]]>>
both = lambda u: 2 * avg(u)
L = (inner(q, curl(cross(ubar, test))) * dx -
inner(both(Upwind * q), both(cross(n, cross(ubar, test)))) * dS_)
else:
perp = state.perp
if state.on_sphere:
outward_normals = CellNormal(state.mesh)
perp_u_upwind = lambda q: Upwind('+') * cross(
outward_normals('+'), q('+')) + Upwind('-') * cross(
outward_normals('-'), q('-'))
else:
perp_u_upwind = lambda q: Upwind('+') * perp(q('+')) + Upwind(
'-') * perp(q('-'))
if ibp == IntegrateByParts.ONCE:
L = (-inner(perp(grad(inner(test, perp(ubar)))), q) * dx -
inner(jump(inner(test, perp(ubar)), n), perp_u_upwind(q)) *
dS_)
else:
L = ((-inner(test,
div(perp(q)) * perp(ubar))) * dx -
inner(jump(inner(test, perp(ubar)), n), perp_u_upwind(q)) *
dS_ + jump(inner(test, perp(ubar)) * perp(q), n) * dS_)
L -= 0.5 * div(test) * inner(q, ubar) * dx
form = transporting_velocity(L, ubar)
return transport(form, TransportEquationType.vector_invariant)
def set_no_boundary(self):
"""
Declare that the problem will have no boundaries.
Raises:
AttributeError: If boundaries have already been added.
"""
# Explicitly check for True to be consistent with other method.
if self._has_boundary is True:
raise AttributeError('Cannot set no boundary after boundaries have'
' been set')
norm = FacetNormal(self.mesh)
self.a += -1 * self.v * dot(self.q, norm) * ds
self._has_boundary = False
def __init__(self, state, V, *, ibp="once", solver_params=None):
self.state = state
self.V = V
self.ibp = ibp
# set up functions required for forms
self.ubar = Function(state.spaces("HDiv"))
self.test = TestFunction(V)
self.trial = TrialFunction(V)
# find out if we are CG
nvertex = V.ufl_domain().ufl_cell().num_vertices()
entity_dofs = V.finat_element.entity_dofs()
# If there are as many dofs on vertices as there are vertices,
# assume a continuous space.
try:
self.is_cg = sum(map(len, entity_dofs[0].values())) == nvertex
except KeyError:
self.is_cg = sum(map(len, entity_dofs[(0, 0)].values())) == nvertex
# DG, embedded DG and hybrid SUPG methods need surface measures,
# n and un
if self.is_cg:
self.dS = None
self.ds = None
else:
if V.extruded:
self.dS = (dS_h + dS_v)
self.ds = (ds_b + ds_t + ds_v)
else:
self.dS = dS
self.ds = ds
self.n = FacetNormal(state.mesh)
self.un = 0.5 * (dot(self.ubar, self.n) +
abs(dot(self.ubar, self.n)))
if solver_params:
self.solver_parameters = solver_params
# default solver options
else:
self.solver_parameters = {
'ksp_type': 'cg',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
}
if state.output.log_level == DEBUG:
self.solver_parameters["ksp_monitor_true_residual"] = True
def linear_continuity_form(state, test, qbar, facet_term=False):
Vu = state.spaces("HDiv")
ubar = Function(Vu)
L = qbar * test * div(ubar) * dx
if facet_term:
n = FacetNormal(state.mesh)
Vu = state.spaces("HDiv")
dS_ = (dS_v + dS_h) if Vu.extruded else dS
L += jump(ubar * test, n) * avg(qbar) * dS_
form = transporting_velocity(L, ubar)
return transport(form, TransportEquationType.conservative)
def advection_term(self, q):
n = FacetNormal(self.state.mesh)
Upwind = 0.5 * (sign(dot(self.ubar, n)) + 1)
if self.state.mesh.topological_dimension() == 3:
# <w,curl(u) cross ubar + grad( u.ubar)>
# =<curl(u),ubar cross w> - <div(w), u.ubar>
# =<u,curl(ubar cross w)> -
# <<u_upwind, [[n cross(ubar cross w)cross]]>>
both = lambda u: 2 * avg(u)
L = (inner(q, curl(cross(self.ubar, self.test))) * dx -
inner(both(Upwind * q),
both(cross(n, cross(self.ubar, self.test)))) * self.dS)
else:
perp = self.state.perp
if self.state.on_sphere:
outward_normals = CellNormal(self.state.mesh)
perp_u_upwind = lambda q: Upwind('+') * cross(
outward_normals('+'), q('+')) + Upwind('-') * cross(
outward_normals('-'), q('-'))
else:
perp_u_upwind = lambda q: Upwind('+') * perp(q('+')) + Upwind(
'-') * perp(q('-'))
if self.ibp == IntegrateByParts.ONCE:
L = (-inner(perp(grad(inner(self.test, perp(self.ubar)))), q) *
dx - inner(jump(inner(self.test, perp(self.ubar)), n),
perp_u_upwind(q)) * self.dS)
else:
L = ((-inner(self.test,
div(perp(q)) * perp(self.ubar))) * dx -
inner(jump(inner(self.test, perp(self.ubar)), n),
perp_u_upwind(q)) * self.dS +
jump(inner(self.test, perp(self.ubar)) * perp(q), n) *
self.dS)
L -= 0.5 * div(self.test) * inner(q, self.ubar) * dx
return L
def vector_manifold_advection_form(state,
test,
q,
ibp=IntegrateByParts.ONCE,
outflow=False):
L = advection_form(state, test, q, ibp, outflow)
Vu = state.spaces("HDiv")
dS_ = (dS_v + dS_h) if Vu.extruded else dS
ubar = Function(Vu)
n = FacetNormal(state.mesh)
un = 0.5 * (dot(ubar, n) + abs(dot(ubar, n)))
L += un('+') * inner(test('-'),
n('+') + n('-')) * inner(q('+'), n('+')) * dS_
L += un('-') * inner(test('+'),
n('+') + n('-')) * inner(q('-'), n('-')) * dS_
return L
def setup(self, state):
if not self._initialised:
mesh_dim = state.mesh.geometric_dimension()
try:
field_dim = state.fields(self.fname).ufl_shape[0]
except IndexError:
field_dim = 1
shape = (mesh_dim, ) * field_dim
space = TensorFunctionSpace(state.mesh, "CG", 1, shape=shape)
super().setup(state, space=space)
f = state.fields(self.fname)
test = TestFunction(space)
trial = TrialFunction(space)
n = FacetNormal(state.mesh)
a = inner(test, trial)*dx
L = -inner(div(test), f)*dx
if space.extruded:
L += dot(dot(test, n), f)*(ds_t + ds_b)
prob = LinearVariationalProblem(a, L, self.field)
self.solver = LinearVariationalSolver(prob)
def central_facet_flux(F, v):
r"""Create the weak form of the central numerical flux through cell facets
This numerical flux, by itself, is unstable. The full right-hand side of
the problem requires an additional facet flux term for stability.
Parameters
----------
F : ufl.Expr
A symbolic expression for the flux
v : firedrake.TestFunction
A test function from the state space
Returns
-------
f : firedrake.Form
A 1-form that discretizes the residual of the flux
"""
mesh = v.ufl_domain()
n = FacetNormal(mesh)
return inner(avg(F), outer(v('+'), n('+')) + outer(v('-'), n('-'))) * dS
def vector_manifold_continuity_form(state,
test,
q,
ibp=IntegrateByParts.ONCE,
outflow=False):
L = continuity_form(state, test, q, ibp, outflow)
Vu = state.spaces("HDiv")
dS_ = (dS_v + dS_h) if Vu.extruded else dS
ubar = Function(Vu)
n = FacetNormal(state.mesh)
un = 0.5 * (dot(ubar, n) + abs(dot(ubar, n)))
L += un('+') * inner(test('-'),
n('+') + n('-')) * inner(q('+'), n('+')) * dS_
L += un('-') * inner(test('+'),
n('+') + n('-')) * inner(q('-'), n('-')) * dS_
form = transporting_velocity(L, ubar)
return transport(form)
def __init__(self, rho, solution):
self.rho = rho
self.solution = solution
self.fs = self.solution.function_space()
self.mesh = self.fs.mesh()
self.n = FacetNormal(self.mesh)
test = TestFunction(self.fs)
tri = TrialFunction(self.fs)
vert_dim = self.mesh.geometric_dimension() - 1
a = test * tri * dx
L = -Dx(test, vert_dim) * self.rho * dx + test * self.n[
vert_dim] * self.rho * ds_tb #+ avg(rho) * jump(gradrho_test, n[dim]) * dS_h (this is zero because jump(phi,n) = 0 for continuous P1 test function!)
prob = LinearVariationalProblem(a,
L,
self.solution,
constant_jacobian=True)
self.weak_grad_solver = LinearVariationalSolver(
prob) # #, solver_parameters=solver_parameters)
def __init__(self, A, pyt_grad_op, pyt_op, queue, kappa):
"""
:arg kappa: The wave number
"""
self.queue = queue
self.k = kappa
self.pyt_op = pyt_op
self.pyt_grad_op = pyt_grad_op
self.A = A
# {{{ Create some functions needed for multing
self.x_fntn = Function(fspace)
self.potential_int = Function(fspace)
self.potential_int.dat.data[:] = 0.0
self.grad_potential_int = Function(vfspace)
self.grad_potential_int.dat.data[:] = 0.0
self.pyt_result = Function(fspace)
self.n = FacetNormal(mesh)
self.v = TestFunction(fspace)
def pressure_gradient_term(self):
u0, rho0, theta0 = split(self.x0)
cp = self.state.parameters.cp
n = FacetNormal(self.state.mesh)
Vtheta = self.state.spaces("HDiv_v")
# introduce new theta so it can be changed by moisture
theta = theta0
# 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 = theta / (1 + water_t)
pi = thermodynamics.pi(self.state.parameters, rho0, theta0)
L = (+cp * div(theta * self.test) * pi * dx -
cp * jump(self.test * theta, n) * avg(pi) * dS_v)
return L
def create_dirichlet_bounds(mesh, V, T, v, k, g, boundary=[1, 2, 3, 4, 5, 6]):
"""
Create the dirichlet boundary conditions:
u = g on boundary
mesh: Mesh, The mesh to define the bound for
V: FunctionSpace, The function space for the boundary
T: Function, The function to be calculated
v: Function, The test function
k: Function, The conductivity of the problem
g: float, Used in above formula
Returns: bcs, R and b, the defining functions of the bound
Type: list<DirichletBC>, Function, Function
"""
norm = FacetNormal(mesh)
bcs = [DirichletBC(V, g, boundary)]
R = 0
b = k*v*dot(grad(T), norm)*ds
return bcs, R, b
def advection(self):
n = FacetNormal(self.mesh)
if len(self.wind) > 1:
element = self.V._ufl_element
degree = element.degree(0)
space_type = repr(element)[0]
element_type = element.family()
mesh = self.V.mesh()
V = VectorFunctionSpace(mesh, element_type, degree)
else:
V = self.V
if self.opt['form'] == 'linear':
w_ = Expression(self.wind, V)
u_ = self.sol
elif self.opt['form'] == 'bilinear':
w_ = Expression(self.wind, V)
u_ = self.u
elif self.opt['form'] == 'nonlinear':
w_ = self.sol
u_ = self.sol
return inner(dot(w_, nabla_grad(u_)), self.v) * dx
def setup(self, state):
if not self._initialised:
space = state.spaces("DG0", state.mesh, "DG", 0)
super().setup(state, space=space)
rain = state.fields('rain')
rho = state.fields('rho')
v = state.fields('rainfall_velocity')
self.phi = TestFunction(space)
flux = TrialFunction(space)
n = FacetNormal(state.mesh)
un = 0.5 * (dot(v, n) + abs(dot(v, n)))
self.flux = Function(space)
a = self.phi * flux * dx
L = self.phi * rain * un * rho
if space.extruded:
L = L * (ds_b + ds_t + ds_v)
else:
L = L * ds
# setup solver
problem = LinearVariationalProblem(a, L, self.flux)
self.solver = LinearVariationalSolver(problem)