def sphere_to_cartesian(mesh, u_zonal, u_merid):
theta, lamda = latlon_coords(mesh)
cartesian_u_expr = -u_zonal*sin(lamda) - u_merid*sin(theta)*cos(lamda)
cartesian_v_expr = u_zonal*cos(lamda) - u_merid*sin(theta)*sin(lamda)
cartesian_w_expr = u_merid*cos(theta)
return as_vector((cartesian_u_expr, cartesian_v_expr, cartesian_w_expr))
def sphere_to_cartesian(mesh, u_zonal, u_merid):
"""Reformulate vector function in theta, lamda directions
as function in x,y,z directions"""
theta, lamda = latlon_coords(mesh)
cartesian_u_expr = -u_zonal*sin(lamda) - u_merid*sin(theta)*cos(lamda)
cartesian_v_expr = u_zonal*cos(lamda) - u_merid*sin(theta)*sin(lamda)
cartesian_w_expr = u_merid*cos(theta)
return as_vector((cartesian_u_expr, cartesian_v_expr, cartesian_w_expr))
def __init__(self):
super().__init__()
X = self.X
# volume and boundary integrals
self.name = "LevelSet Example 2"
self.J = (sin(X[0]) * cos(X[1]) * dx +
pow(1.3 + X[0], 4.2) * pow(1.4 + X[1], 3.3) * dx +
exp(sin(X[0]) + cos(X[1])) * dx +
ln(5 + sin(X[0]) + cos(X[1])) * ds)
self.set_quadrature(8)
def __init__(self):
super().__init__()
X = self.X
# volume and boundary integrals
self.name = "LevelSet Example 3"
V = FunctionSpace(self.mesh, "CG", 1)
u = interpolate(sin(X[0]) * cos(X[1])**2, V)
n = FacetNormal(self.mesh)
self.J = (sin(X[0]) * cos(X[1]) * dx +
pow(1.3 + X[0], 4.2) * pow(1.4 + X[1], 3.3) * dx +
exp(sin(X[0]) + cos(X[1])) * dx +
ln(5 + sin(X[0]) + cos(X[1])) * ds + inner(grad(u), n) * ds)
self.set_quadrature(8)
def __init__(self):
super().__init__()
u, v, X = self.u, self.v, self.X
self.name = "nonlinear PDE constraint"
f = sin(X[1]) * cos(X[0])
g = exp(f)
self.F = (u * v + (1 + u**2) * inner(grad(u), grad(v)) -
f * v) * dx + g * v * ds
self.J = u * dx + pow(1 + u * u, 2.5) * ds
self.set_quadrature(10)
def __init__(self):
super().__init__()
u, v, X = self.u, self.v, self.X
# nonhomogeneous Neumann bc and nonlinear functional on bdry
self.name = "PDE constrained Example 2"
f = sin(X[1]) * cos(X[0])
g = exp(f)
self.F = (u * v + inner(grad(u), grad(v)) - f * v) * dx + g * v * ds
self.J = u * dx + pow(1 + u * u, 2.5) * ds
self.set_quadrature(10)
def __init__(self):
super().__init__()
self.name = "PDE-constraint with DirBC"
u, v, X = self.u, self.v, self.X
f = sin(X[1]) * cos(X[0])
g = f / 2.
self.g = g
self.with_DirBC = 1 # special case for DirichletBC in TaylorTest
self.bc = DirichletBC(self.V, 0., "on_boundary")
self.F = (inner(grad(u + g), grad(v)) + (u + g) * v - f * v) * dx
self.J = (u + g) * (u + g) * dx + pow(1 + u * u, 2.5) * ds
self.set_quadrature(10)
def initial_conditions(self):
# old: ic1 = fd.project( fd.Expression([0.,0.,0.]), self.V )
# old: ic2 = fd.project( fd.Expression(["0.1*(1-cos(pi*x[2]/Lz/2.))",0.,0.], Lz=self.Lz), self.V )
ic1 = fd.project(fd.Constant([0., 0., 0.]), self.V)
x = fd.SpatialCoordinate(self.mesh)
ic2 = fd.project(
fd.as_vector(
[0.1 * (1 - fd.cos(fd.pi * x[2] / self.Lz / 2.)), 0., 0.]),
self.V)
# ic3 = fd.project( fd.Expression(["0.1*x[2]/Lz_B", 0., 0.], Lz_B=self.Lz), self.V )
self.X.assign(ic2)
# self.static_solver()
self.U.assign(ic1)
def regularization_form(r):
mesh = UnitSquareMesh(2 ** r, 2 ** r)
x = SpatialCoordinate(mesh)
S = VectorFunctionSpace(mesh, "CG", 1)
beta = 4.0
reg_solver = RegularizationSolver(S, mesh, beta=beta, gamma=0.0, dx=dx)
# Exact solution with free Neumann boundary conditions for this domain
u_exact = Function(S)
u_exact_component = cos(x[0] * pi * 2) * cos(x[1] * pi * 2)
u_exact.interpolate(as_vector((u_exact_component, u_exact_component)))
f = Function(S)
theta = TestFunction(S)
f_component = (1 + beta * 8 * pi * pi) * u_exact_component
f.interpolate(as_vector((f_component, f_component)))
rhs_form = inner(f, theta) * dx
velocity = Function(S)
rhs = assemble(rhs_form)
reg_solver.solve(velocity, rhs)
File("solution_vel_unitsquare.pvd").write(velocity)
return norm(project(u_exact - velocity, S))
def setup_fallout(dirname):
# declare grid shape, with length L and height H
L = 10.
H = 10.
nlayers = 10
ncolumns = 10
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
x = SpatialCoordinate(mesh)
dt = 0.1
output = OutputParameters(dirname=dirname + "/fallout",
dumpfreq=10,
dumplist=['rain_mixing_ratio'])
parameters = CompressibleParameters()
diagnostic_fields = [Precipitation()]
state = State(mesh,
dt=dt,
output=output,
parameters=parameters,
diagnostic_fields=diagnostic_fields)
Vrho = state.spaces("DG1_equispaced")
problem = [(AdvectionEquation(state, Vrho, "rho", ufamily="CG",
udegree=1), ForwardEuler(state))]
state.fields("rho").assign(1.)
physics_list = [Fallout(state)]
rain0 = state.fields("rain_mixing_ratio")
# set up rain
xc = L / 2
zc = H / 2
rc = H / 4
r = sqrt((x[0] - xc)**2 + (x[1] - zc)**2)
rain_expr = conditional(r > rc, 0., 1e-3 * (cos(pi * r / (rc * 2)))**2)
rain0.interpolate(rain_expr)
# build time stepper
stepper = PrescribedTransport(state, problem, physics_list=physics_list)
return stepper, 10.0
def test_2D_dirichlet_regions():
# Can't do mesh refinement because MeshHierarchy ruins the domain tags
mesh = Mesh("./2D_mesh.msh")
dim = mesh.geometric_dimension()
x = SpatialCoordinate(mesh)
S = VectorFunctionSpace(mesh, "CG", 1)
beta = 1.0
reg_solver = RegularizationSolver(
S, mesh, beta=beta, gamma=0.0, dx=dx, design_domain=0
)
# 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)
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)
+ 4.0 * pi * pi * cos(2 * pi * x[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)
assert norm(project(domainify(u_exact, x) - velocity, S)) < 1e-3
theta_b = Function(Vt).interpolate(Constant(Tsurf))
# Calculate hydrostatic fields
compressible_hydrostatic_balance(state, theta_b, rho0, solve_for_rho=True)
# make mean fields
rho_b = Function(Vr).assign(rho0)
# define perturbation
xc = L / 2
zc = 2000.
rc = 2000.
Tdash = 2.0
theta_pert = Function(Vt).interpolate(conditional(sqrt((x[0] - xc) ** 2 + (x[1] - zc) ** 2) > rc,
0.0, Tdash *
(cos(pi * sqrt(((x[0] - xc) / rc) ** 2 + ((x[1] - zc) / rc) ** 2) / 2.0))
** 2))
# define initial theta
theta0.assign(theta_b * (theta_pert / 300.0 + 1.0))
# find perturbed rho
gamma = TestFunction(Vr)
rho_trial = TrialFunction(Vr)
lhs = gamma * rho_trial * dx
rhs = gamma * (rho_b * theta_b / theta0) * dx
rho_problem = LinearVariationalProblem(lhs, rhs, rho0)
rho_solver = LinearVariationalSolver(rho_problem)
rho_solver.solve()
# initialise fields
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
# Background temperature
# Tbexp = Ts*(p/ps)**kappa/(Ts/G*((p/ps)**kappa - 1) + 1)
Tbexp = G*(1 - exp(N**2*z/g)) + Ts*exp(N**2*z/g)
Tb = Function(W_CG1).interpolate(Tbexp)
# Background potential temperature
thetabexp = Tb*(p_0/p)**kappa
thetab = Function(W_CG1).interpolate(thetabexp)
theta_b = Function(state.V[2]).interpolate(thetab)
rho_b = Function(state.V[1])
sin_tmp = sin(lat) * sin(phi_c)
cos_tmp = cos(lat) * cos(phi_c)
r = a*acos(sin_tmp + cos_tmp*cos(lon-lamda_c))
s = (d**2)/(d**2 + r**2)
theta_pert = deltaTheta*s*sin(2*np.pi*z/L_z)
theta0.interpolate(theta_b)
# Compute the balanced density
compressible_hydrostatic_balance(state, theta_b, rho_b, top=False,
pi_boundary=(p/p_0)**kappa)
theta0.interpolate(theta_pert)
theta0 += theta_b
rho0.assign(rho_b)
# Calculate hydrostatic fields
compressible_hydrostatic_balance(state, theta_b, rho0, solve_for_rho=True)
# make mean fields
rho_b = Function(Vr).assign(rho0)
# define perturbation
xc = L / 2
zc = 2000.
rc = 2000.
Tdash = 2.0
r = sqrt((x - xc)**2 + (z - zc)**2)
theta_pert = Function(Vt).interpolate(
conditional(r > rc, 0.0,
Tdash * (cos(pi * r / (2.0 * rc)))**2))
# define initial theta
theta0.assign(theta_b * (theta_pert / 300.0 + 1.0))
# find perturbed rho
gamma = TestFunction(Vr)
rho_trial = TrialFunction(Vr)
lhs = gamma * rho_trial * dx
rhs = gamma * (rho_b * theta_b / theta0) * dx
rho_problem = LinearVariationalProblem(lhs, rhs, rho0)
rho_solver = LinearVariationalSolver(rho_problem)
rho_solver.solve()
state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
return cosh(x) / sinh(x)
def Z(z):
return Bu * ((z / H) - 0.5)
def n():
return Bu**(-1) * sqrt(
(Bu * 0.5 - tanh(Bu * 0.5)) * (coth(Bu * 0.5) - Bu * 0.5))
a = -4.5
Bu = 0.5
b_exp = a * sqrt(Nsq) * (
-(1. - Bu * 0.5 * coth(Bu * 0.5)) * sinh(Z(z)) * cos(pi * (x - L) / L) -
n() * Bu * cosh(Z(z)) * sin(pi * (x - L) / L))
b_pert = Function(Vb).interpolate(b_exp)
# set total buoyancy
b0.project(b_b + b_pert)
# calculate hydrostatic pressure
p_b = Function(Vp)
incompressible_hydrostatic_balance(state, b_b, p_b)
incompressible_hydrostatic_balance(state, b0, p0)
# set x component of velocity
dbdy = parameters.dbdy
u = -dbdy / f * (z - H / 2)
def from_latlon(latitude,
longitude,
force_zone_number=None,
zone_info=False,
coords=None):
"""
Convert latitude-longitude coordinates to UTM, courtesy of Tobias Bieniek, 2012.
:arg latitude: northward anglular position, origin at the Equator.
:arg longitude: eastward angular position, with origin at the Greenwich Meridian.
:param force_zone_number: force coordinates to fall within a particular UTM zone.
:param zone_info: output zone letter and number.
:param coords: coordinate field of mesh (used to check validity of coordinates).
:return: UTM coordinate 4-tuple.
"""
lat_msg = 'latitude out of range (must be between 80 deg S and 84 deg N)'
lon_msg = 'longitude out of range (must be between 180 deg W and 180 deg E)'
if isinstance(latitude, ufl.indexed.Indexed):
from firedrake import sin, cos, sqrt
if coords is None:
if os.environ.get('WARNINGS', '0') != '0':
print_output("WARNING: Cannot check validity of coordinates.")
else:
minval, maxval = coords.dat.data[:, 0].min(
), coords.dat.data[:, 0].max()
if not (-80.0 <= minval and maxval <= 84.0):
raise OutOfRangeError(lon_msg)
minval, maxval = coords.dat.data[:, 1].min(
), coords.dat.data[:, 1].max()
if not (-180.0 <= minval and maxval <= 180.0):
raise OutOfRangeError(lat_msg)
elif isinstance(latitude, np.ndarray):
from numpy import sin, cos, sqrt
minval, maxval = longitude.min(), longitude.max()
if not (-180.0 <= minval and maxval <= 180.0):
raise OutOfRangeError(lon_msg)
minval, maxval = latitude.min(), latitude.max()
if not (-80.0 <= minval and maxval <= 84.0):
raise OutOfRangeError(lat_msg)
else:
from math import sin, cos, sqrt
if not -180.0 <= longitude <= 180.0:
raise OutOfRangeError(lon_msg)
if not -80.0 <= latitude <= 84.0:
raise OutOfRangeError(lat_msg)
lat_rad = radians(latitude)
lat_sin = sin(lat_rad)
lat_cos = cos(lat_rad)
lat_tan = lat_sin / lat_cos
lat_tan2 = lat_tan * lat_tan
lat_tan4 = lat_tan2 * lat_tan2
if force_zone_number is None:
zone_number = latlon_to_zone_number(latitude, longitude)
else:
zone_number = force_zone_number
lon_rad = radians(longitude)
central_lon_rad = radians(zone_number_to_central_longitude(zone_number))
n = R / sqrt(1 - E * lat_sin**2)
c = E_P2 * lat_cos**2
a = lat_cos * (lon_rad - central_lon_rad)
a2 = a * a
a3 = a2 * a
a4 = a3 * a
a5 = a4 * a
a6 = a5 * a
m = R * (M1 * lat_rad - M2 * sin(2 * lat_rad) + M3 * sin(4 * lat_rad) -
M4 * sin(6 * lat_rad))
easting = K0 * n * (
a + a3 / 6 * (1 - lat_tan2 + c) + a5 / 120 *
(5 - 18 * lat_tan2 + lat_tan4 + 72 * c - 58 * E_P2)) + 500000
northing = K0 * (m + n * lat_tan *
(a2 / 2 + a4 / 24 *
(5 - lat_tan2 + 9 * c + 4 * c**2) + a6 / 720 *
(61 - 58 * lat_tan2 + lat_tan4 + 600 * c - 330 * E_P2)))
if isinstance(latitude, ufl.indexed.Indexed):
if coords.dat.data[:, 1].min() < 0:
northing += 10000000
elif isinstance(latitude, np.ndarray):
if latitude.min() < 0:
northing += 10000000
else:
if latitude < 0:
northing += 10000000
if zone_info:
return easting, northing, zone_number, latitude_to_zone_letter(
latitude)
else:
return easting, northing
def to_latlon(easting,
northing,
zone_number,
zone_letter=None,
northern=None,
force_longitude=False,
coords=None):
"""
Convert UTM coordinates to latitude-longitude, courtesy of Tobias Bieniek, 2012 (with some
minor edits).
:arg easting: eastward-measured Cartesian geographic distance.
:arg northing: northward-measured Cartesian geographic distance.
:arg zone_number: UTM zone number (increasing eastward).
:param zone_letter: UTM zone letter (increasing alphabetically northward).
:param northern: specify northern or southern hemisphere.
:param coords: coordinate field of mesh (used to check validity of coordinates).
:return: latitude-longitude coordinate pair.
"""
if not zone_letter and northern is None:
raise ValueError('either zone_letter or northern needs to be set')
elif zone_letter and northern is not None:
raise ValueError('set either zone_letter or northern, but not both')
if not force_longitude:
if not 100000 <= easting < 1000000:
raise OutOfRangeError(
'easting {:f} out of range (must be between 100,000 m and 999,999 m)'
.format(easting))
msg = 'northing out of range (must be between 0 m and 10,000,000 m)'
if isinstance(northing, ufl.indexed.Indexed):
from firedrake import sin, cos, sqrt
if coords is None:
if os.environ.get('WARNINGS', '0') != '0':
print_output("WARNING: Cannot check validity of coordinates.")
else:
minval, maxval = coords.dat.data[:, 1].min(
), coords.dat.data[:, 1].max()
if not (0 <= minval and maxval <= 10000000):
raise OutOfRangeError(msg)
elif isinstance(northing, np.ndarray):
from numpy import sin, cos, sqrt
minval, maxval = northing.min(), northing.max()
if not (0 <= minval and maxval <= 10000000):
raise OutOfRangeError(msg)
else:
from math import sin, cos, sqrt
if not 0 <= northing <= 10000000:
raise OutOfRangeError(msg)
if not 1 <= zone_number <= 60:
raise OutOfRangeError(
'zone number out of range (must be between 1 and 60)')
if zone_letter:
zone_letter = zone_letter.upper()
if not 'C' <= zone_letter <= 'X' or zone_letter in ['I', 'O']:
raise OutOfRangeError(
'zone letter out of range (must be between C and X)')
northern = zone_letter >= 'N'
x = easting - 500000
y = northing
if not northern:
y -= 10000000
m = y / K0
mu = m / R / M1
p_rad = (mu + P2 * sin(2 * mu) + P3 * sin(4 * mu) + P4 * sin(6 * mu) +
P5 * sin(8 * mu))
p_sin = sin(p_rad)
p_sin2 = p_sin * p_sin
p_cos = cos(p_rad)
p_tan = p_sin / p_cos
p_tan2 = p_tan * p_tan
p_tan4 = p_tan2 * p_tan2
ep_sin = 1 - E * p_sin2
ep_sin_sqrt = sqrt(1 - E * p_sin2)
n = R / ep_sin_sqrt
r = (1 - E) / ep_sin
c = _E * p_cos**2
c2 = c * c
d = x / n / K0
d2 = d * d
d3 = d2 * d
d4 = d3 * d
d5 = d4 * d
d6 = d5 * d
latitude = (
p_rad - p_tan / r * (d2 / 2 - d4 / 24 *
(5 + 3 * p_tan2 + 10 * c - 4 * c2 - 9 * E_P2)) +
d6 / 720 *
(61 + 90 * p_tan2 + 298 * c + 45 * p_tan4 - 252 * E_P2 - 3 * c2))
longitude = (
d - d3 / 6 * (1 + 2 * p_tan2 + c) + d5 / 120 *
(5 - 2 * c + 28 * p_tan2 - 3 * c2 + 8 * E_P2 + 24 * p_tan4)) / p_cos
return degrees(latitude), degrees(
longitude) + zone_number_to_central_longitude(zone_number)
def Ctheta(theta):
return 0.25*(K**2)*(cos(theta)**8)*(5*cos(theta)**2 - 6)
def setup_tracer(dirname, hybridization):
# declare grid shape, with length L and height H
L = 1000.
H = 1000.
nlayers = int(H / 100.)
ncolumns = int(L / 100.)
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=10.0, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname+"/tracer",
dumpfreq=1,
dumplist=['u'],
perturbation_fields=['theta', 'rho'])
parameters = CompressibleParameters()
state = State(mesh, vertical_degree=1, horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist,
diagnostic_fields=[Difference('theta', 'tracer')])
# declare initial fields
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
# declare tracer field and a background field
tracer0 = state.fields("tracer", Vt)
# Isentropic background state
Tsurf = Constant(300.)
theta_b = Function(Vt).interpolate(Tsurf)
rho_b = Function(Vr)
# Calculate initial rho
compressible_hydrostatic_balance(state, theta_b, rho_b,
solve_for_rho=True)
# set up perturbation to theta
xc = 500.
zc = 350.
rc = 250.
x = SpatialCoordinate(mesh)
r = sqrt((x[0]-xc)**2 + (x[1]-zc)**2)
theta_pert = conditional(r > rc, 0., 0.25*(1. + cos((pi/rc)*r)))
theta0.interpolate(theta_b + theta_pert)
rho0.interpolate(rho_b)
tracer0.interpolate(theta0)
state.initialise([('u', u0),
('rho', rho0),
('theta', theta0),
('tracer', tracer0)])
state.set_reference_profiles([('rho', rho_b),
('theta', theta_b)])
# set up advection schemes
ueqn = EulerPoincare(state, Vu)
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
thetaeqn = SUPGAdvection(state, Vt,
supg_params={"dg_direction": "horizontal"},
equation_form="advective")
# build advection dictionary
advected_fields = []
advected_fields.append(("u", ThetaMethod(state, u0, ueqn)))
advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn)))
advected_fields.append(("theta", SSPRK3(state, theta0, thetaeqn)))
advected_fields.append(("tracer", SSPRK3(state, tracer0, thetaeqn)))
# Set up linear solver
if hybridization:
linear_solver = HybridizedCompressibleSolver(state)
else:
linear_solver = CompressibleSolver(state)
compressible_forcing = CompressibleForcing(state)
# build time stepper
stepper = CrankNicolson(state, advected_fields, linear_solver,
compressible_forcing)
return stepper, 100.0
# Get coordinates to pass to Dval function
D0 = Function(W0)
W = VectorFunctionSpace(mesh, D0.ufl_element())
X = interpolate(mesh.coordinates, W)
D0.dat.data[:] = Dval(X.dat.data_ro)
# Adjust mean value of initial D
C = Function(D0.function_space()).assign(Constant(1.0))
area = assemble(C * dx)
Dmean = assemble(D0 * dx) / area
D0 -= Dmean
D0 += Constant(D_mean)
# Dpert
D_p = Function(W0)
Dpert = D_bump * cos(theta) * exp(-(lamda / a)**2) * exp(-(
(theta_2 - theta) / b)**2)
D_p.interpolate(Dpert)
# Dexpr = Dbar + Dpert
Dexpr = D0 + D_p
bexpr = Constant(0)
# Set up functions
xn = Function(M)
un, Dn = xn.split()
un.rename('u')
Dn.rename('D')
fields = {'u': un, 'D': Dn}
# Energy field for output
En = Function(W0, name='Energy')
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])
new_coords = Function(Vc).interpolate(xexpr)
mesh = Mesh(new_coords)
# sponge function
W_DG = FunctionSpace(mesh, "DG", 2)
x, z = SpatialCoordinate(mesh)
zc = H-10000.
mubar = 0.15/dt
mu_top = conditional(z <= zc, 0.0, mubar*sin((pi/2.)*(z-zc)/(H-zc))**2)
mu = Function(W_DG).interpolate(mu_top)
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt)
theta_b = Function(Vt).interpolate(Tsurf)
rho_b = Function(Vr)
# Calculate hydrostatic Pi
compressible_hydrostatic_balance(state, theta_b, rho_b, solve_for_rho=True)
x = SpatialCoordinate(mesh)
a = 5.0e3
deltaTheta = 1.0e-2
xc = 0.5 * L
xr = 4000.
zc = 3000.
zr = 2000.
r = sqrt(((x[0] - xc) / xr)**2 + ((x[1] - zc) / zr)**2)
theta_pert = conditional(r > 1., 0., -7.5 * (1. + cos(pi * r)))
theta0.interpolate(theta_b + theta_pert)
water0.interpolate(theta_pert)
rho0.assign(rho_b)
state.initialise([('u', u0), ('rho', rho0), ('theta', theta0),
('water', water0)])
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")
watereqn = SUPGAdvection(state, Vt, equation_form="advective")
def setup_condens(dirname):
# declare grid shape, with length L and height H
L = 1000.
H = 1000.
nlayers = int(H / 100.)
ncolumns = int(L / 100.)
# make mesh
m = PeriodicIntervalMesh(ncolumns, L)
mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers))
x = SpatialCoordinate(mesh)
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=1.0, maxk=4, maxi=1)
output = OutputParameters(dirname=dirname + "/condens",
dumpfreq=1,
dumplist=['u'],
perturbation_fields=['theta', 'rho'])
parameters = CompressibleParameters()
state = State(mesh,
vertical_degree=1,
horizontal_degree=1,
family="CG",
timestepping=timestepping,
output=output,
parameters=parameters,
fieldlist=fieldlist,
diagnostic_fields=[Sum('water_v', 'water_c')])
# declare initial fields
u0 = state.fields("u")
rho0 = state.fields("rho")
theta0 = state.fields("theta")
# spaces
Vpsi = FunctionSpace(mesh, "CG", 2)
Vt = theta0.function_space()
Vr = rho0.function_space()
# make a gradperp
gradperp = lambda u: as_vector([-u.dx(1), u.dx(0)])
# declare tracer field and a background field
water_v0 = state.fields("water_v", Vt)
water_c0 = state.fields("water_c", Vt)
# Isentropic background state
Tsurf = Constant(300.)
theta_b = Function(Vt).interpolate(Tsurf)
rho_b = Function(Vr)
# Calculate initial rho
compressible_hydrostatic_balance(state, theta_b, rho_b, solve_for_rho=True)
# set up water_v
xc = 500.
zc = 350.
rc = 250.
r = sqrt((x[0] - xc)**2 + (x[1] - zc)**2)
w_expr = conditional(r > rc, 0., 0.25 * (1. + cos((pi / rc) * r)))
# set up velocity field
u_max = 10.0
psi_expr = ((-u_max * L / pi) * sin(2 * pi * x[0] / L) *
sin(pi * x[1] / L))
psi0 = Function(Vpsi).interpolate(psi_expr)
u0.project(gradperp(psi0))
theta0.interpolate(theta_b)
rho0.interpolate(rho_b)
water_v0.interpolate(w_expr)
state.initialise([('u', u0), ('rho', rho0), ('theta', theta0),
('water_v', water_v0), ('water_c', water_c0)])
state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
# set up advection schemes
rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity")
thetaeqn = SUPGAdvection(state,
Vt,
supg_params={"dg_direction": "horizontal"},
equation_form="advective")
# build advection dictionary
advected_fields = []
advected_fields.append(("u", NoAdvection(state, u0, None)))
advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn)))
advected_fields.append(("theta", SSPRK3(state, theta0, thetaeqn)))
advected_fields.append(("water_v", SSPRK3(state, water_v0, thetaeqn)))
advected_fields.append(("water_c", SSPRK3(state, water_c0, thetaeqn)))
physics_list = [Condensation(state)]
# build time stepper
stepper = AdvectionDiffusion(state,
advected_fields,
physics_list=physics_list)
return stepper, 5.0
def 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 main():
# If can't import firedrake, do nothing
#
# filename MUST include "firedrake" (i.e. match *firedrake*.py) in order
# to be run during CI
try:
import firedrake # noqa : F401
except ImportError:
return 0
from meshmode.interop.firedrake import build_connection_from_firedrake
from firedrake import (UnitSquareMesh, FunctionSpace, SpatialCoordinate,
Function, cos)
# Create a firedrake mesh and interpolate cos(x+y) onto it
fd_mesh = UnitSquareMesh(10, 10)
fd_fspace = FunctionSpace(fd_mesh, "DG", 2)
spatial_coord = SpatialCoordinate(fd_mesh)
fd_fntn = Function(fd_fspace).interpolate(cos(sum(spatial_coord)))
# Make connections
cl_ctx = cl.create_some_context()
queue = cl.CommandQueue(cl_ctx)
actx = PyOpenCLArrayContext(queue)
fd_connection = build_connection_from_firedrake(actx, fd_fspace)
fd_bdy_connection = \
build_connection_from_firedrake(actx,
fd_fspace,
restrict_to_boundary="on_boundary")
# Plot the meshmode meshes that the connections connect to
import matplotlib.pyplot as plt
from meshmode.mesh.visualization import draw_2d_mesh
fig, (ax1, ax2) = plt.subplots(1, 2)
ax1.set_title("FiredrakeConnection")
plt.sca(ax1)
draw_2d_mesh(fd_connection.discr.mesh,
draw_vertex_numbers=False,
draw_element_numbers=False,
set_bounding_box=True)
ax2.set_title("FiredrakeConnection 'on_boundary'")
plt.sca(ax2)
draw_2d_mesh(fd_bdy_connection.discr.mesh,
draw_vertex_numbers=False,
draw_element_numbers=False,
set_bounding_box=True)
plt.show()
# Plot fd_fntn using unrestricted FiredrakeConnection
from meshmode.discretization.visualization import make_visualizer
discr = fd_connection.discr
vis = make_visualizer(actx, discr, discr.groups[0].order + 3)
field = fd_connection.from_firedrake(fd_fntn, actx=actx)
fig = plt.figure()
ax1 = fig.add_subplot(1, 2, 1, projection="3d")
ax1.set_title("cos(x+y) in\nFiredrakeConnection")
vis.show_scalar_in_matplotlib_3d(field, do_show=False)
# Now repeat using FiredrakeConnection restricted to "on_boundary"
bdy_discr = fd_bdy_connection.discr
bdy_vis = make_visualizer(actx, bdy_discr, bdy_discr.groups[0].order + 3)
bdy_field = fd_bdy_connection.from_firedrake(fd_fntn, actx=actx)
ax2 = fig.add_subplot(1, 2, 2, projection="3d")
plt.sca(ax2)
ax2.set_title("cos(x+y) in\nFiredrakeConnection 'on_boundary'")
bdy_vis.show_scalar_in_matplotlib_3d(bdy_field, do_show=False)
import matplotlib.cm as cm
fig.colorbar(cm.ScalarMappable())
plt.show()
diagnostics=diagnostics,
fieldlist=fieldlist,
diagnostic_fields=diagnostic_fields)
# interpolate initial conditions
# Initial/current conditions
u0 = state.fields("u")
D0 = state.fields("D")
omega = 7.848e-6 # note lower-case, not the same as Omega
K = 7.848e-6
g = parameters.g
Omega = parameters.Omega
theta, lamda = latlon_coords(mesh)
u_zonal = R*omega*cos(theta) + R*K*(cos(theta)**3)*(4*sin(theta)**2 - cos(theta)**2)*cos(4*lamda)
u_merid = -R*K*4*(cos(theta)**3)*sin(theta)*sin(4*lamda)
uexpr = sphere_to_cartesian(mesh, u_zonal, u_merid)
def Atheta(theta):
return 0.5*omega*(2*Omega + omega)*cos(theta)**2 + 0.25*(K**2)*(cos(theta)**8)*(5*cos(theta)**2 + 26 - 32/(cos(theta)**2))
def Btheta(theta):
return (2*(Omega + omega)*K/30)*(cos(theta)**4)*(26 - 25*cos(theta)**2)
def Ctheta(theta):
return 0.25*(K**2)*(cos(theta)**8)*(5*cos(theta)**2 - 6)
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
def Atheta(theta):
return 0.5*omega*(2*Omega + omega)*cos(theta)**2 + 0.25*(K**2)*(cos(theta)**8)*(5*cos(theta)**2 + 26 - 32/(cos(theta)**2))
def run_profliler(hybridization,
model_degree,
model_family,
mesh_degree,
cfl,
refinements,
layers,
debug,
rtol,
flexsolver=True,
stronger_smoother=False,
suppress_data_output=False):
nlayers = layers # Number of vertical layers
refinements = refinements # Number of horiz. cells = 20*(4^refinements)
hybrid = bool(hybridization)
# Set up problem parameters
parameters = CompressibleParameters()
a_ref = 6.37122e6 # Radius of the Earth (m)
X = 1.0 # Reduced-size Earth reduction factor
a = a_ref / X # Scaled radius of planet (m)
g = parameters.g # Acceleration due to gravity (m/s^2)
N = parameters.N # Brunt-Vaisala frequency (1/s)
p_0 = parameters.p_0 # Reference pressure (Pa, not hPa)
c_p = parameters.cp # SHC of dry air at const. pressure (J/kg/K)
R_d = parameters.R_d # Gas constant for dry air (J/kg/K)
kappa = parameters.kappa # R_d/c_p
T_eq = 300.0 # Isothermal atmospheric temperature (K)
p_eq = 1000.0 * 100.0 # Reference surface pressure at the equator
u_0 = 20.0 # Maximum amplitude of the zonal wind (m/s)
d = 5000.0 # Width parameter for Theta'
lamda_c = 2.0 * np.pi / 3.0 # Longitudinal centerpoint of Theta'
phi_c = 0.0 # Lat. centerpoint of Theta' (equator)
deltaTheta = 1.0 # Maximum amplitude of Theta' (K)
L_z = 20000.0 # Vert. wave length of the Theta' perturb.
gamma = (1 - kappa) / kappa
cs = sqrt(c_p * T_eq / gamma) # Speed of sound in an air parcel
if model_family == "RTCF":
# Cubed-sphere mesh
m = CubedSphereMesh(radius=a,
refinement_level=refinements,
degree=mesh_degree)
elif model_family == "RT" or model_family == "BDFM":
m = IcosahedralSphereMesh(radius=a,
refinement_level=refinements,
degree=mesh_degree)
else:
raise ValueError("Unknown family: %s" % model_family)
cell_vs = interpolate(CellVolume(m), FunctionSpace(m, "DG", 0))
a_max = fmax(cell_vs)
dx_max = sqrt(a_max)
u_max = u_0
PETSc.Sys.Print("\nDetermining Dt from specified horizontal CFL: %s" % cfl)
dt = int(cfl * (dx_max / cs))
# Height position of the model top (m)
z_top = 1.0e4
deltaz = z_top / nlayers
vertical_cfl = dt * (cs / deltaz)
PETSc.Sys.Print("""
Problem parameters:\n
Profiling linear solver for the compressible Euler equations.\n
Speed of sound in compressible atmosphere: %s,\n
Hybridized compressible solver: %s,\n
Model degree: %s,\n
Model discretization: %s,\n
Mesh degree: %s,\n
Horizontal refinements: %s,\n
Vertical layers: %s,\n
Dx (max, m): %s,\n
Dz (m): %s,\n
Dt (s): %s,\n
horizontal CFL: %s,\n
vertical CFL: %s.
""" % (cs, hybrid, model_degree, model_family, mesh_degree, refinements,
nlayers, dx_max, deltaz, dt, cfl, vertical_cfl))
# Build volume mesh
mesh = ExtrudedMesh(m,
layers=nlayers,
layer_height=deltaz,
extrusion_type="radial")
x = SpatialCoordinate(mesh)
# Create polar coordinates:
# Since we use a CG1 field, this is constant on layers
W_Q1 = FunctionSpace(mesh, "CG", 1)
z_expr = sqrt(x[0] * x[0] + x[1] * x[1] + x[2] * x[2]) - a
z = Function(W_Q1).interpolate(z_expr)
lat_expr = asin(x[2] / sqrt(x[0] * x[0] + x[1] * x[1] + x[2] * x[2]))
lat = Function(W_Q1).interpolate(lat_expr)
lon = Function(W_Q1).interpolate(atan_2(x[1], x[0]))
fieldlist = ['u', 'rho', 'theta']
timestepping = TimesteppingParameters(dt=dt, maxk=1, maxi=1)
dirname = 'meanflow_ref'
if hybrid:
dirname += '_hybridization'
# No output
output = OutputParameters(dumpfreq=3600,
dirname=dirname,
perturbation_fields=['theta', 'rho'],
dump_vtus=False,
dump_diagnostics=False,
checkpoint=False,
log_level='INFO')
diagnostics = Diagnostics(*fieldlist)
state = State(mesh,
vertical_degree=model_degree,
horizontal_degree=model_degree,
family=model_family,
timestepping=timestepping,
output=output,
parameters=parameters,
diagnostics=diagnostics,
fieldlist=fieldlist)
# Initial conditions
u0 = state.fields.u
theta0 = state.fields.theta
rho0 = state.fields.rho
# spaces
Vu = u0.function_space()
Vt = theta0.function_space()
Vr = rho0.function_space()
x = SpatialCoordinate(mesh)
# Random velocity field
CG2 = VectorFunctionSpace(mesh, "CG", 2)
urand = Function(CG2)
urand.dat.data[:] += np.random.randn(*urand.dat.data.shape)
u0.project(urand)
# Surface temperature
G = g**2 / (N**2 * c_p)
Ts_expr = G + (T_eq - G) * exp(-(u_max * N**2 / (4 * g * g)) * u_max *
(cos(2.0 * lat) - 1.0))
Ts = Function(W_Q1).interpolate(Ts_expr)
# Surface pressure
ps_expr = p_eq * exp((u_max / (4.0 * G * R_d)) * u_max *
(cos(2.0 * lat) - 1.0)) * (Ts / T_eq)**(1.0 / kappa)
ps = Function(W_Q1).interpolate(ps_expr)
# Background pressure
p_expr = ps * (1 + G / Ts * (exp(-N**2 * z / g) - 1))**(1.0 / kappa)
p = Function(W_Q1).interpolate(p_expr)
# Background temperature
Tb_expr = G * (1 - exp(N**2 * z / g)) + Ts * exp(N**2 * z / g)
Tb = Function(W_Q1).interpolate(Tb_expr)
# Background potential temperature
thetab_expr = Tb * (p_0 / p)**kappa
thetab = Function(W_Q1).interpolate(thetab_expr)
theta_b = Function(theta0.function_space()).interpolate(thetab)
rho_b = Function(rho0.function_space())
sin_tmp = sin(lat) * sin(phi_c)
cos_tmp = cos(lat) * cos(phi_c)
r = a * acos(sin_tmp + cos_tmp * cos(lon - lamda_c))
s = (d**2) / (d**2 + r**2)
theta_pert = deltaTheta * s * sin(2 * np.pi * z / L_z)
theta0.interpolate(theta_pert)
# Compute the balanced density
PETSc.Sys.Print("Computing balanced density field...\n")
# Use vert. hybridization preconditioner for initialization
pi_params = {
'ksp_type': 'preonly',
'pc_type': 'python',
'mat_type': 'matfree',
'pc_python_type': 'gusto.VerticalHybridizationPC',
'vert_hybridization': {
'ksp_type': 'gmres',
'pc_type': 'gamg',
'pc_gamg_sym_graph': True,
'ksp_rtol': 1e-12,
'ksp_atol': 1e-12,
'mg_levels': {
'ksp_type': 'richardson',
'ksp_max_it': 3,
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
}
}
}
if debug:
pi_params['vert_hybridization']['ksp_monitor_true_residual'] = None
compressible_hydrostatic_balance(state,
theta_b,
rho_b,
top=False,
pi_boundary=(p / p_0)**kappa,
solve_for_rho=False,
params=pi_params)
# Random potential temperature perturbation
theta0.assign(0.0)
theta0.dat.data[:] += np.random.randn(len(theta0.dat.data))
# Random density field
rho0.assign(0.0)
rho0.dat.data[:] += np.random.randn(len(rho0.dat.data))
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")
thetaeqn = SUPGAdvection(state, Vt, equation_form="advective")
advected_fields = []
advected_fields.append(("u", ThetaMethod(state, u0, ueqn)))
advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn, subcycles=2)))
advected_fields.append(
("theta", SSPRK3(state, theta0, thetaeqn, subcycles=2)))
# Set up linear solver
if hybrid:
outer_solver_type = "Hybrid_SCPC"
PETSc.Sys.Print("""
Setting up hybridized solver on the traces.""")
if flexsolver:
inner_solver_type = "fgmres_gamg_gmres_smoother"
inner_parameters = {
'ksp_type': 'fgmres',
'ksp_rtol': rtol,
'ksp_max_it': 500,
'ksp_gmres_restart': 30,
'pc_type': 'gamg',
'pc_gamg_sym_graph': None,
'mg_levels': {
'ksp_type': 'gmres',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu',
'ksp_max_it': 3
}
}
else:
inner_solver_type = "fgmres_ml_richardson"
inner_parameters = {
'ksp_type': 'fgmres',
'ksp_rtol': rtol,
'ksp_max_it': 500,
'ksp_gmres_restart': 30,
'pc_type': 'ml',
'pc_mg_cycles': 1,
'pc_ml_maxNlevels': 25,
'mg_levels': {
'ksp_type': 'richardson',
'ksp_richardson_scale': 0.8,
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu',
'ksp_max_it': 3
}
}
if stronger_smoother:
inner_parameters['mg_levels']['ksp_max_it'] = 5
inner_solver_type += "_stronger"
if debug:
PETSc.Sys.Print("""Debugging on.""")
inner_parameters['ksp_monitor_true_residual'] = None
PETSc.Sys.Print("Inner solver: %s" % inner_solver_type)
# Use Firedrake static condensation interface
solver_parameters = {
'mat_type': 'matfree',
'pmat_type': 'matfree',
'ksp_type': 'preonly',
'pc_type': 'python',
'pc_python_type': 'firedrake.SCPC',
'pc_sc_eliminate_fields': '0, 1',
'condensed_field': inner_parameters
}
linear_solver = HybridizedCompressibleSolver(
state,
solver_parameters=solver_parameters,
overwrite_solver_parameters=True)
else:
outer_solver_type = "gmres_SchurPC"
PETSc.Sys.Print("""
Setting up GCR fieldsplit solver with Schur complement PC.""")
solver_parameters = {
'pc_type': 'fieldsplit',
'pc_fieldsplit_type': 'schur',
'ksp_type': 'fgmres',
'ksp_max_it': 100,
'ksp_rtol': rtol,
'pc_fieldsplit_schur_fact_type': 'FULL',
'pc_fieldsplit_schur_precondition': 'selfp',
'fieldsplit_0': {
'ksp_type': 'preonly',
'pc_type': 'bjacobi',
'sub_pc_type': 'ilu'
},
'fieldsplit_1': {
'ksp_type': 'preonly',
'ksp_max_it': 30,
'ksp_monitor_true_residual': None,
'pc_type': 'hypre',
'pc_hypre_type': 'boomeramg',
'pc_hypre_boomeramg_max_iter': 1,
'pc_hypre_boomeramg_agg_nl': 0,
'pc_hypre_boomeramg_coarsen_type': 'Falgout',
'pc_hypre_boomeramg_smooth_type': 'Euclid',
'pc_hypre_boomeramg_eu_bj': 1,
'pc_hypre_boomeramg_interptype': 'classical',
'pc_hypre_boomeramg_P_max': 0,
'pc_hypre_boomeramg_agg_nl': 0,
'pc_hypre_boomeramg_strong_threshold': 0.25,
'pc_hypre_boomeramg_max_levels': 25,
'pc_hypre_boomeramg_no_CF': False
}
}
inner_solver_type = "hypre"
if debug:
solver_parameters['ksp_monitor_true_residual'] = None
linear_solver = CompressibleSolver(state,
solver_parameters=solver_parameters,
overwrite_solver_parameters=True)
# Set up forcing
compressible_forcing = CompressibleForcing(state)
param_info = ParameterInfo(dt=dt,
deltax=dx_max,
deltaz=deltaz,
horizontal_courant=cfl,
vertical_courant=vertical_cfl,
family=model_family,
model_degree=model_degree,
mesh_degree=mesh_degree,
solver_type=outer_solver_type,
inner_solver_type=inner_solver_type)
# Build profiler
profiler = Profiler(parameterinfo=param_info,
state=state,
advected_fields=advected_fields,
linear_solver=linear_solver,
forcing=compressible_forcing,
suppress_data_output=suppress_data_output)
PETSc.Sys.Print("Starting profiler...\n")
profiler.run(t=0, tmax=dt)
def Btheta(theta):
return (2*(Omega + omega)*K/30)*(cos(theta)**4)*(26 - 25*cos(theta)**2)