Skip to content

Commit

Permalink
Merge pull request #12 from repepo/main
Browse files Browse the repository at this point in the history 
Cleanup plotting
  • Loading branch information
@jrekier
jrekier authored Jul 11, 2024
2 parents 78e2e99 + f1e294b commit e65e125
Show file tree
Hide file tree
Showing 2 changed files with 41 additions and 34 deletions.
13 changes: 6 additions & 7 deletions koreviz/kmode.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
import matplotlib.pyplot as plt
import numpy.polynomial.chebyshev as ch
from glob import glob
#import cmasher as cmr
from .plotlib import add_colorbar,default_cmap,radContour,merContour,eqContour
from .libkoreviz import spec2spat_vec,spec2spat_scal
import sys
Expand Down Expand Up @@ -117,31 +116,31 @@ def __init__(self, vort=False,datadir='.',field='u', solnum=0,
sol = spec2spat_vec(self,ut,par,chx,a,b,vsymm,nthreads,
vort=vort,transform=transform)

self.Q,self.S,self.T = sol[:3]
self.Qlm,self.Slm,self.Plm,self.Tlm = sol[:4]

if transform:
if not vort:
if field == 'u':
self.ur, self.utheta, self.uphi = sol[3:]
self.ur, self.utheta, self.uphi = sol[4:]
if par.anelastic: #Comment the block out if you want momentum/mass flux
self.rho = rap.density(self.r)
for irho in range(self.nr):
self.ur[irho,...] /= self.rho[irho]
self.utheta[irho,...] /= self.rho[irho]
self.uphi[irho,...] /= self.rho[irho]
elif field == 'b':
self.br, self.btheta, self.bphi = sol[3:]
self.br, self.btheta, self.bphi = sol[4:]
else:
if field == 'u':
self.vort_r, self.vort_t, self.vort_p = sol[3:]
self.vort_r, self.vort_t, self.vort_p = sol[4:]
elif field == 'b':
self.jr, self.jtheta, self.jphi = sol[3:]
self.jr, self.jtheta, self.jphi = sol[4:]

del sol

else:
sol = spec2spat_scal(self,ut,par,chx,a,b,vsymm,nthreads,transform=transform)
self.Q = sol[0]
self.Qlm = sol[0]
if transform:
scal = sol[1]
exec('self.'+field + '= scal')
Expand Down
62 changes: 35 additions & 27 deletions koreviz/libkoreviz.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,34 +92,41 @@ def spec2spat_vec(M,ut,par,chx,a,b,vsymm,nthreads,
# (this is the SHTns way with m=mres when m is not zero)

lmax2 = int( M.lmax + 1 - np.sign(M.m) ) # the true max value of l
nlm = ( np.sign(M.m)+1 ) * (lmax2+1) - M.m
Q = np.zeros([M.nr, nlm], dtype=complex)
S = np.zeros([M.nr, nlm], dtype=complex)
T = np.zeros([M.nr, nlm], dtype=complex)
# nlm = ( np.sign(M.m)+1 ) * (lmax2+1) - M.m

if M.m == 0 : #pad with zeros for the l=0 component
ql = np.r_[ np.zeros((1,M.nr)) ,Qlr ]
sl = np.r_[ np.zeros((1,M.nr)) ,Slr ]
tl = np.r_[ np.zeros((1,M.nr)) ,Tlr ]
pl = np.r_[ np.zeros((1,M.nr)) ,Plr ]
else :
ql = Qlr
sl = Slr
tl = Tlr

Q[:, np.sign(M.m)*(lmax2+1):] = ql.T
S[:, np.sign(M.m)*(lmax2+1):] = sl.T
T[:, np.sign(M.m)*(lmax2+1):] = tl.T
M.ell = np.arange(M.m,M.lmax+1)
pl = Plr #Only for output, not needed for transforms


# SHTns init
#norm = shtns.sht_schmidt | shtns.SHT_NO_CS_PHASE
norm = shtns.sht_schmidt
norm = shtns.sht_schmidt | shtns.SHT_NO_CS_PHASE
mmax = int( np.sign(M.m) )
mres = max(1,M.m)
sh = shtns.sht( lmax2, mmax=mmax, mres=mres, norm=norm, nthreads=nthreads )
M.sh = sh

# Create spectral arrays in SHTns style

Q = np.zeros([M.nr, sh.nlm], dtype=complex)
S = np.zeros([M.nr, sh.nlm], dtype=complex)
T = np.zeros([M.nr, sh.nlm], dtype=complex)
P = np.zeros([M.nr, sh.nlm], dtype=complex)

mask = sh.m == M.m

Q[:,mask] = ql.T
S[:,mask] = sl.T
T[:,mask] = tl.T
P[:,mask] = pl.T

if transform:
ntheta, nphi = sh.set_grid( M.ntheta+M.ntheta%2, M.nphi, polar_opt=1e-10)
M.theta = np.arccos(sh.cos_theta)
Expand All @@ -136,9 +143,9 @@ def spec2spat_vec(M,ut,par,chx,a,b,vsymm,nthreads,
for ir in range(M.nr):
ur[ir,...], utheta[ir,...], uphi[ir,...] = sh.synth( Q[ir,:], S[ir,:], T[ir,:])

return Q,S,T,ur,utheta,uphi
return Q,S,P,T,ur,utheta,uphi
else:
return Q,S,T
return Q,S,P,T


def spec2spat_scal(M,ut,par,chx,a,b,vsymm,nthreads,transform=True):
Expand Down Expand Up @@ -170,28 +177,29 @@ def spec2spat_scal(M,ut,par,chx,a,b,vsymm,nthreads,transform=True):
elif M.m == 0:
idp = np.arange( s , lm1, 2)

# populate Plr and Tlr
# populate Plr
Plr[idp,:] = np.matmul( Plj, chx.T)

# (this is the SHTns way with m=mres when m is not zero)

lmax2 = int( M.lmax + 1 - np.sign(M.m) ) # the true max value of l
nlm = ( np.sign(M.m)+1 ) * (lmax2+1) - M.m
Q = np.zeros([M.nr, nlm], dtype=complex)
# nlm = ( np.sign(M.m)+1 ) * (lmax2+1) - M.m

if M.m == 0 : #pad with zeros for the l=0 component
ql = np.r_[ np.zeros((1,M.nr)) ,Plr ]
else :
ql = Plr

Q[:, np.sign(M.m)*(lmax2+1):] = ql.T
M.ell = np.arange(M.m,M.lmax+1)
norm = shtns.sht_schmidt
norm = shtns.sht_schmidt | shtns.SHT_NO_CS_PHASE
mmax = int( np.sign(M.m) )
mres = max(1,M.m)
sh = shtns.sht( lmax2, mmax=mmax, mres=mres, norm=norm, nthreads=nthreads )
M.sh = sh

Q = np.zeros([M.nr, sh.nlm], dtype=complex)
mask = sh.m == M.m
Q[:,mask] = ql.T

if transform:
# SHTns init
#norm = shtns.sht_schmidt | shtns.SHT_NO_CS_PHASE
Expand All @@ -216,8 +224,8 @@ def get_ang_momentum(M,epsilon_cmb):
l = M.sh.l
m = M.sh.m
r = M.r
Slm = M.S[0,:]
Tlm = M.T[0,:]
Slm = M.Slm[0,:]
Tlm = M.Tlm[0,:]

Gamma_tor = np.zeros(M.sh.nlm,dtype=np.complex128)

Expand All @@ -237,7 +245,7 @@ def get_ang_momentum(M,epsilon_cmb):
-clm2[k] * Tlm[M.sh.idx(ell-1,mm)] )
Gamma_tor[k] *= np.conjugate(epsilon_cmb[k])

torq_pollm = np.real( 4*np.pi/(2*l+1) * (Gamma_pol)) # elementwise (array) multiplication
torq_pollm = np.real( 4*np.pi/(2*l+1) * (Gamma_pol)) # elementwise (array) multiplication
torq_torlm = np.real( 4*np.pi/(2*l+1) * (Gamma_tor))
mask = M.sh.m == 0
torq_pollm[~mask] *= 2
Expand All @@ -252,9 +260,9 @@ def get_coriolis_torque(M,epsilon_cmb):
l = M.sh.l
m = M.sh.m
r = M.r
Qlm = M.Q[0,:]
Slm = M.S[0,:]
Tlm = M.T[0,:]
Qlm = M.Qlm[0,:]
Slm = M.Slm[0,:]
Tlm = M.Tlm[0,:]

Gamma_rad = np.zeros(M.sh.nlm,dtype=np.complex128)
Gamma_con = np.zeros(M.sh.nlm,dtype=np.complex128)
Expand All @@ -271,7 +279,7 @@ def get_coriolis_torque(M,epsilon_cmb):
Gamma_rad[k] = M.rcmb * ( clm_rad_p2[k] * Qlm[M.sh.idx(ell+2,mm)]
+clm_rad_0[k] * Qlm[M.sh.idx(ell,mm)] )
elif ell >= M.lmax-1:
Gamma_rad[k] = M.rcmb * ( clm_rad_m2[k] * Qlm[M.sh.idx(ell-2,mm)]
Gamma_rad[k] = M.rcmb * ( clm_rad_m2[k] * Qlm[M.sh.idx(ell-2,mm)]
+clm_rad_0[k] * Qlm[M.sh.idx(ell,mm)] )
else:
Gamma_rad[k] = M.rcmb * ( clm_rad_p2[k] * Qlm[M.sh.idx(ell+2,mm)]
Expand All @@ -290,7 +298,7 @@ def get_coriolis_torque(M,epsilon_cmb):
Gamma_con[k] = M.rcmb * ( clm_con_p2[k] * Slm[M.sh.idx(ell+2,mm)]
+clm_con_0[k] * Slm[M.sh.idx(ell,mm)] )
elif ell >= M.lmax-1:
Gamma_con[k] = M.rcmb * ( clm_con_m2[k] * Slm[M.sh.idx(ell-2,mm)]
Gamma_con[k] = M.rcmb * ( clm_con_m2[k] * Slm[M.sh.idx(ell-2,mm)]
+clm_con_0[k] * Slm[M.sh.idx(ell,mm)] )
else:
Gamma_con[k] = M.rcmb * ( clm_con_p2[k] * Slm[M.sh.idx(ell+2,mm)]
Expand Down

0 comments on commit e65e125

Please sign in to comment.