ATM 240 - Python, Fall 2022 "Cheat Sheet"

Cheat sheet

This page is http://www.atmos.albany.edu/facstaff/rfovell/ATM240/csheet.html.
Class web page is This page is http://www.atmos.albany.edu/facstaff/rfovell/ATM240/.

This is a living document and will be revised often.

Python version

import sys
print("You are using Python {}.{}.".format(sys.version_info.major, sys.version_info.minor))

Openers [I copy and paste these lines into new notebooks]

%matplotlib inline
 
import matplotlib.pyplot as plt
import matplotlib.cm as cm 
from matplotlib.cm import get_cmap
from matplotlib.colors import from_levels_and_colors
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.ticker import MultipleLocator
import numpy as np 
from numpy import genfromtxt 
import math

from pyproj import Proj 

import xarray as xr 

import cartopy.crs as ccrs 
import cartopy.feature as cfeat
import cartopy.geodesic
import shapely

import datetime
from datetime import datetime

import pandas as pd

# MetPy
from siphon.simplewebservice.wyoming import WyomingUpperAir   # upper air data repository
from metpy.units import units
import metpy.plots as plots
from metpy.plots import Hodograph
import metpy.calc as mpcalc
from mpl_toolkits.axes_grid1.inset_locator import inset_axes  # for hodograph inset

# for WRF
from netCDF4 import Dataset
from wrf import (getvar, to_np, get_cartopy, latlon_coords, vertcross, smooth2d,
                 cartopy_xlim, cartopy_ylim, interpline, CoordPair, destagger, ALL_TIMES)

# smoothing
import scipy.ndimage # for contour smoothing
from scipy.ndimage import gaussian_filter

from scipy import stats
from scipy.stats import norm, invgauss, skew, kurtosis
from scipy.interpolate import interpn # histogram
from scipy.interpolate import splrep, splev # smoothing normal curve


SMALL_SIZE = 12
MEDIUM_SIZE = 14
BIGGER_SIZE = 16

plt.rc('font', size=SMALL_SIZE)          # controls default text sizes
plt.rc('axes', titlesize=SMALL_SIZE)     # fontsize of the axes title
plt.rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels
plt.rc('xtick', labelsize=MEDIUM_SIZE)    # fontsize of the tick labels
plt.rc('ytick', labelsize=MEDIUM_SIZE)    # fontsize of the tick labels
plt.rc('legend', fontsize=MEDIUM_SIZE)    # legend fontsize
plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title

Some list stuff

• given a list 'mylist'

# forcibly replace oneth mylist element with new string
mylist[1] = 'goodbye'           

# insert 'everyone' after the twoth (i.e., third) element
mylist.insert(2,'everyone')

# remove the list element matching '2' .. NOT the twoth element!!
mylist.remove(2)

# append provided value to the end of the list
mylist.append(4.5)

• type of variable 'x'

type(x)

n=20
listofzeros = [0] * n
print(listofzeros)

Numpy and array stuff

• creating an empty array of specified length

# "empty" arrays can contain unwanted junk.  Use zeros instead.
cat = np.empty(8)

a = np.arange(6)          # creates 1D array with values 0,1,2,3,4, and 5
b = np.arange(0,100,20)   # creates 1D array with values 0,20,40,60,80

x = np.array([1.,2.,4.,5.,5.])
y = np.array([0.,2.,4.5,5.,6.])
qi_levels = np.array([0.1,0.2,0.4,0.8,1.6,3.2,6.4,12.8]) # g/kg

mylist = [[2,3,-5],[21,-2,1]]

a = np.array(mylist,dtype='d')  # double precision floating point
b = np.array(mylist,dtype='f')  # single precision floating point
c = np.array(mylist,dtype='i')  # integer (long int no longer exists)

a = np.zeros((3,2),dtype='d')   # a 3x2 array filled with double precision zeroes

print(" shape of array a = ",np.shape(a))

print(x[3:5])     # prints out array entries 4 and 5, but only those two
print(x[3:8])     # does not matter if ask for more entries than exist
print(x[-1])      # print last entry in array
print(x[-2])      # print second to last entry in array 

# key point is '(( ))'
all_year=np.concatenate((np.array(old_year),np.array(recent_year)))
year=np.concatenate((old_year,recent_year))

mean = np.mean(vals)  # mean 
max = np.max(vals)    # max 
min = np.min(vals)    # min 
std = np.std(vals)    # standard deviation  

# smallest to largest
idx = yyy.argsort()
yyy = yyy[idx]

# largest to smallest
idx = (-yyy).argsort()
yyy = yyy[idx]

lfr=np.around(lf,2)    # rounds numpy array 'lf' to 2 decimal places

ff = genfromtxt('/spare11/atm240/hrrr_swex_area_data.txt',skip_header=1) 
lat=ff[:,0]   # latitude  
lon=ff[:,1]   # longitude  

• Make 1D array 3D (clone)

# make pressure 3D array pres
# https://stackoverflow.com/questions/60044087/expand-and-copy-1d-numpy-array-to-3d

pres = np.tile(p_levels[:, np.newaxis, np.newaxis], (1, len(lat), len(lon)))
print(np.shape(pres))

compos_reflec=np.where(compos_reflec>5,compos_reflec,np.nan) # replaces reflec < 5 w/ nan so won't plot

• datetime, datetime64

# datetime variables look like 2022-01-01 00:00:00
# numpy datetime64 variables look like 2022-01-01T00:00:00.000000
# to convert a datetime to datetime64:

xday=np.datetime64(datetime(int(year),1,1,0,0))

Some Pandas stuff

# read a CSV dataset and let Pandas parse the date-like variables
df = pd.read_csv('WHO_flu_data_22Nov2020.csv', low_memory=False, parse_dates=True)

# read in the date-like variable, then convert it to datetime
date_temp = df['Start date']
date = pd.to_datetime(date_temp)

# read in space-delimited dataset without header, so provide column names separately
df = pd.read_csv('all_deviations_nysm_july2017.txt',delim_whitespace=True) # if NO HEADER need header=None
df.columns =['station', 'index', 'dwind', 'mean'] 

# Godawful way of reading in pandas month,day,year variables and constructing a datetime variable
year=df['year'].to_numpy().astype(int)
month=df['month'].to_numpy().astype(int)
day=df['day'].to_numpy().astype(int)

date=[]
for i in range(len(year)):
    date.append(datetime(year[i],month[i],day[i]))
#
print(date)

Metpy

# u = west-east wind speed
# v = north-south wind speed

fwdir = mpcalc.wind_direction(u,v) # calc wind direction wdir

# example with attaching units
# this gives temperature irma_tmp units of Kelvin, relative humidity irma_rh units of percent
# function result would be Celsius.  The final piece converts to Kelvin

irma_tdew = mpcalc.dewpoint_from_relative_humidity(irma_tmp * units.K, irma_rh * units.percent).to(units('kelvin'))

Loops and conditionals

• if, elif, else

if(mon == 'FEB'):
	monx = 2
elif(mon == 'MAR'):
	monx = 3
else:
	print(" not february or march")

if(fhr >= start and fhr <= end>):
    [do something]
#

ww = np.copy(ogust)
zz = np.copy(LT)

# or
w = ww[(zzz <= 6) | (zzz >= 18)]
z = zz[(zzz <= 6) | (zzz >= 18)]

# and
w = ww[(zzz >= 6) & (zzz <= 18)]
z = zz[(zzz >= 6) & (zzz <= 18)]

for i in range(5):      # this operates on 0, 1, 2, 3, and 4, but not 5
    print(i)
#

a = 1
while a < 10:
    print(a)
    a = a + 1

Matplotlib

• Basic plots

fig = plt.figure(figsize=(8,8))                              # plot width, plot height
plt.plot(x,y,c='lightgrey',linewidth=2,label='mylabel')      # line plot
plt.scatter(x,y,c='k',edgecolors='None',s=20)                # scatterplot
plt.bar(x,y,color='lightgrey',linewidth=2,label='mylabel')   # bar chart
plt.legend(loc='upper left')                                 # legend
plt.grid()                                                   # puts a grid on the plot
plt.grid(axis='both', which='both')   # vert grid too        # on multi-y plots, have to do it before cloning x-axis
plt.errorbar(theta_mean-theta2_mean,z_mean,xerr=theta_diff_std,c='grey',linewidth=0.5,zorder=1)

levels=np.arange(968.,1044.,4.)                              # max level not included!                           
cs = ax.contour(lon,lat,slp,colors='k', levels=levels, linewidths=1, linestyles=':',
                transform=ccrs.PlateCarree())                # transform presumed map-projected data
fmt = '%r '                                                  # format for contour labels
fmt = '%1.1f'												 # float, one decimal place
ax.clabel(cs, cs.levels[::2], inline=True, fmt=fmt, fontsize=10)  # [::2] only every other contour

cmap = plt.cm.seismic                                        # color table
norm = plt.Normalize(-5.,5)                                  # values bounds
cm = ax.pcolormesh(lon,lat,snowd,shading='auto',cmap=cmap, norm=norm, 
                   transform=ccrs.PlateCarree())             # transform presumed map-projected data
plt.colorbar(cm, orientation="vertical",shrink=0.5, label='label')  # add colorbar (label optional)

# for orientation="horizontal", use pad=0.01 to reduce gap between plot bottom and colorbar

#####################################################################################
# putting a color bar on only ONE plot without changing plot size
# used in plot_HRRR_NAM_OBS_maxgusts_BOULDER_V2.ipynb
#####################################################################################

# first panel
ax = fig.add_subplot(1,2,1,projection=proj) 

colors = ['blue', 'white', 'red']
cmap = LinearSegmentedColormap.from_list('name', colors)
norm=plt.Normalize(0,50)
cm = ax.pcolormesh(lon_hrrr,lat_hrrr,maxgust_hrrr,shading='auto',alpha=0.5,cmap=cmap, norm=norm, transform=ccrs.PlateCarree())

# second panel, using same cmap and norm
ax = fig.add_subplot(1,2,2,projection=proj) 
cm = ax.pcolormesh(lon_nam,lat_nam,maxgust_nam,shading='auto',alpha=0.5,cmap=cmap, norm=norm, transform=ccrs.PlateCarree())

# now add the colorbar to a NEW axis

cb_ax = fig.add_axes([.91,.124,.02,.754]) # third number controls WIDTH of colorbar
fig.colorbar(cm,orientation='vertical',cax=cb_ax,label="maximum gust (m/s)")#,fraction=0.1)

# colormap
cmap = ListedColormap(['green','green','green','green','green','green','brown','brown','green','yellow','yellow','yellow','lightgreen','red','yellow','white','white',
         'blue','blue','blue','blue'])
# reverses the colormap
cmap_r = ListedColormap(cmap.colors[::-1])
cm = ax.pcolormesh(lonwrf,latwrf,greenfrac[3,:,:],shading='flat',norm=norm,cmap=cmap_r, transform=ccrs.PlateCarree())

# this bit of magic forces the plot to be square, appropriate because we have a square domain and no map projection
plt.gca().set_aspect('equal', adjustable='box')

# but this works too
# make sure plot is square - appropriate as we have square domain and no map projection
plt.axis('square')

# dashed 1:1 line.  I know the plot extent is 0-30 in x and y directions
plt.plot([0, 40], [0, 40], c='lightgrey', lw=2, linestyle='dashed')

import matplotlib.pyplot as plt
import numpy as np
%matplotlib inline
 
fig = plt.figure(figsize=(10,10))
ax = plt.axes()
 
# Create linspace for color filling
x = np.linspace(-100, 100, 100)
 
# Fill color above/below y=x line
plt.fill_between(x, x, 100, facecolor='r', alpha=0.07, zorder=0)
plt.fill_between(x, x, -100, facecolor='b', alpha=0.07, zorder=0)
 
# Plot dashed lines every 2 [units]
for k in range(-10,10,2):
    plt.plot(x, x+k, c='0.6', ls='--', lw=1, zorder=1)
 
# Create thick dashed lines at x=0 and y=0
ax.axvline(0, color='k', ls='--', lw=1.5, zorder=2)
ax.axhline(0, color='k', ls='--', lw=1.5, zorder=2)
 
# Set axis limits, labels, and tick labels
plt.xlim(-10,10)
plt.ylim(-10,10)
plt.xlabel('X Variable', fontsize=14)
plt.ylabel('Y Variable', fontsize=14)
plt.xticks(np.arange(-10,10,2), fontsize=12)
plt.yticks(np.arange(-10,10,2), fontsize=12)
 
plt.title('Sample Scatterplot Background', fontsize=16)


# Barbs with specified spacing (q), color (cmap), range (norm), with C = windspeed, and linewidth adj
q=4
uunow=uu10[timex-1,:,:]
vvnow=vv10[timex-1,:,:]
C=np.hypot(uunow,vvnow)
colors = ['blue', 'red']
norm=plt.Normalize(0,40)
cmap = LinearSegmentedColormap.from_list('name', colors)
plt.barbs(to_np(lonwrf[::q,::q]),to_np(latwrf[::q,::q]),to_np(uu10[timex-1,::q,::q]),
          to_np(vv10[timex-1,::q,::q]),C[::q,::q],cmap=cmap,norm=norm,lw=1.5,transform=ccrs.PlateCarree())


# x-axis labels for vertical cross-sections that would otherwise have meaningless labels. 
# labels are forced and so need to be checked for accuracy but avoids MultipleLocator base crap
# first get x dimensions of plot, and calc spacing between ticks (here 4)
xmin, xmax = ax1.get_xlim()
print(xmin,xmax)
plt.xlim(xmin,xmax)
xdif=(xmax-xmin)/4
print(" X ",xmax,xmin,xdif)

# this works for forcing desired x-axis labels
# I know the plot goes from -107 to -103 longitude for fixed latitude
# I know I want 5 labels and I know what xmin, xmax are from get_xlim
labels = ['-107','-106','-105','-104','-103']
xl = np.arange(xmin,xmax+xdif,xdif)
print(" XL ",xl)
plt.xticks(xl,labels)

# controlling label sizes for entire notebook
SMALL_SIZE = 12
MEDIUM_SIZE = 14
BIGGER_SIZE = 16

plt.rc('font', size=SMALL_SIZE)          # controls default text sizes
plt.rc('axes', titlesize=SMALL_SIZE)     # fontsize of the axes title
plt.rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels
plt.rc('xtick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
plt.rc('ytick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
plt.rc('legend', fontsize=SMALL_SIZE)    # legend fontsize
plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title

#

plt.title('This is a test scatterplot')         # plot title
plt.xlabel('x',fontsize='medium')				# x axis label
plt.ylabel('y',fontsize='medium')				# y axis label
plt.xticks(np.arange(0,12,2),fontsize='small') 			# ticks 
plt.yticks(np.arange(0,12,2),fontsize='small') 			# ticks
plt.xlim([0,10])   						# limit range of x axis
plt.ylim([0,10])   						# limit range of y axis

plt.xticks([])                          # hides/suppresses x axis labels and ticks

# Annotating numbers/labels on plots

# gcount, cat, and gcc_3 are 1D arrays with number of items, x, and y axis values, respectively
# the multiplicative factors shift the labels
# the '%.0f' strips off decimal place
for i, label in enumerate(gcount):
    plt.annotate('%.0f'%label, (cat[i]*1.03, gcc_3[i]*1.05))

# Another example of annotating labels
# 'state' is a vector of labels w/ same length as x and y

plt.scatter(x,y,s=25,c='k')

for i,label in enumerate(state):
    plt.annotate(label, (x[i], y[i]),size=10,textcoords="offset points",xytext=(5,5)) # xytext pushes label up and to right a little



# control decimal places on x, y axis labels

from matplotlib.ticker import FormatStrFormatter
ax.yaxis.set_major_formatter(FormatStrFormatter('%.1f')) # one decimal place for y-axis

# control tick mark interval

from matplotlib.ticker import MultipleLocator
plt.xlim(datemin,datemax)
ax.xaxis.set_major_locator(MultipleLocator(10)) # Change major ticks to show every 10.

# this might work better for controlling spacing of date labels
spacing=2
for label in ax.xaxis.get_ticklabels()[::spacing]:
    label.set_visible(False)

# title for a multi-panel plot

plt.suptitle(" this is my title ")

# this can control space between top of plot and suptitle title. But it shifts figtext locations too

fig.subplots_adjust(top=0.95)

# make the title BOLD font

plt.title(" WHO Influenza data for the United States 2019-2020",fontweight='bold')

# add text on or near plot,  Coordinates are in range [0,1]
plt.figtext(0.735,0.84,'Text to be added',size=12)

# this version makes font italic
# adding ha='left' does left justify for horizontal arrangement (ha)
# adding loc='left' starts title at left edge of plot
plt.figtext(0.135,0.85,'Data from https://who.int/toolkits/flunet',size=12,fontstyle='italic')

# adding text including floating point variable information, controlling number of decimals printed
plt.figtext(0.28,0.84,f' CAPE {cape:.1f}',size=12)  

# printing a datetime object as a string, cutting out the ugly fractions of second, and adding a bounding box
plt.figtext(0.24,0.15,'Plot valid '+ str(timenow.values)[0:19],size=12,
    bbox=dict(facecolor='white', edgecolor='black', boxstyle='round,pad=1'))

# transform=ax.transAxes sets location of figtext in frame of the individual frame of multiplot figure [axis called 'ax']
plt.figtext(0.97,0.09,str(z0min)+' <= LU <= '+str(z0thresh),size=12,ha='right',transform=ax.transAxes)

# draw a vertical line at x=0 with color lightgrey
plt.axvline(x=0, lw=2,c='lightgrey')

# draw a horizontal line at y=0 with color lightgrey
plt.axhline(y=0, lw=2,c='lightgrey')

# save a figure, and cut out some whitespace
# facecolor='white', transparent=False  removes transparency
plt.savefig("xplot.png", bbox_inches = 'tight',dpi=300,facecolor='white', transparent=False)

datemin=datetime(2019,1,1)
datemax=datetime(2020,11,16)
ax1.xaxis_date()
plt.xlim(datemin,datemax) 


datemin=datetime(2019,4,1,0)
datemax=datetime(2019,5,1,0)
plt.xlim(datemin,datemax)
# Change major ticks to show every 10.
from matplotlib.ticker import MultipleLocator
ax.xaxis.set_major_locator(MultipleLocator(10))

ax2.invert_yaxis()   # inverts y-axis
plt.subplots_adjust(hspace=0.00)  # squeeze out vertical space between plots

# the twinx command clones the ax axis object as ax2, so our new plot shares the same x-axis
ax2 = ax.twinx()

# now we plot wind direction, and ask for a different, more reasonable y-axis plot limits
# this plot shows up on the secondary y-axis at right
plt.scatter(date,MTIC1_wdir,c='red',s=7)    # those wind direction dots were too large so we're downsizing
plt.scatter(date1,MOIC1_wdir,c='blue',s=7)    # those wind direction dots were too large so we're downsizing
plt.ylim([0,360])

List of named colors

• Linear regression

a, b = np.polyfit(x, y, deg=1)       # fits the linear regression (slope = a, intercept = b)
correlation = np.corrcoef(x,y)[0,1]  # fetches off-diagonal of 2x2 correlation matrix
rsq = correlation*correlation        # computes the R-squared 
f = a*x + b                          # computes the predicted values
plt.scatter(x,y)
plt.plot(x,f,c='red')

# alternative
from sklearn.linear_model import LinearRegression
plt.scatter(x,y)
model = LinearRegression()
model.fit(x.reshape(-1, 1), y)

x_new = np.linspace(np.min(x), np.max(x), 100)   # for least squares line
y_new = model.predict(x_new[:, np.newaxis])
plt.plot(x_new,y_new,c='orange',lw=2)

# np.polyfit wants 2 1-D vectors, so can use np.ravel to flatten 2D arrays
# x and y start out (1059, 1799)
x=np.ravel(wspd[mlevel-1,:,:].values-wspd_hrrr_drift[mlevel-1,:,:])
y=np.ravel(wspd_hrrr[mlevel-1,:,:]-wspd_hrrr_drift[mlevel-1,:,:])
# x and y end up (1905141,)

x = np.copy(GF) # 2D array
x = np.ravel(x) # make it 1D from 2D
print(" shape of x ",np.shape(x))

specified_bins=np.arange(1.0,10,0.1)
specified_midpoints=np.arange(1.1,10,0.1) # another try
print(len(specified_bins))
print(len(specified_midpoints))

sigma = 3 # smoothing

n, bins, patches = plt.hist(x, bins=specified_bins, facecolor='blue', edgecolor='k', alpha=0.5,label="Forecasts",zorder=1)
print(" max n fwind ",np.max(n)," max fwind ",np.max(fwind),len(x))
n = gaussian_filter(n, sigma)
plt.plot(specified_midpoints[5:],n[5:],c='blue',lw=3)

# hfxland, luland are 2D arrays (249x249)
# first, xform into 1D arrays for histogram

xx=np.ravel(hfxland)
yy=np.ravel(luland)
#print(np.shape(xx),np.shape(yy))

# set bins for xx (=qq) and yy (=rr)

qq=np.arange(-100,500,25) # hfx
rr=np.arange(1,17,1) # lu

# 2D histogram

data, x_e, y_e = np.histogram2d( xx, yy, bins = (qq,rr), density = False )
#print(" data ", data)
#print(" x_e ",x_e)
#print(" y_e ", y_e)

# creates the colors to be assigned to the points

z = scipy.interpolate.interpn( ( 0.5*(x_e[1:] + x_e[:-1]) , 0.5*(y_e[1:]+y_e[:-1]) ) , data , np.vstack([xx,yy]).T , method = "splinef2d", bounds_error = False)
#To be sure to plot all data
z[np.where(np.isnan(z))] = 0.0

# Sort the points by density, so that the densest points are plotted last
idx = z.argsort()
xx, yy, z = xx[idx], yy[idx], z[idx]
plt.scatter(xx,yy,c=z,edgecolors='None',s=20)

rainnc = getvar(ncfile1, "RAINNC",timeidx=ALL_TIMES) # get all times, otherwise just gets last one

latwrf = getvar(ncfile1, "XLAT",timeidx=-1) # just need one time
lonwrf = getvar(ncfile1, "XLONG",timeidx=-1)

wspd = getvar(ncfile1, "WSPD",timeidx=ALL_TIMES) # (times,lats,lons)
GF = np.max(wspd,axis=0)/np.mean(wspd,axis=0)    # (lats,lons)

# Read in a series of individual wrfout files, effectively combining them into single dataset
ncfile=[Dataset(location + "wrfout_d01_2007-09-21_06:00:00"),
        Dataset(location + "wrfout_d01_2007-09-21_07:00:00"),
        Dataset(location + "wrfout_d01_2007-09-21_08:00:00"),
        Dataset(location + "wrfout_d01_2007-09-21_09:00:00"),
        Dataset(location + "wrfout_d01_2007-09-21_10:00:00"),
        Dataset(location + "wrfout_d01_2007-09-21_11:00:00"),
        Dataset(location + "wrfout_d01_2007-09-21_12:00:00"),
        Dataset(location + "wrfout_d01_2007-09-21_13:00:00"),
        Dataset(location + "wrfout_d01_2007-09-21_14:00:00"),
        Dataset(location + "wrfout_d01_2007-09-21_15:00:00"),
        Dataset(location + "wrfout_d01_2007-09-21_16:00:00"),
        Dataset(location + "wrfout_d01_2007-09-21_17:00:00"),
        Dataset(location + "wrfout_d01_2007-09-21_18:00:00"),
        Dataset(location + "wrfout_d01_2007-09-21_19:00:00"),
        Dataset(location + "wrfout_d01_2007-09-21_20:00:00"),
        Dataset(location + "wrfout_d01_2007-09-21_21:00:00"),
        Dataset(location + "wrfout_d01_2007-09-21_22:00:00"),
        Dataset(location + "wrfout_d01_2007-09-21_23:00:00"),
        Dataset(location + "wrfout_d01_2007-09-22_00:00:00"),
        Dataset(location + "wrfout_d01_2007-09-22_01:00:00"),
        Dataset(location + "wrfout_d01_2007-09-22_02:00:00"),
        Dataset(location + "wrfout_d01_2007-09-22_03:00:00"),
        Dataset(location + "wrfout_d01_2007-09-22_04:00:00"),
        Dataset(location + "wrfout_d01_2007-09-22_05:00:00"),
        Dataset(location + "wrfout_d01_2007-09-22_06:00:00")]

lat_0 = 43
lon_0 = -75

fig = plt.figure(figsize=(10,8)) 
ax = fig.add_subplot(1,1,1,projection=ccrs.LambertConformal(central_latitude=lat_0, central_longitude=lon_0)) 
ax.set_extent([-80,-71.5,40.2,45.25])

ax.add_feature(cfeat.LAND)  				# adds colorshading [yellow] to land areas
ax.add_feature(cfeat.OCEAN) 				# adds colorshading [blue] to ocean areas
ax.add_feature(cfeat.COASTLINE,alpha=0.5)
ax.add_feature(cfeat.BORDERS,linestyle=':') # colon asks for dotted line
ax.add_feature(cfeat.LAKES,alpha=0.5) 		# alpha sets transparency level
ax.add_feature(cfeat.RIVERS)
ax.add_feature(cfeat.STATES,linestyle=':') 
ax.gridlines()

ax.scatter(-73.8,42.7,marker='*',s=400,c='k',transform=ccrs.PlateCarree())

# add a circle about a lat/lon location

radius_in_meters = 1000.
circle_points = cartopy.geodesic.Geodesic().circle(lon=lons[0], lat=lats[0], radius=radius_in_meters, endpoint=False)
geom = shapely.geometry.Polygon(circle_points)
ax.add_geometries((geom,), crs=ccrs.PlateCarree(), facecolor='none', edgecolor='red', linewidth=1)

# adding on a Google-derived relief map as background

from cartopy.io.img_tiles import GoogleTiles

# for google tiles relief map
# https://stackoverflow.com/questions/37423997/cartopy-shaded-relief

class ShadedReliefESRI(GoogleTiles):
    # shaded relief
    def _image_url(self, tile):
        x, y, z = tile
        url = ('https://server.arcgisonline.com/ArcGIS/rest/services/' \
               'World_Shaded_Relief/MapServer/tile/{z}/{y}/{x}.jpg').format(
               z=z, y=y, x=x)
        return url


ax = plt.axes(projection=ShadedReliefESRI().crs)
ax.set_extent([-112.64,-112.02,39.13,39.62])
ax.add_image(ShadedReliefESRI(), 8) # 12 is higher-res and sharper
plt.show()

# A map with an inset, as used in plot_GFS_GRIB_analyses.ipynb
fig = plt.figure(figsize=(16, 10))

# select our central latitude and longitude
centlon = -97.5
centlat = 40.

# Select a projection
proj=ccrs.LambertConformal(central_latitude=centlat, central_longitude=centlon)

# create our axis object 'ax', specifying Lambert Conformal projection
ax = fig.add_subplot(1,1,1,projection=proj) 

# The features function we defined adds some cartopy features on.  We pass it 'ax'
features(ax)

plot_bounds = [-135,-75,20,55]
ax.set_extent(plot_bounds)

# pcolormesh
colors = ['skyblue', 'white', 'red']
cmap = LinearSegmentedColormap.from_list('name', colors)
norm = plt.Normalize(0, 60)
cm = ax.pcolormesh(lon_hrrr,lat_hrrr,wspd_diff,shading='auto',cmap=cmap, norm=norm, transform=ccrs.PlateCarree())
plt.colorbar(cm, orientation="vertical",shrink=0.5,label="wind speed (m/s)") 

q=10
plt.barbs(to_np(lon_hrrr[::q]),to_np(lat_hrrr[::q]),to_np(uu[::q,::q]),to_np(vv[::q,::q]),color='black',zorder=40,transform=ccrs.PlateCarree())

# and add a title
plt.title(" Level "+str(p_levels[mlevel-1])+" mb GFS wind speed analysis (m/s) run "+str(filename)[0:14]); # semicolon suppresses unwanted text printed to screen

# add inset locator box to main map
# https://stackoverflow.com/questions/55303911/add-polygon-box-to-cartopy-python
geom = geometry.box(minx=-110,maxx=-100,miny=36,maxy=43)
ax.add_geometries([geom], crs=cartopy.crs.PlateCarree(), linewidth=3, facecolor='none', edgecolor='k')

##############################################################################################################
# inset plot
##############################################################################################################

# https://stackoverflow.com/questions/45527584/how-to-easily-add-a-sub-axes-with-proper-position-and-size-in-matplotlib-and-car
# https://scipython.com/blog/inset-plots-in-matplotlib/

# select our central latitude and longitude for INSET
centlon = -105.27055
centlat = 40.01499
proj2=ccrs.LambertConformal(central_latitude=centlat, central_longitude=centlon)
ax2 = plt.axes([0, 0, 1, 1], projection=proj2)
ax2.set_global()

# Manually set the position and relative size of the inset axes within ax1
ip = InsetPosition(ax, [0.68,0.02,0.35,0.35]) #[0.4,0.2,0.5,0.5]
ax2.set_axes_locator(ip)

plot_bounds = [-110,-100,36,43]
ax2.set_extent(plot_bounds)
norm = plt.Normalize(0, 60)
cm = ax2.pcolormesh(lon_hrrr,lat_hrrr,wspd_diff,shading='auto',cmap=cmap, norm=norm, zorder=1,transform=ccrs.PlateCarree())
features(ax2)
ax2.add_feature(cfeat.STATES.with_scale('50m'),linestyle='dotted',linewidth=2)
ax2.scatter(-105.27055,40.01499,marker='*',s=200,c='w',edgecolors='k',transform=ccrs.PlateCarree(),zorder=50) # Boulder
ax2.scatter(-106.29,40.01,marker='o',s=125,c='w',edgecolors='k',transform=ccrs.PlateCarree(),zorder=50) # upwind

qlon=7
qlat=4
print(np.shape(lon_hrrr[::qlon]),np.shape(lat_hrrr[::qlat]),np.shape(uu[::qlon,::qlat]))
plt.barbs(to_np(lon_hrrr[::qlon]),to_np(lat_hrrr[::qlat]),to_np(uu[::qlat,::qlon]),to_np(vv[::qlat,::qlon]),color='black',zorder=40,transform=ccrs.PlateCarree())

########## END ############

# create figure and skew plot
fig = plt.figure(figsize=(10,10))
skew = plots.SkewT(fig)

# plot T and Td
skew.plot(p, T, 'red')
skew.plot(p, Td, 'green')

# now the wind barbs, at selected vertical interval
interval = np.arange(100,1000,50) * units.hPa  # remember it does not include the last value of range!!!
idx = mpcalc.resample_nn_1d(p, interval)  # resample, using nearest-neighbor (nn), in 1 dimension (1d). Returns indices
skew.plot_barbs(p[idx], u[idx], v[idx])  # Now only for indices defined in last cell

# set axis limits
skew.ax.set_ylim(1000,100) # limit y-axis to 1000 to 100 mb
skew.ax.set_xlim(-50,30)   # limit x-axis to -50 to +30 C

# plot the reference curves/lines
skew.plot_dry_adiabats()
skew.plot_moist_adiabats()
skew.plot_mixing_lines()

# define parcel
parcel_path = mpcalc.parcel_profile(p, T[0], Td[0])

# calculate CAPE and CIN
cape, cin = mpcalc.cape_cin(p, T, Td, parcel_path)

# shade CAPE and CIN
skew.shade_cape(p, T, parcel_path)
skew.shade_cin(p, T, parcel_path)

# calculate parameters for LFC and EL defaulting to surface parcel
# and indicate on plot if they exist
lfc_p, lfc_t = mpcalc.lfc(p, T, Td)
el_p, el_t = mpcalc.el(p, T, Td)
if lfc_p:
    skew.ax.axhline(lfc_p)   
if el_p:
    skew.ax.axhline(el_p)

# print CAPE and CIN on skew-t, in a box with black border and white background
plt.figtext(0.29,0.16,f' CAPE {cape:.1f} \n CIN {cin:.1f} ',size=12,
            bbox=dict(facecolor='white', edgecolor='black', boxstyle='round,pad=1'))

# make the hodograph inset
# ax_hod = inset_axes(skew.ax, '35%', '35%', loc=1) # default but may obscure wind barbs
ax_hod = inset_axes(skew.ax, '35%', '35%', loc='upper right', bbox_transform=skew.ax.transAxes, 
                    bbox_to_anchor=(-0.05, 0., 1, 1))  # shifts hodo a bit to the left
h = Hodograph(ax_hod, component_range=80.)
h.add_grid(increment=20)
h.plot_colormapped(u, v, height);

ds = xr.open_dataset("/spare11/atm240/EXAMPLE02/nam.t00z.conusnest.hiresf06.tm00.grib2",engine='pynio') 


# get catalogue of variable names (could yield LOTS of output)
names =  ds.variables.keys()
print("Variable names:", names) 

# ------------------------------------------------------------------------------------------------- 
# some fields to extract from GRIB2 data. We're plotting 500 mb height and SLP here. 
# -------------------------------------------------------------------------------------------------  
hgt = ds.HGT_P0_L100_GLC0.sel(lv_ISBL0 = 50000.0)        # 500 mb height  
u10m = ds.UGRD_P0_L103_GLC0.sel(lv_HTGL8 = 10.0)         # 10 m wind speed (u)  
v10m = ds.VGRD_P0_L103_GLC0.sel(lv_HTGL8 = 10.0)         # 10 m wind speed (v) 
hgt0 = ds.HGT_P0_L1_GLC0                                 # terrain height 
u80m = ds.UGRD_P0_L103_GLC0.sel(lv_HTGL8 = 80.0)         # 80 m wind speed (u) 
v80m = ds.VGRD_P0_L103_GLC0.sel(lv_HTGL8 = 80.0)         # 80 m wind speed (v) 
slp_pa  = ds.PRMSL_P0_L101_GLC0                          # sea-level pressure (SLP) 

# lat/lon in degrees
x1 = math.radians(n_lat) # location 1
y1 = math.radians(n_lon)
x2 = math.radians(target_lat) # location 2
y2 = math.radians(target_lon)
        
# Great circle distance in radians
angle1 = math.acos(math.sin(x1) * math.sin(x2) 
				+ math.cos(x1) * math.cos(x2) * math.cos(y1 - y2))

# Convert back to degrees.
angle1 = math.degrees(angle1)

# Each degree on a great circle of Earth is 60 nautical miles.
distance1 = 60.0 * angle1
        
# Convert nautical miles to miles
gcd= distance1*1.15078

filter_by_keys={'typeOfLevel': 'meanSea'}
    filter_by_keys={'typeOfLevel': 'hybrid'}
    filter_by_keys={'typeOfLevel': 'atmosphere'}
    filter_by_keys={'typeOfLevel': 'surface'}
    filter_by_keys={'typeOfLevel': 'planetaryBoundaryLayer'}
    filter_by_keys={'typeOfLevel': 'isobaricInPa'}
    filter_by_keys={'typeOfLevel': 'isobaricInhPa'}
    filter_by_keys={'typeOfLevel': 'heightAboveGround'}
    filter_by_keys={'typeOfLevel': 'depthBelowLandLayer'}
    filter_by_keys={'typeOfLevel': 'heightAboveSea'}
    filter_by_keys={'typeOfLevel': 'atmosphereSingleLayer'}
    
    filter_by_keys={'typeOfLevel': 'lowCloudLayer'}
    filter_by_keys={'typeOfLevel': 'middleCloudLayer'}
    filter_by_keys={'typeOfLevel': 'highCloudLayer'}
    filter_by_keys={'typeOfLevel': 'cloudCeiling'}
    filter_by_keys={'typeOfLevel': 'convectiveCloudBottom'}
    filter_by_keys={'typeOfLevel': 'lowCloudBottom'}
    filter_by_keys={'typeOfLevel': 'middleCloudBottom'}
    filter_by_keys={'typeOfLevel': 'highCloudBottom'}
    filter_by_keys={'typeOfLevel': 'convectiveCloudTop'}
    filter_by_keys={'typeOfLevel': 'lowCloudTop'}
    filter_by_keys={'typeOfLevel': 'middleCloudTop'}
    filter_by_keys={'typeOfLevel': 'highCloudTop'}
    filter_by_keys={'typeOfLevel': 'convectiveCloudLayer'}
    filter_by_keys={'typeOfLevel': 'boundaryLayerCloudLayer'}
    filter_by_keys={'typeOfLevel': 'nominalTop'}
    
    filter_by_keys={'typeOfLevel': 'heightAboveGroundLayer'}
    
    filter_by_keys={'typeOfLevel': 'tropopause'}
    filter_by_keys={'typeOfLevel': 'maxWind'}
    filter_by_keys={'typeOfLevel': 'isothermZero'}
    filter_by_keys={'typeOfLevel': 'highestTroposphericFreezing'}
    filter_by_keys={'typeOfLevel': 'pressureFromGroundLayer'}
    filter_by_keys={'typeOfLevel': 'sigmaLayer'}
    filter_by_keys={'typeOfLevel': 'sigma'}
    filter_by_keys={'typeOfLevel': 'potentialVorticity'}

ds_hybrid = xr.open_dataset(location+filename,engine='cfgrib',filter_by_keys={'typeOfLevel': 'hybrid'})
ds_pressure = xr.open_dataset(location+filename,engine='cfgrib',filter_by_keys={'typeOfLevel': 'isobaricInhPa'})
ds_sfc = xr.open_dataset(location+filename_sfc,engine='cfgrib',
                                filter_by_keys={'level':10,'stepType': 'instant','typeOfLevel': 'heightAboveGround'})

theta_contours = ax1.contour(x, cross.gh[:,:], cross.theta[:,:],
                                     levels= thlevs, colors = 'k')
            
theta_contours.clabel(theta_contours.levels[1::2], fontsize=8, colors='k',
                              inline=1,inline_spacing=8, fmt='%i', rightside_up=True,
                              use_clabeltext=True) 

import urllib.request

url = 'ftp://ftp2.psl.noaa.gov/Datasets/ncep.reanalysis2.dailyavgs/gaussian_grid/air.2m.gauss.1979.nc'

req = urllib.request.Request(url)

with urllib.request.urlopen(req) as resp:
    ds = xr.open_dataset(io.BytesIO(resp.read()))

ds

# Used minisom to create an object called 'som'. This command
# prints out all of its attributes
print(dir(som))