Skip to content

Commit

Permalink
Some vis plotting fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
Jeff Candy committed May 6, 2024
1 parent c95ba13 commit a7d1e34
Show file tree
Hide file tree
Showing 4 changed files with 75 additions and 83 deletions.
106 changes: 53 additions & 53 deletions f2py/pygacode/cgyro/data_plot.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,7 @@ def plot_ky_freq(self,xin):
ax.set_xlabel(self.kystr)
ax.set_ylabel(self.fstr[0])

y1 = self.fnorm[0,:,-1]
y1 = self.fnorm[0,:,-1]
ax.plot(ky,y1,color='blue')
ax.plot(ky,y1,"o",color='k')
if len(ky) > 1:
Expand Down Expand Up @@ -186,7 +186,7 @@ def plot_geo(self,xin):
ax.plot(theta,y,'o',color='k',markersize=2)
ax.set_xlim([-1,1])

# captheta
# captheta
if len(self.geo[0,:]) > 12:
y = self.geo[:,12]
ax = fig.add_subplot(gs[0,3])
Expand Down Expand Up @@ -221,14 +221,14 @@ def plot_geo(self,xin):

x = rmaj+r*np.cos(arg)
y = zmaj+k*r*np.sin(t)

ax.plot(x,y,'k')

fig.tight_layout(pad=0.3)

return


def plot_ball(self,xin):

itime = xin['itime']
Expand Down Expand Up @@ -275,14 +275,14 @@ def plot_ball(self,xin):

ax.plot(x,y1,'-o',color='black',markersize=2,label=r'$\mathrm{Re}$')
ax.plot(x,y2,'-o',color='red',markersize=2,label=r'$\mathrm{Im}$')

ax.legend()

fig.tight_layout(pad=0.3)

return 'ang Re(f) Im(f)',x,y1,y2


def plot_ky_phi(self,xin):

# Plot fields versus time for each ky
Expand All @@ -295,7 +295,7 @@ def plot_ky_phi(self,xin):
moment = xin['moment']
spec = xin['spec']

t = self.getnorm(xin['norm'])
t = self.getnorm(xin['norm'])

if xin['fig'] is None:
fig = plt.figure(MYDIR,figsize=(xin['lx'],xin['ly']))
Expand Down Expand Up @@ -355,14 +355,14 @@ def plot_phi(self,xin):
if self.n_n < 2:
print('ERROR: (plot_phi) This plot not suitable for a single toroidal mode.')
return [None,]*4
t = self.getnorm(xin['norm'])

t = self.getnorm(xin['norm'])

if xin['fig'] is None:
fig = plt.figure(MYDIR,figsize=(xin['lx'],xin['ly']))

self.getbigfield()

f,ft = self.kxky_select(theta,field,moment,spec,gbnorm=True)

#======================================
Expand All @@ -377,15 +377,15 @@ def plot_phi(self,xin):

# Get index for average window
imin,imax=time_index(t,w)

# n=0 intensity (gyroBohm)
y0 = np.sum(abs(f[:,0,:]),axis=0)
s0 = np.sum(abs(f[:,0,:])**2,axis=0)
s0 = np.sum(abs(f[:,0,:])**2,axis=0)

# finite-n intensity (gyroBohm)
yn = np.sum(abs(f[:,1:,:]),axis=(0,1))
sn = np.sum(abs(f[:,1:,:])**2,axis=(0,1))

s = np.ones(imax-imin+1)
rstr = self.rhoi
if absnorm == 0:
Expand Down Expand Up @@ -441,13 +441,13 @@ def plot_flux(self,xin):
norm = xin['norm']
theta = xin['theta']


if xin['fig'] is None and ftype != 'nox':
fig = plt.figure(MYDIR,figsize=(xin['lx'],xin['ly']))

if moment == 'phi':
moment = 'e'

t = self.getnorm(xin['norm'])
usec = self.getflux(cflux)
ns = self.n_species
Expand Down Expand Up @@ -561,14 +561,14 @@ def plot_ky_flux(self,xin):
theta = xin['theta']

t = self.getnorm(xin['norm'])

if self.n_n == 1:
print('ERROR: (plot_ky_flux) Plot not available with a single mode.')
return

if moment == 'phi':
moment = 'e'

ns = self.n_species
nn = self.n_n

Expand All @@ -582,7 +582,7 @@ def plot_ky_flux(self,xin):
fig = plt.figure(MYDIR,figsize=(xin['ly']*ncol,xin['ly']*nrow))

usec = self.getflux(cflux)

ky = self.kynorm

if fc == 0:
Expand Down Expand Up @@ -625,7 +625,7 @@ def plot_ky_flux(self,xin):
y[1,:,:] = np.sum(abs(f[:,:,:]),axis=0)
else:
raise ValueError('(plot_ky_flux.py) Invalid moment.')

color = ['magenta','k','blue','cyan','red','green']

imin,imax=time_index(t,w)
Expand All @@ -635,7 +635,7 @@ def plot_ky_flux(self,xin):
ky = -ky

dk = ky[1]-ky[0]

ave = np.zeros((self.n_n,ns))
for ispec in range(ns):
ave[:,ispec] = time_average(y[ispec,:,:],t,imin,imax)
Expand All @@ -645,7 +645,7 @@ def plot_ky_flux(self,xin):
u = specmap(self.mass[ispec],self.z[ispec])
if not ftype == 'nox':
ax = fig.add_subplot(nrow,ncol,ispec+1)

ax.set_xlabel(self.kystr)
ax.set_ylabel(r'$'+mtag+'_\mathrm{'+u+'}/'+mtag+self.gbnorm+'$',color='k')

Expand All @@ -663,12 +663,12 @@ def plot_ky_flux(self,xin):
ax.set_yscale('log')
ax.plot(ky[1:],ave[1:,ispec],'o',color='k',markersize=2,alpha=0.6)
ax.plot(ky[1:],ave[1:,ispec],color=color[ispec])

if ymax != 'auto':
ax.set_ylim(top=float(ymax))
if ymin != 'auto':
ax.set_ylim(bottom=float(ymin))

# Maximum
j = np.argmax(ave[:,ispec])
if j < len(ky)-1 and j > 0:
Expand All @@ -679,8 +679,8 @@ def plot_ky_flux(self,xin):

if not ftype == 'nox':
fig.tight_layout(pad=0.4)


def plot_rcorr_phi(self,xin):

# Should be checked by Chris after updated form of kxky_select to n_radial-1 points
Expand Down Expand Up @@ -898,8 +898,8 @@ def plot_shift(self,xin):

if xin['fig'] is None:
fig = plt.figure(MYDIR,figsize=(xin['lx'],xin['ly']))
t = self.getnorm(xin['norm'])

t = self.getnorm(xin['norm'])
self.getbigfield()

# Get index for average window
Expand Down Expand Up @@ -1073,16 +1073,16 @@ def plot_xflux(self,xin):

#---------------------------------
# Global flux versus x
g = np.zeros(na) ; g[:] = self.lky_xr[0,ispec]
g = np.ones(na)*self.lky_xr[0,ispec]
for l in range(1,ng):
g[:] = g[:]+2*(np.cos(l*a)*self.lky_xr[l,ispec]-np.sin(l*a)*self.lky_xi[l,ispec])
ax.plot(a/(2*np.pi),g,color=color[ispec])
g = np.zeros(na) ; g[:] = self.lky_xr[0,ispec]
g = np.ones(na)*self.lky_xr[0,ispec]
for l in range(1,ng):
g[:] = g[:]+2*(np.cos(l*ah)*self.lky_xr[l,ispec]-np.sin(l*ah)*self.lky_xi[l,ispec])
ax.plot(a/(2*np.pi),g,color=color[ispec],linestyle='--')
#---------------------------------

#---------------------------------
# Flux partial average over "positive" [-e,e] interval
g0 = self.lky_flux_ave[ispec,0]
Expand All @@ -1096,7 +1096,7 @@ def plot_xflux(self,xin):
ax.plot([0.5-e,0.5],[g1,g1],'o--',color=color[ispec],alpha=0.2,linewidth=3)
ax.plot([-0.5,-0.5+e],[g1,g1],'o--',color=color[ispec],alpha=0.2,linewidth=3)
#---------------------------------

#---------------------------------
# Flux domain average
ga = self.lky_xr[0,ispec]
Expand Down Expand Up @@ -1128,7 +1128,7 @@ def plot_xflux(self,xin):

fig.tight_layout(pad=0.3)


def plot_kxky_phi(self,xin):

norm = xin['norm']
Expand All @@ -1142,7 +1142,7 @@ def plot_kxky_phi(self,xin):
y0 = max(abs(self.ky))

xl = -x0 ; xr = x0

if xin['kxmin'] != 'auto':
xl = float(xin['kxmin'])
if xin['kxmax'] != 'auto':
Expand All @@ -1152,7 +1152,7 @@ def plot_kxky_phi(self,xin):

if xin['fig'] is None:
fig = plt.figure(MYDIR,figsize=(xin['lx'],asp*xin['lx']))

self.getbigfield()
t = self.getnorm(norm)

Expand All @@ -1170,7 +1170,7 @@ def plot_kxky_phi(self,xin):

imin,imax=time_index(t,w)
y = np.sum(abs(f[:,:,imin:imax+1]),axis=2)

# Fix (0,0)
i0 = nx//2-1
y[i0,0] = y[i0+1,0]
Expand Down Expand Up @@ -1211,12 +1211,12 @@ def plot_kx_phi(self,xin):
fig = plt.figure(MYDIR,figsize=(xin['lx'],xin['ly']))

self.getbigfield()

kx = self.kxnorm
nx = self.n_radial-1
nx = self.n_radial-1

imin,imax=time_index(t,w)

dk = kx[1]-kx[0]
x0 = kx[-1]+dk

Expand All @@ -1241,7 +1241,7 @@ def plot_kx_phi(self,xin):
ax.step(kx+dk/2,ave,color='m')
else:
ax.plot(kx,ave,color='m')

else:
nvec = str2list(nstr)
print('INFO: (plot_kx_phi) n = '+str(nvec))
Expand All @@ -1259,7 +1259,7 @@ def plot_kx_phi(self,xin):
ax.plot(kx[nx//2+1:],ave[nx//2+1:],color='k')
else:
ax.plot(kx,ave,label=num)

if self.n_n > 16:
ax.legend(loc=4, ncol=5, prop={'size':12})
else:
Expand Down Expand Up @@ -1343,7 +1343,7 @@ def plot_kx_shift(self,xin):
for k in range(2):
if k == 1:
phi_T = np.roll(phi_T,nx//2,axis=0)

fint = phi_T[nx//4:3*nx//4,:]
phi = np.fft.fftshift(np.fft.fft(fint,axis=0))

Expand Down Expand Up @@ -1404,7 +1404,7 @@ def plot_poly_phi(self,xin):

color = ['m','k','b','c']
xlabel=r'$k$'

ax.set_title(r'$\mathrm{Average~fluctuation~intensity~(Legendre)} \quad $'+mwin)
ax.set_xlabel(xlabel)

Expand All @@ -1427,8 +1427,8 @@ def plot_poly_phi(self,xin):
mat1[k,:] = sp.spherical_jn(k,z)
mat2[k,:] = mat1[k,:]*(-1)**pvec[:]

ai = 1j**kvec
ak = 2*kvec+1
ai = 1j**kvec
ak = 2*kvec+1

if nstr == 'null':
nvec = np.arange(nn)
Expand All @@ -1437,7 +1437,7 @@ def plot_poly_phi(self,xin):
nvec = str2list(nstr)
print('INFO: (plot_kx_phi) n = '+str(nvec))
nsum = False

for n in nvec:
for jt in np.arange(imin,imax+1):

Expand Down Expand Up @@ -1484,12 +1484,12 @@ def plot_poly_phi(self,xin):
ax.set_xlim([0,nk])
ax.set_xlim([0,nk])
ax.set_yscale('log')

if ymin != 'auto':
ax.set_ylim(bottom=float(ymin))
if ymax != 'auto':
ax.set_ylim(top=float(ymax))

fig.tight_layout(pad=0.3)

return
Expand Down Expand Up @@ -1687,7 +1687,7 @@ def plot_hbcut(self,xin):
hpb = np.array(func[1,n0,spec,:,:,itime])
p0 = np.zeros(self.n_energy)
p1 = np.zeros(self.n_energy)

ax.plot(np.sqrt(self.energy),hpa[ix,:],'-o',color='black',markersize=2)
ax.plot(np.sqrt(self.energy),hpb[ix,:],'-o',color='blue',markersize=2)
#========================================================
Expand Down Expand Up @@ -1721,7 +1721,7 @@ def plot_hball(self,xin):

if xin['fig'] is None:
fig = plt.figure(MYDIR,figsize=(xin['lx'],xin['ly']))

if itime > self.n_time-1:
itime = self.n_time-1

Expand All @@ -1741,7 +1741,7 @@ def plot_hball(self,xin):
y = np.array(self.hb[:,:,spec,:,ie,itime])
except:
raise ValueError('(plot_hball.py) Need to run with H_PRINT_FLAG=1')

xp,wp = np.polynomial.legendre.leggauss(self.n_xi)
c = np.zeros(self.n_xi)
alr = np.zeros([len(x)])
Expand All @@ -1753,7 +1753,7 @@ def plot_hball(self,xin):
for j in range(len(x)):
alr[j] = (l+0.5)*np.sum(pl[:]*y[0,j,:])
ali[j] = (l+0.5)*np.sum(pl[:]*y[1,j,:])

ax.plot(x,alr,'-',color=cvec[l%6],label=r'$\ell='+str(l)+'$')
ax.plot(x,ali,'--',color=cvec[l%6])

Expand Down
Loading

0 comments on commit a7d1e34

Please sign in to comment.