diff --git a/kaipy/cdaweb_utils.py b/kaipy/cdaweb_utils.py index a42823e..b60e393 100644 --- a/kaipy/cdaweb_utils.py +++ b/kaipy/cdaweb_utils.py @@ -434,7 +434,7 @@ def fetch_helio_spacecraft_HGS_trajectory(spacecraft, t_start, t_end, mjdc): """ # Read the CDAWeb spacecraft database. spacecraft_data_file = os.path.join( - os.environ["KAIJUHOME"], "kaipy", "satcomp", "sc_helio.json" + os.environ["KAIPYHOME"], "kaipy", "satcomp", "sc_helio.json" ) sc_info = scutils.getScIds(spacecraft_data_file=spacecraft_data_file) @@ -508,7 +508,7 @@ def fetch_helio_spacecraft_trajectory(sc_id, t_start, t_end): """ # Read the CDAWeb spacecraft database. sc_metadata_path = os.path.join( - os.environ["KAIJUHOME"], "kaipy", "satcomp", "sc_helio.json" + os.environ["KAIPYHOME"], "kaipy", "satcomp", "sc_helio.json" ) sc_metadata = scutils.getScIds(spacecraft_data_file=sc_metadata_path) @@ -546,7 +546,7 @@ def ingest_helio_spacecraft_trajectory(sc_id, sc_data, MJDc): """ # Read the CDAWeb spacecraft database. sc_metadata_path = os.path.join( - os.environ["KAIJUHOME"], "kaipy", "satcomp", "sc_helio.json" + os.environ["KAIPYHOME"], "kaipy", "satcomp", "sc_helio.json" ) sc_metadata = scutils.getScIds(spacecraft_data_file=sc_metadata_path) diff --git a/kaipy/gamera/gampp.py b/kaipy/gamera/gampp.py index 65c35a7..5eaf52f 100644 --- a/kaipy/gamera/gampp.py +++ b/kaipy/gamera/gampp.py @@ -221,7 +221,7 @@ class GameraPipe(object): None """ - uID = kh5.PullAtt(f0, "UnitsID") + uID = kh5.PullAtt(f0, "UnitsID",a0="CODE") #Setting default #with h5py.File(f0, 'r') as hf: # uID = hf.attrs.get("UnitsID", "CODE") if not isinstance(uID, str): diff --git a/kaipy/gamhelio/ConfigScripts/startupOH.config b/kaipy/gamhelio/ConfigScripts/startupOH.config deleted file mode 100644 index fd8d36e..0000000 --- a/kaipy/gamhelio/ConfigScripts/startupOH.config +++ /dev/null @@ -1,33 +0,0 @@ -#Modify if needed paths to a grid file, output file and WSA fits file - -[Gamera] -gameraGridFile = heliogrid.h5 -GridDir = /glade/u/home/elenap/gameraHelio/OHelio/grid -gameraIbcFile = innerbc.h5 -IbcDir = /glade/u/home/elenap/gameraHelio/OHelio/ibc - -[Grid] -tMin = 0.1 -tMax = 0.9 -Rin = 1. -Rout = 10. -Ni = 256 -Nj = 128 -Nk = 256 - -[WSA] -wsafile = /glade/u/home/elenap/gameraHelio/OHelio/1AUfiles/CR2095/wsa.h5 -density_temperature_infile = no -gauss_smooth_width = 0 ; 8 -plots = yes -normalized = yes - -[Constants] -gamma = 1.5 ;1.05 -NO2 = 4 - -[Normalization] ; should be normalization from inner helio -Xscale = 1.496e8 !in km -B0 = 1.e-3 ; 5 nT = 5.e-5 Gs -n0 = 200. ; cm-3 - diff --git a/kaipy/gamhelio/ConfigScripts/startupTD.config b/kaipy/gamhelio/ConfigScripts/startupTD.config deleted file mode 100644 index 1abfc52..0000000 --- a/kaipy/gamhelio/ConfigScripts/startupTD.config +++ /dev/null @@ -1,41 +0,0 @@ -[Grid] -gameraGridFile = heliogrid.h5 -GridDir = /glade/u/home/elenap/gameraHelio/grid -gameraIbcFile = innerbcTD.h5 -IbcDir = /glade/u/home/elenap/gameraHelio/ibcTD -tMin = 0.1 -tMax = 0.9 -Rin = 21.5 -Rout = 221.5 -Ni = 256 -Nj = 128 -Nk = 256 - -[Output] -Prefix = 2012 - -[Normalization] -Xscale = 6.955e10 - -[Magnetic polarity] -Monopole : no - -[Constants] -gamma = 1.5 -NO2 = 4 -Tsolar = 25.38 - -[ADAPT] -adaptdir = /glade/u/home/elenap/gameraHelio/WSAfiles/CR2210 -wild_card = vel*2018*R007*.fits -#cadence in days -- make sure to retain units throughout (not Tsolar above) -cadence = 1 -density_temperature_infile = no -gauss_smooth_width = 0 -plots = yes -normalized = no - - -[DUMPS] -init = no -BC = yes diff --git a/kaipy/gamhelio/helioViz.py b/kaipy/gamhelio/helioViz.py index 2f37883..986e905 100644 --- a/kaipy/gamhelio/helioViz.py +++ b/kaipy/gamhelio/helioViz.py @@ -795,9 +795,9 @@ def PlotMerBrNorm( shading="auto") # Plot the heliospheric current sheet. - Ax.contour(np.sqrt(xr_c**2 + yr_c**2), zr_c, Br_r, [0.], + Ax.contour(np.sqrt(xr_c**2 + yr_c**2), zr_c, Br_l, [0.], colors='black') - Ax.contour(-np.sqrt(xl_c**2 + yl_c**2), zl_c, Br_l, [0.], + Ax.contour(-np.sqrt(xl_c**2 + yl_c**2), zl_c, Br_r, [0.], colors='black') else: diff --git a/kaipy/gamhelio/lib/ace.py b/kaipy/gamhelio/lib/ace.py deleted file mode 100644 index 678c0fe..0000000 --- a/kaipy/gamhelio/lib/ace.py +++ /dev/null @@ -1,83 +0,0 @@ -import numpy as np -import datetime - -def cdaweb_mfi(ACE_mf_file): - # Mag. field - ACE_mf_time=[] - ACE_bx=[] - ACE_by=[] - ACE_bz=[] - - i=1 - for line in open(ACE_mf_file,'r'): - i+=1 - if 'dd-mm-yyyy' in line: break - start = i - - i = 0 - for line in open(ACE_mf_file,'r'): - i+=1 - if i 1.e20] = np.nan - ACE_by[np.abs(np.array(ACE_by)) > 1.e20] = np.nan - ACE_bz[np.abs(np.array(ACE_bz)) > 1.e20] = np.nan - - return (ACE_mf_time,ACE_bx,ACE_by,ACE_bz) - -def cdaweb_swe(ACE_plasma_file): - # Plasma - ACE_plasma_time=[] - ACE_v=[] - ACE_n=[] - - i=1 - for line in open(ACE_plasma_file,'r'): - i+=1 - if 'dd-mm-yyyy' in line: break - start = i - - i = 0 - for line in open(ACE_plasma_file,'r'): - i+=1 - if i 1.e20] = np.nan - ACE_n[np.abs(np.array(ACE_n)) > 1.e20] = np.nan - - - return (ACE_plasma_time,ACE_v,ACE_n) - diff --git a/kaipy/gamhelio/lib/cspice.py b/kaipy/gamhelio/lib/cspice.py deleted file mode 100644 index 22899f0..0000000 --- a/kaipy/gamhelio/lib/cspice.py +++ /dev/null @@ -1,29 +0,0 @@ -import numpy as np -import datetime - -def cspice(fname): - dati=[] - x=[] - y=[] - z=[] - r=[] - phi=[] - theta=[] - - for line in open(fname,'r'): - fields = line.strip().split() - year = int(fields[0].strip().split('-')[0]) - doy = int(fields[0].strip().split('-')[1]) - hr = int(fields[1].strip().split(':')[0]) - mn = int(fields[1].strip().split(':')[1]) - dati.append(datetime.datetime(year,1,1,hr,mn)+datetime.timedelta(days=doy-1)) - - x.append(float(fields[2])) - y.append(float(fields[3])) - z.append(float(fields[4])) - - r.append(float(fields[5])) - theta.append(float(fields[6])) - phi.append(float(fields[7])) - - return(dati,x,y,z,r,theta,phi) diff --git a/kaipy/gamhelio/lib/lfmhlib.py b/kaipy/gamhelio/lib/lfmhlib.py deleted file mode 100644 index ded933a..0000000 --- a/kaipy/gamhelio/lib/lfmhlib.py +++ /dev/null @@ -1,293 +0,0 @@ -from pyhdf.SD import SD, SDC -from numpy import sqrt,arccos,arctan2,cos,sin,pi,roll -import time - -def read(filename): - """ - Returns R,theta,phi,Rc,thetac,phic,br,btheta,bphi,vr,vtheta,vphi,rho,cs - """ - - hdffile = SD(filename,SDC.READ) - - x=hdffile.select('X_grid').get() - y=hdffile.select('Y_grid').get() - z=hdffile.select('Z_grid').get() - bx=hdffile.select('bx_').get()[:-1,:-1,:-1] - by=hdffile.select('by_').get()[:-1,:-1,:-1] - bz=hdffile.select('bz_').get()[:-1,:-1,:-1] - vx=hdffile.select('vx_').get()[:-1,:-1,:-1] - vy=hdffile.select('vy_').get()[:-1,:-1,:-1] - vz=hdffile.select('vz_').get()[:-1,:-1,:-1] - rho=hdffile.select('rho_').get()[:-1,:-1,:-1] - cs=hdffile.select('c_').get()[:-1,:-1,:-1] - - t=hdffile.time - hdffile.end() - - # =========== Cell centers ============== - xc=0.125*( x[:-1,:-1,:-1]+x[1:,:-1,:-1]+x[:-1,1:,:-1]+x[:-1,:-1,1:]+ - x[1:,1:,:-1]+x[1:,:-1,1:]+x[:-1,1:,1:]+x[1:,1:,1:] ) - yc=0.125*( y[:-1,:-1,:-1]+y[1:,:-1,:-1]+y[:-1,1:,:-1]+y[:-1,:-1,1:]+ - y[1:,1:,:-1]+y[1:,:-1,1:]+y[:-1,1:,1:]+y[1:,1:,1:] ) - zc=0.125*( z[:-1,:-1,:-1]+z[1:,:-1,:-1]+z[:-1,1:,:-1]+z[:-1,:-1,1:]+ - z[1:,1:,:-1]+z[1:,:-1,1:]+z[:-1,1:,1:]+z[1:,1:,1:] ) - # ======================================= - - R=sqrt(x**2+y**2+z**2) - theta=arccos(z/R) - phi=arctan2(y,x) - phi[phi<0]+=2*pi - - Rc=sqrt(xc**2+yc**2+zc**2) - thetac=arccos(zc/Rc) - phic=arctan2(yc,xc) - phic[phic<0]+=2*pi - - - br = bx*cos(phic)*sin(thetac) + by*sin(phic)*sin(thetac) + bz*cos(thetac) - btheta = bx*cos(phic)*cos(thetac) + by*sin(phic)*cos(thetac) - bz*sin(thetac) - bphi =-bx*sin(phic) + by*cos(phic) - - vr = vx*cos(phic)*sin(thetac) + vy*sin(phic)*sin(thetac) + vz*cos(thetac) - vtheta = vx*cos(phic)*cos(thetac) + vy*sin(phic)*cos(thetac) - vz*sin(thetac) - vphi =-vx*sin(phic) + vy*cos(phic) - - return(t,R,theta,phi,Rc,thetac,phic,br,btheta,bphi,vr,vtheta,vphi,rho,cs) - -def r_theta_phi_uniform(filename): - """ - Return R, theta, phi 1-d arrays, assuming uniform grid - """ - - hdffile = SD(filename,SDC.READ) - - # note, minimal data are read from file - - ni=hdffile.select('X_grid').ni-1 - nj=hdffile.select('X_grid').nj-1 - nk=hdffile.select('X_grid').nk-1 - - - # first get R; in principle, could just take z along the axis - # but are allowing for the possibility that the axis is cut out - x=hdffile.select('X_grid').get(start=(0,0,0),count=(1,1,ni+1)).squeeze() - y=hdffile.select('Y_grid').get(start=(0,0,0),count=(1,1,ni+1)).squeeze() - z=hdffile.select('Z_grid').get(start=(0,0,0),count=(1,1,ni+1)).squeeze() - R = sqrt(x**2+y**2+z**2) - - z=hdffile.select('Z_grid').get(start=(0,0,0),count=(1,nj+1,1)).squeeze() - theta = arccos(z/R[0]) - - x=hdffile.select('X_grid').get(start=(0,nj/2,0),count=(nk+1,1,1)).squeeze() - y=hdffile.select('Y_grid').get(start=(0,nj/2,0),count=(nk+1,1,1)).squeeze() - phi=arctan2(y,x) - phi[phi<0]+=2*pi - phi[-1]=phi[0]+2*pi - - hdffile.end() - - R/=6.96e10 # FIX ME HARD CODED UNITS - Rc = 0.5*(R[1:]+R[:-1]) - thetac = 0.5*(theta[1:]+theta[:-1]) - phic = 0.5*(phi[1:]+phi[:-1]) - - return R,theta,phi,Rc,thetac,phic - -def read_var(filename,var_name): - hdffile = SD(filename,SDC.READ) - if var_name not in ['br','btheta','bphi','vr','vtheta','vphi']: - var=hdffile.select(var_name+'_').get()[:-1,:-1,:-1] - else: - R,theta,phi=r_theta_phi_uniform(filename) - thetac = 0.5*(theta[1:]+theta[:-1]) - phic = 0.5*(phi[1:]+phi[:-1]) - - phic = phic[:,None,None] - thetac = thetac[None,:,None] - - if var_name in ['br','btheta','bphi']: - bx=hdffile.select('bx_').get()[:-1,:-1,:-1] - by=hdffile.select('by_').get()[:-1,:-1,:-1] - bz=hdffile.select('bz_').get()[:-1,:-1,:-1] - - if var_name=='br': - var = bx*cos(phic)*sin(thetac) + by*sin(phic)*sin(thetac) + bz*cos(thetac) - elif var_name=='btheta': - var = bx*cos(phic)*cos(thetac) + by*sin(phic)*cos(thetac) - bz*sin(thetac) - else: - var =-bx*sin(phic) + by*cos(phic) - else: - vx=hdffile.select('vx_').get()[:-1,:-1,:-1] - vy=hdffile.select('vy_').get()[:-1,:-1,:-1] - vz=hdffile.select('vz_').get()[:-1,:-1,:-1] - - if var_name=='vr': - var = vx*cos(phic)*sin(thetac) + vy*sin(phic)*sin(thetac) + vz*cos(thetac) - elif var_name=='vtheta': - var = vx*cos(phic)*cos(thetac) + vy*sin(phic)*cos(thetac) - vz*sin(thetac) - else: - var =-vx*sin(phic) + vy*cos(phic) - hdffile.end() - return(var) - -def read_var_point(filename,var_name,i,j,k,thetac,phic): - thetac = thetac[j] - phic = phic[k] - - - hdffile = SD(filename,SDC.READ) - if var_name not in ['br','btheta','bphi','vr','vtheta','vphi']: - var=hdffile.select(var_name+'_').get(start=(k,j,i),count=(1,1,1)).squeeze() - else: -# R,theta,phi=r_theta_phi_uniform(filename) - - if var_name in ['br','btheta','bphi']: - bx=hdffile.select('bx_').get(start=(k,j,i),count=(1,1,1)).squeeze() - by=hdffile.select('by_').get(start=(k,j,i),count=(1,1,1)).squeeze() - bz=hdffile.select('bz_').get(start=(k,j,i),count=(1,1,1)).squeeze() - - if var_name=='br': - var = bx*cos(phic)*sin(thetac) + by*sin(phic)*sin(thetac) + bz*cos(thetac) - elif var_name=='btheta': - var = bx*cos(phic)*cos(thetac) + by*sin(phic)*cos(thetac) - bz*sin(thetac) - else: - var =-bx*sin(phic) + by*cos(phic) - else: - vx=hdffile.select('vx_').get(start=(k,j,i),count=(1,1,1)).squeeze() - vy=hdffile.select('vy_').get(start=(k,j,i),count=(1,1,1)).squeeze() - vz=hdffile.select('vz_').get(start=(k,j,i),count=(1,1,1)).squeeze() - - if var_name=='vr': - var = vx*cos(phic)*sin(thetac) + vy*sin(phic)*sin(thetac) + vz*cos(thetac) - elif var_name=='vtheta': - var = vx*cos(phic)*cos(thetac) + vy*sin(phic)*cos(thetac) - vz*sin(thetac) - else: - var =-vx*sin(phic) + vy*cos(phic) - hdffile.end() - return(var) - -def read_var_islice(filename,var_name,i,thetac,phic): - nk = phic.size - nj = thetac.size - phic = phic[:,None] - thetac = thetac[None,:] - - hdffile = SD(filename,SDC.READ) - if var_name not in ['br','btheta','bphi','vr','vtheta','vphi']: - var=hdffile.select(var_name+'_').get(start=(0,0,i),count=(nk,nj,1)).squeeze() - else: - if var_name in ['br','btheta','bphi']: - bx=hdffile.select('bx_').get(start=(0,0,i),count=(nk,nj,1)).squeeze() - by=hdffile.select('by_').get(start=(0,0,i),count=(nk,nj,1)).squeeze() - bz=hdffile.select('bz_').get(start=(0,0,i),count=(nk,nj,1)).squeeze() - - if var_name=='br': - var = bx*cos(phic)*sin(thetac) + by*sin(phic)*sin(thetac) + bz*cos(thetac) - elif var_name=='btheta': - var = bx*cos(phic)*cos(thetac) + by*sin(phic)*cos(thetac) - bz*sin(thetac) - else: - var =-bx*sin(phic) + by*cos(phic) - else: - vx=hdffile.select('vx_').get(start=(0,0,i),count=(nk,nj,1)).squeeze() - vy=hdffile.select('vy_').get(start=(0,0,i),count=(nk,nj,1)).squeeze() - vz=hdffile.select('vz_').get(start=(0,0,i),count=(nk,nj,1)).squeeze() - - if var_name=='vr': - var = vx*cos(phic)*sin(thetac) + vy*sin(phic)*sin(thetac) + vz*cos(thetac) - elif var_name=='vtheta': - var = vx*cos(phic)*cos(thetac) + vy*sin(phic)*cos(thetac) - vz*sin(thetac) - else: - var =-vx*sin(phic) + vy*cos(phic) - hdffile.end() - return(var) - -def read_var_jslice(filename,var_name,j,thetac,phic): - hdffile = SD(filename,SDC.READ) - ni = hdffile.select('X_grid').ni-1 - nk = hdffile.select('X_grid').nk-1 - phic = phic[:,None] - thetac = thetac[j] - - if var_name not in ['br','btheta','bphi','vr','vtheta','vphi']: - if var_name in ['X_grid','Y_grid','Z_grid']: - var=hdffile.select(var_name).get(start=(0,j,0),count=(nk,1,ni)).squeeze() - else: - var=hdffile.select(var_name+'_').get(start=(0,j,0),count=(nk,1,ni)).squeeze() - else: - if var_name in ['br','btheta','bphi']: - bx=hdffile.select('bx_').get(start=(0,j,0),count=(nk,1,ni)).squeeze() - by=hdffile.select('by_').get(start=(0,j,0),count=(nk,1,ni)).squeeze() - bz=hdffile.select('bz_').get(start=(0,j,0),count=(nk,1,ni)).squeeze() - - if var_name=='br': - var = bx*cos(phic)*sin(thetac) + by*sin(phic)*sin(thetac) + bz*cos(thetac) - elif var_name=='btheta': - var = bx*cos(phic)*cos(thetac) + by*sin(phic)*cos(thetac) - bz*sin(thetac) - else: - var =-bx*sin(phic) + by*cos(phic) - else: - vx=hdffile.select('vx_').get(start=(0,j,0),count=(nk,1,ni)).squeeze() - vy=hdffile.select('vy_').get(start=(0,j,0),count=(nk,1,ni)).squeeze() - vz=hdffile.select('vz_').get(start=(0,j,0),count=(nk,1,ni)).squeeze() - - if var_name=='vr': - var = vx*cos(phic)*sin(thetac) + vy*sin(phic)*sin(thetac) + vz*cos(thetac) - elif var_name=='vtheta': - var = vx*cos(phic)*cos(thetac) + vy*sin(phic)*cos(thetac) - vz*sin(thetac) - else: - var =-vx*sin(phic) + vy*cos(phic) - hdffile.end() - return(var) - -def read_var_ikslice(filename,var_name,i,k,tc,pc): - hdffile = SD(filename,SDC.READ) - ni = hdffile.select('X_grid').ni-1 - nj = hdffile.select('X_grid').nj-1 - - if var_name not in ['br','btheta','bphi','vr','vtheta','vphi','bt','bp','vt','vp']: - var=hdffile.select(var_name+'_').get(start=(k,0,i),count=(1,nj,1)).squeeze() - else: - if var_name in ['br','btheta','bphi','bt','bp']: - bx=hdffile.select('bx_').get(start=(k,0,i),count=(1,nj,1)).squeeze() - by=hdffile.select('by_').get(start=(k,0,i),count=(1,nj,1)).squeeze() - bz=hdffile.select('bz_').get(start=(k,0,i),count=(1,nj,1)).squeeze() - - if var_name=='br': - var = bx*cos(pc[k])*sin(tc) + by*sin(pc[k])*sin(tc) + bz*cos(tc) - elif (var_name=='btheta' or var_name=='bt'): - var = bx*cos(pc[k])*cos(tc) + by*sin(pc[k])*cos(tc) - bz*sin(tc) - else: - var =-bx*sin(pc[k]) + by*cos(pc[k]) - else: - vx=hdffile.select('vx_').get(start=(k,0,i),count=(1,nj,1)).squeeze() - vy=hdffile.select('vy_').get(start=(k,0,i),count=(1,nj,1)).squeeze() - vz=hdffile.select('vz_').get(start=(k,0,i),count=(1,nj,1)).squeeze() - - if var_name=='vr': - var = vx*cos(pc[k])*sin(tc) + vy*sin(pc[k])*sin(tc) + vz*cos(tc) - elif (var_name=='vtheta' or var_name=='vt'): - var = vx*cos(pc[k])*cos(tc) + vy*sin(pc[k])*cos(tc) - vz*sin(tc) - else: - var =-vx*sin(pc[k]) + vy*cos(pc[k]) - hdffile.end() - return(var) - -def get_time(filename): - hdffile = SD(filename, SDC.READ) - t=hdffile.time - hdffile.end() - return(t) - -def time_shift(t,var,Tsolar=25.38*24*3600.): - """ - Shift variables to the beginning of the CR. - - Parameters: - t: current time - var: variable -- assumed to be a 2D i-slice (nk,nj) of the LFM data - Tsolar: solar rotation period in units of time used in the code. Default to 25.38 days and seconds as time units. - """ - time_shift= (t % Tsolar)/Tsolar - nshift = int(time_shift*(var.shape[0]+1)) # note, using nk+1 here, which gives us full rotation - return roll(var,-nshift,axis=0) - diff --git a/kaipy/gamhelio/lib/mas.py b/kaipy/gamhelio/lib/mas.py deleted file mode 100644 index 9209851..0000000 --- a/kaipy/gamhelio/lib/mas.py +++ /dev/null @@ -1,138 +0,0 @@ -from pyhdf.SD import SD,SDC -import os,sys -import ConfigParser - -mas_units = {'length':6.96e10, - 'time':1445.87, - 'vr':481.3711*1.e5, - 'vt':481.3711*1.e5, - 'vp':481.3711*1.e5, - 'n':1.e8, - 'rho':1.6726e-16, - 'p':0.3875717, - 't':2.807067e7, - 'br':2.2068908, - 'bt':2.2068908, - 'bp':2.2068908} - -mas_var_names = ['t','rho','vt','vp','vr','bt','bp','br'] - -def read_var(fname,varname,normalized=False): - f = SD(fname,SDC.READ) - phi = f.select('fakeDim0')[:] - theta = f.select('fakeDim1')[:] - r = f.select('fakeDim2')[:] - var = f.select('Data-Set-2')[:] - f.end() - - if normalized: - return(phi,theta,r,var) - else: - return(phi,theta,r*mas_units['length'],var*mas_units[varname]) - -def units(): - return(mas_units) - - -def read_all_vars(dir,timeLabel,normalized=False): - vars = {} - - for var_name in mas_var_names: - vars[var_name]={} - (vars[var_name]['phi'], - vars[var_name]['theta'], - vars[var_name]['r'], - vars[var_name]['data']) = read_var(os.path.join(dir,var_name+'%03d.hdf'%timeLabel),var_name,normalized) - - return(vars) - -def set_plot_limits(v,plotConfigFile = 'plot.config'): - plotConfig = ConfigParser.ConfigParser() - plotConfig.read(plotConfigFile) - - # assume here v has the necessary keys - # bascially, this needs to be called after mas.read_all_varaibles - try: - v['bt']['lims'] = plotConfig.getfloat('MAS','bt_lim_lo'),plotConfig.getfloat('MAS','bt_lim_hi') - v['bt']['fmt'] = plotConfig.get('MAS','bt_fmt') - - v['bp']['lims'] = plotConfig.getfloat('MAS','bp_lim_lo'),plotConfig.getfloat('MAS','bp_lim_hi') - v['bp']['fmt'] = plotConfig.get('MAS','bp_fmt') - - v['br']['lims'] = plotConfig.getfloat('MAS','br_lim_lo'),plotConfig.getfloat('MAS','br_lim_hi') - v['br']['fmt'] = plotConfig.get('MAS','br_fmt') - - v['vr']['lims'] = plotConfig.getfloat('MAS','vr_lim_lo'),plotConfig.getfloat('MAS','vr_lim_hi') - v['vr']['fmt'] = plotConfig.get('MAS','vr_fmt') - - v['vt']['lims'] = plotConfig.getfloat('MAS','vt_lim_lo'),plotConfig.getfloat('MAS','vt_lim_hi') - v['vt']['fmt'] = plotConfig.get('MAS','vt_fmt') - - v['vp']['lims'] = plotConfig.getfloat('MAS','vp_lim_lo'),plotConfig.getfloat('MAS','vp_lim_hi') - v['vp']['fmt'] = plotConfig.get('MAS','vp_fmt') - - v['rho']['lims'] = plotConfig.getfloat('MAS','rho_lim_lo'),plotConfig.getfloat('MAS','rho_lim_hi') - v['rho']['fmt'] = plotConfig.get('MAS','rho_fmt') - - v['t']['lims'] = plotConfig.getfloat('MAS','t_lim_lo'),plotConfig.getfloat('MAS','t_lim_hi') - v['t']['fmt'] = plotConfig.get('MAS','t_fmt') - except KeyError: - print('Some keys are missing from the variable dictionary.') - - -def mas2h5(fname,vname): - from numpy import meshgrid,sin,cos - - if vname not in mas_var_names: - sys.exit('Wrong variable name') - - fout = os.path.splitext(fname)[0]+'.h5' - - import mas - phi,theta,r,var=mas.read_var(fname,vname) - import h5py - f=h5py.File(fout,'w') - P,T,R = meshgrid(phi,theta,r,indexing='ij') - x = R*sin(T)*cos(P)/mas_units['length'] - y = R*sin(T)*sin(P)/mas_units['length'] - z = R*cos(T)/mas_units['length'] - f.create_dataset('x',data= x.astype(dtype='float32')) - f.create_dataset('y',data= y.astype(dtype='float32')) - f.create_dataset('z',data= z.astype(dtype='float32')) - f.create_dataset(vname,data= var.astype(dtype='float32')) - f.close() - - # write xmf file - mas2xdmf(fout,vname,phi.size,theta.size,r.size) - -def mas2xdmf(h5name,varname,n1,n2,n3): - fname=os.path.splitext(h5name)[0]+'.xmf' - dsname = os.path.basename(h5name) - - f = open(fname,'w') - f.write('' + '\n'+ - '' + '\n'+ - '' + '\n'+ - ' ' + '\n'+ - ' ' + '\n'+ - ' '%(n1,n2,n3) + '\n'+ - ' ' + '\n'+ - ' ' %(n1,n2,n3)+ '\n'+ - ' %s:/x'%dsname + '\n'+ - ' ' + '\n'+ - ' ' %(n1,n2,n3)+ '\n'+ - ' %s:/y'%dsname + '\n'+ - ' ' + '\n'+ - ' '%(n1,n2,n3) + '\n'+ - ' %s:/z'%dsname + '\n'+ - ' ' + '\n'+ - ' ' + '\n'+ - ' '%varname + '\n'+ - ' '%(n1,n2,n3) + '\n'+ - ' %s:/%s'%(dsname,varname) + '\n'+ - ' ' + '\n'+ - ' ' + '\n'+ - ' ' + '\n'+ - ' ' + '\n'+ - '' ) - f.close() diff --git a/kaipy/gamhelio/lib/msgr.py b/kaipy/gamhelio/lib/msgr.py deleted file mode 100644 index c0b1139..0000000 --- a/kaipy/gamhelio/lib/msgr.py +++ /dev/null @@ -1,51 +0,0 @@ -import numpy as np -import datetime - -def msgr(msgr_file,ave_1h=True): - r=[] - lat=[] - lon=[] - br=[] - bt=[] - bn=[] - dati=[] - for line in open(msgr_file,'r'): - fields = line.strip().split() - - year = int(fields[0]) - doy = int(fields[1]) - hr = int(fields[2]) - mn = int(fields[3]) - r.append(float(fields[7])) - lat.append(float(fields[8])) - lon.append(float(fields[9])) - br.append(float(fields[10])) - bt.append(float(fields[11])) - bn.append(float(fields[12])) - - dati.append(datetime.datetime(year,1,1,hr,mn)+datetime.timedelta(days=doy-1)) - - - if ave_1h: - # assuming 1-min data. Average to 1-h - n=len(lat)-len(lat)%60 - br_1h=np.mean(np.array(br)[:n].reshape(-1,60),1) - bt_1h=np.mean(np.array(bt)[:n].reshape(-1,60),1) - bn_1h=np.mean(np.array(bn)[:n].reshape(-1,60),1) - - return (dati[30:n+30:60],br_1h,bt_1h,bn_1h) - else: - return (dati,np.array(br),np.array(bt),np.array(bn)) - -def msgr_v1(msgr_file,year): - data = np.loadtxt(msgr_file) - - start_dati = datetime.datetime(year,1,1,0) - dati = [start_dati+datetime.timedelta(days=d-1) for d in data[:,0]] - - br=data[:,4] - bt=data[:,5] - bn=data[:,6] - - return (dati,br,bt,bn) - diff --git a/kaipy/gamhelio/lib/poisson.py b/kaipy/gamhelio/lib/poisson.py deleted file mode 100644 index a356bfe..0000000 --- a/kaipy/gamhelio/lib/poisson.py +++ /dev/null @@ -1,128 +0,0 @@ -from scipy.optimize import newton_krylov,anderson -from numpy import ediff1d,zeros_like,cos,sin,zeros_like,arange - -class poisson(): - def __init__(self,theta,phi): - self.theta = theta - self.phi = phi - self.nj = theta.shape[0] - self.nk = phi.shape[0] - self.dtheta = theta[1]-theta[0] # generalize to the case of arbitrary grid spacing later: ediff1d(theta) - self.dphi = phi[1]-phi[0] # ediff1d(phi) - - def setRHS(self,RHS): - self.RHS = RHS - - def residual_xsin2(self,P,lhs=False): - dphi = self.dphi - dtheta = self.dtheta - theta = self.theta - nk = self.nk - - - d2x = zeros_like(P) - d2y = zeros_like(P) - dy = zeros_like(P) - - d2x[:,1:-1] = (P[:,2:] - 2*P[:,1:-1] + P[:,:-2]) /dphi**2 - d2y[1:-1,:] = (P[2:,:] - 2*P[1:-1,:] + P[:-2,:]) /dtheta**2*sin(theta[1:-1,None])**2 - - dy[1:-1,:] = 0.5*(P[2:,:] - P[:-2,:]) /dtheta*sin(theta[1:-1,None])*cos(theta[1:-1,None]) - - # periodic boundary - d2x[:,0] = (P[:,1]-2*P[:,0]+P[:,-1]) /dphi**2 - d2x[:,-1] = (P[:,0]-2*P[:,-1]+P[:,-2]) /dphi**2 - - result = self.RHS*sin(theta[:,None])**2 - - # Pole boundary condition. Note, the below does not mean that d/dtheta=0 literally. - # We set sin(theta)*dPsi/dtheta=0 at the pole (see notes for details), - # and d2y here represents 1/sin(theta)*d/dtheta (sin(theta)*dPsi/dtheta) in that approxiamtion. - # NOTE, BELOW ASSUMES CONSTANT SPACING IN THETA!!! (0.5 FACTOR IS BASED ON THAT). - dy[0,:]=0. - dy[-1,:]=0. - - d2y[0,:] = 0.5*(P[1,:]-P[0,:]) - d2y[-1,:] = -0.5*(P[-1,:]-P[-2,:]) - - ################################################################################################### - if not lhs: - return d2x + d2y + dy - result - else: - return d2x + d2y + dy - - def residual(self,P,lhs=False): - dphi = self.dphi - dtheta = self.dtheta - theta = self.theta - nk = self.nk - - - d2x = zeros_like(P) - d2y = zeros_like(P) - dy = zeros_like(P) - - d2x[:,1:-1] = (P[:,2:] - 2*P[:,1:-1] + P[:,:-2]) /dphi**2/sin(theta[:,None])**2 - d2y[1:-1,:] = (P[2:,:] - 2*P[1:-1,:] + P[:-2,:]) /dtheta**2 - - dy[1:-1,:] = 0.5*(P[2:,:] - P[:-2,:]) /dtheta/sin(theta[1:-1,None])*cos(theta[1:-1,None]) - - # periodic boundary - d2x[:,0] = (P[:,1]-2*P[:,0]+P[:,-1]) /dphi**2/sin(theta)**2 - d2x[:,-1] = (P[:,0]-2*P[:,-1]+P[:,-2]) /dphi**2/sin(theta)**2 - - result = self.RHS - - # Pole boundary condition. Note, the below does not mean that d/dtheta=0 literally. - # We set sin(theta)*dPsi/dtheta=0 at the pole (see notes for details), - # and d2y here represents 1/sin(theta)*d/dtheta (sin(theta)*dPsi/dtheta) in that approxiamtion. - # NOTE, BELOW ASSUMES CONSTANT SPACING IN THETA!!! (0.5 FACTOR IS BASED ON THAT). - dy[0,:]=0. - dy[-1,:]=0. - - d2y[0,:] = 0.5*(P[1,:]-P[0,:])/sin(theta[0])**2 # Psin=0; 0.5*( 0.5*(P[1,:]+P[0,:]) - Psin) - d2y[-1,:] = -0.5*(P[-1,:]-P[-2,:])/sin(theta[-1])**2 #-0.5*( 0.5*(P[-2,:]+P[-1,:]) - P[-1,:].mean()) - - ################################################################################################### - if not lhs: - return d2x + d2y + dy - result - else: - return d2x + d2y + dy - -if __name__=='__main__': - from scipy.optimize import newton_krylov,anderson - from pylab import * - - import wsa - (phi_wsa_v,theta_wsa_v,phi,theta,br0,v_wsa,n_wsa,T_wsa)=wsa.read('/Users/merkivg1/work/LFM-helio_2.0/WSA/fits/2010_2008_WSA_ADAPT/vel_201001022300R001_ans.fits',False) - - (phi_wsa_v,theta_wsa_v,phi,theta,br1,v_wsa,n_wsa,T_wsa)=wsa.read('/Users/merkivg1/work/LFM-helio_2.0/WSA/fits/2010_2008_WSA_ADAPT/vel_201001032300R001_ans.fits',False) - - omega=2*pi/27.27 - phi_prime=phi-omega*1. - ind0=where(phi_prime>0)[0][0] - phi_prime=roll(phi_prime,-ind0) - phi_prime[phi_prime<0]+=2*pi - - import scipy - from scipy.interpolate import interpolate - br1_rolled=roll(br1,-ind0,axis=1) - fbi = scipy.interpolate.RectBivariateSpline(theta,phi,br1_rolled,kx=1,ky=1) - br1_new=fbi(theta,phi_prime) - - import poisson - pois = poisson.poisson(theta,phi) - pois.setRHS(br1_new - br0) - - guess=zeros_like(br1_new) - sol = newton_krylov(pois.residual,guess, method='lgmres', verbose=1,f_rtol=1.e-6) #iter=100 - print('Residual: %g' % abs(pois.residual(sol)).max()) - - figure();pcolormesh(sol);colorbar();title('Solution') - - # check curl - Etheta = -1/sin(theta[:,None])*diff(sol,axis=1) - Ephi = diff(sol,axis=0) - figure();pcolormesh(Etheta);colorbar();title('Etheta') - figure();pcolormesh(Ephi);colorbar();title('Ephi') - diff --git a/kaipy/gamhelio/lib/util.py b/kaipy/gamhelio/lib/util.py deleted file mode 100644 index e97ee71..0000000 --- a/kaipy/gamhelio/lib/util.py +++ /dev/null @@ -1,24 +0,0 @@ -import numpy as np - -def sph2cart(a,theta,phi): - ar,at,ap = a - - ax = ar*cos(phi)*sin(theta)+at*cos(phi)*cos(theta)-ap*sin(phi) - ay = ar*sin(phi)*sin(theta)+at*sin(phi)*cos(theta)+ap*cos(phi) - az = ar*cos(theta)-at*sin(theta) - - return(ax,ay,az) - -def radial_interp(vars_from,var_name,radii_to,coef_label='r_interp_coefs',ind_label='r_interp_above_ind'): - q=[] - ind=[] - for r_to in radii_to: # interpolating to LFM cell centers - r_from = vars_from[var_name]['r'] - ind_r_a = np.flatnonzero(r_from>=r_to)[0] # MAS radius above - r_a = r_from[ind_r_a] - r_b = r_from[ind_r_a-1] - q.append( (r_a-r_to)/(r_a-r_b) ) - ind.append(ind_r_a) - - vars_from[var_name][coef_label]=np.asarray(q,dtype='float') - vars_from[var_name][ind_label]=np.asarray(ind,dtype='int') diff --git a/kaipy/gamhelio/lib/wsa2h5.py b/kaipy/gamhelio/lib/wsa2h5.py deleted file mode 100755 index 13285f4..0000000 --- a/kaipy/gamhelio/lib/wsa2h5.py +++ /dev/null @@ -1,31 +0,0 @@ -#!/usr/bin/env python - -# vgm, Sep 2019: Script to convert WSA fits files into simple h5 files - -from astropy.io import fits -import h5py -import glob - -wsaFiles = glob.glob('./*.fits') - -for wsaFile in wsaFiles: - h5file = wsaFile[:-4]+'h5' - - with fits.open(wsaFile) as hdul: - with h5py.File(h5file,'w') as hf: - # take care of the comments - # removing the last ones since we're incorporating them into dataset attributes directly - for key in hdul[0].header[:-4]: - # stores header data as attributes of the file - hf.attrs[key] = hdul[0].header[key] - hf.attrs[key+" COMMENT"] = str(hdul[0].header.comments[key]) - - # take care of the datasets - br = hf.create_dataset("br",data=hdul[0].data[0]) - br.attrs['Units'] = 'nT' - br.attrs['Comment'] = 'Coronal field at RADOUT' - - vr = hf.create_dataset("vr",data=hdul[0].data[1]) - vr.attrs['Units'] = 'km/s' - vr.attrs['Comment'] = 'Radial velocity at RADOUT' - diff --git a/kaipy/gamhelio/wsa2TDgamera/__init__.py b/kaipy/gamhelio/wsa2TDgamera/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/kaipy/gamhelio/wsa2TDgamera/params.py b/kaipy/gamhelio/wsa2TDgamera/params.py deleted file mode 100644 index 2804207..0000000 --- a/kaipy/gamhelio/wsa2TDgamera/params.py +++ /dev/null @@ -1,55 +0,0 @@ -import configparser - -class params(): - def __init__(self,ConfigFileName): - config = configparser.ConfigParser() - config.read(ConfigFileName) - - - #self.initLFMfile = config.get('Output','Prefix')+'_%dx%dx%d_mhd_0000000.hdf'%(self.ni,self.nj,self.nk) - - self.scale = config.getfloat('Normalization','Xscale') - self.gameraGridFile = config.get('Grid','gameraGridFile') - self.GridDir = config.get('Grid','GridDir') - self.gameraIbcFile = config.get('Grid','gameraIbcFile') - self.IbcDir = config.get('Grid','IbcDir') - self.tMin = config.getfloat('Grid','tMin') - self.tMax = config.getfloat('Grid','tMax') - self.Rin = config.getfloat('Grid','Rin') - self.Rout = config.getfloat('Grid','Rout') - self.Ni = config.getint('Grid','Ni') - self.Nj = config.getint('Grid','Nj') - self.Nk = config.getint('Grid','Nk') - - - self.gamma = config.getfloat('Constants','gamma') - self.NO2 = config.getint('Constants','NO2') - self.Tsolar = config.getfloat('Constants','Tsolar') - - if config.has_section('MAS'): - self.masdir = config.get('MAS','masdir') - self.masTimeLabel = config.get('MAS','masTimeLabel') - self.masFrame = config.get('MAS','masFrame') - self.masFakeRotation = config.getboolean('MAS','masFakeRotation') - - if config.has_section('WSA'): - self.wsaFile = config.get('WSA','wsafile') - self.gaussSmoothWidth = config.getint('WSA','gauss_smooth_width') - self.plots = config.getboolean('WSA','plots') - self.densTempInfile = config.getboolean('WSA','density_temperature_infile') - self.normalized = config.getboolean('WSA','normalized') - - if config.has_section('ADAPT'): - self.adaptdir = config.get('ADAPT','adaptdir') - self.adaptWildCard = config.get('ADAPT','wild_card') - self.adaptCadence = config.getfloat('ADAPT','cadence') - self.gaussSmoothWidth = config.getint('ADAPT','gauss_smooth_width') - self.plots = config.getboolean('ADAPT','plots') - self.densTempInfile = config.getboolean('ADAPT','density_temperature_infile') - self.normalized = config.getboolean('ADAPT','normalized') - - - self.dumpInit = config.getboolean('DUMPS','init') - self.dumpBC = config.getboolean('DUMPS','BC') - - diff --git a/kaipy/kaiH5.py b/kaipy/kaiH5.py index b20641c..167b4a2 100644 --- a/kaipy/kaiH5.py +++ b/kaipy/kaiH5.py @@ -585,7 +585,7 @@ def PullVar(fname, vID, s0=None, slice=()): return V #Get attribute data from Step#s0 or root (s0=None) -def PullAtt(fname, vID, s0=None): +def PullAtt(fname, vID, s0=None, a0=None): ''' Retrieve an attribute from an HDF5 file. @@ -604,9 +604,16 @@ def PullAtt(fname, vID, s0=None): CheckOrDie(fname) with h5py.File(fname, 'r') as hf: if s0 is None: - Q = hf.attrs[vID] + hfA = hf.attrs else: gID = "/Step#%d" % (s0) - Q = hf[gID].attrs[vID] + hfA = hf[gID].attrs + + if (vID not in hfA.keys()): + #Use default + Q = a0 + else: + Q = hfA[vID] + return Q diff --git a/kaipy/scripts/checkkaipypath.py b/kaipy/scripts/checkkaipypath.py old mode 100644 new mode 100755 diff --git a/kaipy/scripts/datamodel/helioSatComp.py b/kaipy/scripts/datamodel/helioSatComp.py index 7e66bd1..81af758 100755 --- a/kaipy/scripts/datamodel/helioSatComp.py +++ b/kaipy/scripts/datamodel/helioSatComp.py @@ -56,7 +56,7 @@ default_path = os.getcwd() # Path to file of heliospheric spacecraft metadata. helio_sc_metadata_path = os.path.join( - os.environ["KAIJUHOME"], "kaipy", "satcomp", "sc_helio.json" + os.environ["KAIPYHOME"], "kaipy", "satcomp", "sc_helio.json" ) diff --git a/kaipy/scripts/postproc/calcdb.xml.template b/kaipy/scripts/postproc/calcdb.xml.template deleted file mode 100644 index b5890a2..0000000 --- a/kaipy/scripts/postproc/calcdb.xml.template +++ /dev/null @@ -1,9 +0,0 @@ - - - - - - diff --git a/kaipy/scripts/postproc/genXDMF.py b/kaipy/scripts/postproc/genXDMF.py index a2ae073..9a35fca 100755 --- a/kaipy/scripts/postproc/genXDMF.py +++ b/kaipy/scripts/postproc/genXDMF.py @@ -8,7 +8,7 @@ import xml.etree.ElementTree as et import xml.dom.minidom import numpy as np -presets = {"gam", "mhdrcm_eq", "mhdrcm_bmin"} +presets = {"gam", "mhdrcm_eq", "mhdrcm_bmin", 'voltSG'} def getDimInfo(h5fname,s0IDstr,preset): @@ -20,7 +20,6 @@ def getDimInfo(h5fname,s0IDstr,preset): gDims = np.asarray(h5f[s0IDstr][gridVars[0]].shape) result['gridVars'] = gridVars result['gDims'] = gDims - result['vDims'] = gDims result['Nd'] = len(gDims) result['geoStr'] = "X_Y" result['topoStr'] = "2DSMesh" @@ -29,25 +28,22 @@ def getDimInfo(h5fname,s0IDstr,preset): gridVars = ['xMin','yMin','zMin'] with h5.File(h5fname, 'r') as h5f: gDims = np.asarray(h5f[s0IDstr][gridVars[0]].shape) - #gDims = np.append(gDims, 1) result['gridVars'] = gridVars result['gDims'] = gDims - result['vDims'] = gDims result['Nd'] = len(gDims) result['geoStr'] = "X_Y_Z" result['topoStr'] = "3DSMesh" result['doAppendStep'] = True - elif preset=="rcm3D": - gridVars = ["rcmxmin_kji", "rcmymin_kji", "rcmalamc_kji"] + elif preset=="voltSG": + gridVars=['X','Y','Z'] with h5.File(h5fname, 'r') as h5f: - gDims = np.asarray(h5f[s0IDstr][gridVars[0]].shape) + gDims = np.asarray(h5f[gridVars[0]].shape) result['gridVars'] = gridVars result['gDims'] = gDims - result['vDims'] = gDims result['Nd'] = len(gDims) result['geoStr'] = "X_Y_Z" result['topoStr'] = "3DSMesh" - result['doAppendStep'] = True + result['doAppendStep'] = False else: # gam, mhdrcm_iono, etc. #Get root-level XY(Z) dimensions #First check to see if they exist @@ -67,7 +63,6 @@ def getDimInfo(h5fname,s0IDstr,preset): result['gridVars'] = gridVars result['gDims'] = gDims - result['vDims'] = gDims - 1 result['Nd'] = len(gDims) result['geoStr'] = '_'.join(gridVars) result['topoStr'] = topoStr @@ -80,9 +75,9 @@ def addRCMVars(Grid, dimInfo, rcmInfo, sID): sIDstr = "Step#" + str(sID) - mr_vDims = dimInfo['vDims'] # mhdrcm var dims + mr_vDims = dimInfo['gDims'] # mhdrcm var dims mr_vDimStr = ' '.join([str(v) for v in mr_vDims]) - mr_nDims = len(vDims) + mr_nDims = len(mr_vDims) rcmh5fname = rcmInfo['rcmh5fname'] rcmVars = rcmInfo['rcmVars'] # List of rcm.h5 variables we want in mhdrcm.xmf rcmKs = rcmInfo['rcmKs'] # List if rcm.h5 k values for 3d rcm.h5 vars @@ -179,13 +174,13 @@ if __name__ == "__main__": dimInfo = getDimInfo(h5fname, s0str, preset) gridVars = dimInfo['gridVars'] gDims = dimInfo['gDims'] - vDims = dimInfo['vDims'] Nd = dimInfo['Nd'] geoStr = dimInfo['geoStr'] topoStr = dimInfo['topoStr'] doAppendStep = dimInfo['doAppendStep'] gDimStr = ' '.join([str(v) for v in gDims]) - vDimStr = ' '.join([str(v) for v in vDims]) + vDimStr_corner = ' '.join([str(v) for v in gDims]) + vDimStr_cc = ' '.join([str(v-1) for v in gDims]) doAddRCMVars = False #Prep to include some rcmh5 vars in mhdrcm.xmf file @@ -283,10 +278,12 @@ if __name__ == "__main__": #-------------------------------- #Step variables for v in range(Nv): + vDimStr = vDimStr_corner if vLocs[v]=="Node" else vDimStr_cc kxmf.AddData(Grid,fNames_link[n],vIds[v],vLocs[v],vDimStr,steps_link[n]) #-------------------------------- #Base grid variables for v in range(Nrv): + vDimStr = vDimStr_corner if rvLocs[v]=="Node" else vDimStr_cc kxmf.AddData(Grid,fNames_link[n],rvIds[v],rvLocs[v],vDimStr) if doAddRCMVars: diff --git a/kaipy/scripts/postproc/run_calcdb.py b/kaipy/scripts/postproc/run_calcdb.py new file mode 100755 index 0000000..cb61909 --- /dev/null +++ b/kaipy/scripts/postproc/run_calcdb.py @@ -0,0 +1,594 @@ +#!/usr/bin/env python + + +"""Set up and run calcdb.x. + +Run calcdb.x to compute ground magnetic field perturbations for a MAGE +magnetosphere simulation. This is all done by a linked set of PBS jobs running +in the MAGE results directory. + +NOTE: Before running this script, the user must run the setupEnvironment.sh +scripts for kaipy and kaiju. + +NOTE: If the calcdb.x binary was built with a module set other than those +listed below, change the module set in the PBS scripts appropriately. + +Author +------ +Eric Winter (eric.winter@jhuapl.edu) + +""" + + +# Import standard modules. +import argparse +import copy +import os +import pathlib +import re +import shutil +import sys + +# Import 3rd-party modules. +from jinja2 import Template + +# Import project-specific modules. +from kaipy import kaiH5 +from kaipy import kaiTools + + +# Program constants and defaults + +# Program description. +DESCRIPTION = "Run calcdb.x to compute ground delta-B values." + +# Default values for command-line arguments. +DEFAULT_ARGUMENTS = { + "calcdb": "calcdb.x", + "debug": False, + "dt": 60.0, + "hpc": "pleiades", + "parintime": 1, + "pbs_account": None, + "verbose": False, + "mage_results_path": None, +} + +# Valid values for command-line arguments. +VALID_HPC = ["derecho", "pleiades"] + +# Directory containing this script. +SCRIPT_DIRECTORY = pathlib.Path(__file__).parent.resolve() +TEMPLATE_DIRECTORY = os.path.join(SCRIPT_DIRECTORY, "templates") + +# Location of template XML file for calcdb.x. +CALCDB_XML_TEMPLATE = os.path.join(TEMPLATE_DIRECTORY, "calcdb-template.xml") + +# # Locations of template PBS file for running calcdb.x. +CALCDB_PBS_TEMPLATE = os.path.join(TEMPLATE_DIRECTORY, "calcdb-template.pbs") + +# Location of template PBS file for pitmerge.py. +PITMERGE_PBS_TEMPLATE = os.path.join( + TEMPLATE_DIRECTORY, "pitmerge-template.pbs" +) + + +def create_command_line_parser(): + """Create the command-line parser. + + Create the parser for the command line. + + Parameters + ---------- + None + + Returns + ------- + parser : argparse.ArgumentParser + Command-line parser for this script. + + Raises + ------ + None + """ + parser = argparse.ArgumentParser(description=DESCRIPTION) + parser.add_argument( + "--calcdb", type=str, default=DEFAULT_ARGUMENTS["calcdb"], + help="Path to calcdb.x binary (default: %(default)s)." + ) + parser.add_argument( + "--debug", "-d", default=DEFAULT_ARGUMENTS["debug"], + action="store_true", + help="Print debugging output (default: %(default)s)." + ) + parser.add_argument( + "--dt", type=float, default=DEFAULT_ARGUMENTS["dt"], + help="Time interval for delta B computation (simulated seconds) " + "(default: %(default)s)." + ) + parser.add_argument( + "--hpc", type=str, default=DEFAULT_ARGUMENTS["hpc"], + choices=VALID_HPC, + help="HPC system to run analysis (default: %(default)s)." + ) + parser.add_argument( + "--parintime", type=int, default=DEFAULT_ARGUMENTS["parintime"], + help="Split the calculation into this many parallel chunks of kaiju " + "simulation steps, one chunk per node (default: %(default)s)." + ) + parser.add_argument( + "--pbs_account", type=str, default=DEFAULT_ARGUMENTS["pbs_account"], + help="PBS account to use for job accounting (default: %(default)s)." + ) + parser.add_argument( + "--verbose", "-v", default=DEFAULT_ARGUMENTS["verbose"], + action="store_true", + help="Print verbose output (default: %(default)s)." + ) + parser.add_argument( + "mage_results_path", type=str, + default=DEFAULT_ARGUMENTS["mage_results_path"], + help="Path to a result file for a MAGE magnetosphere run." + ) + return parser + + +def mage_filename_to_runid(filename: str): + """Parse the runid from a MAGE results file name. + + Parse the runid from a MAGE results file name. + + For a result file from an MPI run, the runid is all text before the + underscore before the set of 6 underscore-separated 4-digit sequences at + the end of the name and the terminal .gam.h5 string. + + For a result file from a serial run, the runid is the name of the file, + less the .gam.h5 extension. + + Parameters + ---------- + filename : str + Name of MAGE results file. + + Returns + ------- + runid : str + The MAGE runid for the file. + + Raises + ------ + None + """ + # Check to see if the result file is for an MPI run. If not, it must be + # for a serial run. + mpi_pattern = ( + r"^(.+)_(\d{4})_(\d{4})_(\d{4})_(\d{4})_(\d{4})_(\d{4})\.gam.h5$" + ) + serial_pattern = r"^(.+)\.gam.h5$" + mpi_re = re.compile(mpi_pattern) + serial_re = re.compile(serial_pattern) + m = mpi_re.match(filename) + if not m: + m = serial_re.match(filename) + runid = m.groups()[0] + return runid + + +def create_calcdb_xml_file(runid: str, args: dict): + """Create the XML input file for calcdb.x from a template. + + Create the XML input file for calcdb.x from a template. The file is + created in the current directory, which much contain the specified MAGE + results file. Note that the contents of the file depends on whether + parintime is 1 (use a single multi-threaded job to run calcdb.x on all + simulation steps) or > 1 (use multiple jobs to split the calcdb.x work + into an array of multi-threaded jobs which run on multiple nodes in + parallel). + + Parameters + ---------- + runid : str + Run ID for MAGE results file to use in calcdb.x calculations. + args : dict + Dictionary of command-line options. + + Returns + ------- + xml_file : str + Name of XML file. + + Raises + ------ + AssertionError + If the MAGE result file contains no time steps. + """ + # Local convenience variables. + debug = args["debug"] + verbose = args["verbose"] + + # Fetch run information from the MAGE result file. + if verbose: + print(f"Fetching MAGE run information for runid {runid}.") + filename, isMPI, Ri, Rj, Rk = kaiTools.getRunInfo(".", runid) + if debug: + print(f"filename = {filename}") + print(f"isMPI = {isMPI}") + print(f"Ri = {Ri}") + print(f"Rj = {Rj}") + print(f"Rk = {Rk}") + + # Get the number of steps and the step IDs from the MAGE results file. + if verbose: + print(f"Counting time steps for run {runid}.") + nSteps, sIds = kaiH5.cntSteps(filename) + if debug: + print(f"nSteps = {nSteps}") + print(f"sIds = {sIds}") + assert nSteps > 0 + + # Find the time for the last step. + if verbose: + print(f"Finding time for last step for run {runid}.") + tFin = kaiH5.tStep(filename, sIds[-1], aID="time") + if debug: + print(f"tFin = {tFin}") + + # Read the template XML file. + if verbose: + print("Reading calcdb XML template.") + with open(CALCDB_XML_TEMPLATE, "r", encoding="utf-8") as f: + template_content = f.read() + if debug: + print(f"template_content = {template_content}") + template = Template(template_content) + if debug: + print(f"template = {template}") + + # Fill in the template options. + options = { + "runid": runid, + "dt": args["dt"], + "tFin": tFin, + "ebfile": runid, + "isMPI": isMPI, + "Ri": Ri, + "Rj": Rj, + "Rk": Rk, + "parintime": args["parintime"], + } + if debug: + print(f"options = {options}") + + # Render the template. + xml_file = f"calcdb-{runid}.xml" + if verbose: + print("Rendering template.") + if verbose: + print(f"Creating {xml_file}.") + xml_content = template.render(options) + with open(xml_file, "w", encoding="utf-8") as f: + f.write(xml_content) + + # Return the name of the XML file. + return xml_file + + +def create_calcdb_pbs_script(runid: str, calcdb_xml_file: str, args: dict): + """Create the PBS script to run calcdb.x from a template. + + Create the PBS script to run calcdb.x from a template. The PBS script is + created in the current directory, which must contain the MAGE results. + + Parameters + ---------- + runid : str + Run ID for MAGE results to use in calcdb.x calculations. + calcdb_xml_file : str + Path to XML input file for calcdb.x. + args : dict + Dictionary of command-line and other options. + + Returns + ------- + calcdb_pbs_script : str + Path to PBS script. + + Raises + ------ + None + """ + # Local convenience variables. + debug = args["debug"] + verbose = args["verbose"] + + # Copy the calcdb.x binary to the results directory, then make it + # executable. + if verbose: + print(f"Copying {args['calcdb']} to results directory.") + shutil.copyfile(args["calcdb"], "calcdb.x") + os.chmod("calcdb.x", 0o755) + + # Read the PBS script template for running calcdb.x. + if verbose: + print("Reading calcdb.x PBS script template.") + with open(CALCDB_PBS_TEMPLATE, "r", encoding="utf-8") as f: + template_content = f.read() + if debug: + print(f"template_content = {template_content}") + template = Template(template_content) + if debug: + print(f"template = {template}") + + # Fill in the template options. + options = { + "hpc": args["hpc"], + "job_name": f"calcdb-{runid}", + "pbs_account": args["pbs_account"], + "conda_prefix": os.environ["CONDA_PREFIX"], + "kaipyhome": os.environ["KAIPYHOME"], + "kaijuhome": os.environ["KAIJUHOME"], + "parintime": args["parintime"], + "calcdb_xml_file": calcdb_xml_file, + } + if debug: + print(f"options = {options}") + + # Render the template. + if verbose: + print("Rendering template.") + calcdb_pbs_script = f"calcdb-{runid}.pbs" + calcdb_pbs_content = template.render(options) + with open(calcdb_pbs_script, "w", encoding="utf-8") as f: + f.write(calcdb_pbs_content) + + # Return the name of the PBS script. + return calcdb_pbs_script + + +def create_pitmerge_pbs_script(runid: str, args: dict): + """Create the PBS script for pitmerge.py to stitch calcdb output. + + Create the PBS script for stitching together calcdb output using the + pitmerge.py script. + + Parameters + ---------- + runid : str + Run ID for MAGE results used in calcdb.x calculations. + args : dict + Dictionary of command-line and other options. + + Returns + ------- + stitching_pbs_script : str + Path to PBS script. + + Raises + ------ + None + """ + # Local convenience variables. + debug = args["debug"] + verbose = args["verbose"] + + # Read the PBS script template for pitmerge.py. + if verbose: + print("Reading pitmerge.py PBS template.") + with open(PITMERGE_PBS_TEMPLATE, "r", encoding="utf-8") as f: + template_content = f.read() + if debug: + print(f"template_content = {template_content}") + template = Template(template_content) + if debug: + print(f"template = {template}") + + # Fill in the template options. + options = { + "hpc": args["hpc"], + "job_name": f"pitmerge-{runid}", + "pbs_account": args["pbs_account"], + "conda_prefix": os.environ["CONDA_PREFIX"], + "kaipyhome": os.environ["KAIPYHOME"], + "kaijuhome": os.environ["KAIJUHOME"], + "runid": runid, + } + if debug: + print(f"options = {options}") + + # Render the template. + if verbose: + print("Rendering template.") + pitmerge_pbs_script = f"pitmerge-{runid}.pbs" + pitmerge_pbs_content = template.render(options) + with open(pitmerge_pbs_script, "w", encoding="utf-8") as f: + f.write(pitmerge_pbs_content) + + # Return the name of the PBS script. + return pitmerge_pbs_script + + +def create_submit_script( + runid: str, calcdb_pbs_script: str, pitmerge_pbs_script: str, + args: dict): + """Create the bash script to submit all of the PBS jobs. + + Create the bash script to submit all of the PBS jobs. The submit script + submits the job to run calcdb.x, then the job to run pitmerge.py on the + results (if needed), then the analysis job. Each job only runs if the + previous job completes successfully. + + Parameters + ---------- + runid : str + Run ID for MAGE results used in calcdb.x calculations. + calcdb_pbs_script : str + Path to calcdb.x PBS script. + pitmerge_pbs_script : str + Path to pitmerge.py PBS script (or None if not needed). + args : dict + Dictionary of command-line and other options. + + Returns + ------- + submit_script : str + Path to bash script to submit PBS jobs. + + Raises + ------ + None + """ + # Submit the scripts in dependency order. + submit_script = f"submit-{runid}.sh" + with open(submit_script, "w", encoding="utf-8") as f: + + # calcdb.x job + if args["parintime"] > 1: + cmd = (f"job_id=`qsub -J 1-{args['parintime']} " + f"{calcdb_pbs_script}`\n") + else: + cmd = f"job_id=`qsub {calcdb_pbs_script}`\n" + f.write(cmd) + cmd = "echo $job_id\n" + f.write(cmd) + cmd = "old_job_id=$job_id\n" + f.write(cmd) + + # pitmerge.py job (if needed). + if pitmerge_pbs_script is not None: + cmd = ( + "job_id=`qsub -W depend=afterok:$old_job_id " + f"{pitmerge_pbs_script}`\n" + ) + f.write(cmd) + cmd = "echo $job_id\n" + f.write(cmd) + cmd = "old_job_id=$job_id\n" + f.write(cmd) + + # Return the name of the PBS script. + return submit_script + + +def run_calcdb(args: dict): + """Compute ground delta-B values for a MAGE magnetosphere run. + + Compute ground delta-B values for a MAGE magnetosphere run. + + Parameters + ---------- + args: dict + Dictionary of command-line options. + + Returns + ------- + int 0 on success + + Raises + ------ + AssertionError + If an invalid argument is provided. + """ + # Set defaults for command-line options, then update with values passed + # from the caller. + local_args = copy.deepcopy(DEFAULT_ARGUMENTS) + local_args.update(args) + args = local_args + + # Local convenience variables. + debug = args["debug"] + verbose = args["verbose"] + + # Validate arguments. + assert len(args["calcdb"]) > 0 + assert args["dt"] > 0.0 + assert args["hpc"] in VALID_HPC + assert args["parintime"] > 0 + assert len(args["mage_results_path"]) > 0 + + # ------------------------------------------------------------------------ + + # Split the MAGE results path into a directory and a file. + (mage_results_dir, mage_results_file) = os.path.split( + args["mage_results_path"] + ) + if debug: + print(f"mage_results_dir = {mage_results_dir}") + print(f"mage_results_file = {mage_results_file}") + + # Move to the results directory. + if verbose: + print(f"Moving to MAGE results directory {mage_results_dir}.") + os.chdir(mage_results_dir) + + # Compute the runid from the file name. + if verbose: + print(f"Computing runid from MAGE results file {mage_results_file}") + runid = mage_filename_to_runid(mage_results_file) + if debug: + print(f"runid = {runid}") + + # Create the XML input file for calcdb.x. + if verbose: + print(f"Creating XML file for calcdb.x for runid {runid}.") + calcdb_xml_file = create_calcdb_xml_file(runid, args) + if debug: + print(f"calcdb_xml_file = {calcdb_xml_file}") + + # Create the PBS script to run calcdb.x. + if verbose: + print(f"Creating PBS script to run calcdb.x for run {runid}.") + calcdb_pbs_script = create_calcdb_pbs_script(runid, calcdb_xml_file, args) + if debug: + print(f"calcdb_pbs_script = {calcdb_pbs_script}") + + # Create the PBS script to stitch together the output from calcdb.x + # using pitmerge.py. This script is only needed if parintime > 1. + if args["parintime"] > 1: + if verbose: + print("Creating PBS script to stitch together the calcdb.x output " + f" for MAGE runid {runid}") + pitmerge_pbs_script = create_pitmerge_pbs_script(runid, args) + else: + pitmerge_pbs_script = None + if debug: + print(f"pitmerge_pbs_script = {pitmerge_pbs_script}") + + # Create the bash script to submit the PBS scripts in the proper order. + if verbose: + print("Creating bash script to submit the PBS jobs.") + submit_script = create_submit_script( + runid, calcdb_pbs_script, pitmerge_pbs_script, args + ) + if debug: + print(f"submit_script = {submit_script}") + + if verbose: + print(f"Please run {submit_script} (in the MAGE result directory) to " + "submit the PBS jobs to perform the MAGE-SuperMag comparison. " + "If you are not using PBS, then you can manually run the " + f"scripts individually in the order listed in {submit_script}.") + + # Return normally. + return 0 + + +def main(): + """Driver for command-line version of code.""" + # Set up the command-line parser. + parser = create_command_line_parser() + + # Parse the command-line arguments. + args = parser.parse_args() + if args.debug: + print(f"args = {args}") + + # Convert the arguments from Namespace to dict. + args = vars(args) + + # Pass the command-line arguments to the main function as a dict. + return_code = run_calcdb(args) + sys.exit(return_code) + + +if __name__ == "__main__": + main() diff --git a/kaipy/scripts/postproc/run_ground_deltaB_analysis.py b/kaipy/scripts/postproc/run_ground_deltaB_analysis.py new file mode 100755 index 0000000..6f71f74 --- /dev/null +++ b/kaipy/scripts/postproc/run_ground_deltaB_analysis.py @@ -0,0 +1,709 @@ +#!/usr/bin/env python + + +"""Analyze the ground delta-B values from a MAGE run. + +Perform an analysis of ground magnetic field perturbations computed +for a MAGE magnetosphere simulation. The computed values are compared +with measured data from SuperMag. Then a set of SuperMAGE ground +delta-B maps is created. This is all done by a linked set of PBS jobs +running the MAGE results directory. + +NOTE: Before running this script, the user must run the setupEnvironment.sh +scripts for kaipy and kaiju. + +NOTE: If the calcdb.x binary was built with a module set other than those +listed below, change the module set in the PBS scripts appropriately. + +Author +------ +Eric Winter (eric.winter@jhuapl.edu) + +""" + + +# Import standard modules. +import argparse +import copy +import os +import pathlib +import re +import shutil +import sys + +# Import 3rd-party modules. +from jinja2 import Template + +# Import project-specific modules. +from kaipy import kaiH5 +from kaipy import kaiTools + + +# Program constants and defaults + +# Program description. +DESCRIPTION = ( + "Compare MAGE ground delta-B to SuperMag measurements, and create\n" + "SuperMAGE analysis maps." +) + +# Default values for command-line arguments. +DEFAULT_ARGUMENTS = { + "calcdb": "calcdb.x", + "debug": False, + "dt": 60.0, + "hpc": "pleiades", + "parintime": 1, + "pbs_account": None, + "smuser": None, + "verbose": False, + "mage_results_path": None, +} + +# Valid values for command-line arguments. +VALID_HPC = ["derecho", "pleiades"] + +# Directory containing this script. +SCRIPT_DIRECTORY = pathlib.Path(__file__).parent.resolve() +TEMPLATE_DIRECTORY = os.path.join(SCRIPT_DIRECTORY, "templates") + +# Location of template XML file for calcdb.x. +CALCDB_XML_TEMPLATE = os.path.join(TEMPLATE_DIRECTORY, "calcdb-template.xml") + +# # Locations of template PBS file for running calcdb.x. +CALCDB_PBS_TEMPLATE = os.path.join(TEMPLATE_DIRECTORY, "calcdb-template.pbs") + +# Location of template PBS file for pitmerge.py. +PITMERGE_PBS_TEMPLATE = os.path.join( + TEMPLATE_DIRECTORY, "pitmerge-template.pbs" +) + +# Location of template PBS file for running the ground delta-B analysis. +GROUND_DELTAB_ANALYSIS_PBS_TEMPLATE = os.path.join( + TEMPLATE_DIRECTORY, "ground_deltaB_analysis-template.pbs" +) + + +def create_command_line_parser(): + """Create the command-line parser. + + Create the parser for the command line. + + Parameters + ---------- + None + + Returns + ------- + parser : argparse.ArgumentParser + Command-line parser for this script. + + Raises + ------ + None + """ + parser = argparse.ArgumentParser(description=DESCRIPTION) + parser.add_argument( + "--calcdb", type=str, default=DEFAULT_ARGUMENTS["calcdb"], + help="Path to calcdb.x binary (default: %(default)s)." + ) + parser.add_argument( + "--debug", "-d", default=DEFAULT_ARGUMENTS["debug"], + action="store_true", + help="Print debugging output (default: %(default)s)." + ) + parser.add_argument( + "--dt", type=float, default=DEFAULT_ARGUMENTS["dt"], + help="Time interval for delta-B computation (seconds) " + "(default: %(default)s)." + ) + parser.add_argument( + "--hpc", type=str, default=DEFAULT_ARGUMENTS["hpc"], + choices=VALID_HPC, + help="HPC system to run analysis (default: %(default)s)." + ) + parser.add_argument( + "--parintime", type=int, default=DEFAULT_ARGUMENTS["parintime"], + help="Split the calculation into this many parallel chunks of MAGE " + "simulation steps, one chunk per node (default: %(default)s)." + ) + parser.add_argument( + "--pbs_account", type=str, default=DEFAULT_ARGUMENTS["pbs_account"], + help="PBS account to use for job accounting (default: %(default)s)." + ) + parser.add_argument( + "--smuser", type=str, default=DEFAULT_ARGUMENTS["smuser"], + help="Account name for SuperMag database access " + "(default: %(default)s)." + ) + parser.add_argument( + "--verbose", "-v", default=DEFAULT_ARGUMENTS["verbose"], + action="store_true", + help="Print verbose output (default: %(default)s)." + ) + parser.add_argument( + "mage_results_path", type=str, + default=DEFAULT_ARGUMENTS["mage_results_path"], + help="Path to a result file for a MAGE magnetosphere run." + ) + return parser + + +def mage_filename_to_runid(filename: str): + """Parse the runid from a MAGE results file name. + + Parse the runid from a MAGE results file name. + + For a result file from an MPI run, the runid is all text before the + underscore before the set of 6 underscore-separated 4-digit sequences at + the end of the name and the terminal .gam.h5 string. + + For a result file from a serial run, the runid is the name of the file, + less the .gam.h5 extension. + + Parameters + ---------- + filename : str + Name of MAGE results file. + + Returns + ------- + runid : str + The MAGE runid for the file. + + Raises + ------ + None + """ + # Check to see if the result file is for an MPI run. If not, it must be + # for a serial run. + mpi_pattern = ( + r"^(.+)_(\d{4})_(\d{4})_(\d{4})_(\d{4})_(\d{4})_(\d{4})\.gam.h5$" + ) + serial_pattern = r"^(.+)\.gam.h5$" + mpi_re = re.compile(mpi_pattern) + serial_re = re.compile(serial_pattern) + m = mpi_re.match(filename) + if not m: + m = serial_re.match(filename) + runid = m.groups()[0] + return runid + + +def create_calcdb_xml_file(runid: str, args: dict): + """Create the XML input file for calcdb.x from a template. + + Create the XML input file for calcdb.x from a template. The file is + created in the current directory, which much contain the specified MAGE + results file. Note that the contents of the file depends on wheter + parintime is 1 (use a single multi-threaded job to run calcdb.x on all + simulation steps) or > 1 (use multiple jobs to split the calcdb.x work + into an array of multi-threaded jobs which run on multiple nodes in + parallel). + + Parameters + ---------- + runid : str + Run ID for MAGE results file to use in calcdb.x calculations. + args : dict + Dictionary of command-line options. + + Returns + ------- + xml_file : str + Name of XML file. + + Raises + ------ + AssertionError + If the MAGE result file contains no time steps. + """ + # Local convenience variables. + debug = args["debug"] + verbose = args["verbose"] + + # Fetch run information from the MAGE result file. + if verbose: + print(f"Fetching MAGE run information for runid {runid}.") + filename, isMPI, Ri, Rj, Rk = kaiTools.getRunInfo(".", runid) + if debug: + print(f"filename = {filename}") + print(f"isMPI = {isMPI}") + print(f"Ri = {Ri}") + print(f"Rj = {Rj}") + print(f"Rk = {Rk}") + + # Get the number of steps and the step IDs from the MAGE results file. + if verbose: + print(f"Counting time steps for run {runid}.") + nSteps, sIds = kaiH5.cntSteps(filename) + if debug: + print(f"nSteps = {nSteps}") + print(f"sIds = {sIds}") + assert nSteps > 0 + + # Find the time for the last step. + if verbose: + print(f"Finding time for last step for run {runid}.") + tFin = kaiH5.tStep(filename, sIds[-1], aID="time") + if debug: + print(f"tFin = {tFin}") + + # Read the template XML file. + if verbose: + print("Reading calcdb XML template.") + with open(CALCDB_XML_TEMPLATE, "r", encoding="utf-8") as f: + template_content = f.read() + if debug: + print(f"template_content = {template_content}") + template = Template(template_content) + if debug: + print(f"template = {template}") + + # Fill in the template options. + options = { + "runid": runid, + "dt": args["dt"], + "tFin": tFin, + "ebfile": runid, + "isMPI": isMPI, + "Ri": Ri, + "Rj": Rj, + "Rk": Rk, + "parintime": args["parintime"], + } + if debug: + print(f"options = {options}") + + # Render the template. + xml_file = f"calcdb-{runid}.xml" + if verbose: + print("Rendering template.") + if verbose: + print(f"Creating {xml_file}.") + xml_content = template.render(options) + with open(xml_file, "w", encoding="utf-8") as f: + f.write(xml_content) + + # Return the name of the XML file. + return xml_file + + +def create_calcdb_pbs_script(runid: str, args: dict): + """Create the PBS script to run calcdb.x from a template. + + Create the PBS script to run calcdb.x from a template. The PBS script is + created in the current directory, which must contain the MAGE results. The + script will set up and run calcdb.x on the specified set of MAGE results, + in the current directory. + + Parameters + ---------- + runid : str + Run ID for MAGE results to use in calcdb.x calculations. + args : dict + Dictionary of command-line and other options. + + Returns + ------- + calcdb_pbs_script : str + Path to PBS script. + + Raises + ------ + None + """ + # Local convenience variables. + debug = args["debug"] + verbose = args["verbose"] + + # Create the XML input file for calcdb.x. + if verbose: + print(f"Creating XML file for calcdb.x for runid {runid}.") + calcdb_xml_file = create_calcdb_xml_file(runid, args) + if debug: + print(f"calcdb_xml_file = {calcdb_xml_file}") + + # Copy the calcdb.x binary to the results directory, then make it + # executable. + if verbose: + print(f"Copying {args['calcdb']} to results directory.") + shutil.copyfile(args["calcdb"], "calcdb.x") + os.chmod("calcdb.x", 0o755) + + # Read the PBS script template for running calcdb.x. + if verbose: + print("Reading calcdb.x PBS template.") + with open(CALCDB_PBS_TEMPLATE, "r", encoding="utf-8") as f: + template_content = f.read() + if debug: + print(f"template_content = {template_content}") + template = Template(template_content) + if debug: + print(f"template = {template}") + + # Fill in the template options. + options = { + "hpc": args["hpc"], + "job_name": f"calcdb-{runid}", + "pbs_account": args["pbs_account"], + "conda_prefix": os.environ["CONDA_PREFIX"], + "kaipyhome": os.environ["KAIPYHOME"], + "kaijuhome": os.environ["KAIJUHOME"], + "parintime": args["parintime"], + "calcdb_xml_file": calcdb_xml_file, + } + if debug: + print(f"options = {options}") + + # Render the template. + if verbose: + print("Rendering template.") + calcdb_pbs_script = f"calcdb-{runid}.pbs" + calcdb_pbs_content = template.render(options) + with open(calcdb_pbs_script, "w", encoding="utf-8") as f: + f.write(calcdb_pbs_content) + + # Return the name of the PBS script. + return calcdb_pbs_script + + +def create_pitmerge_pbs_script(runid: str, args: dict): + """Create the PBS script for pitmerge.py to stitch calcdb output. + + Create the PBS script for stitching together calcdb output using the + pitmerge.py script. + + Parameters + ---------- + runid : str + Run ID for MAGE results used in calcdb.x calculations. + args : dict + Dictionary of command-line and other options. + + Returns + ------- + stitching_pbs_script : str + Path to PBS script. + + Raises + ------ + None + """ + # Local convenience variables. + debug = args["debug"] + verbose = args["verbose"] + + # Read the PBS script template for pitmerge.py. + if verbose: + print("Reading pitmerge.py PBS template.") + with open(PITMERGE_PBS_TEMPLATE, "r", encoding="utf-8") as f: + template_content = f.read() + if debug: + print(f"template_content = {template_content}") + template = Template(template_content) + if debug: + print(f"template = {template}") + + # Fill in the template options. + options = { + "hpc": args["hpc"], + "job_name": f"pitmerge-{runid}", + "pbs_account": args["pbs_account"], + "conda_prefix": os.environ["CONDA_PREFIX"], + "kaipyhome": os.environ["KAIPYHOME"], + "kaijuhome": os.environ["KAIJUHOME"], + "runid": runid, + } + if debug: + print(f"options = {options}") + + # Render the template. + if verbose: + print("Rendering template.") + pitmerge_pbs_script = f"pitmerge-{runid}.pbs" + pitmerge_pbs_content = template.render(options) + with open(pitmerge_pbs_script, "w", encoding="utf-8") as f: + f.write(pitmerge_pbs_content) + + # Return the name of the PBS script. + return pitmerge_pbs_script + + +def create_ground_deltab_analysis_pbs_script(runid: str, args: dict): + """Create the PBS script for the ground delta-B analysis from a template. + + Create the PBS script for the ground delta-B analysis from a template. + + Parameters + ---------- + runid : str + Run ID for MAGE results used in calcdb.x calculations. + args : dict + Dictionary of command-line and other options. + + Returns + ------- + ground_deltab_analysis_pbs_script : str + Path to PBS script. + + Raises + ------ + None + """ + # Local convenience variables. + debug = args["debug"] + verbose = args["verbose"] + + # Read the PBS script template for ground_deltab_analysis.py. + if verbose: + print("Reading ground delta-B analysis PBS template.") + with open(GROUND_DELTAB_ANALYSIS_PBS_TEMPLATE, "r", encoding="utf-8") as f: + template_content = f.read() + if debug: + print(f"template_content = {template_content}") + template = Template(template_content) + if debug: + print(f"template = {template}") + + # Fill in the template options. + options = { + "hpc": args["hpc"], + "job_name": f"ground_deltaB_analysis-{runid}", + "pbs_account": args["pbs_account"], + "conda_prefix": os.environ["CONDA_PREFIX"], + "kaipyhome": os.environ["KAIPYHOME"], + "kaijuhome": os.environ["KAIJUHOME"], + "runid": runid, + "smuser": args["smuser"], + "calcdb_results_path": f"{runid}.deltab.h5", + "mage_results_path": args["mage_results_path"], + } + if debug: + print(f"options = {options}") + + # Render the template. + if verbose: + print("Rendering template.") + ground_deltab_analysis_pbs_script = f"ground_deltab_analysis-{runid}.pbs" + xml_content = template.render(options) + with open(ground_deltab_analysis_pbs_script, "w", encoding="utf-8") as f: + f.write(xml_content) + + # Return the name of the PBS script. + return ground_deltab_analysis_pbs_script + + +def create_submit_script( + runid: str, calcdb_pbs_script: str, pitmerge_pbs_script: str, + ground_deltab_analysis_pbs_script: str, args: dict): + """Create the bash script to submit all of the PBS jobs. + + Create the bash script to submit all of the PBS jobs. The submit script + submits the job to run calcdb.x, then the job to run pitmerge.py on the + results (if needed), then the analysis job. Each job only runs if the + previous job completes successfully. + + Parameters + ---------- + runid : str + Run ID for MAGE results used in calcdb.x calculations. + calcdb_pbs_script : str + Path to calcdb.x PBS script. + pitmerge_pbs_script : str + Path to pitmerge.py PBS script (or None if not needed). + ground_deltab_analysis_pbs_script : str + Path to analysis PBS script. + args : dict + Dictionary of command-line and other options. + + Returns + ------- + submit_script : str + Path to bash script to submit PBS jobs. + + Raises + ------ + None + """ + # Submit the scripts in dependency order. + submit_script = f"submit-{runid}.sh" + with open(submit_script, "w", encoding="utf-8") as f: + + # calcdb.x job + if args["parintime"] > 1: + cmd = (f"job_id=`qsub -J 1-{args['parintime']} " + f"{calcdb_pbs_script}`\n") + else: + cmd = f"job_id=`qsub {calcdb_pbs_script}`\n" + f.write(cmd) + cmd = "echo $job_id\n" + f.write(cmd) + cmd = "old_job_id=$job_id\n" + f.write(cmd) + + # pitmerge.py job (if needed). + if pitmerge_pbs_script is not None: + cmd = ( + "job_id=`qsub -W depend=afterok:$old_job_id " + f"{pitmerge_pbs_script}`\n" + ) + f.write(cmd) + cmd = "echo $job_id\n" + f.write(cmd) + cmd = "old_job_id=$job_id\n" + f.write(cmd) + + # Analysis job. + cmd = ( + "job_id=`qsub -W depend=afterok:$old_job_id " + f"{ground_deltab_analysis_pbs_script}`\n" + ) + f.write(cmd) + cmd = "echo $job_id\n" + f.write(cmd) + + # Return the name of the PBS script. + return submit_script + + +def run_ground_deltab_analysis(args: dict): + """Compute and analyze ground delta-B values for a MAGE run. + + Compute and analyze ground delta-B values for a MAGE run. + + Parameters + ---------- + args: dict + Dictionary of command-line options. + + Returns + ------- + int 0 on success + + Raises + ------ + AssertionError + If an invalid argument is provided. + """ + # Set defaults for command-line options, then update with values passed + # from the caller. + local_args = copy.deepcopy(DEFAULT_ARGUMENTS) + local_args.update(args) + args = local_args + + # Local convenience variables. + debug = args["debug"] + verbose = args["verbose"] + + # Validate arguments. + assert len(args["calcdb"]) > 0 + assert args["dt"] > 0.0 + assert args["hpc"] in VALID_HPC + assert args["parintime"] > 0 + assert len(args["mage_results_path"]) > 0 + + # ------------------------------------------------------------------------ + + # Split the MAGE results path into a directory and a file. + (mage_results_dir, mage_results_file) = os.path.split( + args["mage_results_path"] + ) + if debug: + print(f"mage_results_dir = {mage_results_dir}") + print(f"mage_results_file = {mage_results_file}") + + # Save the current directory. + start_directory = os.getcwd() + if debug: + print(f"start_directory = {start_directory}") + + # Move to the results directory. + if verbose: + print(f"Moving to MAGE results directory {mage_results_dir}.") + os.chdir(mage_results_dir) + + # Compute the runid from the file name. + if verbose: + print(f"Computing runid from MAGE results file {mage_results_file}") + runid = mage_filename_to_runid(mage_results_file) + if debug: + print(f"runid = {runid}") + + # Create the PBS script to run calcdb.x. + if verbose: + print("Creating PBS script to run calcdb.x to compute ground delta-B " + f"from MAGE results for runid {runid}.") + calcdb_pbs_script = create_calcdb_pbs_script(runid, args) + if debug: + print(f"calcdb_pbs_script = {calcdb_pbs_script}") + + # Create the PBS script to stitch together the output from calcdb.x + # using pitmerge.py. This script is only needed if parintime > 1. + if args["parintime"] > 1: + if verbose: + print("Creating PBS script to stitch together the calcdb.x output " + f" for MAGE runid {runid}") + pitmerge_pbs_script = create_pitmerge_pbs_script(runid, args) + else: + pitmerge_pbs_script = None + if debug: + print(f"pitmerge_pbs_script = {pitmerge_pbs_script}") + + # Create the PBS script to compare the calcdb.x results with SuperMag + # data. + if verbose: + print("Creating PBS script to run the ground delta-B analysis " + f"for MAGE runid {runid}.") + ground_deltab_analysis_pbs_script = ( + create_ground_deltab_analysis_pbs_script(runid, args) + ) + if debug: + print( + "ground_deltab_analysis_pbs_script = " + f"{ground_deltab_analysis_pbs_script}" + ) + + # Create the bash script to submit the PBS scripts in the proper order. + if verbose: + print("Creating bash script to submit the PBS jobs.") + submit_script = create_submit_script( + runid, calcdb_pbs_script, pitmerge_pbs_script, + ground_deltab_analysis_pbs_script, args + ) + if debug: + print(f"submit_script = {submit_script}") + + if verbose: + print(f"Please run {submit_script} (in the MAGE result directory) to " + "submit the PBS jobs to perform the MAGE-SuperMag comparison. " + "If you are not using PBS, then you can manually run the " + f"scripts individually in the order listed in {submit_script}.") + + # Move back to the start directory. + os.chdir(start_directory) + + # Return normally. + return 0 + + +def main(): + """Driver for command-line version of code.""" + # Set up the command-line parser. + parser = create_command_line_parser() + + # Parse the command-line arguments. + args = parser.parse_args() + if args.debug: + print(f"args = {args}") + + # Convert the arguments from Namespace to dict. + args = vars(args) + + # Pass the command-line arguments to the main function as a dict. + return_code = run_ground_deltab_analysis(args) + sys.exit(return_code) + + +if __name__ == "__main__": + main() diff --git a/kaipy/scripts/postproc/run_supermag_comparison.py b/kaipy/scripts/postproc/run_supermag_comparison.py deleted file mode 100755 index 39df78f..0000000 --- a/kaipy/scripts/postproc/run_supermag_comparison.py +++ /dev/null @@ -1,321 +0,0 @@ -#!/usr/bin/env python - -"""Run a SuperMag comparison for a MAGE magnetosphere run. - -Perform a comparison of ground magnetic field perturbations computed for a -MAGE magnetosphere simulation with measured data from SuperMag. - -Author ------- -Eric Winter (eric.winter@jhuapl.edu) -""" - - -# Import standard modules. -import argparse -import os -import subprocess -from xml.etree import ElementTree - -# Import 3rd-party modules. -import matplotlib as mpl -import matplotlib.pyplot as plt - -# Import project-specific modules. -import kaipy.supermage as sm - - -# Program constants and defaults - -# Program description. -DESCRIPTION = "Compare MAGE ground delta-B to SuperMag measurements." - -# Default SuperMag user name for queries. -DEFAULT_SUPERMAG_USER = "ewinter" - -# Location of template XML file. -XML_TEMPLATE = os.path.join( - os.environ["KAIJUHOME"], "scripts", "postproc", "calcdb.xml.template" -) - -# Name of XML file read by calcdb.x. -XML_FILENAME_TEMPLATE = "calcdb_RUNID.xml" - -# Number of microseconds in a second. -MICROSECONDS_PER_SECOND = 1e6 - -# Number of seconds in a day. -SECONDS_PER_DAY = 86400 - -# Location of SuperMag cache folder. -SUPERMAG_CACHE_FOLDER = os.path.join(os.environ["HOME"], "supermag") - - -def create_command_line_parser(): - """Create the command-line argument parser. - - Create the parser for command-line arguments. - - Parameters - ---------- - None - - Returns - ------- - parser : argparse.ArgumentParser - Command-line argument parser for this script. - """ - parser = argparse.ArgumentParser(description=DESCRIPTION) - parser.add_argument( - "-d", "--debug", action="store_true", default=False, - help="Print debugging output (default: %(default)s)." - ) - parser.add_argument( - "--mpixml", type=str, default=None, - help="If results from an MPI run, provide XML filename for run " - "(default: %(default)s)." - ) - parser.add_argument( - "--smuser", type=str, default=DEFAULT_SUPERMAG_USER, - help="SuperMag user ID to use for SuperMag queries " - "(default: %(default)s)." - ) - parser.add_argument( - "-v", "--verbose", action="store_true", default=False, - help="Print verbose output (default: %(default)s)." - ) - parser.add_argument( - "mage_results_path", - help="Path to a result file for a MAGE magnetosphere run." - ) - return parser - - -def get_mpi_decomposition(filename): - """Determine the MPI decomposition for the current MPI results. - - Determine the MPI decomposition for the current MPI results. - - The MPI decomposition is the triplet of values of the "N" attribute of - the elements (iPdir, jPdir, kPdir) in the element of the XML - file for the MAGE run. - - Parameters - ---------- - filename : str - Name of XML file. - - Returns - ------- - iPdir, jPdir, kPdir : int - Values of "N" attribute of elements (iPdir, jPdir, kPdir). - """ - # Parse the XML file. - iPdir = jPdir = kPdir = None - tree = ElementTree.parse(filename) - Kaiju_el = tree.getroot() - iPdir = int(Kaiju_el.findall("Gamera/iPdir")[0].attrib["N"]) - jPdir = int(Kaiju_el.findall("Gamera/jPdir")[0].attrib["N"]) - kPdir = int(Kaiju_el.findall("Gamera/kPdir")[0].attrib["N"]) - return iPdir, jPdir, kPdir - - -def filename_to_runid(filename): - """Parse the runid from a MAGE results file name. - - Parse the runid from a MAGE results file name. - - The runid is all text before the first period or underscore in the name. - - Parameters - ---------- - filename : str - Name of MAGE results file. - - Returns - ------- - runid : str - The MAGE runid for the file. - """ - parts = filename.split(".") - # parts = parts[0].split("_") - runid = parts[0] - return runid - - -def create_xml_file(runid, mpixml=None): - """Create the XML input file for calcdb.x from a template. - - Create the XML input file for calcdb.x from a template. - - Parameters - ---------- - runid : str - runid for MAGE results file. - mpixml : str, default None - If results are from an MPI run of MAGE, name of XML run file in results - directory. - - Returns - ------- - xml_file : str - Name of XML file. - """ - # Read the template file. - with open(XML_TEMPLATE) as t: - lines = t.readlines() - - # Process the template here. - # - # This should be done with a proper templating package. - lines[3] = lines[3].replace("RUNID", runid) - lines[5] = lines[5].replace("EBFILE", runid) - lines[5] = lines[5].replace("ISMPI", "true") - if mpixml is not None: - iPdir, jPdir, kPdir = get_mpi_decomposition(mpixml) - lines[6] = lines[6].replace("RI", str(iPdir)) - lines[6] = lines[6].replace("RJ", str(jPdir)) - lines[6] = lines[6].replace("RK", str(kPdir)) - else: - lines[6] = "\n" - # - - # Write out the processed XML. - xml_file = XML_FILENAME_TEMPLATE.replace("RUNID", runid) - with open(xml_file, "w") as f: - f.writelines(lines) - - return xml_file - - -def compute_ground_delta_B(runid, mpixml=None): - """Compute ground delta B values for a MAGE run. - - Compute ground delta B values for a MAGE run. The computation is done with - the program calcdb.x. - - Parameters - ---------- - runid : str - runid for MAGE results file. - mpixml : str, default None - If results are from an MPI run of MAGE, name of XML run file in results - directory. - - Returns - ------- - delta_B_file : str - Name of file containing calcdb.x results. - """ - # Create the XML file for calcdb.x from the template. - xml_file = create_xml_file(runid, mpixml) - - # Run the command to compute ground delta B values. - cmd = "calcdb.x" - args = [xml_file] - subprocess.run([cmd] + args) - - # Compute the name of the file containing the delta B values. - delta_B_file = runid + ".deltab.h5" - return delta_B_file - - -if __name__ == "__main__": - """Begin main program.""" - - # Set up the command-line parser. - parser = create_command_line_parser() - - # Parse the command-line arguments. - args = parser.parse_args() - debug = args.debug - mpixml = args.mpixml - smuser = args.smuser - verbose = args.verbose - mage_results_path = args.mage_results_path - if debug: - print("args = %s" % args) - - # Split the MAGE results path into a directory and a file. - (mage_results_dir, mage_results_file) = os.path.split(mage_results_path) - if debug: - print("mage_results_dir = %s" % mage_results_dir) - print("mage_results_file = %s" % mage_results_file) - - # Compute the runid from the file name. - runid = filename_to_runid(mage_results_file) - if debug: - print("runid = %s" % runid) - - # Move to the results directory. - if verbose: - print("Moving to results directory %s." % mage_results_dir) - os.chdir(mage_results_dir) - - # Compute the ground delta B values for this run. - if verbose: - print("Computing ground delta B values.") - delta_B_file = compute_ground_delta_B(runid, mpixml, mage_results_dir) - if debug: - print("delta_B_file = %s" % delta_B_file) - - # Read the delta B values. - SIM = sm.ReadSimData(delta_B_file) - if debug: - print("SIM = %s" % SIM) - - # Fetch the SuperMag indices for the desired time range. - - # Fetch the start time (as a datetime object) of simulation data. - start = SIM["td"][0] - if debug: - print("start = %s" % start) - - # Compute the duration of the simulated data, in seconds, then days. - duration = SIM["td"][-1] - SIM["td"][0] - duration_seconds = duration.seconds + duration.microseconds/MICROSECONDS_PER_SECOND - numofdays = duration_seconds/SECONDS_PER_DAY - if debug: - print("duration = %s" % duration) - print("duration_seconds = %s" % duration_seconds) - print("numofdays = %s" % numofdays) - - # Fetch the SuperMag indices for this time period. - if verbose: - print("Fetching SuperMag indices.") - SMI = sm.FetchSMIndices(smuser, start, numofdays) - if debug: - print("SMI = %s" % SMI) - - # Fetch the SuperMag data for this time period. - if verbose: - print("Fetching SuperMag data.") - SM = sm.FetchSMData(smuser, start, numofdays, - savefolder=SUPERMAG_CACHE_FOLDER) - if debug: - print("SM = %s" % SM) - - # Interpolate the simulated delta B to the measurement times from SuperMag. - if verbose: - print("Interpolating simulated data to SuperMag times.") - SMinterp = sm.InterpolateSimData(SIM, SM) - if debug: - print("SMinterp = %s" % SMinterp) - - - # Create the plots in memory. - mpl.use("Agg") - - # Make the indices plot. - if verbose: - print("Creating indices comparison plot.") - sm.MakeIndicesPlot(SMI, SMinterp, fignumber=1) - comparison_plot_file = runid + "_indices.png" - plt.savefig(comparison_plot_file) - - # Make the contour plots. - if verbose: - print("Creating contour plots.") - sm.MakeContourPlots(SM, SMinterp, maxx = 1000, fignumber=2) - contour_plot_file = runid + "_contours.png" - plt.savefig(contour_plot_file) diff --git a/kaipy/scripts/postproc/supermag_comparison.py b/kaipy/scripts/postproc/supermag_comparison.py new file mode 100755 index 0000000..f1947d9 --- /dev/null +++ b/kaipy/scripts/postproc/supermag_comparison.py @@ -0,0 +1,255 @@ +#!/usr/bin/env python + + +"""Create plots comparing ground delta B from a MAGE run to SuperMag data. + +Create plots comparing ground delta B from a MAGE run to SuperMag data. + +Author +------ +Eric Winter (eric.winter@jhuapl.edu) +""" + + +# Import standard modules. +import argparse +import copy +import os +import sys + +# Import 3rd-party modules. +import matplotlib as mpl +import matplotlib.pyplot as plt + +# Import project-specific modules. +import kaipy.supermage as sm + + +# Program constants and defaults + +# Program description. +DESCRIPTION = "Create MAGE-SuperMag comparison plots." + +# Default values for command-line arguments. +DEFAULT_ARGUMENTS = { + "debug": False, + "smuser": "", + "verbose": False, + "calcdb_results_path": None, +} + +# Number of seconds in a day. +SECONDS_PER_DAY = 86400 + +# Location of SuperMag cache folder. +SUPERMAG_CACHE_FOLDER = os.path.join(os.environ["HOME"], "supermag") + + +def create_command_line_parser(): + """Create the command-line parser. + + Create the parser for the command-line. + + Parameters + ---------- + None + + Returns + ------- + parser : argparse.ArgumentParser + Command-line argument parser for this script. + + Raises + ------ + None + """ + parser = argparse.ArgumentParser(description=DESCRIPTION) + parser.add_argument( + "--debug", "-d", default=DEFAULT_ARGUMENTS["debug"], + action="store_true", + help="Print debugging output (default: %(default)s)." + ) + parser.add_argument( + "--smuser", type=str, + default=DEFAULT_ARGUMENTS["smuser"], + help="SuperMag user ID to use for SuperMag queries " + "(default: %(default)s)." + ) + parser.add_argument( + "--verbose", "-v", default=DEFAULT_ARGUMENTS["verbose"], + action="store_true", + help="Print verbose output (default: %(default)s)." + ) + parser.add_argument( + "calcdb_results_path", + default=DEFAULT_ARGUMENTS["calcdb_results_path"], + help="Path to a (possibly merged) result file from calcdb.x." + ) + return parser + + +def supermag_comparison(args: dict): + """Create plots comparing MAGE ground delta-B to SuperMag data. + + Create plots comparing MAGE ground-delta B to SuperMag data. + + Parameters + ---------- + args: dict + Dictionary of command-line options. + + Returns + ------- + int 0 on success + + Raises + ------ + AssertionError + If an invalid argument is provided. + """ + # Set defaults for command-line options, then update with values passed + # from the caller. + local_args = copy.deepcopy(DEFAULT_ARGUMENTS) + local_args.update(args) + args = local_args + + # Local convenience variables. + debug = args["debug"] + verbose = args["verbose"] + + # Validate arguments. + assert len(args["calcdb_results_path"]) > 0 + + # ------------------------------------------------------------------------ + + # Split the calcdb results path into a directory and a file. + (calcdb_results_dir, calcdb_results_file) = os.path.split( + args["calcdb_results_path"] + ) + if calcdb_results_dir == "": + calcdb_results_dir = "." + if debug: + print(f"calcdb_results_dir = {calcdb_results_dir}") + print(f"calcdb_results_file = {calcdb_results_file}") + + # Save the current directory. + start_directory = os.getcwd() + if debug: + print(f"start_directory = {start_directory}") + + # Move to the results directory. + if verbose: + print(f"Moving to calcdb results directory {calcdb_results_dir}.") + os.chdir(calcdb_results_dir) + + # Extract the runid. + if verbose: + print("Computing runid from calcdb results file " + f"{calcdb_results_file}") + p = calcdb_results_file.index(".deltab.h5") + runid = calcdb_results_file[:p] + if debug: + print(f"runid = {runid}") + + # Read the delta B values. + if verbose: + print("Reading MAGE-derived ground delta-B values from " + f"{calcdb_results_file}.") + SIM = sm.ReadSimData(calcdb_results_file) + if debug: + print(f"SIM = {SIM}") + + # Fetch the start and stop time (as datetime objects) of simulation data. + date_start = SIM["td"][0] + date_stop = SIM["td"][-1] + if debug: + print(f"date_start = {date_start}") + print(f"date_stop = {date_stop}") + + # Compute the duration of the simulated data, in seconds, then days. + duration = date_stop - date_start + duration_seconds = duration.total_seconds() + duration_days = duration_seconds/SECONDS_PER_DAY + if debug: + print(f"duration = {duration}") + print(f"duration_seconds = {duration_seconds}") + print(f"duration_days = {duration_days}") + + # Fetch the SuperMag indices for this time period. + if verbose: + print(f"Fetching SuperMag indices for the period {date_start} to " + f"{date_stop}.") + SMI = sm.FetchSMIndices(args["smuser"], date_start, duration_days) + if debug: + print(f"SMI = {SMI}") + + # Fetch the SuperMag data for this time period. + if verbose: + print(f"Fetching SuperMag data for the period {date_start} to " + f"{date_stop}.") + SM = sm.FetchSMData(args["smuser"], date_start, duration_days, + savefolder=SUPERMAG_CACHE_FOLDER) + if debug: + print(f"SM = {SM}") + + # Abort if no data was found. + if len(SM["td"]) == 0: + raise TypeError("No SuperMag data found for requested time period, " + " aborting.") + + # Interpolate the simulated delta B to the measurement times from + # SuperMag. + if verbose: + print("Interpolating MAGE data to SuperMag times.") + SMinterp = sm.InterpolateSimData(SIM, SM) + if debug: + print("SMinterp = %s" % SMinterp) + + # ------------------------------------------------------------------------ + + # Create the plots in memory. + mpl.use("Agg") + + # ------------------------------------------------------------------------ + + # Make the indices plot. + if verbose: + print("Creating indices comparison plot.") + sm.MakeIndicesPlot(SMI, SMinterp, fignumber=1) + comparison_plot_file = "indices.png" + plt.savefig(comparison_plot_file) + + # Make the contour plots. + if verbose: + print("Creating contour plots.") + sm.MakeContourPlots(SM, SMinterp, maxx=1000, fignumber=2) + contour_plot_file = "contours.png" + plt.savefig(contour_plot_file) + + # Move back to the start directory. + os.chdir(start_directory) + + # Return normally. + return 0 + + +def main(): + """Driver for command-line version of code.""" + # Set up the command-line parser. + parser = create_command_line_parser() + + # Parse the command line. + args = parser.parse_args() + if args.debug: + print(f"args = {args}") + + # Convert the arguments from Namespace to dict. + args = vars(args) + + # Call the main program code. + return_code = supermag_comparison(args) + sys.exit(return_code) + + +if __name__ == "__main__": + main() diff --git a/kaipy/scripts/postproc/supermage_analysis.py b/kaipy/scripts/postproc/supermage_analysis.py new file mode 100755 index 0000000..1019271 --- /dev/null +++ b/kaipy/scripts/postproc/supermage_analysis.py @@ -0,0 +1,240 @@ +#!/usr/bin/env python + + +"""Create SuperMAGE plots for ground delta B from a MAGE run. + +Create SuperMAGE plots for ground delta B from a MAGE run. + +Author +------ +Eric Winter (eric.winter@jhuapl.edu) +""" + + +# Import standard modules. +import argparse +import copy +import os +import subprocess +import re +import sys + +# Import 3rd-party modules. + +# Import project-specific modules. + + +# Program constants and defaults + +# Program description. +DESCRIPTION = "Create SuperMAGE analysis plots." + +# Default values for command-line arguments. +DEFAULT_ARGUMENTS = { + "debug": False, + "verbose": False, + "mage_results_path": "", +} + + +def create_command_line_parser(): + """Create the command-line parser. + + Create the parser for the command-line. + + Parameters + ---------- + None + + Returns + ------- + parser : argparse.ArgumentParser + Command-line argument parser for this script. + + Raises + ------ + None + """ + parser = argparse.ArgumentParser(description=DESCRIPTION) + parser.add_argument( + "--debug", "-d", default=DEFAULT_ARGUMENTS["debug"], + action="store_true", + help="Print debugging output (default: %(default)s)." + ) + parser.add_argument( + "--verbose", "-v", default=DEFAULT_ARGUMENTS["verbose"], + action="store_true", + help="Print verbose output (default: %(default)s)." + ) + parser.add_argument( + "mage_results_path", + default=DEFAULT_ARGUMENTS["mage_results_path"], + help="Path to a MAGE result file." + ) + return parser + + +def mage_filename_to_runid(filename: str): + """Parse the runid from a MAGE results file name. + + Parse the runid from a MAGE results file name. + + For a result file from an MPI run, the runid is all text before the + underscore before the set of 6 underscore-separated 4-digit sequences at + the end of the name and the terminal .gam.h5 string. + + For a result file from a serial run, the runid is the name of the file, + less the .gam.h5 extension. + + Parameters + ---------- + filename : str + Name of MAGE results file. + + Returns + ------- + runid : str + The MAGE runid for the file. + + Raises + ------ + None + """ + # Check to see if the result file is for an MPI run. If not, it must be + # for a serial run. + mpi_pattern = ( + r"^(.+)_(\d{4})_(\d{4})_(\d{4})_(\d{4})_(\d{4})_(\d{4})\.gam.h5$" + ) + serial_pattern = r"^(.+)\.gam.h5$" + mpi_re = re.compile(mpi_pattern) + serial_re = re.compile(serial_pattern) + m = mpi_re.match(filename) + if not m: + m = serial_re.match(filename) + runid = m.groups()[0] + return runid + + +def create_dbpic_plots(runid: str, args: dict): + """Create the Mercator and polar plots of the dB values. + + Create the Mercator and polar plots of the dB values using dbpic.py. + + Parameters + ---------- + runid : str + Run ID for MAGE results. + args: dict + Dictionary of command-line options. + + Returns + ------- + None + + Raises + ------ + None + """ + # Run dbpic.py. + cmd = f"dbpic.py -d . -id {runid} --projection=both" + subprocess.run(cmd, shell=True, check=True) + return "dbpic.png" + + +def supermage_analysis(args: dict): + """Create SuperMAGE analysis plots for MAGE results. + + Create SuperMAGE analysis plots for MAGE results. + + Parameters + ---------- + args: dict + Dictionary of command-line options. + + Returns + ------- + int 0 on success + + Raises + ------ + None + """ + # Set defaults for command-line options, then update with values passed + # from the caller. + local_args = copy.deepcopy(DEFAULT_ARGUMENTS) + local_args.update(args) + args = local_args + + # Local convenience variables. + debug = args["debug"] + verbose = args["verbose"] + + # Validate arguments. + assert len(args["mage_results_path"]) > 0 + + # ------------------------------------------------------------------------ + + # Split the MAGE results path into a directory and a file. + (mage_results_dir, mage_results_file) = os.path.split( + args["mage_results_path"] + ) + if mage_results_dir == "": + mage_results_dir = "." + if debug: + print(f"mage_results_dir = {mage_results_dir}") + print(f"mage_results_file = {mage_results_file}") + + # Save the current directory. + start_directory = os.getcwd() + if debug: + print(f"start_directory = {start_directory}") + + # Move to the results directory. + if verbose: + print(f"Moving to MAGE results directory {mage_results_dir}.") + os.chdir(mage_results_dir) + + # Compute the runid from the file name. + if verbose: + print(f"Computing runid from MAGE results file {mage_results_file}") + runid = mage_filename_to_runid(mage_results_file) + if debug: + print(f"runid = {runid}") + + # ------------------------------------------------------------------------ + + # Make the dbpic.py plots (Mercator and polar projection). + if verbose: + print("Creating Mercator and polar plots of MAGE ground delta-B " + "values.") + dbpic_plot = create_dbpic_plots(runid, args) + if debug: + print(f"dbpic_plot = {dbpic_plot}") + + # Move back to the start directory. + os.chdir(start_directory) + + # Return normally. + return 0 + + +def main(): + """Driver for command-line version of code.""" + # Set up the command-line parser. + parser = create_command_line_parser() + + # Parse the command line. + args = parser.parse_args() + if args.debug: + print(f"args = {args}") + + # Convert the arguments from Namespace to dict. + args = vars(args) + + # Call the main program code. + return_code = supermage_analysis(args) + sys.exit(return_code) + + +if __name__ == "__main__": + main() diff --git a/kaipy/scripts/postproc/templates/calcdb-template.pbs b/kaipy/scripts/postproc/templates/calcdb-template.pbs new file mode 100644 index 0000000..001d239 --- /dev/null +++ b/kaipy/scripts/postproc/templates/calcdb-template.pbs @@ -0,0 +1,73 @@ +#!/bin/bash + +# This script runs calcdb.x to compute ground delta-B values from the +# results of a MAGE simulation. + +# This script was generated to run on {{ hpc }}. + +#PBS -N {{ job_name }} +{%- if hpc == "derecho" %} +#PBS -A {{ pbs_account }} +#PBS -q main +#PBS -l job_priority=economy +#PBS -l select=1:ncpus=128:ompthreads=128 +#PBS -l walltime=12:00:00 +{% elif hpc == "pleiades" %} +#PBS -q normal +#PBS -l select=1:ncpus=28:ompthreads=28:model=bro +#PBS -l walltime=08:00:00 +{% endif -%} +#PBS -j oe +#PBS -m abe + +echo "Job ${PBS_JOBID} started at `date` on `hostname` in directory `pwd`." + +source $HOME/.bashrc + +echo 'Loading modules.' +module --force purge +{%- if hpc == "derecho" %} +module load ncarenv/23.06 +module load craype/2.7.20 +module load intel/2023.0.0 +module load ncarcompilers/1.0.0 +module load cray-mpich/8.1.25 +module load hdf5-mpi/1.12.2 +{% elif hpc == "pleiades" %} +module load nas +module load pkgsrc/2022Q1-rome +module load nas +module load comp-intel/2020.4.304 +module load mpi-hpe/mpt.2.23 +module load hdf5/1.8.18_mpt +{% endif -%} +echo 'The currently loaded modules are:' +module list + +echo 'Loading python environment.' +conda activate {{ conda_prefix }} +echo "The current conda environment is ${CONDA_PREFIX}." + +echo 'Setting up kaipy environment.' +source {{ kaipyhome }}/kaipy/scripts/setupEnvironment.sh +echo "The kaipy software is located at ${KAIPYHOME}." + +echo 'Setting up MAGE environment.' +source {{ kaijuhome }}/scripts/setupEnvironment.sh +echo "The kaiju software is located at ${KAIJUHOME}." + +echo 'Setting environment variables.' +export OMP_NUM_THREADS={% if hpc == "derecho" %}128{%elif hpc == "pleiades"%}28{% endif %} +export KMP_STACKSIZE=128M +{% if parintime > 1 %}export JNUM=${PBS_ARRAY_INDEX:-0}{% endif %} +echo 'The active environment variables are:' +printenv + +# Compute the ground delta B values. +log_file="{% if parintime > 1 %}calcdb.out.${JNUM}{% else %}calcdb.out{% endif %}" +cmd="./calcdb.x {{ calcdb_xml_file }} {% if parintime > 1 %}${JNUM}{%- endif %} >& ${log_file}" +echo "calcdb.x run command is:" +echo $cmd +eval $cmd + +echo "Job ${PBS_JOBID} ended at `date` on `hostname` in directory `pwd`." diff --git a/kaipy/scripts/postproc/templates/calcdb-template.xml b/kaipy/scripts/postproc/templates/calcdb-template.xml new file mode 100644 index 0000000..12a3a87 --- /dev/null +++ b/kaipy/scripts/postproc/templates/calcdb-template.xml @@ -0,0 +1,10 @@ + + + + + + diff --git a/kaipy/scripts/postproc/templates/ground_deltaB_analysis-template.pbs b/kaipy/scripts/postproc/templates/ground_deltaB_analysis-template.pbs new file mode 100644 index 0000000..cebb144 --- /dev/null +++ b/kaipy/scripts/postproc/templates/ground_deltaB_analysis-template.pbs @@ -0,0 +1,76 @@ +#!/bin/bash + +# This script compares ground delta-B values computed from MAGE +# results with values from SuperMag, and creates a set of SuperMAGE +# comparison maps. + +# This script was generated to run on {{ hpc }}. + +#PBS -N {{ job_name }} +{%- if hpc == "derecho" %} +#PBS -A {{ pbs_account }} +#PBS -q main +#PBS -l job_priority=economy +#PBS -l select=1:ncpus=128 +#PBS -l walltime=12:00:00 +{% elif hpc == "pleiades" %} +#PBS -q normal +#PBS -l select=1:ncpus=28:model=bro +#PBS -l walltime=08:00:00 +{% endif -%} +#PBS -j oe +#PBS -m abe + +echo "Job ${PBS_JOBID} started at `date` on `hostname` in directory `pwd`." + +source $HOME/.bashrc + +echo 'Loading modules.' +module --force purge +{%- if hpc == "derecho" %} +module load ncarenv/23.06 +module load craype/2.7.20 +module load intel/2023.0.0 +module load ncarcompilers/1.0.0 +module load cray-mpich/8.1.25 +module load hdf5-mpi/1.12.2 +{% elif hpc == "pleiades" %} +module load nas +module load pkgsrc/2022Q1-rome +module load nas +module load comp-intel/2020.4.304 +module load mpi-hpe/mpt.2.23 +module load hdf5/1.8.18_mpt +{% endif -%} +echo 'The currently loaded modules are:' +module list + +echo 'Loading python environment.' +conda activate {{ conda_prefix }} +echo "The current conda environment is ${CONDA_PREFIX}." + +echo 'Setting up kaipy environment.' +source {{ kaipyhome }}/kaipy/scripts/setupEnvironment.sh +echo "The kaipy software is located at ${KAIPYHOME}." + +echo 'Setting up MAGE environment.' +source {{ kaijuhome }}/scripts/setupEnvironment.sh +echo "The kaiju software is located at ${KAIJUHOME}." + +echo 'Setting environment variables.' +echo 'The active environment variables are:' +printenv + +# Perform the SuperMag comparison. +cmd="supermag_comparison.py --verbose --smuser={{ smuser }} {{ calcdb_results_path }} >& supermag_comparison.out" +echo "Analysis command is:" +echo $cmd +eval $cmd + +# Perform the SuperMage analysis. +cmd="supermage_analysis.py --verbose {{ mage_results_path }} >& supermage_analysis.out" +echo "Analysis command is:" +echo $cmd +eval $cmd + +echo "Job $PBS_JOBID ended at `date` on `hostname` in directory `pwd`." diff --git a/kaipy/scripts/postproc/templates/pitmerge-template.pbs b/kaipy/scripts/postproc/templates/pitmerge-template.pbs new file mode 100644 index 0000000..c5358bb --- /dev/null +++ b/kaipy/scripts/postproc/templates/pitmerge-template.pbs @@ -0,0 +1,69 @@ +#!/bin/bash + +# This script runs pitmerge.py to merge ground delta-B values created +# by calcdb.x. + +# This script was generated to run on {{ hpc }}. + +#PBS -N {{ job_name }} +{%- if hpc == "derecho" %} +#PBS -A {{ pbs_account }} +#PBS -q main +#PBS -l job_priority=economy +#PBS -l select=1:ncpus=128 +#PBS -l walltime=12:00:00 +{% elif hpc == "pleiades" %} +#PBS -q normal +#PBS -l select=1:ncpus=28:model=bro +#PBS -l walltime=08:00:00 +{% endif -%} +#PBS -j oe +#PBS -m abe + +echo "Job ${PBS_JOBID} started at `date` on `hostname` in directory `pwd`." + +source $HOME/.bashrc + +echo 'Loading modules.' +module --force purge +{%- if hpc == "derecho" %} +module load ncarenv/23.06 +module load craype/2.7.20 +module load intel/2023.0.0 +module load ncarcompilers/1.0.0 +module load cray-mpich/8.1.25 +module load hdf5-mpi/1.12.2 +{% elif hpc == "pleiades" %} +module load nas +module load pkgsrc/2022Q1-rome +module load nas +module load comp-intel/2020.4.304 +module load mpi-hpe/mpt.2.23 +module load hdf5/1.8.18_mpt +{% endif -%} +echo 'The currently loaded modules are:' +module list + +echo 'Loading python environment.' +conda activate {{ conda_prefix }} +echo "The current conda environment is ${CONDA_PREFIX}." + +echo 'Setting up kaipy environment.' +source {{ kaipyhome }}/kaipy/scripts/setupEnvironment.sh +echo "The kaipy software is located at ${KAIPYHOME}." + +echo 'Setting up MAGE environment.' +source {{ kaijuhome }}/scripts/setupEnvironment.sh +echo "The kaiju software is located at ${KAIJUHOME}." + +echo 'Setting environment variables.' +echo 'The active environment variables are:' +printenv + +# Merge the ground delta B values. +cmd="pitmerge.py -runid {{ runid }} >& pitmerge.out" +echo "Stitching command is:" +echo $cmd +eval $cmd + +echo "Job ${PBS_JOBID} ended at `date` on `hostname` in directory `pwd`." diff --git a/kaipy/scripts/quicklook/dbpic.py b/kaipy/scripts/quicklook/dbpic.py index 650f29b..1daa2d7 100755 --- a/kaipy/scripts/quicklook/dbpic.py +++ b/kaipy/scripts/quicklook/dbpic.py @@ -14,8 +14,9 @@ Eric Winter (eric.winter@jhuapl.edu) # Import standard modules. import argparse -import datetime +import copy import os +import sys # Import 3rd-party modules. import cartopy.crs as ccrs @@ -36,181 +37,234 @@ import kaipy.kaiViz as kv # Program constants and defaults # Program description. -description = "Plot the ground magnetic field perturbations for a MAGE magnetosphere run." +DESCRIPTION = ( + "Plot the ground magnetic field perturbations for a MAGE magnetosphere " + "run." +) -# Default identifier for results to read. -default_runid = "msphere" +# Default values for command-line arguments. +DEFAULT_ARGUMENTS = { + "d": os.getcwd(), + "debug": False, + "id": "msphere", + "Jr": False, + "k0": 0, + "n": -1, + "projection": "mercator", + "spacecraft": None, + "verbose": False, +} -# Plot the last step by default. -default_step = -1 +# Default output filenames. +MERCATOR_PLOT_NAME = "qkdbpic.png" +POLAR_PLOT_NAME = "qkdbpole.png" -# Default vertical layer to plot. -default_k0 = 0 - -# Do not show anomalous currents by default. -default_Jr = False - -# Default output filename. -default_output_filename = "qkdbpic.png" - -# Size of figure in inches (width x height). -figSz = (12, 6) +# Size of figures in inches (width x height). +MERCATOR_FIGURE_SIZE = (12, 6) +POLAR_FIGURE_SIZE = (8, 8) # Color to use for magnetic footprint positions. -FOOTPRINT_COLOR = 'red' +FOOTPRINT_COLOR = "red" def create_command_line_parser(): - """Create the command-line argument parser. - - Create the parser for command-line arguments. + """Create the command-line parser. - Returns: - argparse.ArgumentParser: Command-line argument parser for this script. + Create the parser for the command-line. + + Parameters + ---------- + None + + Returns + ------- + parser : argparse.ArgumentParser + Command-line argument parser for this script. + + Raises + ------ + None """ - parser = argparse.ArgumentParser(description=description) + parser = argparse.ArgumentParser(description=DESCRIPTION) parser.add_argument( - "--debug", action="store_true", default=False, - help="Print debugging output (default: %(default)s)." - ) - parser.add_argument( - "-d", type=str, metavar="directory", default=os.getcwd(), + "-d", type=str, metavar="directory", default=DEFAULT_ARGUMENTS["d"], help="Directory containing data to read (default: %(default)s)" ) parser.add_argument( - "-n", type=int, metavar="step", default=default_step, - help="Time slice to plot (default: %(default)s)" + "--debug", default=DEFAULT_ARGUMENTS["debug"], + action="store_true", + help="Print debugging output (default: %(default)s)." ) parser.add_argument( - "-id", type=str, metavar="runid", default=default_runid, + "-id", type=str, metavar="runid", default=DEFAULT_ARGUMENTS["id"], help="Run ID of data (default: %(default)s)" ) parser.add_argument( - "-Jr", action="store_true", default=default_Jr, - help="Show radial component of anomalous current (default: %(default)s)." + "-Jr", action="store_true", default=DEFAULT_ARGUMENTS["Jr"], + help="Show radial component of anomalous current " + "(default: %(default)s)." ) parser.add_argument( - '-k0', type=int, metavar="layer", default=default_k0, + "-k0", type=int, metavar="layer", default=DEFAULT_ARGUMENTS["k0"], help="Vertical layer to plot (default: %(default)s)") parser.add_argument( - "--spacecraft", type=str, metavar="spacecraft", default=None, - help="Names of spacecraft to plot magnetic footprints, separated by commas (default: %(default)s)" + "-n", type=int, metavar="step", default=DEFAULT_ARGUMENTS["n"], + help="Time slice to plot (default: %(default)s)" ) parser.add_argument( - "-v", "--verbose", action="store_true", default=False, + "--projection", type=str, metavar="projection", + default=DEFAULT_ARGUMENTS["projection"], + help="Map projection to use for plots (mercator|polar|both)" + " (default: %(default)s) NOTE: 'mercator' is actually Plate-Carree" + ) + parser.add_argument( + "--quiver", "-q", action="store_true", default=False, + help="Use quiver for polar plot of dB (default: %(default)s)." + ) + parser.add_argument( + "--spacecraft", type=str, metavar="spacecraft", + default=DEFAULT_ARGUMENTS["spacecraft"], + help="Names of spacecraft to plot magnetic footprints, separated by " + "commas (default: %(default)s)" + ) + parser.add_argument( + "--verbose", "-v", default=DEFAULT_ARGUMENTS["verbose"], + action="store_true", help="Print verbose output (default: %(default)s)." ) return parser -if __name__ == "__main__": - """Plot the ground magnetic field perturbations.""" +def create_mercator_plot(args: dict): + """Make a Mercator plot of the ground magnetic field perturbations. - # Set up the command-line parser. - parser = create_command_line_parser() + Make a Mercator plot of the ground magnetic field perturbations. - # Parse the command-line arguments. - args = parser.parse_args() - debug = args.debug - fdir = args.d - runid = args.id - nStp = args.n - k0 = args.k0 - doJr = args.Jr - spacecraft = args.spacecraft - verbose = args.verbose + NOTE: This is actually a Plate-Carree projection, not Metcator. + + Parameters + ---------- + args: dict + Dictionary of command-line options. + + Returns + ------- + fig : matplotlib.Figure + Figure containing plot. + + Raises + ------ + None + """ + # Local convenience variables. + fdir = args["d"] + debug = args["debug"] + runid = args["id"] + doJr = args["Jr"] + k0 = args["k0"] + nStp = args["n"] + verbose = args["verbose"] + + # ------------------------------------------------------------------------ + + # Compute the tag of the file containing the ground magnetic field + # perturbations. + ftag = f"{runid}.deltab" if debug: - print("args = %s" % args) - - # Fetch constants. - bLin = dbViz.dbLin - bMag = dbViz.dbMag - jMag = dbViz.jMag - - # Compute the name of the file containing the ground magnetic field perturbations. - ftag = runid + ".deltab" - if debug: - print("ftag = %s" % ftag) + print(f"ftag = {ftag}") # Read the ground magnetic field perturbations. - fname = os.path.join(fdir, ftag + ".h5") - if debug: - print("fname = %s" % fname) + fname = os.path.join(fdir, f"{ftag}.h5") + if verbose: + print(f"Reading ground magnetic field perturbations from {fname}.") dbdata = gampp.GameraPipe(fdir, ftag) if debug: - print("dbdata = %s" % dbdata) + print(f"dbdata = {dbdata}") print("---") # Get the ID of the coordinate system, and the Earth radius. + if verbose: + print(f"Fetching coordinate system and Earth radius from {fname}.") CoordID, Re = dbViz.GetCoords(fname) - print("Found %s coordinate data ..." % CoordID) + print(f"Found {CoordID} coordinate data ...") if debug: - print("CoordID = %s" % CoordID) - print("Re = %s" % Re) + print(f"CoordID = {CoordID}") + print(f"Re = {Re}") # If the last simulation step was requested, get the step number. if nStp < 0: nStp = dbdata.sFin - print("Using Step %d" % nStp) + print(f"Using Step {nStp}") # Check the vertical level. + if verbose: + print("Checking vertical level.") Z0 = dbViz.CheckLevel(dbdata, k0, Re) if debug: - print("Z0 = %s" % Z0) + print(f"Z0 = {Z0}") # If currents were requested, read them. Otherwise, read the ground # magnetic field perturbations. - if (doJr): + if verbose: + print("Reading currents.") + if doJr: print("Reading Jr ...") Jr = dbdata.GetVar("dbJ", nStp, doVerb=False)[:, :, k0] Q = Jr else: dBn = dbdata.GetVar("dBn", nStp, doVerb=True)[:, :, k0] Q = dBn + if debug: + print(f"Q - {Q}") # Convert MJD to UT. + if verbose: + print(f"Reading MJD of step {nStp}.") MJD = kh5.tStep(fname, nStp, aID="MJD") if debug: - print("MJD = %s" % MJD) + print(f"MJD = {MJD}") utS = ktools.MJD2UT([MJD]) if debug: - print("utS = %s" % utS) - utDT= utS[0] + print(f"utS = {utS}") + utDT = utS[0] if debug: - print("utDT = %s" % utDT) + print(f"utDT = {utDT}") # Create the mapping grid. + if verbose: + print("Creating mapping grid.") crs = ccrs.PlateCarree() if debug: - print("ccrs = %s" % ccrs) + print(f"crs = {crs}") LatI, LonI, LatC, LonC = dbViz.GenUniformLL(dbdata, k0) if debug: - print("LatI = %s" % LatI) - print("LonI = %s" % LonI) - print("LatC = %s" % LatC) - print("LonC = %s" % LonC) + print(f"LatI = {LatI}") + print(f"LonI = {LonI}") + print(f"LatC = {LatC}") + print(f"LonC = {LonC}") # Fetch the color map. + if verbose: + print("Fetching color map.") cmap = kmaps.cmDiv if debug: - print("cmap = %s" % cmap) + print(f"cmap = {cmap}") # Determine color bar settings. - if (doJr): - vQ = kv.genNorm(jMag) + if verbose: + print("Determining color bar settings.") + if doJr: + vQ = kv.genNorm(dbViz.jMag) cbStr = "Anomalous current" else: - vQ = kv.genNorm(bMag, doSymLog=True, linP=bLin) + vQ = kv.genNorm(dbViz.dbMag, doSymLog=True, linP=dbViz.dbLin) cbStr = r"$\Delta B_N$ [nT]" if debug: - print("vQ = %s" % vQ) - print("cbStr = %s" % cbStr) - - # Create plot in memory. - mpl.use("Agg") + print(f"vQ = {vQ}") + print(f"cbStr = {cbStr}") # Create the figure to hold the plot. - fig = plt.figure(figsize=figSz) + fig = plt.figure(figsize=MERCATOR_FIGURE_SIZE) # Specify the grid for the subplots. gs = gridspec.GridSpec(3, 1, height_ratios=[20, 1.0, 1.0], hspace=0.025) @@ -223,11 +277,11 @@ if __name__ == "__main__": AxM.pcolormesh(LonI, LatI, Q, norm=vQ, cmap=cmap) # If requested, overlay the spacecraft magnetic footprints. - if spacecraft: - print("Overplotting magnetic footprints of %s." % spacecraft) + if args["spacecraft"]: + print(f"Overplotting magnetic footprints of {args['spacecraft']}") # Split the list into individual spacecraft names. - spacecraft = spacecraft.split(',') + spacecraft = args["spacecraft"].split(",") # Fetch the position of each footprint pair from CDAWeb. for sc in spacecraft: @@ -237,31 +291,31 @@ if __name__ == "__main__": sc, utS[0] ) if debug: - print("fp_nlat, fp_nlon = %s, %s" % (fp_nlat, fp_nlon)) + print(f"fp_nlat, fp_nlon = {fp_nlat}, {fp_nlon}") # Fetch the southern footprint position. fp_slat, fp_slon = cdaweb_utils.fetch_satellite_magnetic_southern_footprint_position( sc, utS[0] ) if debug: - print("fp_slat, fp_slon = %s, %s" % (fp_slat, fp_slon)) + print(f"fp_slat, fp_slon = {fp_slat}, {fp_slon}") # Plot a labelled dot at the location of each footprint. # Skip if no footprint position found. if fp_nlon is not None: - AxM.plot(fp_nlon, fp_nlat, 'o', c=FOOTPRINT_COLOR) + AxM.plot(fp_nlon, fp_nlat, "o", c=FOOTPRINT_COLOR) lon_nudge = 2.0 lat_nudge = 2.0 - AxM.text(fp_nlon + lon_nudge, fp_nlat + lat_nudge, sc + ' (N)') + AxM.text(fp_nlon + lon_nudge, fp_nlat + lat_nudge, sc + " (N)") else: - print("No northern footprint found for spacecraft %s." % sc) + print(f"No northern footprint found for spacecraft {sc}.") if fp_slon is not None: - AxM.plot(fp_slon, fp_slat, 'o', c=FOOTPRINT_COLOR) + AxM.plot(fp_slon, fp_slat, "o", c=FOOTPRINT_COLOR) lon_nudge = 2.0 lat_nudge = 2.0 - AxM.text(fp_slon + lon_nudge, fp_slat + lat_nudge, sc + ' (S)') + AxM.text(fp_slon + lon_nudge, fp_slat + lat_nudge, sc + " (S)") else: - print("No southern footprint found for spacecraft %s." % sc) + print(f"No southern footprint found for spacecraft {sc}.") # Make the colorbar. kv.genCB(AxCB, vQ, cbStr, cM=cmap) @@ -269,11 +323,254 @@ if __name__ == "__main__": # Add labels and other decorations. tStr = dbViz.GenTStr(AxM, fname, nStp) if debug: - print("tStr = %s" % tStr) + print(f"tStr = {tStr}") dbViz.DecorateDBAxis(AxM, crs, utDT) # Save the figure. - fOut = default_output_filename + path = os.path.join(args["d"], MERCATOR_PLOT_NAME) if debug: - print("fOut = %s" % fOut) - kv.savePic(fOut) + print(f"path = {path}") + kv.savePic(path, saveFigure=fig) + + +def create_polar_plot(args: dict): + """Make a polar plot of the ground magnetic field perturbations. + + Make a polar plot of the ground magnetic field perturbations. + + Parameters + ---------- + args: dict + Dictionary of command-line options. + + Returns + ------- + None + + Raises + ------ + None + """ + # Local convenience variables. + fdir = args["d"] + debug = args["debug"] + runid = args["id"] + doJr = args["Jr"] + k0 = args["k0"] + nStp = args["n"] + quiver = args["quiver"] + verbose = args["verbose"] + + # ------------------------------------------------------------------------ + + # Compute the tag of the file containing the ground magnetic field + # perturbations. + ftag = f"{runid}.deltab" + if debug: + print(f"ftag = {ftag}") + + # Read the ground magnetic field perturbations. + fname = os.path.join(fdir, f"{ftag}.h5") + if verbose: + print(f"Reading ground mgnetic field perturbations from {fname}.") + dbdata = gampp.GameraPipe(fdir, ftag) + if debug: + print(f"dbdata = {dbdata}") + print("---") + + # Get the ID of the coordinate system, and the Earth radius. + if verbose: + print(f"Fetching coordinate system and Earth radius from {fname}.") + CoordID, Re = dbViz.GetCoords(fname) + print(f"Found {CoordID} coordinate data ...") + if debug: + print(f"CoordID = {CoordID}") + print(f"Re = {Re}") + + # If the last simulation step was requested, get the step number. + if nStp < 0: + nStp = dbdata.sFin + print(f"Using Step {nStp}") + + # Check the vertical level. + if verbose: + print("Checking vertical level.") + Z0 = dbViz.CheckLevel(dbdata, k0, Re) + if debug: + print(f"Z0 = {Z0}") + + # If currents were requested, read them. Otherwise, read the ground + # magnetic field perturbations. + if verbose: + print("Reading currents.") + if doJr: + print("Reading Jr ...") + Jr = dbdata.GetVar("dbJ", nStp, doVerb=False)[:, :, k0] + Q = Jr + else: + dBn = dbdata.GetVar("dBn", nStp, doVerb=True)[:, :, k0] + Q = dBn + if debug: + print(f"Q - {Q}") + + # Convert MJD to UT. + if verbose: + print(f"Reading MJD of step {nStp}.") + MJD = kh5.tStep(fname, nStp, aID="MJD") + if debug: + print(f"MJD = {MJD}") + utS = ktools.MJD2UT([MJD]) + if debug: + print(f"utS = {utS}") + utDT = utS[0] + if debug: + print(f"utDT = {utDT}") + + # Create the mapping grids. + if verbose: + print("Creating mapping grids.") + crs = ccrs.PlateCarree() + toplate = ccrs.PlateCarree() + topole = ccrs.Orthographic(-90.0, 90.0) + if debug: + print(f"crs = {crs}") + print(f"toplate = {toplate}") + print(f"topole = {topole}") + LatI, LonI, LatC, LonC = dbViz.GenUniformLL(dbdata, k0) + if debug: + print(f"LatI = {LatI}") + print(f"LonI = {LonI}") + print(f"LatC = {LatC}") + print(f"LonC = {LonC}") + + # Fetch the color map. + if verbose: + print("Fetching color map.") + cmap = kmaps.cmDiv + if debug: + print(f"cmap = {cmap}") + + # Determine color bar settings. + if verbose: + print("Determining color bar settings.") + if doJr: + vQ = kv.genNorm(dbViz.jMag) + cbStr = "Anomalous current" + else: + vQ = kv.genNorm(dbViz.dbMag, doSymLog=True, linP=dbViz.dbLin) + cbStr = r"$\Delta B_N$ [nT]" + if debug: + print(f"vQ = {vQ}") + print(f"cbStr = {cbStr}") + + # Create the figure to hold the plot. + fig = plt.figure(figsize=POLAR_FIGURE_SIZE) + + # Specify the grid for the subplots. + gs = gridspec.GridSpec(3, 1, height_ratios=[20, 1.0, 1.0], hspace=0.025) + + # Create the subplots. + AxM = fig.add_subplot(gs[0, 0], projection=topole) + AxCB = fig.add_subplot(gs[-1, 0]) + + # Make the plot. + AxM.pcolormesh(LonI, LatI, Q, norm=vQ, cmap=cmap, transform=toplate) + + if quiver: + dB_p = dbdata.GetVar("dBp", nStp, doVerb=True)[:, :, k0] + dB_t = dbdata.GetVar("dBt", nStp, doVerb=True)[:, :, k0] + # U,V should be eastward/northward for cartopy + U = -dB_p + V = -dB_t + nSk = 2 + AxM.quiver(LonC[::nSk], LatC[::nSk], + U[::nSk, ::nSk], V[::nSk, ::nSk], + alpha=0.5, transform=toplate) + + # Make the colorbar. + kv.genCB(AxCB, vQ, cbStr, cM=cmap) + + # Add labels and other decorations. + tStr = dbViz.GenTStr(AxM, fname, nStp) + if debug: + print(f"tStr = {tStr}") + dbViz.DecorateDBAxis(AxM, crs, utDT) + + # Save the figure. + path = os.path.join(args["d"], POLAR_PLOT_NAME) + if debug: + print(f"path = {path}") + kv.savePic(path, saveFigure=fig) + + +def dbpic(args: dict): + """Plot the ground magnetic field perturbations from a MAGE run. + + Plot the ground magnetic field perturbations from a MAGE run. + + Parameters + ---------- + args: dict + Dictionary of command-line options. + + Returns + ------- + int 0 on success + + Raises + ------ + None + """ + # Set defaults for command-line options, then update with values passed + # from the caller. + local_args = copy.deepcopy(DEFAULT_ARGUMENTS) + local_args.update(args) + args = local_args + + # Local convenience variables. + # debug = args["debug"] + verbose = args["verbose"] + + # ------------------------------------------------------------------------ + + # Create plots in memory. + mpl.use("Agg") + + # If requested, create and save the Mercator plot. + if args["projection"] in ["mercator", "both"]: + if verbose: + print("Creating Mercator plot.") + create_mercator_plot(args) + + # If requested, create the polar plot. + if args["projection"] in ["polar", "both"]: + if verbose: + print("Creating polar plot.") + create_polar_plot(args) + + # ------------------------------------------------------------------------ + + # Return normally. + return 0 + + +def main(): + """Driver for command-line version of code.""" + # Set up the command-line parser. + parser = create_command_line_parser() + + # Parse the command line. + args = parser.parse_args() + if args.debug: + print(f"args = {args}") + + # Convert the arguments from Namespace to dict. + args = vars(args) + + # Call the main program code. + return_code = dbpic(args) + sys.exit(return_code) + + +if __name__ == "__main__": + main() diff --git a/kaipy/scripts/quicklook/gamerrVid.py b/kaipy/scripts/quicklook/gamerrVid.py index b699d8c..0d2b23e 100755 --- a/kaipy/scripts/quicklook/gamerrVid.py +++ b/kaipy/scripts/quicklook/gamerrVid.py @@ -113,6 +113,30 @@ def makeImage(i,gsph1,gsph2,tOut,doVerb,xyBds,fnList,oDir,errTimes,errListRel,er # end of sequential region AxB.set_xlabel('Time (min)') AxB.set_ylabel('Per-Cell Mean Relative Error',color=relColor) + # values and thresholds for background coloring + bgAlpha=0.2 + greenYellow=1.0e-6 + yellowRed=1.0e-2 + # coloring background messes with y-axis autoscaling + if AxB.get_ylim()[1] < greenYellow: + AxB.axhspan(AxB.get_ylim()[0],AxB.get_ylim()[1],color='green',alpha=bgAlpha) + elif AxB.get_ylim()[0] < greenYellow: + AxB.axhspan(AxB.get_ylim()[0],greenYellow,color='green',alpha=bgAlpha) + + if AxB.get_ylim()[0] > greenYellow and AxB.get_ylim()[1] < yellowRed: + AxB.axhspan(AxB.get_ylim()[0],AxB.get_ylim()[1],color='yellow',alpha=bgAlpha) + elif AxB.get_ylim()[0] < greenYellow and AxB.get_ylim()[1] > yellowRed: + AxB.axhspan(greenYellow,yellowRed,color='yellow',alpha=bgAlpha) + elif AxB.get_ylim()[0] < greenYellow: + AxB.axhspan(greenYellow,AxB.get_ylim()[1],color='yellow',alpha=bgAlpha) + elif AxB.get_ylim()[1] > yellowRed: + AxB.axhspan(AxB.get_ylim()[0],yellowRed,color='yellow',alpha=bgAlpha) + + if AxB.get_ylim()[0] > yellowRed: + AxB.axhspan(AxB.get_ylim()[0],AxB.get_ylim()[1],color='red',alpha=bgAlpha) + elif AxB.get_ylim()[1] > yellowRed: + AxB.axhspan(yellowRed,AxB.get_ylim()[1],color='red',alpha=bgAlpha) + AxB.tick_params(axis='y',which='both',colors=relColor,left=True,right=True,labelleft=True,labelright=False) AxB2.set_ylabel('Per-Cell Mean Absolute Error',color=absColor) #AxB2.yaxis.tick_right() diff --git a/kaipy/scripts/quicklook/heliomovie.py b/kaipy/scripts/quicklook/heliomovie.py index 8a77628..95f1707 100755 --- a/kaipy/scripts/quicklook/heliomovie.py +++ b/kaipy/scripts/quicklook/heliomovie.py @@ -107,7 +107,7 @@ valid_movie_formats = ("gif", "mp4") # Path to spacecraft metadata file. sc_metadata_path = os.path.join( - os.environ["KAIJUHOME"], "kaipy", "satcomp", "sc_helio.json" + os.environ["KAIPYHOME"], "kaipy", "satcomp", "sc_helio.json" ) # Figure sizes by plot type, (width, height) in inches. diff --git a/kaipy/scripts/quicklook/heliopic.py b/kaipy/scripts/quicklook/heliopic.py index 517d8dc..8a848c0 100755 --- a/kaipy/scripts/quicklook/heliopic.py +++ b/kaipy/scripts/quicklook/heliopic.py @@ -117,6 +117,9 @@ DEFAULT_PICTYPE = "pic1" # Colors to use for spacecraft position symbols. SPACECRAFT_COLORS = list(mpl.colors.TABLEAU_COLORS.keys()) +# Names of valid spacecraft +VALID_SPACECRAFT = "ACE,STEREO_A,STEREO_B,Parker_Solar_Probe,Solar_Orbiter" + def create_command_line_parser(): parser = argparse.ArgumentParser(description=DESCRIPTION) @@ -172,8 +175,9 @@ def create_command_line_parser(): ) parser.add_argument( "--spacecraft", type=str, metavar="spacecraft", default=None, - help="Names of spacecraft to plot positions, separated by commas " - "(default: %(default)s)" + help="Names of spacecraft to plot positions, separated by commas" + f" (valid spacecraft are {VALID_SPACECRAFT})" + " (default: %(default)s)" ) parser.add_argument( "--verbose", "-v", action="store_true", diff --git a/kaipy/scripts/quicklook/rcmDataProbe.py b/kaipy/scripts/quicklook/rcmDataProbe.py index fdb33ca..5463a0f 100755 --- a/kaipy/scripts/quicklook/rcmDataProbe.py +++ b/kaipy/scripts/quicklook/rcmDataProbe.py @@ -50,7 +50,7 @@ def create_command_line_parser(): numSamples = 6 varStr = "specFlux" - MainS = """Creates a plot of the differential flux for RCM ions or electrons in units of cm^-2 keV^-1 str^-1 + MainS = """Plots RCM data at a given equatorial x,y,z point. """ parser = argparse.ArgumentParser(description=MainS, formatter_class=RawTextHelpFormatter) parser.add_argument('-id',type=str,metavar="runid",default=ftag,help="RunID of data (default: %(default)s)") diff --git a/kaipy/scripts/quicklook/rcmPrecipSpecFlux.py b/kaipy/scripts/quicklook/rcmPrecipSpecFlux.py deleted file mode 100755 index 6721ba8..0000000 --- a/kaipy/scripts/quicklook/rcmPrecipSpecFlux.py +++ /dev/null @@ -1,171 +0,0 @@ -#!/usr/bin/env python -""" -this code computes number flux vs energy distribution of RCM diffuse precipitation in units of cm^-2 eV^-1 str^-1 s^-1 -""" -import matplotlib as mpl -mpl.use('Agg') -import matplotlib.pyplot as plt -import kaipy.kaiViz as kv -import matplotlib.gridspec as gridspec -import numpy as np -import kaipy.kaiH5 as kh5 - -from matplotlib import rcParams, cycler - -import argparse -from argparse import RawTextHelpFormatter -def create_command_line_parser(): - """Create the command-line argument parser. - - Create the parser for command-line arguments. - - Returns: - argparse.ArgumentParser: Command-line argument parser for this script. - """ - #Defaults - ftag = "msphere" - timeStr = "" - locStr = "-5,0,0" - numSamples = 6 - - MainS = """Creates a plot of the differential flux for RCM ions or electrons in units of cm^-2 keV^-1 str^-1 - """ - parser = argparse.ArgumentParser(description=MainS, formatter_class=RawTextHelpFormatter) - parser.add_argument('-id',type=str,metavar="runid",default=ftag,help="RunID of data (default: %(default)s)") - parser.add_argument('-n',type=int,metavar="numSamples",default=numSamples,help="Number of evenly-spaced time samples (default: %(default)s)") - parser.add_argument('-t',type=str,metavar="times",default=timeStr,help="Comma-separated times (in hours) to plot (example: 1,2,4,6). Ignores numSamples") - parser.add_argument('-l',type=str,metavar="loc",default=locStr,help="Comma-separated x,y,z values for equatorial location (default: %(default)s)") - parser.add_argument('-i',action='store_true',default=False,help="Flag to plot ions instead of electrons (default: %(default)s)") - - return parser - -if __name__ == "__main__": - #Defaults - cmap = plt.cm.plasma - # Set up the command-line parser. - parser = create_command_line_parser() - #Finalize parsing - args = parser.parse_args() - ftag = args.id - numSamples = args.n - timeStr = args.t - locStr = args.l - doIons = args.i - - fIn = ftag+".rcm.h5" - x0,y0,z0 = [float(x) for x in locStr.split(',')] - -#--Time stuff-- - #Get rcm data timesteps - nDataSteps,sIds = kh5.cntSteps(fIn) - dataTimes = kh5.getTs(fIn,sIds,"time")/3600 # [Hrs] - if timeStr == "": # If user didn't specify the times they wanted, do evenly spaced samples using numSamples - tStart = np.amin(np.abs(dataTimes)) - tEnd = np.amax(dataTimes) - nHrs = np.linspace(tStart, tEnd, num=numSamples, endpoint=True) - else: # Otherwise use user-specified times - nHrs = [float(t) for t in timeStr.split(',')] - NumStps = len(nHrs) - rcParams['axes.prop_cycle'] = cycler(color=cmap(np.linspace(0, 1, NumStps))) - - fSz = (14,7) - fig = plt.figure(figsize=fSz) - # conversion factor - massi = 1.67e-27 # mass of ions in kg - ev = 1.607e-19 # 1ev in J - nT2T = 1.e-9 - m2cm = 1.e2 - re = 6.380e6 # radius of earth in m - - Legs = [] - - eos = False #Flag for reaching end of steps - for n in range(NumStps): - - nStpIdx = np.abs(dataTimes - nHrs[n]).argmin() - nStp = nStpIdx + sIds.min() - legStr = "T = +%2.1f Hours"%(dataTimes[nStpIdx]) - if nStpIdx == nDataSteps: # Note that we're about to plot the last timestep and there's no need to plot any more afterwards - eos=True - elif nHrs[n] > dataTimes[-1]: - if not eos: # If this is the first time we've exceeded the end of the available timesteps, plot the last timestep anyways - print("%2.1f [hrs] out of time range (%2.1f, %2.1f), using last step time (%2.1f)"%(nHrs[n],dataTimes[0],dataTimes[-1],dataTimes[nStpIdx])) - eos = True - else: - print("%2.1f [hrs] out of time range (%2.1f, %2.1f), skipping."%(nHrs[n],dataTimes[0],dataTimes[-1])) - continue - - #Pull 3D data - deleeta = kh5.PullVar(fIn,"rcmdeleeta",nStp) - vm = kh5.PullVar(fIn,"rcmvm" ,nStp) - lamc = kh5.PullVar(fIn,"alamc" ,nStp) - bir = kh5.PullVar(fIn,"bir" ,nStp) - sini = kh5.PullVar(fIn,"sini" ,nStp) - xeq = kh5.PullVar(fIn,"rcmxmin" ,nStp) - yeq = kh5.PullVar(fIn,"rcmymin" ,nStp) - zeq = kh5.PullVar(fIn,"rcmzmin" ,nStp) - #Pull dt in s - dtCpl = kh5.tStep(fIn,nStp,aID="dtCpl",aDef=15.0) - - #Do grid calculations on first time - if (n==0): - #Set interfaces - Nlat,Nlon,Nk = deleeta.shape - kion = (lamc>0).argmax() - - if doIons: - kStart=kion - kEnd=Nk - else: - kStart=1 # Skip plasmasphere channel - kEnd=kion - - Nks = kEnd-kStart - ilamc = lamc[kStart:kEnd] - ilami = np.zeros(Nks+1) - for m in range(0,Nks-1): - nk = m+kStart - ilami[m+1] = 0.5*(lamc[nk]+lamc[nk+1]) - ilami[Nks] = lamc[-1] + 0.5*(lamc[-1]-lamc[-2]) - - #Ensure positive energies in case of electrons - ilami = np.abs(ilami) - ilamc = np.abs(ilamc) - - #Find nearest point (everytime) - dR = np.sqrt( (xeq-x0)**2.0 + (yeq-y0)**2.0 + (zeq-z0)**2.0 ) - i0,j0 = np.unravel_index(dR.argmin(), dR.shape) - - #Get energy bins in eV - Ki = vm[i0,j0]*ilami - Kc = vm[i0,j0]*ilamc - #Bin interval in eV - ekscl = vm[i0,j0]*np.diff(ilami) - - #Get differential flux in 1/(eV cm^2 str s) - ijEta = deleeta[i0,j0,kStart:kEnd]*bir[i0,j0]*sini[i0,j0]/ekscl/dtCpl*nT2T/(m2cm**2)/(4*np.pi) - Legs.append(legStr) - print(legStr) - print("\tMin/Mean/Max K = %f / %f / %f"%(Kc.min(),Kc.mean(),Kc.max())) - k0 =ijEta.argmax() - print("\tMax @ K = %f"%(Kc[k0])) - print("\tVM = %f"%(vm[i0,j0])) - plt.loglog(Kc,ijEta) - - sName = "Ions" if doIons else "Electrons" - titStr = "%s @ (x,y,z) = (%5.2f,%5.2f,%5.2f)"%(sName,x0,y0,z0) - Ax = plt.gca() - plt.legend(Legs,prop={'family': 'monospace'},loc='lower left') - - Ax.legend(Legs,loc='lower left') - Ax.set_xlabel("Energy [eV]") - Ax.set_ylabel("differential energy flux /cm^2/eV/str/s") - Ax.set_title(titStr) - Ax.grid() - - #Ax.set_ylim(1.0e+10,1.0e+19) - sTag = "_i" if doIons else "_e" - kv.savePic("qkRcmPrecipSpec%s.png"%sTag) - - #plt.show() - # diff --git a/kaipy/scripts/quicklook/rcmSpecFlux.py b/kaipy/scripts/quicklook/rcmSpecFlux.py deleted file mode 100755 index 2b95c1e..0000000 --- a/kaipy/scripts/quicklook/rcmSpecFlux.py +++ /dev/null @@ -1,173 +0,0 @@ -#!/usr/bin/env python -""" -this code computes differential energy flux in units of cm^-2 keV^-1 str^-1 -""" -import matplotlib as mpl -mpl.use('Agg') -import matplotlib.pyplot as plt -import kaipy.kaiViz as kv -import matplotlib.gridspec as gridspec -import numpy as np -import kaipy.kaiH5 as kh5 - -from matplotlib import rcParams, cycler - -import argparse -from argparse import RawTextHelpFormatter - -def create_command_line_parser(): - """Create the command-line argument parser. - - Create the parser for command-line arguments. - - Returns: - argparse.ArgumentParser: Command-line argument parser for this script. - """ - #Defaults - ftag = "msphere" - timeStr = "" - locStr = "-5,0,0" - numSamples = 6 - - MainS = """Creates a plot of the differential flux for RCM ions or electrons in units of cm^-2 keV^-1 str^-1 - """ - parser = argparse.ArgumentParser(description=MainS, formatter_class=RawTextHelpFormatter) - parser.add_argument('-id',type=str,metavar="runid",default=ftag,help="RunID of data (default: %(default)s)") - parser.add_argument('-n',type=int,metavar="numSamples",default=numSamples,help="Number of evenly-spaced time samples (default: %(default)s)") - parser.add_argument('-t',type=str,metavar="times",default=timeStr,help="Comma-separated times (in hours) to plot (example: 1,2,4,6). Ignores numSamples") - parser.add_argument('-l',type=str,metavar="loc",default=locStr,help="Comma-separated x,y,z values for equatorial location (default: %(default)s)") - parser.add_argument('-e',action='store_true',default=False,help="Flag to plot electrons instead of ions (default: %(default)s)") - - return parser - - -if __name__ == "__main__": - #Defaults - cmap = plt.cm.plasma - # Set up the command-line parser. - parser = create_command_line_parser() - #Finalize parsing - args = parser.parse_args() - ftag = args.id - numSamples = args.n - timeStr = args.t - locStr = args.l - doElectrons = args.e - - fIn = ftag+".rcm.h5" - x0,y0,z0 = [float(x) for x in locStr.split(',')] - -#--Time stuff-- - #Get rcm data timesteps - nDataSteps,sIds = kh5.cntSteps(fIn) - dataTimes = kh5.getTs(fIn,sIds,"time")/3600 # [Hrs] - if timeStr == "": # If user didn't specify the times they wanted, do evenly spaced samples using numSamples - tStart = np.amin(np.abs(dataTimes)) - tEnd = np.amax(dataTimes) - nHrs = np.linspace(tStart, tEnd, num=numSamples, endpoint=True) - else: # Otherwise use user-specified times - nHrs = [float(t) for t in timeStr.split(',')] - NumStps = len(nHrs) - - rcParams['axes.prop_cycle'] = cycler(color=cmap(np.linspace(0, 1, NumStps))) - - fSz = (14,7) - fig = plt.figure(figsize=fSz) - # conversion factor - massi = 1.67e-27 # mass of ions in kg - ev = 1.607e-19 # 1ev in J - nt = 1.0e-9 # nt - re = 6.380e6 # radius of earth in m - # this converts to units/cm^2/keV/str - conversion_factor = 1/np.pi/np.sqrt(8)*np.sqrt(ev/massi)*nt/re/1.0e1 - print('conversion factor=',conversion_factor) - - Legs = [] - - eos = False #Flag for reaching end of steps - for n in range(NumStps): - - nStpIdx = np.abs(dataTimes - nHrs[n]).argmin() - nStp = nStpIdx + sIds.min() - legStr = "T = +%2.1f Hours"%(dataTimes[nStpIdx]) - if nStpIdx == nDataSteps: # Note that we're about to plot the last timestep and there's no need to plot any more afterwards - eos=True - elif nHrs[n] > dataTimes[-1]: - if not eos: # If this is the first time we've exceeded the end of the available timesteps, plot the last timestep anyways - print("%2.1f [hrs] out of time range (%2.1f, %2.1f), using last step time (%2.1f)"%(nHrs[n],dataTimes[0],dataTimes[-1],dataTimes[nStpIdx])) - eos = True - else: - print("%2.1f [hrs] out of time range (%2.1f, %2.1f), skipping."%(nHrs[n],dataTimes[0],dataTimes[-1])) - continue - - #Pull 3D data - eeta = kh5.PullVar(fIn,"rcmeeta",nStp) - vm = kh5.PullVar(fIn,"rcmvm" ,nStp) - lamc = kh5.PullVar(fIn,"alamc" ,nStp) - xeq = kh5.PullVar(fIn,"rcmxmin" ,nStp) - yeq = kh5.PullVar(fIn,"rcmymin" ,nStp) - zeq = kh5.PullVar(fIn,"rcmzmin" ,nStp) - - #Do grid calculations on first time - if (n==0): - #Set interfaces - Nlat,Nlon,Nk = eeta.shape - kion = (lamc>0).argmax() - - if doElectrons: - kStart=1 # Skip plasmasphere channel - kEnd=kion - else: - kStart=kion - kEnd=Nk - - Nks = kEnd-kStart - ilamc = lamc[kStart:kEnd] - ilami = np.zeros(Nks+1) - for m in range(0,Nks-1): - nk = m+kStart - ilami[m+1] = 0.5*(lamc[nk]+lamc[nk+1]) - ilami[Nks] = lamc[-1] + 0.5*(lamc[-1]-lamc[-2]) - - #Ensure positive energies in case of electrons - ilami = np.abs(ilami) - ilamc = np.abs(ilamc) - - lamscl = np.diff(ilami)*np.sqrt(ilamc) - - #Find nearest point (everytime) - dR = np.sqrt( (xeq-x0)**2.0 + (yeq-y0)**2.0 + (zeq-z0)**2.0 ) - i0,j0 = np.unravel_index(dR.argmin(), dR.shape) - - #Get energy bins in keV - Ki = vm[i0,j0]*ilami*1.0e-3 - Kc = vm[i0,j0]*ilamc*1.0e-3 - - #ijEta = eeta[i0,j0,kion:]/lamscl - # 1e3 to convert back to eV - ijEta = conversion_factor*1.0e3*Kc*eeta[i0,j0,kStart:kEnd]/lamscl - Legs.append(legStr) - print(legStr) - print("\tMin/Mean/Max K = %f / %f / %f"%(Kc.min(),Kc.mean(),Kc.max())) - k0 =ijEta.argmax() - print("\tMax @ K = %f"%(Kc[k0])) - print("\tVM = %f"%(vm[i0,j0])) - plt.loglog(Kc,ijEta) - - sName = "Electrons" if doElectrons else "Ions" - titStr = "%s @ (x,y,z) = (%5.2f,%5.2f,%5.2f)"%(sName,x0,y0,z0) - Ax = plt.gca() - plt.legend(Legs,prop={'family': 'monospace'},loc='lower left') - - Ax.legend(Legs,loc='lower left') - Ax.set_xlabel("Energy [keV]") - Ax.set_ylabel("differential energy flux /cm^2/keV/str") - Ax.set_title(titStr) - Ax.grid() - - #Ax.set_ylim(1.0e+10,1.0e+19) - sTag = "_e" if doElectrons else "_i" - kv.savePic("qkrcmspec%s.png"%sTag) - - #plt.show() - # diff --git a/kaipy/scripts/setupEnvironment.sh b/kaipy/scripts/setupEnvironment.sh old mode 100755 new mode 100644 diff --git a/kaipy/supermage.py b/kaipy/supermage.py index 22b5eb8..b6b1d53 100644 --- a/kaipy/supermage.py +++ b/kaipy/supermage.py @@ -675,15 +675,17 @@ def MakeContourPlots(SM, SMinterp, maxx = 'default', fignumber = 1): plt.pcolormesh(SM['td'], hourbins, SMUbigreal, vmin = -1*maxx, vmax = maxx, cmap = cmapp, shading = 'auto') plt.xlim([SMinterp['td'][0], SMinterp['td'][-1]]) plt.grid(True) - plt.text(xxx, yyy, "Real SMU", horizontalalignment='center', verticalalignment='center', transform=ax1.transAxes, fontsize = qqq) + plt.text(xxx, yyy, "Real SMU", horizontalalignment='center', verticalalignment='center', transform=ax1.transAxes) plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%H:%M')) + plt.gca().xaxis.set_tick_params(labelsize=6) ax2 = plt.subplot(3, 2, 2) plt.pcolormesh(SM['td'], hourbins, SMLbigreal, vmin = -1*maxx, vmax = maxx, cmap = cmapp, shading = 'auto') plt.xlim([SMinterp['td'][0], SMinterp['td'][-1]]) plt.grid(True) - plt.text(xxx, yyy, "Real SML", horizontalalignment='center', verticalalignment='center', transform=ax2.transAxes, fontsize = qqq) + plt.text(xxx, yyy, "Real SML", horizontalalignment='center', verticalalignment='center', transform=ax2.transAxes) plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%H:%M')) + plt.gca().xaxis.set_tick_params(labelsize=6) ####################### @@ -691,16 +693,18 @@ def MakeContourPlots(SM, SMinterp, maxx = 'default', fignumber = 1): plt.pcolormesh(SMinterp['td'], hourbins, SMUbigsim, vmin = -1*maxx, vmax = maxx, cmap = cmapp, shading = 'auto') plt.xlim([SMinterp['td'][0], SMinterp['td'][-1]]) plt.grid(True) - plt.text(xxx, yyy, "Interpolated SMU", horizontalalignment='center', verticalalignment='center', transform=ax3.transAxes, fontsize = qqq) + plt.text(xxx, yyy, "Interpolated SMU", horizontalalignment='center', verticalalignment='center', transform=ax3.transAxes) plt.ylabel("MLT", fontsize = qqq) plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%H:%M')) + plt.gca().xaxis.set_tick_params(labelsize=6) ax4 = plt.subplot(3, 2, 4) plt.pcolormesh(SMinterp['td'], hourbins, SMLbigsim, vmin = -1*maxx, vmax = maxx, cmap = cmapp, shading = 'auto') plt.xlim([SMinterp['td'][0], SMinterp['td'][-1]]) plt.grid(True) - plt.text(xxx, yyy, "Interpolated SML", horizontalalignment='center', verticalalignment='center', transform=ax4.transAxes, fontsize = qqq) + plt.text(xxx, yyy, "Interpolated SML", horizontalalignment='center', verticalalignment='center', transform=ax4.transAxes) plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%H:%M')) + plt.gca().xaxis.set_tick_params(labelsize=6) ####################### @@ -708,15 +712,17 @@ def MakeContourPlots(SM, SMinterp, maxx = 'default', fignumber = 1): plt.pcolormesh(SMinterp['td'], hourbins, SMUbigsimR, vmin = -1*maxx, vmax = maxx, cmap = cmapp, shading = 'auto') plt.xlim([SMinterp['td'][0], SMinterp['td'][-1]]) plt.grid(True) - plt.text(xxx, yyy, "Super-SMU", horizontalalignment='center', verticalalignment='center', transform=ax5.transAxes, fontsize = qqq) + plt.text(xxx, yyy, "Super-SMU", horizontalalignment='center', verticalalignment='center', transform=ax5.transAxes) plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%H:%M')) + plt.gca().xaxis.set_tick_params(labelsize=6) ax6 = plt.subplot(3, 2, 6) im = plt.pcolormesh(SMinterp['td'], hourbins, SMLbigsimR, vmin = -1*maxx, vmax = maxx, cmap = cmapp, shading = 'auto') plt.xlim([SMinterp['td'][0], SMinterp['td'][-1]]) plt.grid(True) - plt.text(xxx, yyy, "Super-SML", horizontalalignment='center', verticalalignment='center', transform=ax6.transAxes, fontsize = qqq) + plt.text(xxx, yyy, "Super-SML", horizontalalignment='center', verticalalignment='center', transform=ax6.transAxes) plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%H:%M')) + plt.gca().xaxis.set_tick_params(labelsize=6) plt.tight_layout() @@ -745,10 +751,12 @@ def MakeIndicesPlot(SMI, SMinterp, fignumber=1): plt.plot(SMinterp['td'], SMinterp['SME'], 'r', label='interpolated') plt.plot(SMinterp['td'], SMinterp['SMU'], 'r') plt.plot(SMinterp['td'], SMinterp['SML'], 'r') + plt.gca().xaxis.set_tick_params(labelsize=6) plt.plot(SMinterp['td'], SMinterp['superSME'], 'g', label='super') plt.plot(SMinterp['td'], SMinterp['superSMU'], 'g') plt.plot(SMinterp['td'], SMinterp['superSML'], 'g') + plt.gca().xaxis.set_tick_params(labelsize=6) plt.plot(SMI['td'], SMI['SME'], 'b', label='real') plt.plot(SMI['td'], SMI['SMU'], 'b') @@ -756,7 +764,8 @@ def MakeIndicesPlot(SMI, SMinterp, fignumber=1): plt.grid(True) plt.legend() plt.xlim([SMinterp['td'][0], SMinterp['td'][-1]]) - plt.ylabel("SME/U/L", fontsize=20) + plt.ylabel("SME/U/L") + plt.gca().xaxis.set_tick_params(labelsize=6) ax2 = plt.subplot(312) plt.plot(SMinterp['td'], SMinterp['SMR'], 'r', label='interpolated') @@ -765,28 +774,33 @@ def MakeIndicesPlot(SMI, SMinterp, fignumber=1): plt.grid(True) plt.legend() plt.xlim([SMinterp['td'][0], SMinterp['td'][-1]]) - plt.ylabel("SMR", fontsize=20) + plt.ylabel("SMR") + plt.gca().xaxis.set_tick_params(labelsize=6) ax3 = plt.subplot(313) plt.plot(SMinterp['td'], SMinterp['SMRbins'][0], 'r', label='interpolated') plt.plot(SMinterp['td'], SMinterp['SMRbins'][1], 'r') plt.plot(SMinterp['td'], SMinterp['SMRbins'][2], 'r') plt.plot(SMinterp['td'], SMinterp['SMRbins'][3], 'r') + plt.gca().xaxis.set_tick_params(labelsize=6) plt.plot(SMinterp['td'], SMinterp['superSMRbins'][0], 'g', label='super') plt.plot(SMinterp['td'], SMinterp['superSMRbins'][1], 'g') plt.plot(SMinterp['td'], SMinterp['superSMRbins'][2], 'g') plt.plot(SMinterp['td'], SMinterp['superSMRbins'][3], 'g') + plt.gca().xaxis.set_tick_params(labelsize=6) plt.plot(SMI['td'], SMI['smr00'], 'b', label='real') plt.plot(SMI['td'], SMI['smr06'], 'b') plt.plot(SMI['td'], SMI['smr12'], 'b') plt.plot(SMI['td'], SMI['smr18'], 'b') + plt.gca().xaxis.set_tick_params(labelsize=6) plt.xlim([SMinterp['td'][0], SMinterp['td'][-1]]) plt.legend() plt.grid(True) - plt.ylabel("SMR 6-hour bins", fontsize=20) + plt.ylabel("SMR 6-hour bins") + plt.gca().xaxis.set_tick_params(labelsize=6) plt.show() def Z_Tensor_1D(resistivities, thicknesses, frequencies): diff --git a/pyproject.toml b/pyproject.toml index c4868b4..1c65d65 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,11 +7,11 @@ include-package-data = true [project] name = "kaipy" -version = "1.0.5" +version = "1.0.8" authors = [ {name = "Eric Winter", email = "eric.winter@jhuapl.edu"}, ] -description = "Python software for CGS MAGE and other Kaiu models" +description = "Python software for CGS MAGE and other Kaiju models" readme = "README.md" requires-python = ">=3.8,<=3.10" keywords = ["CGS", "MAGE", "KAIJU"] @@ -33,4 +33,4 @@ dependencies = [ "pytest", "spacepy", "sunpy", -] +] \ No newline at end of file diff --git a/setup.py b/setup.py index 69f958b..28ea62b 100644 --- a/setup.py +++ b/setup.py @@ -2,17 +2,18 @@ from setuptools import setup, find_packages setup( name='kaipy', - version='1.0.5', + version='1.0.8', description='Python software for CGS MAGE and other Kaiju models', author='Kaiju team', author_email='wiltbemj@ucar.edu', url='https://bitbucket.org/aplkaiju/kaipy/src/master/', packages=find_packages(), + include_package_data=True, + package_data={'kaipy': ['scripts/*', 'scripts/*/*']}, python_requires='>=3.8,<=3.10', install_requires=[ 'alive_progress', 'astropy', - # 'cartopy', 'cdasws', 'configparser', 'h5py',