def __init__(self, state, accretion=True, accumulation=True):
super().__init__(state)
# obtain our fields
self.water_c = state.fields('cloud_liquid_mixing_ratio')
self.rain = state.fields('rain_mixing_ratio')
# declare function space
Vt = self.water_c.function_space()
# define some parameters as attributes
dt = state.dt
k_1 = Constant(0.001) # accretion rate in 1/s
k_2 = Constant(2.2) # accumulation rate in 1/s
a = Constant(0.001) # min cloud conc in kg/kg
b = Constant(0.875) # power for rain in accumulation
# make default rates to be zero
accr_rate = Constant(0.0)
accu_rate = Constant(0.0)
if accretion:
accr_rate = k_1 * (self.water_c - a)
if accumulation:
accu_rate = k_2 * self.water_c * self.rain ** b
# make coalescence rate function, that needs to be the same for all updates in one time step
coalesce_rate = Function(Vt)
# adjust coalesce rate using min_value so negative cloud concentration doesn't occur
self.lim_coalesce_rate = Interpolator(conditional(self.rain < 0.0, # if rain is negative do only accretion
conditional(accr_rate < 0.0,
0.0,
min_value(accr_rate, self.water_c / dt)),
# don't turn rain back into cloud
conditional(accr_rate + accu_rate < 0.0,
0.0,
# if accretion rate is negative do only accumulation
conditional(accr_rate < 0.0,
min_value(accu_rate, self.water_c / dt),
min_value(accr_rate + accu_rate, self.water_c / dt)))),
coalesce_rate)
# tell the prognostic fields what to update to
self.water_c_new = Interpolator(self.water_c - dt * coalesce_rate, Vt)
self.rain_new = Interpolator(self.rain + dt * coalesce_rate, Vt)
def PeriodicIntervalMesh(ncells, length, comm=COMM_WORLD):
"""Generate a periodic mesh of an interval.
:arg ncells: The number of cells over the interval.
:arg length: The length the interval.
:kwarg comm: Optional communicator to build the mesh on (defaults to
COMM_WORLD).
"""
if ncells < 3:
raise ValueError("1D periodic meshes with fewer than 3 \
cells are not currently supported")
m = CircleManifoldMesh(ncells, comm=comm)
coord_fs = VectorFunctionSpace(m, 'DG', 1, dim=1)
old_coordinates = m.coordinates
new_coordinates = Function(coord_fs)
periodic_kernel = """double Y,pi;
Y = 0.5*(old_coords[0][1]-old_coords[1][1]);
pi=3.141592653589793;
for(int i=0;i<2;i++){
new_coords[i][0] = atan2(old_coords[i][1],old_coords[i][0])/pi/2;
if(new_coords[i][0]<0.) new_coords[i][0] += 1;
if(new_coords[i][0]==0 && Y<0.) new_coords[i][0] = 1.0;
new_coords[i][0] *= L[0];
}"""
cL = Constant(length)
par_loop(
periodic_kernel, dx, {
"new_coords": (new_coordinates, WRITE),
"old_coords": (old_coordinates, READ),
"L": (cL, READ)
})
return mesh.Mesh(new_coordinates)
def __init__(self, *args, **kwargs):
"""
Initialiser for TimeMixin.
Add M/dt to a and L.
"""
super().__init__(*args, **kwargs)
self._add_function('T')
self._add_function('a')
self._add_function('T_')
self._add_function('C')
self._add_function('_delT')
self.max_t = 1e-9
self.dt = 1e-10
self._dt_invc = Constant(0)
self.steps = self.max_t / self.dt
self.steady_state = True
self.a += self._M()
def setup(self, **kwargs):
for name, field in kwargs.items():
if name in self._fields.keys():
self._fields[name].assign(field)
else:
self._fields[name] = utilities.copy(field)
# Create homogeneous BCs for the Dirichlet part of the boundary
u = self._fields.get('velocity', self._fields.get('u'))
V = u.function_space()
# NOTE: This will have to change when we do Stokes!
bcs = firedrake.DirichletBC(V, Constant((0, 0)), self._dirichlet_ids)
if not self._dirichlet_ids:
bcs = None
# Find the numeric IDs for the ice front
boundary_ids = u.ufl_domain().exterior_facets.unique_markers
ice_front_ids_comp = set(self._dirichlet_ids + self._side_wall_ids)
ice_front_ids = list(set(boundary_ids) - ice_front_ids_comp)
# Create the action and scale functionals
_kwargs = {
'side_wall_ids': self._side_wall_ids,
'ice_front_ids': ice_front_ids
}
action = self._model.action(**self._fields, **_kwargs)
scale = self._model.scale(**self._fields, **_kwargs)
# Set up a minimization problem and solver
quadrature_degree = self._model.quadrature_degree(**self._fields)
params = {'quadrature_degree': quadrature_degree}
problem = MinimizationProblem(action, scale, u, bcs, params)
self._solver = NewtonSolver(problem,
self._tolerance,
solver_parameters=self._solver_parameters,
max_iterations=self._max_iterations)
def test_vector_diffusion(tmpdir, DG, tracer_setup):
setup = tracer_setup(tmpdir, geometry="slice", blob=True)
state = setup.state
f_init = setup.f_init
tmax = setup.tmax
tol = 3.e-2
kappa = 1.
f_end_expr = (1 / (1 + 4 * tmax)) * f_init**(1 / (1 + 4 * tmax))
kappa = Constant([[kappa, 0.], [0., kappa]])
if DG:
V = VectorFunctionSpace(state.mesh, "DG", 1)
else:
V = state.spaces("HDiv", "CG", 1)
f_init = as_vector([f_init, 0.])
f_end_expr = as_vector([f_end_expr, 0.])
mu = 5.
diffusion_params = DiffusionParameters(kappa=kappa, mu=mu)
eqn = DiffusionEquation(state,
V,
"f",
diffusion_parameters=diffusion_params)
if DG:
state.fields("f").interpolate(f_init)
else:
state.fields("f").project(f_init)
diffusion_scheme = [(eqn, BackwardEuler(state))]
f_end = run(state, diffusion_scheme, tmax)
assert errornorm(f_end_expr, f_end) < tol
def setup_saturated(dirname, recovered):
# set up grid and time stepping parameters
dt = 1.
tmax = 3.
deltax = 400.
L = 2000.
H = 10000.
nlayers = int(H / deltax)
ncolumns = int(L / deltax)
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
# option to easily change between recovered and not if necessary
# default should be to use lowest order set of spaces
degree = 0 if recovered else 1
output = OutputParameters(dirname=dirname + '/saturated_balance',
dumpfreq=1,
dumplist=['u'])
parameters = CompressibleParameters()
diagnostic_fields = [Theta_e()]
state = State(mesh,
dt=dt,
output=output,
parameters=parameters,
diagnostic_fields=diagnostic_fields)
tracers = [WaterVapour(), CloudWater()]
if recovered:
u_transport_option = "vector_advection_form"
else:
u_transport_option = "vector_invariant_form"
eqns = CompressibleEulerEquations(state,
"CG",
degree,
u_transport_option=u_transport_option,
active_tracers=tracers)
# Initial conditions
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
water_v0 = state.fields("vapour_mixing_ratio")
water_c0 = state.fields("cloud_liquid_mixing_ratio")
moisture = ['vapour_mixing_ratio', 'cloud_liquid_mixing_ratio']
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
# Isentropic background state
Tsurf = Constant(300.)
total_water = Constant(0.02)
theta_e = Function(Vt).interpolate(Tsurf)
water_t = Function(Vt).interpolate(total_water)
# Calculate hydrostatic exner
saturated_hydrostatic_balance(state, theta_e, water_t)
water_c0.assign(water_t - water_v0)
state.set_reference_profiles([('rho', rho0), ('theta', theta0)])
# Set up transport schemes
if recovered:
VDG1 = state.spaces("DG1_equispaced")
VCG1 = FunctionSpace(mesh, "CG", 1)
Vt_brok = FunctionSpace(mesh, BrokenElement(Vt.ufl_element()))
Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element())
Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1)
u_opts = RecoveredOptions(embedding_space=Vu_DG1,
recovered_space=Vu_CG1,
broken_space=Vu,
boundary_method=Boundary_Method.dynamics)
rho_opts = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=Vr,
boundary_method=Boundary_Method.dynamics)
theta_opts = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=Vt_brok)
wv_opts = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=Vt_brok)
wc_opts = RecoveredOptions(embedding_space=VDG1,
recovered_space=VCG1,
broken_space=Vt_brok)
else:
rho_opts = None
theta_opts = EmbeddedDGOptions()
wv_opts = EmbeddedDGOptions()
wc_opts = EmbeddedDGOptions()
transported_fields = [
SSPRK3(state, 'rho', options=rho_opts),
SSPRK3(state, 'theta', options=theta_opts),
SSPRK3(state, 'vapour_mixing_ratio', options=wv_opts),
SSPRK3(state, 'cloud_liquid_mixing_ratio', options=wc_opts)
]
if recovered:
transported_fields.append(SSPRK3(state, 'u', options=u_opts))
else:
transported_fields.append(ImplicitMidpoint(state, 'u'))
linear_solver = CompressibleSolver(state, eqns, moisture=moisture)
# add physics
physics_list = [Condensation(state)]
# build time stepper
stepper = CrankNicolson(state,
eqns,
transported_fields,
linear_solver=linear_solver,
physics_list=physics_list)
return stepper, tmax
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
def nonlocal_integral_eq(
mesh,
scatterer_bdy_id,
outer_bdy_id,
wave_number,
options_prefix=None,
solver_parameters=None,
fspace=None,
vfspace=None,
true_sol_grad_expr=None,
actx=None,
dgfspace=None,
dgvfspace=None,
meshmode_src_connection=None,
qbx_kwargs=None,
):
r"""
see run_method for descriptions of unlisted args
args:
gamma and beta are used to precondition
with the following equation:
\Delta u - \kappa^2 \gamma u = 0
(\partial_n - i\kappa\beta) u |_\Sigma = 0
"""
# make sure we get outer bdy id as tuple in case it consists of multiple ids
if isinstance(outer_bdy_id, int):
outer_bdy_id = [outer_bdy_id]
outer_bdy_id = tuple(outer_bdy_id)
# away from the excluded region, but firedrake and meshmode point
# into
pyt_inner_normal_sign = -1
ambient_dim = mesh.geometric_dimension()
# {{{ Build src and tgt
# build connection meshmode near src boundary -> src boundary inside meshmode
from meshmode.discretization.poly_element import \
InterpolatoryQuadratureSimplexGroupFactory
from meshmode.discretization.connection import make_face_restriction
factory = InterpolatoryQuadratureSimplexGroupFactory(
dgfspace.finat_element.degree)
src_bdy_connection = make_face_restriction(actx,
meshmode_src_connection.discr,
factory, scatterer_bdy_id)
# source is a qbx layer potential
from pytential.qbx import QBXLayerPotentialSource
disable_refinement = (fspace.mesh().geometric_dimension() == 3)
qbx = QBXLayerPotentialSource(src_bdy_connection.to_discr,
**qbx_kwargs,
_disable_refinement=disable_refinement)
# get target indices and point-set
target_indices, target = get_target_points_and_indices(
fspace, outer_bdy_id)
# }}}
# build the operations
from pytential import bind, sym
r"""
..math:
x \in \Sigma
grad_op(x) =
\nabla(
\int_\Gamma(
u(y) \partial_n H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
)
"""
grad_op = pyt_inner_normal_sign * sym.grad(
ambient_dim,
sym.D(HelmholtzKernel(ambient_dim),
sym.var("u"),
k=sym.var("k"),
qbx_forced_limit=None))
r"""
..math:
x \in \Sigma
op(x) =
i \kappa \cdot
\int_\Gamma(
u(y) \partial_n H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
"""
op = pyt_inner_normal_sign * 1j * sym.var("k") * (sym.D(
HelmholtzKernel(ambient_dim),
sym.var("u"),
k=sym.var("k"),
qbx_forced_limit=None))
# bind the operations
pyt_grad_op = bind((qbx, target), grad_op)
pyt_op = bind((qbx, target), op)
# }}}
class MatrixFreeB(object):
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 mult(self, mat, x, y):
# Copy function data into the fivredrake function
self.x_fntn.dat.data[:] = x[:]
# Transfer the function to meshmode
self.meshmode_src_connection.from_firedrake(project(
self.x_fntn, dgfspace),
out=self.x_mm_fntn)
# Restrict to boundary
x_mm_fntn_on_bdy = src_bdy_connection(self.x_mm_fntn)
# Apply the operation
potential_int_mm = self.pyt_op(self.actx,
u=x_mm_fntn_on_bdy,
k=self.k)
grad_potential_int_mm = self.pyt_grad_op(self.actx,
u=x_mm_fntn_on_bdy,
k=self.k)
# Store in firedrake
self.potential_int.dat.data[target_indices] = potential_int_mm.get(
)
for dim in range(grad_potential_int_mm.shape[0]):
self.grad_potential_int.dat.data[
target_indices, dim] = grad_potential_int_mm[dim].get()
# Integrate the potential
r"""
Compute the inner products using firedrake. Note this
will be subtracted later, hence appears off by a sign.
.. math::
\langle
n(x) \cdot \nabla(
\int_\Gamma(
u(y) \partial_n H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
), v
\rangle_\Sigma
- \langle
i \kappa \cdot
\int_\Gamma(
u(y) \partial_n H_0^{(1)}(\kappa |x - y|)
)d\gamma(y), v
\rangle_\Sigma
"""
self.pyt_result = assemble(
inner(inner(self.grad_potential_int, self.n), self.v) *
ds(outer_bdy_id) -
inner(self.potential_int, self.v) * ds(outer_bdy_id))
# y <- Ax - evaluated potential
self.A.mult(x, y)
with self.pyt_result.dat.vec_ro as ep:
y.axpy(-1, ep)
# {{{ Compute normal helmholtz operator
u = TrialFunction(fspace)
v = TestFunction(fspace)
r"""
.. math::
\langle
\nabla u, \nabla v
\rangle
- \kappa^2 \cdot \langle
u, v
\rangle
- i \kappa \langle
u, v
\rangle_\Sigma
"""
a = inner(grad(u), grad(v)) * dx \
- Constant(wave_number**2) * inner(u, v) * dx \
- Constant(1j * wave_number) * inner(u, v) * ds(outer_bdy_id)
# get the concrete matrix from a general bilinear form
A = assemble(a).M.handle
# }}}
# {{{ Setup Python matrix
B = PETSc.Mat().create()
# build matrix context
Bctx = MatrixFreeB(A, pyt_grad_op, pyt_op, actx, wave_number)
# set up B as same size as A
B.setSizes(*A.getSizes())
B.setType(B.Type.PYTHON)
B.setPythonContext(Bctx)
B.setUp()
# }}}
# {{{ Create rhs
# Remember f is \partial_n(true_sol)|_\Gamma
# so we just need to compute \int_\Gamma\partial_n(true_sol) H(x-y)
sigma = sym.make_sym_vector("sigma", ambient_dim)
r"""
..math:
x \in \Sigma
grad_op(x) =
\nabla(
\int_\Gamma(
f(y) H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
)
"""
grad_op = pyt_inner_normal_sign * \
sym.grad(ambient_dim, sym.S(HelmholtzKernel(ambient_dim),
sym.n_dot(sigma),
k=sym.var("k"), qbx_forced_limit=None))
r"""
..math:
x \in \Sigma
op(x) =
i \kappa \cdot
\int_\Gamma(
f(y) H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
)
"""
op = 1j * sym.var("k") * pyt_inner_normal_sign * \
sym.S(HelmholtzKernel(ambient_dim),
sym.n_dot(sigma),
k=sym.var("k"),
qbx_forced_limit=None)
rhs_grad_op = bind((qbx, target), grad_op)
rhs_op = bind((qbx, target), op)
# Transfer to meshmode
metadata = {'quadrature_degree': 2 * fspace.ufl_element().degree()}
dg_true_sol_grad = project(true_sol_grad_expr,
dgvfspace,
form_compiler_parameters=metadata)
true_sol_grad_mm = meshmode_src_connection.from_firedrake(dg_true_sol_grad,
actx=actx)
true_sol_grad_mm = src_bdy_connection(true_sol_grad_mm)
# Apply the operations
f_grad_convoluted_mm = rhs_grad_op(actx,
sigma=true_sol_grad_mm,
k=wave_number)
f_convoluted_mm = rhs_op(actx, sigma=true_sol_grad_mm, k=wave_number)
# Transfer function back to firedrake
f_grad_convoluted = Function(vfspace)
f_convoluted = Function(fspace)
f_grad_convoluted.dat.data[:] = 0.0
f_convoluted.dat.data[:] = 0.0
for dim in range(f_grad_convoluted_mm.shape[0]):
f_grad_convoluted.dat.data[target_indices,
dim] = f_grad_convoluted_mm[dim].get()
f_convoluted.dat.data[target_indices] = f_convoluted_mm.get()
r"""
\langle
f, v
\rangle_\Gamma
+ \langle
i \kappa \cdot \int_\Gamma(
f(y) H_0^{(1)}(\kappa |x - y|)
)d\gamma(y), v
\rangle_\Sigma
- \langle
n(x) \cdot \nabla(
\int_\Gamma(
f(y) H_0^{(1)}(\kappa |x - y|)
)d\gamma(y)
), v
\rangle_\Sigma
"""
rhs_form = inner(inner(true_sol_grad_expr, FacetNormal(mesh)),
v) * ds(scatterer_bdy_id, metadata=metadata) \
+ inner(f_convoluted, v) * ds(outer_bdy_id) \
- inner(inner(f_grad_convoluted, FacetNormal(mesh)),
v) * ds(outer_bdy_id)
rhs = assemble(rhs_form)
# {{{ set up a solver:
solution = Function(fspace, name="Computed Solution")
# {{{ Used for preconditioning
if 'gamma' in solver_parameters or 'beta' in solver_parameters:
gamma = complex(solver_parameters.pop('gamma', 1.0))
import cmath
beta = complex(solver_parameters.pop('beta', cmath.sqrt(gamma)))
p = inner(grad(u), grad(v)) * dx \
- Constant(wave_number**2 * gamma) * inner(u, v) * dx \
- Constant(1j * wave_number * beta) * inner(u, v) * ds(outer_bdy_id)
P = assemble(p).M.handle
else:
P = A
# }}}
# Set up options to contain solver parameters:
ksp = PETSc.KSP().create()
if solver_parameters['pc_type'] == 'pyamg':
del solver_parameters['pc_type'] # We are using the AMG preconditioner
pyamg_tol = solver_parameters.get('pyamg_tol', None)
if pyamg_tol is not None:
pyamg_tol = float(pyamg_tol)
pyamg_maxiter = solver_parameters.get('pyamg_maxiter', None)
if pyamg_maxiter is not None:
pyamg_maxiter = int(pyamg_maxiter)
ksp.setOperators(B)
ksp.setUp()
pc = ksp.pc
pc.setType(pc.Type.PYTHON)
pc.setPythonContext(
AMGTransmissionPreconditioner(wave_number,
fspace,
A,
tol=pyamg_tol,
maxiter=pyamg_maxiter,
use_plane_waves=True))
# Otherwise use regular preconditioner
else:
ksp.setOperators(B, P)
options_manager = OptionsManager(solver_parameters, options_prefix)
options_manager.set_from_options(ksp)
import petsc4py.PETSc
petsc4py.PETSc.Sys.popErrorHandler()
with rhs.dat.vec_ro as b:
with solution.dat.vec as x:
ksp.solve(b, x)
# }}}
return ksp, solution
def heat(butcher_tableau):
N = 4
dt = Constant(1.0 / N)
t = Constant(0.0)
msh = UnitSquareMesh(N, N)
deg = 2
V = FunctionSpace(msh, "CG", deg)
x, y = SpatialCoordinate(msh)
uexact = t * (x + y)
rhs = expand_derivatives(diff(uexact, t)) - div(grad(uexact))
sols = []
luparams = {
"mat_type": "aij",
"snes_type": "ksponly",
"ksp_type": "preonly",
"pc_type": "lu"
}
ranaLD = {
"mat_type": "aij",
"snes_type": "ksponly",
"ksp_type": "gmres",
"ksp_monitor": None,
"pc_type": "python",
"pc_python_type": "irksome.RanaLD",
"aux": {
"pc_type": "fieldsplit",
"pc_fieldsplit_type": "multiplicative"
}
}
per_field = {"ksp_type": "preonly", "pc_type": "gamg"}
for s in range(butcher_tableau.num_stages):
ranaLD["fieldsplit_%s" % (s, )] = per_field
ranaDU = {
"mat_type": "aij",
"snes_type": "ksponly",
"ksp_type": "gmres",
"ksp_monitor": None,
"pc_type": "python",
"pc_python_type": "irksome.RanaDU",
"aux": {
"pc_type": "fieldsplit",
"pc_fieldsplit_type": "multiplicative"
}
}
for s in range(butcher_tableau.num_stages):
ranaLD["fieldsplit_%s" % (s, )] = per_field
params = [luparams, ranaLD, ranaDU]
for solver_parameters in params:
F, u, bc = Fubc(V, uexact, rhs)
stepper = TimeStepper(F,
butcher_tableau,
t,
dt,
u,
bcs=bc,
solver_parameters=solver_parameters)
stepper.advance()
sols.append(u)
errs = [errornorm(sols[0], uu) for uu in sols[1:]]
return numpy.max(errs)
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)
def setup_balance(dirname):
# set up grid and time stepping parameters
dt = 1.
tmax = 5.
deltax = 400
L = 2000.
H = 10000.
nlayers = int(H / deltax)
ncolumns = int(L / deltax)
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers)
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname + '/dry_balance',
dumpfreq=10,
dumplist=['u'])
parameters = CompressibleParameters()
diagnostics = Diagnostics(*fieldlist)
diagnostic_fields = []
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist,
diagnostic_fields=diagnostic_fields)
# Initial conditions
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
# Isentropic background state
Tsurf = Constant(300.)
theta0.interpolate(Tsurf)
# Calculate hydrostatic Pi
compressible_hydrostatic_balance(state, theta0, rho0, solve_for_rho=True)
state.initialise([('u', u0), ('rho', rho0), ('theta', theta0)])
state.set_reference_profiles([('rho', rho0), ('theta', theta0)])
# Set up advection schemes
ueqn = VectorInvariant(state, Vu)
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
thetaeqn = EmbeddedDGAdvection(state,
Vt,
equation_form="advective",
options=EmbeddedDGOptions())
advected_fields = [("u", ThetaMethod(state, u0, ueqn)),
("rho", SSPRK3(state, rho0, rhoeqn)),
("theta", SSPRK3(state, theta0, thetaeqn))]
# Set up linear solver
linear_solver = CompressibleSolver(state)
# Set up forcing
compressible_forcing = CompressibleForcing(state)
# build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver,
compressible_forcing)
return stepper, tmax
def streamplot(function,
resolution=None,
min_length=None,
max_time=None,
start_width=0.5,
end_width=1.5,
tolerance=3e-3,
loc_tolerance=1e-10,
seed=None,
complex_component="real",
**kwargs):
r"""Create a streamline plot of a vector field
Similar to matplotlib :func:`streamplot <matplotlib.pyplot.streamplot>`
:arg function: the Firedrake :class:`~.Function` to plot
:arg resolution: minimum spacing between streamlines (defaults to domain size / 20)
:arg min_length: minimum length of a streamline (defaults to 4x resolution)
:arg max_time: maximum time to integrate a streamline
:arg start_width: line width at beginning of streamline
:arg end_width: line width at end of streamline, to convey direction
:arg tolerance: dimensionless tolerance for adaptive ODE integration
:arg loc_tolerance: point location tolerance for :meth:`~firedrake.functions.Function.at`
:kwarg complex_component: If plotting complex data, which
component? (``'real'`` or ``'imag'``). Default is ``'real'``.
:kwarg kwargs: same as for matplotlib :class:`~matplotlib.collections.LineCollection`
"""
import randomgen
if function.ufl_shape != (2, ):
raise ValueError("Streamplot only defined for 2D vector fields!")
axes = kwargs.pop("axes", None)
if axes is None:
figure = plt.figure()
axes = figure.add_subplot(111)
mesh = function.ufl_domain()
if resolution is None:
coords = toreal(mesh.coordinates.dat.data_ro, "real")
resolution = (coords.max(axis=0) - coords.min(axis=0)).max() / 20
if min_length is None:
min_length = 4 * resolution
if max_time is None:
area = assemble(Constant(1) * dx(mesh))
average_speed = np.sqrt(
assemble(inner(function, function) * dx) / area)
max_time = 50 * min_length / average_speed
streamplotter = Streamplotter(function,
resolution,
min_length,
max_time,
tolerance,
loc_tolerance,
complex_component=complex_component)
# TODO: better way of seeding start points
shape = streamplotter._grid.shape
xmin = streamplotter._grid_point((0, 0))
xmax = streamplotter._grid_point((shape[0] - 2, shape[1] - 2))
X, Y = np.meshgrid(np.linspace(xmin[0], xmax[0], shape[0] - 2),
np.linspace(xmin[1], xmax[1], shape[1] - 2))
start_points = np.vstack((X.ravel(), Y.ravel())).T
# Randomly shuffle the start points
generator = randomgen.MT19937(seed).generator
for x in generator.permutation(np.array(start_points)):
streamplotter.add_streamline(x)
# Colors are determined by the speed, thicknesses by arc length
speeds = []
widths = []
for streamline in streamplotter.streamlines:
velocity = toreal(
np.array(function.at(streamline, tolerance=loc_tolerance)),
complex_component)
speed = np.sqrt(np.sum(velocity**2, axis=1))
speeds.extend(speed[:-1])
delta = np.sqrt(np.sum(np.diff(streamline, axis=0)**2, axis=1))
arc_length = np.cumsum(delta)
length = arc_length[-1]
s = arc_length / length
linewidth = (1 - s) * start_width + s * end_width
widths.extend(linewidth)
points = []
for streamline in streamplotter.streamlines:
pts = streamline.reshape(-1, 1, 2)
points.extend(np.hstack((pts[:-1], pts[1:])))
speeds = np.array(speeds)
widths = np.array(widths)
points = np.asarray(points)
vmin = kwargs.pop("vmin", speeds.min())
vmax = kwargs.pop("vmax", speeds.max())
norm = kwargs.pop("norm", matplotlib.colors.Normalize(vmin=vmin,
vmax=vmax))
cmap = plt.get_cmap(kwargs.pop("cmap", None))
collection = LineCollection(points, cmap=cmap, norm=norm, linewidth=widths)
collection.set_array(speeds)
axes.add_collection(collection)
_autoscale_view(axes, function.ufl_domain().coordinates.dat.data_ro)
return collection
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)]
else:
# Fixed quadrature with 64 points gives absolute errors below 1e-13
# for a quantity of order 1e-3.
v, _ = integrate.fixed_quad(D_integrand, theta_0, angles[ii], n=64)
val[ii] = -(R / g) * v
return val
# Get coordinates to pass to Dval function
D0 = Function(W0)
W = VectorFunctionSpace(mesh, D0.ufl_element())
X = interpolate(mesh.coordinates, W)
D0.dat.data[:] = Dval(X.dat.data_ro)
# Adjust mean value of initial D
C = Function(D0.function_space()).assign(Constant(1.0))
area = assemble(C * dx)
Dmean = assemble(D0 * dx) / area
D0 -= Dmean
D0 += Constant(D_mean)
# Dpert
D_p = Function(W0)
Dpert = D_bump * cos(theta) * exp(-(lamda / a)**2) * exp(-(
(theta_2 - theta) / b)**2)
D_p.interpolate(Dpert)
# Dexpr = Dbar + Dpert
Dexpr = D0 + D_p
bexpr = Constant(0)
# Set up functions
def build_initial_conditions(prognostic_variables, simulation_parameters):
"""
Initialises the prognostic variables based on the
initial condition string.
:arg prognostic_variables: a PrognosticVariables object.
:arg simulation_parameters: a dictionary containing the simulation parameters.
"""
mesh = simulation_parameters['mesh'][-1]
ic = simulation_parameters['ic'][-1]
alphasq = simulation_parameters['alphasq'][-1]
c0 = simulation_parameters['c0'][-1]
gamma = simulation_parameters['gamma'][-1]
x, = SpatialCoordinate(mesh)
Ld = simulation_parameters['Ld'][-1]
deltax = Ld / simulation_parameters['resolution'][-1]
w = simulation_parameters['peak_width'][-1]
epsilon = 1
ic_dict = {
'two_peaks':
(0.2 * 2 /
(exp(x - 403. / 15. * 40. / Ld) + exp(-x + 403. / 15. * 40. / Ld)) +
0.5 * 2 /
(exp(x - 203. / 15. * 40. / Ld) + exp(-x + 203. / 15. * 40. / Ld))),
'gaussian':
0.5 * exp(-((x - 10.) / 2.)**2),
'gaussian_narrow':
0.5 * exp(-((x - 10.) / 1.)**2),
'gaussian_wide':
0.5 * exp(-((x - 10.) / 3.)**2),
'peakon':
conditional(x < Ld / 2., exp((x - Ld / 2) / sqrt(alphasq)),
exp(-(x - Ld / 2) / sqrt(alphasq))),
'one_peak':
0.5 * 2 /
(exp(x - 203. / 15. * 40. / Ld) + exp(-x + 203. / 15. * 40. / Ld)),
'proper_peak':
0.5 * 2 / (exp(x - Ld / 4) + exp(-x + Ld / 4)),
'new_peak':
0.5 * 2 / (exp((x - Ld / 4) / w) + exp((-x + Ld / 4) / w)),
'flat':
Constant(2 * pi**2 / (9 * 40**2)),
'fast_flat':
Constant(0.1),
'coshes':
Constant(2000) * cosh((2000**0.5 / 2) * (x - 0.75))**(-2) +
Constant(1000) * cosh(1000**0.5 / 2 * (x - 0.25))**(-2),
'd_peakon':
exp(-sqrt((x - Ld / 2)**2 + epsilon * deltax**2) / sqrt(alphasq)),
'zero':
Constant(0.0),
'two_peakons':
conditional(
x < Ld / 4,
exp((x - Ld / 4) / sqrt(alphasq)) -
exp(-(x + Ld / 4) / sqrt(alphasq)),
conditional(
x < 3 * Ld / 4,
exp(-(x - Ld / 4) / sqrt(alphasq)) - exp(
(x - 3 * Ld / 4) / sqrt(alphasq)),
exp((x - 5 * Ld / 4) / sqrt(alphasq)) -
exp(-(x - 3 * Ld / 4) / sqrt(alphasq)))),
'twin_peakons':
conditional(
x < Ld / 4,
exp((x - Ld / 4) / sqrt(alphasq)) + 0.5 * exp(
(x - Ld / 2) / sqrt(alphasq)),
conditional(
x < Ld / 2,
exp(-(x - Ld / 4) / sqrt(alphasq)) + 0.5 * exp(
(x - Ld / 2) / sqrt(alphasq)),
conditional(
x < 3 * Ld / 4,
exp(-(x - Ld / 4) / sqrt(alphasq)) +
0.5 * exp(-(x - Ld / 2) / sqrt(alphasq)),
exp((x - 5 * Ld / 4) / sqrt(alphasq)) +
0.5 * exp(-(x - Ld / 2) / sqrt(alphasq))))),
'periodic_peakon': (conditional(
x < Ld / 2, 0.5 / (1 - exp(-Ld / sqrt(alphasq))) *
(exp((x - Ld / 2) / sqrt(alphasq)) +
exp(-Ld / sqrt(alphasq)) * exp(-(x - Ld / 2) / sqrt(alphasq))),
0.5 / (1 - exp(-Ld / sqrt(alphasq))) *
(exp(-(x - Ld / 2) / sqrt(alphasq)) +
exp(-Ld / sqrt(alphasq)) * exp((x - Ld / 2) / sqrt(alphasq))))),
'cos_bell':
conditional(x < Ld / 4, (cos(pi * (x - Ld / 8) / (2 * Ld / 8)))**2,
0.0),
'antisymmetric':
1 / (exp((x - Ld / 4) / Ld) + exp((-x + Ld / 4) / Ld)) - 1 / (exp(
(Ld - x - Ld / 4) / Ld) + exp((Ld + x + Ld / 4) / Ld))
}
ic_expr = ic_dict[ic]
if prognostic_variables.scheme in ['upwind', 'LASCH']:
VCG5 = FunctionSpace(mesh, "CG", 5)
smooth_condition = Function(VCG5).interpolate(ic_expr)
prognostic_variables.u.project(as_vector([smooth_condition]))
# need to find initial m by solving helmholtz problem
CG1 = FunctionSpace(mesh, "CG", 1)
u0 = prognostic_variables.u
p = TestFunction(CG1)
m_CG = Function(CG1)
ones = Function(prognostic_variables.Vu).project(
as_vector([Constant(1.)]))
Lm = (p * m_CG - p * dot(ones, u0) -
alphasq * p.dx(0) * dot(ones, u0.dx(0))) * dx
mprob0 = NonlinearVariationalProblem(Lm, m_CG)
msolver0 = NonlinearVariationalSolver(mprob0,
solver_parameters={
'ksp_type': 'preonly',
'pc_type': 'lu'
})
msolver0.solve()
prognostic_variables.m.interpolate(m_CG)
if prognostic_variables.scheme == 'LASCH':
prognostic_variables.Eu.assign(prognostic_variables.u)
prognostic_variables.Em.assign(prognostic_variables.m)
elif prognostic_variables.scheme in ('conforming', 'hydrodynamic', 'test',
'LASCH_hydrodynamic',
'LASCH_hydrodynamic_m',
'no_gradient'):
if ic == 'peakon':
Vu = prognostic_variables.Vu
# delta = Function(Vu)
# middle_index = int(len(delta.dat.data[:]) / 2)
# delta.dat.data[middle_index] = 1
# u0 = prognostic_variables.u
# phi = TestFunction(Vu)
#
# eqn = phi * u0 * dx + alphasq * phi.dx(0) * u0.dx(0) * dx - phi * delta * dx
# prob = NonlinearVariationalProblem(eqn, u0)
# solver = NonlinearVariationalSolver(prob)
# solver.solve()
# W = MixedFunctionSpace((Vu, Vu))
# psi, phi = TestFunctions(W)
# w = Function(W)
# u, F = w.split()
# u.interpolate(ic_expr)
# u, F = split(w)
#
# eqn = (psi * u * dx - psi * (0.5 * u * u + F) * dx
# + phi * F * dx + alphasq * phi.dx(0) * F.dx(0) * dx
# - phi * u * u * dx - 0.5 * alphasq * phi * u.dx(0) * u.dx(0) * dx)
#
# u, F = w.split()
#
# prob = NonlinearVariationalProblem(eqn, w)
# solver = NonlinearVariationalSolver(prob)
# solver.solve()
# prognostic_variables.u.assign(u)
prognostic_variables.u.project(ic_expr)
# prognostic_variables.u.interpolate(ic_expr)
else:
VCG5 = FunctionSpace(mesh, "CG", 5)
smooth_condition = Function(VCG5).interpolate(ic_expr)
prognostic_variables.u.project(smooth_condition)
if prognostic_variables.scheme in [
'LASCH_hydrodynamic', 'LASCH_hydrodynamic_m'
]:
prognostic_variables.Eu.assign(prognostic_variables.u)
else:
raise NotImplementedError('Other schemes not yet implemented.')
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 ufl.algorithms.map_integrands import map_integrand_dags
from firedrake import (FunctionSpace, TrialFunction, TrialFunctions,
TestFunction, Function, BrokenElement,
MixedElement, FacetNormal, Constant,
DirichletBC, Projector)
from firedrake.assemble import (allocate_matrix,
create_assembly_callable)
from firedrake.formmanipulation import ArgumentReplacer, split_form
# Extract the problem context
prefix = pc.getOptionsPrefix()
_, P = pc.getOperators()
context = P.getPythonContext()
test, trial = context.a.arguments()
V = test.function_space()
if V.mesh().cell_set._extruded:
# TODO: Merge FIAT branch to support TPC trace elements
raise NotImplementedError("Not implemented on extruded meshes.")
# Break the function spaces and define fully discontinuous spaces
broken_elements = [BrokenElement(Vi.ufl_element()) for Vi in V]
elem = MixedElement(broken_elements)
V_d = FunctionSpace(V.mesh(), elem)
arg_map = {test: TestFunction(V_d), trial: TrialFunction(V_d)}
# Replace the problems arguments with arguments defined
# on the new discontinuous spaces
replacer = ArgumentReplacer(arg_map)
new_form = map_integrand_dags(replacer, context.a)
# Create the space of approximate traces.
# The vector function space will have a non-empty value_shape
W = next(v for v in V if bool(v.ufl_element().value_shape()))
if W.ufl_element().family() in ["Raviart-Thomas", "RTCF"]:
tdegree = W.ufl_element().degree() - 1
else:
tdegree = W.ufl_element().degree()
# NOTE: Once extruded is ready, we will need to be aware of this
# and construct the appropriate trace space for the HDiv element
TraceSpace = FunctionSpace(V.mesh(), "HDiv Trace", tdegree)
# NOTE: For extruded, we will need to add "on_top" and "on_bottom"
trace_conditions = [
DirichletBC(TraceSpace, Constant(0.0), "on_boundary")
]
# Set up the functions for the original, hybridized
# and schur complement systems
self.broken_solution = Function(V_d)
self.broken_rhs = Function(V_d)
self.trace_solution = Function(TraceSpace)
self.unbroken_solution = Function(V)
self.unbroken_rhs = Function(V)
# Create the symbolic Schur-reduction
Atilde = Tensor(new_form)
gammar = TestFunction(TraceSpace)
n = FacetNormal(V.mesh())
# Vector trial function will have a non-empty ufl_shape
sigma = next(f for f in TrialFunctions(V_d) if bool(f.ufl_shape))
# NOTE: Once extruded is ready, this will change slightly
# to include both horizontal and vertical interior facets
K = Tensor(gammar('+') * ufl.dot(sigma, n) * ufl.dS)
# 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_rhs,
tensor=self.schur_rhs,
form_compiler_parameters=context.fc_params)
schur_comp = K * Atilde.inv * K.T
self.S = allocate_matrix(schur_comp,
bcs=trace_conditions,
form_compiler_parameters=context.fc_params)
self._assemble_S = create_assembly_callable(
schur_comp,
tensor=self.S,
bcs=trace_conditions,
form_compiler_parameters=context.fc_params)
self._assemble_S()
self.S.force_evaluation()
Smat = self.S.petscmat
# Nullspace for the multiplier problem
nullsp = P.getNullSpace()
if nullsp.handle != 0:
new_vecs = get_trace_nullspace_vecs(K * Atilde.inv, nullsp, V, V_d,
TraceSpace)
tr_nullsp = PETSc.NullSpace().create(vectors=new_vecs,
comm=pc.comm)
Smat.setNullSpace(tr_nullsp)
# Set up the KSP for the system of Lagrange multipliers
ksp = PETSc.KSP().create(comm=pc.comm)
ksp.setOptionsPrefix(prefix + "trace_")
ksp.setTolerances(rtol=1e-13)
ksp.setOperators(Smat)
ksp.setUp()
ksp.setFromOptions()
self.ksp = ksp
# Now we construct the local tensors for the reconstruction stage
# TODO: Add support for mixed tensors and these variables
# become unnecessary
split_forms = split_form(new_form)
A = Tensor(next(sf.form for sf in split_forms if sf.indices == (0, 0)))
B = Tensor(next(sf.form for sf in split_forms if sf.indices == (1, 0)))
C = Tensor(next(sf.form for sf in split_forms if sf.indices == (1, 1)))
trial = TrialFunction(
FunctionSpace(V.mesh(), BrokenElement(W.ufl_element())))
K_local = Tensor(gammar('+') * ufl.dot(trial, n) * ufl.dS)
# Split functions and reconstruct each bit separately
sigma_h, u_h = self.broken_solution.split()
g, f = self.broken_rhs.split()
# Pressure reconstruction
M = B * A.inv * B.T + C
u_sol = M.inv * f + M.inv * (
B * A.inv * K_local.T * self.trace_solution - B * A.inv * g)
self._assemble_pressure = create_assembly_callable(
u_sol, tensor=u_h, form_compiler_parameters=context.fc_params)
# Velocity reconstruction
sigma_sol = A.inv * g + A.inv * (B.T * u_h -
K_local.T * self.trace_solution)
self._assemble_velocity = create_assembly_callable(
sigma_sol,
tensor=sigma_h,
form_compiler_parameters=context.fc_params)
# Set up the projector for projecting the broken solution
# into the unbroken finite element spaces
# NOTE: Tolerance here matters!
sigma_b, _ = self.broken_solution.split()
sigma_u, _ = self.unbroken_solution.split()
self.projector = Projector(sigma_b,
sigma_u,
solver_parameters={
"ksp_type": "cg",
"ksp_rtol": 1e-13
})
interpolate,
max_value,
assemble,
norm,
Constant,
UnitDiskMesh,
)
import icepack
from icepack.constants import ice_density as ρ_I, glen_flow_law as n, gravity as g
import numpy as np
# Test our numerical solvers against this analytical solution.
R = 100e3 # Radius of ice sheet in meters
R_mesh = R * 0.75 # Radius of mesh (set multiplier to <1 to
# ignore interpolation errors at the edge)
alpha = Constant(R)
T = Constant(254.15)
A = icepack.rate_factor(T)
A0 = 2 * A * (ρ_I * g)**n / (n + 2)
def make_mesh(R_mesh, refinement):
mesh = UnitDiskMesh(refinement)
mesh.coordinates.dat.data[:] *= R_mesh
return mesh
# Pick a geometry in order to have an exact solution. We'll use the
# Bueler profile from section 5.6.3 in Greve and Blatter (2009).
# Citation: Greve, Ralf, and Heinz Blatter. Dynamics of ice sheets
# and glaciers. Springer Science & Business Media, 2009.
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')
family="BDM",
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist)
# interpolate initial conditions
u0 = state.fields("u")
D0 = state.fields("D")
u_max = 2 * pi * R / (12 * day
) # Maximum amplitude of the zonal wind (m/s)
uexpr = as_vector([-u_max * x[1] / R, u_max * x[0] / R, 0.0])
Omega = parameters.Omega
g = parameters.g
Dexpr = Constant(1.0)
bexpr = H - ((R * Omega * u_max + u_max * u_max / 2.0) * (x[2] * x[2] /
(R * R))) / g
# Coriolis expression
fexpr = 2 * Omega * x[2] / R
V = FunctionSpace(mesh, "CG", 1)
f = state.fields("coriolis", V)
f.interpolate(fexpr) # Coriolis frequency (1/s)
b = state.fields("topography", D0.function_space())
b.interpolate(bexpr)
u0.project(uexpr)
D0.interpolate(Dexpr)
state.initialise([('u', u0), ('D', D0)])
euler_poincare = False
if recovered:
VDG1 = FunctionSpace(mesh, "DG", 1)
VCG1 = FunctionSpace(mesh, "CG", 1)
Vt_brok = FunctionSpace(mesh, BrokenElement(Vt.ufl_element()))
Vu_DG1 = VectorFunctionSpace(mesh, "DG", 1)
Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1)
Vu_brok = FunctionSpace(mesh, BrokenElement(Vu.ufl_element()))
u_spaces = (Vu_DG1, Vu_CG1, Vu_brok)
rho_spaces = (VDG1, VCG1, Vr)
theta_spaces = (VDG1, VCG1, Vt_brok)
# Define constant theta_e and water_t
Tsurf = 300.0
theta_b = Function(Vt).interpolate(Constant(Tsurf))
# Calculate hydrostatic fields
compressible_hydrostatic_balance(state, theta_b, rho0, solve_for_rho=True)
# make mean fields
rho_b = Function(Vr).assign(rho0)
# define perturbation
xc = L / 2
zc = 2000.
rc = 2000.
Tdash = 2.0
theta_pert = Function(Vt).interpolate(conditional(sqrt((x[0] - xc) ** 2 + (x[1] - zc) ** 2) > rc,
0.0, Tdash *
(cos(pi * sqrt(((x[0] - xc) / rc) ** 2 + ((x[1] - zc) / rc) ** 2) / 2.0))
def setup_condens(dirname):
# declare grid shape, with length L and height H
L = 1000.
H = 1000.
nlayers = int(H / 100.)
ncolumns = int(L / 100.)
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
x = SpatialCoordinate(mesh)
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=1.0, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname + "/condens",
dumpfreq=1,
dumplist=['u'],
perturbation_fields=['theta', 'rho'])
parameters = CompressibleParameters()
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist,
diagnostic_fields=[Sum('water_v', 'water_c')])
# declare initial fields
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vpsi = FunctionSpace(mesh, "CG", 2)
Vt = theta0.function_space()
Vr = rho0.function_space()
# make a gradperp
gradperp = lambda u: as_vector([-u.dx(1), u.dx(0)])
# declare tracer field and a background field
water_v0 = state.fields("water_v", Vt)
water_c0 = state.fields("water_c", Vt)
# Isentropic background state
Tsurf = Constant(300.)
theta_b = Function(Vt).interpolate(Tsurf)
rho_b = Function(Vr)
# Calculate initial rho
compressible_hydrostatic_balance(state, theta_b, rho_b, solve_for_rho=True)
# set up water_v
xc = 500.
zc = 350.
rc = 250.
r = sqrt((x[0] - xc)**2 + (x[1] - zc)**2)
w_expr = conditional(r > rc, 0., 0.25 * (1. + cos((pi / rc) * r)))
# set up velocity field
u_max = 10.0
psi_expr = ((-u_max * L / pi) * sin(2 * pi * x[0] / L) *
sin(pi * x[1] / L))
psi0 = Function(Vpsi).interpolate(psi_expr)
u0.project(gradperp(psi0))
theta0.interpolate(theta_b)
rho0.interpolate(rho_b)
water_v0.interpolate(w_expr)
state.initialise([('u', u0), ('rho', rho0), ('theta', theta0),
('water_v', water_v0), ('water_c', water_c0)])
state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
# set up advection schemes
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
thetaeqn = SUPGAdvection(state,
Vt,
supg_params={"dg_direction": "horizontal"},
equation_form="advective")
# build advection dictionary
advected_fields = []
advected_fields.append(("u", NoAdvection(state, u0, None)))
advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn)))
advected_fields.append(("theta", SSPRK3(state, theta0, thetaeqn)))
advected_fields.append(("water_v", SSPRK3(state, water_v0, thetaeqn)))
advected_fields.append(("water_c", SSPRK3(state, water_c0, thetaeqn)))
physics_list = [Condensation(state)]
# build time stepper
stepper = AdvectionDiffusion(state,
advected_fields,
physics_list=physics_list)
return stepper, 5.0
def setup_tracer(dirname, hybridization):
# declare grid shape, with length L and height H
L = 1000.
H = 1000.
nlayers = int(H / 100.)
ncolumns = int(L / 100.)
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=10.0, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname+"/tracer",
dumpfreq=1,
dumplist=['u'],
perturbation_fields=['theta', 'rho'])
parameters = CompressibleParameters()
state = State(mesh, vertical_degree=1, horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist,
diagnostic_fields=[Difference('theta', 'tracer')])
# declare initial fields
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
# declare tracer field and a background field
tracer0 = state.fields("tracer", Vt)
# Isentropic background state
Tsurf = Constant(300.)
theta_b = Function(Vt).interpolate(Tsurf)
rho_b = Function(Vr)
# Calculate initial rho
compressible_hydrostatic_balance(state, theta_b, rho_b,
solve_for_rho=True)
# set up perturbation to theta
xc = 500.
zc = 350.
rc = 250.
x = SpatialCoordinate(mesh)
r = sqrt((x[0]-xc)**2 + (x[1]-zc)**2)
theta_pert = conditional(r > rc, 0., 0.25*(1. + cos((pi/rc)*r)))
theta0.interpolate(theta_b + theta_pert)
rho0.interpolate(rho_b)
tracer0.interpolate(theta0)
state.initialise([('u', u0),
('rho', rho0),
('theta', theta0),
('tracer', tracer0)])
state.set_reference_profiles([('rho', rho_b),
('theta', theta_b)])
# set up advection schemes
ueqn = EulerPoincare(state, Vu)
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
thetaeqn = SUPGAdvection(state, Vt,
supg_params={"dg_direction": "horizontal"},
equation_form="advective")
# build advection dictionary
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(("tracer", SSPRK3(state, tracer0, thetaeqn)))
# Set up linear solver
if hybridization:
linear_solver = HybridizedCompressibleSolver(state)
else:
linear_solver = CompressibleSolver(state)
compressible_forcing = CompressibleForcing(state)
# build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver,
compressible_forcing)
return stepper, 100.0
#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
k = Constant(dt)
f = Constant((0, 0))
# Tentative velocity step
F1 = (1.0/k)*inner(u - u0, v)*dx + inner(grad(u0)*u0, v)*dx + \
(1.0/Re)*inner(grad(u), grad(v))*dx - inner(f, v)*dx
a1 = lhs(F1)
L1 = rhs(F1)
# Pressure update
a2 = inner(grad(p), grad(q)) * dx
L2 = -(1 / k) * div(u1) * q * dx
# Velocity update
a3 = inner(u, v) * dx
def compute(self, state):
u = state.fields("u")
dt = Constant(state.timestepping.dt)
return self.field.project(sqrt(dot(u, u))/sqrt(self.area)*dt)
def __init__(self,
salinity,
temperature,
pressure_perturbation,
z,
GammaTfunc=None,
velocity=None,
ice_heat_flux=True,
HJ99Gamma=False,
f=None):
super().__init__(salinity, temperature, pressure_perturbation, z)
if velocity is None:
# Use constant turbulent thermal and salinity exchange values given in Holland and Jenkins 1999.
gammaT = self.gammaT
gammaS = self.gammaS
else:
u = velocity
if HJ99Gamma:
# Holland and Jenkins 1999 turbulent heat/salt exchange velocity as a function
# of friction velocity. This should be the same as in MITgcm.
# Holland, D.M. and Jenkins, A., 1999. Modeling thermodynamic ice–ocean interactions at
# the base of an ice shelf. Journal of Physical Oceanography, 29(8), pp.1787-1800.
# Calculate friction velocity
if isinstance(u, float):
print("Input velocity:", u)
u_bounded = max(u, 1e-3)
print("Bounded velocity:", u_bounded)
else:
u_bounded = conditional(u > 1e-3, u, 1e-3)
u_star = pow(self.C_d * pow(u_bounded, 2), 0.5)
# Calculate turbulent component of thermal and salinity exchange velocity (Eq 15)
# N.b MITgcm sets etastar = 1 so that the first part of Gamma_turb term below is constant.
# eta_star = (1 + (zeta_N * u_star) / (f * L0 * Rc))^-1/2 (Eq 18)
# In H&J99 eta_star is set to 1 when the Obukhov length is negative (i.e the buoyancy flux
# is destabilising. This occurs during freezing. (Melting is stabilising)
# So it does seem a bit odd because then eta_star is tuned to freezing conditions
# need to work this out...?
Gamma_Turb = 1.0 / (2.0 * self.zeta_N *
self.eta_star) - 1.0 / self.k
if f is not None:
# Add extra term if using coriolis term
# Calculate viscous sublayer thickness (Eq 17)
h_nu = 5.0 * self.nu / u_star
Gamma_Turb += ln(
u_star * self.zeta_N * pow(self.eta_star, 2) /
(abs(f) * h_nu)) / self.k
# Calculate molecular components of thermal exchange velocity (Eq 16)
GammaT_Mole = 12.5 * pow(self.Pr, 2.0 / 3.0) - 6.0
# Calculate molecular component of salinity exchange velocity (Eq 16)
GammaS_Mole = 12.5 * pow(self.Sc, 2.0 / 3.0) - 6.0
# Calculate thermal and salinity exchange velocity. (Eq 14)
# Do we need to catch -ve gamma? could have -ve Gamma_Turb?
gammaT = u_star / (Gamma_Turb + GammaT_Mole)
gammaS = u_star / (Gamma_Turb + GammaS_Mole)
# print exchange velocities if testing when input velocity is a float.
if isinstance(gammaT, float) or isinstance(gammaS, float):
print("gammaT = ", gammaT)
print("gammaS = ", gammaS)
else:
# ISOMIP+ based on Jenkins et al 2010. Measurement of basal rates beneath Ronne Ice Shelf
u_tidal = 0.01
u_star = pow(self.C_d * (pow(u, 2) + pow(u_tidal, 2)), 0.5)
if GammaTfunc is None:
gammaT = self.GammaT * u_star
gammaS = self.GammaS * u_star
else:
gammaT = GammaTfunc * u_star
gammaS = (GammaTfunc / 35.0) * u_star
# print exchange velocities if testing when input velocity is a float.
if isinstance(gammaT, float) or isinstance(gammaS, float):
print("gammaT = ", gammaT)
print("gammaS = ", gammaS)
b_plus_cPb = (self.b + self.c * self.P_full
) # save calculating this each time...
# Calculate coefficients in quadratic equation for salinity at ice-ocean boundary.
# Aa.Sb^2 + Bb.Sb + Cc = 0
Aa = self.c_p_m * gammaT * self.a
Bb = -gammaS * self.Lf
Bb -= self.c_p_m * gammaT * self.T
Bb += self.c_p_m * gammaT * b_plus_cPb
Cc = gammaS * self.S * self.Lf
if ice_heat_flux:
Aa -= gammaS * self.c_p_i * self.a
Bb += gammaS * self.S * self.c_p_i * self.a
Bb -= gammaS * self.c_p_i * b_plus_cPb
Bb += gammaS * self.c_p_i * self.T_ice
Cc += gammaS * self.S * self.c_p_i * b_plus_cPb
Cc -= gammaS * self.S * self.c_p_i * self.T_ice
S1 = (-Bb + pow(Bb**2 - 4.0 * Aa * Cc, 0.5)) / (2.0 * Aa)
S2 = (-Bb - pow(Bb**2 - 4.0 * Aa * Cc, 0.5)) / (2.0 * Aa)
print(type(S1))
print(S1)
if isinstance(S1, (float, sympy.core.numbers.Float)):
# Print statements for testing
print("S1 = ", S1)
print("S2 = ", S2)
if S1 > 0:
self.Sb = S1
print("Choose S1")
else:
self.Sb = S2
print("Choose S2")
elif isinstance(S1, sympy.core.add.Add):
self.Sb = S2
else:
self.Sb = conditional(S1 > 0.0, S1, S2)
self.Tb = self.a * self.Sb + self.b + self.c * self.P_full
self.wb = gammaS * (self.S - self.Sb) / self.Sb
if ice_heat_flux:
self.Q_ice = -self.rho0 * (self.T_ice -
self.Tb) * self.c_p_i * self.wb
else:
if isinstance(S1, float):
self.Q_ice = 0.0
else:
self.Q_ice = Constant(0.0)
self.Q_mixed = -self.rho0 * self.c_p_m * gammaT * (self.Tb - self.T)
self.Q_latent = self.Q_ice - self.Q_mixed
self.QS_mixed = -self.rho0 * gammaS * (self.Sb - self.S)
self.T_flux_bc = -(self.wb + gammaT) * (self.Tb - self.T)
self.S_flux_bc = -(self.wb + gammaS) * (self.Sb - self.S)
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)}')
fieldlist=fieldlist,
diagnostic_fields=diagnostic_fields)
# Initial conditions
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
water0 = state.fields("water", theta0.function_space())
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
# Isentropic background state
Tsurf = Constant(300.)
theta_b = Function(Vt).interpolate(Tsurf)
rho_b = Function(Vr)
# Calculate hydrostatic Pi
compressible_hydrostatic_balance(state, theta_b, rho_b, solve_for_rho=True)
x = SpatialCoordinate(mesh)
a = 5.0e3
deltaTheta = 1.0e-2
xc = 0.5 * L
xr = 4000.
zc = 3000.
zr = 2000.
r = sqrt(((x[0] - xc) / xr)**2 + ((x[1] - zc) / zr)**2)
no_normal_flow_bc_ids=[1, 2])
# Initial conditions
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vu = state.spaces("HDiv")
Vt = state.spaces("theta")
Vr = state.spaces("DG")
x, z = SpatialCoordinate(mesh)
# Define constant theta_e and water_t
Tsurf = 300.0
theta_b = Function(Vt).interpolate(Constant(Tsurf))
# Calculate hydrostatic fields
compressible_hydrostatic_balance(state, theta_b, rho0, solve_for_rho=True)
# make mean fields
rho_b = Function(Vr).assign(rho0)
# define perturbation
xc = L / 2
zc = 2000.
rc = 2000.
Tdash = 2.0
r = sqrt((x - xc)**2 + (z - zc)**2)
theta_pert = Function(Vt).interpolate(
conditional(r > rc, 0.0,
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, density_field, factor=1.):
super().__init__(required_fields=(density_field, "u_gradient"))
self.density_field = density_field
self.factor = Constant(factor)
