def gll_points(n):
"""GLL points and weights
:param n: Number of points
:returns: (x, w)
:rtype:
"""
assert n>=2
if n==2:
x = np.array([-1.0, 1.0])
w = np.array([ 1.0, 1.0])
return x, w
# See Nodal Discontinuous Galerkin Methods Appendix A for x and
# the Mathworld page on Lobatto Quadrature for w
x = j_roots(n-2, 1, 1)[0]
L = eval_legendre(n-1, x)
w1 = 2.0/(n*(n-1))
w = 2.0/(n*(n-1)*L*L)
x = np.hstack([-1.0, x, 1.0])
w = np.hstack([w1, w, w1])
return x, w
def gll_points(n):
"""GLL points and weights
:param n: Number of points
:returns: (x, w)
:rtype:
"""
assert n >= 2
if n == 2:
x = np.array([-1.0, 1.0])
w = np.array([1.0, 1.0])
return x, w
# See Nodal Discontinuous Galerkin Methods Appendix A for x and
# the Mathworld page on Lobatto Quadrature for w
x = j_roots(n - 2, 1, 1)[0]
L = eval_legendre(n - 1, x)
w1 = 2.0 / (n * (n - 1))
w = 2.0 / (n * (n - 1) * L * L)
x = np.hstack([-1.0, x, 1.0])
w = np.hstack([w1, w, w1])
return x, w
def _nodes_LGL(self, n):
"""
Legendre-Gauss-Lobatto(LGL) points
"""
roots, weight = special.j_roots(n - 2, 1, 1)
nodes = np.hstack((-1, roots, 1))
return nodes
def poisson_beta_pmf(k, k_on, k_off, k_syn, n_roots=50):
assert (k_on > 0)
assert (k_off > 0)
roots, weights = special.j_roots(n_roots, alpha=k_off - 1, beta=k_on - 1)
mus = k_syn * (roots + 1) / 2
assert (max(mus) < 1e6)
gs = np.sum(weights * stats.poisson.pmf(k.reshape(-1, 1), mus), axis=1)
probabilities = 1 / special.beta(k_on, k_off) * 2**(1 - k_on - k_off) * gs
return probabilities
def dBP(at, alpha, bet, lam):
at.shape = (len(at), 1)
np.repeat(at, 50, axis=1)
def fun(at, m):
if (max(m) < 1e6):
return (poisson.pmf(at, m))
else:
return (norm.pdf(at, loc=m, scale=sqrt(m)))
x, w = j_roots(50, alpha=bet - 1, beta=alpha - 1)
gs = np.sum(w * fun(at, m=lam * (1 + x) / 2), axis=1) #m*s_i
prob = 1 / beta_fun(alpha, bet) * 2**(-alpha - bet + 1) * gs
return (prob)
def calc_solver_matrix(nr, nz, ne):
'''
nr: number of radial points
nz: number of axial points
ne: number of axial segments
'''
solver_data = {'nc': nr, 'mc': nz}
def dspoly(ia, nd, nc, x):
q = np.ones((nd, nc))
y = np.zeros(2 * nc)
c = np.zeros((nd, nc))
d = np.zeros((nd, nc))
f = np.zeros(nd)
loopq = [2 * (i + 1) - 3 for i in range(1, nc)]
loopc = [2 * (i + 1) - 4 for i in range(1, nc)]
cfac = np.array([2 * i for i in range(1, nc)])
loopd = loopq[:-1] #[2*(i+1)-5 for i in range(2,nc)]
dfac = np.array([(2 * i - 2) * (2 * i - 4 + ia)
for i in range(3, nc + 1)])
for j in range(nd):
y = np.array([x[j]**n for n in range(1, 2 * nc + 1)])
q[j, 1:] = np.array([y[i] for i in loopq])
c[j, 1:] = cfac * np.array([y[i] for i in loopc]) #2. * y[0]
d[j][1] = 2. * ia
d[j, 2:] = dfac * np.array([y[i] for i in loopd])
f[j] = 1. / float(2. * (j + 1) - 2. + ia)
return q, f, c, d
def dupoly(nd, nc, x):
q = np.ones((nd, nc))
y = np.zeros(2 * nc)
c = np.zeros((nd, nc))
d = np.zeros((nd, nc))
for j in range(nd):
y[0] = 1.
for i in range(1, nc):
y[i] = y[i - 1] * (2. * x[j] - 1.)
q[j, :] = np.array([y[i] for i in range(nc)])
c[j][1] = 2.
for i in range(2, nc):
c[j][i] = 2. * i * y[i - 1]
d[j][i] = 4. * i * (i - 1.) * y[i - 2]
return q, c, d
def build_colloc(n_pts):
'''
levih method
'''
n_pts_m = n_pts - 2 #number of interior points
roots = np.zeros([n_pts])
roots[-1] = 1 #right boundary
roots[1:-1] = special.roots_sh_legendre(n_pts_m)[0]
Qmat = np.zeros([n_pts, n_pts])
Cmat = np.zeros([n_pts, n_pts])
Dmat = np.zeros([n_pts, n_pts])
for i in range(n_pts):
for j in range(n_pts):
Qmat[i, j] = roots[i]**j
for j in range(1, n_pts):
Cmat[i, j] = (j) * roots[i]**[j - 1]
for j in range(2, n_pts):
Dmat[i, j] = (j) * (j - 1) * roots[i]**[j - 2]
Qinv = np.linalg.inv(Qmat)
Amat = np.dot(Cmat, Qinv)
Bmat = np.dot(Dmat, Qinv)
return (roots, Amat, Bmat)
def advect_operator(ne, nz):
nz_n = ne * nz - (ne - 1) #total number of axial points
f = ne #Element width adjustment to derivatives
roots, A, _ = build_colloc(nz) # construct first derivative operator
### Calculation location of overall gridpoints
Xvals = np.zeros(ne * nz)
for k in range(ne):
Xvals[k * nz:(k + 1) * nz] = (roots / f + k / f) #/f
to_delete = [nz * (k + 1) for k in range(ne - 1)]
Xvals = np.delete(Xvals, to_delete)
Q = A[-1, -1] - A[0, 0]
#adjusted element advection operator
Z = np.zeros([nz, nz])
Z[:, :] = f * A
#construct overall advection operator
Adv_Op = np.zeros([nz_n, nz_n])
#fill in blocks for element interiors
for k in range(ne):
ii = k * (nz - 1)
iii = (k + 1) * nz - k
Adv_Op[ii:iii, ii:iii] = Z[:, :]
#fill in continuation points where elements meet
for k in range(ne - 1):
idx = k * nz - k
ii = (k + 1) * nz - (k + 1)
iii = (k + 2) * nz - (k + 1)
Adv_Op[ii, :] = 0 # zero out continuation
CC1 = Z[-1, :-1] - Z[-1, -1] / Q * A[-1, :-1]
CC2 = Z[-1, -1] / Q * A[0, 1:]
Adv_Op[ii, idx:ii] = CC1
Adv_Op[ii, ii + 1:iii] = CC2
Adv_Op[0, :] = 0 # Constant inlet boundary
return Adv_Op, Xvals
ngeor = 3 # spherical
alfar = 1.
betar = 0.5
# radial matrices
rr = betar * (special.j_roots(nr - 1, alfar, betar)[0] + alfar)
rr = np.append(rr, [1]) #[0],rr)
r = np.sqrt(rr)
q, f, c, d = dspoly(ngeor, nr, nr, r)
qi2 = np.linalg.inv(q)
wr = f.dot(qi2) #.dot(f)
br = d.dot(qi2)
# axial
if ne == 1:
z = special.roots_sh_legendre(nz - 2)[0]
z = np.append(np.append([0], z), [1])
q, c, d = dupoly(nz, nz, z)
qi = np.linalg.inv(q)
az = c.dot(qi)
elif ne > 1:
az, _ = advect_operator(ne, nz)
solver_data['mc'] = ne * nz - (ne - 1)
#save data to export object
solver_data['wr'] = wr
solver_data['az'] = az
solver_data['br'] = br
return solver_data
def getZIntegrator(state, var="vortZ", nNs=40):
"""
compute integrator
Input:
- var: variable to integrate vortZ, uS, or uPhi
"""
nL = state.specRes.L
nM = state.specRes.M
nN = state.specRes.N
E = state.parameters.E
nNmax, nLmax, nMmax, nNs = nN, nL, nM, nNs
N_z = nNmax + nLmax // 2 # function of Nmax and Lmax
# Jacobi polynomial roots and weights
# Legendre alpha=0, beta=0
alpha = 0
beta = 0
roots, w = j_roots(N_z, alpha, beta)
d = 10. * np.sqrt(E)
ekmanR = 1. - d
grid_s = np.linspace(0, ekmanR, nNs + 1)
grid_s = 0.5 * (grid_s[0:-1] + grid_s[1:])
grid_z = np.zeros((nNs, N_z))
grid_cost = np.zeros_like(grid_z)
grid_sint = np.zeros_like(grid_z)
grid_r = np.zeros_like(grid_z)
# Generating (s,z) grids
for id_s, s in enumerate(grid_s):
# Quadrature points
#z_local = roots * sqrt(1.0 - s * s) * (1.0 - d) # from -1 to 1
z_local = roots * np.sqrt(
(1.0 - d) * (1.0 - d) - s * s) # from -1 to 1
grid_cost[id_s, :] = sz2ct(s, z_local) # from -1 to 1
grid_sint[id_s, :] = sz2st(s, z_local)
grid_r[id_s, :] = sz2r(s, z_local)
grid_z[id_s, :] = z_local
#vortz_tor = np.zeros((nNs, nNmax, nLmax * (nLmax + 1) // 2), dtype=complex)
vortz_tor = np.zeros(
(nNs, nNmax, nMmax * (nMmax + 1) // 2 + nMmax * (nLmax - nMmax)),
dtype=complex)
vortz_pol = np.zeros_like(vortz_tor)
#ur_pol = np.zeros_like(vortz_tor)
uphi_tor = np.zeros_like(vortz_tor)
uphi_pol = np.zeros_like(vortz_tor)
#uth_tor = np.zeros_like(vortz_tor)
#uth_pol = np.zeros_like(vortz_tor)
us_tor = np.zeros_like(vortz_tor)
us_pol = np.zeros_like(vortz_tor)
tr0 = time.time()
for id_s in range(0, nNs):
#print('id_s=', id_s)
print('id_s: ', id_s, end=" ") # ' timestep: ', timestep)
t0 = time.time()
for n in range(0, nNmax): # debug
# print('n=', n)
ind = 0
for l in range(0, nLmax):
#for m in range(0, min(l, Mmax) + 1):
for m in range(0, min(l, nMmax - 1) + 1):
# pdb.set_trace()
# Leo: compute on the fly instead of using memory
r_local = grid_r[id_s, :]
cos_local = grid_cost[id_s, :]
sin_local = grid_sint[id_s, :]
# polynomials for reconstruction
plm_w1 = W(n, l, r_local)
plm_lapw1 = laplacianW(n, l, r_local)
plm_diffw1 = diffW(n, l, r_local)
plm_leg = legSchmidtM(l, m, cos_local)
plm_diffLeg = diffLegM(l, m, cos_local)
if var == "vortZ":
# computing contributions from toroidal and poloidal vorticity onto Z
vort_z = 1j * m * plm_lapw1 * plm_leg # times P
vortz_tor[id_s, n, ind] = sum(
w * vort_z) * np.sqrt(1 - grid_s[id_s]**2) * (1.0 -
d)
vort_z = l * (l + 1.0) * plm_w1 * plm_leg * cos_local / r_local \
- plm_diffw1 * plm_diffLeg / r_local # times P
vortz_pol[id_s, n, ind] = sum(
w * vort_z) * np.sqrt(1 - grid_s[id_s]**2) * (1.0 -
d)
#sum(w*vort_z)*sqrt(1-grid_s[id_s]**2)*(1.0-d)
elif var == "uPhi":
#uth_tor[id_s, n, ind, :] = 1j * m * plm_w1 * plm_leg / sin_local
u_phi = -plm_w1 * plm_diffLeg / sin_local # -\partial_\theta T
uphi_tor[id_s, n, ind] = sum(
w * u_phi) * np.sqrt(1 - grid_s[id_s]**2) * (1.0 -
d)
#ur_pol[id_s, n, ind, :] = l * (l + 1.0) * plm_w1 * plm_leg / r_local # 1/r * L2(P)
#uth_pol[id_s, n, ind, :] = plm_diffw1 * plm_diffLeg / (r_local * sin_local) # ???
# 1/(r sin(theta)) \partial_r * r * \partial_\phi P
u_phi = 1j * m * plm_diffw1 * plm_leg / (r_local *
sin_local)
uphi_pol[id_s, n, ind] = sum(
w * u_phi) * np.sqrt(1 - grid_s[id_s]**2) * (1.0 -
d)
elif var == "uS":
# u_s = u_r * sin(theta) + u_theta * cos(theta)
u_s = l * (
l + 1.0
) * plm_w1 * plm_leg * sin_local / r_local + plm_diffw1 * plm_diffLeg * cos_local / (
r_local * sin_local)
us_pol[id_s, n, ind] = sum(
w * u_s) * np.sqrt(1 - grid_s[id_s]**2) * (1.0 - d)
u_s = 1j * m * plm_w1 * plm_leg * cos_local / sin_local
us_tor[id_s, n, ind] = sum(
w * u_s) * np.sqrt(1 - grid_s[id_s]**2) * (1.0 - d)
# u_z = u_r * cos(theta) - u_theta * sin(theta)
#uz_pol[id_s, n, ind, :] = l * (
# l + 1.0) * plm_w1 * plm_leg * cos_local / r_local - plm_diffw1 * plm_diffLeg / r_local
#uz_tor[id_s, n, ind, :] = 1j * m * plm_w1 * plm_leg
# h = u_z * vort_z
ind = ind + 1
#Make sure dimensions idx[l,m] agree
#assert(ind == vortz_tor.shape[2]), "size of integrator and ind=idx[l,m] don't agree"
#assert(ind == vortz_pol.shape[2]), "size of integrator and ind=idx[l,m] don't agree"
#assert(ind == us_tor.shape[2]), "size of integrator and ind=idx[l,m] don't agree"
#assert(ind == us_pol.shape[2]), "size of integrator and ind=idx[l,m] don't agree"
#assert(ind == uphi_tor.shape[2]), "size of integrator and ind=idx[l,m] don't agree"
#assert(ind == uphi_pol.shape[2]), "size of integrator and ind=idx[l,m] don't agree"
t1 = time.time()
print("time: ", t1 - t0)
tr1 = time.time()
print("Total time:", tr1 - tr0)
#resolution
res = (nNmax, nLmax, nMmax, nNs)
if var == "vortZ":
return Integrator(grid_s, vortz_tor, vortz_pol, res, var)
elif var == "uPhi":
return Integrator(grid_s, uphi_tor, uphi_pol, res, var)
elif var == "uS":
return Integrator(grid_s, us_tor, us_pol, res, var)
else:
print("output variable must be defined omegaZ, uPhi, or uS")
def computeUniformVorticity(state, rmax):
"""
function to compute uniform vorticity from a state file from 0 to rmax
Input:
- state: input state to compute uniform vorticity
- rmax : maximum radius to compute vorticity
"""
#rot.projectWf(dataT, E, nN, nL, nM, "rotation")
#Projection of spectral coefficients onto the T10, T11 modes
#Computing Wf with n=0,1,2,... nN-1
#dataT: Toroidal data
#E: Ekman number
#nN: number of radial modes
#nL: number of spherical harmonic order
#nM: number of azimuthal wave number
#scale: 'rotation' or 'viscous' timescale
dataT = state.fields.velocity_tor
nN = state.specRes.N
nL = state.specRes.L
nM = state.specRes.M
E = state.parameters.E
#2*nX + 1 = 2*NN+l + 3
l = 1
nX = (nN + l // 2 + 1)
#Jacobi polynomial roots and weights
#Legendre alpha=0, beta=0
alpha = 0
beta = 0
roots, w = j_roots(nX, alpha, beta)
#ekmanR=1.0
#ekmanR=1.-10.*np.sqrt(E)
ekmanR = rmax
#x = [-a,b]
#z = [0,1]
grid_r = (roots + 1.) / (2.) * ekmanR
poly = np.empty((nX, nN), dtype='float64', order='F')
quicc_pybind.wnl(poly, l, grid_r)
#index table
idx = idxlm(nL, nM)
Y11R = 0
Y11I = 0
Y10 = 0
#compute integral
for n in range(0, nN):
#print('nN', nN, 'n', n, 'dataT.shape:', dataT.shape)
#breakpoint()
#print("idx:",idx[1,1])
fi_wi = -2 * grid_r**3 * poly[:, n] * w
#print(grid_r**3*poly[:, n]*w)
#print(fi_wi.sum()*ekmanR/2)
Y11R = Y11R + dataT[idx[1, 1], n].real * fi_wi.sum() * (
ekmanR / 2.) * 5. / (ekmanR**5)
#print(dataT[idx[1,1], n,0])
fi_wi = 2 * grid_r**3 * poly[:, n] * w
Y11I = Y11I + dataT[idx[1, 1], n].imag * fi_wi.sum() * (
ekmanR / 2.) * 5. / (ekmanR**5)
#print(dataT[idx[1,1], n, 1])
#print(grid_r**3*poly[:, n]*w)
fi_wi = grid_r**3 * poly[:, n] * w
#print(fi_wi.sum()*ekmanR/2)
#print(dataT[idx[1,0], n, 0])
Y10 = Y10 + dataT[idx[1, 0], n].real * fi_wi.sum() * (
ekmanR / 2.) * 5. / (ekmanR**5)
#print("projectWf(n=",n,")", Y11R, Y11I, Y10)
omegaF = np.zeros(3)
omegaF[0] = Y11R
omegaF[1] = Y11I
omegaF[2] = Y10
return omegaF