def heat(n, deg, time_stages, stage_type="deriv", splitting=IA):
N = 2**n
msh = UnitIntervalMesh(N)
params = {
"snes_type": "ksponly",
"ksp_type": "preonly",
"mat_type": "aij",
"pc_type": "lu"
}
V = FunctionSpace(msh, "CG", deg)
x, = SpatialCoordinate(msh)
t = Constant(0.0)
dt = Constant(2.0 / N)
uexact = exp(-t) * sin(pi * x)
rhs = expand_derivatives(diff(uexact, t)) - div(grad(uexact))
butcher_tableau = GaussLegendre(time_stages)
u = project(uexact, V)
v = TestFunction(V)
F = (inner(Dt(u), v) * dx + inner(grad(u), grad(v)) * dx -
inner(rhs, v) * dx)
bc = DirichletBC(V, Constant(0), "on_boundary")
stepper = TimeStepper(F,
butcher_tableau,
t,
dt,
u,
bcs=bc,
solver_parameters=params,
stage_type=stage_type,
splitting=splitting)
while (float(t) < 1.0):
if (float(t) + float(dt) > 1.0):
dt.assign(1.0 - float(t))
stepper.advance()
t.assign(float(t) + float(dt))
return errornorm(uexact, u) / norm(uexact)
def wave(n, deg, butcher_tableau, splitting=AI):
N = 2**n
msh = UnitIntervalMesh(N)
params = {
"snes_type": "ksponly",
"ksp_type": "preonly",
"mat_type": "aij",
"pc_type": "lu"
}
V = FunctionSpace(msh, "CG", deg)
W = FunctionSpace(msh, "DG", deg - 1)
Z = V * W
x, = SpatialCoordinate(msh)
t = Constant(0.0)
dt = Constant(2.0 / N)
up = project(as_vector([0, sin(pi * x)]), Z)
u, p = split(up)
v, w = TestFunctions(Z)
F = (inner(Dt(u), v) * dx + inner(u.dx(0), w) * dx + inner(Dt(p), w) * dx -
inner(p, v.dx(0)) * dx)
E = 0.5 * (inner(u, u) * dx + inner(p, p) * dx)
stepper = TimeStepper(F,
butcher_tableau,
t,
dt,
up,
solver_parameters=params,
splitting=splitting)
energies = []
while (float(t) < 1.0):
if (float(t) + float(dt) > 1.0):
dt.assign(1.0 - float(t))
stepper.advance()
t.assign(float(t) + float(dt))
energies.append(assemble(E))
return np.array(energies)
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()
bcs = None if len(self.state.bcs) == 0 else self.state.bcs
u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs)
solver_parameters = {}
if logger.isEnabledFor(DEBUG):
solver_parameters["ksp_monitor_true_residual"] = None
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
class Diffusion2D(Application):
"""
Application class containing the description of the diffusion problem.
The spatial domain is a 10x10 square with
periodic boundary conditions in each direction.
The initial condition is a Gaussian in the centre of the domain.
The spatial discretisation is P1 DG (piecewise linear discontinous
elements) and uses an interior penalty method which penalises jumps
at element interfaces.
"""
def __init__(self,
mesh: object,
kappa: float,
comm_space: MPI.Comm,
mu: float = 5.,
*args,
**kwargs):
"""
Constructor
:param mesh: spatial domain
:param kappa: diffusion coefficient
:param mu: penalty weighting function
"""
super(Diffusion2D, self).__init__(*args, **kwargs)
# Spatial domain and function space
self.mesh = mesh
V = FunctionSpace(self.mesh, "DG", 1)
self.function_space = V
self.comm_space = comm_space
# Placeholder for time step - will be updated in the update method
self.dt = Constant(0.)
# Things we need for the form
gamma = TestFunction(V)
phi = TrialFunction(V)
self.f = Function(V)
n = FacetNormal(mesh)
# Set up the rhs and bilinear form of the equation
a = (inner(gamma, phi) * dx + self.dt *
(inner(grad(gamma),
grad(phi) * kappa) * dx -
inner(2 * avg(outer(phi, n)), avg(grad(gamma) * kappa)) * dS -
inner(avg(grad(phi) * kappa), 2 * avg(outer(gamma, n))) * dS +
mu * inner(2 * avg(outer(phi, n)),
2 * avg(outer(gamma, n) * kappa)) * dS))
rhs = inner(gamma, self.f) * dx
# Function to hold the solution
self.soln = Function(V)
# Setup problem and solver
prob = LinearVariationalProblem(a, rhs, self.soln)
self.solver = NonlinearVariationalSolver(prob)
# Set the data structure for any user-defined time point
self.vector_template = VectorDiffusion2D(size=len(self.function_space),
comm_space=self.comm_space)
# Set initial condition:
# Setting up a Gaussian blob in the centre of the domain.
self.vector_t_start = VectorDiffusion2D(size=len(self.function_space),
comm_space=self.comm_space)
x = SpatialCoordinate(self.mesh)
initial_tracer = exp(-((x[0] - 5)**2 + (x[1] - 5)**2))
tmp = Function(self.function_space)
tmp.interpolate(initial_tracer)
self.vector_t_start.set_values(np.copy(tmp.dat.data))
def step(self, u_start: VectorDiffusion2D, t_start: float,
t_stop: float) -> VectorDiffusion2D:
"""
Time integration routine for 2D diffusion problem:
Backward Euler
:param u_start: approximate solution for the input time t_start
:param t_start: time associated with the input approximate solution u_start
:param t_stop: time to evolve the input approximate solution to
:return: approximate solution at input time t_stop
"""
# Time-step size
self.dt.assign(t_stop - t_start)
# Get data from VectorDiffusion2D object u_start
# and copy to Firedrake Function object tmp
tmp = Function(self.function_space)
for i in range(len(u_start.values)):
tmp.dat.data[i] = u_start.values[i]
self.f.assign(tmp)
# Take Backward Euler step
self.solver.solve()
# Copy data from Firedrake Function object to VectorDiffusion2D object
ret = VectorDiffusion2D(size=len(self.function_space),
comm_space=self.comm_space)
ret.set_values(np.copy(self.soln.dat.data))
return ret
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 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)
def solver_CG(mesh, el, space, deg, T, dt=0.001, warm_up=False):
"""Solve the scalar wave equation on a unit square/cube using a
CG FEM formulation with several different element types.
Parameters
----------
mesh: Firedrake.mesh
A utility mesh from the Firedrake package
el: string
The type of element either "tria" or "quad".
`tria` in 3d implies tetrahedra and
`quad` in 3d implies hexahedral elements.
space: string
The space of the FEM. Available options are:
"CG": Continuous Galerkin Finite Elements,
"KMV": Kong-Mulder-Veldhuzien higher-order mass lumped elements
"S" (for Serendipity) (NB: quad/hexs only)
"spectral": spectral elements using GLL quad
points (NB: quads/hexs only)
deg: int
The spatial polynomial degree.
T: float
The simulation duration in simulation seconds.
dt: float, optional
Simulation timestep
warm_up: boolean, optional
Warm up symbolics by running one timestep and shutting down.
Returns
-------
u_n: Firedrake.Function
The solution at time `T`
"""
sd = mesh.geometric_dimension()
V = _build_space(mesh, el, space, deg)
quad_rule1, quad_rule2 = _build_quad_rule(el, V, space)
params = _select_params(space)
# DEBUG
# outfile = fd.File(os.getcwd() + "/results/simple_shots.pvd")
# END DEBUG
tot_dof = COMM_WORLD.allreduce(V.dof_dset.total_size, op=MPI.SUM)
# if COMM_WORLD.rank == 0:
# print("------------------------------------------")
# print("The problem has " + str(tot_dof) + " degrees of freedom.")
# print("------------------------------------------")
nt = int(T / dt) # number of timesteps
u = fd.TrialFunction(V)
v = fd.TestFunction(V)
u_np1 = fd.Function(V) # n+1
u_n = fd.Function(V) # n
u_nm1 = fd.Function(V) # n-1
# constant speed
c = Constant(1.5)
m = (
(1.0 / (c * c))
* (u - 2.0 * u_n + u_nm1)
/ Constant(dt * dt)
* v
* dx(rule=quad_rule1)
) # mass-like matrix
a = dot(grad(u_n), grad(v)) * dx(rule=quad_rule2) # stiffness matrix
# injection of source into mesh
ricker = Constant(0.0)
source = Constant([0.5] * sd)
coords = fd.SpatialCoordinate(mesh)
F = m + a - delta_expr(source, *coords) * ricker * v * dx(rule=quad_rule2)
a, r = fd.lhs(F), fd.rhs(F)
A, R = fd.assemble(a), fd.assemble(r)
solver = fd.LinearSolver(A, solver_parameters=params, options_prefix="")
# timestepping loop
results = []
t = 0.0
for step in range(nt):
with PETSc.Log.Stage("{el}{deg}".format(el=el, deg=deg)):
ricker.assign(RickerWavelet(t, freq=6))
R = fd.assemble(r, tensor=R)
solver.solve(u_np1, R)
snes = _get_time("SNESSolve")
ksp = _get_time("KSPSolve")
pcsetup = _get_time("PCSetUp")
pcapply = _get_time("PCApply")
jac = _get_time("SNESJacobianEval")
residual = _get_time("SNESFunctionEval")
sparsity = _get_time("CreateSparsity")
results.append(
[tot_dof, snes, ksp, pcsetup, pcapply, jac, residual, sparsity]
)
if warm_up:
# Warm up symbolics/disk cache
solver.solve(u_np1, R)
sys.exit("Warming up...")
u_nm1.assign(u_n)
u_n.assign(u_np1)
t = step * float(dt)
# if step % 10 == 0:
# outfile.write(u_n)
# print("Time is " + str(t), flush=True)
results = np.asarray(results)
if mesh.comm.rank == 0:
with open(
"data/scalar_wave.{el}.{deg}.{space}.csv".format(
el=el, deg=deg, space=space
),
"w",
) as f:
np.savetxt(
f,
results,
fmt=["%d"] + ["%e"] * 7,
delimiter=",",
header="tot_dof,SNESSolve,KSPSolve,PCSetUp,PCApply,SNESJacobianEval,SNESFunctionEval,CreateSparsity",
comments="",
)
return u_n
class DiagnosticEquations(object):
"""
An object setting up diagnostic equations and solvers for
the stochastic Camassa-Holm equation.
:arg diagnostic_variables: a DiagnosticVariables object.
:arg prognostic_variables: a PrognosticVariables object.
:arg outputting: an Outputting object.
:arg simulation_parameters: a dictionary storing the simulation parameters.
"""
def __init__(self, diagnostic_variables, prognostic_variables, outputting,
simulation_parameters):
self.diagnostic_variables = diagnostic_variables
self.prognostic_variables = prognostic_variables
self.outputting = outputting
self.simulation_parameters = simulation_parameters
Dt = Constant(simulation_parameters['dt'][-1])
Ld = simulation_parameters['Ld'][-1]
u = self.prognostic_variables.u
Xi = self.prognostic_variables.dXi
Vu = u.function_space()
vector_u = True if Vu.ufl_element() == VectorElement else False
ones = Function(
VectorFunctionSpace(self.prognostic_variables.mesh, "CG",
1)).project(as_vector([Constant(1.0)]))
self.to_update_constants = False
self.interpolators = []
self.projectors = []
self.solvers = []
mesh = u.function_space().mesh()
x, = SpatialCoordinate(mesh)
alphasq = simulation_parameters['alphasq'][-1]
periodic = simulation_parameters['periodic'][-1]
# do peakon data checks here
true_peakon_data = simulation_parameters['true_peakon_data'][-1]
if true_peakon_data is not None:
self.true_peakon_file = Dataset(
'results/' + true_peakon_data + '/data.nc', 'r')
# check length of file is correct
ndump = simulation_parameters['ndump'][-1]
tmax = simulation_parameters['tmax'][-1]
dt = simulation_parameters['dt'][-1]
if len(self.true_peakon_file['time'][:]) != int(tmax /
(ndump * dt)) + 1:
raise ValueError(
'If reading in true peakon data, the dump frequency must be the same as that used for the true peakon data.'
+
' Length of true peakon data as %i, but proposed length is %i'
% (len(self.true_peakon_file['time'][:]),
int(tmax / (ndump * dt)) + 1))
if self.true_peakon_file['p'][:].shape != (int(tmax /
(ndump * dt)) +
1, ):
raise ValueError(
'True peakon data shape %i must be the same shape as proposed data %i'
% ((int(tmax / (ndump * dt)) + 1, ),
self.true_peakon_file['p'][:].shape))
# do peakon data checks here
true_mean_peakon_data = simulation_parameters['true_mean_peakon_data'][
-1]
if true_mean_peakon_data is not None:
self.true_mean_peakon_file = Dataset(
'results/' + true_mean_peakon_data + '/data.nc', 'r')
# check length of file is correct
ndump = simulation_parameters['ndump'][-1]
tmax = simulation_parameters['tmax'][-1]
dt = simulation_parameters['dt'][-1]
if len(self.true_mean_peakon_file['time'][:]) != int(tmax /
(ndump * dt)):
raise ValueError(
'If reading in true peakon data, the dump frequency must be the same as that used for the true peakon data.'
)
if self.true_mean_peakon_file['p'][:].shape != (int(
tmax / (ndump * dt)), ):
raise ValueError(
'True peakon data must have same shape as proposed data!')
for key, value in self.diagnostic_variables.fields.items():
if key == 'uscalar':
uscalar = self.diagnostic_variables.fields['uscalar']
u_interpolator = Interpolator(dot(ones, u), uscalar)
self.interpolators.append(u_interpolator)
elif key == 'Euscalar':
Eu = self.prognostic_variables.Eu
Euscalar = self.diagnostic_variables.fields['Euscalar']
Eu_interpolator = Interpolator(dot(ones, Eu), Euscalar)
self.interpolators.append(Eu_interpolator)
elif key == 'Xiscalar':
Xi = self.prognostic_variables.dXi
Xiscalar = self.diagnostic_variables.fields['Xiscalar']
Xi_interpolator = Interpolator(dot(ones, Xi), Xiscalar)
self.interpolators.append(Xi_interpolator)
elif key == 'du':
if type(u.function_space().ufl_element()) == VectorElement:
u_to_project = self.diagnostic_variables.fields['uscalar']
else:
u_to_project = u
du = self.diagnostic_variables.fields['du']
du_projector = Projector(u_to_project.dx(0), du)
self.projectors.append(du_projector)
elif key == 'jump_du':
du = self.diagnostic_variables.fields['du']
jump_du = self.diagnostic_variables.fields['jump_du']
V = jump_du.function_space()
jtrial = TrialFunction(V)
psi = TestFunction(V)
Lj = psi('+') * abs(jump(du)) * dS
aj = psi('+') * jtrial('+') * dS
jprob = LinearVariationalProblem(aj, Lj, jump_du)
jsolver = LinearVariationalSolver(jprob)
self.solvers.append(jsolver)
elif key == 'du_smooth':
du = self.diagnostic_variables.fields['du']
du_smooth = self.diagnostic_variables.fields['du_smooth']
projector = Projector(du, du_smooth)
self.projectors.append(projector)
elif key == 'u2_flux':
gamma = simulation_parameters['gamma'][-1]
u2_flux = self.diagnostic_variables.fields['u2_flux']
xis = self.prognostic_variables.pure_xi_list
xis_x = []
xis_xxx = []
CG1 = FunctionSpace(mesh, "CG", 1)
psi = TestFunction(CG1)
for xi in xis:
xis_x.append(Function(CG1).project(xi.dx(0)))
for xi_x in xis_x:
xi_xxx = Function(CG1)
form = (psi * xi_xxx + psi.dx(0) * xi_x.dx(0)) * dx
prob = NonlinearVariationalProblem(form, xi_xxx)
solver = NonlinearVariationalSolver(prob)
solver.solve()
xis_xxx.append(xi_xxx)
flux_expr = 0.0 * x
for xi, xi_x, xi_xxx in zip(xis, xis_x, xis_xxx):
flux_expr += (6 * u.dx(0) * xi + 12 * u * xi_x + gamma *
xi_xxx) * (6 * u.dx(0) * xi + 24 * u * xi_x +
gamma * xi_xxx)
projector = Projector(flux_expr, u2_flux)
self.projectors.append(projector)
elif key == 'a':
# find 6 * u_x * Xi + gamma * Xi_xxx
mesh = u.function_space().mesh()
gamma = simulation_parameters['gamma'][-1]
a_flux = self.diagnostic_variables.fields['a']
xis = self.prognostic_variables.pure_xis
xis_x = []
xis_xxx = []
CG1 = FunctionSpace(mesh, "CG", 1)
psi = TestFunction(CG1)
for xi in xis:
xis_x.append(Function(CG1).project(xi.dx(0)))
for xi_x in xis_x:
xi_xxx = Function(CG1)
form = (psi * xi_xxx + psi.dx(0) * xi_x.dx(0)) * dx
prob = NonlinearVariationalProblem(form, xi_xxx)
solver = NonlinearVariationalSolver(prob)
solver.solve()
xis_xxx.append(xi_xxx)
x, = SpatialCoordinate(mesh)
a_expr = 0.0 * x
for xi, xi_x, xi_xxx in zip(xis, xis_x, xis_xxx):
a_expr += 6 * u.dx(0) * xi + gamma * xi_xxx
projector = Projector(a_expr, a_flux)
self.projectors.append(projector)
elif key == 'b':
# find 12 * u * Xi_x
mesh = u.function_space().mesh()
gamma = simulation_parameters['gamma'][-1]
b_flux = self.diagnostic_variables.fields['b']
xis = self.prognostic_variables.pure_xis
x, = SpatialCoordinate(mesh)
b_expr = 0.0 * x
for xi, xi_x, xi_xxx in zip(xis, xis_x, xis_xxx):
b_expr += 12 * u * xi.dx(0)
projector = Projector(b_expr, b_flux)
self.projectors.append(projector)
elif key == 'kdv_1':
# find the first part of the kdv form
u0 = prognostic_variables.u0
uh = (u + u0) / 2
us = Dt * uh + sqrt(Dt) * Xi
psi = TestFunction(Vu)
du_1 = self.diagnostic_variables.fields['kdv_1']
eqn = psi * du_1 * dx - 6 * psi.dx(0) * uh * us * dx
prob = NonlinearVariationalProblem(eqn, du_1)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
elif key == 'kdv_2':
# find the second part of the kdv form
u0 = prognostic_variables.u0
uh = (u + u0) / 2
us = Dt * uh + sqrt(Dt) * Xi
psi = TestFunction(Vu)
du_2 = self.diagnostic_variables.fields['kdv_2']
eqn = psi * du_2 * dx + 6 * psi * uh * us.dx(0) * dx
prob = NonlinearVariationalProblem(eqn, du_2)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
elif key == 'kdv_3':
# find the third part of the kdv form
u0 = prognostic_variables.u0
uh = (u + u0) / 2
us = Dt * uh + sqrt(Dt) * Xi
du_3 = self.diagnostic_variables.fields['kdv_3']
gamma = simulation_parameters['gamma'][-1]
phi = TestFunction(Vu)
F = Function(Vu)
eqn = (phi * F * dx + phi.dx(0) * us.dx(0) * dx)
prob = NonlinearVariationalProblem(eqn, F)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
self.projectors.append(Projector(-gamma * F.dx(0), du_3))
# nu = TestFunction(Vu)
# back_eqn = nu * du_3 * dx - gamma * nu.dx(0) * F * dx
# back_prob = NonlinearVariationalProblem(back_eqn, du_3)
# back_solver = NonlinearVariationalSolver(back_prob)
# self.solvers.append(solver)
elif key == 'm':
m = self.diagnostic_variables.fields['m']
phi = TestFunction(Vu)
eqn = phi * m * dx - phi * u * dx - alphasq * phi.dx(0) * u.dx(
0) * dx
prob = NonlinearVariationalProblem(eqn, m)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
elif key == 'u_xx':
u_xx = self.diagnostic_variables.fields['u_xx']
phi = TestFunction(Vu)
eqn = phi * u_xx * dx + phi.dx(0) * u_xx.dx(0) * dx
prob = NonlinearVariationalProblem(eqn, u_xx)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
elif key == 'u_sde':
self.to_update_constants = True
self.Ld = Ld
self.alphasq = alphasq
self.p = Constant(1.0 * 0.5 * (1 + exp(-Ld / sqrt(alphasq))) /
(1 - exp(-Ld / sqrt(alphasq))))
self.q = Constant(Ld / 2)
u_sde = self.diagnostic_variables.fields['u_sde']
if periodic:
expr = conditional(
x < self.q - Ld / 2,
self.p * ((exp(-(x - self.q + Ld) / sqrt(alphasq)) +
exp(-Ld / sqrt(alphasq)) * exp(
(x - self.q + Ld) / sqrt(alphasq))) /
(1 - exp(-Ld / sqrt(alphasq)))),
conditional(
x < self.q + Ld / 2,
self.p * ((exp(-sqrt((self.q - x)**2 / alphasq)) +
exp(-Ld / sqrt(alphasq)) *
exp(sqrt((self.q - x)**2 / alphasq))) /
(1 - exp(-Ld / sqrt(alphasq)))),
self.p *
((exp(-(self.q + Ld - x) / sqrt(alphasq)) +
exp(-Ld / sqrt(alphasq) * exp(
(self.q + Ld - x) / sqrt(alphasq)))) /
(1 - exp(-Ld / sqrt(alphasq))))))
else:
expr = conditional(
x < self.q - Ld / 2,
self.p * exp(-(x - self.q + Ld) / sqrt(alphasq)),
conditional(
x < self.q + Ld / 2,
self.p * exp(-sqrt((self.q - x)**2 / alphasq)),
self.p * exp(-(self.q + Ld - x) / sqrt(alphasq))))
self.interpolators.append(Interpolator(expr, u_sde))
elif key == 'u_sde_weak':
u_sde = self.diagnostic_variables.fields['u_sde']
u_sde_weak = self.diagnostic_variables.fields['u_sde_weak']
psi = TestFunction(Vu)
eqn = psi * u_sde_weak * dx - psi * (u - u_sde) * dx
prob = NonlinearVariationalProblem(eqn, u_sde_weak)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
elif key == 'u_sde_mean':
self.to_update_constants = True
self.p = Constant(1.0)
self.q = Constant(Ld / 2)
if periodic:
raise NotImplementedError(
'u_sde_mean not yet implemented for periodic peakon')
u_sde = self.diagnostic_variables.fields['u_sde_mean']
expr = conditional(
x < self.q - Ld / 2,
self.p * exp(-(x - self.q + Ld) / sqrt(alphasq)),
conditional(
x < self.q + Ld / 2,
self.p * exp(-sqrt((self.q - x)**2 / alphasq)),
self.p * exp(-(self.q + Ld - x) / sqrt(alphasq))))
self.interpolators.append(Interpolator(expr, u_sde))
elif key == 'u_sde_weak_mean':
u_sde = self.diagnostic_variables.fields['u_sde_mean']
u_sde_weak = self.diagnostic_variables.fields[
'u_sde_weak_mean']
psi = TestFunction(Vu)
eqn = psi * u_sde_weak * dx - psi * (u - u_sde) * dx
prob = NonlinearVariationalProblem(eqn, u_sde_weak)
solver = NonlinearVariationalSolver(prob)
self.solvers.append(solver)
elif key == 'pure_xi':
pure_xi = 0.0 * x
for xi in self.prognostic_variables.pure_xi_list:
if vector_u:
pure_xi += dot(ones, xi)
else:
pure_xi += xi
Xiscalar = self.diagnostic_variables.fields['pure_xi']
Xi_interpolator = Interpolator(pure_xi, Xiscalar)
self.interpolators.append(Xi_interpolator)
elif key == 'pure_xi_x':
pure_xi_x = 0.0 * x
for xix in self.prognostic_variables.pure_xi_x_list:
if vector_u:
pure_xi_x += dot(ones, xix)
else:
pure_xi_x += xix
Xiscalar = self.diagnostic_variables.fields['pure_xi_x']
Xi_interpolator = Interpolator(pure_xi_x, Xiscalar)
self.interpolators.append(Xi_interpolator)
elif key == 'pure_xi_xx':
pure_xi_xx = 0.0 * x
for xixx in self.prognostic_variables.pure_xi_xx_list:
if vector_u:
pure_xi_xx += dot(ones, xixx)
else:
pure_xi_xx += xixx
Xiscalar = self.diagnostic_variables.fields['pure_xi_xx']
Xi_interpolator = Interpolator(pure_xi_xx, Xiscalar)
self.interpolators.append(Xi_interpolator)
elif key == 'pure_xi_xxx':
pure_xi_xxx = 0.0 * x
for xixxx in self.prognostic_variables.pure_xi_xxx_list:
if vector_u:
pure_xi_xxx += dot(ones, xixxx)
else:
pure_xi_xxx += xixxx
Xiscalar = self.diagnostic_variables.fields['pure_xi_xxx']
Xi_interpolator = Interpolator(pure_xi_xxx, Xiscalar)
self.interpolators.append(Xi_interpolator)
elif key == 'pure_xi_xxxx':
pure_xi_xxxx = 0.0 * x
for xixxxx in self.prognostic_variables.pure_xi_xx_list:
if vector_u:
pure_xi_xxxx += dot(ones, xixxxx)
else:
pure_xi_xxxx += xixxxx
Xiscalar = self.diagnostic_variables.fields['pure_xi_xxxx']
Xi_interpolator = Interpolator(pure_xi_xxxx, Xiscalar)
self.interpolators.append(Xi_interpolator)
else:
raise NotImplementedError('Diagnostic %s not yet implemented' %
key)
def update_constants(self):
"""
Update p and q constants from true peakon data.
"""
self.p.assign(self.true_peakon_file['p'][self.outputting.t_idx] * 0.5 *
(1 + exp(-self.Ld / sqrt(self.alphasq))) /
(1 - exp(-self.Ld / sqrt(self.alphasq))))
self.q.assign(self.true_peakon_file['q'][self.outputting.t_idx])
def solve(self):
"""
Do interpolations, projections and solves.
"""
if self.to_update_constants:
self.update_constants()
for interpolator in self.interpolators:
interpolator.interpolate()
for projector in self.projectors:
projector.project()
for solver in self.solvers:
solver.solve()
class TimeMixin:
"""
Class to extend Problems as a mixin.
This mixin extends the problem to allow time stepping by adding the dT/dt
term.
To use, define a new class with this in the inheritance chain.
i.e::
class NewProblem(TimeMixin, Problem):
pass
Attributes:
T (firedrake.Function):
The trial function for the problem.
v (firedrake.Function):
The test function for the problem.
a (firedrake.Function):
The section containing the combination of terms involving both T
and v.
T_ (firedrake.Function):
A function used to hold the previous value for T.
C (firedrake.Function):
A function used to hold the heat capacity of the mesh.
_delT (firedrake.Function):
A function used as a placeholder for the time derivative.
max_t (float):
The time to iterate up to.
dt (float):
The size of each time step.
_dt_invc (firedrake.Constant):
Utility function to pass time step information to firedrake.
steps (int):
Number of iterations.
steady_state (bool):
Toggle whether solving steady state (dT/dt = 0) or time dependent
problem.
"""
# pylint: disable=too-many-instance-attributes, no-member
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 set_method(self, method='BackwardEuler', **kwargs):
"""
Replace T with the correct substitution for the method, then replace
the delT placeholder with the simple finite difference approximation:
dT/dt ~ (T - T_)/delta_t
Args:
method (str, optional):
The method to use. Defaults to 'BackwardEuler'.
"""
T = self.T
iter_method = IterationMethod(self)
substitution = iter_method.get_substitution(method, **kwargs)
self._update_func('T', substitution)
self.T = T
delT = (self.T - self.T_) * self._dt_invc
self._update_func('_delT', delT)
def remove_timescale(self):
"""
Set 1/dt to 0 so that the M terms vanish.
"""
self.steady_state = True
self._dt_invc.assign(0)
def set_timescale(self, max_t=None, dt=None, steps=None):
"""
Set the time stepping variables (max_t, dt, and number of steps).
This method is designed to be given 2 variables and calculate the
third. If 3 variables are given it will check that they are consistent.
Args:
max_t (float, optional):
The time to iterate until. Defaults to None.
dt (float, optional):
The increase in time for each iteration. Defaults to None.
steps (int, optional):
The number of steps to take. Defaults to None.
Raises:
ValueError: If less than 2 variables are provided.
ValueError: If 3 values are provided which are inconsistent.
RuntimeWarning: If steps is not an int.
"""
num_var_none = 0
if max_t is None:
num_var_none += 1
if dt is None:
num_var_none += 1
if steps is None:
num_var_none += 1
if num_var_none > 1:
raise ValueError('Must specify at least 2 of max_t, dt, and steps')
if num_var_none == 0:
calc_steps = int(max_t / dt)
if calc_steps != max_t / dt:
calc_steps += 1
if steps != calc_steps:
raise ValueError('Conflicting arguments. Try specifying only 2'
' of max_t, dt, and steps.')
if max_t is None:
max_t = dt * steps
if dt is None:
dt = max_t / steps
if steps is None:
steps = max_t / dt
if steps != int(steps):
steps = int(steps + 1)
LOGGER.warning("steps is not an integer, rounding up.")
self.max_t = max_t
self.dt = dt
self._dt_invc.assign(1 / dt)
self.steps = steps
self.steady_state = False
def _M(self):
"""
Create the mass matrix section.
Returns:
Function: The complete mass matrix section using delT.
"""
return self.C * self._delT * self.v * dx
