Exemple #1
1
    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
Exemple #3
0
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)

######################################################
Exemple #4
0
        # 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.)
Exemple #5
0
    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)
Exemple #6
0
    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
Exemple #7
0
    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
Exemple #8
0
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")
Exemple #9
0
    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)
Exemple #10
0
# 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)