AtHeartEngineer/kaipy
1
0
Fork 0
mirror of https://github.com/JHUAPL/kaipy.git synced 2026-01-07 21:43:57 -05:00
Code Issues Packages Projects Releases Wiki Activity

Merged in mndf (pull request #12)

Browse Source 
Mndf

Approved-by: Michael Wiltberger
This commit is contained in:
Dong Lin
2025-07-31 14:54:51 +00:00
parent c74f7fdc28 b05090ec2d
commit 7910fd31c3
2 changed files with 118 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

119
kaipy/remix/remix.py
View File

@@ -7,6 +7,8 @@ import h5py
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np
import kaipy.kaiViz as kv
import cmasher as cmr
# Kaipy modules
from kaipy.kdefs import RionE, REarth
@@ -19,6 +21,7 @@ mu0o4pi = 1.e-7 # mu0=4pi*10^-7 => mu0/4pi=10^-7
facMax = 1.5
facCM = cm.RdBu_r
flxCM = cm.inferno
eflxCM = plt.get_cmap('cmr.seaweed')
class remix:
"""
@@ -51,7 +54,7 @@ class remix:
calcFaceAreas(self, x, y)
Calculates the area of each face in a quad mesh.
plot(self, varname, ncontours=16, addlabels={}, gs=None, doInset=False, doCB=True, doCBVert=True, doGTYPE=False, doPP=False)
plot(self, varname, ncontours=16, addlabels={}, gs=None, doInset=False, doCB=True, doCBVert=True, doGTYPE=False, nPP=0, doMNDF=False)
Plots the specified variable.
efield(self, returnDeltas=False, ri=Ri*1e3)
@@ -115,7 +118,10 @@ class remix:
'Deflux': {'min': 0, 'max': 10.},
'Penergy': {'min': 0, 'max': 240},
'Pflux': {'min': 0, 'max': 1.e7},
'Peflux': {'min': 0, 'max': 1.}
'Peflux': {'min': 0, 'max': 1.},
'beta': {'min':0.0, 'max':0.5},
'deltae': {'min':0, 'max':20},
'amtype': {'min':0, 'max':3}
}
def get_data(self, h5file, step):
@@ -171,6 +177,9 @@ class remix:
self.variables['sigmah']['data'] = self.ion['Hall conductance ' + h]
self.variables['energy']['data'] = self.ion['Average energy ' + h]
self.variables['flux']['data'] = self.ion['Number flux ' + h]
self.variables['beta']['data'] = self.ion['RCM beta '+h]
self.variables['deltae']['data'] = self.ion['Mono potential drop '+h]
self.variables['amtype']['data'] = self.ion['Auroral model type '+h]
if 'RCM grid type ' + h in self.ion.keys():
self.variables['gtype']['data'] = self.ion['RCM grid type ' + h]
if 'RCM plasmasphere density ' + h in self.ion.keys():
@@ -201,6 +210,9 @@ class remix:
self.variables['sigmah']['data'] = self.ion['Hall conductance ' + h][:, ::-1]
self.variables['energy']['data'] = self.ion['Average energy ' + h][:, ::-1]
self.variables['flux']['data'] = self.ion['Number flux ' + h][:, ::-1]
self.variables['beta']['data'] = self.ion['RCM beta '+h][:, ::-1]
self.variables['deltae']['data'] = self.ion['Mono potential drop '+h][:, ::-1]
self.variables['amtype']['data'] = self.ion['Auroral model type '+h][:, ::-1]
if 'RCM grid type ' + h in self.ion.keys():
self.variables['gtype']['data'] = self.ion['RCM grid type ' + h][:, ::-1]
if 'RCM plasmasphere density ' + h in self.ion.keys():
@@ -305,7 +317,7 @@ class remix:
# TODO: define and consolidate allowed variable names
def plot(self, varname, ncontours=16, addlabels={}, gs=None, doInset=False, doCB=True, doCBVert=True, doGTYPE=False, doPP=False):
def plot(self, varname, ncontours=16, addlabels={}, gs=None, doInset=False, doCB=True, doCBVert=True, doGTYPE=False, nPP=0, doMNDF=False):
"""
Plot the specified variable on a polar grid.
@@ -318,7 +330,8 @@ class remix:
doCB (bool): Whether to show the colorbar (default: True).
doCBVert (bool): Whether to show the colorbar vertically (default: True).
doGTYPE (bool): Whether to overplot grid type contours (default: False).
doPP (bool): Whether to overplot polar cap boundary contours (default: False).
nPP (real): plasmaspheric number density contour (default: 0 and not shown).
doMNDF (bool): Whether to plot mono/diffuse using gree/blue colorbar (default: False).
Returns:
ax (AxesSubplot): The AxesSubplot object containing the plot.
@@ -412,7 +425,13 @@ class remix:
LW = 0.75
alpha = 1
tOff = np.pi / 2.
ax.contour(tc + tOff, rc, tmp, levels=con_level, colors=con_color, linewidths=LW, alpha=alpha)
contour = ax.contour(tc + tOff, rc, tmp, levels=con_level, colors=con_color, linewidths=LW, alpha=alpha)
if con_name=='npsp':
cl_con = 'red'
ax.clabel(contour, inline=False, fontsize=8, fmt=lambda val: f"Npsph={val:.1f}", colors=cl_con, use_clabeltext=True)
if con_name=='gtype':
cl_con = 'purple'
ax.clabel(contour, inline=False, fontsize=8, fmt=lambda val: f"GIMAG={val:.2f}", colors=cl_con, use_clabeltext=True)
if not self.Initialized:
sys.exit("Variables should be initialized for the specific hemisphere (call init_var) prior to plotting.")
@@ -431,6 +450,10 @@ class remix:
'energy' : r'Energy [keV]',
'flux' : r'Flux [1/cm$^2$s]',
'eflux' : r'Energy flux [erg/cm$^2$s]',
'gtype' : r'IMAG grid type',
'beta' : r'$\beta$',
'deltae' : r'$\Delta$E [kV]',
'amtype' : r'Auroral model type',
'ephi' : r'$E_\phi$ [mV/m]',
'etheta' : r'$E_\theta$ [mV/m]',
'efield' : r'|E| [mV/m]',
@@ -461,7 +484,7 @@ class remix:
upper = self.variables[varname]['data'].max()
# define number format string for min/max labels
if varname !='flux' and varname !='Pflux':
if varname not in ['flux','Dflux','Mflux','Pflux']:
if varname == 'jped':
format_str = '%.2f'
else:
@@ -472,7 +495,10 @@ class remix:
if (varname == 'potential') or (varname == 'current'):
cmap=facCM
elif varname in ['flux','energy','eflux','joule','Mflux','Menergy','Meflux','Dflux','Denergy','Deflux','Pflux','Penergy','Peflux']:
cmap=flxCM
if not doMNDF:
cmap=flxCM
else:
cmap=eflxCM
latlblclr = 'white'
elif (varname == 'velocity'):
cmap=cm.YlOrRd
@@ -480,6 +506,12 @@ class remix:
cmap=None #cm.RdBu_r#None # cm.hsv
elif (varname == 'magpert'):
cmap = cm.YlOrRd
elif (varname == 'beta'):
cmap=cm.inferno
elif (varname == 'gtype'):
cmap=cm.magma_r
elif (varname == 'deltae'):
cmap=cm.viridis
else:
cmap=None # default is used
@@ -518,7 +550,7 @@ class remix:
if varname == 'joule':
self.variables[varname]['data'] = self.joule()*1.e3 # convert to mW/m^2
variable = self.variables[varname]['data']
# variable = self.variables[varname]['data']
fig = plt.gcf()
@@ -528,7 +560,39 @@ class remix:
else:
ax=fig.add_subplot(polar=True)
p=ax.pcolormesh(theta+tOff,r,variable,cmap=cmap,vmin=lower,vmax=upper)
# p=ax.pcolormesh(theta+tOff,r,variable,cmap=cmap,vmin=lower,vmax=upper)
if not doMNDF:
# traditional colorbar limits and maps
variable = self.variables[varname]['data']
p=ax.pcolormesh(theta+tOff,r,variable,cmap=cmap,vmin=lower,vmax=upper)
else:
# for mono/diffuse, use different colorbar limits and maps
efxAmax = 10 # limit in the Mono-diffuse asymmetric colorbar.
numAmax = 1e9
engAmax = 20
idiff = self.variables['amtype']['data']<=2
if varname in ['eflux','Meflux','Deflux','thmeflux']:
variable = self.variables[varname]['data']
variable[idiff] = -abs(variable[idiff])
# use linear scale if below 0.1 mW/m^2
# set color bar limit at efxAmax mW/m^2
vQ = kv.genNorm(efxAmax,doSymLog=False,linP=0.01*efxAmax)
p=ax.pcolormesh(theta+tOff,r,variable,cmap=cmap,norm=vQ)
elif varname in ['flux','Mflux','Dflux','thmflux']:
variable = self.variables[varname]['data']
variable[idiff] = -abs(variable[idiff])
# use linear scale if below 1e7 #/cm^2/s
# set color bar limit at numAmax #/cm^2/s
vQ = kv.genNorm(numAmax,doSymLog=False,linP=0.01*numAmax)
p=ax.pcolormesh(theta+tOff,r,variable,cmap=cmap,norm=vQ)
elif varname in ['energy','Menergy','Denergy']:
variable = self.variables[varname]['data']
variable[idiff] = -abs(variable[idiff])
# use linear scale if below 0.1 keV
# set color bar limit at engAmax keV
vQ = kv.genNorm(engAmax,doSymLog=False,linP=0.01*engAmax)
p=ax.pcolormesh(theta+tOff,r,variable,cmap=cmap,norm=vQ)
if (not doInset):
if (doCB):
@@ -538,8 +602,21 @@ class remix:
cb=plt.colorbar(p,ax=ax,pad=0.1,shrink=0.85,orientation="horizontal")
cb.set_label(cblabels[varname])
ax.text(-75.*np.pi/180.,1.2*r.max(),('min: '+format_str+'\nmax: ' +format_str) %
if not doMNDF:
ax.text(-75.*np.pi/180.,1.2*r.max(),('min: '+format_str+'\nmax: ' +format_str) %
(variable.min() ,variable.max()))
else:
if varname in ['flux','Mflux','Dflux']:
cticks = cb.get_ticks()
cticklab = [str('{:6.1e}'.format(abs(tick))) for tick in cticks]
cb.set_ticklabels(cticklab)
if varname in ['eflux','Meflux','Deflux','energy','Menergy','Denergy']:
cticks = cb.get_ticks()
cticklab = [str('{:.1f}'.format(abs(tick))) for tick in cticks]
cb.set_ticklabels(cticklab)
if varname in ['gtype']:
cb.set_ticks([0,1])
cb.set_ticklabels(['MHD','IMAG'])
lines, labels = plt.thetagrids((0.,90.,180.,270.),hour_labels)
lines, labels = plt.rgrids(circles,lbls,fontsize=8,color=latlblclr)
@@ -550,8 +627,18 @@ class remix:
# ax.axis([0,2*np.pi,0,r.max()],'tight')
ax.axis([0,2*np.pi,0,r.max()])
if (not doInset):
ax.text(-75.*np.pi/180.,1.2*r.max(),('min: '+format_str+'\nmax: ' +format_str) %
(variable.min() ,variable.max()))
if doMNDF and varname in ['eflux','Meflux','Deflux','flux','Mflux','Dflux','energy','Menergy','Denergy']:
mono=variable.copy()
diff=variable.copy()
mono[ idiff]=0.0
diff[~idiff]=0.0
# print('mono min/max:',mono.min(),mono.max())
# print('diff min/max:',diff.min(),diff.max())
ax.text(-77.*np.pi/180.,1.1*r.max(),('Diff max: '+format_str) % (-diff.min()))
ax.text( 75.*np.pi/180.,1.1*r.max(),('Mono max: '+format_str) % ( mono.max()))
else:
ax.text(-75.*np.pi/180.,1.2*r.max(),('min: '+format_str+'\nmax: ' +format_str) % (variable.min() ,variable.max()))
if varname in ['eflux','Meflux','Deflux','Peflux']:
ax.text(75.*np.pi/180.,1.3*r.max(),('GW: '+format_str) % power)
if (varname == 'current'):
@@ -561,9 +648,11 @@ class remix:
if varname=='current':
potential_overplot(doInset)
if doGTYPE and varname=='eflux':
boundary_overplot('gtype',[0.01,0.99],'green',doInset)
if doPP and varname=='eflux':
boundary_overplot('npsp',[10],'cyan',doInset)
cl_con = 'purple'
boundary_overplot('gtype',[0.5],cl_con,doInset)
if nPP>0.0 and varname=='eflux':
cl_con = 'red'
boundary_overplot('npsp',[nPP],cl_con,doInset)
return ax
# mpl.rcParams['contour.negative_linestyle'] = 'solid'

22
kaipy/scripts/quicklook/mixpic.py
View File

@@ -122,8 +122,12 @@ def create_command_line_parser():
help="Show RCM grid type in the eflx plot (default: %(default)s)"
)
parser.add_argument(
'-PP', action='store_true', default=False,
help="Show plasmapause (10/cc) in the eflx/nflx plot (default: %(default)s)"
'--nPP', type=float, metavar="nPP", default=0,
help="Plasmasphere density contour to show (default: %(default)/cc not shown)"
)
parser.add_argument(
'-MNDF', action='store_true', default=False,
help="Show mono-diffuse eflx/nflx in linear/log green/blue colormap (default: %(default)s)"
)
parser.add_argument(
'-vid', action='store_true', default=False,
@@ -156,7 +160,8 @@ def makePlot(i, remixFile, nStp, args, varDict):
spacecraft = args.spacecraft
verbose = args.verbose
do_GTYPE = args.GTYPE
do_PP = args.PP
nPP = args.nPP
do_MNDF = args.MNDF
do_vid = args.vid
do_overwrite = args.overwrite
do_hash = not args.nohash
@@ -233,12 +238,12 @@ def makePlot(i, remixFile, nStp, args, varDict):
axs[0] = ion.plot('current', gs=gs[0, 0])
axs[1] = ion.plot('sigmap', gs=gs[0, 1])
axs[2] = ion.plot('sigmah', gs=gs[0, 2])
axs[3] = ion.plot('joule', gs=gs[1, 0])
axs[4] = ion.plot('energy', gs=gs[1, 1])
if do_nflux:
axs[5] = ion.plot('flux', gs=gs[1, 2],doGTYPE=do_GTYPE,doPP=do_PP)
axs[3] = ion.plot('flux', gs=gs[1, 0],doGTYPE=do_GTYPE,nPP=nPP,doMNDF=do_MNDF)
else:
axs[5] = ion.plot('eflux', gs=gs[1, 2],doGTYPE=do_GTYPE,doPP=do_PP)
axs[3] = ion.plot('joule', gs=gs[1, 0])
axs[4] = ion.plot('energy', gs=gs[1, 1],doGTYPE=do_GTYPE,nPP=nPP,doMNDF=do_MNDF)
axs[5] = ion.plot('eflux', gs=gs[1, 2],doGTYPE=do_GTYPE,nPP=nPP,doMNDF=do_MNDF)
# If requested, plot the magnetic footprints for the specified
# spacecraft.
@@ -314,7 +319,8 @@ def main():
spacecraft = args.spacecraft
verbose = args.verbose
do_GTYPE = args.GTYPE
do_PP = args.PP
nPP = args.nPP
do_MNDF = args.MNDF
do_vid = args.vid
do_overwrite = args.overwrite
do_hash = not args.nohash