class InteriorPenalty(Diffusion):
"""
Interior penalty diffusion method
:arg state: :class:`.State` object.
:arg V: Function space of diffused field
:arg direction: list containing directions in which function space
is discontinuous: 1 corresponds to vertical, 2 to horizontal.
:arg params: dictionary containing the interior penalty parameters
:mu and kappa where mu is the penalty weighting function, which is
:recommended to be proportional to 1/dx
"""
def __init__(self, state, V, direction=[1,2], params=None):
super(InteriorPenalty, self).__init__(state)
dt = state.timestepping.dt
kappa = params['kappa']
mu = params['mu']
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
if 1 in direction:
a += dt*get_flux_form(dS_v, kappa)
if 2 in direction:
a += dt*get_flux_form(dS_h, kappa)
L = inner(gamma,phi)*dx
problem = LinearVariationalProblem(a, action(L,self.phi1), self.phi1)
self.solver = LinearVariationalSolver(problem)
def apply(self, x_in, x_out):
self.phi1.assign(x_in)
self.solver.solve()
x_out.assign(self.phi1)
class ShallowWaterForcing(Forcing):
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 apply(self, scaling, x_in, x_nl, x_out, **kwargs):
self.x0.assign(x_nl)
self.u_forcing_solver.solve() # places forcing in self.uF
self.uF *= scaling
uF, _ = x_out.split()
x_out.assign(x_in)
uF += self.uF
class CompressibleForcing(Forcing):
"""
Forcing class for compressible Euler equations.
"""
def __init__(self, state, linear=False):
self.state = state
self._build_forcing_solver(linear)
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 apply(self, scaling, x_in, x_nl, x_out, **kwargs):
self.x0.assign(x_nl)
self.scaling.assign(scaling)
if "mu_alpha" in kwargs and kwargs["mu_alpha"] is not None:
self.mu_scaling.assign(kwargs["mu_alpha"])
self.u_forcing_solver.solve() # places forcing in self.uF
u_out, _, _ = x_out.split()
x_out.assign(x_in)
u_out += self.uF
class SUPGAdvection(Advection):
"""
An SUPG advection scheme that can apply DG upwinding (in the direction
specified by the direction arg) if the function space is only
partially continuous.
:arg state: :class:`.State` object.
:arg V:class:`.FunctionSpace` object. The advected field function space.
:arg direction: list containing the directions in which the function
space is discontinuous. 1 corresponds to the vertical direction, 2 to
the horizontal direction
:arg supg_params: dictionary containing SUPG parameters tau for each
direction. If not supplied tau is set to dt/sqrt(15.)
"""
def __init__(self, state, V, direction=[], supg_params=None):
super(SUPGAdvection, self).__init__(state)
dt = state.timestepping.dt
params = supg_params.copy() if supg_params else {}
params.setdefault('a0', dt/sqrt(15.))
params.setdefault('a1', dt/sqrt(15.))
gamma = TestFunction(V)
theta = TrialFunction(V)
self.theta0 = Function(V)
# make SUPG test function
taus = [params["a0"], params["a1"]]
for i in direction:
taus[i] = 0.0
tau = Constant(((taus[0], 0.), (0., taus[1])))
dgamma = dot(dot(self.ubar, tau), grad(gamma))
gammaSU = gamma + dgamma
n = FacetNormal(state.mesh)
un = 0.5*(dot(self.ubar, n) + abs(dot(self.ubar, n)))
a_mass = gammaSU*theta*dx
arhs = a_mass - dt*gammaSU*dot(self.ubar, grad(theta))*dx
if 1 in direction:
arhs -= (
dt*dot(jump(gammaSU), (un('+')*theta('+')
- un('-')*theta('-')))*dS_v
- dt*(gammaSU('+')*dot(self.ubar('+'), n('+'))*theta('+')
+ gammaSU('-')*dot(self.ubar('-'), n('-'))*theta('-'))*dS_v
)
if 2 in direction:
arhs -= (
dt*dot(jump(gammaSU), (un('+')*theta('+')
- un('-')*theta('-')))*dS_h
- dt*(gammaSU('+')*dot(self.ubar('+'), n('+'))*theta('+')
+ gammaSU('-')*dot(self.ubar('-'), n('-'))*theta('-'))*dS_h
)
self.theta1 = Function(V)
self.dtheta = Function(V)
problem = LinearVariationalProblem(a_mass, action(arhs,self.theta1), self.dtheta)
self.solver = LinearVariationalSolver(problem,
options_prefix='SUPGAdvection')
def apply(self, x_in, x_out):
# SSPRK Stage 1
self.theta1.assign(x_in)
self.solver.solve()
self.theta1.assign(self.dtheta)
# SSPRK Stage 2
self.solver.solve()
self.theta1.assign(0.75*x_in + 0.25*self.dtheta)
# SSPRK Stage 3
self.solver.solve()
x_out.assign((1.0/3.0)*x_in + (2.0/3.0)*self.dtheta)
class IncompressibleForcing(Forcing):
"""
Forcing class for incompressible Euler Boussinesq equations.
"""
def __init__(self, state, linear=False):
self.state = state
self._build_forcing_solver(linear)
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 apply(self, scaling, x_in, x_nl, x_out, **kwargs):
self.x0.assign(x_nl)
self.scaling.assign(scaling)
if "mu_alpha" in kwargs and kwargs["mu_alpha"] is not None:
self.mu_scaling.assign(kwargs["mu_alpha"])
self.u_forcing_solver.solve() # places forcing in self.uF
u_out, p_out, _ = x_out.split()
x_out.assign(x_in)
u_out += self.uF
if "incompressible" in kwargs and kwargs["incompressible"]:
self.divergence_solver.solve()
p_out.assign(self.divu)
class DGAdvection(Advection):
"""
DG 3 step SSPRK advection scheme that can be applied to a scalar
or vector field
:arg state: :class:`.State` object.
:arg V: function space of advected field - should be DG
:arg continuity: optional boolean.
If ``True``, the advection equation is of the form:
:math: `D_t +\nabla \cdot(uD) = 0`.
If ``False``, the advection equation is of the form:
:math: `D_t + (u \cdot \nabla)D = 0`.
"""
def __init__(self, state, V, continuity=False):
super(DGAdvection, self).__init__(state)
element = V.fiat_element
assert element.entity_dofs() == element.entity_closure_dofs(), "Provided space is not discontinuous"
dt = state.timestepping.dt
if V.extruded:
surface_measure = (dS_h + dS_v)
else:
surface_measure = dS
phi = TestFunction(V)
D = TrialFunction(V)
self.D1 = Function(V)
self.dD = Function(V)
n = FacetNormal(state.mesh)
# ( dot(v, n) + |dot(v, n)| )/2.0
un = 0.5*(dot(self.ubar, n) + abs(dot(self.ubar, n)))
a_mass = inner(phi,D)*dx
if continuity:
a_int = -inner(grad(phi), outer(D, self.ubar))*dx
else:
a_int = -inner(div(outer(phi,self.ubar)),D)*dx
a_flux = (dot(jump(phi), un('+')*D('+') - un('-')*D('-')))*surface_measure
arhs = a_mass - dt*(a_int + a_flux)
DGproblem = LinearVariationalProblem(a_mass, action(arhs,self.D1),
self.dD)
self.DGsolver = LinearVariationalSolver(DGproblem,
solver_parameters={
'ksp_type':'preonly',
'pc_type':'bjacobi',
'sub_pc_type': 'ilu'},
options_prefix='DGAdvection')
def apply(self, x_in, x_out):
# SSPRK Stage 1
self.D1.assign(x_in)
self.DGsolver.solve()
self.D1.assign(self.dD)
# SSPRK Stage 2
self.DGsolver.solve()
self.D1.assign(0.75*x_in + 0.25*self.dD)
# SSPRK Stage 3
self.DGsolver.solve()
x_out.assign((1.0/3.0)*x_in + (2.0/3.0)*self.dD)
class EulerPoincareForm(Advection):
def __init__(self, state, V):
super(EulerPoincareForm, self).__init__(state)
dt = state.timestepping.dt
w = TestFunction(V)
u = TrialFunction(V)
self.u0 = Function(V)
ustar = 0.5*(self.u0 + u)
n = FacetNormal(state.mesh)
Upwind = 0.5*(sign(dot(self.ubar, n))+1)
if state.mesh.geometric_dimension() == 3:
if V.extruded:
surface_measure = (dS_h + dS_v)
else:
surface_measure = dS
# <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)
Eqn = (
inner(w, u-self.u0)*dx
+ dt*inner(ustar, curl(cross(self.ubar, w)))*dx
- dt*inner(both(Upwind*ustar),
both(cross(n, cross(self.ubar, w))))*surface_measure
- dt*div(w)*inner(ustar, self.ubar)*dx
)
# define surface measure and terms involving perp differently
# for slice (i.e. if V.extruded is True) and shallow water
# (V.extruded is False)
else:
if V.extruded:
surface_measure = (dS_h + dS_v)
perp = lambda u: as_vector([-u[1], u[0]])
perp_u_upwind = Upwind('+')*perp(ustar('+')) + Upwind('-')*perp(ustar('-'))
else:
surface_measure = dS
outward_normals = CellNormal(state.mesh)
perp = lambda u: cross(outward_normals, u)
perp_u_upwind = Upwind('+')*cross(outward_normals('+'),ustar('+')) + Upwind('-')*cross(outward_normals('-'),ustar('-'))
Eqn = (
(inner(w, u-self.u0)
- dt*inner(w, div(perp(ustar))*perp(self.ubar))
- dt*div(w)*inner(ustar, self.ubar))*dx
- dt*inner(jump(inner(w, perp(self.ubar)), n),
perp_u_upwind)*surface_measure
+ dt*jump(inner(w,
perp(self.ubar))*perp(ustar), n)*surface_measure
)
a = lhs(Eqn)
L = rhs(Eqn)
self.u1 = Function(V)
uproblem = LinearVariationalProblem(a, L, self.u1)
self.usolver = LinearVariationalSolver(uproblem,
options_prefix='EPAdvection')
def apply(self, x_in, x_out):
self.u0.assign(x_in)
self.usolver.solve()
x_out.assign(self.u1)
def run(steady=False):
"""
solve CdT/dt = S + div(k*grad(T))
=> C*v*(dT/dt)/k*dx - S*v/k*dx + grad(v)*grad(T)*dx = v*dot(grad(T), n)*ds
"""
steps = 250
dt = 1e-10
timescale = (0, steps * dt)
if steady:
print('Running steady state.')
else:
print(f'Running with time step {dt:.2g}s on time interval: '
f'{timescale[0]:.2g}s - {timescale[1]:.2g}s')
dt_invc = Constant(1 / dt)
extent = [40e-6, 40e-6, 40e-6]
mesh = BoxMesh(20, 20, 20, *extent)
V = FunctionSpace(mesh, 'CG', 1)
print(V.dim())
T = Function(V) # temperature at time i+1 (electron for now)
T_ = Function(V) # temperature at time i
v = TestFunction(V) # test function
S = create_S(mesh, V, extent)
C = create_heat_capacity(mesh, V, extent)
k = create_conductivity(mesh, V, T)
set_initial_value(mesh, T_, extent)
# Mass matrix section
M = C * T * dt_invc * v * dx
M_ = C * T_ * dt_invc * v * dx
# Stiffness matrix section
A = k * dot(grad(T), grad(v)) * dx
# function section
f = S * v * dx
# boundaries
bcs, R, b = create_dirichlet_bounds(mesh,
V,
T,
v,
k,
g=100,
boundary=[1, 2, 3, 4, 5, 6])
# bcs += create_dirichlet_bounds(mesh, V, T, v, k, 500, [6])[0]
# bcs, R, b = create_robin_bounds(mesh, T, v, k, 1e8/(100), 1e8)
if steady:
steps = 1
a = A + R
L = f + b
else:
a = M + A + R
L = M_ + f + b
prob = NonlinearVariationalProblem(a - L, T, bcs=bcs)
solver = NonlinearVariationalSolver(prob, solver_parameters=SOLVE_PARAMS)
T.assign(T_)
timestamp = datetime.now().strftime("%d-%b-%Y-%H-%M-%S")
outfile = File(f'{timestamp}/first_output.pvd')
outfile.write(T_, target_degree=1, target_continuity=H1)
last_perc = 0
for i in range(steps):
solver.solve()
perc = int(100 * (i + 1) / steps)
if perc > last_perc:
print(f'{perc}%')
last_perc = perc
T_.assign(T)
outfile.write(T_, target_degree=1, target_continuity=H1)
def _ad_copy(self):
from firedrake import Function
r = Function(self.function_space())
r.assign(self)
return r
def _ad_add(self, other):
from firedrake import Function
r = Function(self.function_space())
Function.assign(r, self + other)
return r
def _ad_mul(self, other):
from firedrake import Function
r = Function(self.function_space())
r.assign(self * other)
return r
class Forcing(object, metaclass=ABCMeta):
"""
Base class for forcing terms for Gusto.
:arg state: x :class:`.State` object.
:arg euler_poincare: if True then the momentum equation is in Euler
Poincare form and we need to add 0.5*grad(u^2) to the forcing term.
If False then this term is not added.
:arg linear: if True then we are solving a linear equation so nonlinear
terms (namely the Euler Poincare term) should not be added.
:arg extra_terms: extra terms to add to the u component of the forcing
term - these will be multiplied by the appropriate test function.
"""
def __init__(self,
state,
euler_poincare=True,
linear=False,
extra_terms=None,
moisture=None):
self.state = state
if linear:
self.euler_poincare = False
logger.warning(
'Setting euler_poincare to False because you have set linear=True'
)
else:
self.euler_poincare = euler_poincare
# set up functions
self.Vu = state.spaces("HDiv")
# this is the function that the forcing term is applied to
self.x0 = Function(state.W)
self.test = TestFunction(self.Vu)
self.trial = TrialFunction(self.Vu)
# this is the function that contains the result of solving
# <test, trial> = <test, F(x0)>, where F is the forcing term
self.uF = Function(self.Vu)
# find out which terms we need
self.extruded = self.Vu.extruded
self.coriolis = state.Omega is not None or hasattr(
state.fields, "coriolis")
self.sponge = state.mu is not None
self.hydrostatic = state.hydrostatic
self.topography = hasattr(state.fields, "topography")
self.extra_terms = extra_terms
self.moisture = moisture
# some constants to use for scaling terms
self.scaling = Constant(1.)
self.impl = Constant(1.)
self._build_forcing_solvers()
def mass_term(self):
return inner(self.test, self.trial) * dx
def coriolis_term(self):
u0 = split(self.x0)[0]
return -inner(self.test, cross(2 * self.state.Omega, u0)) * dx
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 euler_poincare_term(self):
u0 = split(self.x0)[0]
return -0.5 * div(self.test) * inner(self.state.h_project(u0), u0) * dx
def hydrostatic_term(self):
u0 = split(self.x0)[0]
return inner(u0, self.state.k) * inner(self.test, self.state.k) * dx
@abstractmethod
def pressure_gradient_term(self):
pass
def forcing_term(self):
L = self.pressure_gradient_term()
if self.extruded:
L += self.gravity_term()
if self.coriolis:
L += self.coriolis_term()
if self.euler_poincare:
L += self.euler_poincare_term()
if self.topography:
L += self.topography_term()
if self.extra_terms is not None:
L += inner(self.test, self.extra_terms) * dx
# scale L
L = self.scaling * L
# sponge term has a separate scaling factor as it is always implicit
if self.sponge:
L -= self.impl * self.state.timestepping.dt * self.sponge_term()
# hydrostatic term has no scaling factor
if self.hydrostatic:
L += (2 * self.impl - 1) * self.hydrostatic_term()
return L
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 apply(self, scaling, x_in, x_nl, x_out, **kwargs):
"""
Function takes x as input, computes F(x_nl) and returns
x_out = x + scale*F(x_nl)
as output.
:arg x_in: :class:`.Function` object
:arg x_nl: :class:`.Function` object
:arg x_out: :class:`.Function` object
:arg implicit: forcing stage for sponge and hydrostatic terms, if present
"""
self.scaling.assign(scaling)
self.x0.assign(x_nl)
implicit = kwargs.get("implicit")
if implicit is not None:
self.impl.assign(int(implicit))
self.u_forcing_solver.solve() # places forcing in self.uF
uF = x_out.split()[0]
x_out.assign(x_in)
uF += self.uF
