def latlon_coords(mesh):
x0, y0, z0 = SpatialCoordinate(mesh)
unsafe = z0 / sqrt(x0 * x0 + y0 * y0 + z0 * z0)
safe = Min(Max(unsafe, -1.0), 1.0) # avoid silly roundoff errors
theta = asin(safe) # latitude
lamda = atan_2(y0, x0) # longitude
return theta, lamda
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
# Coriolis
fexpr = 2*Omega*x[2]/R
V = FunctionSpace(mesh, "CG", 1)
f = state.fields("coriolis", V)
f.interpolate(fexpr) # Coriolis frequency (1/s)
b = state.fields("topography", D0.function_space())
b.interpolate(bexpr)
u0.project(uexpr)
D0.interpolate(Dexpr)
state.initialise([('u', u0),
def latlon_coords(mesh):
"""Compute latitude-longitude coordinates given Cartesian ones"""
x0, y0, z0 = SpatialCoordinate(mesh)
unsafe = z0 / sqrt(x0 * x0 + y0 * y0 + z0 * z0)
safe = Min(Max(unsafe, -1.0), 1.0) # avoid silly roundoff errors
theta = asin(safe) # latitude
lamda = atan_2(y0, x0) # longitude
return theta, lamda
theta, lamda = latlon_coords(mesh)
R0, lamda_c, theta_c = pi / 9., -pi / 2., pi / 6.
R0sq, Rsq = R0**2, R**2
lsq, thsq = (lamda - lamda_c)**2, (theta - theta_c)**2
r = sqrt(Min(R0sq, lsq + thsq))
bexpr = 2000 * (1 - r / R0)
uexpr = u_0 * as_vector([-x[1], x[0], 0.0]) / R
Dexpr = H - ((R * Omega * u_0 + 0.5 * u_0**2) * x[2]**2 / Rsq) / g - bexpr
# Build function spaces
degree = 2
family = ("DG", "BDM", "CG")
W0 = FunctionSpace(mesh, family[0], degree - 1, family[0])
W1_elt = FiniteElement(family[1], triangle, degree)
W1 = FunctionSpace(mesh, W1_elt, name="HDiv")
W2 = FunctionSpace(mesh, family[2], degree + 1)
M = MixedFunctionSpace((W1, W0))
# Set up functions