def __init__(self, state, V, *, ibp="once", solver_params=None):
super().__init__(state=state,
V=V,
ibp=ibp,
solver_params=solver_params)
self.Upwind = 0.5 * (sign(dot(self.ubar, self.n)) + 1)
if self.state.mesh.topological_dimension() == 2:
self.perp = state.perp
if state.on_sphere:
outward_normals = CellNormal(state.mesh)
self.perp_u_upwind = lambda q: self.Upwind('+') * cross(
outward_normals('+'), q('+')) + self.Upwind('-') * cross(
outward_normals('-'), q('-'))
else:
self.perp_u_upwind = lambda q: self.Upwind('+') * state.perp(
q('+')) + self.Upwind('-') * state.perp(q('-'))
self.gradperp = lambda u: state.perp(grad(u))
elif self.state.mesh.topological_dimension() == 3:
if self.ibp == "twice":
raise NotImplementedError(
"ibp=twice is not implemented for 3d problems")
else:
raise RuntimeError("topological mesh dimension must be 2 or 3")
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 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 setup(self, state, vorticity_type=None):
"""Solver for vorticity.
:arg state: The state containing model.
:arg vorticity_type: must be "relative", "absolute" or "potential"
"""
if not self._initialised:
vorticity_types = ["relative", "absolute", "potential"]
if vorticity_type not in vorticity_types:
raise ValueError("vorticity type must be one of %s, not %s" % (vorticity_types, vorticity_type))
try:
space = state.spaces("CG")
except AttributeError:
dgspace = state.spaces("DG")
cg_degree = dgspace.ufl_element().degree() + 2
space = FunctionSpace(state.mesh, "CG", cg_degree)
super().setup(state, space=space)
u = state.fields("u")
gamma = TestFunction(space)
q = TrialFunction(space)
if vorticity_type == "potential":
D = state.fields("D")
a = q*gamma*D*dx
else:
a = q*gamma*dx
if state.on_sphere:
cell_normals = CellNormal(state.mesh)
gradperp = lambda psi: cross(cell_normals, grad(psi))
L = (- inner(gradperp(gamma), u))*dx
else:
raise NotImplementedError("The vorticity diagnostics have only been implemented for 2D spherical geometries.")
if vorticity_type != "relative":
f = state.fields("coriolis")
L += gamma*f*dx
problem = LinearVariationalProblem(a, L, self.field)
self.solver = LinearVariationalSolver(problem, solver_parameters={"ksp_type": "cg"})
def __init__(self,
mesh,
dt,
output=None,
parameters=None,
diagnostics=None,
diagnostic_fields=None):
if output is None:
raise RuntimeError(
"You must provide a directory name for dumping results")
else:
self.output = output
self.parameters = parameters
if diagnostics is not None:
self.diagnostics = diagnostics
else:
self.diagnostics = Diagnostics()
if diagnostic_fields is not None:
self.diagnostic_fields = diagnostic_fields
else:
self.diagnostic_fields = []
# The mesh
self.mesh = mesh
self.spaces = SpaceCreator(mesh)
if self.output.dumplist is None:
self.output.dumplist = []
self.fields = StateFields(*self.output.dumplist)
self.dumpdir = None
self.dumpfile = None
self.to_pickup = None
# figure out if we're on a sphere
try:
self.on_sphere = (mesh._base_mesh.geometric_dimension() == 3
and mesh._base_mesh.topological_dimension() == 2)
except AttributeError:
self.on_sphere = (mesh.geometric_dimension() == 3
and mesh.topological_dimension() == 2)
# build the vertical normal and define perp for 2d geometries
dim = mesh.topological_dimension()
if self.on_sphere:
x = SpatialCoordinate(mesh)
R = sqrt(inner(x, x))
self.k = interpolate(x / R, mesh.coordinates.function_space())
if dim == 2:
outward_normals = CellNormal(mesh)
self.perp = lambda u: cross(outward_normals, u)
else:
kvec = [0.0] * dim
kvec[dim - 1] = 1.0
self.k = Constant(kvec)
if dim == 2:
self.perp = lambda u: as_vector([-u[1], u[0]])
# setup logger
logger.setLevel(output.log_level)
set_log_handler(mesh.comm)
if parameters is not None:
logger.info("Physical parameters that take non-default values:")
logger.info(", ".join("%s: %s" % (k, float(v))
for (k, v) in vars(parameters).items()))
# Constant to hold current time
self.t = Constant(0.0)
if type(dt) is Constant:
self.dt = dt
elif type(dt) in (float, int):
self.dt = Constant(dt)
else:
raise TypeError(
f'dt must be a Constant, float or int, not {type(dt)}')
Dn.interpolate(Dexpr)
b.interpolate(bexpr)
if COMM_WORLD.Get_rank() == 0:
print("Finished setting up functions at", ctime())
# Build perp and upwind perp (including 2D version for test purposes)
n = FacetNormal(mesh)
s = lambda u: 0.5 * (sign(dot(u, n)) + 1)
uw = lambda u, v: (s(u)('+') * v('+') + s(u)('-') * v('-'))
if mesh.geometric_dimension() == 2:
perp = lambda u: as_vector([-u[1], u[0]])
p_uw = lambda u, v: perp(uw(u, v))
else:
perp = lambda u: cross(CellNormal(mesh), u)
out_n = CellNormal(mesh)
p_uw = lambda u, v: (s(u)('+') * cross(out_n('+'), v('+')) + s(u)
('-') * cross(out_n('-'), v('-')))
# Initial solve for vorticity for output
eta = TestFunction(W2)
q_ = TrialFunction(W2)
q_eqn = eta * q_ * Dn * dx + inner(perp(grad(eta)), un) * dx - eta * f * dx
q_p = LinearVariationalProblem(lhs(q_eqn), rhs(q_eqn), qn)
q_solver = LinearVariationalSolver(q_p,
solver_parameters={
"ksp_type": "preonly",
"pc_type": "lu"
})
q_solver.solve()
def __init__(self,
mesh,
vertical_degree=None,
horizontal_degree=1,
family="RT",
Coriolis=None,
sponge_function=None,
hydrostatic=None,
timestepping=None,
output=None,
parameters=None,
diagnostics=None,
fieldlist=None,
diagnostic_fields=None,
u_bc_ids=None):
self.family = family
self.vertical_degree = vertical_degree
self.horizontal_degree = horizontal_degree
self.Omega = Coriolis
self.mu = sponge_function
self.hydrostatic = hydrostatic
self.timestepping = timestepping
if output is None:
raise RuntimeError(
"You must provide a directory name for dumping results")
else:
self.output = output
self.parameters = parameters
if fieldlist is None:
raise RuntimeError(
"You must provide a fieldlist containing the names of the prognostic fields"
)
else:
self.fieldlist = fieldlist
if diagnostics is not None:
self.diagnostics = diagnostics
else:
self.diagnostics = Diagnostics(*fieldlist)
if diagnostic_fields is not None:
self.diagnostic_fields = diagnostic_fields
else:
self.diagnostic_fields = []
if u_bc_ids is not None:
self.u_bc_ids = u_bc_ids
else:
self.u_bc_ids = []
# The mesh
self.mesh = mesh
# Build the spaces
self._build_spaces(mesh, vertical_degree, horizontal_degree, family)
# Allocate state
self._allocate_state()
if self.output.dumplist is None:
self.output.dumplist = fieldlist
self.fields = FieldCreator(fieldlist, self.xn, self.output.dumplist)
# set up bcs
V = self.fields('u').function_space()
self.bcs = []
if V.extruded:
self.bcs.append(DirichletBC(V, 0.0, "bottom"))
self.bcs.append(DirichletBC(V, 0.0, "top"))
for id in self.u_bc_ids:
self.bcs.append(DirichletBC(V, 0.0, id))
self.dumpfile = None
# figure out if we're on a sphere
try:
self.on_sphere = (mesh._base_mesh.geometric_dimension() == 3
and mesh._base_mesh.topological_dimension() == 2)
except AttributeError:
self.on_sphere = (mesh.geometric_dimension() == 3
and mesh.topological_dimension() == 2)
# build the vertical normal and define perp for 2d geometries
dim = mesh.topological_dimension()
if self.on_sphere:
x = SpatialCoordinate(mesh)
R = sqrt(inner(x, x))
self.k = interpolate(x / R, mesh.coordinates.function_space())
if dim == 2:
outward_normals = CellNormal(mesh)
self.perp = lambda u: cross(outward_normals, u)
else:
kvec = [0.0] * dim
kvec[dim - 1] = 1.0
self.k = Constant(kvec)
if dim == 2:
self.perp = lambda u: as_vector([-u[1], u[0]])
# project test function for hydrostatic case
if self.hydrostatic:
self.h_project = lambda u: u - self.k * inner(u, self.k)
else:
self.h_project = lambda u: u
# Constant to hold current time
self.t = Constant(0.0)
# setup logger
logger.setLevel(output.log_level)
set_log_handler(mesh.comm)
logger.info("Timestepping parameters that take non-default values:")
logger.info(", ".join("%s: %s" % item
for item in vars(timestepping).items()))
if parameters is not None:
logger.info("Physical parameters that take non-default values:")
logger.info(", ".join("%s: %s" % item
for item in vars(parameters).items()))