import matplotlib
# matplotlib.use("TkAgg")
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from matplotlib.colors import BoundaryNorm
import os
import numpy as np
import xarray as xr
import argparse
from scipy.interpolate import griddata
from scipy import interpolate
import matplotlib.dates as mdates
import pandas as pd
[docs]
def plot_profile(filename, pdate, lat, lon):
""" This creates a simple plot showing the 2D fields"""
DATA = xr.open_dataset(filename)
(c_y, c_x) = naive_fast(DATA.lat.values, DATA.lon.values, lat, lon)
DATA = DATA.sel(time=pdate,west_east=c_x,south_north=c_y)
plt.figure(figsize=(20, 12))
depth = np.append(0,np.cumsum(DATA.LAYER_HEIGHT.values))
rho = np.append(DATA.LAYER_RHO[0],DATA.LAYER_RHO.values)
plt.step(rho,depth)
ax = plt.gca()
ax.invert_yaxis()
plt.show()
[docs]
def plot_profile_1D_timeseries(filename, var, domainy=None, start=None, end=None, lat=None, lon=None, stake_file=None, pit_name=None):
# Get dataset
ds = xr.open_dataset(filename)
if ((start is not None) & (end is not None)):
ds = ds.sel(time=slice(start,end))
# Select location
if ((lat is not None) & (lon is not None)):
ds = ds.sel(lat=lat, lon=lon, method='nearest')
# Get first layer height
fl = ds.attrs['First_layer_height_log_profile']
# Get data
if var=='T':
if ((lat is None) & (lon is None)):
V = ds.LAYER_T[:,0,0,:].values
else:
V = ds.LAYER_T[:,:].values
cmap = plt.get_cmap('YlGnBu_r')
barLabel = 'Temperature [K]'
if var=='RHO':
if ((lat is None) & (lon is None)):
V = ds.LAYER_RHO[:,0,0,:].values
else:
V = ds.LAYER_RHO[:,:].values
cmap = plt.get_cmap('YlGnBu_r')
barLabel = 'Density [kg m^-3]'
if var=='IF':
if ((lat is None) & (lon is None)):
V = ds.LAYER_ICE_FRACTION[:,0,0,:].values
else:
V = ds.LAYER_ICE_FRACTION[:,:].values
cmap = plt.get_cmap('YlGnBu_r')
barLabel = 'Ice fraction [-]'
if var=='REF':
if ((lat is None) & (lon is None)):
V = ds.LAYER_REFREEZE[:,0,0,:].values
else:
V = ds.LAYER_REFREEZE[:,:].values
cmap = plt.get_cmap('YlGnBu_r')
barLabel = 'Refreezing [m w.e.]'
if var=='LWC':
if ((lat is None) & (lon is None)):
V = ds.LAYER_LWC[:,0,0,:].values
else:
V = ds.LAYER_LWC[:,:].values
cmap = plt.get_cmap('YlGnBu_r')
barLabel = 'Liquid Water Content [-]'
if var=='POR':
if ((lat is None) & (lon is None)):
V = ds.LAYER_POROSITY[:,0,0,:].values
else:
V = ds.LAYER_POROSITY[:,:].values
cmap = plt.get_cmap('YlGnBu_r')
barLabel = 'Air Porosity [-]'
if var=='DEPTH':
if ((lat is None) & (lon is None)):
V = ds.LAYER_HEIGHT[:,0,0,:].values.cumsum(axis=1)
else:
V = ds.LAYER_HEIGHT[:,:].values.cumsum(axis=1)
cmap = plt.get_cmap('YlGnBu_r')
barLabel = 'Depth [m]'
if var=='HEIGHT':
if ((lat is None) & (lon is None)):
V = ds.LAYER_HEIGHT[:,0,0,:].values
else:
V = ds.LAYER_HEIGHT[:,:].values
cmap = plt.get_cmap('YlGnBu_r')
barLabel = 'Layer height [m]'
if ((lat is None) & (lon is None)):
D = ds.LAYER_HEIGHT[:,0,0,:].values.cumsum(axis=1)
else:
D = ds.LAYER_HEIGHT[:,:].values.cumsum(axis=1)
# Get dimensions
time = np.arange(ds.dims['time'])
if ((lat is None) & (lon is None)):
depth = ds.SNOWHEIGHT[:,0,0].values
else:
depth = ds.SNOWHEIGHT[:].values
# Calc plotting domain height
Dn = (np.int(np.floor(ds.SNOWHEIGHT.max()))+1)
## Create new grid
xi = time
if domainy is None:
domainy=0.0
yi = np.arange(domainy, Dn, fl)
X, Y = np.meshgrid(xi,yi)
data = np.full_like(X, np.nan, dtype=np.double)
# Re-calc depth data top=zero
D = (-(D.transpose()-depth).transpose())
def find_nearest(array, values):
array = np.asarray(array)
# the last dim must be 1 to broadcast in (array - values) below.
values = np.expand_dims(values, axis=-1)
indices = np.nanargmin(np.abs(array - values),axis=-1)
dist = np.nanmin(np.abs(array - values), axis=-1)
return indices,dist
for i in range(len(xi)):
sel = np.where(yi<depth[i])
idx,dist = find_nearest(D[i,:],yi[sel])
data[sel,i] = V[i,idx]
fig, ax = plt.subplots(figsize=(20, 10))
CS = ax.pcolormesh(X,Y,data,cmap=cmap)
#CS = ax.pcolormesh(X,Y,data, vmin=0, vmax=0.1)
N = pd.date_range(ds.time[0].values, ds.time[-1].values, freq='m')
M = pd.date_range(ds.time[0].values, ds.time[-1].values, freq='H')
if ((stake_file!=None) & (pit_name!=None)):
df = pd.read_csv(stake_file, sep='\t',index_col='TIMESTAMP')
df = df[df[pit_name]!=-9999]
for index, row in df.iterrows():
res = (M == pd.Timestamp(index)).argmax()
if res!=0:
plt.scatter(res,row[pit_name])
labIdx = []
label = []
for q in N:
o = np.where(M==q)
labIdx.append(o[0][0])
label.append(q.strftime('%Y-%m-%d'))
plt.xticks(labIdx, label,rotation=45,fontsize=16,weight='normal')
plt.ylabel('Depth [m]',fontsize=16, weight='normal')
plt.xlabel('Date',fontsize=16,weight='normal')
plt.title(var+'-Profile',fontsize=16, weight='normal')
# Make a colorbar for the ContourSet returned by the contourf call.
cbar = plt.colorbar(CS)
cbar.ax.set_ylabel(barLabel,fontsize=16,fontname = "Helvetica", weight='normal')
cbar.ax.set_yticks(barLabel,fontname = "Helvetica", weight='normal')
plt.show()
[docs]
def plot_profile_1D(filename, pdate, d=None, lat=None, lon=None):
""" This creates a simple plot showing the 2D fields"""
DATA = xr.open_dataset(filename)
DATA = DATA.sel(time=pdate)
if ((lat is not None) & (lon is not None)):
DATA = DATA.sel(lat=lat, lon=lon, method='nearest')
plt.figure(figsize=(5, 5))
depth = np.append(0,np.cumsum(DATA.LAYER_HEIGHT.values))
if ((lat is None) & (lon is None)):
rho = np.append(DATA.LAYER_RHO[:,:,0],DATA.LAYER_RHO.values)
t = np.append(DATA.LAYER_T[:,:,0],DATA.LAYER_T.values)
else:
rho = np.append(DATA.LAYER_RHO[0],DATA.LAYER_RHO.values)
t = np.append(DATA.LAYER_T[0],DATA.LAYER_T.values)
print('Date: %s' % (pdate))
print('T2: %.2f \t RH: %.2f \t U: %.2f \t G: %.2f' % (DATA.T2,DATA.RH2,DATA.U2,DATA.G))
if (d is not None):
for dmeas in d:
#idx, val = find_nearest(depth,d)
idx, val = find_nearest(depth,dmeas)
print('nearest depth: %.3f \t density: %.2f \t T: %.2f' % (val,rho[idx],t[idx]))
plt.step(rho,depth)
ax1 = plt.gca()
ax1.invert_yaxis()
ax1.set_ylabel('Depth [m]')
ax1.tick_params(axis='x', labelcolor='blue')
ax1.set_xlabel('Density [kg m^-3]', color='blue')
ax2 = ax1.twiny()
ax2.plot(t,depth, color='red')
ax2.set_xlabel('Temperature [K]', color='red')
ax2.tick_params(axis='x', labelcolor='red')
plt.show()
[docs]
def find_nearest(array, value):
array = np.asarray(array)
idx = (np.nanargmin(np.abs(array - value)))
return (idx,array[idx])
[docs]
def naive_fast(latvar,lonvar,lat0,lon0):
# Read latitude and longitude from file into numpy arrays
latvals = latvar[:]
lonvals = lonvar[:]
dist_sq = (latvals-lat0)**2 + (lonvals-lon0)**2
minindex_flattened = dist_sq.argmin() # 1D index of min element
iy_min,ix_min = np.unravel_index(minindex_flattened, latvals.shape)
return iy_min,ix_min
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Quick plot of the results file.')
parser.add_argument('-f', '--file', dest='file', help='Path to the result file')
parser.add_argument('-d', '--date', dest='pdate', help='Date of the profile plot')
parser.add_argument('-v', '--var', dest='var', default='RHO', help='Which variable to plot (e.g. T, RHO, etc.)')
parser.add_argument('-n', '--lat', dest='lat', default=None, help='Latitude value in case of 2D simulation', type=float)
parser.add_argument('-m', '--lon', dest='lon', default=None, help='Longitude value in case of 2D simulation', type=float)
parser.add_argument('-s', '--start', dest='start', default=None, help='Start date for the time plot')
parser.add_argument('-e', '--end', dest='end', default=None, help='End date for the time plot')
parser.add_argument('--stake-file', dest='stake_file', default=None, help='Path to the stake data file')
parser.add_argument('--pit', dest='pit_name', default=None, help='Name of the pit in the stake data file')
parser.add_argument('--depth', dest='d', nargs='+', default=None, help='An array with depth values for which the corresponding values are to be displayed', type=float)
args = parser.parse_args()
if (args.pdate is None):
plot_profile_1D_timeseries(args.file, args.var, args.d, args.start, args.end, stake_file=args.stake_file, pit_name=args.pit_name)
else:
plot_profile_1D(args.file, args.pdate, args.d, args.lat, args.lon)