def test(self): from netCDF4 import Dataset as NetCDF timeidx = 0 in_wrfnc = NetCDF(wrf_in) if (varname == "interplevel"): ref_ht_850 = _get_refvals(referent, "interplevel", repeat, multi) hts = getvar(in_wrfnc, "z", timeidx=timeidx) p = getvar(in_wrfnc, "pressure", timeidx=timeidx) hts_850 = interplevel(hts, p, 850) nt.assert_allclose(to_np(hts_850), ref_ht_850) elif (varname == "vertcross"): ref_ht_cross = _get_refvals(referent, "vertcross", repeat, multi) hts = getvar(in_wrfnc, "z", timeidx=timeidx) p = getvar(in_wrfnc, "pressure", timeidx=timeidx) pivot_point = CoordPair(hts.shape[-1] // 2, hts.shape[-2] // 2) ht_cross = vertcross(hts, p, pivot_point=pivot_point, angle=90.) nt.assert_allclose(to_np(ht_cross), ref_ht_cross, rtol=.01) elif (varname == "interpline"): ref_t2_line = _get_refvals(referent, "interpline", repeat, multi) t2 = getvar(in_wrfnc, "T2", timeidx=timeidx) pivot_point = CoordPair(t2.shape[-1] // 2, t2.shape[-2] // 2) t2_line1 = interpline(t2, pivot_point=pivot_point, angle=90.0) nt.assert_allclose(to_np(t2_line1), ref_t2_line) elif (varname == "vinterp"): # Tk to theta fld_tk_theta = _get_refvals(referent, "vinterp", repeat, multi) fld_tk_theta = np.squeeze(fld_tk_theta) tk = getvar(in_wrfnc, "temp", timeidx=timeidx, units="k") interp_levels = [200, 300, 500, 1000] field = vinterp(in_wrfnc, field=tk, vert_coord="theta", interp_levels=interp_levels, extrapolate=True, field_type="tk", timeidx=timeidx, log_p=True) tol = 5 / 100. atol = 0.0001 field = np.squeeze(field) nt.assert_allclose(to_np(field), fld_tk_theta, tol, atol)
def cross_2d_vars(cross_line, d2_vars, ds, end_point, start_point, ncfile)->xr.Dataset: line_cross = xr.Dataset() for v in d2_vars: line_cross[v] = wrf.interpline(ds[v], wrfin=ncfile, start_point=start_point, end_point=end_point, latlon=True, meta=True) line_cross = line_cross.rename(**{'line_idx': cross_line}) return line_cross
print(start_i, start_j) print(lat[start_i, start_j], lon[start_i, start_j]) print(end_i, end_j) print(lat[end_i, end_j], lon[end_i, end_j]) start = (start_j, start_i) end = (end_j, end_i) ## 2. Compute the vertical cross-section interpolation xy = xy(PS, start_point=start, end_point=end) rh_cross = interp2dxy(RH, xy, meta=False) print(np.nanmin(rh_cross), np.nanmax(rh_cross)) ## 3. Interpolate LAT-LON along the cross section and build LAT-LON pairs (to put them on the X-axis in the final plot) lat_interp = interpline(lat, start_point=CoordPair(x=start_j, y=start_i), end_point=CoordPair(x=end_j, y=end_i), latlon=True, meta=False) lon_interp = interpline(lon, start_point=CoordPair(x=start_j, y=start_i), end_point=CoordPair(x=end_j, y=end_i), latlon=False, meta=False) latlon_pairs = np.dstack((lat_interp, lon_interp)) ps_interp = interpline(PS, start_point=CoordPair(x=start_j, y=start_i), end_point=CoordPair(x=end_j, y=end_i), latlon=False, meta=False) ######################################################
# proximos a la topografia y los NaN para que # quede mas lindo el grafico. var_cross_filled = np.ma.copy(to_np(var_cross)) w_cross_filled = np.ma.copy(to_np(w_cross)) for i in range(var_cross_filled.shape[-1]): column_vals = var_cross_filled[:,i] first_idx = int(np.transpose((column_vals > -200).nonzero())[0]) var_cross_filled[0:first_idx, i] = var_cross_filled[first_idx, i] for i in range(w_cross_filled.shape[-1]): column_vals = w_cross_filled[:,i] first_idx = int(np.transpose((column_vals > -200).nonzero())[0]) w_cross_filled[0:first_idx, i] = w_cross_filled[first_idx, i] # Para sacar las alturas del terreno del corte ter_line = interpline(ter, wrfin=wrf_file, start_point=cross_start, end_point=cross_end) # Sacar las lat/lon lats, lons = latlon_coords(var) # Sacar la projeccion de cartopy cart_proj = get_cartopy(var) ############# GRAFICADO ############# # Figura fig = plt.figure(figsize=(8,6)) ax_cross = plt.axes() ## si se quiere tener un intervalo para los datos #var_levels = np.arange(5, 100., 10.)
def test(self): try: from netCDF4 import Dataset as NetCDF except: pass try: from PyNIO import Nio except: pass if not multi: timeidx = 0 if not pynio: in_wrfnc = NetCDF(wrf_in) else: # Note: Python doesn't like it if you reassign an outer scoped # variable (wrf_in in this case) if not wrf_in.endswith(".nc"): wrf_file = wrf_in + ".nc" else: wrf_file = wrf_in in_wrfnc = Nio.open_file(wrf_file) else: timeidx = None if not pynio: nc = NetCDF(wrf_in) in_wrfnc = [nc for i in xrange(repeat)] else: if not wrf_in.endswith(".nc"): wrf_file = wrf_in + ".nc" else: wrf_file = wrf_in nc = Nio.open_file(wrf_file) in_wrfnc = [nc for i in xrange(repeat)] if (varname == "interplevel"): ref_ht_500 = _get_refvals(referent, "z_500", repeat, multi) hts = getvar(in_wrfnc, "z", timeidx=timeidx) p = getvar(in_wrfnc, "pressure", timeidx=timeidx) hts_500 = interplevel(hts, p, 500) nt.assert_allclose(to_np(hts_500), ref_ht_500) elif (varname == "vertcross"): ref_ht_cross = _get_refvals(referent, "ht_cross", repeat, multi) ref_p_cross = _get_refvals(referent, "p_cross", repeat, multi) hts = getvar(in_wrfnc, "z", timeidx=timeidx) p = getvar(in_wrfnc, "pressure", timeidx=timeidx) pivot_point = CoordPair(hts.shape[-1] / 2, hts.shape[-2] / 2) ht_cross = vertcross(hts, p, pivot_point=pivot_point, angle=90.) nt.assert_allclose(to_np(ht_cross), ref_ht_cross, rtol=.01) # Test opposite p_cross1 = vertcross(p, hts, pivot_point=pivot_point, angle=90.0) nt.assert_allclose(to_np(p_cross1), ref_p_cross, rtol=.01) # Test point to point start_point = CoordPair(0, hts.shape[-2] / 2) end_point = CoordPair(-1, hts.shape[-2] / 2) p_cross2 = vertcross(p, hts, start_point=start_point, end_point=end_point) nt.assert_allclose(to_np(p_cross1), to_np(p_cross2)) elif (varname == "interpline"): ref_t2_line = _get_refvals(referent, "t2_line", repeat, multi) t2 = getvar(in_wrfnc, "T2", timeidx=timeidx) pivot_point = CoordPair(t2.shape[-1] / 2, t2.shape[-2] / 2) t2_line1 = interpline(t2, pivot_point=pivot_point, angle=90.0) nt.assert_allclose(to_np(t2_line1), ref_t2_line) # Test point to point start_point = CoordPair(0, t2.shape[-2] / 2) end_point = CoordPair(-1, t2.shape[-2] / 2) t2_line2 = interpline(t2, start_point=start_point, end_point=end_point) nt.assert_allclose(to_np(t2_line1), to_np(t2_line2)) elif (varname == "vinterp"): # Tk to theta fld_tk_theta = _get_refvals(referent, "fld_tk_theta", repeat, multi) fld_tk_theta = np.squeeze(fld_tk_theta) tk = getvar(in_wrfnc, "temp", timeidx=timeidx, units="k") interp_levels = [200, 300, 500, 1000] field = vinterp(in_wrfnc, field=tk, vert_coord="theta", interp_levels=interp_levels, extrapolate=True, field_type="tk", timeidx=timeidx, log_p=True) tol = 5 / 100. atol = 0.0001 field = np.squeeze(field) #print (np.amax(np.abs(to_np(field) - fld_tk_theta))) nt.assert_allclose(to_np(field), fld_tk_theta, tol, atol) # Tk to theta-e fld_tk_theta_e = _get_refvals(referent, "fld_tk_theta_e", repeat, multi) fld_tk_theta_e = np.squeeze(fld_tk_theta_e) interp_levels = [200, 300, 500, 1000] field = vinterp(in_wrfnc, field=tk, vert_coord="theta-e", interp_levels=interp_levels, extrapolate=True, field_type="tk", timeidx=timeidx, log_p=True) tol = 3 / 100. atol = 50.0001 field = np.squeeze(field) #print (np.amax(np.abs(to_np(field) - fld_tk_theta_e)/fld_tk_theta_e)*100) nt.assert_allclose(to_np(field), fld_tk_theta_e, tol, atol) # Tk to pressure fld_tk_pres = _get_refvals(referent, "fld_tk_pres", repeat, multi) fld_tk_pres = np.squeeze(fld_tk_pres) interp_levels = [850, 500] field = vinterp(in_wrfnc, field=tk, vert_coord="pressure", interp_levels=interp_levels, extrapolate=True, field_type="tk", timeidx=timeidx, log_p=True) field = np.squeeze(field) #print (np.amax(np.abs(to_np(field) - fld_tk_pres))) nt.assert_allclose(to_np(field), fld_tk_pres, tol, atol) # Tk to geoht_msl fld_tk_ght_msl = _get_refvals(referent, "fld_tk_ght_msl", repeat, multi) fld_tk_ght_msl = np.squeeze(fld_tk_ght_msl) interp_levels = [1, 2] field = vinterp(in_wrfnc, field=tk, vert_coord="ght_msl", interp_levels=interp_levels, extrapolate=True, field_type="tk", timeidx=timeidx, log_p=True) field = np.squeeze(field) #print (np.amax(np.abs(to_np(field) - fld_tk_ght_msl))) nt.assert_allclose(to_np(field), fld_tk_ght_msl, tol, atol) # Tk to geoht_agl fld_tk_ght_agl = _get_refvals(referent, "fld_tk_ght_agl", repeat, multi) fld_tk_ght_agl = np.squeeze(fld_tk_ght_agl) interp_levels = [1, 2] field = vinterp(in_wrfnc, field=tk, vert_coord="ght_agl", interp_levels=interp_levels, extrapolate=True, field_type="tk", timeidx=timeidx, log_p=True) field = np.squeeze(field) #print (np.amax(np.abs(to_np(field) - fld_tk_ght_agl))) nt.assert_allclose(to_np(field), fld_tk_ght_agl, tol, atol) # Hgt to pressure fld_ht_pres = _get_refvals(referent, "fld_ht_pres", repeat, multi) fld_ht_pres = np.squeeze(fld_ht_pres) z = getvar(in_wrfnc, "height", timeidx=timeidx, units="m") interp_levels = [500, 50] field = vinterp(in_wrfnc, field=z, vert_coord="pressure", interp_levels=interp_levels, extrapolate=True, field_type="ght", timeidx=timeidx, log_p=True) field = np.squeeze(field) #print (np.amax(np.abs(to_np(field) - fld_ht_pres))) nt.assert_allclose(to_np(field), fld_ht_pres, tol, atol) # Pressure to theta fld_pres_theta = _get_refvals(referent, "fld_pres_theta", repeat, multi) fld_pres_theta = np.squeeze(fld_pres_theta) p = getvar(in_wrfnc, "pressure", timeidx=timeidx) interp_levels = [200, 300, 500, 1000] field = vinterp(in_wrfnc, field=p, vert_coord="theta", interp_levels=interp_levels, extrapolate=True, field_type="pressure", timeidx=timeidx, log_p=True) field = np.squeeze(field) #print (np.amax(np.abs(to_np(field) - fld_pres_theta))) nt.assert_allclose(to_np(field), fld_pres_theta, tol, atol) # Theta-e to pres fld_thetae_pres = _get_refvals(referent, "fld_thetae_pres", repeat, multi) fld_thetae_pres = np.squeeze(fld_thetae_pres) eth = getvar(in_wrfnc, "eth", timeidx=timeidx) interp_levels = [850, 500, 5] field = vinterp(in_wrfnc, field=eth, vert_coord="pressure", interp_levels=interp_levels, extrapolate=True, field_type="theta-e", timeidx=timeidx, log_p=True) field = np.squeeze(field) #print (np.amax(np.abs(to_np(field) - fld_thetae_pres))) nt.assert_allclose(to_np(field), fld_thetae_pres, tol, atol)
w1_cross = w.vertcross(w1, h, wrfin=wrflist, start_point=start_point, end_point=end_point, latlon=True, meta=True) qv1_cross = w.vertcross(qv1, h, wrfin=wrflist, start_point=start_point, end_point=end_point, latlon=True, meta=True) ter_line = w.interpline(ter, wrfin=wrflist, start_point=start_point, end_point=end_point) #%% fig, ax = plt.subplots(figsize=(15, 3), dpi=100) t = pd.Timestamp(f'20170610 12:00:00') x_labels, y = xy(u1_cross) xs = np.arange(0, u1_cross.shape[-1], 1) # add NamCo station # label_color = 'white' # ax.plot(NamCo_ll[1], 4760, c=label_color, markersize=5, marker='^') # ax.text(NamCo_ll[1], 4580, 'NamCo', c=label_color, ha='center',weight='bold', fontsize=8) # QV var = fill_gap(qv1_cross.sel(Time=t)) * 1000
def test(self): try: from netCDF4 import Dataset as NetCDF except ImportError: pass try: import Nio except ImportError: pass timeidx = 0 if not multi else None pat = os.path.join(dir, pattern) wrf_files = glob.glob(pat) wrf_files.sort() wrfin = [] for fname in wrf_files: if not pynio: f = NetCDF(fname) try: f.set_always_mask(False) except AttributeError: pass wrfin.append(f) else: if not fname.endswith(".nc"): _fname = fname + ".nc" else: _fname = fname f = Nio.open_file(_fname) wrfin.append(f) if (varname == "interplevel"): ref_ht_500 = _get_refvals(referent, "z_500", multi) ref_p_5000 = _get_refvals(referent, "p_5000", multi) ref_ht_multi = _get_refvals(referent, "z_multi", multi) ref_p_multi = _get_refvals(referent, "p_multi", multi) ref_ht2_500 = _get_refvals(referent, "z2_500", multi) ref_p2_5000 = _get_refvals(referent, "p2_5000", multi) ref_ht2_multi = _get_refvals(referent, "z2_multi", multi) ref_p2_multi = _get_refvals(referent, "p2_multi", multi) ref_p_lev2d = _get_refvals(referent, "p_lev2d", multi) hts = getvar(wrfin, "z", timeidx=timeidx) p = getvar(wrfin, "pressure", timeidx=timeidx) wspd_wdir = getvar(wrfin, "wspd_wdir", timeidx=timeidx) # Make sure the numpy versions work first hts_500 = interplevel(to_np(hts), to_np(p), 500) hts_500 = interplevel(hts, p, 500) # Note: the '*2*' versions in the reference are testing # against the new version of interplevel in NCL nt.assert_allclose(to_np(hts_500), ref_ht_500) nt.assert_allclose(to_np(hts_500), ref_ht2_500) # Make sure the numpy versions work first p_5000 = interplevel(to_np(p), to_np(hts), 5000) p_5000 = interplevel(p, hts, 5000) nt.assert_allclose(to_np(p_5000), ref_p_5000) nt.assert_allclose(to_np(p_5000), ref_p2_5000) hts_multi = interplevel(to_np(hts), to_np(p), [1000., 850., 500., 250.]) hts_multi = interplevel(hts, p, [1000., 850., 500., 250.]) nt.assert_allclose(to_np(hts_multi), ref_ht_multi) nt.assert_allclose(to_np(hts_multi), ref_ht2_multi) p_multi = interplevel(to_np(p), to_np(hts), [500., 2500., 5000., 10000.]) p_multi = interplevel(p, hts, [500., 2500., 5000., 10000.]) nt.assert_allclose(to_np(p_multi), ref_p_multi) nt.assert_allclose(to_np(p_multi), ref_p2_multi) pblh = getvar(wrfin, "PBLH", timeidx=timeidx) p_lev2d = interplevel(to_np(p), to_np(hts), to_np(pblh)) p_lev2d = interplevel(p, hts, pblh) nt.assert_allclose(to_np(p_lev2d), ref_p_lev2d) # Just make sure these run below wspd_wdir_500 = interplevel(to_np(wspd_wdir), to_np(p), 500) wspd_wdir_500 = interplevel(wspd_wdir, p, 500) wspd_wdir_multi = interplevel(to_np(wspd_wdir), to_np(p), [1000, 500, 250]) wdpd_wdir_multi = interplevel(wspd_wdir, p, [1000, 500, 250]) wspd_wdir_pblh = interplevel(to_np(wspd_wdir), to_np(hts), pblh) wspd_wdir_pblh = interplevel(wspd_wdir, hts, pblh) if multi: wspd_wdir_pblh_2 = interplevel(to_np(wspd_wdir), to_np(hts), pblh[0, :]) wspd_wdir_pblh_2 = interplevel(wspd_wdir, hts, pblh[0, :]) # Since PBLH doesn't change in this case, it should match # the 0 time from previous computation. Note that this # only works when the data has 1 time step that is repeated. # If you use a different case with multiple times, # this will probably fail. nt.assert_allclose(to_np(wspd_wdir_pblh_2[:, 0, :]), to_np(wspd_wdir_pblh[:, 0, :])) nt.assert_allclose(to_np(wspd_wdir_pblh_2[:, -1, :]), to_np(wspd_wdir_pblh[:, 0, :])) elif (varname == "vertcross"): ref_ht_cross = _get_refvals(referent, "ht_cross", multi) ref_p_cross = _get_refvals(referent, "p_cross", multi) ref_ht_vertcross1 = _get_refvals(referent, "ht_vertcross1", multi) ref_ht_vertcross2 = _get_refvals(referent, "ht_vertcross2", multi) ref_ht_vertcross3 = _get_refvals(referent, "ht_vertcross3", multi) hts = getvar(wrfin, "z", timeidx=timeidx) p = getvar(wrfin, "pressure", timeidx=timeidx) pivot_point = CoordPair(hts.shape[-1] / 2, hts.shape[-2] / 2) # Beginning in wrf-python 1.3, users can select number of levels. # Originally, for pressure, dz was 10, so let's back calculate # the number of levels. p_max = np.floor(np.amax(p) / 10) * 10 # bottom value p_min = np.ceil(np.amin(p) / 10) * 10 # top value p_autolevels = int((p_max - p_min) / 10) # Make sure the numpy versions work first ht_cross = vertcross(to_np(hts), to_np(p), pivot_point=pivot_point, angle=90., autolevels=p_autolevels) ht_cross = vertcross(hts, p, pivot_point=pivot_point, angle=90., autolevels=p_autolevels) try: nt.assert_allclose(to_np(ht_cross), ref_ht_cross, atol=.0001) except AssertionError: absdiff = np.abs(to_np(ht_cross) - ref_ht_cross) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise lats = hts.coords["XLAT"] lons = hts.coords["XLONG"] # Test the manual projection method with lat/lon # Only do this for the non-multi case, since the domain # might be moving if not multi: if lats.ndim > 2: # moving nest lats = lats[0, :] lons = lons[0, :] ll_point = ll_points(lats, lons) pivot = CoordPair(lat=lats[int(lats.shape[-2] / 2), int(lats.shape[-1] / 2)], lon=lons[int(lons.shape[-2] / 2), int(lons.shape[-1] / 2)]) v1 = vertcross(hts, p, wrfin=wrfin, pivot_point=pivot_point, angle=90.0) v2 = vertcross(hts, p, projection=hts.attrs["projection"], ll_point=ll_point, pivot_point=pivot_point, angle=90.) try: nt.assert_allclose(to_np(v1), to_np(v2), atol=.0001) except AssertionError: absdiff = np.abs(to_np(v1) - to_np(v2)) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise # Test opposite p_cross1 = vertcross(p, hts, pivot_point=pivot_point, angle=90.0) try: nt.assert_allclose(to_np(p_cross1), ref_p_cross, atol=.0001) except AssertionError: absdiff = np.abs(to_np(p_cross1) - ref_p_cross) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise # Test point to point start_point = CoordPair(0, hts.shape[-2] / 2) end_point = CoordPair(-1, hts.shape[-2] / 2) p_cross2 = vertcross(p, hts, start_point=start_point, end_point=end_point) try: nt.assert_allclose(to_np(p_cross1), to_np(p_cross2), atol=.0001) except AssertionError: absdiff = np.abs(to_np(p_cross1) - to_np(p_cross2)) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise # Check the new vertcross routine pivot_point = CoordPair(hts.shape[-1] / 2, hts.shape[-2] / 2) ht_cross = vertcross(hts, p, pivot_point=pivot_point, angle=90., latlon=True) try: nt.assert_allclose(to_np(ht_cross), to_np(ref_ht_vertcross1), atol=.005) except AssertionError: absdiff = np.abs(to_np(ht_cross) - to_np(ref_ht_vertcross1)) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise levels = [1000., 850., 700., 500., 250.] ht_cross = vertcross(hts, p, pivot_point=pivot_point, angle=90., levels=levels, latlon=True) try: nt.assert_allclose(to_np(ht_cross), to_np(ref_ht_vertcross2), atol=.0001) except AssertionError: absdiff = np.abs(to_np(ht_cross) - to_np(ref_ht_vertcross2)) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise idxs = (0, slice(None)) if lats.ndim > 2 else (slice(None), ) start_lat = np.amin( lats[idxs]) + .25 * (np.amax(lats[idxs]) - np.amin(lats[idxs])) end_lat = np.amin( lats[idxs]) + .65 * (np.amax(lats[idxs]) - np.amin(lats[idxs])) start_lon = np.amin( lons[idxs]) + .25 * (np.amax(lons[idxs]) - np.amin(lons[idxs])) end_lon = np.amin( lons[idxs]) + .65 * (np.amax(lons[idxs]) - np.amin(lons[idxs])) start_point = CoordPair(lat=start_lat, lon=start_lon) end_point = CoordPair(lat=end_lat, lon=end_lon) ll_point = ll_points(lats[idxs], lons[idxs]) ht_cross = vertcross(hts, p, start_point=start_point, end_point=end_point, projection=hts.attrs["projection"], ll_point=ll_point, latlon=True, autolevels=1000) try: nt.assert_allclose(to_np(ht_cross), to_np(ref_ht_vertcross3), atol=.01) except AssertionError: absdiff = np.abs(to_np(ht_cross) - to_np(ref_ht_vertcross3)) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise if multi: ntimes = hts.shape[0] for t in range(ntimes): hts = getvar(wrfin, "z", timeidx=t) p = getvar(wrfin, "pressure", timeidx=t) ht_cross = vertcross(hts, p, start_point=start_point, end_point=end_point, wrfin=wrfin, timeidx=t, latlon=True, autolevels=1000) refname = "ht_vertcross_t{}".format(t) ref_ht_vertcross = _get_refvals(referent, refname, False) try: nt.assert_allclose(to_np(ht_cross), to_np(ref_ht_vertcross), atol=.01) except AssertionError: absdiff = np.abs( to_np(ht_cross) - to_np(ref_ht_vertcross)) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise elif (varname == "interpline"): ref_t2_line = _get_refvals(referent, "t2_line", multi) ref_t2_line2 = _get_refvals(referent, "t2_line2", multi) ref_t2_line3 = _get_refvals(referent, "t2_line3", multi) t2 = getvar(wrfin, "T2", timeidx=timeidx) pivot_point = CoordPair(t2.shape[-1] / 2, t2.shape[-2] / 2) # Make sure the numpy version works t2_line1 = interpline(to_np(t2), pivot_point=pivot_point, angle=90.0) t2_line1 = interpline(t2, pivot_point=pivot_point, angle=90.0) nt.assert_allclose(to_np(t2_line1), ref_t2_line) # Test the new NCL wrf_user_interplevel result try: nt.assert_allclose(to_np(t2_line1), ref_t2_line2) except AssertionError: absdiff = np.abs(to_np(t2_line1) - ref_t2_line2) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise # Test the manual projection method with lat/lon lats = t2.coords["XLAT"] lons = t2.coords["XLONG"] if multi: if lats.ndim > 2: # moving nest lats = lats[0, :] lons = lons[0, :] ll_point = ll_points(lats, lons) pivot = CoordPair(lat=lats[int(lats.shape[-2] / 2), int(lats.shape[-1] / 2)], lon=lons[int(lons.shape[-2] / 2), int(lons.shape[-1] / 2)]) l1 = interpline(t2, wrfin=wrfin, pivot_point=pivot_point, angle=90.0) l2 = interpline(t2, projection=t2.attrs["projection"], ll_point=ll_point, pivot_point=pivot_point, angle=90.) try: nt.assert_allclose(to_np(l1), to_np(l2)) except AssertionError: absdiff = np.abs(to_np(l1) - to_np(l2)) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise # Test point to point start_point = CoordPair(0, t2.shape[-2] / 2) end_point = CoordPair(-1, t2.shape[-2] / 2) t2_line2 = interpline(t2, start_point=start_point, end_point=end_point) try: nt.assert_allclose(to_np(t2_line1), to_np(t2_line2)) except AssertionError: absdiff = np.abs(to_np(t2_line1) - t2_line2) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise # Now test the start/end with lat/lon points start_lat = float( np.amin(lats) + .25 * (np.amax(lats) - np.amin(lats))) end_lat = float( np.amin(lats) + .65 * (np.amax(lats) - np.amin(lats))) start_lon = float( np.amin(lons) + .25 * (np.amax(lons) - np.amin(lons))) end_lon = float( np.amin(lons) + .65 * (np.amax(lons) - np.amin(lons))) start_point = CoordPair(lat=start_lat, lon=start_lon) end_point = CoordPair(lat=end_lat, lon=end_lon) t2_line3 = interpline(t2, wrfin=wrfin, timeidx=0, start_point=start_point, end_point=end_point, latlon=True) try: nt.assert_allclose(to_np(t2_line3), ref_t2_line3, atol=.0001) except AssertionError: absdiff = np.abs(to_np(t2_line3) - ref_t2_line3) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise # Test all time steps if multi: refnc = NetCDF(referent) ntimes = t2.shape[0] for t in range(ntimes): t2 = getvar(wrfin, "T2", timeidx=t) line = interpline(t2, wrfin=wrfin, timeidx=t, start_point=start_point, end_point=end_point, latlon=True) refname = "t2_line_t{}".format(t) refline = refnc.variables[refname][:] try: nt.assert_allclose(to_np(line), to_np(refline), atol=.0001) except AssertionError: absdiff = np.abs(to_np(line) - refline) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise refnc.close() # Test NCLs single time case if not multi: refnc = NetCDF(referent) ref_t2_line4 = refnc.variables["t2_line4"][:] t2 = getvar(wrfin, "T2", timeidx=0) line = interpline(t2, wrfin=wrfin, timeidx=0, start_point=start_point, end_point=end_point, latlon=True) try: nt.assert_allclose(to_np(line), to_np(ref_t2_line4), atol=.0001) except AssertionError: absdiff = np.abs(to_np(line) - ref_t2_line4) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise finally: refnc.close() elif (varname == "vinterp"): # Tk to theta fld_tk_theta = _get_refvals(referent, "fld_tk_theta", multi) fld_tk_theta = np.squeeze(fld_tk_theta) tk = getvar(wrfin, "temp", timeidx=timeidx, units="k") interp_levels = [200, 300, 500, 1000] # Make sure the numpy version works field = vinterp(wrfin, field=to_np(tk), vert_coord="theta", interp_levels=interp_levels, extrapolate=True, field_type="tk", timeidx=timeidx, log_p=True) field = vinterp(wrfin, field=tk, vert_coord="theta", interp_levels=interp_levels, extrapolate=True, field_type="tk", timeidx=timeidx, log_p=True) tol = 0 / 100. atol = 0.0001 field = np.squeeze(field) try: nt.assert_allclose(to_np(field), fld_tk_theta) except AssertionError: absdiff = np.abs(to_np(field) - fld_tk_theta) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise # Tk to theta-e fld_tk_theta_e = _get_refvals(referent, "fld_tk_theta_e", multi) fld_tk_theta_e = np.squeeze(fld_tk_theta_e) interp_levels = [200, 300, 500, 1000] field = vinterp(wrfin, field=tk, vert_coord="theta-e", interp_levels=interp_levels, extrapolate=True, field_type="tk", timeidx=timeidx, log_p=True) field = np.squeeze(field) try: nt.assert_allclose(to_np(field), fld_tk_theta_e, atol=.0001) except AssertionError: absdiff = np.abs(to_np(field) - fld_tk_theta_e) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise # Tk to pressure fld_tk_pres = _get_refvals(referent, "fld_tk_pres", multi) fld_tk_pres = np.squeeze(fld_tk_pres) interp_levels = [850, 500] field = vinterp(wrfin, field=tk, vert_coord="pressure", interp_levels=interp_levels, extrapolate=True, field_type="tk", timeidx=timeidx, log_p=True) field = np.squeeze(field) try: nt.assert_allclose(to_np(field), fld_tk_pres, atol=.0001) except AssertionError: absdiff = np.abs(to_np(field) - fld_tk_pres) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise # Tk to geoht_msl fld_tk_ght_msl = _get_refvals(referent, "fld_tk_ght_msl", multi) fld_tk_ght_msl = np.squeeze(fld_tk_ght_msl) interp_levels = [1, 2] field = vinterp(wrfin, field=tk, vert_coord="ght_msl", interp_levels=interp_levels, extrapolate=True, field_type="tk", timeidx=timeidx, log_p=True) field = np.squeeze(field) try: nt.assert_allclose(to_np(field), fld_tk_ght_msl, atol=.0001) except AssertionError: absdiff = np.abs(to_np(field) - fld_tk_ght_msl) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise # Tk to geoht_agl fld_tk_ght_agl = _get_refvals(referent, "fld_tk_ght_agl", multi) fld_tk_ght_agl = np.squeeze(fld_tk_ght_agl) interp_levels = [1, 2] field = vinterp(wrfin, field=tk, vert_coord="ght_agl", interp_levels=interp_levels, extrapolate=True, field_type="tk", timeidx=timeidx, log_p=True) field = np.squeeze(field) try: nt.assert_allclose(to_np(field), fld_tk_ght_agl, atol=.0001) except AssertionError: absdiff = np.abs(to_np(field) - fld_tk_ght_agl) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise # Hgt to pressure fld_ht_pres = _get_refvals(referent, "fld_ht_pres", multi) fld_ht_pres = np.squeeze(fld_ht_pres) z = getvar(wrfin, "height", timeidx=timeidx, units="m") interp_levels = [500, 50] field = vinterp(wrfin, field=z, vert_coord="pressure", interp_levels=interp_levels, extrapolate=True, field_type="ght", timeidx=timeidx, log_p=True) field = np.squeeze(field) try: nt.assert_allclose(to_np(field), fld_ht_pres, atol=.0001) except AssertionError: absdiff = np.abs(to_np(field) - fld_ht_pres) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise # Pressure to theta fld_pres_theta = _get_refvals(referent, "fld_pres_theta", multi) fld_pres_theta = np.squeeze(fld_pres_theta) p = getvar(wrfin, "pressure", timeidx=timeidx) interp_levels = [200, 300, 500, 1000] field = vinterp(wrfin, field=p, vert_coord="theta", interp_levels=interp_levels, extrapolate=True, field_type="pressure", timeidx=timeidx, log_p=True) field = np.squeeze(field) try: nt.assert_allclose(to_np(field), fld_pres_theta, atol=.0001) except AssertionError: absdiff = np.abs(to_np(field) - fld_pres_theta) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise # Theta-e to pres fld_thetae_pres = _get_refvals(referent, "fld_thetae_pres", multi) fld_thetae_pres = np.squeeze(fld_thetae_pres) eth = getvar(wrfin, "eth", timeidx=timeidx) interp_levels = [850, 500, 5] field = vinterp(wrfin, field=eth, vert_coord="pressure", interp_levels=interp_levels, extrapolate=True, field_type="theta-e", timeidx=timeidx, log_p=True) field = np.squeeze(field) try: nt.assert_allclose(to_np(field), fld_thetae_pres, atol=.0001) except AssertionError: absdiff = np.abs(to_np(field) - fld_thetae_pres) maxdiff = np.amax(absdiff) print(maxdiff) print(np.argwhere(absdiff == maxdiff)) raise
def make_result_file(opts): infilename = os.path.expanduser( os.path.join(opts.outdir, "ci_test_file.nc")) outfilename = os.path.expanduser( os.path.join(opts.outdir, "ci_result_file.nc")) with Dataset(infilename) as infile, Dataset(outfilename, "w") as outfile: for varname in WRF_DIAGS: var = getvar(infile, varname) add_to_ncfile(outfile, var, varname) for interptype in INTERP_METHS: if interptype == "interpline": hts = getvar(infile, "z") p = getvar(infile, "pressure") hts_850 = interplevel(hts, p, 850.) add_to_ncfile(outfile, hts_850, "interplevel") if interptype == "vertcross": hts = getvar(infile, "z") p = getvar(infile, "pressure") pivot_point = CoordPair(hts.shape[-1] // 2, hts.shape[-2] // 2) ht_cross = vertcross(hts, p, pivot_point=pivot_point, angle=90.) add_to_ncfile(outfile, ht_cross, "vertcross") if interptype == "interpline": t2 = getvar(infile, "T2") pivot_point = CoordPair(t2.shape[-1] // 2, t2.shape[-2] // 2) t2_line = interpline(t2, pivot_point=pivot_point, angle=90.0) add_to_ncfile(outfile, t2_line, "interpline") if interptype == "vinterp": tk = getvar(infile, "temp", units="k") interp_levels = [200, 300, 500, 1000] field = vinterp(infile, field=tk, vert_coord="theta", interp_levels=interp_levels, extrapolate=True, field_type="tk", log_p=True) add_to_ncfile(outfile, field, "vinterp") for latlonmeth in LATLON_METHS: if latlonmeth == "xy": # Hardcoded values from other unit tests lats = [-55, -60, -65] lons = [25, 30, 35] xy = ll_to_xy(infile, lats[0], lons[0]) add_to_ncfile(outfile, xy, "xy") else: # Hardcoded from other unit tests i_s = np.asarray([10, 100, 150], int) - 1 j_s = np.asarray([10, 100, 150], int) - 1 ll = xy_to_ll(infile, i_s[0], j_s[0]) add_to_ncfile(outfile, ll, "ll")
def test(self): try: from netCDF4 import Dataset as NetCDF except: pass try: from PyNIO import Nio except: pass if not multi: timeidx = 0 if not pynio: in_wrfnc = NetCDF(wrf_in) else: # Note: Python doesn't like it if you reassign an outer scoped # variable (wrf_in in this case) if not wrf_in.endswith(".nc"): wrf_file = wrf_in + ".nc" else: wrf_file = wrf_in in_wrfnc = Nio.open_file(wrf_file) else: timeidx = None if not pynio: nc = NetCDF(wrf_in) in_wrfnc = [nc for i in xrange(repeat)] else: if not wrf_in.endswith(".nc"): wrf_file = wrf_in + ".nc" else: wrf_file = wrf_in nc = Nio.open_file(wrf_file) in_wrfnc = [nc for i in xrange(repeat)] if (varname == "interplevel"): ref_ht_500 = _get_refvals(referent, "z_500", repeat, multi) hts = getvar(in_wrfnc, "z", timeidx=timeidx) p = getvar(in_wrfnc, "pressure", timeidx=timeidx) # Make sure the numpy versions work first hts_500 = interplevel(to_np(hts), to_np(p), 500) hts_500 = interplevel(hts, p, 500) nt.assert_allclose(to_np(hts_500), ref_ht_500) elif (varname == "vertcross"): ref_ht_cross = _get_refvals(referent, "ht_cross", repeat, multi) ref_p_cross = _get_refvals(referent, "p_cross", repeat, multi) hts = getvar(in_wrfnc, "z", timeidx=timeidx) p = getvar(in_wrfnc, "pressure", timeidx=timeidx) pivot_point = CoordPair(hts.shape[-1] / 2, hts.shape[-2] / 2) #ht_cross = vertcross(to_np(hts), p, pivot_point=pivot_point, # angle=90., latlon=True) # Make sure the numpy versions work first ht_cross = vertcross(to_np(hts), to_np(p), pivot_point=pivot_point, angle=90.) ht_cross = vertcross(hts, p, pivot_point=pivot_point, angle=90.) # Note: Until the bug is fixed in NCL, the wrf-python cross # sections will contain one extra point nt.assert_allclose(to_np(ht_cross)[..., 0:-1], ref_ht_cross, rtol=.01) # Test the manual projection method with lat/lon lats = hts.coords["XLAT"] lons = hts.coords["XLONG"] ll_point = ll_points(lats, lons) pivot = CoordPair(lat=lats[int(lats.shape[-2] / 2), int(lats.shape[-1] / 2)], lon=lons[int(lons.shape[-2] / 2), int(lons.shape[-1] / 2)]) v1 = vertcross(hts, p, wrfin=in_wrfnc, pivot_point=pivot_point, angle=90.0) v2 = vertcross(hts, p, projection=hts.attrs["projection"], ll_point=ll_point, pivot_point=pivot_point, angle=90.) nt.assert_allclose(to_np(v1), to_np(v2), rtol=.01) # Test opposite p_cross1 = vertcross(p, hts, pivot_point=pivot_point, angle=90.0) nt.assert_allclose(to_np(p_cross1)[..., 0:-1], ref_p_cross, rtol=.01) # Test point to point start_point = CoordPair(0, hts.shape[-2] / 2) end_point = CoordPair(-1, hts.shape[-2] / 2) p_cross2 = vertcross(p, hts, start_point=start_point, end_point=end_point) nt.assert_allclose(to_np(p_cross1), to_np(p_cross2)) elif (varname == "interpline"): ref_t2_line = _get_refvals(referent, "t2_line", repeat, multi) t2 = getvar(in_wrfnc, "T2", timeidx=timeidx) pivot_point = CoordPair(t2.shape[-1] / 2, t2.shape[-2] / 2) #t2_line1 = interpline(to_np(t2), pivot_point=pivot_point, # angle=90.0, latlon=True) # Make sure the numpy version works t2_line1 = interpline(to_np(t2), pivot_point=pivot_point, angle=90.0) t2_line1 = interpline(t2, pivot_point=pivot_point, angle=90.0) # Note: After NCL is fixed, remove the slice. nt.assert_allclose(to_np(t2_line1)[..., 0:-1], ref_t2_line) # Test the manual projection method with lat/lon lats = t2.coords["XLAT"] lons = t2.coords["XLONG"] ll_point = ll_points(lats, lons) pivot = CoordPair(lat=lats[int(lats.shape[-2] / 2), int(lats.shape[-1] / 2)], lon=lons[int(lons.shape[-2] / 2), int(lons.shape[-1] / 2)]) l1 = interpline(t2, wrfin=in_wrfnc, pivot_point=pivot_point, angle=90.0) l2 = interpline(t2, projection=t2.attrs["projection"], ll_point=ll_point, pivot_point=pivot_point, angle=90.) nt.assert_allclose(to_np(l1), to_np(l2), rtol=.01) # Test point to point start_point = CoordPair(0, t2.shape[-2] / 2) end_point = CoordPair(-1, t2.shape[-2] / 2) t2_line2 = interpline(t2, start_point=start_point, end_point=end_point) nt.assert_allclose(to_np(t2_line1), to_np(t2_line2)) elif (varname == "vinterp"): # Tk to theta fld_tk_theta = _get_refvals(referent, "fld_tk_theta", repeat, multi) fld_tk_theta = np.squeeze(fld_tk_theta) tk = getvar(in_wrfnc, "temp", timeidx=timeidx, units="k") interp_levels = [200, 300, 500, 1000] # Make sure the numpy version works field = vinterp(in_wrfnc, field=to_np(tk), vert_coord="theta", interp_levels=interp_levels, extrapolate=True, field_type="tk", timeidx=timeidx, log_p=True) field = vinterp(in_wrfnc, field=tk, vert_coord="theta", interp_levels=interp_levels, extrapolate=True, field_type="tk", timeidx=timeidx, log_p=True) tol = 5 / 100. atol = 0.0001 field = np.squeeze(field) #print (np.amax(np.abs(to_np(field) - fld_tk_theta))) nt.assert_allclose(to_np(field), fld_tk_theta, tol, atol) # Tk to theta-e fld_tk_theta_e = _get_refvals(referent, "fld_tk_theta_e", repeat, multi) fld_tk_theta_e = np.squeeze(fld_tk_theta_e) interp_levels = [200, 300, 500, 1000] field = vinterp(in_wrfnc, field=tk, vert_coord="theta-e", interp_levels=interp_levels, extrapolate=True, field_type="tk", timeidx=timeidx, log_p=True) tol = 3 / 100. atol = 50.0001 field = np.squeeze(field) #print (np.amax(np.abs(to_np(field) - fld_tk_theta_e)/fld_tk_theta_e)*100) nt.assert_allclose(to_np(field), fld_tk_theta_e, tol, atol) # Tk to pressure fld_tk_pres = _get_refvals(referent, "fld_tk_pres", repeat, multi) fld_tk_pres = np.squeeze(fld_tk_pres) interp_levels = [850, 500] field = vinterp(in_wrfnc, field=tk, vert_coord="pressure", interp_levels=interp_levels, extrapolate=True, field_type="tk", timeidx=timeidx, log_p=True) field = np.squeeze(field) #print (np.amax(np.abs(to_np(field) - fld_tk_pres))) nt.assert_allclose(to_np(field), fld_tk_pres, tol, atol) # Tk to geoht_msl fld_tk_ght_msl = _get_refvals(referent, "fld_tk_ght_msl", repeat, multi) fld_tk_ght_msl = np.squeeze(fld_tk_ght_msl) interp_levels = [1, 2] field = vinterp(in_wrfnc, field=tk, vert_coord="ght_msl", interp_levels=interp_levels, extrapolate=True, field_type="tk", timeidx=timeidx, log_p=True) field = np.squeeze(field) #print (np.amax(np.abs(to_np(field) - fld_tk_ght_msl))) nt.assert_allclose(to_np(field), fld_tk_ght_msl, tol, atol) # Tk to geoht_agl fld_tk_ght_agl = _get_refvals(referent, "fld_tk_ght_agl", repeat, multi) fld_tk_ght_agl = np.squeeze(fld_tk_ght_agl) interp_levels = [1, 2] field = vinterp(in_wrfnc, field=tk, vert_coord="ght_agl", interp_levels=interp_levels, extrapolate=True, field_type="tk", timeidx=timeidx, log_p=True) field = np.squeeze(field) #print (np.amax(np.abs(to_np(field) - fld_tk_ght_agl))) nt.assert_allclose(to_np(field), fld_tk_ght_agl, tol, atol) # Hgt to pressure fld_ht_pres = _get_refvals(referent, "fld_ht_pres", repeat, multi) fld_ht_pres = np.squeeze(fld_ht_pres) z = getvar(in_wrfnc, "height", timeidx=timeidx, units="m") interp_levels = [500, 50] field = vinterp(in_wrfnc, field=z, vert_coord="pressure", interp_levels=interp_levels, extrapolate=True, field_type="ght", timeidx=timeidx, log_p=True) field = np.squeeze(field) #print (np.amax(np.abs(to_np(field) - fld_ht_pres))) nt.assert_allclose(to_np(field), fld_ht_pres, tol, atol) # Pressure to theta fld_pres_theta = _get_refvals(referent, "fld_pres_theta", repeat, multi) fld_pres_theta = np.squeeze(fld_pres_theta) p = getvar(in_wrfnc, "pressure", timeidx=timeidx) interp_levels = [200, 300, 500, 1000] field = vinterp(in_wrfnc, field=p, vert_coord="theta", interp_levels=interp_levels, extrapolate=True, field_type="pressure", timeidx=timeidx, log_p=True) field = np.squeeze(field) #print (np.amax(np.abs(to_np(field) - fld_pres_theta))) nt.assert_allclose(to_np(field), fld_pres_theta, tol, atol) # Theta-e to pres fld_thetae_pres = _get_refvals(referent, "fld_thetae_pres", repeat, multi) fld_thetae_pres = np.squeeze(fld_thetae_pres) eth = getvar(in_wrfnc, "eth", timeidx=timeidx) interp_levels = [850, 500, 5] field = vinterp(in_wrfnc, field=eth, vert_coord="pressure", interp_levels=interp_levels, extrapolate=True, field_type="theta-e", timeidx=timeidx, log_p=True) field = np.squeeze(field) #print (np.amax(np.abs(to_np(field) - fld_thetae_pres))) nt.assert_allclose(to_np(field), fld_thetae_pres, tol, atol)
# Make a copy of the z cross data. Let's use regular numpy arrays for this. dbz_cross_filled = np.ma.copy(to_np(dbz_cross)) # For each cross section column, find the first index with non-missing # values and copy these to the missing elements below. for i in range(dbz_cross_filled.shape[-1]): column_vals = dbz_cross_filled[:, i] # Let's find the lowest index that isn't filled. The nonzero function # finds all unmasked values greater than 0. Since 0 is a valid value # for dBZ, let's change that threshold to be -200 dBZ instead. first_idx = int(np.transpose((column_vals > -200).nonzero())[0]) dbz_cross_filled[0:first_idx, i] = dbz_cross_filled[first_idx, i] # Get the terrain heights along the cross section line ter_line = interpline(ter, wrfin=nc, start_point=start_point, end_point=end_point) # Get the lat/lon points lats, lons = latlon_coords(dbz) # Get the cartopy projection object cart_proj = get_cartopy(dbz) # Create the figure fig = plt.figure(figsize=(8, 6)) ax_cross = plt.axes() dbz_levels = np.arange(5., 75., 5.) # Create the color table found on NWS pages.
def cross_section(variable_name, date, time, start_lat, start_lon, end_lat, end_lon, save=False): '''This function plots a vertical cross section of the chosen variable. Supported variables for plotting procedure are vertical_velocity, rh, omega, absolute_vorticity, theta_e and reflectivity.''' ### Predefine some variables ### # Define data filename data_dir = '/scratch3/thomasl/work/data/casestudy_baden/' filename = '{}wrfout_d02_{}_{}:00'.format(data_dir, date, time) # Define save directory save_dir = '/scratch3/thomasl/work/retrospective_part/' 'casestudy_baden/cross_section/' # Create the start point and end point for the cross section start_point = CoordPair(lat=start_lat, lon=start_lon) end_point = CoordPair(lat=end_lat, lon=end_lon) ### Start plotting procedure ### # Open NetCDF file ncfile = Dataset(filename) # Extract the model height, terrain height and variables ht = getvar(ncfile, 'z') / 1000 # change to km ter = getvar(ncfile, 'ter') / 1000 if variable_name == 'vertical_velocity': variable = getvar(ncfile, 'wa', units='kt') title_name = 'Vertical Velocity' colorbar_label = 'Vertical Velocity [$kn$]' variable_min = -2 variable_max = 2 elif variable_name == 'rh': variable = getvar(ncfile, 'rh') title_name = 'Relative Humidity' colorbar_label = 'Relative Humidity [$pct$]' variable_min = 0 variable_max = 100 elif variable_name == 'omega': variable = getvar(ncfile, 'omega') title_name = 'Vertical Motion (Omega)' colorbar_label = 'Omega [$Pa$ $s^-$$^1$]' variable_min = -5 variable_max = 5 elif variable_name == 'absolute_vorticity': variable = getvar(ncfile, 'avo') title_name = 'Absolute Vorticity' colorbar_label = 'Absolute Vorticity [$10^{-5}$' '$s^{-1}$]' variable_min = -50 variable_max = 100 elif variable_name == 'theta_e': variable = getvar(ncfile, 'theta_e') title_name = 'Theta-E' colorbar_label = 'Theta-E [$K$]' variable_min = 315 variable_max = 335 elif variable_name == 'reflectivity': variable = getvar(ncfile, 'dbz') #, timeidx=-1 title_name = 'Reflectivity' colorbar_label = 'Reflectivity [$dBZ$]' variable_min = 5 variable_max = 75 # Linear Z for interpolation Z = 10**(variable / 10) # Compute the vertical cross-section interpolation z_cross = vertcross(Z, ht, wrfin=ncfile, start_point=start_point, end_point=end_point, latlon=True, meta=True) # Convert back after interpolation variable_cross = 10.0 * np.log10(z_cross) # Make a copy of the z cross data variable_cross_filled = np.ma.copy(to_np(variable_cross)) # For each cross section column, find the first index with # non-missing values and copy these to the missing elements below for i in range(variable_cross_filled.shape[-1]): column_vals = variable_cross_filled[:, i] first_idx = int(np.transpose((column_vals > -200).nonzero())[0]) variable_cross_filled[0:first_idx, i] = variable_cross_filled[first_idx, i] ter_line = interpline(ter, wrfin=ncfile, start_point=start_point, end_point=end_point) # Get latitude and longitude points lats, lons = latlon_coords(variable) # Create the figure fig = plt.figure(figsize=(15, 10)) ax = plt.axes() ys = to_np(variable_cross.coords['vertical']) xs = np.arange(0, variable_cross.shape[-1], 1) # Make contour plot if variable_name == 'reflectivity': levels = np.arange(variable_min, variable_max, 5) # Create the dbz color table found on NWS pages. dbz_rgb = np.array( [[4, 233, 231], [1, 159, 244], [3, 0, 244], [2, 253, 2], [1, 197, 1], [0, 142, 0], [253, 248, 2], [229, 188, 0], [253, 149, 0], [253, 0, 0], [212, 0, 0], [188, 0, 0], [248, 0, 253], [152, 84, 198]], np.float32) / 255.0 dbz_cmap, dbz_norm = from_levels_and_colors(levels, dbz_rgb, extend='max') else: levels = np.linspace(variable_min, variable_max, 11) if variable_name == 'omega' or variable_name == 'vertical_velocity': def truncate_colormap(cmap, minval=0.0, maxval=1.0, n=100): new_cmap = colors.LinearSegmentedColormap.from_list( 'trunc({n},{a:.2f},{b:.2f})'.format(n=cmap.name, a=minval, b=maxval), cmap(np.linspace(minval, maxval, n))) return new_cmap old_cmap = plt.get_cmap('RdYlBu') cmap = truncate_colormap(old_cmap, 0.05, 0.9) norm = colors.DivergingNorm(vmin=variable_min, vcenter=0, vmax=variable_max) else: cmap = ListedColormap(sns.cubehelix_palette(20, start=.5, rot=-.75)) if variable_name == 'omega' or variable_name == 'vertical_velocity': variable_contours = ax.contourf(xs, ys, to_np(variable_cross_filled), levels=levels, cmap=cmap, extend='both', norm=norm) elif variable_name == 'rh': variable_contours = ax.contourf(xs, ys, to_np(variable_cross_filled), levels=levels, cmap=cmap, extend='max') elif variable_name == 'reflectivity': variable_contours = ax.contourf(xs, ys, to_np(variable_cross_filled), levels=levels, cmap=dbz_cmap, norm=dbz_norm, extend='both') else: variable_contours = ax.contourf(xs, ys, to_np(variable_cross_filled), levels=levels, cmap=cmap, extend='both') # Plot wind barbs if variable_name == 'vertical_velocity': u = getvar(ncfile, 'ua', units='kt') U = 10**(u / 10) u_cross = vertcross(U, ht, wrfin=ncfile, start_point=start_point, end_point=end_point, latlon=True, meta=True) u_cross = 10.0 * np.log10(u_cross) u_cross_filled = np.ma.copy(to_np(u_cross)) for i in range(u_cross_filled.shape[-1]): column_vals = u_cross_filled[:, i] first_idx = int(np.transpose((column_vals > -200).nonzero())[0]) u_cross_filled[0:first_idx, i] = u_cross_filled[first_idx, i] ax.barbs(xs[::3], ys[::3], to_np(u_cross_filled[::3, ::3]), to_np(variable_cross_filled[::3, ::3]), length=7, zorder=1) if variable_name == 'omega': u = getvar(ncfile, 'ua', units='kt') U = 10**(u / 10) u_cross = vertcross(U, ht, wrfin=ncfile, start_point=start_point, end_point=end_point, latlon=True, meta=True) u_cross = 10.0 * np.log10(u_cross) u_cross_filled = np.ma.copy(to_np(u_cross)) for i in range(u_cross_filled.shape[-1]): column_vals = u_cross_filled[:, i] first_idx = int(np.transpose((column_vals > -200).nonzero())[0]) u_cross_filled[0:first_idx, i] = u_cross_filled[first_idx, i] w = getvar(ncfile, 'wa', units='kt') W = 10**(w / 10) w_cross = vertcross(W, ht, wrfin=ncfile, start_point=start_point, end_point=end_point, latlon=True, meta=True) w_cross = 10.0 * np.log10(w_cross) w_cross_filled = np.ma.copy(to_np(w_cross)) for i in range(w_cross_filled.shape[-1]): column_vals = w_cross_filled[:, i] first_idx = int(np.transpose((column_vals > -200).nonzero())[0]) w_cross_filled[0:first_idx, i] = w_cross_filled[first_idx, i] ax.barbs(xs[::3], ys[::3], to_np(u_cross_filled[::3, ::3]), to_np(w_cross_filled[::3, ::3]), length=7, zorder=1, color='grey') # Add color bar cbar = mpu.colorbar(variable_contours, ax, orientation='vertical', aspect=40, shrink=.05, pad=0.05) cbar.set_label(colorbar_label, fontsize=15) cbar.set_ticks(levels) # Set x-ticks to use latitude and longitude labels coord_pairs = to_np(variable_cross.coords['xy_loc']) x_ticks = np.arange(coord_pairs.shape[0]) x_labels = [ pair.latlon_str(fmt='{:.2f}, {:.2f}') for pair in to_np(coord_pairs) ] # Set desired number of x ticks below num_ticks = 5 thin = int((len(x_ticks) / num_ticks) + .5) ax.set_xticks(x_ticks[::thin]) ax.set_xticklabels(x_labels[::thin], rotation=45, fontsize=10) ax.set_xlim(x_ticks[0], x_ticks[-1]) # Set y-ticks and limit the height ax.set_yticks(np.linspace(0, 12, 13)) ax.set_ylim(0, 12) # Set x-axis and y-axis labels ax.set_xlabel('Latitude, Longitude', fontsize=12.5) ax.set_ylabel('Height [$km$]', fontsize=12.5) # Fill in mountian area ht_fill = ax.fill_between(xs, 0, to_np(ter_line), facecolor='saddlebrown', zorder=2) # Make nicetime xr_file = xr.open_dataset(filename) nicetime = pd.to_datetime(xr_file.QVAPOR.isel(Time=0).XTIME.values) # Add title ax.set_title('Vertical Cross-Section of {}'.format(title_name), fontsize=20, loc='left') ax.set_title('Valid time: {} {}'.format( nicetime.strftime('%Y-%m-%d %H:%M'), 'UTC'), fontsize=15, loc='right') # Add grid for y axis ax.grid(axis='y', linestyle='--', color='grey') plt.show() ### Save figure ### if save == True: fig.savefig('{}cross_section_{}.png'.format(save_dir, variable_name), bbox_inches='tight', dpi=300)