def _compute_Sq_with_reduced_lmax( self, new_lmax):
"""
model: SphericalModel object
pr: PhaseRetriver object
"""
new_sh_obj = shtns.sht(new_lmax)
old_sh_obj = self.sh
new_S_q = np.zeros_like( self.S_q )
old_all_slm = self.all_slm
new_sh_obj.set_grid( self.n_theta, self.n_phi )
n_q = self.n_q
new_all_slm = np.zeros((self.n_q, int((new_lmax+2)*(new_lmax+1)/2) )
, dtype=np.complex128)
for qid in range(n_q):
new_slm = new_sh_obj.spec_array()
for lid in range(new_lmax+1):
for mid in range(lid+1):
new_slm[new_sh_obj.idx(lid, mid)] = old_all_slm[qid][old_sh_obj.idx(lid,mid)]
new_S_q[qid] = new_sh_obj.synth(new_slm)
new_all_slm[qid, :] = new_slm
return new_S_q, new_all_slm, new_sh_obj
def get_Sq_with_reduced_lmax(new_lmax, model=None, pr=None):
"""
model: SphericalModel object
pr: PhaseRetriver object
"""
new_sh_obj = shtns.sht(new_lmax)
if model is None:
old_sh_obj = pr.sh
new_S_q = np.zeros_like(pr.I_guess)
old_all_slm = pr.all_slm_guess
new_sh_obj.set_grid(pr.n_theta, pr.n_phi)
n_q = pr.n_q
else:
old_sh_obj = model.sh
new_S_q = np.zeros_like(model.S_q)
old_all_slm = model.all_slm
new_sh_obj.set_grid(model.n_theta, model.n_phi)
n_q = model.n_q
for qid in range(n_q):
new_slm = new_sh_obj.spec_array()
for lid in range(new_lmax + 1):
for mid in range(lid + 1):
new_slm[new_sh_obj.idx(lid,
mid)] = old_all_slm[qid][old_sh_obj.idx(
lid, mid)]
new_S_q[qid] = new_sh_obj.synth(new_slm)
return new_S_q
def power_spectrum(Br, nphi, ntheta, l_trunc, r=ro):
'''
Obtain power spectrum of core surface field to degree l_trunc (default at the CMB)
'''
# SH transform from spatial to spectral space
m_max = l_trunc
# by default 'flag = sht_quick_init' uses gaussian grid
sh = shtns.sht(l_trunc, m_max)
nlat, nlon = sh.set_grid(nphi=nphi, nlat=ntheta)
vr = Br.T.astype(
'float64') # NOTE: array has to be dtype='float64' and not 'float32'
clm = sh.analys(vr) # spatial to spectral
# Construct Gauss coefficients
glm = np.zeros(sh.nlm)
hlm = np.zeros(sh.nlm)
for l in range(l_trunc + 1):
fac = 1 / (l + 1
) # cmb_radius/(l+1) # from Br potential to full potential
for m in range(l + 1):
glm[sh.idx(l, m)] = fac * np.real(clm[sh.idx(l, m)])
hlm[sh.idx(l, m)] = fac * np.imag(clm[sh.idx(l, m)])
# Power spectrum
degrees = np.arange(1, l_trunc + 1, 1)
spectrum = np.zeros(l_trunc + 1)
for l in range(l_trunc + 1):
fac = (l + 1) * (earth_radius / cmb_radius)**(
2 * l + 4) # TODO: check factors using time average
sum_m = 0
for m in range(l + 1):
sum_m += glm[sh.idx(l, m)]**2 + hlm[sh.idx(l, m)]**2
spectrum[l] = fac * sum_m
spectrum = spectrum[1:] # ignore g00
return degrees, spectrum, clm, glm, hlm
def setup(self, file_info, anti_aliasing=False):
import shtns
import numpy as np
if file_info['modes_m_max'] != file_info['modes_m_max']:
raise Exception("Only num_lon == num_lat supported")
ntrunc = file_info['modes_n_max']
self._shtns = shtns.sht(ntrunc, ntrunc, 1,
shtns.sht_orthonormal + shtns.SHT_NO_CS_PHASE)
nlons = (ntrunc + 1) * 2
nlats = (ntrunc + 1)
if anti_aliasing:
if nlons & 1:
raise Exception(
"Only even numbers of longitudinal coordinates allowed for anti-aliasing"
)
if nlats & 1:
raise Exception(
"Only even numbers of latitudinal coordinates allowed for anti-aliasing"
)
print("Anti-aliasing:")
print(" + old lon/lat: ", nlons, nlats)
nlons += nlons // 2
nlats += nlats // 2
print(" + new lon/lat: ", nlons, nlats)
if file_info['grid_type'] == 'GAUSSIAN':
#self._shtns.set_grid(nlats,nlons,shtns.sht_gauss_fly|shtns.SHT_PHI_CONTIGUOUS, 1.e-10)
self._shtns.set_grid(
nlats, nlons, shtns.sht_quick_init | shtns.SHT_PHI_CONTIGUOUS,
0)
elif file_info['grid_type'] == 'REGULAR':
#self._shtns.set_grid(nlats,nlons,shtns.sht_reg_dct|shtns.SHT_PHI_CONTIGUOUS, 1.e-10)
self._shtns.set_grid(nlats, nlons,
shtns.sht_reg_dct | shtns.SHT_PHI_CONTIGUOUS,
0)
else:
raise Exception("Grid type '" + file_info['grid_type'] +
"' not supported!")
self.lats = np.arcsin(self._shtns.cos_theta)
self.lons = (2. * np.pi / nlons) * np.arange(nlons)
self.nlons = nlons
self.nlats = nlats
self.ntrunc = ntrunc
self.nlm = self._shtns.nlm
self.degree = self._shtns.l
self.lap = -self.degree * (self.degree + 1.0).astype(np.complex)
self.invlap = np.zeros(self.lap.shape, self.lap.dtype)
self.invlap[1:] = 1. / self.lap[1:]
self.lap = self.lap / self.rsphere**2
self.invlap = self.invlap * self.rsphere**2
def __init__(self,nlons,nlats,ntrunc,rsphere,gridtype='gaussian'):
"""
initialize
nlons: number of longitudes
nlats: number of latitudes
"""
self._shtns = shtns.sht(ntrunc, ntrunc, 1, \
shtns.sht_orthonormal+shtns.SHT_NO_CS_PHASE)
if gridtype == 'gaussian':
#self._shtns.set_grid(nlats,nlons,shtns.sht_gauss_fly|shtns.SHT_PHI_CONTIGUOUS,1.e-10)
self._shtns.set_grid(nlats,nlons,shtns.sht_quick_init|shtns.SHT_PHI_CONTIGUOUS,1.e-10)
elif gridtype == 'regular':
self._shtns.set_grid(nlats,nlons,shtns.sht_reg_dct|shtns.SHT_PHI_CONTIGUOUS,1.e-10)
self.lats = np.arcsin(self._shtns.cos_theta)
self.lons = (2.*np.pi/nlons)*np.arange(nlons)
self.nlons = nlons
self.nlats = nlats
self.ntrunc = ntrunc
self.nlm = self._shtns.nlm
self.degree = self._shtns.l
self.lap = -self.degree*(self.degree+1.0).astype(np.complex)
self.invlap = np.zeros(self.lap.shape, self.lap.dtype)
self.invlap[1:] = 1./self.lap[1:]
self.rsphere = rsphere
self.lap = self.lap/rsphere**2
self.invlap = self.invlap*rsphere**2
def write_model_h5(model, lmax, fname, varpath):
# assumes the model is sampled on a regular grid where the poles are not
# included
# prepare the transform
mmax = lmax
sh = shtns.sht(lmax, mmax)
npts_ph, npts_th = model.shape
grid_typ = shtns.sht_reg_fast | shtns.SHT_THETA_CONTIGUOUS
grid_typ = grid_typ | shtns.SHT_LOAD_SAVE_CFG
polar_opt_threshold = 1.0e-10
sh.set_grid(npts_th,
npts_ph,
flags=grid_typ,
polar_opt=polar_opt_threshold)
# apply transform
ylm = sh.analys(model)
with File(fname, 'r+') as f:
path_l = os.path.split(varpath)[0] + '/l'
if not path_l in f:
f.create_dataset(path_l, data=sh.l)
path_m = os.path.split(varpath)[0] + '/m'
if not path_m in f:
f.create_dataset(path_m, data=sh.m)
f.create_dataset(varpath, data=ylm)
return
def __init__(self, nlons, nlats, ntrunc, rsphere, gridtype='gaussian'):
"""initialize
nlons: number of longitudes
nlats: number of latitudes"""
self._shtns = shtns.sht(ntrunc, ntrunc, 1,
shtns.sht_orthonormal + shtns.SHT_NO_CS_PHASE)
if gridtype == 'gaussian':
self._shtns.set_grid(
nlats, nlons, shtns.sht_quick_init | shtns.SHT_PHI_CONTIGUOUS,
0)
elif gridtype == 'regular':
self._shtns.set_grid(nlats, nlons,
shtns.sht_reg_dct | shtns.SHT_PHI_CONTIGUOUS,
0)
self.lats = np.arcsin(self._shtns.cos_theta)
self.lons = (2. * np.pi / nlons) * np.arange(nlons)
self.nlons = nlons
self.nlats = nlats
self.ntrunc = ntrunc
self.nlm = self._shtns.nlm
self.degree = self._shtns.l
self.lap = -self.degree * (self.degree + 1.0).astype(complex)
self.invlap = np.zeros(self.lap.shape, self.lap.dtype)
self.invlap[1:] = 1. / self.lap[1:]
self.rsphere = rsphere
self.lap = self.lap / rsphere**2
self.invlap = self.invlap * rsphere**2
print("N: " + str(self.nlons) + ", " + str(self.nlats))
print("Mtrunc: " + str(self.ntrunc))
print("Nlm: " + str(self.nlm))
def __init__(self, Lmax, Mmax, fmetric=None, fricci=None):
self.grid = shtns.sht(Lmax,Mmax, norm=shtns.SHT_REAL_NORM, nthreads=4)
self.Lmax = Lmax
self.nTheta, self.nPhi = self.grid.set_grid()
self.nlm = self.grid.nlm
self.numTerms = (self.Lmax+1)*(self.Lmax+1)
self.extents = np.array([self.nTheta,self.nPhi])
self.l, self.m = YlmIndex(np.arange(self.numTerms))
self.index = np.array([self.grid.idx(int(self.l[i]),int(abs(self.m[i]))) for i in xrange(self.numTerms)])
#grid coordinates and properties:
self.theta, self.phi = np.meshgrid(
np.arccos(np.polynomial.legendre.leggauss(self.nTheta)[0])[::-1],
np.linspace(0,2*pi*(1-1./self.nPhi),self.nPhi),
sparse=False)
self.theta, self.phi = self.theta.T, self.phi.T
self.gaussian_weights = np.polynomial.legendre.leggauss(self.nTheta)[1]
self.costheta = np.cos(self.theta)
self.sintheta = np.sin(self.theta)
#Metric on grid points:
if fmetric == None:
self.gthth = np.zeros(self.extents)
self.gphph = np.zeros(self.extents)
self.gthph = np.zeros(self.extents)
else:
self.gthth, self.gphph, self.gthph = fmetric(self.theta, self.phi)
self.UpdateMetric()
# Compute ricci scalar, or not...
if fricci==None:
self.ricci = None
else:
self.ricci = fricci(self.theta, self.phi)
def write_model_ylm(model, lmax, fname):
# assumes the model is sampled on a regular grid where the poles are not
# included
# prepare the transform
mmax = lmax
sh = shtns.sht(lmax, mmax)
npts_ph, npts_th = model.shape
grid_typ = shtns.sht_reg_fast | shtns.SHT_THETA_CONTIGUOUS
grid_typ = grid_typ | shtns.SHT_LOAD_SAVE_CFG
polar_opt_threshold = 1.0e-10
sh.set_grid(npts_th,
npts_ph,
flags=grid_typ,
polar_opt=polar_opt_threshold)
# apply transform
ylm = sh.analys(model)
with open(fname, 'w') as f:
for l in zip(sh.l, sh.m, np.real(ylm), np.imag(ylm)):
f.write('%5d %5d %14e %14e\n' % l)
return
def __init__(self,nlons,nlats,ntrunc,rsphere,gridtype='gaussian'):
"""initialize
nlons: number of longitudes
nlats: number of latitudes
ntrunc: spectral truncation
rsphere: sphere radius (m)
gridtype: 'gaussian' (default) or 'regular'"""
self._shtns = shtns.sht(ntrunc, ntrunc, 1,\
shtns.sht_fourpi|shtns.SHT_NO_CS_PHASE)
if gridtype == 'gaussian':
self._shtns.set_grid(nlats,nlons,shtns.sht_quick_init|shtns.SHT_PHI_CONTIGUOUS,1.e-8)
elif gridtype == 'regular':
self._shtns.set_grid(nlats,nlons,shtns.sht_reg_dct|shtns.SHT_PHI_CONTIGUOUS,1.e-8)
#self._shtns.print_info()
self.lats = np.arcsin(self._shtns.cos_theta)
self.lons = (2.*np.pi/nlons)*np.arange(nlons)
self.nlons = nlons
self.nlats = nlats
self.ntrunc = ntrunc
self.nlm = self._shtns.nlm
self.degree = self._shtns.l
self.order = self._shtns.m
if gridtype == 'gaussian':
self.gauwts =\
np.concatenate((self._shtns.gauss_wts(),self._shtns.gauss_wts()[::-1]))
else:
self.gauwts = None
self.gridtype = gridtype
self.lap = -self.degree*(self.degree+1.0).astype(np.complex)
self.invlap = np.zeros(self.lap.shape, self.lap.dtype)
self.invlap[1:] = 1./self.lap[1:]
self.rsphere = rsphere
self.lap = self.lap/rsphere**2
self.invlap = self.invlap*rsphere**2
def __init__(self, l_max, force_python=False):
"""initialize a spherical harmonic transform object"""
self._l_max = l_max
if not force_python:
try:
import shtns
self._shtns = shtns.sht(l_max, l_max)
except ImportError:
LOG.debug('Could not import shtns')
def __init__(self, l_max, force_python=False):
"""initialize a spherical harmonic transform object"""
self._l_max = l_max
if not force_python:
try:
import shtns
self._shtns = shtns.sht(l_max, l_max)
except ImportError:
LOG.debug('Could not import shtns')
def fcf(run_ID, directory, l_trunc=8):
'''
Flux concentration factor used in measure of semblance (Christensen et al. 2010). Note that we truncate at degree and order 8
'''
St_file = list_St_files(run_ID, directory) # Find all St_no in folder
n = len(St_file)
# Loop over time
FCF_s = []
i = 0
for file in St_file:
print('Loading {} ({}/{})'.format(file, i + 1, n))
filename = '{}/{}'.format(directory, file)
(_, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, ntheta, nphi, _, _,
theta, phi, _, _, Br, _) = surfaceload(filename)
# Truncate degree 8
m_max = l_trunc
sh = shtns.sht(l_trunc, m_max)
sh = shtns.sht(l_trunc, m_max)
nlat, nlon = sh.set_grid(nphi=nphi, nlat=ntheta)
# NOTE: array has to be dtype='float64' and not 'float32'
vr = Br.T.astype('float64')
clm = sh.analys(vr) # spatial to spectral
Br_f = sh.synth(clm) # spectral to spatial
Br_4 = (Br_f**4).mean()
Br_2 = (Br_f**2).mean()
FCF = (Br_4 - Br_2**2) / Br_2**2
# Append
FCF_s.append(np.mean(FCF))
i += 1
# Time average (should really divide by dt but n is good enough)
FCF_out = sum(FCF_s) / n
# Save
with h5py.File('{}/fcf'.format(directory), 'w') as f:
f.create_dataset('FCF_out', data=FCF_out)
print('{}/fcf saved'.format(directory))
return FCF_out
def __init__(self,nlons,nlats,ntrunc,rsphere,gridtype='gaussian'):
"""initialize
nlons: number of longitudes
nlats: number of latitudes
"""
self._shtns = shtns.sht(ntrunc, ntrunc, 1, shtns.sht_orthonormal+shtns.SHT_NO_CS_PHASE)
if gridtype == 'gaussian':
self._shtns.set_grid(nlats,nlons,shtns.sht_quick_init|shtns.SHT_PHI_CONTIGUOUS,0)
elif gridtype == 'regular':
self._shtns.set_grid(nlats,nlons,shtns.sht_reg_dct|shtns.SHT_PHI_CONTIGUOUS,0)
self.lats = np.arcsin(self._shtns.cos_theta)
self.lons = (2.*np.pi/nlons)*np.arange(nlons)
def filter_field(Br, nphi, ntheta, l_trunc):
'''Use SHTns to truncate field to degree l_trunc'''
m_max = l_trunc
sh = shtns.sht(
l_trunc,
m_max) # by default 'flag = sht_quick_init' uses gaussian grid
nlat, nlon = sh.set_grid(
nphi=nphi,
nlat=ntheta) # NOTE: array has to be dtype='float64' and not 'float32'
vr = Br.T.astype('float64')
ylm = sh.analys(vr) # spatial to spectral
Br_f = sh.synth(ylm) # spectral to spatial
return Br_f
def extrapot(lmax,rcmb,brcmb,rout):
nphi, ntheta = brcmb.shape
nrout = len(rout)
polar_opt = 1e-10
lmax = int(nphi/3)
mmax = lmax
try:
import shtns
except ImportError:
print("Potential extrapolation requires the SHTns library")
print("It can be obtained here: https://bitbucket.org/nschaeff/shtns")
norm=shtns.sht_orthonormal | shtns.SHT_NO_CS_PHASE
sh = shtns.sht(lmax,mmax=mmax,norm=norm)
ntheta, nphi = sh.set_grid(ntheta, nphi, polar_opt=polar_opt)
L = sh.l * (sh.l + 1)
brlm = sh.analys(brcmb.T)
bpolcmb = np.zeros_like(brlm)
bpolcmb[1:] = rcmb**2 * brlm[1:]/L[1:]
btor = np.zeros_like(brlm)
brout = np.zeros([ntheta,nphi,nrout])
btout = np.zeros([ntheta,nphi,nrout])
bpout = np.zeros([ntheta,nphi,nrout])
for k,radius in enumerate(rout):
print(("%d/%d" %(k,nrout)))
radratio = rcmb/radius
bpol = bpolcmb * radratio**(sh.l)
brlm = bpol * L/radius**2
brout[...,k] = sh.synth(brlm)
dbpoldr = -sh.l/radius * bpol
slm = dbpoldr
btout[...,k], bpout[...,k] = sh.synth(slm,btor)
brout = np.transpose(brout,(1,0,2))
btout = np.transpose(btout,(1,0,2))
bpout = np.transpose(bpout,(1,0,2))
return brout, btout, bpout
def setup(self, file_info, data):
import shtns
import numpy as np
if file_info['modes_m_max'] != file_info['modes_m_max']:
raise Exception("Only num_lon == num_lat supported")
ntrunc = file_info['modes_n_max']
self._shtns = shtns.sht(
ntrunc, ntrunc, 1,
shtns.sht_orthonormal + shtns.SHT_NO_CS_PHASE)
nlons = (ntrunc + 1) * 2
nlats = (ntrunc + 1)
if file_info['grid_type'] == 'GAUSSIAN':
#self._shtns.set_grid(nlats,nlons,shtns.sht_gauss_fly|shtns.SHT_PHI_CONTIGUOUS, 1.e-10)
self._shtns.set_grid(
nlats, nlons,
shtns.sht_quick_init | shtns.SHT_PHI_CONTIGUOUS, 0)
elif file_info['grid_type'] == 'REGULAR':
#self._shtns.set_grid(nlats,nlons,shtns.sht_reg_dct|shtns.SHT_PHI_CONTIGUOUS, 1.e-10)
self._shtns.set_grid(
nlats, nlons, shtns.sht_reg_dct | shtns.SHT_PHI_CONTIGUOUS,
0)
else:
raise Exception("Grid type '" + file_info['grid_type'] +
"' not supported!")
self.lats = np.arcsin(self._shtns.cos_theta)
self.lons = (2. * np.pi / nlons) * np.arange(nlons)
self.nlons = nlons
self.nlats = nlats
self.ntrunc = ntrunc
self.nlm = self._shtns.nlm
self.degree = self._shtns.l
self.lap = -self.degree * (self.degree + 1.0).astype(np.complex)
self.invlap = np.zeros(self.lap.shape, self.lap.dtype)
self.invlap[1:] = 1. / self.lap[1:]
self.lap = self.lap / self.rsphere**2
self.invlap = self.invlap * self.rsphere**2
sh_data = np.frombuffer(data, dtype=np.complex128)
return sh_data
def simulate( self, q_values,
n_theta, n_phi,
lmax, n_psi,
dont_rotate = True,
make_positive = True):
"""
Simulates S(q), projects into spherical harmonics, and computes correlation
Parameters
----------
q_values : np.array, q magnitudes to be simulated
n_theta : int, number of thetas in q space, range = [0, pi]
This is not the same as Bragg theta; this theta = Bragg theta + pi/2
n_phi: int, number of phis, range = [0, 2pi]
lmax: int, maximum order of spherical harmonics projection
dont_rotate: default True, do not rotate the traj when simulating S(q)
"""
self.n_theta = n_theta
self.n_phi = n_phi
self.q_values = q_values
self.n_q = q_values.shape[0]
self.lmax = lmax
self.n_psi = n_psi
self.sh = shtns.sht(lmax)
self.sh.set_grid( n_theta, n_phi )
# simulate S_q
self._compute_Sq( self.model, q_values, dont_rotate)
# make sure that the inverse spherical transformation always gets a positive intensity map
if make_positive:
self._make_Sq_positive()
# project into spherical harmonics
self._sph_harm_project()
# sum slm to get cl
self._leg_coefs_from_sph_harm()
# compute correlations
self._sph_coefs_to_corr()
def __init__(self,
R,
lmax,
nphi=None,
ntheta=None,
delta_phi=1.,
verbose=False):
self.R = R
self.lmax = lmax
# prepare the spherical harmonic transform
self.nphi = nphi or self.lmax * 4
self.ntheta = ntheta or self.lmax * 2
self.mmax = self.lmax
self.sh = shtns.sht(self.lmax, self.mmax)
grid_typ = shtns.sht_gauss | shtns.SHT_THETA_CONTIGUOUS | \
shtns.SHT_LOAD_SAVE_CFG
polar_opt_threshold = 1.0e-10
self.sh.set_grid(self.ntheta,
self.nphi,
flags=grid_typ,
polar_opt=polar_opt_threshold)
self.theta = np.arccos(self.sh.cos_theta)
self.phi = np.linspace(0, 2 * np.pi, self.nphi)
self.workspace = self.sh.spec_array()
self.delta_phi = np.deg2rad(delta_phi)
self.group_velocity = None
self.phase_velocity = None
self.topo = None
self.group_velocity_ylm = None
self.phase_velocity_ylm = None
self.topo_ylm = None
self.group_velocity_notrot = None
self.phase_velocity_notrot = None
self.topo_notrot = None
self.group_velocity_ylm_notrot = None
self.phase_velocity_ylm_notrot = None
self.topo_ylm_notrot = None
def __init__(self, nlons, nlats, ntrunc, rsphere, gridtype='gaussian'):
"""initialize
nlons: number of longitudes
nlats: number of latitudes
ntrunc: spectral truncation
rsphere: sphere radius (m)
gridtype: 'gaussian' (default) or 'regular'"""
self._shtns = shtns.sht(ntrunc, ntrunc, 1,\
shtns.sht_fourpi|shtns.SHT_NO_CS_PHASE)
if gridtype == 'gaussian':
self._shtns.set_grid(
nlats, nlons, shtns.sht_quick_init | shtns.SHT_PHI_CONTIGUOUS,
1.e-8)
elif gridtype == 'regular':
self._shtns.set_grid(nlats, nlons,
shtns.sht_reg_dct | shtns.SHT_PHI_CONTIGUOUS,
1.e-8)
#self._shtns.print_info()
self.lats = np.arcsin(self._shtns.cos_theta)
self.lons = (2. * np.pi / nlons) * np.arange(nlons)
self.nlons = nlons
self.nlats = nlats
self.ntrunc = ntrunc
self.nlm = self._shtns.nlm
self.degree = self._shtns.l
self.order = self._shtns.m
if gridtype == 'gaussian':
self.gauwts =\
np.concatenate((self._shtns.gauss_wts(),self._shtns.gauss_wts()[::-1]))
else:
self.gauwts = None
self.gridtype = gridtype
self.lap = -self.degree * (self.degree + 1.0).astype(np.complex)
self.invlap = np.zeros(self.lap.shape, self.lap.dtype)
self.invlap[1:] = 1. / self.lap[1:]
self.rsphere = rsphere
self.lap = self.lap / rsphere**2
self.invlap = self.invlap * rsphere**2
def setup(self, file_info, anti_aliasing):
import shtns
import numpy as np
if file_info['modes_m_max'] != file_info['modes_m_max']:
raise Exception("Only num_lon == num_lat supported")
ntrunc = file_info['modes_n_max']
self._shtns = shtns.sht(ntrunc, ntrunc, 1,
shtns.sht_orthonormal + shtns.SHT_NO_CS_PHASE)
nlons = (ntrunc + 1) * 2
nlats = (ntrunc + 1)
if anti_aliasing:
if nlons & 1:
raise Exception(
"Only even numbers of longitudes coordinates allowed for anti-aliasing"
)
if nlats & 1:
raise Exception(
"Only even numbers of latitudinal coordinates allowed for anti-aliasing"
)
print("Anti-aliasing:")
print(" + old lon/lat: ", nlons, nlats)
nlons += nlons // 2
nlats += nlats // 2
print(" + new lon/lat: ", nlons, nlats)
self._shtns.set_grid(nlats, nlons,
shtns.sht_quick_init | shtns.SHT_PHI_CONTIGUOUS,
0)
self.lats = np.arcsin(self._shtns.cos_theta)
self.lons = (2. * np.pi / nlons) * np.arange(nlons)
def gauss_coeffs(Br, nphi, ntheta, l_trunc):
'''
Obtain Gauss coefficients from core surface field
'''
# SH transform from spatial to spectral space
m_max = l_trunc
# by default 'flag = sht_quick_init' uses gaussian grid
sh = shtns.sht(l_trunc, m_max)
nlat, nlon = sh.set_grid(nphi=nphi, nlat=ntheta)
vr = Br.T.astype(
'float64') # NOTE: array has to be dtype='float64' and not 'float32'
clm = sh.analys(vr) # spatial to spectral
# Construct Gauss coefficients
glm = np.zeros(sh.nlm)
hlm = np.zeros(sh.nlm)
for l in range(l_trunc + 1):
fac = 1 / (l + 1
) # cmb_radius/(l+1) # from Br potential to full potential
for m in range(l + 1):
glm[sh.idx(l, m)] = fac * np.real(clm[sh.idx(l, m)])
hlm[sh.idx(l, m)] = fac * np.imag(clm[sh.idx(l, m)])
return glm, hlm
def filter_model_shtns(model,
lmax,
nphi_out=None,
ntheta_out=None,
order=2,
lmax_transform=None):
# assumes the model is sampled on a regular grid where the poles are not
# included
if lmax_transform is None:
lmax_transform = lmax * 2
nphi, ntheta = model.shape
lmax_transform = min(min(ntheta - 2, nphi / 2 - 1), lmax_transform)
# prepare the transform
mmax = lmax_transform
sh = shtns.sht(lmax_transform, mmax)
grid_typ = shtns.sht_reg_fast | shtns.SHT_THETA_CONTIGUOUS
grid_typ = grid_typ | shtns.SHT_LOAD_SAVE_CFG
polar_opt_threshold = 1.0e-10
sh.set_grid(ntheta, nphi, flags=grid_typ, polar_opt=polar_opt_threshold)
# apply transform
ylm = sh.analys(model)
ylm *= 1. / (1. + (sh.l * 1. / lmax)**(2 * order))
if nphi_out and ntheta_out:
sh.set_grid(ntheta_out,
nphi_out,
flags=grid_typ,
polar_opt=polar_opt_threshold)
return sh.synth(ylm)
def load(self, filename):
ff = h5py.File(filename, 'r')
self.S_q = ff['intensity'].value
self.n_theta = self.S_q.shape[1]
self.n_phi = self.S_q.shape[2]
self.lmax = ff['lmax'].value
self.sh = shtns.sht(self.lmax)
# generate the shtns object
self.sh.set_grid( self.n_theta, self.n_phi )
self.Sq_original = ff['intensity_original'].value
self.all_slm =ff['slm'].value
self.corr = ff['corr'].value
self.cl = ff['leg_coefs'].value
self.cospsi = ff['cospsi'].value
self.n_psi = self.cospsi.size
self.q_values = ff['qvalues'].value
self.n_q = self.q_values.size
ff.close()
#zsurf = nfile.variables['zsurf'] # lat, lon
#pres_half = nfile.variables['pres_half'] # time, phalf, lat, lon
#height_half = nfile.variables['height_half'] # time, phalf, lat, lon
#pres_full = nfile.variables['pres_full'] # time, pfull, lat, lon
#height_full = nfile.variables['height'] # time, pfull, lat, lon
field = nfile.variables[fieldname] # time, pfull, lat, lon
#field_window = field[-w:,:,:,:]
field_state = field[-1, :, :, :] # snapshot for now
nlon = field.shape[3]
nlat = field.shape[2]
mmax = min(tnum, (nlon - 1) // 2)
print tnum, mmax
sh = shtns.sht(tnum, mmax) # reduce 2nd argument here as needed
#nlm = sh.nlm # size of combined lm index
sh.set_grid(nlat, nlon) # build default grid (gauss grid, phi-contiguous)
#print np.sctypeDict.keys()
field_k_l_m = np.empty((field.shape[1], tnum + 1, mmax + 1), dtype='complex64')
for k in range(0, field.shape[1]):
field_lm = sh.analys(
field_state[k, :, :].astype('float64')) # sh.spec_array implicit?
for l in range(0, tnum + 1):
for m in range(0, mmax + 1):
field_k_l_m[k, l, m] = field_lm[sh.idx(l, m)] if m <= l else np.nan
uspec = np.abs(field_k_l_m)**2 # pfull, l, m
uspec[:, :,
1:] *= 2 # weight factor for missing m<0 (reality condition); 'optional' if not thinking of as summing displayed components
def __init__(self,
lmax=15,
mmax=None,
mres=1,
norm=shtns.sht_fourpi,
nlat=None,
nlon=None,
flags=(shtns.sht_quick_init | shtns.SHT_PHI_CONTIGUOUS
| shtns.SHT_SOUTH_POLE_FIRST),
polar_opt=1.0e-8,
nl_order=2,
radius=radius_earth):
#print(lmax,mmax,mres,nlat,nlon,flags,polar_opt,nl_order,radius)
#print('flags',flags,shtns.sht_quick_init,shtns.SHT_PHI_CONTIGUOUS,shtns.SHT_SOUTH_POLE_FIRST)
if lmax is None and nlat is None:
raise ValueError('lmax or nlat should be given.')
elif lmax is None:
lmax = compute_lmax(nlat)
#print('lmax',lmax)
self.lmax = int(lmax)
self.radius = float(radius)
#print(self.radius,radius_earth)
if mmax is None and mres == 1:
# triangular truncation
self.mmax = self.lmax
self.mres = 1
else:
self.mmax = int(float(lmax) / mres)
self.mres = mres
#print('norm',norm,self.lmax,self.mmax,self.mres)
self.sh = shtns.sht(self.lmax, self.mmax, self.mres, norm=norm)
bin_flags = bin_int(flags, 14)
#print(bin_flags,bin_flags[-14])
if bin_flags[-14] == '1':
self.order_lat = 'south_to_north'
else:
self.order_lat = 'north_to_south'
# in shtns, 0 means None
if nlat is None:
nlat = 0
if nlon is None:
nlon = 0
self.nlat, self.nlon = self.sh.set_grid(nlat=nlat,
nphi=nlon,
flags=flags,
polar_opt=polar_opt,
nl_order=nl_order)
#print('flags',flags)
self.sin_lats = self.sh.cos_theta
self.lons = np.arange(self.nlon) * 360. / (self.nlon * self.mres)
self.lats = np.arcsin(self.sin_lats) / np.pi * 180.
self.l_idx = self.sh.l
#print('l_idx',self.l_idx)
self.l2_idx = self.l_idx * (self.l_idx + 1)
#print('l2_idx',self.l2_idx)
self.m_idx = self.sh.m
self.nlm = self.sh.nlm
self.delta = 360. / (self.nlon * self.mres)
# create arrays 2D lats et lons
self.LONS, self.LATS = np.meshgrid(self.lons, self.lats)
self.cosLATS = np.cos(self.LATS / 180 * np.pi)
self.lrange = np.arange(self.lmax + 1)
self.l2_l = self.lrange * (self.lrange + 1)
self.kh_l = np.sqrt(self.l2_l) / self.radius
self._complex64_save_netCFD = np.dtype([('real', np.float32),
('imag', np.float32)])
return vx, vy, vz
m = 20
lmax = 20
lUsr = int(sys.argv[1])
mUsr = int(sys.argv[2])
ntheta, nphi = [128, 256]
polar_opt = 1e-10
norm = shtns.sht_orthonormal | shtns.SHT_NO_CS_PHASE
sh = shtns.sht(lmax, mmax=m, norm=norm, nthreads=1)
ntheta, nphi = sh.set_grid(ntheta, nphi, polar_opt=polar_opt)
S = sh.spec_array()
T = sh.spec_array()
T[sh.idx(lUsr, mUsr)] = 1.
utheta, uphi = sh.synth(S, T)
psi = sh.synth(T)
x, y, z, th2D, p2D = get_grid(np.arccos(sh.cos_theta),
np.linspace(0., 2 * np.pi, nphi))
ux, uy, uz = get_cart(utheta, uphi, th2D, p2D)
def timeLongitude(self,
removeMean=True,
lat0=0.,
levels=12,
cm='RdYlBu_r',
deminc=True,
shtns_lib='shtns'):
"""
Plot the time-longitude diagram of Br (input latitude can be chosen)
.. warning:: the python bindings of `SHTns <https://bitbucket.org/bputigny/shtns-magic>`_ are mandatory to use this plotting function!
:param lat0: value of the latitude
:type lat0: float
:param levels: number of contour levels
:type levels: int
:param cm: name of the colormap
:type cm: str
:param deminc: a logical to indicate if one wants do get rid of the
possible azimuthal symmetry
:type deminc: bool
:param shtns_lib: version of shtns library used: can be either 'shtns'
or 'shtns-magic'
:type shtns_lib: char
:param removeMean: remove the time-averaged part when set to True
:type removeMean: bool
"""
# The python bindings of shtns are mandatory to use this function !!!
import shtns
if removeMean:
blmCut = self.blm - self.blm.mean(axis=0)
else:
blmCut = self.blm
# Define shtns setup
sh = shtns.sht(int(self.l_max_cmb),
int(self.m_max_cmb / self.minc),
mres=int(self.minc),
norm=shtns.sht_orthonormal | shtns.SHT_NO_CS_PHASE)
polar_opt_threshold = 1e-10
nlat = max(int(self.l_max_cmb * (3. / 2. / 2.) * 2.), 192)
nphi = 2 * nlat / self.minc
nlat, nphi = sh.set_grid(nlat, nphi, polar_opt=polar_opt_threshold)
th = np.linspace(np.pi / 2., -np.pi / 2., nlat)
lat0 *= np.pi / 180.
mask = np.where(abs(th - lat0) == abs(th - lat0).min(), 1, 0)
idx = np.nonzero(mask)[0][0]
# Transform data on grid space
BrCMB = np.zeros((self.nstep, nphi, nlat), 'Float64')
if deminc:
dat = np.zeros((self.nstep, self.minc * nphi + 1), 'Float64')
else:
dat = np.zeros((self.nstep, nphi), 'Float64')
for k in range(self.nstep):
tmp = sh.synth(blmCut[k, :] * sh.l * (sh.l + 1) / self.rcmb**2)
tmp = tmp.T # Longitude, Latitude
if shtns_lib == 'shtns-magic':
BrCMB[k, ...] = rearangeLat(tmp)
else:
BrCMB[k, ...] = tmp
if deminc:
dat[k, :] = symmetrize(BrCMB[k, :, idx], self.minc)
else:
dat[k, :] = BrCMB[k, :, idx]
th = np.linspace(np.pi / 2., -np.pi / 2., nlat)
if deminc:
phi = np.linspace(-np.pi, np.pi, self.minc * nphi + 1)
else:
phi = np.linspace(-np.pi / self.minc, np.pi / self.minc, nphi)
fig = plt.figure()
ax = fig.add_subplot(111)
vmin = -max(abs(dat.max()), abs(dat.min()))
vmax = -vmin
cs = np.linspace(vmin, vmax, levels)
ax.contourf(phi, self.time, dat, cs, cmap=plt.get_cmap(cm))
ax.set_xlabel('Longitude')
ax.set_ylabel('Time')
w2 = np.fft.fft2(dat)
w2 = abs(w2[1:self.nstep / 2 + 1, 0:self.m_max_cmb + 1])
dw = 2. * np.pi / (self.time[-1] - self.time[0])
omega = dw * np.arange(self.nstep)
omega = omega[1:self.nstep / 2 + 1]
ms = np.arange(self.m_max_cmb + 1)
fig1 = plt.figure()
ax1 = fig1.add_subplot(111)
ax1.contourf(ms, omega, w2, 17, cmap=plt.get_cmap('jet'))
ax1.set_yscale('log')
ax1.set_xlim(0, 13)
ax1.set_xlabel(r'Azimuthal wavenumber')
ax1.set_ylabel(r'Frequency')
def __init__(self,
q_values,
lmax, n_theta, n_phi,
corr = None, cospsi = None,
ref_SphModel = None, auto_only = False,
bark=False):
# q values correponding to correlations, in ascending order!!!
self.q_values = q_values
# can pass a reference Spherical Model, use it for guessing, and comparing
self.ref_SphModel = ref_SphModel
self.lmax = lmax
self.n_q = len(q_values)
self.n_theta = n_theta
self.n_phi = n_phi
self.sh = shtns.sht(lmax)
self.sh.set_grid( n_theta, n_phi )
if corr is None:
self.corr = self.ref_SphModel.corr
self.cospsi = self.ref_SphModel.cospsi
self.auto_only = auto_only
if auto_only:
self.cl = np.zeros_like(self.ref_SphModel.cl)
self.cl[:,range(self.n_q),range(self.n_q)] = self.ref_SphModel.cl[:,range(self.n_q),range(self.n_q)]
else:
self.cl = self.ref_SphModel.cl
else:
if len(corr.shape) == 2:
# that means there are only autocorrelator
self.auto_only = True
if len(cospsi.shape) == 1:
# that means all the q values have the same cospsi, reshape
cospsi = np.array( [cospsi] * self.n_q )
elif len(corr.shape) == 3:
self.auto_only = False
# auto and cross-correlations
if len(cospsi.shape) == 1:
# that means all the q values have the same cospsi, reshape
cospsi = np.array( [[cospsi] * self.n_q] * self.n_q )
else:
print("Error: correlations should be either a 2D or 3D array")
self.corr = corr
self.cospsi = cospsi
try:
assert( self.cospsi.shape[0] == self.n_q)
assert( self.corr.shape[-1] == self.cospsi.shape[-1])
assert( len(self.corr.shape) == len(self.cospsi.shape) )
except:
print ("ERROR: Mismatch in cosines and correlations provided!")
# project leg poly
self._compute_legendre_projection()
if bark:
print ("\
.:##:::.\n \
.:::::/;;\:.\n\
()::::::@::/;;#;|:.\n\
::::##::::|;;##;|::\n\
\':::::::::\;;;/::\'\n\
':::::::::::\n\
|O|O|O|O|O|O\n\
:#:::::::##::.\n\
.:###:::::#:::::.\n\
:::##:::::::::::#:.\n\
::::;:::::::::###::.\n\
\':::;::###::;::#:::::\n\
::::;::#::;::::::::::\n\
:##:;::::::;::::###::: .\n\
.:::::; .:::##::::::::::::::::\n\
::::::; :::::::::::::::::##::\n\
The Phase Retrieving Golden Retriever Puppy!")
def save(self, save_filename):
ff = h5py.File(save_filename,'a')
ff.create_dataset('intensity', data = self.I_guess)
ff.create_dataset('slm',data = self.all_slm_guess)
ff.create_dataset('qvalues', data = self.q_values)
ff.create_dataset('leg_coefs', data = self.cl)
ff.create_dataset('corr', data = self.corr)
ff.create_dataset('deltas', data = self.deltas)
ff.create_dataset('cospsi', data = self.cospsi)
ff.create_dataset('lmax', data = self.lmax)
ff.close()
u2d = u3d[2]
v2d = v3d[2]
hdiv_lm, hrot_lm = ds.oper.hdivrotsh_from_uv(u2d.astype(np.float64),
v2d.astype(np.float64))
uuu, vvv = ds.oper.uv_from_hdivrotsh(hdiv_lm, hrot_lm)
print 'uuu', uuu
print 'u2d', u2d
print 'vvv', vvv
print 'v2d', v2d
print 'shtns'
#does not work :(:(:(:(
sh1 = shtns.sht(127, 127, 1, norm=1)
nla, nlo = sh1.set_grid(nlat=192,
nphi=384,
flags=8708,
polar_opt=1e-08,
nl_order=2)
radius = 6367470.0
l2 = sh1.l * (sh1.l + 1)
hd = ds.oper.create_array_sh()
hr = ds.oper.create_array_sh()
vx1 = -v2d
sh1.spat_to_SHsphtor(vx1, u2d, hd, hr)
def filt(dat_r, dat_t, dat_p, lmax=None, lCut=10, filt_type="high"):
dat_r = float64(dat_r)
dat_t = float64(dat_t)
dat_p = float64(dat_p)
nphi = dat_r.shape[0]
nlat = dat_r.shape[1]
nr = dat_r.shape[2]
# Convert from MagIC to shtns format: nphi,ntheta,nr -> ntheta,nphi,nr
dat_r = transpose(dat_r, (1, 0, 2))
dat_t = transpose(dat_t, (1, 0, 2))
dat_p = transpose(dat_p, (1, 0, 2))
if lmax is None:
lmax = nphi / 3
sh = shtns.sht(lmax=lmax,
mmax=lmax,
norm=shtns.sht_schmidt | shtns.SHT_NO_CS_PHASE)
polar_opt = 1e-10
nlat, nphi = sh.set_grid(nlat=nlat, nphi=nphi, polar_opt=polar_opt)
Q = zeros([nr, sh.nlm], dtype="complex128")
S = zeros([nr, sh.nlm], dtype="complex128")
T = zeros([nr, sh.nlm], dtype="complex128")
Q1 = zeros([nr, sh.nlm], dtype="complex128")
S1 = zeros([nr, sh.nlm], dtype="complex128")
T1 = zeros([nr, sh.nlm], dtype="complex128")
# Forward transform: data -> SH
for i in range(nr):
Q[i, :], S[i, :], T[i, :] = sh.analys(dat_r[..., i], dat_t[..., i],
dat_p[..., i])
# Apply filter
if filt_type == "high":
mask = sh.l >= lCut
elif filt_type == "low":
mask = sh.l <= lCut
else:
print('Filter type must be "high" or "low"')
exit()
Q1[:, mask] = Q[:, mask]
S1[:, mask] = S[:, mask]
T1[:, mask] = T[:, mask]
# Inverse transform: SH -> data
for k in range(nr):
dat_r[..., k], dat_t[...,
k], dat_p[...,
k] = sh.synth(Q1[k, :], S1[k, :],
T1[k, :])
# Convert from shtns to MagIC format: ntheta,nphi,nr -> nphi,ntheta,nr
dat_r = transpose(dat_r, (1, 0, 2))
dat_t = transpose(dat_t, (1, 0, 2))
dat_p = transpose(dat_p, (1, 0, 2))
return dat_r, dat_t, dat_p
def plot_radial_cmplx_wrapper(dir_NM, i_mode, n_lat_grid, i_radius_str = None, show = True, fmt = 'png', transparent = True):
# Define directories.
dir_processed = os.path.join(dir_NM, 'processed')
dir_spectral = os.path.join(dir_processed, 'spectral')
dir_plot = os.path.join(dir_processed, 'plots')
mkdir_if_not_exist(dir_plot)
# Determine if the plot is 'quick' mode or 'full' mode.
if i_radius_str is None:
option = 'quick'
else:
option = 'full'
fig = plt.figure(figsize = (7.0, 5.0))
# Reconstruct the coordinate grid.
n_lon_grid = (2*n_lat_grid) - 1
lon_grid = np.linspace(0.0, 2.0*np.pi, n_lon_grid + 1, endpoint = True)[:-1]
lat_grid = np.linspace(-np.pi/2.0, np.pi/2.0, n_lat_grid, endpoint=True)
# List of radii to sample.
if i_radius_str == 'all':
_, _, _, _, _, header_info = load_vsh_coefficients(dir_NM, i_mode, 0)
n_samples = len(header_info['r_sample'])
i_radius_list = list(range(n_samples))
elif i_radius_str is None:
i_radius_list = [None]
else:
i_radius_list = [int(i_radius_str)]
first_iteration = True
for j_radius in i_radius_list:
# Load the VSH coefficients for the specified mode.
coeffs, header_info, r_sample, i_sample = load_vsh_coefficients(dir_NM, i_mode, j_radius)
Ulm_real, Vlm_real, Wlm_real, Ulm_imag, Vlm_imag, Wlm_imag = coeffs
title = region_int_to_title(i_sample, r_sample, shell_name_path = os.path.join(dir_processed, 'shell_names.txt'))
if first_iteration:
# Infer the maximum l-value used.
#n_coeffs = len(Ulm)
n_coeffs = coeffs.shape[-1]
l_max = (int((np.round(np.sqrt(8*n_coeffs + 1)) - 1))//2) - 1
# Re-construct the shtns calculator.
m_max = l_max
sh_calculator = shtns.sht(l_max, m_max)
grid_type = shtns.sht_reg_fast \
| shtns.SHT_SOUTH_POLE_FIRST \
| shtns.SHT_PHI_CONTIGUOUS
sh_calculator.set_grid(n_lat_grid, n_lon_grid, flags = grid_type)
first_iteration = False
n_c_levels_default = 20
# Expand the VSH into spatial components and plot.
U_real_r, V_real_e, V_real_n, W_real_e, W_real_n = project_from_spherical_harmonics(sh_calculator, Ulm_real, Vlm_real, Wlm_real)
U_imag_r, V_imag_e, V_imag_n, W_imag_e, W_imag_n = project_from_spherical_harmonics(sh_calculator, Ulm_imag, Vlm_imag, Wlm_imag)
U_cplx_r = U_real_r + 1.0j*U_imag_r
#U_abs_r = np.abs(U_cplx_r)
#U_phase_r = np.angle(U_cplx_r)
plot_radial_cplx(lon_grid, lat_grid, U_cplx_r, fig = fig, ax_arr = None, show = False, title = title, n_c_levels = n_c_levels_default)
# Save the plot.
fig_base_name = 'disp_radial_cplx'
if option == 'quick':
name_fig = '{:}_{:}_{:>05d}.{:}'.format(fig_base_name, option, i_mode, fmt)
else:
name_fig = '{:}_{:}_{:>05d}_{:>03d}.{:}'.format(fig_base_name, option, i_mode, j_radius, fmt)
path_fig = os.path.join(dir_plot, name_fig)
print('Saving figure to {:}'.format(path_fig))
save_figure(path_fig, fmt, transparent = transparent)
# Show the plot.
if show:
plt.show()
# Close the plot.
plt.close()
return
def plot_sh_disp_wrapper(dir_NM, i_mode, n_lat_grid, mode_real_or_complex, show = True, fmt = 'pdf', transparent = False, i_radius_str = None, close = True, c_bar_label = 'Default', figsize = None, path_outline = None):
'''
Plots a vector field on the surface of a sphere in terms of the radial, consoidal and toroidal components.
This is a wrapper for plot_sh_disp() which first loads the necessary arrays.
Input:
dir_NM, i_mode, n_lat_grid, fmt
See 'Definitions of variables' in NMPostProcess/process.py.
Output:
None
'''
# Set transparency.
transparent = True
# Define directories.
dir_processed = os.path.join(dir_NM, 'processed')
dir_spectral = os.path.join(dir_processed, 'spectral')
dir_plot = os.path.join(dir_processed, 'plots')
mkdir_if_not_exist(dir_plot)
# Load outline data if specified.
if path_outline is not None:
outline_data = np.loadtxt(path_outline)
else:
outline_data = None
# Determine if the plot is 'quick' mode or 'full' mode.
if i_radius_str is None:
option = 'quick'
else:
option = 'full'
if mode_real_or_complex in ['real', 'complex_imag_only', 'complex_real_only']:
if figsize is None:
figsize = (3.5, 6.0)
two_panel = False
elif mode_real_or_complex == 'complex':
if figsize is None:
figsize = (7.0, 6.0)
two_panel = True
else:
raise ValueError
# Reconstruct the coordinate grid.
n_lon_grid = (2*n_lat_grid) - 1
lon_grid = np.linspace(0.0, 2.0*np.pi, n_lon_grid + 1, endpoint = True)[:-1]
lat_grid = np.linspace(-np.pi/2.0, np.pi/2.0, n_lat_grid, endpoint=True)
if i_radius_str == 'all':
_, header_info, _, _ = load_vsh_coefficients(dir_NM, i_mode, 0)
n_samples = len(header_info['r_sample'])
i_radius_list = list(range(n_samples))
elif i_radius_str is None:
i_radius_list = [None]
else:
i_radius_list = [int(i_radius_str)]
# Identify the eigenvalue list file.
# Load the frequency of the mode.
path_eigenvalues = os.path.join(dir_processed, 'eigenvalue_list.txt')
freq_mHz = read_eigenvalues(path_eigenvalues, i_mode = i_mode)
first_iteration = True
for j_radius in i_radius_list:
fig = plt.figure(figsize = figsize)
# Load the VSH coefficients for the specified mode.
coeffs, header_info, r_sample, i_sample = load_vsh_coefficients(dir_NM, i_mode, j_radius)
if mode_real_or_complex == 'real':
Ulm, Vlm, Wlm = coeffs
else:
Ulm_real, Vlm_real, Wlm_real, Ulm_imag, Vlm_imag, Wlm_imag = coeffs
title_str_mode = 'Mode {:>4d}'.format(i_mode)
title_str_freq = 'freq. {:>7.4f} mHz'.format(freq_mHz)
title_str_sample = region_int_to_title(i_sample, r_sample, shell_name_path = os.path.join(dir_processed, 'shell_names.txt'))
title_str_sample = title_str_sample[0].lower() + title_str_sample[1:]
if two_panel:
title = '{:} ({:}), sampled at {:}'.format(title_str_mode, title_str_freq, title_str_sample)
else:
title = '{:} ({:}),\nsampled at {:}'.format(title_str_mode, title_str_freq, title_str_sample)
if first_iteration:
# Infer the maximum l-value used.
#n_coeffs = len(Ulm)
n_coeffs = coeffs.shape[-1]
l_max = (int((np.round(np.sqrt(8*n_coeffs + 1)) - 1))//2) - 1
# Re-construct the shtns calculator.
m_max = l_max
sh_calculator = shtns.sht(l_max, m_max)
grid_type = shtns.sht_reg_fast \
| shtns.SHT_SOUTH_POLE_FIRST \
| shtns.SHT_PHI_CONTIGUOUS
sh_calculator.set_grid(n_lat_grid, n_lon_grid, flags = grid_type)
first_iteration = False
n_c_levels_default = 20
# Expand the VSH into spatial components and plot.
if mode_real_or_complex == 'real':
U_r, V_e, V_n, W_e, W_n = project_from_spherical_harmonics(sh_calculator, Ulm, Vlm, Wlm)
fig, ax_arr, handles, contour_info = plot_sh_disp_3_comp(lon_grid, lat_grid, U_r, V_e, V_n, W_e, W_n, fig = fig, ax_arr = None, show = False, title = title, n_c_levels = n_c_levels_default)
fields = { 'U' : {'r' : U_r},
'V' : {'e' : V_e, 'n' : V_n},
'W' : {'e' : W_e, 'n' : W_n}}
else:
U_real_r, V_real_e, V_real_n, W_real_e, W_real_n = project_from_spherical_harmonics(sh_calculator, Ulm_real, Vlm_real, Wlm_real)
U_imag_r, V_imag_e, V_imag_n, W_imag_e, W_imag_n = project_from_spherical_harmonics(sh_calculator, Ulm_imag, Vlm_imag, Wlm_imag)
fields = {
'U' : { 'real' : { 'r' : U_real_r},
'imag' : { 'r' : U_imag_r}},
'V' : { 'real' : { 'n' : V_real_n,
'e' : V_real_e},
'imag' : { 'n' : V_imag_n,
'e' : V_imag_e}},
'W' : { 'real' : { 'n' : W_real_n,
'e' : W_real_e},
'imag' : { 'n' : W_imag_n,
'e' : W_imag_e}}}
#U_abs_r = np.sqrt((U_real_r**2.0) + (U_imag_r**2.0))
#V_abs_e = np.sqrt((V_real_e**2.0) + (V_imag_e**2.0))
#V_abs_n = np.sqrt((V_real_n**2.0) + (V_imag_n**2.0))
#W_abs_e = np.sqrt((W_real_e**2.0) + (W_imag_e**2.0))
#W_abs_n = np.sqrt((W_real_n**2.0) + (W_imag_n**2.0))
if c_bar_label == 'Default':
c_bar_label_real = 'Real part'
c_bar_label_imag = 'Imaginary part'
else:
c_bar_label_real = c_bar_label
c_bar_label_imag = c_bar_label
if mode_real_or_complex == 'complex':
#plt.imshow(W_real_e)
#plt.show()
fig, ax_arr_real, handles_real, contour_info_real = plot_sh_disp_3_comp(lon_grid, lat_grid, U_real_r, V_real_e, V_real_n, W_real_e, W_real_n, fig = fig, axes_grid_int = 121, ax_arr = None, show = False, title = title, n_c_levels = n_c_levels_default, c_bar_label = c_bar_label_real, outline_data = outline_data)
fig, ax_arr_imag, handles_imag, contour_info_imag = plot_sh_disp_3_comp(lon_grid, lat_grid, U_imag_r, V_imag_e, V_imag_n, W_imag_e, W_imag_n, fig = fig, axes_grid_int = 122, ax_arr = None, show = False, title = None, n_c_levels = n_c_levels_default, c_bar_label = c_bar_label_imag, outline_data = outline_data)
contour_info = [contour_info_real, contour_info_imag]
ax_arr = [ax_arr_real, ax_arr_imag]
handles = [handles_real, handles_imag]
elif mode_real_or_complex == 'complex_real_only':
fig, ax_arr, handles, contour_info = plot_sh_disp_3_comp(lon_grid, lat_grid, U_real_r, V_real_e, V_real_n, W_real_e, W_real_n, fig = fig, ax_arr = None, show = False, title = title, n_c_levels = n_c_levels_default, c_bar_label = c_bar_label_real, outline_data = outline_data)
elif mode_real_or_complex == 'complex_imag_only':
fig, ax_arr, handles, contour_info = plot_sh_disp_3_comp(lon_grid, lat_grid, U_real_r, V_real_e, V_real_n, W_real_e, W_real_n, fig = fig, ax_arr = None, show = False, title = title, n_c_levels = n_c_levels_default, c_bar_label = c_bar_label_imag, outline_data = outline_data)
else:
raise ValueError
# Save the plot.
if fmt is not None:
real_or_complex_str_dict = {
'real' : 'real',
'complex' : 'cplx',
'complex_real_only' : 'real_only',
'complex_imag_only' : 'imag_only'}
real_or_complex_str = real_or_complex_str_dict[mode_real_or_complex]
if option == 'quick':
name_fig = 'displacement_{:}_{:}_{:>05d}.{:}'.format(option, real_or_complex_str, i_mode, fmt)
else:
name_fig = 'displacement_{:}_{:}_{:>05d}_{:>03d}.{:}'.format(option, real_or_complex_str, i_mode, j_radius, fmt)
#plt.draw()
path_fig = os.path.join(dir_plot, name_fig)
print('Saving figure to {:}'.format(path_fig))
save_figure(path_fig, fmt, transparent = transparent)
# Show the plot.
if show:
plt.show()
# Close the plot.
if close:
plt.close()
return fig, ax_arr, handles, fields, contour_info
def timeLongitude(self, removeMean=True, lat0=0., levels=12, cm='RdYlBu_r',
deminc=True, shtns_lib='shtns'):
"""
Plot the time-longitude diagram of Br (input latitude can be chosen)
.. warning:: the python bindings of `SHTns <https://bitbucket.org/bputigny/shtns-magic>`_ are mandatory to use this plotting function!
:param lat0: value of the latitude
:type lat0: float
:param levels: number of contour levels
:type levels: int
:param cm: name of the colormap
:type cm: str
:param deminc: a logical to indicate if one wants do get rid of the
possible azimuthal symmetry
:type deminc: bool
:param shtns_lib: version of shtns library used: can be either 'shtns'
or 'shtns-magic'
:type shtns_lib: char
:param removeMean: remove the time-averaged part when set to True
:type removeMean: bool
"""
# The python bindings of shtns are mandatory to use this function !!!
import shtns
if removeMean:
blmCut = self.blm-self.blm.mean(axis=0)
else:
blmCut = self.blm
# Define shtns setup
sh = shtns.sht(int(self.l_max_cmb), int(self.m_max_cmb/self.minc),
mres=int(self.minc),
norm=shtns.sht_orthonormal | shtns.SHT_NO_CS_PHASE)
polar_opt_threshold = 1e-10
nlat = max((self.l_max_cmb*(3/2/2)*2),192)
nphi = 2*nlat/self.minc
nlat, nphi = sh.set_grid(nlat, nphi, polar_opt=polar_opt_threshold)
th = np.linspace(np.pi/2., -np.pi/2., nlat)
lat0 *= np.pi/180.
mask = np.where(abs(th-lat0) == abs(th-lat0).min(), 1, 0)
idx = np.nonzero(mask)[0][0]
# Transform data on grid space
BrCMB = np.zeros((self.nstep, nphi, nlat), 'Float64')
if deminc:
dat = np.zeros((self.nstep, self.minc*nphi+1), 'Float64')
else:
dat = np.zeros((self.nstep, nphi), 'Float64')
for k in range(self.nstep):
tmp = sh.synth(blmCut[k, :]*sh.l*(sh.l+1)/self.rcmb**2)
tmp = tmp.T # Longitude, Latitude
if shtns_lib == 'shtns-magic':
BrCMB[k, ...] = rearangeLat(tmp)
else:
BrCMB[k, ...] = tmp
if deminc:
dat[k, :] = symmetrize(BrCMB[k, :, idx], self.minc)
else:
dat[k, :] = BrCMB[k, :, idx]
th = np.linspace(np.pi/2., -np.pi/2., nlat)
if deminc:
phi = np.linspace(-np.pi, np.pi, self.minc*nphi+1)
else:
phi = np.linspace(-np.pi/self.minc, np.pi/self.minc, nphi)
fig = plt.figure()
ax = fig.add_subplot(111)
vmin = -max(abs(dat.max()), abs(dat.min()))
vmax = -vmin
cs = np.linspace(vmin, vmax, levels)
ax.contourf(phi, self.time, dat, cs, cmap=plt.get_cmap(cm))
ax.set_xlabel('Longitude')
ax.set_ylabel('Time')
w2 = np.fft.fft2(dat)
w2 = abs(w2[1:self.nstep/2+1, 0:self.m_max_cmb+1])
dw = 2.*np.pi/(self.time[-1]-self.time[0])
omega = dw*np.arange(self.nstep)
omega = omega[1:self.nstep/2+1]
ms = np.arange(self.m_max_cmb+1)
fig1 = plt.figure()
ax1 = fig1.add_subplot(111)
ax1.contourf(ms, omega, w2, 17, cmap=plt.get_cmap('jet'))
ax1.set_yscale('log')
ax1.set_xlim(0,13)
ax1.set_xlabel(r'Azimuthal wavenumber')
ax1.set_ylabel(r'Frequency')
def movieRad(self, cut=0.5, levels=12, cm='RdYlBu_r', png=False, step=1,
normed=False, dpi=80, bgcolor=None, deminc=True, removeMean=False,
precision='Float64', shtns_lib='shtns', contour=False, mer=False):
"""
Plotting function (it can also write the png files)
.. warning:: the python bindings of `SHTns <https://bitbucket.org/bputigny/shtns-magic>`_ are mandatory to use this plotting function!
:param levels: number of contour levels
:type levels: int
:param cm: name of the colormap
:type cm: str
:param cut: adjust the contour extrema to max(abs(data))*cut
:type cut: float
:param png: save the movie as a series of png files when
set to True
:type png: bool
:param dpi: dot per inch when saving PNGs
:type dpi: int
:param bgcolor: background color of the figure
:type bgcolor: str
:param normed: the colormap is rescaled every timestep when set to True,
otherwise it is calculated from the global extrema
:type normed: bool
:param step: the stepping between two timesteps
:type step: int
:param deminc: a logical to indicate if one wants do get rid of the
possible azimuthal symmetry
:type deminc: bool
:param precision: single or double precision
:type precision: char
:param shtns_lib: version of shtns library used: can be either 'shtns'
or 'shtns-magic'
:type shtns_lib: char
:param contour: also display the solid contour levels when set to True
:type contour: bool
:param mer: display meridians and circles when set to True
:type mer: bool
:param removeMean: remove the time-averaged part when set to True
:type removeMean: bool
"""
# The python bindings of shtns are mandatory to use this function !!!
import shtns
if removeMean:
dataCut = self.wlm-self.wlm.mean(axis=0)
else:
dataCut = self.wlm
# Define shtns setup
sh = shtns.sht(int(self.l_max_r), int(self.m_max_r/self.minc),
mres=int(self.minc),
norm=shtns.sht_orthonormal | shtns.SHT_NO_CS_PHASE)
polar_opt_threshold = 1e-10
nlat = max((self.l_max_r*(3/2/2)*2),192)
nphi = 2*nlat/self.minc
nlat, nphi = sh.set_grid(nlat, nphi, polar_opt=polar_opt_threshold)
# Transform data on grid space
data = np.zeros((self.nstep, nphi, nlat), precision)
for k in range(self.nstep):
tmp = sh.synth(dataCut[k, :]*sh.l*(sh.l+1)/self.radius**2)
tmp = tmp.T # Longitude, Latitude
if shtns_lib == 'shtns-magic':
data[k, ...] = rearangeLat(tmp)
else:
data[k, ...] = tmp
if png:
plt.ioff()
if not os.path.exists('movie'):
os.mkdir('movie')
else:
plt.ion()
if not normed:
vmin = - max(abs(data.max()), abs(data.min()))
vmin = cut * vmin
vmax = -vmin
cs = np.linspace(vmin, vmax, levels)
th = np.linspace(np.pi/2., -np.pi/2., nlat)
if deminc:
phi = np.linspace(-np.pi, np.pi, self.minc*nphi+1)
xxout, yyout = hammer2cart(th, -np.pi)
xxin, yyin = hammer2cart(th, np.pi)
else:
phi = np.linspace(-np.pi/self.minc, np.pi/self.minc, nphi)
xxout, yyout = hammer2cart(th, -np.pi/self.minc)
xxin, yyin = hammer2cart(th, np.pi/self.minc)
ttheta, pphi = np.meshgrid(th, phi)
xx, yy = hammer2cart(ttheta, pphi)
if deminc:
fig = plt.figure(figsize=(8, 4))
else:
fig = plt.figure(figsize=(8/self.minc, 4))
fig.subplots_adjust(top=0.99, right=0.99, bottom=0.01, left=0.01)
ax = fig.add_subplot(111, frameon=False)
if mer:
theta = np.linspace(np.pi/2, -np.pi/2, nlat)
meridians = np.r_[-120, -60, 0, 60, 120]
circles = np.r_[ 60, 30, 0, -30, -60]
for k in range(self.nstep):
if k == 0:
if normed:
vmin = - max(abs(data[k, ...].max()), abs(data[k, ...].min()))
vmin = cut * vmin
vmax = -vmin
cs = np.linspace(vmin, vmax, levels)
if deminc:
dat = symmetrize(data[k, ...], self.minc)
else:
dat = data[k, ...]
im = ax.contourf(xx, yy, dat, cs, cmap=plt.get_cmap(cm), extend='both')
if contour:
ax.contour(xx, yy, dat, cs, linestyles=['-', '-'],
colors=['k', 'k'], linewidths=[0.7, 0.7])
ax.plot(xxout, yyout, 'k-', lw=1.5)
ax.plot(xxin, yyin, 'k-', lw=1.5)
if mer:
for lat0 in circles:
x0, y0 = hammer2cart(lat0*np.pi/180., phi)
ax.plot(x0, y0, 'k:', linewidth=0.7)
for lon0 in meridians:
x0, y0 = hammer2cart(theta, lon0*np.pi/180.)
ax.plot(x0, y0, 'k:', linewidth=0.7)
ax.axis('off')
man = plt.get_current_fig_manager()
man.canvas.draw()
if png:
filename = 'movie/img%05d.png' % k
print('write %s' % filename)
#st = 'echo %i' % ivar + ' > movie/imgmax'
if bgcolor is not None:
fig.savefig(filename, facecolor=bgcolor, dpi=dpi)
else:
fig.savefig(filename, dpi=dpi)
elif k != 0 and k % step == 0:
if not png:
print(k)
plt.cla()
if normed:
vmin = - max(abs(data[k, ...].max()), abs(data[k, ...].min()))
vmin = cut * vmin
vmax = -vmin
cs = np.linspace(vmin, vmax, levels)
if deminc:
dat = symmetrize(data[k, ...], self.minc)
else:
dat = data[k, ...]
im = ax.contourf(xx, yy, dat, cs, cmap=plt.get_cmap(cm), extend='both')
if contour:
ax.contour(xx, yy, dat, cs, colors=['k'],
linestyles=['-', '-'], linewidths=[0.7, 0.7])
ax.plot(xxout, yyout, 'k-', lw=1.5)
ax.plot(xxin, yyin, 'k-', lw=1.5)
if mer:
for lat0 in circles:
x0, y0 = hammer2cart(lat0*np.pi/180., phi)
ax.plot(x0, y0, 'k:', linewidth=0.7)
for lon0 in meridians:
x0, y0 = hammer2cart(theta, lon0*np.pi/180.)
ax.plot(x0, y0, 'k:', linewidth=0.7)
ax.axis('off')
man.canvas.draw()
if png:
filename = 'movie/img%05d.png' % k
print('write %s' % filename)
#st = 'echo %i' % ivar + ' > movie/imgmax'
if bgcolor is not None:
fig.savefig(filename, facecolor=bgcolor, dpi=dpi)
else:
fig.savefig(filename, dpi=dpi)
# The fact that you are presently reading this means that you have had # knowledge of the CeCILL license and that you accept its terms. # ################################### # SHTns Python interface example # ################################### import numpy # numpy for arrays import shtns # shtns module # compiled and installe using ./configure --enable-python && make && make install lmax = 7 # maximum degree of spherical harmonic representation. mmax = 3 # maximum order of spherical harmonic representation. sh = shtns.sht(lmax, mmax) # create sht object with given lmax and mmax (orthonormalized) # sh = shtns.sht(lmax, mmax, mres=2, norm=shtns.sht_schmidt | shtns.SHT_NO_CS_PHASE) # use schmidt semi-normalized harmonics nlat, nphi = sh.set_grid() # build default grid (gauss grid, phi-contiguous) print(sh.nlat, sh.nphi) # displays the latitudinal and longitudinal grid sizes. cost = sh.cos_theta # latitudinal coordinates of the grid as cos(theta) el = sh.l # array of size sh.nlm giving the spherical harmonic degree l for any sh coefficient l2 = el*(el+1) # array l(l+1) that is useful for computing laplacian ### use advanced options to create a regular grid, theta-contiguous, and with south-pole comming first. # nlat = lmax*2 # nphi = mmax*3 # grid_typ = shtns.sht_reg_fast | shtns.SHT_THETA_CONTIGUOUS | shtns.SHT_SOUTH_POLE_FIRST # polar_opt_threshold = 1.0e-10 # sh.set_grid(nlat, nphi, flags=grid_typ, polar_opt=polar_opt_threshold)
# The fact that you are presently reading this means that you have had
# knowledge of the CeCILL license and that you accept its terms.
#
###################################
# SHTns Python interface example #
###################################
import numpy # numpy for arrays
import shtns # shtns module
# compiled and installe using ./configure --enable-python && make && make install
lmax = 7 # maximum degree of spherical harmonic representation.
mmax = 3 # maximum order of spherical harmonic representation.
sh = shtns.sht(
lmax, mmax) # create sht object with given lmax and mmax (orthonormalized)
# sh = shtns.sht(lmax, mmax, mres=2, norm=shtns.sht_schmidt | shtns.SHT_NO_CS_PHASE) # use schmidt semi-normalized harmonics
nlat, nphi = sh.set_grid() # build default grid (gauss grid, phi-contiguous)
print(sh.nlat,
sh.nphi) # displays the latitudinal and longitudinal grid sizes.
cost = sh.cos_theta # latitudinal coordinates of the grid as cos(theta)
el = sh.l # array of size sh.nlm giving the spherical harmonic degree l for any sh coefficient
l2 = el * (el + 1) # array l(l+1) that is useful for computing laplacian
### use advanced options to create a regular grid, theta-contiguous, and with south-pole comming first.
# nlat = lmax*2
# nphi = mmax*3
# grid_typ = shtns.sht_reg_fast | shtns.SHT_THETA_CONTIGUOUS | shtns.SHT_SOUTH_POLE_FIRST
# polar_opt_threshold = 1.0e-10
