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 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 remove_initial_w(u, Vv):
bc = DirichletBC(u.function_space()[0], Constant((0,0)), "bottom")
bc.apply(u)
uv = Function(Vv).project(u)
ustar = Function(u.function_space()).project(uv)
uin = Function(u.function_space()).assign(u - ustar)
u.assign(uin)
def remove_initial_w(u, Vv):
bc = DirichletBC(u.function_space()[0], 0.0, "bottom")
bc.apply(u)
uv = Function(Vv).project(u)
ustar = Function(u.function_space()).project(uv)
uin = Function(u.function_space()).assign(u - ustar)
u.assign(uin)
def createSubMatrix(self, mat, row_is, col_is, target=None):
if target is not None:
# Repeat call, just return the matrix, since we don't
# actually assemble in here.
target.assemble()
return target
from firedrake import DirichletBC
# These are the sets of ISes of which the the row and column
# space consist.
row_ises = self._y.function_space().dof_dset.field_ises
col_ises = self._x.function_space().dof_dset.field_ises
row_inds = find_sub_block(row_is, row_ises)
if row_is == col_is and row_ises == col_ises:
col_inds = row_inds
else:
col_inds = find_sub_block(col_is, col_ises)
asub = ExtractSubBlock().split(self.a,
argument_indices=(row_inds, col_inds))
Wrow = asub.arguments()[0].function_space()
Wcol = asub.arguments()[1].function_space()
row_bcs = []
col_bcs = []
for bc in self.row_bcs:
for i, r in enumerate(row_inds):
if bc.function_space().index == r:
row_bcs.append(DirichletBC(Wrow.split()[i],
bc.function_arg,
bc.sub_domain,
method=bc.method))
if Wrow == Wcol and row_inds == col_inds and self.row_bcs == self.col_bcs:
col_bcs = row_bcs
else:
for bc in self.col_bcs:
for i, c in enumerate(col_inds):
if bc.function_space().index == c:
col_bcs.append(DirichletBC(Wcol.split()[i],
bc.function_arg,
bc.sub_domain,
method=bc.method))
submat_ctx = ImplicitMatrixContext(asub,
row_bcs=row_bcs,
col_bcs=col_bcs,
fc_params=self.fc_params,
appctx=self.appctx)
submat_ctx.on_diag = self.on_diag and row_inds == col_inds
submat = PETSc.Mat().create(comm=mat.comm)
submat.setType("python")
submat.setSizes((submat_ctx.row_sizes, submat_ctx.col_sizes),
bsize=submat_ctx.block_size)
submat.setPythonContext(submat_ctx)
submat.setUp()
return submat
def remove_initial_w(u):
Vu = u.function_space()
Vv = FunctionSpace(Vu._ufl_domain, Vu.ufl_element()._elements[-1])
bc = DirichletBC(Vu[0], 0.0, "bottom")
bc.apply(u)
uv = Function(Vv).project(u)
ustar = Function(u.function_space()).project(uv)
uin = Function(u.function_space()).assign(u - ustar)
u.assign(uin)
def incompressible_hydrostatic_balance(state, b0, p0, top=False, params=None):
# get F
Vu = state.spaces("HDiv")
Vv = FunctionSpace(state.mesh, Vu.ufl_element()._elements[-1])
v = TrialFunction(Vv)
w = TestFunction(Vv)
if top:
bstring = "top"
else:
bstring = "bottom"
bcs = [DirichletBC(Vv, 0.0, bstring)]
a = inner(w, v) * dx
L = inner(state.k, w) * b0 * dx
F = Function(Vv)
solve(a == L, F, bcs=bcs)
# define mixed function space
VDG = state.spaces("DG")
WV = (Vv) * (VDG)
# get pprime
v, pprime = TrialFunctions(WV)
w, phi = TestFunctions(WV)
bcs = [DirichletBC(WV[0], zero(), bstring)]
a = (inner(w, v) + div(w) * pprime + div(v) * phi) * dx
L = phi * div(F) * dx
w1 = Function(WV)
if params is None:
params = {
'ksp_type': 'gmres',
'pc_type': 'fieldsplit',
'pc_fieldsplit_type': 'schur',
'pc_fieldsplit_schur_fact_type': 'full',
'pc_fieldsplit_schur_precondition': 'selfp',
'fieldsplit_1_ksp_type': 'preonly',
'fieldsplit_1_pc_type': 'gamg',
'fieldsplit_1_mg_levels_pc_type': 'bjacobi',
'fieldsplit_1_mg_levels_sub_pc_type': 'ilu',
'fieldsplit_0_ksp_type': 'richardson',
'fieldsplit_0_ksp_max_it': 4,
'ksp_atol': 1.e-08,
'ksp_rtol': 1.e-08
}
solve(a == L, w1, bcs=bcs, solver_parameters=params)
v, pprime = w1.split()
p0.project(pprime)
def _apply_bcs(self):
"""
Set the zero boundary conditions in the velocity.
"""
unp1 = self.state.xnp1.split()[0]
if unp1.function_space().extruded:
M = unp1.function_space()
bcs = [DirichletBC(M, 0.0, "bottom"), DirichletBC(M, 0.0, "top")]
for bc in bcs:
bc.apply(unp1)
def find_domain_boundaries(mesh):
"""
Makes a scalar DG0 function whose values are 0. everywhere except for in
cells on the boundary of the domain, where the values are 1.0.
This allows boundary cells to be identified easily.
:arg mesh: the mesh.
"""
DG0 = FunctionSpace(mesh, "DG", 0)
CG1 = FunctionSpace(mesh, "CG", 1)
on_exterior_DG0 = Function(DG0)
on_exterior_CG1 = Function(CG1)
# we get values in CG1 initially as DG0 will not work for triangular elements
bc_codes = ['on_boundary', 'top', 'bottom']
bcs = [DirichletBC(CG1, Constant(1.0), bc_code) for bc_code in bc_codes]
for bc in bcs:
try:
bc.apply(on_exterior_CG1)
except ValueError:
pass
on_exterior_DG0.interpolate(on_exterior_CG1)
return on_exterior_DG0
def __init__(self, mesh, conditions, timestepping, params, output, solver_params):
super().__init__(mesh, conditions, timestepping, params, output, solver_params)
self.u0 = Function(self.V)
self.u1 = Function(self.V)
self.sigma0 = Function(self.S)
self.sigma1 = Function(self.S)
theta = conditions.theta
uh = (1-theta) * self.u0 + theta * self.u1
a = Function(self.U)
h = Function(self.U)
p = TestFunction(self.V)
q = TestFunction(self.S)
self.initial_condition((self.u0, conditions.ic['u']), (a, conditions.ic['a']), (h, conditions.ic['h']))
ep_dot = self.strain(grad(uh))
zeta = self.zeta(h, a, self.delta(uh))
eta = zeta * params.e ** (-2)
rheology = 2 * eta * ep_dot + (zeta - eta) * tr(ep_dot) * Identity(2) - 0.5 * self.Ice_Strength(h, a) * Identity(2)
self.initial_condition((self.sigma0, rheology),(self.sigma1, self.sigma0))
def sigma_next(timestep, zeta, ep_dot, sigma, P):
A = 1 + 0.25 * (timestep * params.e ** 2) / params.T
B = timestep * 0.125 * (1 - params.e ** 2) / params.T
rhs = (1 - (timestep * params.e ** 2) / (4 * params.T)) * sigma - timestep / params.T * (
0.125 * (1 - params.e ** 2) * tr(sigma) * Identity(2) - 0.25 * P * Identity(2) + zeta * ep_dot)
C = (rhs[0, 0] - rhs[1, 1]) / A
D = (rhs[0, 0] + rhs[1, 1]) / (A + 2 * B)
sigma00 = 0.5 * (C + D)
sigma11 = 0.5 * (D - C)
sigma01 = rhs[0, 1]
sigma = as_matrix([[sigma00, sigma01], [sigma01, sigma11]])
return sigma
s = sigma_next(self.timestep, zeta, ep_dot, self.sigma0, self.Ice_Strength(h, a))
sh = (1-theta) * s + theta * self.sigma0
eqn = self.momentum_equation(h, self.u1, self.u0, p, sh, params.rho, uh, conditions.ocean_curr,
params.rho_a, params.C_a, params.rho_w, params.C_w, conditions.geo_wind, params.cor, self.timestep, ind=self.ind)
tensor_eqn = inner(self.sigma1-s, q) * dx
if conditions.stabilised['state']:
alpha = conditions.stabilised['alpha']
eqn += stabilisation_term(alpha=alpha, zeta=avg(zeta), mesh=mesh, v=uh, test=p)
bcs = DirichletBC(self.V, conditions.bc['u'], "on_boundary")
uprob = NonlinearVariationalProblem(eqn, self.u1, bcs)
self.usolver = NonlinearVariationalSolver(uprob, solver_parameters=solver_params.bt_params)
sprob = NonlinearVariationalProblem(tensor_eqn, self.sigma1)
self.ssolver = NonlinearVariationalSolver(sprob, solver_parameters=solver_params.bt_params)
def __init__(self, mesh, conditions, timestepping, params, output, solver_params):
super().__init__(mesh, conditions, timestepping, params, output, solver_params)
self.u0 = Function(self.V)
self.u1 = Function(self.V)
self.h = Function(self.U)
self.a = Function(self.U)
self.p = TestFunction(self.V)
theta = conditions.theta
self.uh = (1-theta) * self.u0 + theta * self.u1
ep_dot = self.strain(grad(self.uh))
self.initial_condition((self.u0, conditions.ic['u']),(self.u1, self.u0),
(self.a, conditions.ic['a']),(self.h, conditions.ic['h']))
zeta = self.zeta(self.h, self.a, self.delta(self.uh))
eta = zeta * params.e ** -2
sigma = 2 * eta * ep_dot + (zeta - eta) * tr(ep_dot) * Identity(2) - 0.5 * self.Ice_Strength(self.h,self.a) * Identity(2)
self.eqn = self.momentum_equation(self.h, self.u1, self.u0, self.p, sigma, params.rho, self.uh,
conditions.ocean_curr, params.rho_a, params.C_a, params.rho_w,
params.C_w, conditions.geo_wind, params.cor, self.timestep)
if conditions.stabilised['state']:
alpha = conditions.stabilised['alpha']
self.eqn += self.stabilisation_term(alpha=alpha, zeta=avg(zeta), mesh=mesh, v=self.uh, test=self.p)
self.bcs = DirichletBC(self.V, conditions.bc['u'], "on_boundary")
def A(fs):
u = TrialFunction(fs)
v = TestFunction(fs)
a = (dot(grad(v), grad(u))) * dx
bc = DirichletBC(fs, 0., "on_boundary")
A = assemble(a, bcs = bc)
return A
def bcs(self):
if not self._bcs:
self._bcs = [
DirichletBC(self.V.sub(0), Expression(self.uB,
self.V_dict['V']),
'on_boundary')
]
return self._bcs
def Fubc(V, uexact, rhs):
u = interpolate(uexact, V)
v = TestFunction(V)
F = inner(Dt(u), v) * dx + inner(grad(u), grad(v)) * dx - inner(rhs,
v) * dx
bc = DirichletBC(V, uexact, "on_boundary")
return (F, u, bc)
def _build_forcing_solvers(self):
a = self.mass_term()
L = self.forcing_term()
if self.Vu.extruded:
bcs = [
DirichletBC(self.Vu, 0.0, "bottom"),
DirichletBC(self.Vu, 0.0, "top")
]
else:
bcs = None
u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs)
solver_parameters = {}
if self.state.output.log_level == DEBUG:
solver_parameters["ksp_monitor_true_residual"] = True
self.u_forcing_solver = LinearVariationalSolver(
u_forcing_problem,
solver_parameters=solver_parameters,
options_prefix="UForcingSolver")
def setup(self, state):
super(GeostrophicImbalance, self).setup(state)
u = state.fields("u")
b = state.fields("b")
p = state.fields("p")
f = state.parameters.f
Vu = u.function_space()
v = TrialFunction(Vu)
w = TestFunction(Vu)
a = inner(w, v)*dx
L = (div(w)*p+inner(w, as_vector([f*u[1], 0.0, b])))*dx
bcs = [DirichletBC(Vu, 0.0, "bottom"),
DirichletBC(Vu, 0.0, "top")]
self.imbalance = Function(Vu)
imbalanceproblem = LinearVariationalProblem(a, L, self.imbalance, bcs=bcs)
self.imbalance_solver = LinearVariationalSolver(
imbalanceproblem, solver_parameters={'ksp_type': 'cg'})
def test_assemble_firedrake(fs):
"test assemble with a firedrake object"
# assemble firedrake matrix
u = TrialFunction(fs)
v = TestFunction(fs)
a = (dot(grad(v), grad(u))) * dx
bc = DirichletBC(fs, 0., "on_boundary")
A = assemble(a, bcs=bc)
assert isinstance(A, Matrix)
def __init__(self):
super().__init__()
self.name = "PDE-constraint with DirBC"
u, v, X = self.u, self.v, self.X
f = sin(X[1]) * cos(X[0])
g = f / 2.
self.g = g
self.with_DirBC = 1 # special case for DirichletBC in TaylorTest
self.bc = DirichletBC(self.V, 0., "on_boundary")
self.F = (inner(grad(u + g), grad(v)) + (u + g) * v - f * v) * dx
self.J = (u + g) * (u + g) * dx + pow(1 + u * u, 2.5) * ds
self.set_quadrature(10)
def set_no_normal_flow_bcs(self, state, no_normal_flow_bc_ids):
"""
Sets up the no-normal-flow boundary conditions, storing the
:class:`DirichletBC` object at each specified boundary. There must be
a velocity variable named 'u' to apply the boundary conditions to.
:arg state: The `State` object.
:arg normal_flow_bcs_ids: A list of IDs of the domain boundaries at
which no normal flow will be enforced.
"""
if 'u' not in self.field_names:
raise NotImplementedError(
'No-normal-flow boundary conditions can only be applied ' +
'when there is a variable called "u" and none was found')
Vu = state.spaces("HDiv")
if Vu.extruded:
self.bcs['u'].append(DirichletBC(Vu, 0.0, "bottom"))
self.bcs['u'].append(DirichletBC(Vu, 0.0, "top"))
for id in no_normal_flow_bc_ids:
self.bcs['u'].append(DirichletBC(Vu, 0.0, id))
def heat(n, deg, time_stages, stage_type="deriv", splitting=IA):
N = 2**n
msh = UnitIntervalMesh(N)
params = {
"snes_type": "ksponly",
"ksp_type": "preonly",
"mat_type": "aij",
"pc_type": "lu"
}
V = FunctionSpace(msh, "CG", deg)
x, = SpatialCoordinate(msh)
t = Constant(0.0)
dt = Constant(2.0 / N)
uexact = exp(-t) * sin(pi * x)
rhs = expand_derivatives(diff(uexact, t)) - div(grad(uexact))
butcher_tableau = GaussLegendre(time_stages)
u = project(uexact, V)
v = TestFunction(V)
F = (inner(Dt(u), v) * dx + inner(grad(u), grad(v)) * dx -
inner(rhs, v) * dx)
bc = DirichletBC(V, Constant(0), "on_boundary")
stepper = TimeStepper(F,
butcher_tableau,
t,
dt,
u,
bcs=bc,
solver_parameters=params,
stage_type=stage_type,
splitting=splitting)
while (float(t) < 1.0):
if (float(t) + float(dt) > 1.0):
dt.assign(1.0 - float(t))
stepper.advance()
t.assign(float(t) + float(dt))
return errornorm(uexact, u) / norm(uexact)
def __init__(self, mesh, conditions, timestepping, params, output, solver_params):
super().__init__(mesh, conditions, timestepping, params, output, solver_params)
self.w0 = Function(self.W3)
self.w1 = Function(self.W3)
u0, s0, h0, a0 = self.w0.split()
p, q, r, m = TestFunctions(self.W3)
self.initial_condition((u0, conditions.ic['u']), (s0, conditions.ic['s']),
(a0, conditions.ic['a']), (h0, conditions.ic['h']))
self.w1.assign(self.w0)
u1, s1, h1, a1 = split(self.w1)
u0, s0, h0, a0 = split(self.w0)
theta = conditions.theta
uh = (1-theta) * u0 + theta * u1
sh = (1-theta) * s0 + theta * s1
hh = (1-theta) * h0 + theta * h1
ah = (1-theta) * a0 + theta * a1
ep_dot = self.strain(grad(uh))
zeta = self.zeta(hh, ah, self.delta(uh))
rheology = params.e ** 2 * sh + Identity(2) * 0.5 * ((1 - params.e ** 2) * tr(sh) + self.Ice_Strength(hh, ah))
eqn = self.momentum_equation(hh, u1, u0, p, sh, params.rho, uh, conditions.ocean_curr, params.rho_a,
params.C_a, params.rho_w, params.C_w, conditions.geo_wind, params.cor, self.timestep, ind=self.ind)
eqn += self.transport_equation(uh, hh, ah, h1, h0, a1, a0, r, m, self.n, self.timestep)
eqn += inner(self.ind * (s1 - s0) + 0.5 * self.timestep * rheology / params.T, q) * dx
eqn -= inner(q * zeta * self.timestep / params.T, ep_dot) * dx
if conditions.stabilised['state']:
alpha = conditions.stabilised['alpha']
eqn += self.stabilisation_term(alpha=alpha, zeta=avg(zeta), mesh=mesh, v=uh, test=p)
bcs = DirichletBC(self.W3.sub(0), conditions.bc['u'], "on_boundary")
uprob = NonlinearVariationalProblem(eqn, self.w1, bcs)
self.usolver = NonlinearVariationalSolver(uprob, solver_parameters=solver_params.bt_params)
self.u1, self.s0, self.h1, self.a1 = self.w1.split()
def __init__(self, mesh, conditions, timestepping, params, output, solver_params):
super().__init__(mesh, conditions, timestepping, params, output, solver_params)
self.w0 = Function(self.W2)
self.w1 = Function(self.W2)
u0, h0, a0 = self.w0.split()
p, q, r = TestFunctions(self.W2)
self.initial_condition((u0, conditions.ic['u']), (h0, conditions.ic['h']),
(a0, conditions.ic['a']))
self.w1.assign(self.w0)
u1, h1, a1 = split(self.w1)
u0, h0, a0 = split(self.w0)
theta = conditions.theta
uh = (1-theta) * u0 + theta * u1
ah = (1-theta) * a0 + theta * a1
hh = (1-theta) * h0 + theta * h1
ep_dot = self.strain(grad(uh))
zeta = self.zeta(hh, ah, self.delta(uh))
eta = zeta * params.e ** (-2)
sigma = 2 * eta * ep_dot + (zeta - eta) * tr(ep_dot) * Identity(2) - self.Ice_Strength(hh, ah) * 0.5 * Identity(
2)
eqn = self.momentum_equation(hh, u1, u0, p, sigma, params.rho, uh, conditions.ocean_curr, params.rho_a,
params.C_a, params.rho_w, params.C_w, conditions.geo_wind, params.cor, self.timestep)
eqn += self.transport_equation(uh, hh, ah, h1, h0, a1, a0, q, r, self.n, self.timestep)
if conditions.stabilised['state']:
alpha = conditions.stabilised['alpha']
eqn += self.stabilisation_term(alpha=alpha, zeta=avg(zeta), mesh=mesh, v=uh, test=p)
bcs = DirichletBC(self.W2.sub(0), conditions.bc['u'], "on_boundary")
uprob = NonlinearVariationalProblem(eqn, self.w1, bcs)
self.usolver = NonlinearVariationalSolver(uprob, solver_parameters=solver_params.bt_params)
self.u1, self.h1, self.a1 = self.w1.split()
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 __init__(self, mesh, conditions, timestepping, params, output, solver_params):
super().__init__(mesh, conditions, timestepping, params, output, solver_params)
self.w0 = Function(self.W1)
self.w1 = Function(self.W1)
self.a = Function(self.U)
self.h = Function(self.U)
self.u0, self.s0 = self.w0.split()
self.p, self.q = TestFunctions(self.W1)
self.initial_condition((self.u0, conditions.ic['u']), (self.s0, conditions.ic['s']),
(self.a, conditions.ic['a']), (self.h, conditions.ic['h']))
self.w1.assign(self.w0)
u1, s1 = split(self.w1)
u0, s0 = split(self.w0)
theta = conditions.theta
uh = (1-theta) * u0 + theta * u1
sh = (1-theta) * s0 + theta * s1
self.ep_dot = self.strain(grad(uh))
zeta = self.zeta(self.h, self.a, self.delta(uh))
self.rheology = params.e ** 2 * sh + Identity(2) * 0.5 * ((1 - params.e ** 2) * tr(sh) + self.Ice_Strength(self.h, self.a))
self.eqn = self.momentum_equation(self.h, u1, u0, self.p, sh, params.rho, uh, conditions.ocean_curr, params.rho_a,
params.C_a, params.rho_w, params.C_w, conditions.geo_wind, params.cor, self.timestep, ind=self.ind)
self.eqn += inner(self.ind * (s1 - s0) + 0.5 * self.timestep * self.rheology / params.T, self.q) * dx
self.eqn -= inner(self.q * zeta * self.timestep / params.T, self.ep_dot) * dx
if conditions.stabilised['state']:
alpha = conditions.stabilised['alpha']
self.eqn += self.stabilisation_term(alpha=alpha, zeta=avg(zeta), mesh=mesh, v=uh, test=self.p)
self.bcs = DirichletBC(self.W1.sub(0), conditions.bc['u'], "on_boundary")
def initialize(self, pc):
"""Set up the problem context. Take the original
mixed problem and reformulate the problem as a
hybridized mixed system.
A KSP is created for the Lagrange multiplier system.
"""
from firedrake import (FunctionSpace, Function, Constant,
TrialFunction, TrialFunctions, TestFunction,
DirichletBC, assemble)
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.formmanipulation import split_form
from ufl.algorithms.replace import replace
# Extract the problem context
prefix = pc.getOptionsPrefix() + "hybridization_"
_, P = pc.getOperators()
self.cxt = P.getPythonContext()
if not isinstance(self.cxt, ImplicitMatrixContext):
raise ValueError("The python context must be an ImplicitMatrixContext")
test, trial = self.cxt.a.arguments()
V = test.function_space()
mesh = V.mesh()
if len(V) != 2:
raise ValueError("Expecting two function spaces.")
if all(Vi.ufl_element().value_shape() for Vi in V):
raise ValueError("Expecting an H(div) x L2 pair of spaces.")
# Automagically determine which spaces are vector and scalar
for i, Vi in enumerate(V):
if Vi.ufl_element().sobolev_space().name == "HDiv":
self.vidx = i
else:
assert Vi.ufl_element().sobolev_space().name == "L2"
self.pidx = i
# Create the space of approximate traces.
W = V[self.vidx]
if W.ufl_element().family() == "Brezzi-Douglas-Marini":
tdegree = W.ufl_element().degree()
else:
try:
# If we have a tensor product element
h_deg, v_deg = W.ufl_element().degree()
tdegree = (h_deg - 1, v_deg - 1)
except TypeError:
tdegree = W.ufl_element().degree() - 1
TraceSpace = FunctionSpace(mesh, "HDiv Trace", tdegree)
# Break the function spaces and define fully discontinuous spaces
broken_elements = ufl.MixedElement([ufl.BrokenElement(Vi.ufl_element()) for Vi in V])
V_d = FunctionSpace(mesh, broken_elements)
# Set up the functions for the original, hybridized
# and schur complement systems
self.broken_solution = Function(V_d)
self.broken_residual = Function(V_d)
self.trace_solution = Function(TraceSpace)
self.unbroken_solution = Function(V)
self.unbroken_residual = Function(V)
# Set up the KSP for the hdiv residual projection
hdiv_mass_ksp = PETSc.KSP().create(comm=pc.comm)
hdiv_mass_ksp.setOptionsPrefix(prefix + "hdiv_residual_")
# HDiv mass operator
p = TrialFunction(V[self.vidx])
q = TestFunction(V[self.vidx])
mass = ufl.dot(p, q)*ufl.dx
# TODO: Bcs?
M = assemble(mass, bcs=None, form_compiler_parameters=self.cxt.fc_params)
M.force_evaluation()
Mmat = M.petscmat
hdiv_mass_ksp.setOperators(Mmat)
hdiv_mass_ksp.setUp()
hdiv_mass_ksp.setFromOptions()
self.hdiv_mass_ksp = hdiv_mass_ksp
# Storing the result of A.inv * r, where A is the HDiv
# mass matrix and r is the HDiv residual
self._primal_r = Function(V[self.vidx])
tau = TestFunction(V_d[self.vidx])
self._assemble_broken_r = create_assembly_callable(
ufl.dot(self._primal_r, tau)*ufl.dx,
tensor=self.broken_residual.split()[self.vidx],
form_compiler_parameters=self.cxt.fc_params)
# Create the symbolic Schur-reduction:
# Original mixed operator replaced with "broken"
# arguments
arg_map = {test: TestFunction(V_d),
trial: TrialFunction(V_d)}
Atilde = Tensor(replace(self.cxt.a, arg_map))
gammar = TestFunction(TraceSpace)
n = ufl.FacetNormal(mesh)
sigma = TrialFunctions(V_d)[self.vidx]
# We zero out the contribution of the trace variables on the exterior
# boundary. Extruded cells will have both horizontal and vertical
# facets
if mesh.cell_set._extruded:
trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), "on_boundary"),
DirichletBC(TraceSpace, Constant(0.0), "bottom"),
DirichletBC(TraceSpace, Constant(0.0), "top")]
K = Tensor(gammar('+') * ufl.dot(sigma, n) * ufl.dS_h +
gammar('+') * ufl.dot(sigma, n) * ufl.dS_v)
else:
trace_bcs = [DirichletBC(TraceSpace, Constant(0.0), "on_boundary")]
K = Tensor(gammar('+') * ufl.dot(sigma, n) * ufl.dS)
# If boundary conditions are contained in the ImplicitMatrixContext:
if self.cxt.row_bcs:
raise NotImplementedError("Strong BCs not currently handled. Try imposing them weakly.")
# Assemble the Schur complement operator and right-hand side
self.schur_rhs = Function(TraceSpace)
self._assemble_Srhs = create_assembly_callable(
K * Atilde.inv * self.broken_residual,
tensor=self.schur_rhs,
form_compiler_parameters=self.cxt.fc_params)
schur_comp = K * Atilde.inv * K.T
self.S = allocate_matrix(schur_comp, bcs=trace_bcs,
form_compiler_parameters=self.cxt.fc_params)
self._assemble_S = create_assembly_callable(schur_comp,
tensor=self.S,
bcs=trace_bcs,
form_compiler_parameters=self.cxt.fc_params)
self._assemble_S()
self.S.force_evaluation()
Smat = self.S.petscmat
# Nullspace for the multiplier problem
nullspace = create_schur_nullspace(P, -K * Atilde,
V, V_d, TraceSpace,
pc.comm)
if nullspace:
Smat.setNullSpace(nullspace)
# Set up the KSP for the system of Lagrange multipliers
trace_ksp = PETSc.KSP().create(comm=pc.comm)
trace_ksp.setOptionsPrefix(prefix)
trace_ksp.setOperators(Smat)
trace_ksp.setUp()
trace_ksp.setFromOptions()
self.trace_ksp = trace_ksp
split_mixed_op = dict(split_form(Atilde.form))
split_trace_op = dict(split_form(K.form))
# Generate reconstruction calls
self._reconstruction_calls(split_mixed_op, split_trace_op)
# NOTE: The projection stage *might* be replaced by a Fortin
# operator. We may want to allow the user to specify if they
# wish to use a Fortin operator over a projection, or vice-versa.
# In a future add-on, we can add a switch which chooses either
# the Fortin reconstruction or the usual KSP projection.
# Set up the projection KSP
hdiv_projection_ksp = PETSc.KSP().create(comm=pc.comm)
hdiv_projection_ksp.setOptionsPrefix(prefix + 'hdiv_projection_')
# Reuse the mass operator from the hdiv_mass_ksp
hdiv_projection_ksp.setOperators(Mmat)
# Construct the RHS for the projection stage
self._projection_rhs = Function(V[self.vidx])
self._assemble_projection_rhs = create_assembly_callable(
ufl.dot(self.broken_solution.split()[self.vidx], q)*ufl.dx,
tensor=self._projection_rhs,
form_compiler_parameters=self.cxt.fc_params)
# Finalize ksp setup
hdiv_projection_ksp.setUp()
hdiv_projection_ksp.setFromOptions()
self.hdiv_projection_ksp = hdiv_projection_ksp
def __init__(self, equation, alpha):
self.field_name = equation.field_name
implicit_terms = ["incompressibility", "sponge"]
dt = equation.state.dt
W = equation.function_space
self.x0 = Function(W)
self.xF = Function(W)
# set up boundary conditions on the u subspace of W
bcs = [
DirichletBC(W.sub(0), bc.function_arg, bc.sub_domain)
for bc in equation.bcs['u']
]
# drop terms relating to transport and diffusion
residual = equation.residual.label_map(
lambda t: any(t.has_label(transport, diffusion, return_tuple=True)
), drop)
# the lhs of both of the explicit and implicit solvers is just
# the time derivative form
trials = TrialFunctions(W)
a = residual.label_map(lambda t: t.has_label(time_derivative),
replace_subject(trials),
map_if_false=drop)
# the explicit forms are multiplied by (1-alpha) and moved to the rhs
L_explicit = -(1 - alpha) * dt * residual.label_map(
lambda t: t.has_label(time_derivative) or t.get(name) in
implicit_terms or t.get(name) == "hydrostatic_form", drop,
replace_subject(self.x0))
# the implicit forms are multiplied by alpha and moved to the rhs
L_implicit = -alpha * dt * residual.label_map(
lambda t: t.has_label(time_derivative) or t.get(name) in
implicit_terms or t.get(name) == "hydrostatic_form", drop,
replace_subject(self.x0))
# now add the terms that are always fully implicit
if any(t.get(name) in implicit_terms for t in residual):
L_implicit -= dt * residual.label_map(
lambda t: t.get(name) in implicit_terms,
replace_subject(self.x0), drop)
# the hydrostatic equations require some additional forms:
if any([t.has_label(hydrostatic) for t in residual]):
L_explicit += residual.label_map(
lambda t: t.get(name) == "hydrostatic_form",
replace_subject(self.x0), drop)
L_implicit -= residual.label_map(
lambda t: t.get(name) == "hydrostatic_form",
replace_subject(self.x0), drop)
# now we can set up the explicit and implicit problems
explicit_forcing_problem = LinearVariationalProblem(a.form,
L_explicit.form,
self.xF,
bcs=bcs)
implicit_forcing_problem = LinearVariationalProblem(a.form,
L_implicit.form,
self.xF,
bcs=bcs)
solver_parameters = {}
if logger.isEnabledFor(DEBUG):
solver_parameters["ksp_monitor_true_residual"] = None
self.solvers = {}
self.solvers["explicit"] = LinearVariationalSolver(
explicit_forcing_problem,
solver_parameters=solver_parameters,
options_prefix="ExplicitForcingSolver")
self.solvers["implicit"] = LinearVariationalSolver(
implicit_forcing_problem,
solver_parameters=solver_parameters,
options_prefix="ImplicitForcingSolver")
Vt,
equation_form="advective",
options=EmbeddedDGOptions())
advected_fields = []
advected_fields.append(("u", ThetaMethod(state, u0, ueqn)))
advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn)))
advected_fields.append(("theta", SSPRK3(state, theta0, thetaeqn)))
advected_fields.append(("water", SSPRK3(state, water0, watereqn)))
# Set up linear solver
linear_solver = CompressibleSolver(state)
# Set up forcing
compressible_forcing = CompressibleForcing(state)
bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")]
diffused_fields = [("u",
InteriorPenalty(state,
Vu,
kappa=75.,
mu=Constant(10. / delta),
bcs=bcs)),
("theta",
InteriorPenalty(state,
Vt,
kappa=75.,
mu=Constant(10. / delta)))]
# build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver,
compressible_forcing, diffused_fields)
def initialize(self, pc):
"""Set up the problem context. Take the original
mixed problem and reformulate the problem as a
hybridized mixed system.
A KSP is created for the Lagrange multiplier system.
"""
from firedrake import (FunctionSpace, Function, Constant,
TrialFunction, TrialFunctions, TestFunction,
DirichletBC)
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.formmanipulation import split_form
from ufl.algorithms.replace import replace
# Extract the problem context
prefix = pc.getOptionsPrefix() + "hybridization_"
_, P = pc.getOperators()
self.ctx = P.getPythonContext()
if not isinstance(self.ctx, ImplicitMatrixContext):
raise ValueError(
"The python context must be an ImplicitMatrixContext")
test, trial = self.ctx.a.arguments()
V = test.function_space()
mesh = V.mesh()
if len(V) != 2:
raise ValueError("Expecting two function spaces.")
if all(Vi.ufl_element().value_shape() for Vi in V):
raise ValueError("Expecting an H(div) x L2 pair of spaces.")
# Automagically determine which spaces are vector and scalar
for i, Vi in enumerate(V):
if Vi.ufl_element().sobolev_space().name == "HDiv":
self.vidx = i
else:
assert Vi.ufl_element().sobolev_space().name == "L2"
self.pidx = i
# Create the space of approximate traces.
W = V[self.vidx]
if W.ufl_element().family() == "Brezzi-Douglas-Marini":
tdegree = W.ufl_element().degree()
else:
try:
# If we have a tensor product element
h_deg, v_deg = W.ufl_element().degree()
tdegree = (h_deg - 1, v_deg - 1)
except TypeError:
tdegree = W.ufl_element().degree() - 1
TraceSpace = FunctionSpace(mesh, "HDiv Trace", tdegree)
# Break the function spaces and define fully discontinuous spaces
broken_elements = ufl.MixedElement(
[ufl.BrokenElement(Vi.ufl_element()) for Vi in V])
V_d = FunctionSpace(mesh, broken_elements)
# Set up the functions for the original, hybridized
# and schur complement systems
self.broken_solution = Function(V_d)
self.broken_residual = Function(V_d)
self.trace_solution = Function(TraceSpace)
self.unbroken_solution = Function(V)
self.unbroken_residual = Function(V)
shapes = (V[self.vidx].finat_element.space_dimension(),
np.prod(V[self.vidx].shape))
domain = "{[i,j]: 0 <= i < %d and 0 <= j < %d}" % shapes
instructions = """
for i, j
w[i,j] = w[i,j] + 1
end
"""
self.weight = Function(V[self.vidx])
par_loop((domain, instructions),
ufl.dx, {"w": (self.weight, INC)},
is_loopy_kernel=True)
instructions = """
for i, j
vec_out[i,j] = vec_out[i,j] + vec_in[i,j]/w[i,j]
end
"""
self.average_kernel = (domain, instructions)
# Create the symbolic Schur-reduction:
# Original mixed operator replaced with "broken"
# arguments
arg_map = {test: TestFunction(V_d), trial: TrialFunction(V_d)}
Atilde = Tensor(replace(self.ctx.a, arg_map))
gammar = TestFunction(TraceSpace)
n = ufl.FacetNormal(mesh)
sigma = TrialFunctions(V_d)[self.vidx]
if mesh.cell_set._extruded:
Kform = (gammar('+') * ufl.jump(sigma, n=n) * ufl.dS_h +
gammar('+') * ufl.jump(sigma, n=n) * ufl.dS_v)
else:
Kform = (gammar('+') * ufl.jump(sigma, n=n) * ufl.dS)
# Here we deal with boundaries. If there are Neumann
# conditions (which should be enforced strongly for
# H(div)xL^2) then we need to add jump terms on the exterior
# facets. If there are Dirichlet conditions (which should be
# enforced weakly) then we need to zero out the trace
# variables there as they are not active (otherwise the hybrid
# problem is not well-posed).
# If boundary conditions are contained in the ImplicitMatrixContext:
if self.ctx.row_bcs:
# Find all the subdomains with neumann BCS
# These are Dirichlet BCs on the vidx space
neumann_subdomains = set()
for bc in self.ctx.row_bcs:
if bc.function_space().index == self.pidx:
raise NotImplementedError(
"Dirichlet conditions for scalar variable not supported. Use a weak bc"
)
if bc.function_space().index != self.vidx:
raise NotImplementedError(
"Dirichlet bc set on unsupported space.")
# append the set of sub domains
subdom = bc.sub_domain
if isinstance(subdom, str):
neumann_subdomains |= set([subdom])
else:
neumann_subdomains |= set(
as_tuple(subdom, numbers.Integral))
# separate out the top and bottom bcs
extruded_neumann_subdomains = neumann_subdomains & {
"top", "bottom"
}
neumann_subdomains = neumann_subdomains - extruded_neumann_subdomains
integrand = gammar * ufl.dot(sigma, n)
measures = []
trace_subdomains = []
if mesh.cell_set._extruded:
ds = ufl.ds_v
for subdomain in sorted(extruded_neumann_subdomains):
measures.append({
"top": ufl.ds_t,
"bottom": ufl.ds_b
}[subdomain])
trace_subdomains.extend(
sorted({"top", "bottom"} - extruded_neumann_subdomains))
else:
ds = ufl.ds
if "on_boundary" in neumann_subdomains:
measures.append(ds)
else:
measures.extend((ds(sd) for sd in sorted(neumann_subdomains)))
markers = [int(x) for x in mesh.exterior_facets.unique_markers]
dirichlet_subdomains = set(markers) - neumann_subdomains
trace_subdomains.extend(sorted(dirichlet_subdomains))
for measure in measures:
Kform += integrand * measure
trace_bcs = [
DirichletBC(TraceSpace, Constant(0.0), subdomain)
for subdomain in trace_subdomains
]
else:
# No bcs were provided, we assume weak Dirichlet conditions.
# We zero out the contribution of the trace variables on
# the exterior boundary. Extruded cells will have both
# horizontal and vertical facets
trace_subdomains = ["on_boundary"]
if mesh.cell_set._extruded:
trace_subdomains.extend(["bottom", "top"])
trace_bcs = [
DirichletBC(TraceSpace, Constant(0.0), subdomain)
for subdomain in trace_subdomains
]
# Make a SLATE tensor from Kform
K = Tensor(Kform)
# Assemble the Schur complement operator and right-hand side
self.schur_rhs = Function(TraceSpace)
self._assemble_Srhs = create_assembly_callable(
K * Atilde.inv * AssembledVector(self.broken_residual),
tensor=self.schur_rhs,
form_compiler_parameters=self.ctx.fc_params)
mat_type = PETSc.Options().getString(prefix + "mat_type", "aij")
schur_comp = K * Atilde.inv * K.T
self.S = allocate_matrix(schur_comp,
bcs=trace_bcs,
form_compiler_parameters=self.ctx.fc_params,
mat_type=mat_type,
options_prefix=prefix)
self._assemble_S = create_assembly_callable(
schur_comp,
tensor=self.S,
bcs=trace_bcs,
form_compiler_parameters=self.ctx.fc_params,
mat_type=mat_type)
with timed_region("HybridOperatorAssembly"):
self._assemble_S()
Smat = self.S.petscmat
nullspace = self.ctx.appctx.get("trace_nullspace", None)
if nullspace is not None:
nsp = nullspace(TraceSpace)
Smat.setNullSpace(nsp.nullspace(comm=pc.comm))
# Set up the KSP for the system of Lagrange multipliers
trace_ksp = PETSc.KSP().create(comm=pc.comm)
trace_ksp.setOptionsPrefix(prefix)
trace_ksp.setOperators(Smat)
trace_ksp.setUp()
trace_ksp.setFromOptions()
self.trace_ksp = trace_ksp
split_mixed_op = dict(split_form(Atilde.form))
split_trace_op = dict(split_form(K.form))
# Generate reconstruction calls
self._reconstruction_calls(split_mixed_op, split_trace_op)
V = VectorFunctionSpace(mesh, "Lagrange", 2)
Q = FunctionSpace(mesh, "Lagrange", 1)
# Define trial and test functions
u = TrialFunction(V)
p = TrialFunction(Q)
v = TestFunction(V)
q = TestFunction(Q)
# Define time-dependent pressure boundary condition
#p_in = Expression("sin(3.0*t)", t=0.0)
u_in = Expression(('1.0*(x[1]>0.5)+0.5*(x[1]<=0.5)', '0.0'))
# Define boundary conditions
noslip = DirichletBC(V, (0, 0), [3, 4])
#inflow = DirichletBC(Q, p_in, "x[0] < DOLFIN_EPS" )
inflow = DirichletBC(V, u_in, 1)
outflow = DirichletBC(Q, 0, 2)
bcu = [noslip, inflow]
bcp = [outflow]
# Create functions
u0 = Function(V)
u1 = Function(V)
p1 = Function(Q)
# Define coefficients
def pml(mesh,
scatterer_bdy_id,
outer_bdy_id,
wave_number,
options_prefix=None,
solver_parameters=None,
inner_region=None,
fspace=None,
tfspace=None,
true_sol_grad=None,
pml_type=None,
delta=None,
quad_const=None,
speed=None,
pml_min=None,
pml_max=None):
"""
For unlisted arg descriptions, see run_method
:arg inner_region: boundary id of non-pml region
:arg pml_type: Type of pml function, either 'quadratic' or 'bdy_integral'
:arg delta: For :arg:`pml_type` of 'bdy_integral', added to denominator
to prevent 1 / 0 at edge of boundary
:arg quad_const: For :arg:`pml_type` of 'quadratic', a scaling constant
:arg speed: Speed of sound
:arg pml_min: A list, *pml_min[i]* is where to begin pml layer in direction
*i*
:arg pml_max: A list, *pml_max[i]* is where to end pml layer in direction *i*
"""
# Handle defauls
if pml_type is None:
pml_type = 'bdy_integral'
if delta is None:
delta = 1e-3
if quad_const is None:
quad_const = 1.0
if speed is None:
speed = 340.0
pml_types = ['bdy_integral', 'quadratic']
if pml_type not in pml_types:
raise ValueError("PML type of %s is not one of %s" %
(pml_type, pml_types))
xx = SpatialCoordinate(mesh)
# {{{ create sigma functions for PML
sigma = None
if pml_type == 'bdy_integral':
sigma = [
Constant(speed) / (Constant(delta + extent) - abs(coord))
for extent, coord in zip(pml_max, xx)
]
elif pml_type == 'quadratic':
sigma = [
Constant(quad_const) * (abs(coord) - Constant(min_))**2
for min_, coord in zip(pml_min, xx)
]
r"""
Here \kappa is the wave number and c is the speed
..math::
\kappa = \frac{ \omega } { c }
"""
omega = wave_number * speed
# {{{ Set up PML functions
gamma = [
Constant(1.0) + conditional(
abs(real(coord)) >= real(min_),
Constant(1j / omega) * sigma_i, Constant(0.0))
for min_, coord, sigma_i in zip(pml_min, xx, sigma)
]
kappa = [None] * len(gamma)
gamma_prod = 1.0
for i in range(len(gamma)):
gamma_prod *= gamma[i]
tensor_i = [Constant(0.0) for _ in range(len(gamma))]
tensor_i[i] = 1.0
r"""
*i*th entry is
.. math::
\frac{\prod_{j\neq i} \gamma_j}{ \gamma_i }
"""
for j in range(len(gamma)):
if j != i:
tensor_i[i] *= gamma[j]
else:
tensor_i[i] /= gamma[j]
kappa[i] = tensor_i
kappa = as_tensor(kappa)
# }}}
p = TrialFunction(fspace)
q = TestFunction(fspace)
k = wave_number # Just easier to look at
a = (inner(dot(grad(p), kappa), grad(q)) -
Constant(k**2) * gamma_prod * inner(p, q)) * dx
n = FacetNormal(mesh)
L = inner(dot(true_sol_grad, n), q) * ds(scatterer_bdy_id)
bc = DirichletBC(fspace, Constant(0), outer_bdy_id)
solution = Function(fspace)
#solve(a == L, solution, bcs=[bc], options_prefix=options_prefix)
# Create a solver and return the KSP object with the solution so that can get
# PETSc information
# Create problem
problem = vs.LinearVariationalProblem(a, L, solution, [bc], None)
# Create solver and call solve
solver = vs.LinearVariationalSolver(problem,
solver_parameters=solver_parameters,
options_prefix=options_prefix)
solver.solve()
return solver.snes, solution
nx = 101 mesh = UnitSquareMesh(nx - 1, nx - 1) V = FunctionSpace(mesh, "CG", 1) u = TrialFunction(V) v = TestFunction(V) f = Function(V) x = SpatialCoordinate(mesh) f.interpolate((8 * pi * pi) * sin(x[0] * pi * 2) * sin(x[1] * pi * 2)) a = (dot(grad(v), grad(u))) * dx L = f * v * dx bc = DirichletBC(V, 0., "on_boundary") A = assemble(a, bcs=bc) b = assemble(L) u = Function(V) solve(A, u, b) # Create some fake data that is systematically different from the FEM solution. # note that all parameters are on a log scale, so we take the true values # and take the logarithm # forcing covariance parameters (taken to be known) sigma_f = np.log(2.e-2)
