def test_chebder_axis(self):
# check that axis keyword works
c2d = np.random.random((3, 4))
tgt = np.vstack([cheb.chebder(c) for c in c2d.T]).T
res = cheb.chebder(c2d, axis=0)
assert_almost_equal(res, tgt)
tgt = np.vstack([cheb.chebder(c) for c in c2d])
res = cheb.chebder(c2d, axis=1)
assert_almost_equal(res, tgt)
def test_chebder_axis(self):
# check that axis keyword works
c2d = np.random.random((3, 4))
tgt = np.vstack([cheb.chebder(c) for c in c2d.T]).T
res = cheb.chebder(c2d, axis=0)
assert_almost_equal(res, tgt)
tgt = np.vstack([cheb.chebder(c) for c in c2d])
res = cheb.chebder(c2d, axis=1)
assert_almost_equal(res, tgt)
def solid_angle_average_tor(nr, a, b, x1, operation='simple', l=None):
# precondition: assume that the weight of the spectra is already in DCT format
# make copies
x_1 = np.array(x1)
# pad the data
x1 = np.zeros(int(3 / 2 * nr))
x1[:len(x_1)] = x_1
# create the spectral function that describes T_1=x
r1 = np.zeros(int(3 / 2 * nr))
r1[1] = .5
# r matches radius
r = idct(r1, type=2) * a + b
if operation == 'simple':
# bring the functions to physical space
y1 = idct(x1, type=2)
# multiply in physical
y3 = y1**2
# done this way to approximate QuICC at most
elif operation == 'curl':
# prepare vectors for the derivative
x1a = np.array(x1)
x1a[1:] *= 2
dx1 = np.append(chebder(x1), 0.)
d2x1 = np.append(chebder(dx1), 0.)
# prepare derivatives for the DCT
dx1[1:] *= .5
d2x1[1:] *= .5
# transform everything
y1 = idct(x1, type=2)
dy1 = idct(dx1, type=2)
d2y1 = idct(d2x1, type=2)
# compute the 2 pieces of the bilinear operator
Diss1 = d2y1 + 2 * dy1 / r - l * (l + 1) * y1 / r**2
y3 = Diss1**2
return y3 * r**2
def grad(self, point):
'''
Evaluates the gradient of the polynomial at the given point.
Parameters
----------
point : array-like
the point at which to evaluate the polynomial
Returns
-------
out : ndarray
Gradient of the polynomial at the given point.
'''
super(MultiCheb, self).__call__(point)
out = np.empty(self.dim, dtype="complex_")
if self.jac is None:
jac = list()
for i in range(self.dim):
jac.append(cheb.chebder(self.coeff, axis=i))
self.jac = jac
spot = 0
for i in self.jac:
out[spot] = chebvalnd(point, i)
spot += 1
return out
def evaluate_basis_derivative_all(self, x=None, k=0):
if x is None:
x = self.mesh(False, False)
V = self.vandermonde(x)
if k > 0:
D = np.zeros((self.N, self.N))
D[:-k, :] = n_cheb.chebder(np.eye(self.N), k)
V = np.dot(V, D)
return self._composite_basis(V)
def proj_lapl(nr, a, b, x=None):
# evaluate the second derivative of radial basis in radial direction function of degree <nr
# this projection operator corresponds to 1/r**2 d/dr r**2d/dr projecting from spectral to physical
# used for the theta and phi contribution of the poloidal field (S-part of a qst decomposition)
if x is None:
xx, ww = cheb.chebgauss(nr)
else:
xx = np.array(x)
# use chebyshev polynomials normalized for FCT
coeffs = np.eye(nr) * 2
#coeffs = np.eye(nr)
coeffs[0, 0] = 1.
c1 = cheb.chebder(coeffs,1)
c2 = cheb.chebder(coeffs,2)
return np.mat((2*cheb.chebval(xx, c1) /a /(a * xx + b) + cheb.chebval(xx, c2) /a**2).transpose())
def cheby_newton_root(z, f, z0=None, degree=512):
import numpy.polynomial.chebyshev as npcheb
import scipy.optimize as scpop
Lz = np.max(z) - np.min(z)
if z0 is None:
z0 = Lz / 2
def to_x(z, Lz):
# convert back to [-1,1]
return (2 / Lz) * z - 1
def to_z(x, Lz):
# convert back from [-1,1]
return (x + 1) * Lz / 2
logger.info("searching for roots starting from z={}".format(z0))
x = to_x(z, Lz)
x0 = to_x(z0, Lz)
cheb_coeffs = npcheb.chebfit(x, f, degree)
cheb_interp = npcheb.Chebyshev(cheb_coeffs)
cheb_der = npcheb.chebder(cheb_coeffs)
def newton_func(x_newton):
return npcheb.chebval(x_newton, cheb_coeffs)
def newton_derivative_func(x_newton):
return npcheb.chebval(x_newton, cheb_der)
try:
x_root = scpop.newton(newton_func,
x0,
fprime=newton_derivative_func,
tol=1e-10)
z_root = to_z(x_root, Lz)
except:
logger.info("error in root find")
x_root = np.nan
z_root = np.nan
logger.info("newton: found root z={} (x0:{} -> {})".format(
z_root, x0, x_root))
for x0 in x:
print(x0, newton_func(x0))
a = Lz / 4
b = Lz * 3 / 4
logger.info("bisecting between z=[{},{}] (x=[{},{}])".format(
a, b, to_x(a, Lz), to_x(b, Lz)))
logger.info("f(a) = {} and f(b) = {}".format(newton_func(to_x(a, Lz)),
newton_func(to_x(b, Lz))))
x_root_2 = scpop.bisect(newton_func, to_x(a, Lz), to_x(b, Lz))
z_root_2 = to_z(x_root_2, Lz)
logger.info("bisect: found root z={} (x={})".format(z_root_2, x_root_2))
return z_root_2
def evaluate_basis_derivative(self, x=None, i=0, k=0, output_array=None):
if x is None:
x = self.mesh(False, False)
x = np.atleast_1d(x)
v = self.evaluate_basis(x, i, output_array)
if k > 0:
D = np.zeros((self.N, self.N))
D[:-k, :] = n_cheb.chebder(np.eye(self.N), k)
v = np.dot(v, D)
return v
def chebReconstructHessian(freqs, locations, spectrum, numPts):
"""
:param freqs:
:param locations:
:param spectrum:
:param numPts:
:return:
"""
deriv = cheb.chebder(spectrum, m=2, axis=0)
return genericChebVal(locations, deriv)
def chebReconstructDerivative(freqs, locations, spectrum, numPts):
"""
:param freqs:
:param locations:
:param spectrum:
:param numPts:
:return: derivative of iFFT of the spectrum as an array with shape (Npts, Ncomponent) or (Ncomponent,) if 1-d
"""
deriv = cheb.chebder(spectrum, axis=0)
return genericChebVal(locations, deriv)
def test_chebder(self):
# check exceptions
assert_raises(ValueError, ch.chebder, [0], -1)
# check that zeroth deriviative does nothing
for i in range(5):
tgt = [1] + [0] * i
res = ch.chebder(tgt, m=0)
assert_equal(trim(res), trim(tgt))
# check that derivation is the inverse of integration
for i in range(5):
for j in range(2, 5):
tgt = [1] + [0] * i
res = ch.chebder(ch.chebint(tgt, m=j), m=j)
assert_almost_equal(trim(res), trim(tgt))
# check derivation with scaling
for i in range(5):
for j in range(2, 5):
tgt = [1] + [0] * i
res = ch.chebder(ch.chebint(tgt, m=j, scl=2), m=j, scl=.5)
assert_almost_equal(trim(res), trim(tgt))
def test_chebder(self) :
# check exceptions
assert_raises(ValueError, ch.chebder, [0], -1)
# check that zeroth deriviative does nothing
for i in range(5) :
tgt = [1] + [0]*i
res = ch.chebder(tgt, m=0)
assert_equal(trim(res), trim(tgt))
# check that derivation is the inverse of integration
for i in range(5) :
for j in range(2,5) :
tgt = [1] + [0]*i
res = ch.chebder(ch.chebint(tgt, m=j), m=j)
assert_almost_equal(trim(res), trim(tgt))
# check derivation with scaling
for i in range(5) :
for j in range(2,5) :
tgt = [1] + [0]*i
res = ch.chebder(ch.chebint(tgt, m=j, scl=2), m=j, scl=.5)
assert_almost_equal(trim(res), trim(tgt))
def get_vandermonde_basis_derivative(self, V, k=0):
"""Return k'th derivative of basis as a Vandermonde matrix
Parameters
----------
V : array of ndim = 2
Chebyshev Vandermonde matrix
k : int
k'th derivative
"""
assert self.N == V.shape[1]
if k > 0:
D = np.zeros((self.N, self.N))
D[:-k, :] = n_cheb.chebder(np.eye(self.N), k)
V = np.dot(V, D)
return self.get_vandermonde_basis(V)
def proj_dradial_dr(nr, a, b, x = None):
# evaluate the first derivative of radial basis function of degree <nr
# this projection operator corresponds to 1/r d/dr r projecting from spectral to physical
# used for the theta and phi contribution of the poloidal field (S-part of a qst decomposition)
if x is None:
xx, ww = cheb.chebgauss(nr)
else:
xx = np.array(x)
# use chebyshev polynomials normalized for FCT
coeffs = np.eye(nr) * 2
#coeffs = np.eye(nr)
coeffs[0, 0] = 1.
c = cheb.chebder(coeffs)
#
temp = (cheb.chebval(xx,c)/a + cheb.chebval(xx,coeffs)/(a*xx+b)).transpose()
#return spsp.coo_matrix(temp)
return np.mat(temp)
def solid_angle_average_pol_s(nr, a, b, x1):
# precondition: assume that the weight of the spectra is already in DCT format
# take derivative
# first bring back to correct form
x_1a = np.array(x1)
x_1a[1:] *= 2
dx_1 = chebder(x_1a)
# make copies
x_1b = np.array(x1)
# pad the data
x1a = np.zeros(int(3 / 2 * nr))
x1b = np.zeros(int(3 / 2 * nr))
x1a[:len(dx_1)] = dx_1
x1b[:len(x_1b)] = x_1b
# create the spectral function that describes T_1=x
r1 = np.zeros(int(3 / 2 * nr))
r1[1] = .5
# prepare the chebs for DCT
x1a[1:] *= .5
# bring the functions to physical space
y1a = idct(x1a, type=2)
y1b = idct(x1b, type=2)
x = idct(r1, type=2)
r = x * a + b
# multiply in physical
y3 = (y1a / a + y1b / r)**2
return y3 * r**2
def evaluate_mode(self, l, m, *args, **kwargs):
"""
evaluate (l,m) mode and return FieldOut
input:
spectral coefficients: dataT, dataP,
physical grid: r, theta, phi0
kron: 'meridional' or 'equatorial'
outpu: FieldOut
e.g:
evaluate_mode(l, m, FieldOut, dataT[i, :], dataP[i, :], r, theta, kron='meridional', phi0=p)
"""
#print(l, m)
# raise exception if wrong geometry
if not (self.geometry == 'shell' or self.geometry == 'sphere'):
raise NotImplementedError('makeMeridionalSlice is not implemented for the geometry: '+self.geometry)
# prepare the input data
#Evaluated fields
Field_r = args[0][0]
Field_theta = args[0][1]
Field_phi = args[0][2]
#spectral coefficients
modeT = args[1]
modeP = args[2]
#grid points
r = args[3]
theta = args[4]
phi = args[5]
phi0 = kwargs['phi0']
# define factor
factor = 1. if m==0 else 2.
if kwargs['kron'] == 'isogrid' or kwargs['kron'] == 'points':
x = kwargs['x']
if self.geometry == 'shell':
# assume that the mode is weighted like Philippe sets it
modeP[1:] *= 2.
modeT[1:] *= 2.
# prepare the q_part
modeP_r = cheb.chebval(x, modeP)/r
q_part = modeP_r * l*(l+1)
# prepare the s_part
dP = np.zeros_like(modeP)
d_temp = cheb.chebder(modeP)
dP[:len(d_temp)] = d_temp
s_part = modeP_r + cheb.chebval(x, dP)/self.a
# prepare the t_part
t_part = cheb.chebval(x, modeT)
elif self.geometry == 'sphere':
#TODO: Leo
q_part = None
s_part = None
t_part = None
else: # hence either meridional or equatorial
if self.geometry == 'shell':
# prepare the q_part
# q = l*(l+1) * Poloidal / r
# idct: inverse discrete cosine transform
# type = 2: type of transform
# q = Field_radial
# prepare the q_part
# idct = \sum c_n * T_n(xj)
modeP_r = idct(modeP, type = 2)/r
q_part = modeP_r * l*(l+1)
# prepare the s_part
# s_part = orthogonal to Field_radial and Field_Toroidal
# s_part = (Poloidal)/r + d(Poloidal)/dr = 1/r d(Poloidal)/dr
dP = np.zeros_like(modeP)
d_temp = cheb.chebder(modeP)
dP[:len(d_temp)] = d_temp
s_part = modeP_r + idct(dP, type = 2)/self.a
# prepare the t_part
# t_part = Field_Toroidal
t_part = idct(modeT, type = 2)
elif self.geometry == 'sphere':
#TODO: Leo, Where do you get this???
#print('modeP:', modeP.shape, modeP) #32
#print('modeT:', modeT.shape, modeT) #32
#print('r:',r.shape, r) # 64
#print('specRes:',self.specRes)
modeP_r = np.zeros_like(r)
dP = np.zeros_like(r)
t_part = np.zeros_like(r)
for n in range(self.specRes.N):
modeP_r=modeP_r+modeP[n]*wor.W(n, l, r)/r
dP = dP+modeP[n]*wor.diffW(n, l, r)
t_part = t_part + modeT[n]*wor.W(n,l,r)
q_part = modeP_r*l*(l+1)#same stuff
s_part = modeP_r + dP
# depending on the kron type it changes how 2d data are formed
# Mapping from qst to r,theta, phi
# phi0: azimuthal angle of meridional plane. It's also the phase shift of the flow with respect to the mantle frame,
if kwargs['kron'] == 'meridional':
eimp = np.exp(1j * m * phi0)
idx_ = self.idx[l, m]
Field_r += np.real(tools.kron(q_part, self.Plm[idx_, :]) *
eimp) * factor
Field_theta += np.real(tools.kron(s_part, self.dPlm[idx_, :]) * eimp) * 2
Field_theta += np.real(tools.kron(t_part, self.Plm_sin[idx_, :]) * eimp * 1j * m) * 2#factor
Field_phi += np.real(tools.kron(s_part, self.Plm_sin[idx_, :]) * eimp * 1j * m) * 2#factor
Field_phi -= np.real(tools.kron(t_part, self.dPlm[idx_, :]) * eimp) * 2
elif kwargs['kron'] == 'points':
eimp = np.exp(1j * m * phi)
idx_ = self.idx[l, m]
Field_r += np.real(q_part * self.Plm[idx_, :] * eimp) * factor
Field_theta += np.real(s_part * self.dPlm[idx_, :] * eimp) * 2
Field_theta += np.real(t_part * self.Plm_sin[idx_, :] * eimp * 1j * m) * 2#factor
Field_phi += np.real(s_part * self.Plm_sin[idx_, :] * eimp * 1j * m) * 2#factor
Field_phi -= np.real(t_part * self.dPlm[idx_, :] * eimp) * 2
elif kwargs['kron'] == 'equatorial':
"""
eimp = np.exp(1j * m * (phi0 + phi))
idx_ = self.idx[l, m]
Field_r += np.real(tools.kron(q_part, eimp) *
self.Plm[idx_, :][0]) * factor
Field_theta += np.real(tools.kron(s_part, eimp)) * self.dPlm[idx_, :][0] * 2
Field_theta += np.real(tools.kron(t_part, eimp) * 1j * m) * self.Plm_sin[idx_, :][0] * 2#factor
Field_phi += np.real(tools.kron(s_part, eimp) * 1j * m) * self.Plm_sin[idx_, :][0] * 2#factor
Field_phi -= np.real(tools.kron(t_part, eimp)) * self.dPlm[idx_, :][0] * 2
"""
idx_ = self.idx[l, m]
#Field_r += np.real(tools.kron(self.Plm[idx_, :], eimp) * q_part) * factor
Field_r[:, m] += self.Plm[idx_, :] * q_part
#Field_theta += np.real(tools.kron(self.dPlm[idx_, :], eimp) * s_part) * 2
Field_theta[:, m] += self.dPlm[idx_, :] * s_part
#Field_theta += np.real(tools.kron(self.Plm_sin[idx_, :], eimp * 1j * m) * t_part) * 2#factor
Field_theta[:, m] += self.Plm_sin[idx_, :] * 1j * m * t_part
#Field_phi += np.real(tools.kron(self.Plm_sin[idx_, :], eimp * 1j * m) * s_part) * 2#factor
Field_phi[:, m] += self.Plm_sin[idx_, :]* 1j * m * s_part
#Field_phi -= np.real(tools.kron(self.dPlm[idx_, :], eimp) * t_part) * 2
Field_phi[:, m] -= self.dPlm[idx_, :] * t_part
elif kwargs['kron'] == 'isogrid':
#eimp = np.exp(1j * m * (phi0 + phi))
idx_ = self.idx[l, m]
#Field_r += np.real(tools.kron(self.Plm[idx_, :], eimp) * q_part) * factor
Field_r[:, m] += self.Plm[idx_, :] * q_part
#Field_theta += np.real(tools.kron(self.dPlm[idx_, :], eimp) * s_part) * 2
Field_theta[:, m] += self.dPlm[idx_, :] * s_part
#Field_theta += np.real(tools.kron(self.Plm_sin[idx_, :], eimp * 1j * m) * t_part) * 2#factor
Field_theta[:, m] += self.Plm_sin[idx_, :] * 1j * m * t_part
#Field_phi += np.real(tools.kron(self.Plm_sin[idx_, :], eimp * 1j * m) * s_part) * 2#factor
Field_phi[:, m] += self.Plm_sin[idx_, :]* 1j * m * s_part
#Field_phi -= np.real(tools.kron(self.dPlm[idx_, :], eimp) * t_part) * 2
Field_phi[:, m] -= self.dPlm[idx_, :] * t_part
else:
raise ValueError('Kron type not understood: '+kwargs['kron'])
pass
dtype=np.float64)
sub_nperiod[-1] = int(info['sub_nperiod_last'])
middle = sub_nperiod.cumsum() - sub_nperiod / 2
time0 = np.linspace(0, np.float64(info['length']), nperiod)
psr = pm.psr_timing(
pr.psr(info['psr_par']),
te.times(
te.time(
np.float64(info['stt_date']) *
np.ones(nperiod, dtype=np.float64),
np.float64(info['stt_sec']) + time0)), freq_s)
phase1 = psr.phase
phase_start = phase1.date[0]
chebc = nc.chebfit(time0, phase1.date - phase1.date[0] + phase1.second,
int(nperiod / 2) + 1)
chebd = nc.chebder(chebc)
middle_time = np.interp(middle, phase1.date - phase0 + phase1.second,
time0)[sub_s:sub_e]
psr1 = pm.psr_timing(
pr.psr(info['psr_par']),
te.times(
te.time(
np.float64(info['stt_date']) *
np.ones(nsub_new[k], dtype=np.float64),
np.float64(info['stt_sec']) + middle_time)), freq_s)
middle_phase = psr1.phase
data1 = d.period_scrunch(sub_s, sub_e)[chanstart:chanend, :, 0]
if not args.dm_corr:
freq_real = np.linspace(freq_start, freq_end,
nchan + 1)[:-1] * psr.vchange.mean()
if 'best_dm' in info.keys():
def computeZintegralOnTheFly(state, zInt):
"""
returns the z-integral of vorticity and the flow up to
a max spherical harmonic order Mmax and Worland Polynomial (Nmax)
vort_int=computeZintegral(state, vortz_tor, vortz_pol, ur_pol, uphi_tor, uphi_pol, (Nmax, Lmax, Mmax, N_s), w, grid_s)
"""
f = h5py.File(state, 'r')
LL = f['/truncation/spectral/dim2D'].value + 1
MM = f['/truncation/spectral/dim3D'].value + 1
NN = f['/truncation/spectral/dim1D'].value + 1
E = f['/physical/ekman'].value
dataP = f['velocity/velocity_pol'].value
dataT = f['velocity/velocity_tor'].value
eta = f['/physical/rratio'].value
f.close()
# compute the diffeomorfism parameters between Tchebyshev and radius space
a = .5
b = .5 * (1 + eta) / (1 - eta)
# compute boundary layer
#d = 10.*E**.5
#riBoundary = eta/(1-eta)+d
#roBoundary = 1/(1-eta)-d
# TODO: decide if the one needs to import the resolution from an argument
# compute the outside tangent cylinder part
#NsPoints = 20 # NsPoints is the number of points int the radius of the gap
NsPoints, Nmax, Lmax, Mmax = zInt.res
# there are points also on top of the tangent cylinder
# build the cylindrical radial grid
#s = np.linspace(0, roBoundary, 2 * NsPoints + 1)[1::2]
#s1 = s[s>=riBoundary]
#s2 = s[s<riBoundary]
# compute the outside tangent cylinder part
#ss1, no = np.meshgrid(s1, np.ones(2 * NsPoints))
#fs1 = (roBoundary ** 2 - s1 ** 2) ** .5
#x, w = leggauss(NsPoints * 2)
#no, zz1 = np.meshgrid(fs1, x)
#no, ww1 = np.meshgrid(fs1, w)
#zz1 *= fs1
# compute the inside of the tangent cylinder
#ss2, no = np.meshgrid(s2, np.ones(2 * NsPoints))
#fs2 = ((roBoundary ** 2 - s2 ** 2) ** .5 - (riBoundary ** 2 - s2 ** 2) ** .5) / 2
#x, w = leggauss(NsPoints)
#w *= .5
#no, zz2 = np.meshgrid(fs2, x)
#no, ww2 = np.meshgrid(fs2, w)
#zz2 *= fs2
#means = ((roBoundary ** 2 - s2 ** 2) ** .5 + (riBoundary ** 2 - s2 ** 2) ** .5) / 2
#zz2 += means
#zz2 = np.vstack((zz2, -zz2))
#ww2 = np.vstack((ww2, ww2))
# combine the 2 grids together
#ss = np.hstack((ss2, ss1))
#zz = np.hstack((zz2, zz1))
#ww = np.hstack((ww2, ww1))
# prepare the grid for tchebyshev polynomials
#ttheta = np.arctan2(ss, zz)
#rr = (ss**2+zz**2)**.5
#xx = (rr - b)/a
#xx_th = np.cos(ttheta)
#ccos_theta = xx_th
#ssin_theta = np.sin(ttheta)
# prepare the division weight
#fs = np.hstack([fs2*2, fs1]) * 2
zInt.generateGrid(state)
ss = zInt.ss
s = ss[0, :]
zz = zInt.zz
ttheta = zInt.ttheta
ww = zInt.ww
rr = zInt.rr
xx = zInt.xx
xx_th = zInt.ccos_theta
ssin_theta = zInt.ssin_theta
fs = zInt.fs
# prepare resolution for the integrator
Ns = ss.shape[1]
# edit: now Ns == NsPoints
# generate the dictionary that maps the l,m to idx
idx = idxlm(LL, MM)
#idxM = idxlm(Lmax,Mmax)
#Transform into complex data
dataT = dataT[:, :, 0] + dataT[:, :, 1] * 1j
dataP = dataP[:, :, 0] + dataP[:, :, 1] * 1j
#main idea compute for each m the integral
#Int(f(s,z)dz) = sum(weight*vort_z(s,z)) = f(s)
#vort_z function that returns the component of the z-vorticity
#for id_s in range(Ns):
#x = xx[:, id_s]
#w = ww[:, id_s]
#r = rr[:, id_s]
#theta = ttheta[:, id_s]
x = np.reshape(xx, (-1))
w = np.reshape(ww, (-1))
r = np.reshape(rr, (-1))
theta = np.reshape(ttheta, (-1))
sin_theta = np.sin(theta)
cos_theta = np.cos(theta)
#xtheta = xx_th[:, id_s]
xtheta = np.reshape(xx_th, (-1))
#Allocating memory for output
vort_int = np.zeros((len(x), Mmax), dtype=complex)
Ur_int = np.zeros_like(vort_int)
Uth_int = np.zeros_like(vort_int)
Uphi_int = np.zeros_like(vort_int)
Us_int = np.zeros_like(vort_int)
Uz_int = np.zeros_like(vort_int)
H_int = np.zeros_like(vort_int)
# compute the radial matrices
#Tn = shell.proj_radial(Nmax, a, b, x) # evaluate the Tn
#dTndr = shell.proj_dradial_dr(Nmax, a, b, x) # evaluate 1/r d/dr(r Tn)
#Tn_r = shell.proj_radial_r(Nmax, a, b, x) # evaluate 1/r Tn
# produce the mapping for the tri-curl part
#Tn_r2 = shell.proj_radial_r2(Nmax, a, b, x) # evaluate 1/r**2 Tn
#d2Tndr2 = shell.proj_lapl(Nmax, a, b, x) # evaluate 1/r**2 dr r**2 dr
for l in range(0, min(LL, Lmax)):
for m in range(0, min(l + 1, MM, Mmax)):
# the computation is carried out over a vertical grid
# assume that u_r, u_theta and u_phi are vectors
# compute the theta evaluations
plm = spherical.lplm(Lmax, l, m, xtheta)
dplm = spherical.dplm(Lmax, l, m, xtheta)
plm_sin = spherical.lplm_sin(Lmax, l, m, xtheta)
# store mode nicely
modeP = dataP[idx[l, m], :Nmax]
modeT = dataT[idx[l, m], :Nmax]
# use the DCT weighting
modeP[1:] *= 2.
modeT[1:] *= 2.
# compute the transformed radial part
#ur_part = l*(l+1)*dot(Tn_r, dataP[idx[l, m], :Nmax])
P_r = cheb.chebval(x, modeP) / r
ur_part = l * (l + 1) * P_r
#upol_part = dot(dTndr, dataP[idx[l, m], :Nmax])
dmodeP = cheb.chebder(modeP)
tempP = cheb.chebval(x, dmodeP) / a
upol_part = P_r + tempP
#utor_part = dot(Tn, dataT[idx[l, m], :Nmax])
utor_part = cheb.chebval(x, modeT)
T_r = cheb.chebval(x, modeT) / r
#omegar_part = l*(l+1)*dot(Tn_r, dataT[idx[l, m], :Nmax])
omegar_part = l * (l + 1) * T_r
#omegator_part = dot(dTndr, dataT[idx[l, m], :Nmax])
dmodeT = cheb.chebder(modeT)
omegator_part = T_r + cheb.chebval(x, dmodeT) / a
#omegapol_part = -( dot(d2Tndr2, dataP[idx[l, m], :Nmax]) - l*(l+1)*dot(Tn_r2, dataP[idx[l, m], :Nmax]) )
ddmodeP = cheb.chebder(dmodeP)
omegapol_part = -(cheb.chebval(x, ddmodeP) / a**2 + 2 * tempP / r -
l * (l + 1) / r * P_r)
# compute the r-coordinate components
u_r = ur_part * plm
u_theta = dplm * upol_part + 1j * m * plm_sin * utor_part
u_phi = 1j * m * plm_sin * upol_part - dplm * utor_part
omega_r = omegar_part * plm
omega_theta = dplm * omegator_part + 1j * m * plm_sin * omegapol_part
omega_phi = 1j * m * plm_sin * omegator_part - dplm * omegapol_part
# convert in cylindrical coordinate components
u_z = u_r * cos_theta - u_theta * sin_theta
u_s = u_r * sin_theta + u_theta * cos_theta
omega_z = omega_r * cos_theta - omega_theta * sin_theta
#omega_s = omega_r * sin_theta + omega_theta * cos_theta
# For real data:
#vort_z = zInt.vortz_tor[id_s, n, idxM[l, m], :] * dataP[idx[l, m], :] \
# + zInt.vortz_pol[id_s, n, idxM[l, m], :] * dataT[idx[l, m], :]
# u_r = zInt.ur_pol[id_s, n, idxM[l,m], :] * dataP[idx[l,m], n]
#u_th = zInt.uth_tor[id_s, n, idxM[l,m], :] * dataT[idx[l,m], n] \
# + zInt.uth_pol[id_s, n, idxM[l,m], :] * dataP[idx[l,m], n]
#u_phi = zInt.uphi_tor[id_s, n, idxM[l, m], :] * dataT[idx[l, m], n] \
# + zInt.uphi_pol[id_s, n, idxM[l, m], :] * dataP[idx[l, m], n]
#u_s = zInt.us_tor[id_s, n, idxM[l, m], :] * dataT[idx[l, m], n] \
# + zInt.us_pol[id_s, n, idxM[l, m], :] * dataP[idx[l, m], n]
#u_z = zInt.uz_tor[id_s, n, idxM[l,m], :] * dataT[idx[l,m], n] \
# + zInt.uz_pol[id_s, n, idxM[l,m], :] * dataP[idx[l,m], n]
# don't forget to include the factor sqrt(1-s^2) to compute the integral
# For real data:
# Integrate
vort_int[:, m] += w * omega_z #/ fs[id_s]
# Nicolo: modified the weight to fs
#Ur_int[:, m] += w * u_r #/ fs[id_s]
#Uth_int[:, m] += w * u_th #/ fs[id_s]
Uphi_int[:, m] += w * u_phi #/ fs[id_s]
Us_int[:, m] += w * u_s #/ fs[id_s]
Uz_int[:, m] += w * u_z #/ fs[id_s]
#H_int[:, m] += w * u_z * vort_z #/ fs[id_s]
vort_temp = np.reshape(vort_int, (xx.shape[0], xx.shape[1], Mmax))
vort_int = np.sum(vort_temp, axis=0).T
#Ur_temp = np.reshape(Ur_int, (xx.shape[0], xx.shape[1], Mmax) )
#Ur_int = np.sum(Ur_temp, axis=0)
#Uth_temp = np.reshape(Uth_int, (xx.shape[0], xx.shape[1], Mmax) )
#Uth_int = np.sum(Uth_temp, axis=0)
Uphi_temp = np.reshape(Uphi_int, (xx.shape[0], xx.shape[1], Mmax))
Uphi_int = np.sum(Uphi_temp, axis=0).T
Us_temp = np.reshape(Us_int, (xx.shape[0], xx.shape[1], Mmax))
Us_int = np.sum(Us_temp, axis=0).T
Uz_temp = np.reshape(Uz_int, (xx.shape[0], xx.shape[1], Mmax))
Uz_int = np.sum(Uz_temp, axis=0).T
#H_temp = np.reshape(H_int, (xx.shape[0], xx.shape[1], Mmax) )
#H_int = np.sum(H_temp, axis=0)
result = {
's': s,
'm': np.arange(Mmax),
'Omega_z': vort_int,
'U_phi': Uphi_int,
'U_s': Us_int,
'U_z': Uz_int
}
return result
def chebder(cs, m=1, scl=1) :
from numpy.polynomial.chebyshev import chebder
return chebder(cs, m, scl)
def ortho_pol_s(nr, a, b, x1, x2, xmin=-1, xmax=1, I1=None):
# precondition: assume that the weight of the spectra is already in DCT format
# take derivative
# first bring back to correct form
x_1a = np.array(x1)
x_1a[1:] *= 2
x_2a = np.array(x2)
x_2a[1:] *= 2
dx_1 = chebder(x_1a)
dx_2 = chebder(x_2a)
# make copies
x_1b = np.array(x1)
x_2b = np.array(x2)
# pad the data
x1a = np.zeros(int(3 / 2 * nr))
x1b = np.zeros(int(3 / 2 * nr))
x2a = np.zeros(int(3 / 2 * nr))
x2b = np.zeros(int(3 / 2 * nr))
x1a[:len(dx_1)] = dx_1
x1b[:len(x_1b)] = x_1b
x2a[:len(dx_2)] = dx_2
x2b[:len(x_2b)] = x_2b
# create the spectral function that describes T_1=x
r1 = np.zeros(int(3 / 2 * nr))
r1[1] = 1.
# prepare the chebs for DCT
x1a[1:] *= .5
#x1b[1:] *= .5
x2a[1:] *= .5
#x2b[1:] *= .5
r1[1:] *= .5
# bring the functions to physical space
y1a = idct(x1a, type=2)
y1b = idct(x1b, type=2)
y2a = idct(x2a, type=2)
y2b = idct(x2b, type=2)
x = idct(r1, type=2)
# multiply in physical
y3 = (y1a * (x + b / a) + y1b) * (y2a * (x + b / a) + y2b)
# bring back to spectral space
c = dct(y3, type=2) / (2 * len(y3))
c = c[:nr]
# represent the cheb indices in the mathematically correct way
"""
c[1:] *= 2.
# integrate, prepare index
idx = np.arange(len(c))
# and compute integration weight for -1 to +1
w = (1+(-1)**idx)/(1-idx**2)
w[1]=0
#return np.sum(w*c)/4
x,w = chebgauss(len(y3))
return np.sum(w*y3) * a
"""
Ic = I1 * c
temp = chebval([xmax, xmin], Ic)
return temp[0] - temp[1]
def ortho_tor(nr,
a,
b,
x1,
x2,
xmin=-1,
xmax=1,
I1=None,
operation='simple',
l=None):
# precondition: assume that the weight of the spectra is already in DCT format
# make copies
x_1 = np.array(x1)
x_2 = np.array(x2)
# pad the data
x1 = np.zeros(int(3 / 2 * nr))
x2 = np.zeros(int(3 / 2 * nr))
x1[:len(x_1)] = x_1
x2[:len(x_2)] = x_2
# create the spectral function that describes T_1=x
r1 = np.zeros(int(3 / 2 * nr))
r1[1] = 1.
r1[1:] *= .5
if operation == 'simple':
# bring the functions to physical space
y1 = idct(x1, type=2)
y2 = idct(x2, type=2)
r = idct(r1, type=2)
# multiply in physical
y3 = y1 * y2 * (a * r + b)**2
# done this way to approximate QuICC at most
elif operation == 'curl':
# prepare vectors for the derivative
x1a = np.array(x1)
x2a = np.array(x2)
x1a[1:] *= 2
x2a[1:] *= 2
# derivative in chebzshev space
dx1 = np.append(chebder(x1a), 0.)
d2x1 = np.append(chebder(dx1), 0.)
dx2 = np.append(chebder(x2a), 0.)
d2x2 = np.append(chebder(dx2), 0.)
# prepare derivatives for the DCT
dx1[1:] *= .5
d2x1[1:] *= .5
dx2[1:] *= .5
d2x2[1:] *= .5
# transform everything
r = idct(r1, type=2)
r = (a * r + b)
y1 = idct(x1, type=2)
y2 = idct(x2, type=2)
dy1 = idct(dx1, type=2)
dy2 = idct(dx2, type=2)
d2y1 = idct(d2x1, type=2)
d2y2 = idct(d2x2, type=2)
# compute the 2 pieces of the bilinear operator
Diss1 = r * d2y1 + 2 * dy1 - l * (l + 1) / r * y1
Diss2 = r * d2y2 + 2 * dy2 - l * (l + 1) / r * y2
y3 = Diss1 * Diss2
# bring back to spectral space
c = dct(y3, type=2) / (2 * len(y3))
c = c[:nr]
# represent the cheb indices in the mathematically correct way
"""
c[1:] *= 2.
# integrate, prepare index
idx = np.arange(len(c))
# and compute integration weight for integration from -1 to +1
w = (1+(-1)**idx)/(1-idx**2)
w[1]=0
#return np.sum(w*c)/4
x,w = chebgauss(len(y3))
return np.sum(w*y3)*a
"""
Ic = I1 * c
temp = chebval([xmax, xmin], Ic)
return temp[0] - temp[1]
def __init__(self,
Om0,
Ol0,
Ob0,
qs,
basis='Chebyshev',
zmin=0.0,
zmax=2.0,
free_knots=None,
sigma80=None,
H0=67.3,
Tcmb=2.7255,
Neff=3.046,
**kwargs):
IVCDM.__init__(self,
Om0,
Ol0,
Ob0,
q=0.0,
sigma80=sigma80,
H0=H0,
Tcmb=Tcmb,
Neff=Neff)
self.qs = np.asarray(qs) # array of expansion coefficients
self.Nq = self.qs.size
if basis in ['A', 'B', 'C', 'Chebyshev', 'cubic_spline']:
self.basis = str(basis)
else:
self.basis = 'A'
if self.basis == 'A': #1/(1+z)^n
self.cn_n = np.zeros(self.Nq)
self.cn_n[1:] = np.array(
[self.qs[n] / n for n in range(1, self.Nq)])
self.expsumc = np.exp(-np.sum(self.cn_n))
# construct coefficients (gamma) for rho cdm
N = self.Nq - 1 # index of last q parameter
alpha = np.zeros(2 * N + 1)
alpha[:self.Nq] = [
self.qs[k] / (self.qs[0] + k + 3) for k in range(self.Nq)
]
beta = np.zeros_like(alpha)
for k in range(1, self.Nq):
s = np.sum([
self.qs[n] * self.qs[k - n] / (k - n)
for n in range(0, k - 1 + 1)
])
beta[k] = s / (self.qs[0] + k + 3)
for k in range(self.Nq, 2 * self.Nq - 1):
s = np.sum([
self.qs[n] * self.qs[k - n] / (k - n)
for n in range(k - N, N + 1)
])
beta[k] = s / (self.qs[0] + k + 3)
self.gamma = self.expsumc * (alpha + beta)
self.sum_gamma = np.sum(self.gamma)
# different parametrisation of interaction
elif self.basis == 'Chebyshev':
self.zmin = zmin
self.zmax = zmax
self.Delta_z = (zmax - zmin) / 2.0
self.zbar = (zmax + zmin) / 2.0
self.zbarp1 = self.zbar + 1.0
# chebyshev coefficients of dDelta/dx and d^2Delta/dx^2
self.dcheb_coeff = chebyshev.chebder(self.qs, m=1) # 1st deriv
self.ddcheb_coeff = chebyshev.chebder(self.qs, m=2) # 2nd deriv
self.dxdz = 1.0 / self.Delta_z
# q0 - q1 + q2 - q3 + ... (-1)^N qN
alt_ones = np.ones((self.Nq, ))
alt_ones[1::2] = -1.0
self.Delta0 = np.sum(alt_ones * self.qs)
self.Delta0p1 = self.Delta0 + 1.0
elif self.basis == 'B': # (1+z)^n
self.qs_zder1 = polynomial.polyder(
self.qs, m=1) # coefficients of dDelta/dz
self.qs_zder2 = polynomial.polyder(
self.qs, m=2) # coefficients of d2Delta/dz2
self.cn_np3 = np.zeros(self.Nq + 3)
self.cn_np3[3:] = np.array(
[self.qs[n] / (n + 3.0) for n in range(self.Nq)])
self.sum_cn_np3 = np.sum(self.cn_np3)
self.Delta0 = np.sum(self.qs)
self.Delta0p1 = self.Delta0 + 1.0
elif self.basis == 'C': # 1/(1+z)^n = a^n
self.qs_ader1 = polynomial.polyder(
self.qs, m=1) # coefficients of dDelta/da
self.qs_ader2 = polynomial.polyder(
self.qs, m=2) # coefficients of d2Delta/da2
self.cn_nm3 = np.zeros(self.Nq - 3)
self.cn_nm3[1:] = np.array(
[self.qs[n] / (n - 3.0) for n in range(4, self.Nq)])
self.sum_cn_nm3 = np.sum(self.cn_nm3)
self.Delta0 = np.sum(self.qs)
self.Delta0p1 = self.Delta0 + 1.0
elif self.basis == 'cubic_spline':
assert free_knots[0] < free_knots[-1]
assert free_knots.size == self.Nq
assert free_knots[-1] < 1.0
# The _free_ knots are uniformly spaced in units of scale factor.
# The first and last knots are fixed and we insert them here:
a_first = 0.97 * free_knots[0]
a_last = 1.0
self.a_knots = np.concatenate(
([a_first], free_knots, [a_last])) # x
self.q_knots = np.concatenate(([0.0], self.qs, [0.0])) # y
# Construct splines (zero, first and second derivatives):
# 'clamped': first derivative is zero at endpoint
# 'natural': second derivative is zero at endpoint
self.Delta_spl = CubicSpline(x=self.a_knots,
y=self.q_knots,
bc_type=('clamped', 'natural'))
self.dDelta_spl = self.Delta_spl.derivative(nu=1)
self.d2Delta_spl = self.Delta_spl.derivative(nu=2)
self._points = list(1.0 / self.a_knots[::-1] - 1.0)
self.Omega_cdm_spl = None
self.Omega_vac_spl = None # only for cubic spline
def chebyshev_derivative(number, x):
"""Evaluate the derivative of `number`-th Chebyshev polynomial at `x`."""
coefficients = np.zeros(number + 1)
coefficients[number] = 1
return cheb.chebval(x, cheb.chebder(coefficients))
# Chebyshev-Gauss nodes and weights from spectralDNS.shen.shentransform import ChebyshevTransform CT = ChebyshevTransform(quad="GL") points, w = CT.points_and_weights(n+1, CT.quad) #points, w = n_cheb.chebgauss(n+1) #points = -points # Chebyshev Vandermonde matrix V = n_cheb.chebvander(points, n) scl=1.0 # First derivative matrix zero padded to (n+1)x(n+1) D1 = np.zeros((n+1,n+1)) D1[:-1,:] = n_cheb.chebder(np.eye(n+1), 1, scl=scl) # Second derivative matrix zero padded to (n+1)x(n+1) D2 = np.zeros((n+1,n+1)) D2[:-2,:] = n_cheb.chebder(np.eye(n+1), 2, scl=scl) Vx = np.dot(V, D1) Vxx = np.dot(V, D2) # Matrix of trial functions P = np.zeros((n+1,n+1)) P[:,0] = (V[:,0] - V[:,1])/2 P[:,1:-1] = V[:,:-2] - V[:,2:] P[:,-1] = (V[:,0] + V[:,1])/2
def lagrange(N, x0=x0, x1=x1, C_a=1.3, C_f=1.0, grid='cheb0'):
#x = np.linspace(x0, x1, N+1)
if grid == 'cheb0':
dummy = np.zeros(N)
dummy[-1] = 1
x = np.sort(np.array(cheb.chebroots(dummy).tolist() + [-1, 1]))
x = 0.5 * (x1 - x0) * x + 0.5 * (x1 + x0)
if grid == 'cheb1':
dummy = np.zeros(N + 1)
dummy[-1] = 1
x = np.sort(
np.array(cheb.chebroots(cheb.chebder(dummy)).tolist() + [-1, 1]))
x = 0.5 * (x1 - x0) * x + 0.5 * (x1 + x0)
if grid == 'legendre0':
dummy = np.zeros(N)
dummy[-1] = 1
x = np.sort(np.array(leg.legroots(dummy).tolist() + [-1, 1]))
x = 0.5 * (x1 - x0) * x + 0.5 * (x1 + x0)
if grid == 'legendre1':
dummy = np.zeros(N + 1)
dummy[-1] = 1
x = np.sort(
np.array(leg.legroots(leg.legder(dummy)).tolist() + [-1, 1]))
x = 0.5 * (x1 - x0) * x + 0.5 * (x1 + x0)
#x = 0.5*(x1-x0)*np.cos( np.pi / N * np.arange(N,-1,-1) )+ 0.5*(x1+x0)
#x = 0.5*(x1-x0)*np.cos( np.pi / (N+1) * ( np.arange(N+1,0,-1) -0.5) )+ 0.5*(x1+x0)
xx, yy = np.meshgrid(x, x)
AA = yy - xx
AA[AA == 0] = 1
a = np.prod(AA, axis=1)
#mth lagrangian polynomial
def phi(y, m):
"""Evaluates m-th Lagrangian polynomial at point y"""
return np.prod((y - x)[np.arange(len(x)) != m]) / a[m]
np.seterr(all='ignore')
D = 1 / AA
np.seterr(all='raise')
D = D - np.diag(np.diag(D))
D_diag = np.diag(np.sum(D, axis=1))
D = np.multiply(np.multiply(D.T, a).T, 1 / a) + D_diag
#Initialization of nonlinear part
Uxy = np.zeros([N + 1, N + 1, N + 1])
#Integration matrix for integrals of between x_k and x_N
#W1 = np.zeros( [N+1, N+1] )
#Integration matrix for integrals of between x_0 and x_k
for kk in xrange(N + 1):
for ii in xrange(N + 1):
#if kk<=ii:
#W1[kk, ii] = quad(phi, x[kk] , x[-1] , args=(ii,) )[0]
if kk >= ii:
#W2[kk, ii] = quad(phi, x[0] , x[kk] , args=(ii,) )[0]
dummy = x[kk] - x[ii] - xx
dummy[np.diag_indices(N + 1)] = 1
Uxy[kk, ii, :] = np.prod(dummy, axis=1) / a
#Aggregation in
a_ker = np.vectorize(partial(agg_kernel, C_a=C_a))
Ain = 0.5 * a_ker(yy - xx, xx)
Ain[0] = 0
#Aggregation out
Aout = a_ker(yy, xx - yy)
Aout[-1] = 0
#Initialization of fragmentation in
frag = np.vectorize(partial(fragmentation, C_f=C_f))
Fin = gam(yy, xx)
Fin[-1] = 0
Fin = np.multiply(Fin, frag(x))
#Fin[np.diag_indices(N+1)]*=0.5
return x, D, Fin, Ain, Aout, Uxy
C = ChebyshevCL(n)
s = C.get_device_name()
print('Accessing device: ' + s)
# Chebyshev spectral differentiation
x = C.coeff_to_nodal(e1)
f = np.exp(np.cos(np.pi*x))
fhat = C.nodal_to_coeff(f)
Dfhat = np.zeros(n)
Dfhat[:-1] = chebder(fhat)
Df = C.coeff_to_nodal(Dfhat)
fx = C.nodal_diff(f)
fig = plt.figure(1,(12,6))
ax1 = fig.add_subplot(121)
ax1.plot(x,Df,lw=2)
ax1.set_title('Modal differentiation')
ax2 = fig.add_subplot(122)
ax2.plot(x,fx,lw=2)
ax2.set_title('Nodal differentiation')
plt.show()
# Chebyshev-Gauss nodes and weights from spectralDNS.shen.shentransform import ChebyshevTransform CT = ChebyshevTransform(quad="GL") points, w = CT.points_and_weights(n + 1, CT.quad) #points, w = n_cheb.chebgauss(n+1) #points = -points # Chebyshev Vandermonde matrix V = n_cheb.chebvander(points, n) scl = 1.0 # First derivative matrix zero padded to (n+1)x(n+1) D1 = np.zeros((n + 1, n + 1)) D1[:-1, :] = n_cheb.chebder(np.eye(n + 1), 1, scl=scl) # Second derivative matrix zero padded to (n+1)x(n+1) D2 = np.zeros((n + 1, n + 1)) D2[:-2, :] = n_cheb.chebder(np.eye(n + 1), 2, scl=scl) Vx = np.dot(V, D1) Vxx = np.dot(V, D2) # Matrix of trial functions P = np.zeros((n + 1, n + 1)) P[:, 0] = (V[:, 0] - V[:, 1]) / 2 P[:, 1:-1] = V[:, :-2] - V[:, 2:] P[:, -1] = (V[:, 0] + V[:, 1]) / 2
def chebder(cs, m=1, scl=1):
from numpy.polynomial.chebyshev import chebder
return chebder(cs, m, scl)
def chebDiff(A, m=1):
return np_cheb.chebder(A, m=m)
