def setup_DGadvection(element, vector=False):
refinements = 3 # number of horizontal cells = 20*(4^refinements)
R = 1.0
dt = pi / 3 * 0.01
mesh = IcosahedralSphereMesh(radius=R, refinement_level=refinements)
global_normal = Expression(("x[0]", "x[1]", "x[2]"))
mesh.init_cell_orientations(global_normal)
fieldlist = ["u", "D"]
timestepping = TimesteppingParameters(dt=dt)
state = ShallowWaterState(
mesh, vertical_degree=None, horizontal_degree=2, family="BDM", timestepping=timestepping, fieldlist=fieldlist
)
# interpolate initial conditions
u0 = Function(state.V[0], name="velocity")
x = SpatialCoordinate(mesh)
uexpr = as_vector([-x[1], x[0], 0.0])
u0.project(uexpr)
if vector:
if element is "BDM":
BrokenSpace = FunctionSpace(mesh, element, 2)
else:
BrokenSpace = VectorFunctionSpace(mesh, "DG", 1)
Space = VectorFunctionSpace(mesh, "CG", 1)
f = Function(Space, name="f")
fexpr = Expression(("exp(-pow(x[2],2) - pow(x[1],2))", "0.0", "0.0"))
f_end = Function(Space)
f_end_expr = Expression(("exp(-pow(x[2],2) - pow(x[0],2))", "0", "0"))
else:
if element is "BDM":
BrokenSpace = FunctionSpace(mesh, element, 2)
else:
BrokenSpace = FunctionSpace(mesh, "DG", 1)
Space = FunctionSpace(mesh, "CG", 1)
f = Function(Space, name="f")
fexpr = Expression("exp(-pow(x[2],2) - pow(x[1],2))")
f_end = Function(Space)
f_end_expr = Expression("exp(-pow(x[2],2) - pow(x[0],2))")
f.interpolate(fexpr)
f_end.interpolate(f_end_expr)
return state, BrokenSpace, u0, f, f_end
def setup_SUPGadvection(direction):
nlayers = 25 # horizontal layers
columns = 25 # number of columns
L = 1.0
m = PeriodicIntervalMesh(columns, L)
H = 1.0 # Height position of the model top
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers)
# Space for initialising velocity
W_VectorCG1 = VectorFunctionSpace(mesh, "CG", 1)
W_CG1 = FunctionSpace(mesh, "CG", 1)
# vertical coordinate and normal
z = Function(W_CG1).interpolate(Expression("x[1]"))
k = Function(W_VectorCG1).interpolate(Expression(("0.","1.")))
fieldlist = ['u','rho', 'theta']
timestepping = TimesteppingParameters(dt=0.01)
parameters = CompressibleParameters()
state = CompressibleState(mesh, vertical_degree=1, horizontal_degree=1,
family="CG",
z=z, k=k,
timestepping=timestepping,
parameters=parameters,
fieldlist=fieldlist)
# interpolate initial conditions
u0 = Function(state.V[0], name="velocity")
uexpr = as_vector([1.0, 0.0])
u0.project(uexpr)
if len(direction) == 0:
space = W_CG1
else:
space = state.V[2]
f = Function(space, name='f')
x = SpatialCoordinate(mesh)
f_expr = sin(2*pi*x[0])*sin(2*pi*x[1])
f.interpolate(f_expr)
f_end = Function(space)
f_end_expr = sin(2*pi*(x[0]-0.5))*sin(2*pi*x[1])
f_end.interpolate(f_end_expr)
return state, u0, f, f_end
def test_mass_transport_solver_convergence(solver_type):
Lx, Ly = 1.0, 1.0
u0 = 1.0
h_in, dh = 1.0, 0.2
delta_x, error = [], []
mass_transport = solver_type()
for N in range(24, 97, 4):
delta_x.append(Lx / N)
mesh = firedrake.RectangleMesh(N, N, Lx, Ly)
x, y = firedrake.SpatialCoordinate(mesh)
degree = 1
V = firedrake.VectorFunctionSpace(mesh, family='CG', degree=degree)
Q = firedrake.FunctionSpace(mesh, family='CG', degree=degree)
h0 = interpolate(h_in - dh * x / Lx, Q)
a = firedrake.Function(Q)
u = interpolate(firedrake.as_vector((u0, 0)), V)
T = 0.5
δx = 1.0 / N
δt = δx / u0
num_timesteps = int(T / δt)
h = h0.copy(deepcopy=True)
for step in range(num_timesteps):
h = mass_transport.solve(δt, h0=h, a=a, u=u, h_inflow=h0)
z = x - u0 * num_timesteps * δt
h_exact = interpolate(h_in - dh / Lx * firedrake.max_value(0, z), Q)
error.append(norm(h - h_exact) / norm(h_exact))
print(delta_x[-1], error[-1])
log_delta_x = np.log2(np.array(delta_x))
log_error = np.log2(np.array(error))
slope, intercept = np.polyfit(log_delta_x, log_error, 1)
print('log(error) ~= {:g} * log(dx) + {:g}'.format(slope, intercept))
assert slope > degree - 0.1
def test_interpolating_vector_field():
n = 32
array_vx = np.array([[(i + j) / n for j in range(n + 1)] for i in range(n + 1)])
missing = -9999.0
array_vx[0, 0] = missing
array_vx = np.flipud(array_vx)
array_vy = np.array([[(j - i) / n for j in range(n + 1)] for i in range(n + 1)])
array_vy[-1, -1] = -9999.0
array_vy = np.flipud(array_vy)
vx = make_rio_dataset(array_vx, missing)
vy = make_rio_dataset(array_vy, missing)
mesh = make_domain(48, 48, xmin=1 / 4, ymin=1 / 4, width=1 / 2, height=1 / 2)
x, y = firedrake.SpatialCoordinate(mesh)
V = firedrake.VectorFunctionSpace(mesh, "CG", 1)
u = firedrake.interpolate(firedrake.as_vector((x + y, x - y)), V)
v = icepack.interpolate((vx, vy), V)
assert firedrake.norm(u - v) / firedrake.norm(u) < 1e-10
def setup(self, state):
super(GeostrophicImbalance, self).setup(state)
u = state.fields("u")
b = state.fields("b")
p = state.fields("p")
f = state.parameters.f
Vu = u.function_space()
v = TrialFunction(Vu)
w = TestFunction(Vu)
a = inner(w, v) * dx
L = (div(w) * p + inner(w, as_vector([f * u[1], 0.0, b]))) * dx
bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")]
self.imbalance = Function(Vu)
imbalanceproblem = LinearVariationalProblem(a,
L,
self.imbalance,
bcs=bcs)
self.imbalance_solver = LinearVariationalSolver(
imbalanceproblem, solver_parameters={'ksp_type': 'cg'})
def __init__(self, state):
super(Fallout, self).__init__(state)
self.state = state
self.rain = state.fields('rain')
# function spaces
Vt = self.rain.function_space()
Vu = state.fields('u').function_space()
# introduce sedimentation rate
# for now assume all rain falls at terminal velocity
terminal_velocity = 10 # in m/s
self.v = state.fields("rainfall_velocity", Vu)
self.v.project(as_vector([0, -terminal_velocity]))
# sedimentation will happen using a full advection method
advection_equation = EmbeddedDGAdvection(state,
Vt,
equation_form="advective",
outflow=True)
self.advection_method = SSPRK3(state, self.rain, advection_equation)
def test_mass_transport_solver_convergence():
Lx, Ly = 1.0, 1.0
u0 = 1.0
h_in, dh = 1.0, 0.2
delta_x, error = [], []
mass_transport = MassTransport()
for N in range(16, 97, 4):
delta_x.append(Lx / N)
mesh = firedrake.RectangleMesh(N, N, Lx, Ly)
x, y = firedrake.SpatialCoordinate(mesh)
degree = 1
V = firedrake.VectorFunctionSpace(mesh, family='CG', degree=degree)
Q = firedrake.FunctionSpace(mesh, family='CG', degree=degree)
h = interpolate(h_in - dh * x / Lx, Q)
a = firedrake.Function(Q)
u = interpolate(firedrake.as_vector((u0, 0)), V)
T = 0.5
num_timesteps = int(0.5 * N * u0 * T / Lx)
dt = T / num_timesteps
for k in range(num_timesteps):
h = mass_transport.solve(dt, h0=h, a=a, u=u)
z = x - u0 * T
h_exact = interpolate(h_in - dh / Lx * firedrake.max_value(0, z), Q)
error.append(norm(h - h_exact) / norm(h_exact))
print(delta_x[-1], error[-1])
log_delta_x = np.log2(np.array(delta_x))
log_error = np.log2(np.array(error))
slope, intercept = np.polyfit(log_delta_x, log_error, 1)
assert slope > degree - 0.05
print(slope, intercept)
def test_strain_heating():
E_initial = firedrake.interpolate(E_surface + q_bed / α * h * (1 - ζ), Q)
u0 = 100.0
du = 100.0
u_expr = as_vector((u0 + du * x / Lx, 0))
u = firedrake.interpolate(u_expr, V)
w = firedrake.interpolate((-du / Lx + dh / Lx / h * u[0]) * ζ, W)
E_q = E_initial.copy(deepcopy=True)
E_0 = E_initial.copy(deepcopy=True)
model = icepack.models.HeatTransport3D()
from icepack.models.hybrid import (horizontal_strain, vertical_strain,
stresses)
T = model.temperature(E_q)
A = icepack.rate_factor(T)
ε_x, ε_z = horizontal_strain(u, s, h), vertical_strain(u, h)
τ_x, τ_z = stresses(ε_x, ε_z, A)
q = inner(τ_x, ε_x) + inner(τ_z, ε_z)
kwargs = {
'u': u,
'w': w,
'h': h,
's': s,
'q_bed': q_bed,
'E_inflow': E_initial,
'E_surface': Constant(E_surface)
}
dt = 10.0
final_time = Lx / u0
num_steps = int(final_time / dt) + 1
for step in range(num_steps):
E_q.assign(model.solve(dt, E0=E_q, q=q, **kwargs))
E_0.assign(model.solve(dt, E0=E_0, q=Constant(0), **kwargs))
assert assemble(E_q * h * dx) > assemble(E_0 * h * dx)
def initialise_fields(state):
L = 1.e5
H = 1.0e4 # Height position of the model top
# Initial conditions
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vt = theta0.function_space()
Vr = rho0.function_space()
# Thermodynamic constants required for setting initial conditions
# and reference profiles
g = state.parameters.g
N = state.parameters.N
# N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz)
x, z = SpatialCoordinate(state.mesh)
Tsurf = 300.
thetab = Tsurf * exp(N**2 * z / g)
theta_b = Function(Vt).interpolate(thetab)
rho_b = Function(Vr)
# Calculate hydrostatic exner
compressible_hydrostatic_balance(state, theta_b, rho_b)
a = 5.0e3
deltaTheta = 1.0e-2
theta_pert = deltaTheta * sin(pi * z / H) / (1 + (x - L / 2)**2 / a**2)
theta0.interpolate(theta_b + theta_pert)
rho0.assign(rho_b)
u0.project(as_vector([20.0, 0.0]))
state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
def f_g_scattered_plane_wave(self, d):
"""Sets f and g to correspond to the scattering of a plane wave
by a compactly-supported heterogeneous region.
Parameters - d - list of the length of the spatial dimension;
the direction in which the plane wave propagates.
"""
d = fd.as_vector(d)
x = fd.SpatialCoordinate(self.V.mesh())
# Incident wave
u_I = fd.exp(1j * self._k * fd.dot(x, d))
identity = fd.as_matrix([[1.0, 0.0], [0.0, 1.0]])
f = fd.div(fd.dot((identity-self._A),fd.grad(u_I)))\
+ self._k**2.0 * fd.inner((1.0-self._n), u_I)
self.set_f(f)
self.set_g(0.0)
def tracer_slice(tmpdir, degree):
n = 30 if degree == 0 else 15
m = PeriodicIntervalMesh(n, 1.)
mesh = ExtrudedMesh(m, layers=n, layer_height=1. / n)
# Parameters chosen so that dt != 1 and u != 1
# Gaussian is translated by 1.5 times width of domain to demonstrate
# that transport is working correctly
dt = 0.01
tmax = 0.75
output = OutputParameters(dirname=str(tmpdir), dumpfreq=25)
state = State(mesh, dt=dt, output=output)
uexpr = as_vector([2.0, 0.0])
x = SpatialCoordinate(mesh)
width = 1. / 10.
f0 = 0.5
fmax = 2.0
xc_init = 0.25
xc_end = 0.75
r_init = sqrt((x[0] - xc_init)**2 + (x[1] - 0.5)**2)
r_end = sqrt((x[0] - xc_end)**2 + (x[1] - 0.5)**2)
f_init = f0 + (fmax - f0) * exp(-(r_init / width)**2)
f_end = f0 + (fmax - f0) * exp(-(r_end / width)**2)
tol = 0.12
return TracerSetup(state,
tmax,
f_init,
f_end,
"CG",
degree,
uexpr,
tol=tol)
def test_diagnostic_solver_side_friction():
ice_shelf = icepack.models.IceShelf()
opts = {'dirichlet_ids': [1], 'side_wall_ids': [3, 4], 'tol': 1e-12}
mesh = firedrake.RectangleMesh(32, 32, Lx, Ly)
degree = 2
V = firedrake.VectorFunctionSpace(mesh, 'CG', degree)
Q = firedrake.FunctionSpace(mesh, 'CG', degree)
x, y = firedrake.SpatialCoordinate(mesh)
u_initial = interpolate(as_vector((exact_u(x), 0)), V)
h = interpolate(h0 - dh * x / Lx, Q)
A = interpolate(firedrake.Constant(icepack.rate_factor(T)), Q)
# Choose the side wall friction coefficient so that, assuming the ice is
# sliding at the maximum speed for the solution without friction, the
# stress is 10 kPa.
from icepack.constants import weertman_sliding_law as m
τ = 0.01
u_max = icepack.norm(u_initial, norm_type='Linfty')
Cs = firedrake.Constant(τ * u_max**(-1 / m))
u = ice_shelf.diagnostic_solve(u0=u_initial, h=h, A=A, Cs=Cs, **opts)
assert icepack.norm(u) < icepack.norm(u_initial)
def static_solver(self):
def epsilon(u):
return 0.5 * (fd.nabla_grad(u) + fd.nabla_grad(u).T)
#return sym(nabla_grad(u))
def sigma(u):
d = u.geometric_dimension() # space dimension
return self.lam * fd.nabla_div(u) * fd.Identity(
d) + 2 * self.mu * epsilon(u)
# Define variational problem
u = fd.TrialFunction(self.V)
v = fd.TestFunction(self.V)
# f = fd.Constant((0, 0, -self.g)) # body force / rho
f = fd.Constant((0, 0, 0)) # body force / rho
T = self.surface_force()
a = fd.inner(sigma(u), epsilon(v)) * fd.dx
L = fd.dot(f, v) * fd.dx + fd.dot(T, v) * fd.ds(1)
# Compute solution
# old: self.DBC = fd.DirichletBC( self.V, fd.Expression([0.,0.,0.]), self.bottom_id)
self.DBC = fd.DirichletBC(self.V, fd.as_vector([0., 0., 0.]),
self.bottom_id)
fd.solve(a == L, self.X, bcs=self.DBC)
def test_plot_extruded_field():
nx, ny = 32, 32
mesh2d = firedrake.UnitSquareMesh(nx, ny)
mesh3d = firedrake.ExtrudedMesh(mesh2d, layers=1)
x, y, z = firedrake.SpatialCoordinate(mesh3d)
Q = firedrake.FunctionSpace(mesh3d,
family='CG',
degree=2,
vfamily='GL',
vdegree=4)
q = interpolate((x**2 - y**2) * (1 - z**4), Q)
q_contours = icepack.plot.tricontourf(q)
assert q_contours is not None
V = firedrake.VectorFunctionSpace(mesh3d,
dim=2,
family='CG',
degree=2,
vfamily='GL',
vdegree=4)
u = interpolate(as_vector((1 - z**4, 0)), V)
u_contours = icepack.plot.tricontourf(u)
assert u_contours is not None
def _build_forcing_solvers(self):
super(CompressibleEadyForcing, self)._build_forcing_solvers()
# theta_forcing
dthetady = self.state.parameters.dthetady
Vt = self.state.spaces("HDiv_v")
F = TrialFunction(Vt)
gamma = TestFunction(Vt)
self.thetaF = Function(Vt)
u0, _, _ = split(self.x0)
a = gamma * F * dx
L = -self.scaling * gamma * (dthetady *
inner(u0, as_vector([0., 1., 0.]))) * dx
theta_forcing_problem = LinearVariationalProblem(a, L, self.thetaF)
solver_parameters = {}
if logger.isEnabledFor(DEBUG):
solver_parameters["ksp_monitor_true_residual"] = None
self.theta_forcing_solver = LinearVariationalSolver(
theta_forcing_problem,
solver_parameters=solver_parameters,
options_prefix="ThetaForcingSolver")
N = parameters.N
bref = x[1]*(N**2)
# interpolate the expression to the function
b_b = Function(state.V[2]).interpolate(bref)
# setup constants
a = Constant(5.0e3)
deltab = Constant(1.0e-2)
H = Constant(H)
L = Constant(L)
b_pert = deltab*sin(np.pi*x[1]/H)/(1 + (x[0] - L/2)**2/a**2)
# interpolate the expression to the function
b0.interpolate(b_b + b_pert)
# interpolate velocity to vector valued function space
uinit = Function(W_VectorCG1).interpolate(as_vector([20.0,0.0]))
# project to the function space we actually want to use
# this step is purely because it is not yet possible to interpolate to the
# vector function spaces we require for the compatible finite element
# methods that we use
u0.project(uinit)
# pass these initial conditions to the state.initialise method
state.initialise([u0, p0, b0])
# set the background buoyancy
state.set_reference_profiles(b_b)
# we want to output the perturbation buoyancy, so tell the dump method
# which background field to subtract
state.output.meanfields = {'b':state.bbar}
##############################################################################
from gusto import *
from firedrake import Expression, FunctionSpace, as_vector,\
VectorFunctionSpace, PeriodicIntervalMesh, ExtrudedMesh, Constant, SpatialCoordinate, exp
nlayers = 70 # horizontal layers
columns = 180 # number of columns
L = 144000.
m = PeriodicIntervalMesh(columns, L)
# build volume mesh
H = 35000. # Height position of the model top
ext_mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers)
Vc = VectorFunctionSpace(ext_mesh, "DG", 2)
coord = SpatialCoordinate(ext_mesh)
x = Function(Vc).interpolate(as_vector([coord[0],coord[1]]))
H = Constant(H)
a = Constant(1000.)
xc = Constant(L/2.)
new_coords = Function(Vc).interpolate(as_vector([x[0], x[1]+(H-x[1])*a**2/(H*((x[0]-xc)**2+a**2))]))
mesh = Mesh(new_coords)
# Space for initialising velocity
W_VectorCG = VectorFunctionSpace(mesh, "CG", 2)
W_CG = FunctionSpace(mesh, "CG", 2)
W_DG = FunctionSpace(mesh, "DG", 2)
# vertical coordinate and normal
z = Function(W_CG).interpolate(Expression("x[1]"))
k = Function(W_VectorCG).interpolate(Expression(("0.","1.")))
dt = 5.0
def setup_condens(dirname):
# declare grid shape, with length L and height H
L = 1000.
H = 1000.
nlayers = int(H / 100.)
ncolumns = int(L / 100.)
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
x = SpatialCoordinate(mesh)
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=1.0, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname + "/condens",
dumpfreq=1,
dumplist=['u'],
perturbation_fields=['theta', 'rho'])
parameters = CompressibleParameters()
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist,
diagnostic_fields=[Sum('water_v', 'water_c')])
# declare initial fields
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vpsi = FunctionSpace(mesh, "CG", 2)
Vt = theta0.function_space()
Vr = rho0.function_space()
# make a gradperp
gradperp = lambda u: as_vector([-u.dx(1), u.dx(0)])
# declare tracer field and a background field
water_v0 = state.fields("water_v", Vt)
water_c0 = state.fields("water_c", Vt)
# Isentropic background state
Tsurf = Constant(300.)
theta_b = Function(Vt).interpolate(Tsurf)
rho_b = Function(Vr)
# Calculate initial rho
compressible_hydrostatic_balance(state, theta_b, rho_b, solve_for_rho=True)
# set up water_v
xc = 500.
zc = 350.
rc = 250.
r = sqrt((x[0] - xc)**2 + (x[1] - zc)**2)
w_expr = conditional(r > rc, 0., 0.25 * (1. + cos((pi / rc) * r)))
# set up velocity field
u_max = 10.0
psi_expr = ((-u_max * L / pi) * sin(2 * pi * x[0] / L) *
sin(pi * x[1] / L))
psi0 = Function(Vpsi).interpolate(psi_expr)
u0.project(gradperp(psi0))
theta0.interpolate(theta_b)
rho0.interpolate(rho_b)
water_v0.interpolate(w_expr)
state.initialise([('u', u0), ('rho', rho0), ('theta', theta0),
('water_v', water_v0), ('water_c', water_c0)])
state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
# set up advection schemes
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
thetaeqn = SUPGAdvection(state,
Vt,
supg_params={"dg_direction": "horizontal"},
equation_form="advective")
# build advection dictionary
advected_fields = []
advected_fields.append(("u", NoAdvection(state, u0, None)))
advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn)))
advected_fields.append(("theta", SSPRK3(state, theta0, thetaeqn)))
advected_fields.append(("water_v", SSPRK3(state, water_v0, thetaeqn)))
advected_fields.append(("water_c", SSPRK3(state, water_c0, thetaeqn)))
physics_list = [Condensation(state)]
# build time stepper
stepper = AdvectionDiffusion(state,
advected_fields,
physics_list=physics_list)
return stepper, 5.0
coords[:, 2], Delta=Delta, data_array=kx_array)
Ky.dat.data[...] = coords2ijk(coords[:, 0], coords[:, 1],
coords[:, 2], Delta=Delta, data_array=ky_array)
Kz.dat.data[...] = coords2ijk(coords[:, 0], coords[:, 1],
coords[:, 2], Delta=Delta, data_array=kz_array)
print("END: Read in reservoir fields")
# Permeability field harmonic interpolation to facets
Kx_facet = fd.conditional(fd.gt(fd.avg(Kx), 0.0), Kx('+')*Kx('-') / fd.avg(Kx), 0.0)
Ky_facet = fd.conditional(fd.gt(fd.avg(Ky), 0.0), Ky('+')*Ky('-') / fd.avg(Ky), 0.0)
Kz_facet = fd.conditional(fd.gt(fd.avg(Kz), 0.0), Kz('+')*Kz('-') / fd.avg(Kz), 0.0)
# We can now define the bilinear and linear forms for the left and right
dx = fd.dx
KdivU = fd.as_vector((Kx_facet*u.dx(0), Ky_facet*u.dx(1), Kz_facet*u.dx(2)))
a = (fd.dot(KdivU, fd.grad(v))) * dx
m = u * v * phi * dx
# Defining the eigenvalue problem
petsc_a = fd.assemble(a).M.handle
petsc_m = fd.assemble(m).M.handle
num_eigenvalues = 3
# Set solver options
opts = PETSc.Options()
opts.setValue("eps_gen_hermitian", None)
#opts.setValue("st_pc_factor_shift_type", "NONZERO")
opts.setValue("eps_type", "krylovschur")
def domainify(vector, x):
dim = len(x)
return conditional(
x[0] > 0.0, vector, as_vector(tuple(0.0 for _ in range(dim)))
)
def test_3D_dirichlet_regions():
# Can't do mesh refinement because MeshHierarchy ruins the domain tags
mesh = Mesh("./3D_mesh.msh")
dim = mesh.geometric_dimension()
x = SpatialCoordinate(mesh)
solver_parameters_amg = {
"ksp_type": "cg",
"ksp_converged_reason": None,
"ksp_rtol": 1e-7,
"pc_type": "hypre",
"pc_hypre_type": "boomeramg",
"pc_hypre_boomeramg_max_iter": 5,
"pc_hypre_boomeramg_coarsen_type": "PMIS",
"pc_hypre_boomeramg_agg_nl": 2,
"pc_hypre_boomeramg_strong_threshold": 0.95,
"pc_hypre_boomeramg_interp_type": "ext+i",
"pc_hypre_boomeramg_P_max": 2,
"pc_hypre_boomeramg_relax_type_all": "sequential-Gauss-Seidel",
"pc_hypre_boomeramg_grid_sweeps_all": 1,
"pc_hypre_boomeramg_truncfactor": 0.3,
"pc_hypre_boomeramg_max_levels": 6,
}
S = VectorFunctionSpace(mesh, "CG", 1)
beta = 1.0
reg_solver = RegularizationSolver(
S,
mesh,
beta=beta,
gamma=0.0,
dx=dx,
design_domain=0,
solver_parameters=solver_parameters_amg,
)
# Exact solution with free Neumann boundary conditions for this domain
u_exact = Function(S)
u_exact_component = (
(-cos(x[0] * pi * 2) + 1)
* (-cos(x[1] * pi * 2) + 1)
* (-cos(x[2] * pi * 2) + 1)
)
u_exact.interpolate(
as_vector(tuple(u_exact_component for _ in range(dim)))
)
f = Function(S)
f_component = (
-beta
* (
4.0
* pi
* pi
* cos(2 * pi * x[0])
* (-cos(2 * pi * x[1]) + 1)
* (-cos(2 * pi * x[2]) + 1)
+ 4.0
* pi
* pi
* cos(2 * pi * x[1])
* (-cos(2 * pi * x[0]) + 1)
* (-cos(2 * pi * x[2]) + 1)
+ 4.0
* pi
* pi
* cos(2 * pi * x[2])
* (-cos(2 * pi * x[1]) + 1)
* (-cos(2 * pi * x[0]) + 1)
)
+ u_exact_component
)
f.interpolate(as_vector(tuple(f_component for _ in range(dim))))
theta = TestFunction(S)
rhs_form = inner(f, theta) * dx
velocity = Function(S)
rhs = assemble(rhs_form)
reg_solver.solve(velocity, rhs)
error = norm(
project(
domainify(u_exact, x) - velocity,
S,
solver_parameters=solver_parameters_amg,
)
)
assert error < 5e-2
def setup_sw(dirname):
refinements = 3 # number of horizontal cells = 20*(4^refinements)
R = 6371220.
H = 2000.
day = 24.*60.*60.
u_0 = 2*pi*R/(12*day) # Maximum amplitude of the zonal wind (m/s)
mesh = IcosahedralSphereMesh(radius=R,
refinement_level=refinements, degree=3)
global_normal = Expression(("x[0]", "x[1]", "x[2]"))
mesh.init_cell_orientations(global_normal)
fieldlist = ['u', 'D']
timestepping = TimesteppingParameters(dt=3600.)
output = OutputParameters(dirname=dirname+"/sw_linear_w2", steady_state_dump_err={'u':True,'D':True}, dumpfreq=12)
parameters = ShallowWaterParameters(H=H)
diagnostics = Diagnostics(*fieldlist)
state = ShallowWaterState(mesh, vertical_degree=None, horizontal_degree=1,
family="BDM",
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist)
g = parameters.g
Omega = parameters.Omega
# Coriolis expression
R = Constant(R)
Omega = Constant(parameters.Omega)
x = SpatialCoordinate(mesh)
fexpr = 2*Omega*x[2]/R
V = FunctionSpace(mesh, "CG", 1)
state.f = Function(V).interpolate(fexpr) # Coriolis frequency (1/s)
u_max = Constant(u_0)
# interpolate initial conditions
# Initial/current conditions
u0, D0 = Function(state.V[0]), Function(state.V[1])
uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0])
g = Constant(parameters.g)
Dexpr = - ((R * Omega * u_max)*(x[2]*x[2]/(R*R)))/g
u0.project(uexpr)
D0.interpolate(Dexpr)
state.initialise([u0, D0])
advection_dict = {}
advection_dict["u"] = NoAdvection(state)
advection_dict["D"] = NoAdvection(state)
linear_solver = ShallowWaterSolver(state)
# Set up forcing
sw_forcing = ShallowWaterForcing(state, linear=True)
# build time stepper
stepper = Timestepper(state, advection_dict, linear_solver,
sw_forcing)
return stepper, 2*day
def split(self, fields):
from firedrake import replace, as_vector, split
from firedrake import NonlinearVariationalProblem as NLVP
fields = tuple(tuple(f) for f in fields)
splits = self._splits.get(tuple(fields))
if splits is not None:
return splits
splits = []
problem = self._problem
splitter = ExtractSubBlock()
for field in fields:
F = splitter.split(problem.F, argument_indices=(field, ))
J = splitter.split(problem.J, argument_indices=(field, field))
us = problem.u.split()
V = F.arguments()[0].function_space()
# Exposition:
# We are going to make a new solution Function on the sub
# mixed space defined by the relevant fields.
# But the form may refer to the rest of the solution
# anyway.
# So we pull it apart and will make a new function on the
# subspace that shares data.
pieces = [us[i].dat for i in field]
if len(pieces) == 1:
val, = pieces
subu = function.Function(V, val=val)
subsplit = (subu, )
else:
val = op2.MixedDat(pieces)
subu = function.Function(V, val=val)
# Split it apart to shove in the form.
subsplit = split(subu)
# Permutation from field indexing to indexing of pieces
field_renumbering = dict([f, i] for i, f in enumerate(field))
vec = []
for i, u in enumerate(us):
if i in field:
# If this is a field we're keeping, get it from
# the new function. Otherwise just point to the
# old data.
u = subsplit[field_renumbering[i]]
if u.ufl_shape == ():
vec.append(u)
else:
for idx in numpy.ndindex(u.ufl_shape):
vec.append(u[idx])
# So now we have a new representation for the solution
# vector in the old problem. For the fields we're going
# to solve for, it points to a new Function (which wraps
# the original pieces). For the rest, it points to the
# pieces from the original Function.
# IOW, we've reinterpreted our original mixed solution
# function as being made up of some spaces we're still
# solving for, and some spaces that have just become
# coefficients in the new form.
u = as_vector(vec)
F = replace(F, {problem.u: u})
J = replace(J, {problem.u: u})
if problem.Jp is not None:
Jp = splitter.split(problem.Jp, argument_indices=(field, field))
Jp = replace(Jp, {problem.u: u})
else:
Jp = None
bcs = []
for bc in problem.bcs:
Vbc = bc.function_space()
if Vbc.parent is not None and isinstance(Vbc.parent.ufl_element(), VectorElement):
index = Vbc.parent.index
else:
index = Vbc.index
cmpt = Vbc.component
# TODO: need to test this logic
if index in field:
if len(field) == 1:
W = V
else:
W = V.sub(field_renumbering[index])
if cmpt is not None:
W = W.sub(cmpt)
bcs.append(type(bc)(W,
bc.function_arg,
bc.sub_domain,
method=bc.method))
new_problem = NLVP(F, subu, bcs=bcs, J=J, Jp=Jp,
form_compiler_parameters=problem.form_compiler_parameters)
new_problem._constant_jacobian = problem._constant_jacobian
splits.append(type(self)(new_problem, mat_type=self.mat_type, pmat_type=self.pmat_type,
appctx=self.appctx))
return self._splits.setdefault(tuple(fields), splits)
def __init__(self,
mesh,
dt,
output=None,
parameters=None,
diagnostics=None,
diagnostic_fields=None):
if output is None:
raise RuntimeError(
"You must provide a directory name for dumping results")
else:
self.output = output
self.parameters = parameters
if diagnostics is not None:
self.diagnostics = diagnostics
else:
self.diagnostics = Diagnostics()
if diagnostic_fields is not None:
self.diagnostic_fields = diagnostic_fields
else:
self.diagnostic_fields = []
# The mesh
self.mesh = mesh
self.spaces = SpaceCreator(mesh)
if self.output.dumplist is None:
self.output.dumplist = []
self.fields = StateFields(*self.output.dumplist)
self.dumpdir = None
self.dumpfile = None
self.to_pickup = None
# figure out if we're on a sphere
try:
self.on_sphere = (mesh._base_mesh.geometric_dimension() == 3
and mesh._base_mesh.topological_dimension() == 2)
except AttributeError:
self.on_sphere = (mesh.geometric_dimension() == 3
and mesh.topological_dimension() == 2)
# build the vertical normal and define perp for 2d geometries
dim = mesh.topological_dimension()
if self.on_sphere:
x = SpatialCoordinate(mesh)
R = sqrt(inner(x, x))
self.k = interpolate(x / R, mesh.coordinates.function_space())
if dim == 2:
outward_normals = CellNormal(mesh)
self.perp = lambda u: cross(outward_normals, u)
else:
kvec = [0.0] * dim
kvec[dim - 1] = 1.0
self.k = Constant(kvec)
if dim == 2:
self.perp = lambda u: as_vector([-u[1], u[0]])
# setup logger
logger.setLevel(output.log_level)
set_log_handler(mesh.comm)
if parameters is not None:
logger.info("Physical parameters that take non-default values:")
logger.info(", ".join("%s: %s" % (k, float(v))
for (k, v) in vars(parameters).items()))
# Constant to hold current time
self.t = Constant(0.0)
if type(dt) is Constant:
self.dt = dt
elif type(dt) in (float, int):
self.dt = Constant(dt)
else:
raise TypeError(
f'dt must be a Constant, float or int, not {type(dt)}')
thetab = Tsurf*exp(N**2*z/g)
theta_b = Function(state.V[2]).interpolate(thetab)
rho_b = Function(state.V[1])
# Calculate hydrostatic Pi
compressible_hydrostatic_balance(state, theta_b, rho_b)
W_DG1 = FunctionSpace(mesh, "DG", 1)
x = Function(W_DG1).interpolate(Expression("x[0]"))
a = 5.0e3
deltaTheta = 1.0e-2
theta_pert = deltaTheta*sin(np.pi*z/H)/(1 + (x - L/2)**2/a**2)
theta0.interpolate(theta_b + theta_pert)
rho0.assign(rho_b)
u0.project(as_vector([20.0,0.0]))
state.initialise([u0, rho0, theta0])
state.set_reference_profiles(rho_b, theta_b)
state.output.meanfields = {'rho':state.rhobar, 'theta':state.thetabar}
# Set up advection schemes
Vtdg = FunctionSpace(mesh, "DG", 2)
advection_dict = {}
advection_dict["u"] = EulerPoincareForm(state, state.V[0])
advection_dict["rho"] = DGAdvection(state, state.V[1], continuity=True)
advection_dict["theta"] = EmbeddedDGAdvection(state, state.V[2], Vdg=Vtdg, continuity=False)
# Set up linear solver
schur_params = {'pc_type': 'fieldsplit',
'pc_fieldsplit_type': 'schur',
def setup_sw(dirname, euler_poincare):
refinements = 3 # number of horizontal cells = 20*(4^refinements)
mesh = IcosahedralSphereMesh(radius=R, refinement_level=refinements)
x = SpatialCoordinate(mesh)
mesh.init_cell_orientations(x)
fieldlist = ['u', 'D']
timestepping = TimesteppingParameters(dt=1500.)
output = OutputParameters(dirname=dirname + "/sw",
dumplist_latlon=['D', 'D_error'],
steady_state_error_fields=['D', 'u'])
parameters = ShallowWaterParameters(H=H)
diagnostic_fields = [
RelativeVorticity(),
AbsoluteVorticity(),
PotentialVorticity(),
ShallowWaterPotentialEnstrophy('RelativeVorticity'),
ShallowWaterPotentialEnstrophy('AbsoluteVorticity'),
ShallowWaterPotentialEnstrophy('PotentialVorticity'),
Difference('RelativeVorticity', 'AnalyticalRelativeVorticity'),
Difference('AbsoluteVorticity', 'AnalyticalAbsoluteVorticity'),
Difference('PotentialVorticity', 'AnalyticalPotentialVorticity'),
Difference('SWPotentialEnstrophy_from_PotentialVorticity',
'SWPotentialEnstrophy_from_RelativeVorticity'),
Difference('SWPotentialEnstrophy_from_PotentialVorticity',
'SWPotentialEnstrophy_from_AbsoluteVorticity'),
MeridionalComponent('u'),
ZonalComponent('u'),
RadialComponent('u')
]
state = State(mesh,
vertical_degree=None,
horizontal_degree=1,
family="BDM",
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostic_fields=diagnostic_fields,
fieldlist=fieldlist)
# interpolate initial conditions
u0 = state.fields("u")
D0 = state.fields("D")
uexpr = as_vector([-u_max * x[1] / R, u_max * x[0] / R, 0.0])
Omega = parameters.Omega
g = parameters.g
Dexpr = H - ((R * Omega * u_max + u_max * u_max / 2.0) * (x[2] * x[2] /
(R * R))) / g
# Coriolis
fexpr = 2 * Omega * x[2] / R
V = FunctionSpace(mesh, "CG", 1)
f = state.fields("coriolis", Function(V))
f.interpolate(fexpr) # Coriolis frequency (1/s)
u0.project(uexpr)
D0.interpolate(Dexpr)
state.initialise([('u', u0), ('D', D0)])
if euler_poincare:
ueqn = EulerPoincare(state, u0.function_space())
sw_forcing = ShallowWaterForcing(state, euler_poincare=True)
else:
ueqn = VectorInvariant(state, u0.function_space())
sw_forcing = ShallowWaterForcing(state, euler_poincare=False)
Deqn = AdvectionEquation(state,
D0.function_space(),
equation_form="continuity")
advected_fields = []
advected_fields.append(("u", ThetaMethod(state, u0, ueqn)))
advected_fields.append(("D", SSPRK3(state, D0, Deqn)))
linear_solver = ShallowWaterSolver(state)
# build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver, sw_forcing)
vspace = FunctionSpace(state.mesh, "CG", 3)
vexpr = (2 * u_max / R) * x[2] / R
vrel_analytical = state.fields("AnalyticalRelativeVorticity", vspace)
vrel_analytical.interpolate(vexpr)
vabs_analytical = state.fields("AnalyticalAbsoluteVorticity", vspace)
vabs_analytical.interpolate(vexpr + f)
pv_analytical = state.fields("AnalyticalPotentialVorticity", vspace)
pv_analytical.interpolate((vexpr + f) / D0)
return stepper, 0.25 * day
tmax = dt
else:
tmax = 9000.
nlayers = 70 # horizontal layers
columns = 180 # number of columns
L = 144000.
m = PeriodicIntervalMesh(columns, L)
# build volume mesh
H = 35000. # Height position of the model top
ext_mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers)
Vc = VectorFunctionSpace(ext_mesh, "DG", 2)
coord = SpatialCoordinate(ext_mesh)
x = Function(Vc).interpolate(as_vector([coord[0], coord[1]]))
a = 1000.
xc = L/2.
x, z = SpatialCoordinate(ext_mesh)
hm = 1.
zs = hm*a**2/((x-xc)**2 + a**2)
smooth_z = True
dirname = 'nh_mountain'
if smooth_z:
dirname += '_smootherz'
zh = 5000.
xexpr = as_vector([x, conditional(z < zh, z + cos(0.5*pi*z/zh)**6*zs, z)])
else:
xexpr = as_vector([x, z + ((H-z)/H)*zs])
#
#np.save('./spe10z85_kx.npy', kx_array)
#np.save('./spe10z85_ky.npy', ky_array)
#np.save('./spe10z85_phi.npy', phi_array)
kx_geomean = np.exp(np.log(kx_array).sum()/(Nx*Ny))
ky_geomean = np.exp(np.log(ky_array).sum()/(Nx*Ny))
Kx_hom = fd.Constant(kx_geomean)
Ky_hom = fd.Constant(ky_geomean)
phi_hom = fd.Constant(phi_array.mean())
print("END: Read in reservoir fields")
# -----------------------------------------------------------------------------
# Solve homogeneous problem
# -----------------------------------------------------------------------------
print("START: Solve homogeneous problem")
dx = fd.dx
KdivU = fd.as_vector((Kx_hom*u.dx(0), Ky_hom*u.dx(1)))
a = (fd.dot(KdivU, fd.grad(v))) * dx
m = u * v * phi_hom * dx
# Defining the eigenvalue problem
petsc_a = fd.assemble(a).M.handle
petsc_m = fd.assemble(m).M.handle
num_eigenvalues = 30
# Set solver options
opts = PETSc.Options()
opts.setValue("eps_gen_hermitian", None)
opts.setValue("st_pc_factor_shift_type", "NONZERO")
opts.setValue('st_type', 'sinvert')
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt)
dirname = 'sk_hydrostatic'
output = OutputParameters(dirname=dirname,
dumpfreq=50,
dumplist=['u'],
perturbation_fields=['theta', 'rho'],
log_level='INFO')
parameters = CompressibleParameters()
diagnostics = Diagnostics(*fieldlist)
diagnostic_fields = [CourantNumber()]
Omega = as_vector((0., 0., 0.5e-4))
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="RTCF",
Coriolis=Omega,
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist,
diagnostic_fields=diagnostic_fields)
# Initial conditions
u0 = state.fields("u")
# rescaling beta = 1.0 f = f / beta L = beta * L # Construct 2D periodic base mesh m = PeriodicRectangleMesh(columns, 1, 2. * L, 1.e5, quadrilateral=True) # build 3D mesh by extruding the base mesh mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers) ############################################################################## # set up all the other things that state requires ############################################################################## # Coriolis expression Omega = as_vector([0., 0., f * 0.5]) # list of prognostic fieldnames # this is passed to state and used to construct a dictionary, # state.field_dict so that we can access fields by name # u is the 3D velocity # p is the pressure # b is the buoyancy fieldlist = ['u', 'p', 'b'] # class containing timestepping parameters # all values not explicitly set here use the default values provided # and documented in configuration.py timestepping = TimesteppingParameters(dt=dt) # class containing output parameters
def test_3D_cartesian_recovery(geometry, element, mesh, expr):
family = "RTCF" if element == "quadrilateral" else "BDM"
# horizontal base spaces
cell = mesh._base_mesh.ufl_cell().cellname()
u_hori = FiniteElement(family, cell, 1)
w_hori = FiniteElement("DG", cell, 0)
# vertical base spaces
u_vert = FiniteElement("DG", interval, 0)
w_vert = FiniteElement("CG", interval, 1)
# build elements
u_element = HDiv(TensorProductElement(u_hori, u_vert))
w_element = HDiv(TensorProductElement(w_hori, w_vert))
theta_element = TensorProductElement(w_hori, w_vert)
v_element = u_element + w_element
# DG1
DG1_hori = FiniteElement("DG", cell, 1, variant="equispaced")
DG1_vert = FiniteElement("DG", interval, 1, variant="equispaced")
DG1_elt = TensorProductElement(DG1_hori, DG1_vert)
DG1 = FunctionSpace(mesh, DG1_elt)
vec_DG1 = VectorFunctionSpace(mesh, DG1_elt)
# spaces
DG0 = FunctionSpace(mesh, "DG", 0)
CG1 = FunctionSpace(mesh, "CG", 1)
Vt = FunctionSpace(mesh, theta_element)
Vt_brok = FunctionSpace(mesh, BrokenElement(theta_element))
Vu = FunctionSpace(mesh, v_element)
vec_CG1 = VectorFunctionSpace(mesh, "CG", 1)
# our actual theta and rho and v
rho_CG1_true = Function(CG1).interpolate(expr)
theta_CG1_true = Function(CG1).interpolate(expr)
v_CG1_true = Function(vec_CG1).interpolate(as_vector([expr, expr, expr]))
rho_Vt_true = Function(Vt).interpolate(expr)
# make the initial fields by projecting expressions into the lowest order spaces
rho_DG0 = Function(DG0).interpolate(expr)
rho_CG1 = Function(CG1)
theta_Vt = Function(Vt).interpolate(expr)
theta_CG1 = Function(CG1)
v_Vu = Function(Vu).project(as_vector([expr, expr, expr]))
v_CG1 = Function(vec_CG1)
rho_Vt = Function(Vt)
# make the recoverers and do the recovery
rho_recoverer = Recoverer(rho_DG0, rho_CG1, VDG=DG1, boundary_method=Boundary_Method.dynamics)
theta_recoverer = Recoverer(theta_Vt, theta_CG1, VDG=DG1, boundary_method=Boundary_Method.dynamics)
v_recoverer = Recoverer(v_Vu, v_CG1, VDG=vec_DG1, boundary_method=Boundary_Method.dynamics)
rho_Vt_recoverer = Recoverer(rho_DG0, rho_Vt, VDG=Vt_brok, boundary_method=Boundary_Method.physics)
rho_recoverer.project()
theta_recoverer.project()
v_recoverer.project()
rho_Vt_recoverer.project()
rho_diff = errornorm(rho_CG1, rho_CG1_true) / norm(rho_CG1_true)
theta_diff = errornorm(theta_CG1, theta_CG1_true) / norm(theta_CG1_true)
v_diff = errornorm(v_CG1, v_CG1_true) / norm(v_CG1_true)
rho_Vt_diff = errornorm(rho_Vt, rho_Vt_true) / norm(rho_Vt_true)
tolerance = 1e-7
error_message = ("""
Incorrect recovery for {variable} with {boundary} boundary method
on {geometry} 3D Cartesian domain with {element} elements
""")
assert rho_diff < tolerance, error_message.format(variable='rho', boundary='dynamics',
geometry=geometry, element=element)
assert v_diff < tolerance, error_message.format(variable='v', boundary='dynamics',
geometry=geometry, element=element)
assert theta_diff < tolerance, error_message.format(variable='rho', boundary='dynamics',
geometry=geometry, element=element)
assert rho_Vt_diff < tolerance, error_message.format(variable='rho', boundary='physics',
geometry=geometry, element=element)
def __init__(self, state, V):
super(EulerPoincareForm, self).__init__(state)
dt = state.timestepping.dt
w = TestFunction(V)
u = TrialFunction(V)
self.u0 = Function(V)
ustar = 0.5*(self.u0 + u)
n = FacetNormal(state.mesh)
Upwind = 0.5*(sign(dot(self.ubar, n))+1)
if state.mesh.geometric_dimension() == 3:
if V.extruded:
surface_measure = (dS_h + dS_v)
else:
surface_measure = dS
# <w,curl(u) cross ubar + grad( u.ubar)>
# =<curl(u),ubar cross w> - <div(w), u.ubar>
# =<u,curl(ubar cross w)> -
# <<u_upwind, [[n cross(ubar cross w)cross]]>>
both = lambda u: 2*avg(u)
Eqn = (
inner(w, u-self.u0)*dx
+ dt*inner(ustar, curl(cross(self.ubar, w)))*dx
- dt*inner(both(Upwind*ustar),
both(cross(n, cross(self.ubar, w))))*surface_measure
- dt*div(w)*inner(ustar, self.ubar)*dx
)
# define surface measure and terms involving perp differently
# for slice (i.e. if V.extruded is True) and shallow water
# (V.extruded is False)
else:
if V.extruded:
surface_measure = (dS_h + dS_v)
perp = lambda u: as_vector([-u[1], u[0]])
perp_u_upwind = Upwind('+')*perp(ustar('+')) + Upwind('-')*perp(ustar('-'))
else:
surface_measure = dS
outward_normals = CellNormal(state.mesh)
perp = lambda u: cross(outward_normals, u)
perp_u_upwind = Upwind('+')*cross(outward_normals('+'),ustar('+')) + Upwind('-')*cross(outward_normals('-'),ustar('-'))
Eqn = (
(inner(w, u-self.u0)
- dt*inner(w, div(perp(ustar))*perp(self.ubar))
- dt*div(w)*inner(ustar, self.ubar))*dx
- dt*inner(jump(inner(w, perp(self.ubar)), n),
perp_u_upwind)*surface_measure
+ dt*jump(inner(w,
perp(self.ubar))*perp(ustar), n)*surface_measure
)
a = lhs(Eqn)
L = rhs(Eqn)
self.u1 = Function(V)
uproblem = LinearVariationalProblem(a, L, self.u1)
self.usolver = LinearVariationalSolver(uproblem,
options_prefix='EPAdvection')
'ksp_type': 'preonly',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
}
}
compressible_hydrostatic_balance(state, theta_b, rho_b, params=piparams)
PETSc.Sys.Print("Finished computing hydrostatic varaibles...\n")
a = 5.0e3
deltaTheta = 1.0e-2
theta_pert = deltaTheta * sin(np.pi * z / H) / (1 + (x - L / 2)**2 / a**2)
theta0.interpolate(theta_b + theta_pert)
rho0.assign(rho_b)
u0.project(as_vector([20.0, 0.0]))
state.initialise([('u', u0), ('rho', rho0), ('theta', theta0)])
state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
# Set up advection schemes
ueqn = EulerPoincare(state, Vu)
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
supg = True
if supg:
thetaeqn = SUPGAdvection(state, Vt, equation_form="advective")
else:
thetaeqn = EmbeddedDGAdvection(state,
Vt,
equation_form="advective",
g = parameters.g
Omega = parameters.Omega
# Coriolis expression
R = Constant(R)
Omega = Constant(parameters.Omega)
x = SpatialCoordinate(mesh)
fexpr = 2*Omega*x[2]/R
V = FunctionSpace(mesh, "CG", 1)
state.f = Function(V).interpolate(fexpr) # Coriolis frequency (1/s)
u_max = Constant(u_0)
# interpolate initial conditions
# Initial/current conditions
u0, D0 = Function(state.V[0]), Function(state.V[1])
uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0])
g = Constant(parameters.g)
Dexpr = - ((R * Omega * u_max)*(x[2]*x[2]/(R*R)))/g
u0.project(uexpr)
D0.interpolate(Dexpr)
state.initialise([u0, D0])
advection_dict = {}
advection_dict["u"] = NoAdvection(state)
advection_dict["D"] = NoAdvection(state)
linear_solver = ShallowWaterSolver(state)
# Set up forcing
sw_forcing = ShallowWaterForcing(state)
def heat_exchanger_optimization(mu=0.03, n_iters=1000):
output_dir = "2D/"
path = os.path.abspath(__file__)
dir_path = os.path.dirname(path)
mesh = fd.Mesh(f"{dir_path}/2D_mesh.msh")
# Perturb the mesh coordinates. Necessary to calculate shape derivatives
S = fd.VectorFunctionSpace(mesh, "CG", 1)
s = fd.Function(S, name="deform")
mesh.coordinates.assign(mesh.coordinates + s)
# Initial level set function
x, y = fd.SpatialCoordinate(mesh)
PHI = fd.FunctionSpace(mesh, "CG", 1)
phi_expr = sin(y * pi / 0.2) * cos(x * pi / 0.2) - fd.Constant(0.8)
# Avoid recording the operation interpolate into the tape.
# Otherwise, the shape derivatives will not be correct
with fda.stop_annotating():
phi = fd.interpolate(phi_expr, PHI)
phi.rename("LevelSet")
fd.File(output_dir + "phi_initial.pvd").write(phi)
# Physics
mu = fd.Constant(mu) # viscosity
alphamin = 1e-12
alphamax = 2.5 / (2e-4)
parameters = {
"mat_type": "aij",
"ksp_type": "preonly",
"ksp_converged_reason": None,
"pc_type": "lu",
"pc_factor_mat_solver_type": "mumps",
}
stokes_parameters = parameters
temperature_parameters = parameters
u_inflow = 2e-3
tin1 = fd.Constant(10.0)
tin2 = fd.Constant(100.0)
P2 = fd.VectorElement("CG", mesh.ufl_cell(), 2)
P1 = fd.FiniteElement("CG", mesh.ufl_cell(), 1)
TH = P2 * P1
W = fd.FunctionSpace(mesh, TH)
U = fd.TrialFunction(W)
u, p = fd.split(U)
V = fd.TestFunction(W)
v, q = fd.split(V)
epsilon = fd.Constant(10000.0)
def hs(phi, epsilon):
return fd.Constant(alphamax) * fd.Constant(1.0) / (
fd.Constant(1.0) + exp(-epsilon * phi)) + fd.Constant(alphamin)
def stokes(phi, BLOCK_INLET_MOUTH, BLOCK_OUTLET_MOUTH):
a_fluid = mu * inner(grad(u), grad(v)) - div(v) * p - q * div(u)
darcy_term = inner(u, v)
return (a_fluid * dx + hs(phi, epsilon) * darcy_term * dx(0) +
alphamax * darcy_term *
(dx(BLOCK_INLET_MOUTH) + dx(BLOCK_OUTLET_MOUTH)))
# Dirichlet boundary conditions
inflow1 = fd.as_vector([
u_inflow * sin(
((y - (line_sep -
(dist_center + inlet_width))) * pi) / inlet_width),
0.0,
])
inflow2 = fd.as_vector([
u_inflow * sin(((y - (line_sep + dist_center)) * pi) / inlet_width),
0.0,
])
noslip = fd.Constant((0.0, 0.0))
# Stokes 1
bcs1_1 = fd.DirichletBC(W.sub(0), noslip, WALLS)
bcs1_2 = fd.DirichletBC(W.sub(0), inflow1, INLET1)
bcs1_3 = fd.DirichletBC(W.sub(1), fd.Constant(0.0), OUTLET1)
bcs1_4 = fd.DirichletBC(W.sub(0), noslip, INLET2)
bcs1_5 = fd.DirichletBC(W.sub(0), noslip, OUTLET2)
bcs1 = [bcs1_1, bcs1_2, bcs1_3, bcs1_4, bcs1_5]
# Stokes 2
bcs2_1 = fd.DirichletBC(W.sub(0), noslip, WALLS)
bcs2_2 = fd.DirichletBC(W.sub(0), inflow2, INLET2)
bcs2_3 = fd.DirichletBC(W.sub(1), fd.Constant(0.0), OUTLET2)
bcs2_4 = fd.DirichletBC(W.sub(0), noslip, INLET1)
bcs2_5 = fd.DirichletBC(W.sub(0), noslip, OUTLET1)
bcs2 = [bcs2_1, bcs2_2, bcs2_3, bcs2_4, bcs2_5]
# Forward problems
U1, U2 = fd.Function(W), fd.Function(W)
L = inner(fd.Constant((0.0, 0.0, 0.0)), V) * dx
problem = fd.LinearVariationalProblem(stokes(-phi, INMOUTH2, OUTMOUTH2),
L,
U1,
bcs=bcs1)
solver_stokes1 = fd.LinearVariationalSolver(
problem,
solver_parameters=stokes_parameters,
options_prefix="stokes_1")
solver_stokes1.solve()
problem = fd.LinearVariationalProblem(stokes(phi, INMOUTH1, OUTMOUTH1),
L,
U2,
bcs=bcs2)
solver_stokes2 = fd.LinearVariationalSolver(
problem,
solver_parameters=stokes_parameters,
options_prefix="stokes_2")
solver_stokes2.solve()
# Convection difussion equation
ks = fd.Constant(1e0)
cp_value = 5.0e5
cp = fd.Constant(cp_value)
T = fd.FunctionSpace(mesh, "DG", 1)
t = fd.Function(T, name="Temperature")
w = fd.TestFunction(T)
# Mesh-related functions
n = fd.FacetNormal(mesh)
h = fd.CellDiameter(mesh)
u1, p1 = fd.split(U1)
u2, p2 = fd.split(U2)
def upwind(u):
return (dot(u, n) + abs(dot(u, n))) / 2.0
u1n = upwind(u1)
u2n = upwind(u2)
# Penalty term
alpha = fd.Constant(500.0)
# Bilinear form
a_int = dot(grad(w), ks * grad(t) - cp * (u1 + u2) * t) * dx
a_fac = (fd.Constant(-1.0) * ks * dot(avg(grad(w)), jump(t, n)) * dS +
fd.Constant(-1.0) * ks * dot(jump(w, n), avg(grad(t))) * dS +
ks("+") *
(alpha("+") / avg(h)) * dot(jump(w, n), jump(t, n)) * dS)
a_vel = (dot(
jump(w),
cp * (u1n("+") + u2n("+")) * t("+") - cp *
(u1n("-") + u2n("-")) * t("-"),
) * dS + dot(w,
cp * (u1n + u2n) * t) * ds)
a_bnd = (dot(w,
cp * dot(u1 + u2, n) * t) * (ds(INLET1) + ds(INLET2)) +
w * t * (ds(INLET1) + ds(INLET2)) - w * tin1 * ds(INLET1) -
w * tin2 * ds(INLET2) + alpha / h * ks * w * t *
(ds(INLET1) + ds(INLET2)) - ks * dot(grad(w), t * n) *
(ds(INLET1) + ds(INLET2)) - ks * dot(grad(t), w * n) *
(ds(INLET1) + ds(INLET2)))
aT = a_int + a_fac + a_vel + a_bnd
LT_bnd = (alpha / h * ks * tin1 * w * ds(INLET1) +
alpha / h * ks * tin2 * w * ds(INLET2) -
tin1 * ks * dot(grad(w), n) * ds(INLET1) -
tin2 * ks * dot(grad(w), n) * ds(INLET2))
problem = fd.LinearVariationalProblem(derivative(aT, t), LT_bnd, t)
solver_temp = fd.LinearVariationalSolver(
problem,
solver_parameters=temperature_parameters,
options_prefix="temperature",
)
solver_temp.solve()
# fd.solve(eT == 0, t, solver_parameters=temperature_parameters)
# Cost function: Flux at the cold outlet
scale_factor = 4e-4
Jform = fd.assemble(
fd.Constant(-scale_factor * cp_value) * inner(t * u1, n) * ds(OUTLET1))
# Constraints: Pressure drop on each fluid
power_drop = 1e-2
Power1 = fd.assemble(p1 / power_drop * ds(INLET1))
Power2 = fd.assemble(p2 / power_drop * ds(INLET2))
phi_pvd = fd.File("phi_evolution.pvd")
def deriv_cb(phi):
with stop_annotating():
phi_pvd.write(phi[0])
c = fda.Control(s)
# Reduced Functionals
Jhat = LevelSetFunctional(Jform, c, phi, derivative_cb_pre=deriv_cb)
P1hat = LevelSetFunctional(Power1, c, phi)
P1control = fda.Control(Power1)
P2hat = LevelSetFunctional(Power2, c, phi)
P2control = fda.Control(Power2)
Jhat_v = Jhat(phi)
print("Initial cost function value {:.5f}".format(Jhat_v), flush=True)
print("Power drop 1 {:.5f}".format(Power1), flush=True)
print("Power drop 2 {:.5f}".format(Power2), flush=True)
beta_param = 0.08
# Regularize the shape derivatives only in the domain marked with 0
reg_solver = RegularizationSolver(S,
mesh,
beta=beta_param,
gamma=1e5,
dx=dx,
design_domain=0)
tol = 1e-5
dt = 0.05
params = {
"alphaC": 1.0,
"debug": 5,
"alphaJ": 1.0,
"dt": dt,
"K": 1e-3,
"maxit": n_iters,
"maxtrials": 5,
"itnormalisation": 10,
"tol_merit":
5e-3, # new merit can be within 0.5% of the previous merit
# "normalize_tol" : -1,
"tol": tol,
}
solver_parameters = {
"reinit_solver": {
"h_factor": 2.0,
}
}
# Optimization problem
problem = InfDimProblem(
Jhat,
reg_solver,
ineqconstraints=[
Constraint(P1hat, 1.0, P1control),
Constraint(P2hat, 1.0, P2control),
],
solver_parameters=solver_parameters,
)
results = nlspace_solve(problem, params)
return results
# 3.1) Define trial and test functions
w = fd.TestFunction(X)
# 3.2) Set initial conditions
c0 = fd.Function(X, name="conc0")
c = fd.Function(X, name="conc")
# 3.3) set boundary conditions
t_bc = fd.DirichletBC(X, cIn, inlet)
# ======================
# 3.4) Variational Form
# coefficients
# advective velocity
vel = fd.Function(V, name='velocity').interpolate(fd.as_vector([1.0, 0.]))
K = fd.Constant(K)
Diff = fd.Constant(Diff)
Dt = fd.Constant(1e-3)
c_mid = 0.5 * (c + c0) # Crank-Nicolson timestepping
fc = fd.Constant(s)
# weak form (transport)
F1 = w * (c - c0) * fd.dx + Dt * (
w * fd.dot(vel, fd.grad(c_mid)) + Diff * fd.dot(fd.grad(w), fd.grad(c_mid))
+ w * K * c_mid - fc * w) * fd.dx # - Dt*h_n*w*fd.ds(outlet)
# strong form
R = (c - c0) + Dt * (fd.dot(vel, fd.grad(c_mid)) -
Diff * fd.div(fd.grad(c_mid)) + K * c_mid - fc)
def test_line_search():
mesh = fd.UnitSquareMesh(50, 50)
# Shape derivative
S = fd.VectorFunctionSpace(mesh, "CG", 1)
s = fd.Function(S, name="deform")
mesh.coordinates.assign(mesh.coordinates + s)
# Level set
PHI = fd.FunctionSpace(mesh, "CG", 1)
x, y = fd.SpatialCoordinate(mesh)
with fda.stop_annotating():
phi = fd.interpolate(-(x - 0.5), PHI)
phi.rename("original")
solver_parameters = {
"ts_atol": 1e-4,
"ts_rtol": 1e-4,
"ts_dt": 1e-2,
"ts_exact_final_time": "matchstep",
"ts_monitor": None,
}
## Search direction
with fda.stop_annotating():
delta_x = fd.interpolate(fd.as_vector([-100 * x, 0.0]), S)
delta_x.rename("velocity")
# Cost function
J = fd.assemble(hs(-phi) * dx)
# Reduced Functional
c = fda.Control(s)
Jhat = LevelSetFunctional(J, c, phi)
# InfDim Problem
beta_param = 0.08
reg_solver = RegularizationSolver(S,
mesh,
beta=beta_param,
gamma=1e5,
dx=dx)
solver_parameters = {"hj_solver": solver_parameters}
problem = InfDimProblem(Jhat,
reg_solver,
solver_parameters=solver_parameters)
new_phi = fd.Function(PHI, name="new_ls")
orig_phi = fd.Function(PHI)
with fda.stop_annotating():
orig_phi.assign(phi)
problem.delta_x.assign(delta_x)
AJ, AC = 1.0, 1.0
C = np.array([])
merit = merit_eval_new(AJ, J, AC, C)
rtol = 1e-4
new_phi, newJ, newG, newH = line_search(
problem,
orig_phi,
new_phi,
merit_eval_new,
merit,
AJ,
AC,
dt=1.0,
tol_merit=rtol,
maxtrials=20,
)
assert_allclose(newJ, J, rtol=rtol)
def facet_normal_2(mesh):
r"""Compute the horizontal component of the unit outward normal vector
to a mesh"""
ν = firedrake.FacetNormal(mesh)
return firedrake.as_vector((ν[0], ν[1]))
diagnostic_fields = [Sum('D', 'topography')]
state = State(mesh, horizontal_degree=1,
family="BDM",
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostic_fields=diagnostic_fields,
fieldlist=fieldlist)
# interpolate initial conditions
u0 = state.fields('u')
D0 = state.fields('D')
x = SpatialCoordinate(mesh)
u_max = 20. # Maximum amplitude of the zonal wind (m/s)
uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0])
theta, lamda = latlon_coords(mesh)
Omega = parameters.Omega
g = parameters.g
Rsq = R**2
R0 = pi/9.
R0sq = R0**2
lamda_c = -pi/2.
lsq = (lamda - lamda_c)**2
theta_c = pi/6.
thsq = (theta - theta_c)**2
rsq = Min(R0sq, lsq+thsq)
r = sqrt(rsq)
bexpr = 2000 * (1 - r/R0)
Dexpr = H - ((R * Omega * u_max + 0.5*u_max**2)*x[2]**2/Rsq)/g - bexpr
def setup_sw(dirname, dt, u_transport_option):
refinements = 3 # number of horizontal cells = 20*(4^refinements)
mesh = IcosahedralSphereMesh(radius=R,
refinement_level=refinements)
x = SpatialCoordinate(mesh)
mesh.init_cell_orientations(x)
output = OutputParameters(dirname=dirname+"/sw", dumplist_latlon=['D', 'D_error'], steady_state_error_fields=['D', 'u'])
parameters = ShallowWaterParameters(H=H)
diagnostic_fields = [RelativeVorticity(), AbsoluteVorticity(),
PotentialVorticity(),
ShallowWaterPotentialEnstrophy('RelativeVorticity'),
ShallowWaterPotentialEnstrophy('AbsoluteVorticity'),
ShallowWaterPotentialEnstrophy('PotentialVorticity'),
Difference('RelativeVorticity',
'AnalyticalRelativeVorticity'),
Difference('AbsoluteVorticity',
'AnalyticalAbsoluteVorticity'),
Difference('PotentialVorticity',
'AnalyticalPotentialVorticity'),
Difference('SWPotentialEnstrophy_from_PotentialVorticity',
'SWPotentialEnstrophy_from_RelativeVorticity'),
Difference('SWPotentialEnstrophy_from_PotentialVorticity',
'SWPotentialEnstrophy_from_AbsoluteVorticity'),
MeridionalComponent('u'),
ZonalComponent('u'),
RadialComponent('u')]
state = State(mesh,
dt=dt,
output=output,
parameters=parameters,
diagnostic_fields=diagnostic_fields)
Omega = parameters.Omega
fexpr = 2*Omega*x[2]/R
eqns = ShallowWaterEquations(state, family="BDM", degree=1,
fexpr=fexpr,
u_transport_option=u_transport_option)
# interpolate initial conditions
u0 = state.fields("u")
D0 = state.fields("D")
uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0])
g = parameters.g
Dexpr = H - ((R * Omega * u_max + u_max*u_max/2.0)*(x[2]*x[2]/(R*R)))/g
u0.project(uexpr)
D0.interpolate(Dexpr)
vspace = FunctionSpace(state.mesh, "CG", 3)
vexpr = (2*u_max/R)*x[2]/R
f = state.fields("coriolis")
vrel_analytical = state.fields("AnalyticalRelativeVorticity", vspace)
vrel_analytical.interpolate(vexpr)
vabs_analytical = state.fields("AnalyticalAbsoluteVorticity", vspace)
vabs_analytical.interpolate(vexpr + f)
pv_analytical = state.fields("AnalyticalPotentialVorticity", vspace)
pv_analytical.interpolate((vexpr+f)/D0)
return state, eqns
