def __init__(self, mesh, conditions, timestepping, params, output, solver_params):
super().__init__(mesh, conditions, timestepping, params, output, solver_params)
self.u0 = Function(self.V)
self.u1 = Function(self.V)
self.sigma0 = Function(self.S)
self.sigma1 = Function(self.S)
theta = conditions.theta
uh = (1-theta) * self.u0 + theta * self.u1
a = Function(self.U)
h = Function(self.U)
p = TestFunction(self.V)
q = TestFunction(self.S)
self.initial_condition((self.u0, conditions.ic['u']), (a, conditions.ic['a']), (h, conditions.ic['h']))
ep_dot = self.strain(grad(uh))
zeta = self.zeta(h, a, self.delta(uh))
eta = zeta * params.e ** (-2)
rheology = 2 * eta * ep_dot + (zeta - eta) * tr(ep_dot) * Identity(2) - 0.5 * self.Ice_Strength(h, a) * Identity(2)
self.initial_condition((self.sigma0, rheology),(self.sigma1, self.sigma0))
def sigma_next(timestep, zeta, ep_dot, sigma, P):
A = 1 + 0.25 * (timestep * params.e ** 2) / params.T
B = timestep * 0.125 * (1 - params.e ** 2) / params.T
rhs = (1 - (timestep * params.e ** 2) / (4 * params.T)) * sigma - timestep / params.T * (
0.125 * (1 - params.e ** 2) * tr(sigma) * Identity(2) - 0.25 * P * Identity(2) + zeta * ep_dot)
C = (rhs[0, 0] - rhs[1, 1]) / A
D = (rhs[0, 0] + rhs[1, 1]) / (A + 2 * B)
sigma00 = 0.5 * (C + D)
sigma11 = 0.5 * (D - C)
sigma01 = rhs[0, 1]
sigma = as_matrix([[sigma00, sigma01], [sigma01, sigma11]])
return sigma
s = sigma_next(self.timestep, zeta, ep_dot, self.sigma0, self.Ice_Strength(h, a))
sh = (1-theta) * s + theta * self.sigma0
eqn = self.momentum_equation(h, self.u1, self.u0, p, sh, params.rho, uh, conditions.ocean_curr,
params.rho_a, params.C_a, params.rho_w, params.C_w, conditions.geo_wind, params.cor, self.timestep, ind=self.ind)
tensor_eqn = inner(self.sigma1-s, q) * dx
if conditions.stabilised['state']:
alpha = conditions.stabilised['alpha']
eqn += stabilisation_term(alpha=alpha, zeta=avg(zeta), mesh=mesh, v=uh, test=p)
bcs = DirichletBC(self.V, conditions.bc['u'], "on_boundary")
uprob = NonlinearVariationalProblem(eqn, self.u1, bcs)
self.usolver = NonlinearVariationalSolver(uprob, solver_parameters=solver_params.bt_params)
sprob = NonlinearVariationalProblem(tensor_eqn, self.sigma1)
self.ssolver = NonlinearVariationalSolver(sprob, solver_parameters=solver_params.bt_params)
def setup_solver(self, up_init=None):
""" Setup the solvers
"""
self.up0 = Function(self.W)
if up_init is not None:
chk_in = checkpointing.HDF5File(up_init, file_mode='r')
chk_in.read(self.up0, "/up")
chk_in.close()
self.u0, self.p0 = split(self.up0)
self.up = Function(self.W)
if up_init is not None:
chk_in = checkpointing.HDF5File(up_init, file_mode='r')
chk_in.read(self.up, "/up")
chk_in.close()
self.u1, self.p1 = split(self.up)
self.up.sub(0).rename("velocity")
self.up.sub(1).rename("pressure")
v, q = TestFunctions(self.W)
h = CellVolume(self.mesh)
u_norm = sqrt(dot(self.u0, self.u0))
if self.has_nullspace:
nullspace = MixedVectorSpaceBasis(
self.W,
[self.W.sub(0), VectorSpaceBasis(constant=True)])
else:
nullspace = None
tau = ((2.0 / self.dt)**2 + (2.0 * u_norm / h)**2 +
(4.0 * self.nu / h**2)**2)**(-0.5)
# temporal discretization
F = (1.0 / self.dt) * inner(self.u1 - self.u0, v) * dx
# weak form
F += (+inner(dot(self.u0, nabla_grad(self.u1)), v) * dx +
self.nu * inner(grad(self.u1), grad(v)) * dx -
(1.0 / self.rho) * self.p1 * div(v) * dx +
div(self.u1) * q * dx - inner(self.forcing, v) * dx)
# residual form
R = (+(1.0 / self.dt) * (self.u1 - self.u0) +
dot(self.u0, nabla_grad(self.u1)) - self.nu * div(grad(self.u1)) +
(1.0 / self.rho) * grad(self.p1) - self.forcing)
# GLS
F += tau * inner(
+dot(self.u0, nabla_grad(v)) - self.nu * div(grad(v)) +
(1.0 / self.rho) * grad(q), R) * dx
self.problem = NonlinearVariationalProblem(F, self.up, self.bcs)
self.solver = NonlinearVariationalSolver(
self.problem,
nullspace=nullspace,
solver_parameters=self.solver_parameters)
def __init__(self,
F,
butcher_tableau,
t,
dt,
u0,
bcs=None,
solver_parameters=None,
update_solver_parameters=None,
splitting=AI,
nullspace=None,
appctx=None):
self.u0 = u0
self.t = t
self.dt = dt
self.num_fields = len(u0.function_space())
self.num_stages = len(butcher_tableau.b)
self.butcher_tableau = butcher_tableau
Fbig, update_stuff, UU, bigBCs, gblah, nsp = getFormStage(
F, butcher_tableau, u0, t, dt, bcs, splitting=splitting)
self.UU = UU
self.bigBCs = bigBCs
self.bcdat = gblah
self.update_stuff = update_stuff
self.prob = NonlinearVariationalProblem(Fbig, UU, bigBCs)
appctx_irksome = {
"F": F,
"butcher_tableau": butcher_tableau,
"t": t,
"dt": dt,
"u0": u0,
"bcs": bcs,
"nullspace": nullspace
}
if appctx is None:
appctx = appctx_irksome
else:
appctx = {**appctx, **appctx_irksome}
self.solver = NonlinearVariationalSolver(
self.prob,
appctx=appctx,
nullspace=nsp,
solver_parameters=solver_parameters)
unew, Fupdate, update_bcs, update_bcs_gblah = self.update_stuff
self.update_problem = NonlinearVariationalProblem(
Fupdate, unew, update_bcs)
self.update_solver = NonlinearVariationalSolver(
self.update_problem, solver_parameters=update_solver_parameters)
self._update = self._update_general
def __init__(self,
F,
butcher_tableau,
t,
dt,
u0,
bcs=None,
solver_parameters=None):
self.u0 = u0
self.t = t
self.dt = dt
self.num_fields = len(u0.function_space())
self.num_stages = len(butcher_tableau.b)
self.butcher_tableau = butcher_tableau
bigF, stages, bigBCs, bigBCdata = \
getForm(F, butcher_tableau, t, dt, u0, bcs)
self.stages = stages
self.bigBCs = bigBCs
self.bigBCdata = bigBCdata
problem = NLVP(bigF, stages, bigBCs)
self.solver = NLVS(problem, solver_parameters=solver_parameters)
self.ks = stages.split()
def __init__(self, F, butcher_tableau, t, dt, u0, bcs=None,
solver_parameters=None, splitting=AI,
appctx=None, nullspace=None, bc_type="DAE"):
self.u0 = u0
self.t = t
self.dt = dt
self.num_fields = len(u0.function_space())
self.num_stages = len(butcher_tableau.b)
self.butcher_tableau = butcher_tableau
bigF, stages, bigBCs, bigNSP, bigBCdata = \
getForm(F, butcher_tableau, t, dt, u0, bcs, bc_type, splitting, nullspace)
self.stages = stages
self.bigBCs = bigBCs
self.bigBCdata = bigBCdata
problem = NLVP(bigF, stages, bigBCs)
appctx_irksome = {"F": F,
"butcher_tableau": butcher_tableau,
"t": t,
"dt": dt,
"u0": u0,
"bcs": bcs,
"bc_type": bc_type,
"splitting": splitting,
"nullspace": nullspace}
if appctx is None:
appctx = appctx_irksome
else:
appctx = {**appctx, **appctx_irksome}
push_parent(u0.function_space().dm, stages.function_space().dm)
self.solver = NLVS(problem,
appctx=appctx,
solver_parameters=solver_parameters,
nullspace=bigNSP)
pop_parent(u0.function_space().dm, stages.function_space().dm)
if self.num_stages == 1 and self.num_fields == 1:
self.ws = (stages,)
else:
self.ws = stages.split()
A1, A2 = splitting(butcher_tableau.A)
try:
self.updateb = numpy.linalg.solve(A2.T, butcher_tableau.b)
except numpy.linalg.LinAlgError:
raise NotImplementedError("A=A1 A2 splitting needs A2 invertible")
boo = numpy.zeros(self.updateb.shape, dtype=self.updateb.dtype)
boo[-1] = 1
if numpy.allclose(self.updateb, boo):
self._update = self._update_A2Tmb
else:
self._update = self._update_general
def __init__(self, mesh, conditions, timestepping, params, output, solver_params):
super().__init__(mesh, conditions, timestepping, params, output, solver_params)
self.w0 = Function(self.W3)
self.w1 = Function(self.W3)
u0, s0, h0, a0 = self.w0.split()
p, q, r, m = TestFunctions(self.W3)
self.initial_condition((u0, conditions.ic['u']), (s0, conditions.ic['s']),
(a0, conditions.ic['a']), (h0, conditions.ic['h']))
self.w1.assign(self.w0)
u1, s1, h1, a1 = split(self.w1)
u0, s0, h0, a0 = split(self.w0)
theta = conditions.theta
uh = (1-theta) * u0 + theta * u1
sh = (1-theta) * s0 + theta * s1
hh = (1-theta) * h0 + theta * h1
ah = (1-theta) * a0 + theta * a1
ep_dot = self.strain(grad(uh))
zeta = self.zeta(hh, ah, self.delta(uh))
rheology = params.e ** 2 * sh + Identity(2) * 0.5 * ((1 - params.e ** 2) * tr(sh) + self.Ice_Strength(hh, ah))
eqn = self.momentum_equation(hh, u1, u0, p, sh, params.rho, uh, conditions.ocean_curr, params.rho_a,
params.C_a, params.rho_w, params.C_w, conditions.geo_wind, params.cor, self.timestep, ind=self.ind)
eqn += self.transport_equation(uh, hh, ah, h1, h0, a1, a0, r, m, self.n, self.timestep)
eqn += inner(self.ind * (s1 - s0) + 0.5 * self.timestep * rheology / params.T, q) * dx
eqn -= inner(q * zeta * self.timestep / params.T, ep_dot) * dx
if conditions.stabilised['state']:
alpha = conditions.stabilised['alpha']
eqn += self.stabilisation_term(alpha=alpha, zeta=avg(zeta), mesh=mesh, v=uh, test=p)
bcs = DirichletBC(self.W3.sub(0), conditions.bc['u'], "on_boundary")
uprob = NonlinearVariationalProblem(eqn, self.w1, bcs)
self.usolver = NonlinearVariationalSolver(uprob, solver_parameters=solver_params.bt_params)
self.u1, self.s0, self.h1, self.a1 = self.w1.split()
def __init__(self, mesh, conditions, timestepping, params, output, solver_params):
super().__init__(mesh, conditions, timestepping, params, output, solver_params)
self.w0 = Function(self.W2)
self.w1 = Function(self.W2)
u0, h0, a0 = self.w0.split()
p, q, r = TestFunctions(self.W2)
self.initial_condition((u0, conditions.ic['u']), (h0, conditions.ic['h']),
(a0, conditions.ic['a']))
self.w1.assign(self.w0)
u1, h1, a1 = split(self.w1)
u0, h0, a0 = split(self.w0)
theta = conditions.theta
uh = (1-theta) * u0 + theta * u1
ah = (1-theta) * a0 + theta * a1
hh = (1-theta) * h0 + theta * h1
ep_dot = self.strain(grad(uh))
zeta = self.zeta(hh, ah, self.delta(uh))
eta = zeta * params.e ** (-2)
sigma = 2 * eta * ep_dot + (zeta - eta) * tr(ep_dot) * Identity(2) - self.Ice_Strength(hh, ah) * 0.5 * Identity(
2)
eqn = self.momentum_equation(hh, u1, u0, p, sigma, params.rho, uh, conditions.ocean_curr, params.rho_a,
params.C_a, params.rho_w, params.C_w, conditions.geo_wind, params.cor, self.timestep)
eqn += self.transport_equation(uh, hh, ah, h1, h0, a1, a0, q, r, self.n, self.timestep)
if conditions.stabilised['state']:
alpha = conditions.stabilised['alpha']
eqn += self.stabilisation_term(alpha=alpha, zeta=avg(zeta), mesh=mesh, v=uh, test=p)
bcs = DirichletBC(self.W2.sub(0), conditions.bc['u'], "on_boundary")
uprob = NonlinearVariationalProblem(eqn, self.w1, bcs)
self.usolver = NonlinearVariationalSolver(uprob, solver_parameters=solver_params.bt_params)
self.u1, self.h1, self.a1 = self.w1.split()
def wrapper(self, **kwargs):
"""To disable the annotation, just pass :py:data:`annotate=False` to this routine, and it acts exactly like the
Firedrake solve call. This is useful in cases where the solve is known to be irrelevant or diagnostic
for the purposes of the adjoint computation (such as projecting fields to other function spaces
for the purposes of visualisation)."""
annotate = annotate_tape(kwargs)
if annotate:
tape = get_working_tape()
problem = self._ad_problem
sb_kwargs = NonlinearVariationalSolveBlock.pop_kwargs(kwargs)
sb_kwargs.update(kwargs)
block = NonlinearVariationalSolveBlock(
problem._ad_F == 0,
problem._ad_u,
problem._ad_bcs,
problem_J=problem._ad_J,
solver_params=self.parameters,
solver_kwargs=self._ad_kwargs,
**sb_kwargs)
if not self._ad_nlvs:
from firedrake import NonlinearVariationalSolver
self._ad_nlvs = NonlinearVariationalSolver(
self._ad_problem_clone(self._ad_problem,
block.get_dependencies()),
**self._ad_kwargs)
block._ad_nlvs = self._ad_nlvs
tape.add_block(block)
with stop_annotating():
out = solve(self, **kwargs)
if annotate:
block.add_output(
self._ad_problem._ad_u.create_block_variable())
return out
def solve(self, file_path='ttip_result/solution.pvd'):
"""
Setup and solve the nonlinear problem.
Save value to file given.
Any additional keyword arguments are passed to the iteration method.
Args:
file_path (string, optional):
The path to save the pvd file to.
vtk files will be generated in the same directory as the pvd.
It is recommended that this is a separate drectory per run.
Defaults to 'TTiP_result/solution.pvd'.
"""
F = self.problem.a - self.problem.L
steady_state = self.is_steady_state()
if isinstance(self.problem, BoundaryMixin):
var_prob = NonlinearVariationalProblem(F,
self.u,
bcs=self.problem.bcs)
else:
var_prob = NonlinearVariationalProblem(F, self.u)
solver = NonlinearVariationalSolver(problem=var_prob,
solver_parameters=self.params)
outfile = File(file_path)
outfile.write(self.u, target_degree=1, target_continuity=H1)
if steady_state:
solver.solve()
outfile.write(self.u, target_degree=1, target_continuity=H1)
else:
self.problem.T_.assign(self.u)
last_perc = 0
for i in range(self.problem.steps):
solver.solve()
perc = int(100 * (i + 1) / self.problem.steps)
if perc > last_perc:
print(f'{perc}%')
last_perc = perc
self.problem.T_.assign(self.u)
outfile.write(self.u, target_degree=1, target_continuity=H1)
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
def compressible_hydrostatic_balance(state,
theta0,
rho0,
exner0=None,
top=False,
exner_boundary=Constant(1.0),
mr_t=None,
solve_for_rho=False,
params=None):
"""
Compute a hydrostatically balanced density given a potential temperature
profile. By default, this uses a vertically-oriented hybridization
procedure for solving the resulting discrete systems.
:arg state: The :class:`State` object.
:arg theta0: :class:`.Function`containing the potential temperature.
:arg rho0: :class:`.Function` to write the initial density into.
:arg top: If True, set a boundary condition at the top. Otherwise, set
it at the bottom.
:arg exner_boundary: a field or expression to use as boundary data for exner
on the top or bottom as specified.
:arg mr_t: the initial total water mixing ratio field.
"""
# Calculate hydrostatic Pi
VDG = state.spaces("DG")
Vu = state.spaces("HDiv")
Vv = FunctionSpace(state.mesh, Vu.ufl_element()._elements[-1])
W = MixedFunctionSpace((Vv, VDG))
v, exner = TrialFunctions(W)
dv, dexner = TestFunctions(W)
n = FacetNormal(state.mesh)
cp = state.parameters.cp
# add effect of density of water upon theta
theta = theta0
if mr_t is not None:
theta = theta0 / (1 + mr_t)
alhs = ((cp * inner(v, dv) - cp * div(dv * theta) * exner) * dx +
dexner * div(theta * v) * dx)
if top:
bmeasure = ds_t
bstring = "bottom"
else:
bmeasure = ds_b
bstring = "top"
arhs = -cp * inner(dv, n) * theta * exner_boundary * bmeasure
# Possibly make g vary with spatial coordinates?
g = state.parameters.g
arhs -= g * inner(dv, state.k) * dx
bcs = [DirichletBC(W.sub(0), zero(), bstring)]
w = Function(W)
exner_problem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs)
if params is None:
params = {
'ksp_type': 'preonly',
'pc_type': 'python',
'mat_type': 'matfree',
'pc_python_type': 'gusto.VerticalHybridizationPC',
# Vertical trace system is only coupled vertically in columns
# block ILU is a direct solver!
'vert_hybridization': {
'ksp_type': 'preonly',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
}
}
exner_solver = LinearVariationalSolver(exner_problem,
solver_parameters=params,
options_prefix="exner_solver")
exner_solver.solve()
v, exner = w.split()
if exner0 is not None:
exner0.assign(exner)
if solve_for_rho:
w1 = Function(W)
v, rho = w1.split()
rho.interpolate(thermodynamics.rho(state.parameters, theta0, exner))
v, rho = split(w1)
dv, dexner = TestFunctions(W)
exner = thermodynamics.exner_pressure(state.parameters, rho, theta0)
F = ((cp * inner(v, dv) - cp * div(dv * theta) * exner) * dx +
dexner * div(theta0 * v) * dx +
cp * inner(dv, n) * theta * exner_boundary * bmeasure)
F += g * inner(dv, state.k) * dx
rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs)
rhosolver = NonlinearVariationalSolver(rhoproblem,
solver_parameters=params,
options_prefix="rhosolver")
rhosolver.solve()
v, rho_ = w1.split()
rho0.assign(rho_)
else:
rho0.interpolate(thermodynamics.rho(state.parameters, theta0, exner))
def __init__(self, prognostic_variables, simulation_parameters):
mesh = simulation_parameters['mesh'][-1]
x, = SpatialCoordinate(mesh)
Ld = simulation_parameters['Ld'][-1]
self.scheme = simulation_parameters['scheme'][-1]
self.dt = simulation_parameters['dt'][-1]
self.num_Xis = simulation_parameters['num_Xis'][-1]
self.Xi_family = simulation_parameters['Xi_family'][-1]
self.dXi = prognostic_variables.dXi
self.dWs = [Constant(0.0) for dw in range(self.num_Xis)]
self.dW_nums = prognostic_variables.dW_nums
self.Xi_functions = []
self.nXi_updates = simulation_parameters['nXi_updates'][-1]
self.smooth_t = simulation_parameters['smooth_t'][-1]
self.fixed_dW = simulation_parameters['fixed_dW'][-1]
if self.smooth_t is not None and self.nXi_updates > 1:
raise ValueError('Prescribing forcing and including multiple Xi updates are not compatible.')
if self.smooth_t is not None or self.fixed_dW is not None:
print('WARNING: Remember to change sigma to sigma * sqrt(dt) with the prescribed forcing option or the fixed_dW option.')
seed = simulation_parameters['seed'][-1]
np.random.seed(seed)
# make sure sigma is a Constant
if self.num_Xis != 0:
if isinstance(simulation_parameters['sigma'][-1], Constant):
self.sigma = simulation_parameters['sigma'][-1]
else:
self.sigma = Constant(simulation_parameters['sigma'][-1])
else:
self.sigma = Constant(0.0)
self.pure_xi_list = prognostic_variables.pure_xi_list
self.pure_xi_x_list = prognostic_variables.pure_xi_x_list
self.pure_xi_xx_list = prognostic_variables.pure_xi_xx_list
self.pure_xi_xxx_list = prognostic_variables.pure_xi_xxx_list
self.pure_xi_xxxx_list = prognostic_variables.pure_xi_xxxx_list
for xi in range(self.num_Xis):
self.pure_xi_list.append(Function(self.dXi.function_space()))
self.pure_xi_x_list.append(Function(self.dXi.function_space()))
self.pure_xi_xx_list.append(Function(self.dXi.function_space()))
self.pure_xi_xxx_list.append(Function(self.dXi.function_space()))
self.pure_xi_xxxx_list.append(Function(self.dXi.function_space()))
if self.Xi_family == 'sines':
for n in range(self.num_Xis):
if (n+1) % 2 == 1:
self.Xi_functions.append(self.sigma * sin(2*(n+1)*pi*x/Ld))
else:
self.Xi_functions.append(self.sigma * cos(2*(n+1)*pi*x/Ld))
elif self.Xi_family == 'double_sines':
for n in range(self.num_Xis):
if (n+1) % 2 == 1:
self.Xi_functions.append(self.sigma * sin(4*(n+1)*pi*x/Ld))
else:
self.Xi_functions.append(self.sigma * cos(4*(n+1)*pi*x/Ld))
elif self.Xi_family == 'high_freq_sines':
for n in range(self.num_Xis):
if (n+1) % 2 == 1:
self.Xi_functions.append(self.sigma * sin((2*(n+1)+10)*pi*x/Ld))
else:
self.Xi_functions.append(self.sigma * cos((2*(n+1)+10)*pi*x/Ld))
elif self.Xi_family == 'gaussians':
for n in range(self.num_Xis):
self.Xi_functions.append(self.sigma * 0.5*self.num_Xis*exp(-((x-Ld*(n+1)/(self.num_Xis +1.0))/2.)**2))
elif self.Xi_family == 'quadratic':
if self.num_Xis > 1:
raise NotImplementedError('Quadratic Xi not yet implemented for more than one Xi')
else:
self.Xi_functions.append(32/(Ld*Ld)*conditional(x > Ld/4,
conditional(x > 3*Ld/8,
conditional(x > 5*Ld/8,
conditional(x < 3*Ld/4,
self.sigma * (x - 3*Ld/4)**2,
0.0),
(x-Ld/2)**2+Ld**2/32),
(x-Ld/4)**2),
0.0))
elif self.Xi_family == 'proper_peak':
if self.num_Xis > 1:
raise NotImplementedError('Quadratic Xi not yet implemented for more than one Xi')
else:
self.Xi_functions.append(self.sigma * 0.5*2/(exp(x-Ld/2)+exp(-x+Ld/2)))
elif self.Xi_family == 'constant':
if self.num_Xis > 1:
raise NotImplementedError('Constant Xi not yet implemented for more than one Xi')
else:
self.Xi_functions.append(self.sigma * (sin(0*pi*x/Ld)+1))
else:
raise NotImplementedError('Xi_family %s not implemented' % self.Xi_family)
# make lists of functions for xi_x, xi_xx and xi_xxx
if self.scheme in ['hydrodynamic', 'LASCH_hydrodynamic']:
self.dXi_x = prognostic_variables.dXi_x
self.dXi_xx = prognostic_variables.dXi_xx
self.Xi_x_functions = []
self.Xi_xx_functions = []
for Xi_expr in self.Xi_functions:
Xi_x_function = Function(self.dXi_x.function_space())
Xi_xx_function = Function(self.dXi_xx.function_space())
phi_x = TestFunction(self.dXi_x.function_space())
phi_xx = TestFunction(self.dXi_xx.function_space())
Xi_x_eqn = phi_x * Xi_x_function * dx + phi_x.dx(0) * Xi_expr * dx
Xi_xx_eqn = phi_xx * Xi_xx_function * dx + phi_xx.dx(0) * Xi_x_function * dx
Xi_x_problem = NonlinearVariationalProblem(Xi_x_eqn, Xi_x_function)
Xi_xx_problem = NonlinearVariationalProblem(Xi_xx_eqn, Xi_xx_function)
Xi_x_solver = NonlinearVariationalSolver(Xi_x_problem)
Xi_xx_solver = NonlinearVariationalSolver(Xi_xx_problem)
# for some reason these solvers don't work for constant Xi functions
# so just manually make the derivatives be zero
if self.Xi_family == 'constant':
Xi_x_function.interpolate(0.0*x)
Xi_xx_function.interpolate(0.0*x)
else:
Xi_x_solver.solve()
Xi_xx_solver.solve()
self.Xi_x_functions.append(Xi_x_function)
self.Xi_xx_functions.append(Xi_xx_function)
# now make a master xi
Xi_expr = 0.0*x
for dW, Xi_function, pure_xi, pure_xi_x, pure_xi_xx, pure_xi_xxx, pure_xi_xxxx in zip(self.dWs, self.Xi_functions, self.pure_xi_list, self.pure_xi_x_list, self.pure_xi_xx_list, self.pure_xi_xxx_list, self.pure_xi_xxxx_list):
Xi_expr += dW * Xi_function
if self.scheme in ['upwind', 'LASCH']:
pure_xi.interpolate(as_vector([Xi_function]))
pure_xi_x.project(as_vector([Xi_function.dx(0)]))
CG1 = FunctionSpace(mesh, "CG", 1)
psi = TestFunction(CG1)
xixx_scalar = Function(CG1)
xixx_eqn = psi * xixx_scalar * dx + psi.dx(0) * Xi_function.dx(0) * dx
prob = NonlinearVariationalProblem(xixx_eqn, xixx_scalar)
solver = NonlinearVariationalSolver(prob)
solver.solve()
pure_xi_xx.interpolate(as_vector([xixx_scalar]))
else:
pure_xi.interpolate(Xi_function)
# I guess we can't take the gradient of constants
if self.Xi_family != 'constant':
pure_xi_x.project(Xi_function.dx(0))
pure_xi_xx.project(pure_xi_x.dx(0))
pure_xi_xxx.project(pure_xi_xx.dx(0))
pure_xi_xxxx.project(pure_xi_xxx.dx(0))
if self.scheme in ['upwind', 'LASCH']:
self.dXi_interpolator = Interpolator(as_vector([Xi_expr]), self.dXi)
else:
self.dXi_interpolator = Interpolator(Xi_expr, self.dXi)
if self.scheme in ['hydrodynamic', 'LASCH_hydrodynamic']:
# initialise blank expressions
Xi_x_expr = 0.0*x
Xi_xx_expr = 0.0*x
# make full expressions by adding all dW * Xi_xs
for dW, Xi_x_function, Xi_xx_function in zip(self.dWs, self.Xi_x_functions, self.Xi_xx_functions):
Xi_x_expr += dW * Xi_x_function
Xi_xx_expr += dW * Xi_xx_function
self.dXi_x_interpolator = Interpolator(Xi_x_expr, self.dXi_x)
self.dXi_xx_interpolator = Interpolator(Xi_xx_expr, self.dXi_xx)
def solver(self):
# setup solver using lhs and rhs defined in derived class
problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.dq, bcs=self.bcs)
solver_name = self.field_name+self.__class__.__name__
return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name)
def assemble(self, eqn, func, bcs, params):
uprob = NonlinearVariationalProblem(eqn, func, bcs)
self.usolver = NonlinearVariationalSolver(uprob, solver_parameters=params)
class IncNavierStokes3D(object):
def __init__(self, mesh, nu, rho, dt=0.001, verbose=False):
self.verbose = verbose
self.mesh = mesh
self.dt = dt
self.nu = nu
self.rho = rho
self.mu = self.nu * self.rho
self.has_nullspace = False
self.forcing = Constant((0.0, 0.0, 0.0))
self.V = VectorFunctionSpace(self.mesh, "CG", 2)
self.Q = FunctionSpace(self.mesh, "CG", 1)
self.W = self.V * self.Q
self.solver_parameters = {
"mat_type": "aij",
"snes_type": "ksponly",
"ksp_type": "fgmres",
"pc_type": "asm",
"pc_asm_type": "restrict",
"pc_asm_overlap": 2,
"sub_ksp_type": "preonly",
"sub_pc_type": "ilu",
"sub_pc_factor_levels": 1,
}
if self.verbose:
self.solver_parameters["snes_monitor"] = True
self.solver_parameters["ksp_converged_reason"] = True
def setup_solver(self, up_init=None):
""" Setup the solvers
"""
self.up0 = Function(self.W)
if up_init is not None:
chk_in = checkpointing.HDF5File(up_init, file_mode='r')
chk_in.read(self.up0, "/up")
chk_in.close()
self.u0, self.p0 = split(self.up0)
self.up = Function(self.W)
if up_init is not None:
chk_in = checkpointing.HDF5File(up_init, file_mode='r')
chk_in.read(self.up, "/up")
chk_in.close()
self.u1, self.p1 = split(self.up)
self.up.sub(0).rename("velocity")
self.up.sub(1).rename("pressure")
v, q = TestFunctions(self.W)
h = CellVolume(self.mesh)
u_norm = sqrt(dot(self.u0, self.u0))
if self.has_nullspace:
nullspace = MixedVectorSpaceBasis(
self.W,
[self.W.sub(0), VectorSpaceBasis(constant=True)])
else:
nullspace = None
tau = ((2.0 / self.dt)**2 + (2.0 * u_norm / h)**2 +
(4.0 * self.nu / h**2)**2)**(-0.5)
# temporal discretization
F = (1.0 / self.dt) * inner(self.u1 - self.u0, v) * dx
# weak form
F += (+inner(dot(self.u0, nabla_grad(self.u1)), v) * dx +
self.nu * inner(grad(self.u1), grad(v)) * dx -
(1.0 / self.rho) * self.p1 * div(v) * dx +
div(self.u1) * q * dx - inner(self.forcing, v) * dx)
# residual form
R = (+(1.0 / self.dt) * (self.u1 - self.u0) +
dot(self.u0, nabla_grad(self.u1)) - self.nu * div(grad(self.u1)) +
(1.0 / self.rho) * grad(self.p1) - self.forcing)
# GLS
F += tau * inner(
+dot(self.u0, nabla_grad(v)) - self.nu * div(grad(v)) +
(1.0 / self.rho) * grad(q), R) * dx
self.problem = NonlinearVariationalProblem(F, self.up, self.bcs)
self.solver = NonlinearVariationalSolver(
self.problem,
nullspace=nullspace,
solver_parameters=self.solver_parameters)
def get_mixed_fs(self):
return self.W
def set_forcing(self, forcing):
self.forcing = forcing
def set_bcs(self, u_bcs, p_bcs):
self.bcs = list(chain.from_iterable([u_bcs, p_bcs]))
def step(self):
if self.verbose:
printp0("IncNavierStokes")
self.solver.solve()
self.up0.assign(self.up)
return self.up.split()
class StageDerivativeTimeStepper:
"""Front-end class for advancing a time-dependent PDE via a Runge-Kutta
method formulated in terms of stage derivatives.
:arg F: A :class:`ufl.Form` instance describing the semi-discrete problem
F(t, u; v) == 0, where `u` is the unknown
:class:`firedrake.Function and `v` is the
:class:firedrake.TestFunction`.
:arg butcher_tableau: A :class:`ButcherTableau` instance giving
the Runge-Kutta method to be used for time marching.
:arg t: A :class:`firedrake.Constant` instance that always
contains the time value at the beginning of a time step
:arg dt: A :class:`firedrake.Constant` containing the size of the
current time step. The user may adjust this value between
time steps.
:arg u0: A :class:`firedrake.Function` containing the current
state of the problem to be solved.
:arg bcs: An iterable of :class:`firedrake.DirichletBC` containing
the strongly-enforced boundary conditions. Irksome will
manipulate these to obtain boundary conditions for each
stage of the RK method.
:arg bc_type: How to manipulate the strongly-enforced boundary
conditions to derive the stage boundary conditions.
Should be a string, either "DAE", which implements BCs as
constraints in the style of a differential-algebraic
equation, or "ODE", which takes the time derivative of the
boundary data and evaluates this for the stage values
:arg solver_parameters: A :class:`dict` of solver parameters that
will be used in solving the algebraic problem associated
with each time step.
:arg splitting: An callable used to factor the Butcher matrix
:arg appctx: An optional :class:`dict` containing application context.
This gets included with particular things that Irksome will
pass into the nonlinear solver so that, say, user-defined preconditioners
have access to it.
:arg nullspace: A list of tuples of the form (index, VSB) where
index is an index into the function space associated with
`u` and VSB is a :class: `firedrake.VectorSpaceBasis`
instance to be passed to a
`firedrake.MixedVectorSpaceBasis` over the larger space
associated with the Runge-Kutta method
"""
def __init__(self, F, butcher_tableau, t, dt, u0, bcs=None,
solver_parameters=None, splitting=AI,
appctx=None, nullspace=None, bc_type="DAE"):
self.u0 = u0
self.t = t
self.dt = dt
self.num_fields = len(u0.function_space())
self.num_stages = len(butcher_tableau.b)
self.butcher_tableau = butcher_tableau
bigF, stages, bigBCs, bigNSP, bigBCdata = \
getForm(F, butcher_tableau, t, dt, u0, bcs, bc_type, splitting, nullspace)
self.stages = stages
self.bigBCs = bigBCs
self.bigBCdata = bigBCdata
problem = NLVP(bigF, stages, bigBCs)
appctx_irksome = {"F": F,
"butcher_tableau": butcher_tableau,
"t": t,
"dt": dt,
"u0": u0,
"bcs": bcs,
"bc_type": bc_type,
"splitting": splitting,
"nullspace": nullspace}
if appctx is None:
appctx = appctx_irksome
else:
appctx = {**appctx, **appctx_irksome}
push_parent(u0.function_space().dm, stages.function_space().dm)
self.solver = NLVS(problem,
appctx=appctx,
solver_parameters=solver_parameters,
nullspace=bigNSP)
pop_parent(u0.function_space().dm, stages.function_space().dm)
if self.num_stages == 1 and self.num_fields == 1:
self.ws = (stages,)
else:
self.ws = stages.split()
A1, A2 = splitting(butcher_tableau.A)
try:
self.updateb = numpy.linalg.solve(A2.T, butcher_tableau.b)
except numpy.linalg.LinAlgError:
raise NotImplementedError("A=A1 A2 splitting needs A2 invertible")
boo = numpy.zeros(self.updateb.shape, dtype=self.updateb.dtype)
boo[-1] = 1
if numpy.allclose(self.updateb, boo):
self._update = self._update_A2Tmb
else:
self._update = self._update_general
def _update_general(self):
"""Assuming the algebraic problem for the RK stages has been
solved, updates the solution. This will not typically be
called by an end user."""
b = self.updateb
dtc = float(self.dt)
u0 = self.u0
ns = self.num_stages
nf = self.num_fields
ws = self.ws
for s in range(ns):
for i, u0d in enumerate(u0.dat):
u0d.data[:] += dtc * b[s] * ws[nf*s+i].dat.data_ro
def _update_A2Tmb(self):
"""Assuming the algebraic problem for the RK stages has been
solved, updates the solution. This will not typically be
called by an end user. This handles the common but highly
specialized case of `w = Ak` or `A = I A` splitting where
A2^{-T} b = e_{num_stages}"""
dtc = float(self.dt)
u0 = self.u0
ns = self.num_stages
nf = self.num_fields
ws = self.ws
for i, u0d in enumerate(u0.dat):
u0d.data[:] += dtc * ws[nf*(ns-1)+i].dat.data_ro
def advance(self):
"""Advances the system from time `t` to time `t + dt`.
Note: overwrites the value `u0`."""
for gdat, gcur, gmethod in self.bigBCdata:
gmethod(gcur, self.u0)
push_parent(self.u0.function_space().dm, self.stages.function_space().dm)
self.solver.solve()
pop_parent(self.u0.function_space().dm, self.stages.function_space().dm)
self._update()
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 __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))
rho_averaged = Function(Vt)
rho_recoverer = Recoverer(rho0,
rho_averaged,
VDG=FunctionSpace(mesh,
BrokenElement(Vt.ufl_element())),
boundary_method=physics_boundary_method)
rho_recoverer.project()
exner = thermodynamics.exner_pressure(state.parameters, rho_averaged, theta0)
p = thermodynamics.p(state.parameters, exner)
T = thermodynamics.T(state.parameters, theta0, exner, r_v=w_v)
w_sat = thermodynamics.r_sat(state.parameters, T, p)
w_functional = (phi * w_v * dxp - phi * w_sat * dxp)
w_problem = NonlinearVariationalProblem(w_functional, w_v)
w_solver = NonlinearVariationalSolver(w_problem)
w_solver.solve()
water_v0.assign(w_v)
water_c0.assign(water_t - water_v0)
state.set_reference_profiles([('rho', rho_b), ('theta', theta_b),
('vapour_mixing_ratio', water_vb)])
rho_opts = None
theta_opts = EmbeddedDGOptions()
u_transport = ImplicitMidpoint(state, "u")
transported_fields = [
SSPRK3(state, "rho", options=rho_opts),
SSPRK3(state, "theta", options=theta_opts),
def compressible_hydrostatic_balance(state, theta0, rho0, pi0=None,
top=False, pi_boundary=Constant(1.0),
water_t=None,
solve_for_rho=False,
params=None):
"""
Compute a hydrostatically balanced density given a potential temperature
profile.
:arg state: The :class:`State` object.
:arg theta0: :class:`.Function`containing the potential temperature.
:arg rho0: :class:`.Function` to write the initial density into.
:arg top: If True, set a boundary condition at the top. Otherwise, set
it at the bottom.
:arg pi_boundary: a field or expression to use as boundary data for pi on
the top or bottom as specified.
:arg water_t: the initial total water mixing ratio field.
"""
# Calculate hydrostatic Pi
VDG = state.spaces("DG")
Vv = state.spaces("Vv")
W = MixedFunctionSpace((Vv, VDG))
v, pi = TrialFunctions(W)
dv, dpi = TestFunctions(W)
n = FacetNormal(state.mesh)
cp = state.parameters.cp
# add effect of density of water upon theta
theta = theta0
if water_t is not None:
theta = theta0 / (1 + water_t)
alhs = (
(cp*inner(v, dv) - cp*div(dv*theta)*pi)*dx
+ dpi*div(theta*v)*dx
)
if top:
bmeasure = ds_t
bstring = "bottom"
else:
bmeasure = ds_b
bstring = "top"
arhs = -cp*inner(dv, n)*theta*pi_boundary*bmeasure
g = state.parameters.g
arhs -= g*inner(dv, state.k)*dx
bcs = [DirichletBC(W.sub(0), 0.0, bstring)]
w = Function(W)
PiProblem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs)
if params is None:
params = {'pc_type': 'fieldsplit',
'pc_fieldsplit_type': 'schur',
'ksp_type': 'gmres',
'ksp_monitor_true_residual': True,
'ksp_max_it': 100,
'ksp_gmres_restart': 50,
'pc_fieldsplit_schur_fact_type': 'FULL',
'pc_fieldsplit_schur_precondition': 'selfp',
'fieldsplit_0_ksp_type': 'richardson',
'fieldsplit_0_ksp_max_it': 5,
'fieldsplit_0_pc_type': 'gamg',
'fieldsplit_1_pc_gamg_sym_graph': True,
'fieldsplit_1_mg_levels_ksp_type': 'chebyshev',
'fieldsplit_1_mg_levels_ksp_chebyshev_esteig': True,
'fieldsplit_1_mg_levels_ksp_max_it': 5,
'fieldsplit_1_mg_levels_pc_type': 'bjacobi',
'fieldsplit_1_mg_levels_sub_pc_type': 'ilu'}
PiSolver = LinearVariationalSolver(PiProblem,
solver_parameters=params)
PiSolver.solve()
v, Pi = w.split()
if pi0 is not None:
pi0.assign(Pi)
if solve_for_rho:
w1 = Function(W)
v, rho = w1.split()
rho.interpolate(thermodynamics.rho(state.parameters, theta0, Pi))
v, rho = split(w1)
dv, dpi = TestFunctions(W)
pi = thermodynamics.pi(state.parameters, rho, theta0)
F = (
(cp*inner(v, dv) - cp*div(dv*theta)*pi)*dx
+ dpi*div(theta0*v)*dx
+ cp*inner(dv, n)*theta*pi_boundary*bmeasure
)
F += g*inner(dv, state.k)*dx
rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs)
rhosolver = NonlinearVariationalSolver(rhoproblem, solver_parameters=params)
rhosolver.solve()
v, rho_ = w1.split()
rho0.assign(rho_)
else:
rho0.interpolate(thermodynamics.rho(state.parameters, theta0, Pi))
def moist_hydrostatic_balance(state, theta_e, water_t, pi_boundary=Constant(1.0)):
"""
Given a wet equivalent potential temperature, theta_e, and the total moisture
content, water_t, compute a hydrostatically balance virtual potential temperature,
dry density and water vapour profile.
:arg state: The :class:`State` object.
:arg theta_e: The initial wet equivalent potential temperature profile.
:arg water_t: The total water pseudo-mixing ratio profile.
:arg pi_boundary: the value of pi on the lower boundary of the domain.
"""
theta0 = state.fields('theta')
rho0 = state.fields('rho')
water_v0 = state.fields('water_v')
# Calculate hydrostatic Pi
Vt = theta0.function_space()
Vr = rho0.function_space()
Vv = state.fields('u').function_space()
n = FacetNormal(state.mesh)
g = state.parameters.g
cp = state.parameters.cp
R_d = state.parameters.R_d
p_0 = state.parameters.p_0
VDG = state.spaces("DG")
if any(deg > 2 for deg in VDG.ufl_element().degree()):
state.logger.warning("default quadrature degree most likely not sufficient for this degree element")
quadrature_degree = (5, 5)
params = {'ksp_type': 'preonly',
'ksp_monitor_true_residual': True,
'ksp_converged_reason': True,
'snes_converged_reason': True,
'ksp_max_it': 100,
'mat_type': 'aij',
'pc_type': 'lu',
'pc_factor_mat_solver_type': 'mumps'}
theta0.interpolate(theta_e)
water_v0.interpolate(water_t)
Pi = Function(Vr)
epsilon = 0.9 # relaxation constant
# set up mixed space
Z = MixedFunctionSpace((Vt, Vt))
z = Function(Z)
gamma, phi = TestFunctions(Z)
theta_v, w_v = z.split()
# give first guesses for trial functions
theta_v.assign(theta0)
w_v.assign(water_v0)
theta_v, w_v = split(z)
# define variables
T = thermodynamics.T(state.parameters, theta_v, Pi, r_v=w_v)
p = thermodynamics.p(state.parameters, Pi)
w_sat = thermodynamics.r_sat(state.parameters, T, p)
dxp = dx(degree=(quadrature_degree))
# set up weak form of theta_e and w_sat equations
F = (-gamma * theta_e * dxp
+ gamma * thermodynamics.theta_e(state.parameters, T, p, w_v, water_t) * dxp
- phi * w_v * dxp
+ phi * w_sat * dxp)
problem = NonlinearVariationalProblem(F, z)
solver = NonlinearVariationalSolver(problem, solver_parameters=params)
theta_v, w_v = z.split()
Pi_h = Function(Vr).interpolate((p / p_0) ** (R_d / cp))
# solve for Pi with theta_v and w_v constant
# then solve for theta_v and w_v with Pi constant
for i in range(5):
compressible_hydrostatic_balance(state, theta0, rho0, pi0=Pi_h, water_t=water_t)
Pi.assign(Pi * (1 - epsilon) + epsilon * Pi_h)
solver.solve()
theta0.assign(theta0 * (1 - epsilon) + epsilon * theta_v)
water_v0.assign(water_v0 * (1 - epsilon) + epsilon * w_v)
# now begin on Newton solver, setup up new mixed space
Z = MixedFunctionSpace((Vt, Vt, Vr, Vv))
z = Function(Z)
gamma, phi, psi, w = TestFunctions(Z)
theta_v, w_v, pi, v = z.split()
# use previous values as first guesses for newton solver
theta_v.assign(theta0)
w_v.assign(water_v0)
pi.assign(Pi)
theta_v, w_v, pi, v = split(z)
# define variables
T = thermodynamics.T(state.parameters, theta_v, pi, r_v=w_v)
p = thermodynamics.p(state.parameters, pi)
w_sat = thermodynamics.r_sat(state.parameters, T, p)
F = (-gamma * theta_e * dxp
+ gamma * thermodynamics.theta_e(state.parameters, T, p, w_v, water_t) * dxp
- phi * w_v * dxp
+ phi * w_sat * dxp
+ cp * inner(v, w) * dxp
- cp * div(w * theta_v / (1.0 + water_t)) * pi * dxp
+ psi * div(theta_v * v / (1.0 + water_t)) * dxp
+ cp * inner(w, n) * pi_boundary * theta_v / (1.0 + water_t) * ds_b
+ g * inner(w, state.k) * dxp)
bcs = [DirichletBC(Z.sub(3), 0.0, "top")]
problem = NonlinearVariationalProblem(F, z, bcs=bcs)
solver = NonlinearVariationalSolver(problem, solver_parameters=params)
solver.solve()
theta_v, w_v, pi, v = z.split()
# assign final values
theta0.assign(theta_v)
water_v0.assign(w_v)
# find rho
compressible_hydrostatic_balance(state, theta0, rho0, water_t=water_t, solve_for_rho=True)
class Equations(object):
"""
An object setting up equations and solvers for
the stochastic Camassa-Holm equation.
:arg prognostic_variables: a PrognosticVariables object.
:arg simulation_parameters: a dictionary storing the simulation parameters.
"""
def __init__(self, prognostic_variables, simulation_parameters):
mesh = simulation_parameters['mesh'][-1]
self.scheme = simulation_parameters['scheme'][-1]
self.timestepping = simulation_parameters['timestepping'][-1]
alphasq = simulation_parameters['alphasq'][-1]
c0 = simulation_parameters['c0'][-1]
gamma = simulation_parameters['gamma'][-1]
Dt = Constant(simulation_parameters['dt'][-1])
self.solvers = []
if alphasq.values()[0] > 0.0 and gamma.values()[0] == 0.0:
self.setup = 'ch'
if self.scheme == 'upwind' and self.timestepping == 'ssprk3':
Vm = prognostic_variables.Vm
Vu = prognostic_variables.Vu
self.m = prognostic_variables.m
self.u = prognostic_variables.u
self.Xi = prognostic_variables.dXi
self.m0 = Function(Vm).assign(self.m)
# now make problem for the actual problem
psi = TestFunction(Vm)
self.m_trial = Function(Vm)
self.dm = Function(
Vm
) # introduce this as the advection operator for a single step
us = Dt * self.u + self.Xi
nhat = FacetNormal(mesh)
un = 0.5 * (dot(us, nhat) + abs(dot(us, nhat)))
ones = Function(Vu).project(as_vector([Constant(1.)]))
Lm = (psi * self.dm * dx -
psi.dx(0) * self.m_trial * dot(ones, us) * dx +
psi * self.m_trial * dot(ones, us.dx(0)) * dx +
jump(psi) * (un('+') * self.m_trial('+') -
un('-') * self.m_trial('-')) * dS)
mprob = NonlinearVariationalProblem(Lm, self.dm)
self.msolver = NonlinearVariationalSolver(mprob,
solver_parameters={
'ksp_type':
'preonly',
'pc_type':
'bjacobi',
'sub_pc_type':
'ilu'
})
phi = TestFunction(Vu)
Lu = (dot(phi, ones) * self.m * dx - dot(phi, self.u) * dx -
alphasq * dot(self.u.dx(0), phi.dx(0)) * dx)
uprob = NonlinearVariationalProblem(Lu, self.u)
self.usolver = NonlinearVariationalSolver(uprob,
solver_parameters={
'ksp_type':
'preonly',
'pc_type': 'lu'
})
elif self.scheme == 'hydrodynamic' and self.timestepping == 'midpoint':
Vu = prognostic_variables.Vu
self.u = prognostic_variables.u
W = MixedFunctionSpace((Vu, ) * 3)
psi, phi, zeta = TestFunctions(W)
w1 = Function(W)
self.u1, dFh, dGh = split(w1)
uh = (self.u1 + self.u) / 2
dXi = prognostic_variables.dXi
dXi_x = prognostic_variables.dXi_x
dXi_xx = prognostic_variables.dXi_xx
dvh = Dt * uh + dXi
Lu = (psi * (self.u1 - self.u) * dx +
psi * uh.dx(0) * dvh * dx - psi.dx(0) * dFh * dx +
psi * dGh * dx + phi * dFh * dx +
alphasq * phi.dx(0) * dFh.dx(0) * dx -
phi * uh * uh * Dt * dx -
0.5 * alphasq * phi * uh.dx(0) * uh.dx(0) * Dt * dx +
zeta * dGh * dx + alphasq * zeta.dx(0) * dGh.dx(0) * dx -
2 * zeta * uh * dXi_x * dx -
alphasq * zeta * uh.dx(0) * dXi_xx * dx)
self.u1, dFh, dGh = w1.split()
uprob = NonlinearVariationalProblem(Lu, w1)
self.usolver = NonlinearVariationalSolver(uprob,
solver_parameters={
'mat_type':
'aij',
'ksp_type':
'preonly',
'pc_type': 'lu'
})
elif self.scheme == 'no_gradient' and self.timestepping == 'midpoint':
# a version of the hydrodynamic form but without exploiting the gradient
Vu = prognostic_variables.Vu
self.u = prognostic_variables.u
W = MixedFunctionSpace((Vu, ) * 3)
psi, phi, zeta = TestFunctions(W)
w1 = Function(W)
self.u1, dFh, dGh = split(w1)
uh = (self.u1 + self.u) / 2
dXi = prognostic_variables.dXi
dXi_x = prognostic_variables.dXi_x
dXi_xx = prognostic_variables.dXi_xx
dvh = Dt * uh + dXi
Lu = (psi * (self.u1 - self.u) * dx +
psi * uh.dx(0) * dvh * dx + psi * dFh.dx(0) * dx +
psi * dGh * dx + phi * dFh * dx +
alphasq * phi.dx(0) * dFh.dx(0) * dx -
phi * uh * uh * Dt * dx -
0.5 * alphasq * phi * uh.dx(0) * uh.dx(0) * Dt * dx +
zeta * dGh * dx + alphasq * zeta.dx(0) * dGh.dx(0) * dx -
2 * zeta * uh * dXi_x * dx -
alphasq * zeta * uh.dx(0) * dXi_xx * dx)
self.u1, dFh, dGh = w1.split()
uprob = NonlinearVariationalProblem(Lu, w1)
self.usolver = NonlinearVariationalSolver(uprob,
solver_parameters={
'mat_type':
'aij',
'ksp_type':
'preonly',
'pc_type': 'lu'
})
elif self.scheme == 'test' and self.timestepping == 'midpoint':
self.u = prognostic_variables.u
Vu = prognostic_variables.Vu
psi = TestFunction(Vu)
self.u1 = Function(Vu)
uh = (self.u1 + self.u) / 2
dvh = Dt * uh + prognostic_variables.dXi
eqn = (psi * (self.u1 - self.u) * dx -
psi * uh * dvh.dx(0) * dx)
prob = NonlinearVariationalProblem(eqn, self.u1)
self.usolver = NonlinearVariationalSolver(prob,
solver_parameters={
'mat_type':
'aij',
'ksp_type':
'preonly',
'pc_type': 'lu'
})
else:
raise ValueError(
'Scheme %s and timestepping %s either not compatible or not recognised.'
% (self.scheme, self.timestepping))
elif alphasq.values()[0] == 0.0 and gamma.values()[0] > 0.0:
self.setup = 'kdv'
if self.scheme == 'upwind' and self.timestepping == 'ssprk3':
raise NotImplementedError(
'Scheme %s and timestepping %s not yet implemented.' %
(self.scheme, self.timestepping))
elif self.scheme == 'upwind' and self.timestepping == 'midpoint':
raise NotImplementedError(
'Scheme %s and timestepping %s not yet implemented.' %
(self.scheme, self.timestepping))
elif self.scheme == 'hydrodynamic' and self.timestepping == 'midpoint':
raise NotImplementedError(
'Scheme %s and timestepping %s not yet implemented.' %
(self.scheme, self.timestepping))
else:
raise ValueError(
'Scheme %s and timestepping %s either not compatible or not recognised.'
% (self.scheme, self.timestepping))
else:
raise NotImplementedError(
'Schemes for your values of alpha squared %.3f and gamma %.3f are not yet implemented.'
% (alphasq, gamma))
def solve(self):
if self.scheme == 'upwind' and self.timestepping == 'ssprk3':
# do three step RK method for m
self.m_trial.assign(self.m0)
self.msolver.solve()
self.m_trial.assign(self.m0 + self.dm)
self.msolver.solve()
self.m_trial.assign(3. / 4 * self.m0 + 1. / 4 *
(self.m_trial + self.dm))
self.msolver.solve()
self.m.assign(1. / 3 * self.m0 + 2. / 3 * (self.m_trial + self.dm))
self.m0.assign(self.m)
# now solve inverse problem for u
self.usolver.solve()
elif self.scheme == 'upwind' and self.timestepping == 'midpoint':
self.msolver.solve()
self.m.assign(self.m_trial)
self.usolver.solve()
self.m0.assign(self.m)
elif self.scheme == 'hydrodynamic' and self.timestepping == 'midpoint':
self.usolver.solve()
self.u.assign(self.u1)
elif self.scheme == 'no_gradient' and self.timestepping == 'midpoint':
self.usolver.solve()
self.u.assign(self.u1)
elif self.scheme == 'test' and self.timestepping == 'midpoint':
self.usolver.solve()
self.u.assign(self.u1)
else:
raise ValueError(
'Scheme %s and timestepping %s either not compatible or not recognised.'
% (self.scheme, self.timestepping))
class StageValueTimeStepper:
def __init__(self,
F,
butcher_tableau,
t,
dt,
u0,
bcs=None,
solver_parameters=None,
update_solver_parameters=None,
splitting=AI,
nullspace=None,
appctx=None):
self.u0 = u0
self.t = t
self.dt = dt
self.num_fields = len(u0.function_space())
self.num_stages = len(butcher_tableau.b)
self.butcher_tableau = butcher_tableau
Fbig, update_stuff, UU, bigBCs, gblah, nsp = getFormStage(
F, butcher_tableau, u0, t, dt, bcs, splitting=splitting)
self.UU = UU
self.bigBCs = bigBCs
self.bcdat = gblah
self.update_stuff = update_stuff
self.prob = NonlinearVariationalProblem(Fbig, UU, bigBCs)
appctx_irksome = {
"F": F,
"butcher_tableau": butcher_tableau,
"t": t,
"dt": dt,
"u0": u0,
"bcs": bcs,
"nullspace": nullspace
}
if appctx is None:
appctx = appctx_irksome
else:
appctx = {**appctx, **appctx_irksome}
self.solver = NonlinearVariationalSolver(
self.prob,
appctx=appctx,
nullspace=nsp,
solver_parameters=solver_parameters)
unew, Fupdate, update_bcs, update_bcs_gblah = self.update_stuff
self.update_problem = NonlinearVariationalProblem(
Fupdate, unew, update_bcs)
self.update_solver = NonlinearVariationalSolver(
self.update_problem, solver_parameters=update_solver_parameters)
self._update = self._update_general
# Unused for now since null spaces don't seem to work with it.
def _update_riia(self):
u0 = self.u0
UUs = self.UU.split()
nf = self.num_fields
ns = self.num_stages
for i, u0d in enumerate(u0.dat):
u0d.data[:] = UUs[nf * (ns - 1) + i].dat.data_ro[:]
def _update_general(self):
(unew, Fupdate, update_bcs, update_bcs_gblah) = self.update_stuff
for gdat, gcur, gmethod in update_bcs_gblah:
gmethod(gdat, gcur)
self.update_solver.solve()
for u0d, und in zip(self.u0.dat, unew.dat):
u0d.data[:] = und.data_ro[:]
def advance(self):
for gdat, gcur, gmethod in self.bcdat:
gmethod(gdat, gcur)
self.solver.solve()
self._update()
class TimeStepper:
"""Front-end class for advancing a time-dependent PDE via a Runge-Kutta
method.
:arg F: A :class:`ufl.Form` instance describing the semi-discrete problem
F(t, u; v) == 0, where `u` is the unknown
:class:`firedrake.Function and `v` is the
:class:firedrake.TestFunction`.
:arg butcher_tableau: A :class:`ButcherTableau` instance giving
the Runge-Kutta method to be used for time marching.
:arg t: A :class:`firedrake.Constant` instance that always
contains the time value at the beginning of a time step
:arg dt: A :class:`firedrake.Constant` containing the size of the
current time step. The user may adjust this value between
time steps, but see :class:`AdaptiveTimeStepper` for a
method that attempts to do this automatically.
:arg u0: A :class:`firedrake.Function` containing the current
state of the problem to be solved.
:arg bcs: An iterable of :class:`firedrake.DirichletBC` containing
the strongly-enforced boundary conditions. Irksome will
manipulate these to obtain boundary conditions for each
stage of the RK method.
:arg solver_parameters: A :class:`dict` of solver parameters that
will be used in solving the algebraic problem associated
with each time step.
"""
def __init__(self,
F,
butcher_tableau,
t,
dt,
u0,
bcs=None,
solver_parameters=None):
self.u0 = u0
self.t = t
self.dt = dt
self.num_fields = len(u0.function_space())
self.num_stages = len(butcher_tableau.b)
self.butcher_tableau = butcher_tableau
bigF, stages, bigBCs, bigBCdata = \
getForm(F, butcher_tableau, t, dt, u0, bcs)
self.stages = stages
self.bigBCs = bigBCs
self.bigBCdata = bigBCdata
problem = NLVP(bigF, stages, bigBCs)
self.solver = NLVS(problem, solver_parameters=solver_parameters)
self.ks = stages.split()
def _update(self):
"""Assuming the algebraic problem for the RK stages has been
solved, updates the solution. This will not typically be
called by an end user."""
b = self.butcher_tableau.b
dtc = float(self.dt)
u0 = self.u0
ns = self.num_stages
nf = self.num_fields
if nf == 1:
ks = self.ks
for i in range(ns):
u0 += dtc * b[i] * ks[i]
else:
k = self.stages
for s in range(ns):
for i in range(nf):
u0.dat.data[i][:] += dtc * b[s] * k.dat.data[nf * s + i][:]
def advance(self):
"""Advances the system from time `t` to time `t + dt`.
Note: overwrites the value `u0`."""
for gdat, gcur in self.bigBCdata:
gdat.interpolate(gcur)
self.solver.solve()
self._update()
def compressible_hydrostatic_balance(state, theta0, rho0, pi0=None,
top=False, pi_boundary=Constant(1.0),
solve_for_rho=False,
params=None):
"""
Compute a hydrostatically balanced density given a potential temperature
profile.
:arg state: The :class:`State` object.
:arg theta0: :class:`.Function`containing the potential temperature.
:arg rho0: :class:`.Function` to write the initial density into.
:arg top: If True, set a boundary condition at the top. Otherwise, set
it at the bottom.
:arg pi_boundary: a field or expression to use as boundary data for pi on
the top or bottom as specified.
"""
# Calculate hydrostatic Pi
W = MixedFunctionSpace((state.Vv,state.V[1]))
v, pi = TrialFunctions(W)
dv, dpi = TestFunctions(W)
n = FacetNormal(state.mesh)
cp = state.parameters.cp
alhs = (
(cp*inner(v,dv) - cp*div(dv*theta0)*pi)*dx
+ dpi*div(theta0*v)*dx
)
if top:
bmeasure = ds_t
bstring = "bottom"
else:
bmeasure = ds_b
bstring = "top"
arhs = -cp*inner(dv,n)*theta0*pi_boundary*bmeasure
if state.parameters.geopotential:
Phi = state.Phi
arhs += div(dv)*Phi*dx - inner(dv,n)*Phi*bmeasure
else:
g = state.parameters.g
arhs -= g*inner(dv,state.k)*dx
if(state.mesh.geometric_dimension() == 2):
bcs = [DirichletBC(W.sub(0), Expression(("0.", "0.")), bstring)]
elif(state.mesh.geometric_dimension() == 3):
bcs = [DirichletBC(W.sub(0), Expression(("0.", "0.", "0.")), bstring)]
w = Function(W)
PiProblem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs)
if(params is None):
params = {'pc_type': 'fieldsplit',
'pc_fieldsplit_type': 'schur',
'ksp_type': 'gmres',
'ksp_monitor_true_residual': True,
'ksp_max_it': 100,
'ksp_gmres_restart': 50,
'pc_fieldsplit_schur_fact_type': 'FULL',
'pc_fieldsplit_schur_precondition': 'selfp',
'fieldsplit_0_ksp_type': 'richardson',
'fieldsplit_0_ksp_max_it': 5,
'fieldsplit_0_pc_type': 'gamg',
'fieldsplit_1_pc_gamg_sym_graph': True,
'fieldsplit_1_mg_levels_ksp_type': 'chebyshev',
'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues': True,
'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues_random': True,
'fieldsplit_1_mg_levels_ksp_max_it': 5,
'fieldsplit_1_mg_levels_pc_type': 'bjacobi',
'fieldsplit_1_mg_levels_sub_pc_type': 'ilu'}
PiSolver = LinearVariationalSolver(PiProblem,
solver_parameters=params)
PiSolver.solve()
v, Pi = w.split()
if pi0 is not None:
pi0.assign(Pi)
kappa = state.parameters.kappa
R_d = state.parameters.R_d
p_0 = state.parameters.p_0
if solve_for_rho:
w1 = Function(W)
v, rho = w1.split()
rho.interpolate(p_0*(Pi**((1-kappa)/kappa))/R_d/theta0)
v, rho = split(w1)
dv, dpi = TestFunctions(W)
pi = ((R_d/p_0)*rho*theta0)**(kappa/(1.-kappa))
F = (
(cp*inner(v,dv) - cp*div(dv*theta0)*pi)*dx
+ dpi*div(theta0*v)*dx
+ cp*inner(dv,n)*theta0*pi_boundary*bmeasure
)
if state.parameters.geopotential:
F += - div(dv)*Phi*dx + inner(dv,n)*Phi*bmeasure
else:
F += g*inner(dv,state.k)*dx
rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs)
rhosolver = NonlinearVariationalSolver(rhoproblem, solver_parameters=params)
rhosolver.solve()
v, rho_ = w1.split()
rho0.assign(rho_)
else:
rho0.interpolate(p_0*(Pi**((1-kappa)/kappa))/R_d/theta0)
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)
compressible_hydrostatic_balance(state, theta_b, rho_b)
W = MixedFunctionSpace((state.Vv,state.V[1]))
w1 = Function(W)
v, rho = w1.split()
rho.assign(rho_b)
v, rho = split(w1)
dv, dpi = TestFunctions(W)
pi = ((R_d/p_0)*rho*theta_b)**(kappa/(1.-kappa))
F = (
(c_p*inner(v,dv) - c_p*div(dv*theta_b)*pi)*dx
+ dpi*div(theta_b*v)*dx
+ g*inner(dv,k)*dx
+ c_p*inner(dv,n)*theta_b*ds_b # bottom surface value pi = 1.
)
rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs)
rhosolver = NonlinearVariationalSolver(rhoproblem, solver_parameters=params)
rhosolver.solve()
v, rho = w1.split()
rho_b.interpolate(rho)
W_DG1 = FunctionSpace(mesh, "DG", 1)
x = SpatialCoordinate(mesh)
theta_pert = Function(state.V[2]).interpolate(Expression("sqrt(pow(x[0]-xc,2)+pow(x[1]-zc,2)) > rc ? 0.0 : 0.25*(1. + cos((pi/rc)*(sqrt(pow((x[0]-xc),2)+pow((x[1]-zc),2)))))", xc=500., zc=350., rc=250.))
theta0.interpolate(theta_b + theta_pert)
rho0.interpolate(rho_b)
state.initialise([u0, rho0, theta0])
state.set_reference_profiles(rho_b, theta_b)
state.output.meanfields = {'rho':state.rhobar, 'theta':state.thetabar}
def run(steady=False):
"""
solve CdT/dt = S + div(k*grad(T))
=> C*v*(dT/dt)/k*dx - S*v/k*dx + grad(v)*grad(T)*dx = v*dot(grad(T), n)*ds
"""
steps = 250
dt = 1e-10
timescale = (0, steps * dt)
if steady:
print('Running steady state.')
else:
print(f'Running with time step {dt:.2g}s on time interval: '
f'{timescale[0]:.2g}s - {timescale[1]:.2g}s')
dt_invc = Constant(1 / dt)
extent = [40e-6, 40e-6, 40e-6]
mesh = BoxMesh(20, 20, 20, *extent)
V = FunctionSpace(mesh, 'CG', 1)
print(V.dim())
T = Function(V) # temperature at time i+1 (electron for now)
T_ = Function(V) # temperature at time i
v = TestFunction(V) # test function
S = create_S(mesh, V, extent)
C = create_heat_capacity(mesh, V, extent)
k = create_conductivity(mesh, V, T)
set_initial_value(mesh, T_, extent)
# Mass matrix section
M = C * T * dt_invc * v * dx
M_ = C * T_ * dt_invc * v * dx
# Stiffness matrix section
A = k * dot(grad(T), grad(v)) * dx
# function section
f = S * v * dx
# boundaries
bcs, R, b = create_dirichlet_bounds(mesh,
V,
T,
v,
k,
g=100,
boundary=[1, 2, 3, 4, 5, 6])
# bcs += create_dirichlet_bounds(mesh, V, T, v, k, 500, [6])[0]
# bcs, R, b = create_robin_bounds(mesh, T, v, k, 1e8/(100), 1e8)
if steady:
steps = 1
a = A + R
L = f + b
else:
a = M + A + R
L = M_ + f + b
prob = NonlinearVariationalProblem(a - L, T, bcs=bcs)
solver = NonlinearVariationalSolver(prob, solver_parameters=SOLVE_PARAMS)
T.assign(T_)
timestamp = datetime.now().strftime("%d-%b-%Y-%H-%M-%S")
outfile = File(f'{timestamp}/first_output.pvd')
outfile.write(T_, target_degree=1, target_continuity=H1)
last_perc = 0
for i in range(steps):
solver.solve()
perc = int(100 * (i + 1) / steps)
if perc > last_perc:
print(f'{perc}%')
last_perc = perc
T_.assign(T)
outfile.write(T_, target_degree=1, target_continuity=H1)
