def getPPF(dfTors, method='daily', res=10):
r"""
Calculate PPF density from tornado dataframe
Parameters
----------
dfTors : dataframe
method : 'total' or 'daily'
"""
# set up ~80km grid over CONUS
latgrid = np.arange(20, 55, res / 111)
longrid = np.arange(-130, -65, res / 111 / np.cos(35 * np.pi / 180))
interval = int(80 / res)
disk = circle_filter(interval)
dfTors['SPC_time'] = dfTors['UTC_time'] - timedelta(hours=12)
dfTors = dfTors.set_index(['SPC_time'])
groups = dfTors.groupby(pd.Grouper(freq="D"))
aggregate_grid = []
for group in groups:
slon, slat = group[1]['slon'].values, group[1]['slat'].values
elon, elat = group[1]['elon'].values, group[1]['elat'].values
torlons = [
i for x1, x2 in zip(slon, elon) for i in np.linspace(x1, x2, 10)
]
torlats = [
i for y1, y2 in zip(slat, elat) for i in np.linspace(y1, y2, 10)
]
# get grid count
grid, _, _ = np.histogram2d(torlats, torlons, bins=[latgrid, longrid])
grid = (grid > 0) * 1.0
grid = maximum_filter(grid, footprint=disk)
aggregate_grid.append(grid)
if method == 'daily':
grid = np.mean(aggregate_grid, axis=0)
PPF = gfilt(grid, sigma=1.5 * interval) * 100
if method == 'total':
grid = np.sum(aggregate_grid, axis=0)
PPF = gfilt((grid >= 1) * 1.0, sigma=1.5 * interval) * 100
return PPF, .5 * (longrid[:len(longrid) - 1] + longrid[1:]), .5 * (
latgrid[:len(latgrid) - 1] + latgrid[1:])
def subplot_generator(trans_num, taste_num):
for nrn_ind in trange(aligned_windows.shape[2]):
plot_dat = aligned_windows[:,:,nrn_ind]
# Take mean before filtering
mean_dat = np.mean(plot_dat,axis=1)
plot_dat = gfilt(plot_dat, sigma = 10)
plot_dat_long = np.reshape(plot_dat,(-1, *plot_dat.shape[2:]))
#plot_dat_long = zscore(plot_dat_long,axis=-1)
mean_dat = gfilt(mean_dat, sigma = 20)
mean_dat = mean_dat.swapaxes(0,1)
mean_dat = np.reshape(mean_dat, (-1, mean_dat.shape[-1]), order = 'F')
#mean_dat = zscore(mean_dat,axis=-1)
marks = np.arange(0,plot_dat_long.shape[0], plot_dat.shape[1])
trans_num = plot_dat.shape[2]
taste_num = plot_dat.shape[0]
fig = plt.figure(figsize = (5,10))
upper_axes = [fig.add_subplot(2,trans_num,num+1) for num in range(trans_num)]
lower_shape = (taste_num*2, trans_num)
inds = np.array(list(np.ndindex(lower_shape)))
wanted_inds = np.where(inds[:,0] >3)[0]
lower_axes = [fig.add_subplot(lower_shape[0],lower_shape[1],this_ind+1) \
for this_ind in wanted_inds]
for this_trans in range(len(upper_axes)):
upper_axes[this_trans].imshow(plot_dat_long[:,this_trans],
aspect='auto',cmap='viridis')
upper_axes[this_trans].hlines(marks, 0, plot_dat.shape[-1],
linewidth = 2, color = 'red', linestyle = 'dashed')
for this_plot in range(len(lower_axes)):
lower_axes[this_plot].plot(mean_dat[this_plot].T)
for this_ax in np.concatenate([[*upper_axes,*lower_axes]]):
this_ax.axes.get_yaxis().set_visible(False)
#this_ax.axes.get_xaxis().set_visible(False)
fig.savefig(os.path.join(plot_dir,f'nrn_{nrn_ind}_aligned'))
plt.close(fig)
def interpCart(self):
r"""
Interpolates polar storm-centered gridded fields into cartesian coordinates
"""
#Interpolate storm-centered recon data into gridded polar grid (rho, phi and gridded data)
grid_rho, grid_phi, grid_z_pol = self.interpPol()
#Calculate RMW
rmw = grid_rho[0, np.nanargmax(np.mean(grid_z_pol, axis=0))]
#Wraps around the data to ensure no cutoff around 0 degrees
grid_z_pol_wrap = np.concatenate([grid_z_pol] * 3)
#Radially smooth based on RMW - more smoothing farther out from RMW
grid_z_pol_final = np.array([gfilt(grid_z_pol_wrap,(6,3+abs(r-rmw)/10))[:,i] \
for i,r in enumerate(grid_rho[0,:])]).T[len(grid_phi):2*len(grid_phi)]
#Function for interpolating polar cartesian to coordinates
def pol2cart(rho, phi):
x = rho * np.cos(phi)
y = rho * np.sin(phi)
return (x, y)
#Interpolate the rho and phi gridded fields to a 1D cartesian list
pinterp_grid = [
pol2cart(i, j)
for i, j in zip(grid_rho.flatten(), grid_phi.flatten())
]
#Flatten the radially smoothed variable grid to match with the shape of pinterp_grid
pinterp_z = grid_z_pol_final.flatten()
#Setting up the grid in cartesian coordinate space, based on previously specified radial limit
#Grid resolution = 1 kilometer
grid_x, grid_y = np.meshgrid(np.linspace(-self.radlim,self.radlim,self.radlim*2+1),\
np.linspace(-self.radlim,self.radlim,self.radlim*2+1))
grid_z = griddata(pinterp_grid,
pinterp_z, (grid_x, grid_y),
method='linear')
#Return output grid
return grid_x, grid_y, grid_z
import os, re, numpy as np, math
from scipy.ndimage import gaussian_filter as gfilt
import happi
S = happi.Open(["./restart*"], verbose=False)
# COMPARE THE TOTAL NUMBER OF CREATED IONS
nion = S.Scalar.Ntot_ion(timesteps=0).getData()[0]
Validate("Number of ions at iteration 0", nion)
# COMPARE THE Ey FIELD
Ey = S.Field.Field0.Ey(timesteps=800).getData()[0]
Ey = gfilt(Ey, 6) # smoothing
Ey = Ey[50:200:4, 300::4] # zoom on the reflected laser
Validate("Ey field at iteration 800", Ey, 0.02)
format_dict1 = {'color': 'Greys', 'vs':[0,1], 'cbar_title': 'p-value'}
plt.figure()
plot_spearman(corr[0], NAMES, format_dict0, ax0, False)
#plot_spearman(corr[1], NAMES, format_dict1, ax1)
#%%
### exploring current/voltage maps ###
import matplotlib.pyplot as plt
import numpy as np
from home_dataTransforms import ug_to_mol
from scipy.ndimage import gaussian_filter as gfilt
samp = NBL3_2; scan_idx = 3
i = samp['XBIC_maps'][scan_idx][0,:,:-2]
v = samp['XBIV_maps'][scan_idx][0,:,:-2]
mol_cu = ug_to_mol(samp, 'XBIV_maps', scan_idx, elements, 1); molfilt_cu = gfilt(mol_cu,sigma=1)
mol_te = ug_to_mol(samp, 'XBIV_maps', scan_idx, elements, 2); molfilt_te = gfilt(mol_te,sigma=1)
mol_cute = molfilt_cu / molfilt_te
fig, (ax0, ax1) = plt.subplots(1,2)
ax0.imshow(v); ax1.imshow(mol_cute)
plt.figure()
plt.hexbin(i, v, mincnt=1)
#%%
# copy the color space in ImageJ, e.g. make a red colormap
from matplotlib.colors import LinearSegmentedColormap
colors = [(0, 0, 0), (0.5, 0, 0), (1, 0, 0)] # R -> G -> B
cmap_name = 'imgj_reds'
# Create the colormap
cm = LinearSegmentedColormap.from_list(cmap_name, colors, N=255)
def plot_swath(self,storm,Maps,varname,swathfunc,track_dict,radlim=200,\
domain="dynamic",ax=None,return_ax=False,prop={},map_prop={}):
#Set default properties
default_prop={'cmap':'category','levels':None,'left_title':'','right_title':'All storms','pcolor':True}
default_map_prop={'res':'m','land_color':'#FBF5EA','ocean_color':'#EDFBFF','linewidth':0.5,'linecolor':'k','figsize':(14,9),'dpi':200}
#Initialize plot
prop = self.add_prop(prop,default_prop)
map_prop = self.add_prop(map_prop,default_map_prop)
self.plot_init(ax,map_prop)
#Keep record of lat/lon coordinate extrema
max_lat = None
min_lat = None
max_lon = None
min_lon = None
#Retrieve recon data
lats = Maps['center_lat']
lons = Maps['center_lon']
#Retrieve storm data
storm_data = storm.dict
vmax = storm_data['vmax']
styp = storm_data['type']
sdate = storm_data['date']
#Add to coordinate extrema
if max_lat == None:
max_lat = max(lats)+2.5
else:
if max(lats) > max_lat: max_lat = max(lats)
if min_lat == None:
min_lat = min(lats)-2.5
else:
if min(lats) < min_lat: min_lat = min(lats)
if max_lon == None:
max_lon = max(lons)+2.5
else:
if max(lons) > max_lon: max_lon = max(lons)
if min_lon == None:
min_lon = min(lons)-2.5
else:
if min(lons) < min_lon: min_lon = min(lons)
bound_w,bound_e,bound_s,bound_n = self.dynamic_map_extent(min_lon,max_lon,min_lat,max_lat)
distproj = ccrs.LambertConformal()
out = distproj.transform_points(ccrs.PlateCarree(),np.array([bound_w,bound_w,bound_e,bound_e]),\
np.array([bound_s,bound_n,bound_s,bound_n]))
grid_res = 1*1e3 #m
xi = np.arange(int(min(out[:,0])/grid_res)*grid_res,int(max(out[:,0])/grid_res)*grid_res+grid_res,grid_res)
yi = np.arange(int(min(out[:,1])/grid_res)*grid_res,int(max(out[:,1])/grid_res)*grid_res+grid_res,grid_res)
xmgrid,ymgrid = np.meshgrid(xi,yi)
out = distproj.transform_points(ccrs.PlateCarree(),Maps['center_lon'],Maps['center_lat'])
cx = np.rint(gfilt(out[:,0],1)/grid_res)*grid_res
cy = np.rint(gfilt(out[:,1],1)/grid_res)*grid_res
aggregate_grid=np.ones(xmgrid.shape)*np.nan
def nanfunc(func,a,b):
c = np.concatenate([a[None],b[None]])
c = np.ma.array(c, mask=np.isnan(c))
d = func(c,axis=0)
e = d.data
e[d.mask] = np.nan
return e
for t,(x_center,y_center,var) in enumerate(zip(cx,cy,Maps['maps'])):
x_fromc = x_center+Maps['grid_x']*1e3
y_fromc = y_center+Maps['grid_y']*1e3
inrecon = np.where((xmgrid>=np.min(x_fromc)) & (xmgrid<=np.max(x_fromc)) & \
(ymgrid>=np.min(y_fromc)) & (ymgrid<=np.max(y_fromc)))
inmap = np.where((x_fromc>=np.min(xmgrid)) & (x_fromc<=np.max(xmgrid)) & \
(y_fromc>=np.min(ymgrid)) & (y_fromc<=np.max(ymgrid)))
aggregate_grid[inrecon] = nanfunc(swathfunc,aggregate_grid[inrecon],var[inmap])
if prop['levels'] is None:
prop['levels'] = (np.nanmin(aggregate_grid),np.nanmax(aggregate_grid))
cmap,clevs = get_cmap_levels(varname,prop['cmap'],prop['levels'])
out = self.proj.transform_points(distproj,xmgrid,ymgrid)
lons = out[:,:,0]
lats = out[:,:,1]
norm = mlib.colors.BoundaryNorm(clevs, cmap.N)
cbmap = self.ax.contourf(lons,lats,aggregate_grid,cmap=cmap,norm=norm,levels=clevs,transform=ccrs.PlateCarree())
#Storm-centered plot domain
if domain == "dynamic":
bound_w,bound_e,bound_s,bound_n = self.dynamic_map_extent(min_lon,max_lon,min_lat,max_lat)
self.ax.set_extent([bound_w,bound_e,bound_s,bound_n], crs=ccrs.PlateCarree())
#Pre-generated or custom domain
else:
bound_w,bound_e,bound_s,bound_n = self.set_projection(domain)
#Determine number of lat/lon lines to use for parallels & meridians
self.plot_lat_lon_lines([bound_w,bound_e,bound_s,bound_n])
#--------------------------------------------------------------------------------------
#Add left title
type_array = np.array(storm_data['type'])
idx = np.where((type_array == 'SD') | (type_array == 'SS') | (type_array == 'TD') | (type_array == 'TS') | (type_array == 'HU'))
tropical_vmax = np.array(storm_data['vmax'])[idx]
subtrop = classify_subtrop(np.array(storm_data['type']))
peak_idx = storm_data['vmax'].index(np.nanmax(tropical_vmax))
peak_basin = storm_data['wmo_basin'][peak_idx]
storm_type = get_storm_type(np.nanmax(tropical_vmax),subtrop,peak_basin)
dot = u"\u2022"
vartitle = get_recon_title(varname)
self.ax.set_title(f"{storm_type} {storm_data['name']}\n" + 'Recon: '+' '.join(vartitle),loc='left',fontsize=17,fontweight='bold')
#Add right title
#max_ppf = max(PPF)
start_date = dt.strftime(min(Maps['time']),'%H:%M UTC %d %b %Y')
end_date = dt.strftime(max(Maps['time']),'%H:%M UTC %d %b %Y')
self.ax.set_title(f'Start ... {start_date}\nEnd ... {end_date}',loc='right',fontsize=13)
#--------------------------------------------------------------------------------------
#Add legend
#Phantom legend
handles=[]
for _ in range(10):
handles.append(mlines.Line2D([], [], linestyle='-',label='',lw=0))
l = self.ax.legend(handles=handles,loc='upper left',fancybox=True,framealpha=0,fontsize=11.5)
plt.draw()
#Get the bbox
bb = l.legendPatch.get_bbox().inverse_transformed(self.fig.transFigure)
bb_ax = self.ax.get_position()
#Define colorbar axis
cax = self.fig.add_axes([bb.x0+bb.width, bb.y0-.05*bb.height, 0.015, bb.height])
# cbmap = mlib.cm.ScalarMappable(norm=norm, cmap=cmap)
cbar = self.fig.colorbar(cbmap,cax=cax,orientation='vertical',\
ticks=clevs)
cax.tick_params(labelsize=11.5)
cax.yaxis.set_ticks_position('left')
rect_offset = 0.0
if prop['cmap']=='category' and varname=='sfmr':
cax.yaxis.set_ticks(np.linspace(0,1,len(clevs)))
cax.yaxis.set_ticklabels(clevs)
cax2 = cax.twinx()
cax2.yaxis.set_ticks_position('right')
cax2.yaxis.set_ticks((np.linspace(0,1,len(clevs))[:-1]+np.linspace(0,1,len(clevs))[1:])*.5)
cax2.set_yticklabels(['TD','TS','Cat-1','Cat-2','Cat-3','Cat-4','Cat-5'],fontsize=11.5)
cax2.tick_params('both', length=0, width=0, which='major')
cax.yaxis.set_ticks_position('left')
rect_offset = 0.7
rectangle = mpatches.Rectangle((bb.x0,bb.y0-0.1*bb.height),(1.8+rect_offset)*bb.width,1.1*bb.height,\
fc = 'w',edgecolor = '0.8',alpha = 0.8,\
transform=self.fig.transFigure, zorder=2)
self.ax.add_patch(rectangle)
#Add plot credit
text = self.plot_credit()
self.add_credit(text)
def main():
parser = create_arg_parse()
args = parser.parse_args()
tiff = gdal.Open(args.density)
(upper_left_x, x_size, x_rotation, upper_left_y, y_rotation,
y_size) = tiff.GetGeoTransform()
proj = tiff.GetProjection()
red = tiff.GetRasterBand(1).ReadAsArray()
green = tiff.GetRasterBand(2).ReadAsArray()
blue = tiff.GetRasterBand(3).ReadAsArray()
allband = np.dstack((red, green, blue))
ndvi = np.zeros_like(red)
classes = [
{
'color': (0, 0, 0),
'class': 0.
},
{
'color': (0, 80, 255),
'class': 1.
},
{
'color': (0, 150, 255),
'class': 2.
},
{
'color': (0, 255, 255),
'class': 3.
},
{
'color': (0, 255, 150),
'class': 4.
},
{
'color': (0, 255, 80),
'class': 5.
},
{
'color': (0, 200, 0),
'class': 6.
},
{
'color': (150, 255, 0),
'class': 7.
},
{
'color': (255, 255, 0),
'class': 8.
},
{
'color': (255, 150, 0),
'class': 9.
},
{
'color': (255, 0, 0),
'class': 10.
},
]
colors = np.array([c['color'] for c in classes],
dtype=[('R', '<i4'), ('G', '<i4'), ('B', '<i4')])
colarr = colors.view(np.int).reshape(colors.shape + (-1, ))
classval = np.array([c['class'] for c in classes], dtype='d')
order = np.argsort(colors, axis=0, order=('R', 'G', 'B'))
refsort = colarr[order, :]
valorder = classval[order]
vals = allband.reshape(-1, 3)
ndvi = -np.empty(vals.shape[0], dtype="d")
nclass = len(refsort)
for iref in range(nclass):
print("Class: %s/%s" % (iref, nclass - 1))
vo = valorder[iref]
imatch = np.where((vals == refsort[iref, :]).all(axis=1))
ndvi[imatch] = vo
assert np.all(ndvi > -1.)
ndvi = ndvi.reshape(red.shape)
ndvi = gfilt(ndvi, sigma=3, order=0)
gtiff_driver = gdal.GetDriverByName('GTiff')
if gtiff_driver is None:
raise ValueError
fpath_out = 'ndvi_adj2.tif'
ds = gtiff_driver.Create(
fpath_out,
ndvi.shape[1],
ndvi.shape[0],
1,
gdal.GDT_Float32,
)
ds.SetGeoTransform(
[upper_left_x, x_size, x_rotation, upper_left_y, y_rotation, y_size])
ds.SetProjection(proj)
ds.GetRasterBand(1).WriteArray(ndvi)
ds.FlushCache()
def interpRecon(dfRecon, radlim=150):
from scipy.interpolate import griddata
# read in recon data
data = [
k
for i, j, k in zip(dfRecon['xdist'], dfRecon['ydist'], dfRecon['wspd'])
if np.nan not in [i, j, k]
]
path = [
(i, j)
for i, j, k in zip(dfRecon['xdist'], dfRecon['ydist'], dfRecon['wspd'])
if np.nan not in [i, j, k]
]
# polar
def cart2pol(x, y, offset=0):
rho = np.sqrt(x**2 + y**2)
phi = np.arctan2(y, x)
return (rho, phi + offset)
pol_path = [cart2pol(*p) for p in path]
pol_path_wrap = [cart2pol(*p,offset=-2*np.pi) for p in path]+pol_path+\
[cart2pol(*p,offset=2*np.pi) for p in path]
data_wrap = np.concatenate([data] * 3)
grid_rho, grid_phi = np.meshgrid(np.linspace(0, radlim, radlim * 2 + 1),
np.linspace(-np.pi, np.pi, 181))
grid_z_pol = griddata(pol_path_wrap,
data_wrap, (grid_rho, grid_phi),
method='linear')
rmw = grid_rho[0, np.nanargmax(np.mean(grid_z_pol, axis=0))]
print(rmw)
filleye = np.where((grid_rho < rmw) & (np.isnan(grid_z_pol)))
grid_z_pol[filleye] = 0
grid_z_pol_wrap = np.concatenate([grid_z_pol] * 3)
# smooth
grid_z_pol_final = np.array([gfilt(grid_z_pol_wrap,(5,2+r/10))[:,i] \
for i,r in enumerate(grid_rho[0,:])]).T[len(grid_phi):2*len(grid_phi)]
# back to cartesian
def pol2cart(rho, phi):
x = rho * np.cos(phi)
y = rho * np.sin(phi)
return (x, y)
pinterp_grid = [
pol2cart(i, j) for i, j in zip(grid_rho.flatten(), grid_phi.flatten())
]
pinterp_z = grid_z_pol_final.flatten()
grid_x, grid_y = np.meshgrid(np.linspace(-radlim, radlim, radlim * 2 + 1),
np.linspace(-radlim, radlim, radlim * 2 + 1))
grid_z = griddata(pinterp_grid,
pinterp_z, (grid_x, grid_y),
method='linear')
return grid_x, grid_y, grid_z
def getFrontPosition(timestep):
ne = S.Field.Field0.Rho_eon(
timesteps=timestep).getData()[0] # electron density
ne = gfilt(ne, 30) # smoothing
ne = np.abs(ne[:, 490:510].mean(axis=1)) # profile
return np.argmax(ne)
def getFrontPosition(timestep): ne = S.Field.Field0.Rho_eon(timesteps=timestep).getData()[0] # electron density ne = gfilt(ne, 30) # smoothing ne = np.abs(ne[:,490:510].mean(axis=1)) # profile return np.argmax(ne)
pdfssp = pickle.load(
open(cart_out + 'pdfs_refCLUS_last20_{}.p'.format(area), 'rb'))
for filt in [False, True]:
for last20 in [False, True]:
if last20:
ke = 'all_last20'
else:
ke = 'all'
xss = np.linspace(-2500., 2500., 201)
xi_grid, yi_grid = np.meshgrid(xss, xss)
zi = pdfssp[('hist', ke)]
if filt:
zi = gfilt(zi, 10)
zi = zi / np.sum(zi)
levzi = np.linspace(np.percentile(zi, 10), np.percentile(zi, 90),
9)
figs = []
levels = np.linspace(-1e-5, 1e-5, 10)
if filt: levels = np.linspace(-8e-6, 8e-6, 10)
for ssp in ['ssp126', 'ssp245', 'ssp370', 'ssp585']:
zi2 = pdfssp[(ssp, ke)]
if filt:
zi2 = gfilt(zi2, 10)