def _parameter_conversion(self, dataset):
# NOTE We are here once again convert from (x,y,z) parametrization to (r,\theta,\phi) parametrization
for i in [1,2]:
# Magnitude a_i
dataset["a_{}".format(i)] = xp.sqrt(dataset["spin_{}x".format(i)]**2 + dataset["spin_{}y".format(i)]**2 + dataset["spin_{}z".format(i)]**2)
# Tilt angle tilt_i
dataset["tilt_{}".format(i)] = xp.arccos(dataset["spin_{}z".format(i)]/dataset["a_{}".format(i)])
def angle(self, A, B):
fp = A * B
fp = cp.sum(fp)
norm_a = cp.sqrt(cp.sum(A * A))
norm_b = cp.sqrt(cp.sum(B * B))
cos_theta = fp / (norm_a * norm_b)
return cp.arccos(cos_theta)
def acos(x: Array, /) -> Array:
"""
Array API compatible wrapper for :py:func:`np.arccos <numpy.arccos>`.
See its docstring for more information.
"""
if x.dtype not in _floating_dtypes:
raise TypeError("Only floating-point dtypes are allowed in acos")
return Array._new(np.arccos(x._array))
def nvst_ringotf_cal(self, fre):
cp.cuda.Device(0).use
rh0 = fre
if (rh0 < 0) or (rh0 > 1):
return 0.0
if (self.diaratio < 0) or (self.diaratio > 1):
return 0.0
R = 0.5
A = self.diaratio
if (rh0 < 2.0 * R):
c = 2.0 * cp.arccos(
rh0 / (R * 2.0)) * R**2 - rh0 * cp.sqrt(R**2 - rh0**2 / 4.0)
else:
c = 0.0
if (rh0 < 2.0 * R * A):
E = 2.0 * cp.arccos(rh0 /
(A * R * 2.0)) * (R * A)**2 - rh0 * cp.sqrt(
(R * A)**2 - rh0**2 / 4.0)
else:
E = 0.0
if (rh0 <= R + R * A) and (rh0 > R - R * A):
s1 = 0.5 * cp.arccos(((R * A)**2 + rh0**2 - R**2) /
(2.0 * A * R * rh0)) * (R * A)**2
s2 = 0.5 * cp.arccos(
(rh0**2 + R**2 - (R * A)**2) / (2.0 * R * rh0)) * R**2
s3 = 0.5 * cp.sin(
cp.arccos(((R * A)**2 + rh0**2 - R**2) /
(2.0 * A * R * rh0))) * A * R * rh0
D = s1 + s2 - s3
D = D * 2
elif (rh0 <= R - A * R):
D = cp.pi * (A * R)**2
else:
D = 0.0
H = (C + E - 2 * D) / (cp.pi * R**2)
if rh0 == 0:
H = 1 - A**2
return float(H)
def random_in_unit_sphere_list(size: int) -> Vec3List:
u = random_float_list(size)
v = random_float_list(size)
theta = u * 2 * cp.pi
phi = cp.arccos(2 * v - 1)
r = cp.cbrt(random_float_list(size))
sinTheta = cp.sin(theta)
cosTheta = cp.cos(theta)
sinPhi = cp.sin(phi)
cosPhi = cp.cos(phi)
x = r * sinPhi * cosTheta
y = r * sinPhi * sinTheta
z = r * cosPhi
return Vec3List(cp.transpose(cp.array([x, y, z])))
def random_in_unit_sphere():
u = random_float()
v = random_float()
theta = u * 2 * cp.pi
phi = cp.arccos(2 * v - 1)
r = cp.cbrt(random_float())
sinTheta = cp.sin(theta)
cosTheta = cp.cos(theta)
sinPhi = cp.sin(phi)
cosPhi = cp.cos(phi)
x = r * sinPhi * cosTheta
y = r * sinPhi * sinTheta
z = r * cosPhi
return Vec3(x, y, z)
def random_in_unit_sphere() -> Vec3:
"""
This method is modified from:
https://karthikkaranth.me/blog/generating-random-points-in-a-sphere/#better-choice-of-spherical-coordinates
"""
u = random_float()
v = random_float()
theta = u * 2 * cp.pi
phi = cp.arccos(2 * v - 1)
r = cp.cbrt(random_float())
sinTheta = cp.sin(theta)
cosTheta = cp.cos(theta)
sinPhi = cp.sin(phi)
cosPhi = cp.cos(phi)
x = r * sinPhi * cosTheta
y = r * sinPhi * sinTheta
z = r * cosPhi
return Vec3(x, y, z)
def negative_sampling(self, embeddings, is_test):
if is_test:
prepare_func = self.prepare_test_negative_sampling()
else:
prepare_func = self.prepare_train_negative_sampling()
for _ in range(self.model.negative_sample_num):
negative_sampling_matrix = prepare_func
negative_embeddings = negative_sampling_matrix.dot(embeddings)
norm_embeddings = cp.linalg.norm(embeddings, axis=1).reshape(
(embeddings.shape[0], 1))
target_negative_cos = cp.cos(embeddings, negative_embeddings)
new_embeddings = embeddings + (embeddings - negative_embeddings) \
* self.model.target_negative_weight * cp.cos(cp.arccos(target_negative_cos) / 2)
norm_new_embeddings = cp.linalg.norm(new_embeddings,
axis=1).reshape(
(new_embeddings.shape[0],
1))
embeddings = new_embeddings / norm_new_embeddings * norm_embeddings
def IsoScatter(photon_state):
####################################################################################
# This function takes the photons and scatters them off the dust in the atmosphere
# using an isotropic scattering method.
#
# photon_state -- A CuPy array consisting of [x,y,z,phi,theta,lambda] for each photon.
####################################################################################
# Theta is the polar angle with respect to the lab frame z-axis.
# theta = 2*cp.pi*cp.random.rand(len(photon_state))
# Generate a uniform distribution for mu
mu = 2 * cp.random.rand(len(photon_state)) - 1
# Extract the angle theta from that uniform distribution
theta = cp.arccos(mu)
# Phi is the azimuthal angle, scattering should always be isotropic in phi.
phi = 2 * cp.pi * cp.random.rand(len(photon_state))
return theta, phi
def arcsec__(z):
return cp.arccos(1.0/z)
def eig_special_3d(S, full=False):
"""Eigensolution for symmetric real 3-by-3 matrices.
Arguments:
S: array-like
A floating point array with shape ``(6, ...)`` containing structure tensor.
Use float64 to avoid numerical errors. When using lower precision, ensure
that the values of S are not very small/large. Pass cupy.ndarray to avoid copying S.
full: bool, optional
A flag indicating that all three eigenvalues should be returned.
Returns:
val: cupy.ndarray
An array with shape ``(3, ...)`` containing sorted eigenvalues
vec: cupy.ndarray
An array with shape ``(3, ...)`` containing eigenvector corresponding to
the smallest eigenvalue. If full, vec has shape ``(3, 3, ...)`` and contains
all three eigenvectors.
More:
An analytic solution of eigenvalue problem for real symmetric matrix,
using an affine transformation and a trigonometric solution of third
order polynomial. See https://en.wikipedia.org/wiki/Eigenvalue_algorithm
which refers to Smith's algorithm https://dl.acm.org/citation.cfm?id=366316.
Authors: [email protected], 2019; [email protected], 2019-2020
"""
S = np.asarray(S)
# Check data type. Must be floating point.
if not np.issubdtype(S.dtype, np.floating):
raise ValueError('S must be floating point type.')
# Flatten S.
input_shape = S.shape
S = S.reshape(6, -1)
# Create v vector.
v = np.array([[2 * np.pi / 3], [4 * np.pi / 3]], dtype=S.dtype)
# Computing eigenvalues.
# Allocate vec and val. We will use them for intermediate computations as well.
if full:
val = np.empty((3, ) + S.shape[1:], dtype=S.dtype)
vec = np.empty((9, ) + S.shape[1:], dtype=S.dtype)
tmp = np.empty((4, ) + S.shape[1:], dtype=S.dtype)
B03 = val
B36 = vec[:3]
else:
val = np.empty((3, ) + S.shape[1:], dtype=S.dtype)
vec = np.empty((3, ) + S.shape[1:], dtype=S.dtype)
tmp = np.empty((4, ) + S.shape[1:], dtype=S.dtype)
B03 = val
B36 = vec
# Views for B.
B0 = B03[0]
B1 = B03[1]
B2 = B03[2]
B3 = B36[0]
B4 = B36[1]
B5 = B36[2]
# Compute q, mean of diagonal. We need to use q multiple times later.
# Using np.mean has precision issues.
q = np.add(S[0], S[1], out=tmp[0])
q += S[2]
q /= 3
# Compute S minus q. Insert it directly into B where it'll stay.
Sq = np.subtract(S[:3], q, out=B03)
# Compute s, off-diagonal elements. Store in part of B not yet used.
s = np.sum(np.multiply(S[3:], S[3:], out=B36), axis=0, out=tmp[1])
s *= 2
# Compute p.
p = np.sum(np.multiply(Sq, Sq, out=B36), axis=0, out=tmp[2])
del Sq # Last use of Sq.
p += s
p *= 1 / 6
np.sqrt(p, out=p)
# Compute inverse p, while avoiding 0 division.
# Reuse s allocation and delete s to ensure we don't efter it's been reused.
p_inv = s
del s
non_zero_p_mask = p == 0
np.divide(1, p, out=p_inv)
p_inv[non_zero_p_mask] = 0
# Compute B. First part is already filled.
B03 *= p_inv
np.multiply(S[3:], p_inv, out=B36)
# Compute d, determinant of B.
d = np.prod(B03, axis=0, out=tmp[3])
# Reuse allocation for p_inv and delete variable.
d_tmp = p_inv
del p_inv
# Computation of d.
np.multiply(B2, B3, d_tmp)
d_tmp *= B3
d -= d_tmp
np.multiply(B4, B4, out=d_tmp)
d_tmp *= B1
d -= d_tmp
np.prod(B36, axis=0, out=d_tmp)
d_tmp *= 2
d += d_tmp
np.multiply(B5, B5, out=d_tmp)
d_tmp *= B0
d -= d_tmp
d *= 0.5
# Ensure -1 <= d/2 <= 1.
np.clip(d, -1, 1, out=d)
# Compute phi. Beware that we reuse d variable!
phi = d
del d
phi = np.arccos(phi, out=phi)
phi /= 3
# Compute val, ordered eigenvalues. Resuing B allocation.
del B03, B36, B0, B1, B2, B3, B4, B5
np.add(v, phi[np.newaxis], out=val[:2])
val[2] = phi
np.cos(val, out=val)
p *= 2
val *= p
val += q
# Remove all variable using tmp allocation.
del q
del p
del phi
del d_tmp
# Computing eigenvectors -- either only one or all three.
if full:
l = val
vec = vec.reshape(3, 3, -1)
vec_tmp = tmp[:3]
else:
l = val[0]
vec_tmp = tmp[2]
# Compute vec. The tmp variable can be reused.
# u = S[4] * S[5] - (S[2] - l) * S[3]
u = np.subtract(S[2], l, out=vec[0])
np.multiply(u, S[3], out=u)
u_tmp = np.multiply(S[4], S[5], out=tmp[3])
np.subtract(u_tmp, u, out=u)
# Put values of u into vector 2 aswell.
# v = S[3] * S[5] - (S[1] - l) * S[4]
v = np.subtract(S[1], l, out=vec_tmp)
np.multiply(v, S[4], out=v)
v_tmp = np.multiply(S[3], S[5], out=tmp[3])
np.subtract(v_tmp, v, out=v)
# w = S[3] * S[4] - (S[0] - l) * S[5]
w = np.subtract(S[0], l, out=vec[2])
np.multiply(w, S[5], out=w)
w_tmp = np.multiply(S[3], S[4], out=tmp[3])
np.subtract(w_tmp, w, out=w)
vec[1] = u
np.multiply(u, v, out=vec[0])
u = vec[1]
np.multiply(u, w, out=vec[1])
np.multiply(v, w, out=vec[2])
# Remove u, v, w and l variables.
del u
del v
del w
del l
# Normalizing -- depends on number of vectors.
if full:
# vec is [x1 x2 x3, y1 y2 y3, z1 z2 z3]
l = np.sum(np.square(vec), axis=0, out=vec_tmp)[:, np.newaxis]
vec = np.swapaxes(vec, 0, 1)
else:
# vec is [x1 y1 z1] = v1
l = np.sum(np.square(vec, out=tmp[:3]), axis=0, out=vec_tmp)
cp.rsqrt(l, out=l)
vec *= l
return val.reshape(val.shape[:-1] +
input_shape[1:]), vec.reshape(vec.shape[:-1] +
input_shape[1:])
def start_from_function(self, vc, g, c):
latmin = xp.min(g.latCell[:])
latmax = xp.max(g.latCell[:])
lonmin = xp.min(g.lonCell[:])
lonmax = xp.max(g.lonCell[:])
latmid = 0.5 * (latmin + latmax)
latwidth = latmax - latmin
lonmid = 0.5 * (lonmin + lonmax)
lonwidth = lonmax - lonmin
pi = np.pi
sin = np.sin
exp = np.exp
r = c.sphere_radius
if c.test_case == 1:
a = c.sphere_radius
R = a / 3
u0 = 2 * np.pi * a / (12 * 86400)
h0 = 1000.
lon_c = .5 * np.pi
lat_c = 0.
r = a*xp.arccos(xp.sin(lat_c)*xp.sin(g.latCell[:]) + \
xp.cos(lat_c)*xp.cos(g.latCell[:])*xp.cos(g.lonCell[:]-lon_c))
self.thickness[:] = xp.where(
r <= R, 0.25 * h0 * (1 + xp.cos(np.pi * r / R)), 0.) + h0
self.vorticity[:] = 2 * u0 / a * np.sin(g.latCell[:])
self.divergence[:] = 0.
#self.compute_diagnostics(g, c)
self.SS0 = xp.sum((self.thickness + g.bottomTopographyCell) *
g.areaCell) / xp.sum(g.areaCell)
elif c.test_case == 2:
# SWSTC #2, with a stationary analytic solution
a = c.sphere_radius
u0 = 2 * np.pi * a / (12 * 86400)
gh0 = 2.94e4
gh = xp.sin(g.latCell[:])**2
gh = -(a * c.Omega0 * u0 + 0.5 * u0 * u0) * gh + gh0
self.thickness[:] = gh / c.gravity
h0 = gh0 / c.gravity
self.vorticity[:] = 2 * u0 / a * xp.sin(g.latCell[:])
self.divergence[:] = 0.
self.psi_cell[:] = -a * h0 * u0 * xp.sin(g.latCell[:])
self.psi_cell[:] += a * u0 / c.gravity * (
a * c.Omega0 * u0 + 0.5 * u0**2) * (xp.sin(
g.latCell[:]))**3 / 3.
self.psi_cell -= self.psi_cell[0]
self.phi_cell[:] = 0.
self.SS0 = xp.sum((self.thickness + g.bottomTopographyCell) *
g.areaCell) / xp.sum(g.areaCell)
## For debugging ##
if False:
# To check the consistency between psi_cell and vorticity
self.thickness_edge = vc.cell2edge(self.thickness)
self.psi_vertex[:] = vc.cell2vertex(self.psi_cell)
nVelocity = vc.discrete_grad_n(self.phi_cell)
nVelocity -= vc.discrete_grad_td(self.psi_vertex)
nVelocity /= self.thickness_edge
vorticity1 = vc.vertex2cell(vc.discrete_curl_t(nVelocity))
err = vorticity1 - self.vorticity
print("vorticity computed using normal vel.")
print("relative error = %e" % (xp.sqrt(
xp.sum(err**2 * g.areaCell) /
xp.sum(self.vorticity**2 * g.areaCell))))
self.phi_vertex[:] = vc.cell2vertex(self.phi_cell)
tVelocity = vc.discrete_grad_n(self.psi_cell)
tVelocity += vc.discrete_grad_tn(self.phi_vertex)
tVelocity /= self.thickness_edge
vorticity2 = vc.discrete_curl_v(tVelocity)
err = vorticity2 - self.vorticity
print("vorticity computed using tang vel.")
print("relative error = %e" % (xp.sqrt(
xp.sum(err**2 * g.areaCell) /
xp.sum(self.vorticity**2 * g.areaCell))))
err = 0.5 * (vorticity1 + vorticity2) - self.vorticity
print("vorticity computed using both normal and tang vel.")
print("relative error = %e" % (xp.sqrt(
xp.sum(err**2 * g.areaCell) /
xp.sum(self.vorticity**2 * g.areaCell))))
raise ValueError(
"Testing the consistency between streamfunction and vorticity."
)
## End of debugging ##
elif c.test_case == 5:
#SWSTC #5: zonal flow over a mountain topography
a = c.sphere_radius
u0 = 20.
h0 = 5960.
gh = c.gravity * h0 - xp.sin(
g.latCell[:])**2 * (a * c.Omega0 * u0 + 0.5 * u0 * u0)
h = gh / c.gravity
# Define the mountain topography
h_s0 = 2000.
R = np.pi / 9
lat_c = np.pi / 6.
lon_c = -.5 * np.pi
r = xp.sqrt((g.latCell[:] - lat_c)**2 + (g.lonCell[:] - lon_c)**2)
r = xp.where(r < R, r, R)
g.bottomTopographyCell[:] = h_s0 * (1 - r / R)
self.thickness[:] = h[:] - g.bottomTopographyCell[:]
self.vorticity[:] = 2 * u0 / a * xp.sin(g.latCell[:])
self.divergence[:] = 0.
self.SS0 = xp.sum((self.thickness + g.bottomTopographyCell) *
g.areaCell) / xp.sum(g.areaCell)
self.curlWind_cell[:] = 0.
self.divWind_cell[:] = 0.
elif c.test_case == 6:
# Setup shallow water test case 6: Rossby-Haurwitz Wave
#
# Reference: Williamson, D.L., et al., "A Standard Test Set for Numerical
# Approximations to the Shallow Water Equations in Spherical
# Geometry" J. of Comp. Phys., 102, pp. 211--224
a = c.sphere_radius
h0 = 8000.
w = 7.848e-6
K = 7.848e-6
R = 4
cos_lat = xp.cos(g.latCell)
A = 0.5*w*(2*c.Omega0 + w)*cos_lat**2 + \
0.25*K**2*cos_lat**(2*R) * ( \
(R+1)*cos_lat**2 + \
(2*R**2 - R -2) - \
2*R**2 / (cos_lat**2) )
B = 2*(c.Omega0 + w)*K/(R+1)/(R+2)*cos_lat**R * \
(R**2 + 2*R + 2 - (R+1)**2*cos_lat**2)
C = 0.25*K**2*cos_lat**(2*R) * \
((R+1)*cos_lat**2 - (R+2))
# The thickness field
self.thickness[:] = c.gravity * h0 + a**2*A + \
a**2*B*xp.cos(R*g.lonCell) + \
a**2*C*xp.cos(2*R*g.lonCell)
self.thickness[:] /= c.gravity
# Vorticity and divergence fields
self.vorticity[:] = 2*w*xp.sin(g.latCell) - \
K*xp.sin(g.latCell)*cos_lat**R* \
(R**2 + 3*R + 2)*xp.cos(R*g.lonCell)
self.divergence[:] = 0.
self.SS0 = xp.sum((self.thickness + g.bottomTopographyCell) *
g.areaCell) / xp.sum(g.areaCell)
elif c.test_case == 7:
# Setup shallow water test case 7: Height and wind at 500 mb
#
# Reference: Williamson, D.L., et al., "A Standard Test Set for Numerical
# Approximations to the Shallow Water Equations in Spherical
# Geometry" J. of Comp. Phys., 102, pp. 211--224
ini_dat = np.loadtxt('tc7-init-on-%d.dat' % g.nCells)
self.thickness[:] = ini_dat[:, 2]
self.vorticity[:] = ini_dat[:, 3]
self.divergence[:] = ini_dat[:, 4]
self.SS0 = xp.sum((self.thickness + g.bottomTopographyCell) *
g.areaCell) / xp.sum(g.areaCell)
elif c.test_case == 8:
# Setup shallow water test case 8: barotropic instability test case
#
# Reference:
self.divergence[:] = 0.
# Compute the vorticity field
if c.use_gpu:
self.vorticity[:] = xp.asarray(
cmp.compute_swtc8_vort(g.latCell.get()))
else:
self.vorticity[:] = cmp.compute_swtc8_vort(g.latCell)
# Compute the thickness field, and shift it towards 10km average depth
if c.use_gpu:
gh = xp.asarray(cmp.compute_swtc8_gh(g.latCell.get()))
else:
gh = cmp.compute_swtc8_gh(g.latCell)
gh += 10000. * c.gravity - xp.sum(gh * g.areaCell) / xp.sum(
g.areaCell)
self.thickness[:] = gh / c.gravity
# Add a small perturbation to the thickness field
pert = 120 * xp.cos(g.latCell)
if xp.max(g.lonCell) > 1.1 * np.pi:
print("Found the range of lonCell to be [0, 2pi]")
print("Shifting it to [-pi, pi]")
g.lonCell[:] = xp.where(g.lonCell > np.pi,
g.lonCell - 2 * np.pi, g.lonCell)
g.lonVertex[:] = xp.where(g.lonVertex > np.pi,
g.lonVertex - 2 * np.pi, g.lonVertex)
g.lonEdge[:] = np.where(g.lonEdge > np.pi,
g.lonEdge - 2 * np.pi, g.lonEdge)
pert *= xp.exp(-(g.lonCell * 3)**2)
pert *= xp.exp(-((g.latCell - np.pi / 4) * 15)**2)
# self.thickness[:] += pert
self.SS0 = xp.sum((self.thickness + g.bottomTopographyCell) *
g.areaCell) / xp.sum(g.areaCell)
elif c.test_case == 12:
# SWSTC #2, with a stationary analytic solution, modified for the northern hemisphere
a = c.sphere_radius
u0 = 2 * np.pi * a / (12 * 86400)
gh0 = 2.94e4
gh = xp.sin(g.latCell[:])**2
gh = -(a * c.Omega0 * u0 + 0.5 * u0 * u0) * gh + gh0
self.thickness[:] = gh / c.gravity
h0 = gh0 / c.gravity
self.vorticity[:] = 2 * u0 / a * xp.sin(g.latCell[:])
self.divergence[:] = 0.
self.psi_cell[:] = -a * h0 * u0 * xp.sin(g.latCell[:])
self.psi_cell[:] += a * u0 / c.gravity * (
a * c.Omega0 * u0 + 0.5 * u0**2) * (xp.sin(
g.latCell[:]))**3 / 3.
self.phi_cell[:] = 0.
self.SS0 = xp.sum((self.thickness + g.bottomTopographyCell) *
g.areaCell) / xp.sum(g.areaCell)
elif c.test_case == 15:
# Zonal flow over a mountain topography, on the northern hemisphere,
# modeled after SWSTC #5
a = c.sphere_radius
u0 = 20.
h0 = 5960.
gh = c.gravity * h0 - xp.sin(
g.latCell[:])**2 * (a * c.Omega0 * u0 + 0.5 * u0 * u0)
h = gh / c.gravity
# Define the mountain topography
h_s0 = 2000.
R = np.pi / 9
lat_c = np.pi / 6.
lon_c = .5 * np.pi
r = xp.sqrt((g.latCell[:] - lat_c)**2 + (g.lonCell[:] - lon_c)**2)
r = xp.where(r < R, r, R)
g.bottomTopographyCell[:] = h_s0 * (1 - r / R)
self.thickness[:] = h[:] - g.bottomTopographyCell[:]
self.vorticity[:] = 2 * u0 / a * xp.sin(g.latCell[:])
self.divergence[:] = 0.
self.SS0 = xp.sum((self.thickness + g.bottomTopographyCell) *
g.areaCell) / xp.sum(g.areaCell)
self.curlWind_cell[:] = 0.
self.divWind_cell[:] = 0.
elif c.test_case == 21:
# A wind-driven gyre at mid-latitude in the northern hemisphere
tau0 = 1.e-4
latmin = xp.min(g.latCell[:])
latmax = xp.max(g.latCell[:])
lonmin = xp.min(g.lonCell[:])
lonmax = xp.max(g.lonCell[:])
latmid = 0.5 * (latmin + latmax)
latwidth = latmax - latmin
lonmid = 0.5 * (lonmin + lonmax)
lonwidth = lonmax - lonmin
r = c.sphere_radius
self.vorticity[:] = 0.
self.divergence[:] = 0.
self.thickness[:] = 4000.
self.psi_cell[:] = 0.0
self.psi_vertex[:] = 0.0
# Initialize wind
self.curlWind_cell[:] = -tau0 * np.pi/(latwidth*r) * \
xp.sin(np.pi*(g.latCell[:]-latmin) / latwidth)
self.divWind_cell[:] = 0.
self.SS0 = xp.sum((self.thickness + g.bottomTopographyCell) *
g.areaCell) / xp.sum(g.areaCell)
elif c.test_case == 22:
# One gyre with no forcing, for a bounded domain over NA
d = xp.sqrt(32 * (g.latCell[:] - latmid)**2 / latwidth**2 + 4 *
(g.lonCell[:] - (-1.1))**2 / .3**2)
f0 = xp.mean(g.fCell)
self.thickness[:] = 4000.
self.psi_cell[:] = 2 * xp.exp(-d**2) * 0.5 * (1 -
xp.tanh(20 *
(d - 1.5)))
# self.psi_cell[:] -= np.sum(self.psi_cell * g.areaCell) / np.sum(g.areaCell)
self.psi_cell *= c.gravity / f0 * self.thickness
self.phi_cell[:] = 0.
self.vorticity = vc.discrete_laplace_v(self.psi_cell)
self.vorticity /= self.thickness
self.divergence[:] = 0.
# Initialize wind
self.curlWind_cell[:] = 0.
self.divWind_cell[:] = 0.
# Eliminate bottom drag
#c.bottomDrag = 0.
# Eliminate lateral diffusion
#c.delVisc = 0.
#c.del2Visc = 0.
self.SS0 = xp.sum((self.thickness + g.bottomTopographyCell) *
g.areaCell) / xp.sum(g.areaCell)
else:
raise ValueError("Invaid choice for the test case.")
# Set time to zero
self.time = 0.0