Exemple #1
0
    def test_basic_wrf(self):
        import xarray as xr

        ds = sio.open_xr_dataset(get_demo_file("wrf_tip_d1.nc"))

        # this is because read_dataset changes some stuff, let's see if
        # georef still ok
        dsxr = xr.open_dataset(get_demo_file("wrf_tip_d1.nc"))
        assert ds.salem.grid == dsxr.salem.grid

        lon, lat = ds.salem.grid.ll_coordinates
        assert_allclose(lon, ds["XLONG"], atol=1e-4)
        assert_allclose(lat, ds["XLAT"], atol=1e-4)

        # then something strange happened
        assert ds.isel(Time=0).salem.grid == ds.salem.grid
        assert ds.isel(Time=0).T2.salem.grid == ds.salem.grid

        nlon, nlat = ds.isel(Time=0).T2.salem.grid.ll_coordinates
        assert_allclose(nlon, ds["XLONG"], atol=1e-4)
        assert_allclose(nlat, ds["XLAT"], atol=1e-4)

        # the grid should not be missunderstood as lonlat
        t2 = ds.T2.isel(Time=0) - 273.15
        with pytest.raises(RuntimeError):
            g = t2.salem.grid
Exemple #2
0
    def test_basic_wrf(self):
        import xarray as xr

        ds = sio.open_xr_dataset(get_demo_file('wrf_tip_d1.nc')).chunk()

        # this is because read_dataset changes some stuff, let's see if
        # georef still ok
        dsxr = xr.open_dataset(get_demo_file('wrf_tip_d1.nc'))
        assert ds.salem.grid == dsxr.salem.grid

        lon, lat = ds.salem.grid.ll_coordinates
        assert_allclose(lon, ds['XLONG'], atol=1e-4)
        assert_allclose(lat, ds['XLAT'], atol=1e-4)

        # then something strange happened
        assert ds.isel(Time=0).salem.grid == ds.salem.grid
        assert ds.isel(Time=0).T2.salem.grid == ds.salem.grid

        nlon, nlat = ds.isel(Time=0).T2.salem.grid.ll_coordinates
        assert_allclose(nlon, ds['XLONG'], atol=1e-4)
        assert_allclose(nlat, ds['XLAT'], atol=1e-4)

        # the grid should not be missunderstood as lonlat
        t2 = ds.T2.isel(Time=0) - 273.15
        with pytest.raises(RuntimeError):
            g = t2.salem.grid
Exemple #3
0
    def test_lookup_transform(self):

        dsw = sio.open_wrf_dataset(get_demo_file('wrfout_d01.nc'))
        dse = sio.open_xr_dataset(get_demo_file('era_interim_tibet.nc'))
        out = dse.salem.lookup_transform(dsw.T2C.isel(time=0), method=len)
        # qualitative tests (quantitative testing done elsewhere)
        assert out[0, 0] == 0
        assert out.mean() > 1

        dsw = sio.open_wrf_dataset(get_demo_file('wrfout_d01.nc'))
        dse = sio.open_xr_dataset(get_demo_file('era_interim_tibet.nc'))
        _, lut = dse.salem.lookup_transform(dsw.T2C.isel(time=0), method=len,
                                              return_lut=True)
        out2 = dse.salem.lookup_transform(dsw.T2C.isel(time=0), method=len,
                                          lut=lut)
        # qualitative tests (quantitative testing done elsewhere)
        assert_allclose(out, out2)
Exemple #4
0
    def test_transform_logic(self):

        # This is just for the naming and dim logic, the rest is tested elsewh
        ds1 = sio.open_wrf_dataset(get_demo_file('wrfout_d01.nc')).chunk()
        ds2 = sio.open_wrf_dataset(get_demo_file('wrfout_d01.nc')).chunk()

        # 2darray case
        t2 = ds2.T2.isel(time=1)
        with pytest.raises(ValueError):
            ds1.salem.transform_and_add(t2.values, grid=t2.salem.grid)

        ds1.salem.transform_and_add(t2.values,
                                    grid=t2.salem.grid,
                                    name='t2_2darr')
        assert 't2_2darr' in ds1
        assert_allclose(ds1.t2_2darr.coords['south_north'],
                        t2.coords['south_north'])
        assert_allclose(ds1.t2_2darr.coords['west_east'],
                        t2.coords['west_east'])
        assert ds1.salem.grid == ds1.t2_2darr.salem.grid

        # 3darray case
        t2 = ds2.T2
        ds1.salem.transform_and_add(t2.values,
                                    grid=t2.salem.grid,
                                    name='t2_3darr')
        assert 't2_3darr' in ds1
        assert_allclose(ds1.t2_3darr.coords['south_north'],
                        t2.coords['south_north'])
        assert_allclose(ds1.t2_3darr.coords['west_east'],
                        t2.coords['west_east'])
        assert 'time' in ds1.t2_3darr.coords

        # dataarray case
        ds1.salem.transform_and_add(t2, name='NEWT2')
        assert 'NEWT2' in ds1
        assert_allclose(ds1.NEWT2, ds1.T2)
        assert_allclose(ds1.t2_3darr.coords['south_north'],
                        t2.coords['south_north'])
        assert_allclose(ds1.t2_3darr.coords['west_east'],
                        t2.coords['west_east'])
        assert 'time' in ds1.t2_3darr.coords

        # dataset case
        ds1.salem.transform_and_add(ds2[['RAINC', 'RAINNC']],
                                    name={
                                        'RAINC': 'PRCPC',
                                        'RAINNC': 'PRCPNC'
                                    })
        assert 'PRCPC' in ds1
        assert_allclose(ds1.PRCPC, ds1.RAINC)
        assert 'time' in ds1.PRCPNC.coords

        # what happens with external data?
        dse = sio.open_xr_dataset(get_demo_file('era_interim_tibet.nc'))
        out = ds1.salem.transform(dse.t2m, interp='linear')
        assert_allclose(out.coords['south_north'], t2.coords['south_north'])
        assert_allclose(out.coords['west_east'], t2.coords['west_east'])
Exemple #5
0
    def test_prcp(self):

        wf = get_demo_file("wrfout_d01.nc")

        w = sio.open_wrf_dataset(wf)
        nc = sio.open_xr_dataset(wf)

        nc["REF_PRCP_NC"] = nc["RAINNC"] * 0.0
        uns = (
            nc["RAINNC"].isel(Time=slice(1, len(nc.bottom_top_stag))).values
            - nc["RAINNC"].isel(Time=slice(0, -1)).values
        )
        nc["REF_PRCP_NC"].values[1:, ...] = uns * 60 / 180.0  # for three hours
        nc["REF_PRCP_NC"].values[0, ...] = np.NaN

        nc["REF_PRCP_C"] = nc["RAINC"] * 0.0
        uns = (
            nc["RAINC"].isel(Time=slice(1, len(nc.bottom_top_stag))).values - nc["RAINC"].isel(Time=slice(0, -1)).values
        )
        nc["REF_PRCP_C"].values[1:, ...] = uns * 60 / 180.0  # for three hours
        nc["REF_PRCP_C"].values[0, ...] = np.NaN

        nc["REF_PRCP"] = nc["REF_PRCP_C"] + nc["REF_PRCP_NC"]

        for suf in ["_NC", "_C", ""]:

            assert_allclose(w["PRCP" + suf], nc["REF_PRCP" + suf], rtol=1e-5)

            wn = w.isel(time=slice(1, 3))
            ncn = nc.isel(Time=slice(1, 3))
            assert_allclose(wn["PRCP" + suf], ncn["REF_PRCP" + suf], rtol=1e-5)

            wn = w.isel(time=2)
            ncn = nc.isel(Time=2)
            assert_allclose(wn["PRCP" + suf], ncn["REF_PRCP" + suf], rtol=1e-5)

            wn = w.isel(time=1)
            ncn = nc.isel(Time=1)
            assert_allclose(wn["PRCP" + suf], ncn["REF_PRCP" + suf], rtol=1e-5)

            wn = w.isel(time=0)
            self.assertTrue(~np.any(np.isfinite(wn["PRCP" + suf].values)))

            wn = w.isel(time=slice(1, 3), south_north=slice(50, -1))
            ncn = nc.isel(Time=slice(1, 3), south_north=slice(50, -1))
            assert_allclose(wn["PRCP" + suf], ncn["REF_PRCP" + suf], rtol=1e-5)

            wn = w.isel(time=2, south_north=slice(50, -1))
            ncn = nc.isel(Time=2, south_north=slice(50, -1))
            assert_allclose(wn["PRCP" + suf], ncn["REF_PRCP" + suf], rtol=1e-5)

            wn = w.isel(time=1, south_north=slice(50, -1))
            ncn = nc.isel(Time=1, south_north=slice(50, -1))
            assert_allclose(wn["PRCP" + suf], ncn["REF_PRCP" + suf], rtol=1e-5)

            wn = w.isel(time=0, south_north=slice(50, -1))
            self.assertTrue(~np.any(np.isfinite(wn["PRCP" + suf].values)))
Exemple #6
0
 def test_roi(self):
     import xarray as xr
     # Check that all attrs are preserved
     with sio.open_xr_dataset(get_demo_file('era_interim_tibet.nc')) as ds:
         ds.encoding = {'_FillValue': np.NaN}
         ds['t2m'].encoding = {'_FillValue': np.NaN}
         ds_ = ds.salem.roi(roi=np.ones_like(ds.t2m.values[0, ...]))
         xr.testing.assert_identical(ds, ds_)
         assert ds.encoding == ds_.encoding
         assert ds.t2m.encoding == ds_.t2m.encoding
Exemple #7
0
    def test_prcp(self):

        wf = get_demo_file('wrfout_d01.nc')

        w = sio.open_wrf_dataset(wf).chunk()
        nc = sio.open_xr_dataset(wf)

        nc['REF_PRCP_NC'] = nc['RAINNC'] * 0.
        uns = nc['RAINNC'].isel(Time=slice(1, len(nc.bottom_top_stag))).values - \
              nc['RAINNC'].isel(Time=slice(0, -1)).values
        nc['REF_PRCP_NC'].values[1:, ...] = uns * 60 / 180.  # for three hours
        nc['REF_PRCP_NC'].values[0, ...] = np.NaN

        nc['REF_PRCP_C'] = nc['RAINC'] * 0.
        uns = nc['RAINC'].isel(Time=slice(1, len(nc.bottom_top_stag))).values - \
              nc['RAINC'].isel(Time=slice(0, -1)).values
        nc['REF_PRCP_C'].values[1:, ...] = uns * 60 / 180.  # for three hours
        nc['REF_PRCP_C'].values[0, ...] = np.NaN

        nc['REF_PRCP'] = nc['REF_PRCP_C'] + nc['REF_PRCP_NC']

        for suf in ['_NC', '_C', '']:

            assert_allclose(w['PRCP' + suf], nc['REF_PRCP' + suf], rtol=1e-5)

            wn = w.isel(time=slice(1, 3))
            ncn = nc.isel(Time=slice(1, 3))
            assert_allclose(wn['PRCP' + suf], ncn['REF_PRCP' + suf], rtol=1e-5)

            wn = w.isel(time=2)
            ncn = nc.isel(Time=2)
            assert_allclose(wn['PRCP' + suf], ncn['REF_PRCP' + suf], rtol=1e-5)

            wn = w.isel(time=1)
            ncn = nc.isel(Time=1)
            assert_allclose(wn['PRCP' + suf], ncn['REF_PRCP' + suf], rtol=1e-5)

            wn = w.isel(time=0)
            self.assertTrue(~np.any(np.isfinite(wn['PRCP' + suf].values)))

            wn = w.isel(time=slice(1, 3), south_north=slice(50, -1))
            ncn = nc.isel(Time=slice(1, 3), south_north=slice(50, -1))
            assert_allclose(wn['PRCP' + suf], ncn['REF_PRCP' + suf], rtol=1e-5)

            wn = w.isel(time=2, south_north=slice(50, -1))
            ncn = nc.isel(Time=2, south_north=slice(50, -1))
            assert_allclose(wn['PRCP' + suf], ncn['REF_PRCP' + suf], rtol=1e-5)

            wn = w.isel(time=1, south_north=slice(50, -1))
            ncn = nc.isel(Time=1, south_north=slice(50, -1))
            assert_allclose(wn['PRCP' + suf], ncn['REF_PRCP' + suf], rtol=1e-5)

            wn = w.isel(time=0, south_north=slice(50, -1))
            self.assertTrue(~np.any(np.isfinite(wn['PRCP' + suf].values)))
Exemple #8
0
    def test_era(self):

        ds = sio.open_xr_dataset(get_demo_file('era_interim_tibet.nc')).chunk()
        self.assertEqual(ds.salem.x_dim, 'longitude')
        self.assertEqual(ds.salem.y_dim, 'latitude')

        dss = ds.salem.subset(ds=ds)
        self.assertEqual(dss.salem.grid, ds.salem.grid)

        lon = 91.1
        lat = 31.1
        dss = ds.salem.subset(corners=((lon, lat), (lon, lat)), margin=1)

        self.assertEqual(len(dss.latitude), 3)
        self.assertEqual(len(dss.longitude), 3)

        np.testing.assert_almost_equal(dss.longitude, [90.0, 90.75, 91.5])
Exemple #9
0
    def test_era(self):

        ds = sio.open_xr_dataset(get_demo_file("era_interim_tibet.nc"))
        self.assertEqual(ds.salem.x_dim, "longitude")
        self.assertEqual(ds.salem.y_dim, "latitude")

        dss = ds.salem.subset(ds=ds)
        self.assertEqual(dss.salem.grid, ds.salem.grid)

        lon = 91.1
        lat = 31.1
        dss = ds.salem.subset(corners=((lon, lat), (lon, lat)), margin=1)

        self.assertEqual(len(dss.latitude), 3)
        self.assertEqual(len(dss.longitude), 3)

        np.testing.assert_almost_equal(dss.longitude, [90.0, 90.75, 91.5])
Exemple #10
0
    def test_unstagger(self):

        wf = get_demo_file("wrf_cropped.nc")

        w = sio.open_wrf_dataset(wf)
        nc = sio.open_xr_dataset(wf)

        nc["PH_UNSTAGG"] = nc["P"] * 0.0
        uns = (
            nc["PH"].isel(bottom_top_stag=slice(0, -1)).values
            + nc["PH"].isel(bottom_top_stag=slice(1, len(nc.bottom_top_stag))).values
        )
        nc["PH_UNSTAGG"].values = uns * 0.5

        assert_allclose(w["PH"], nc["PH_UNSTAGG"])

        # chunk
        v = w["PH"].chunk((1, 6, 13, 13))
        assert_allclose(v.mean(), nc["PH_UNSTAGG"].mean(), atol=1e-2)

        wn = w.isel(west_east=slice(4, 8))
        ncn = nc.isel(west_east=slice(4, 8))
        assert_allclose(wn["PH"], ncn["PH_UNSTAGG"])

        wn = w.isel(south_north=slice(4, 8), time=1)
        ncn = nc.isel(south_north=slice(4, 8), Time=1)
        assert_allclose(wn["PH"], ncn["PH_UNSTAGG"])

        wn = w.isel(west_east=4)
        ncn = nc.isel(west_east=4)
        assert_allclose(wn["PH"], ncn["PH_UNSTAGG"])

        wn = w.isel(bottom_top=4)
        ncn = nc.isel(bottom_top=4)
        assert_allclose(wn["PH"], ncn["PH_UNSTAGG"])

        wn = w.isel(bottom_top=0)
        ncn = nc.isel(bottom_top=0)
        assert_allclose(wn["PH"], ncn["PH_UNSTAGG"])

        wn = w.isel(bottom_top=-1)
        ncn = nc.isel(bottom_top=-1)
        assert_allclose(wn["PH"], ncn["PH_UNSTAGG"])

        w["PH"].chunk()
Exemple #11
0
    def test_transform_logic(self):

        # This is just for the naming and dim logic, the rest is tested elsewh
        ds1 = sio.open_wrf_dataset(get_demo_file("wrfout_d01.nc"))
        ds2 = sio.open_wrf_dataset(get_demo_file("wrfout_d01.nc"))

        # 2darray case
        t2 = ds2.T2.isel(time=1)
        with pytest.raises(ValueError):
            ds1.salem.transform_and_add(t2.values, grid=t2.salem.grid)

        ds1.salem.transform_and_add(t2.values, grid=t2.salem.grid, name="t2_2darr")
        assert "t2_2darr" in ds1
        assert_allclose(ds1.t2_2darr.coords["south_north"], t2.coords["south_north"])
        assert_allclose(ds1.t2_2darr.coords["west_east"], t2.coords["west_east"])
        assert ds1.salem.grid == ds1.t2_2darr.salem.grid

        # 3darray case
        t2 = ds2.T2
        ds1.salem.transform_and_add(t2.values, grid=t2.salem.grid, name="t2_3darr")
        assert "t2_3darr" in ds1
        assert_allclose(ds1.t2_3darr.coords["south_north"], t2.coords["south_north"])
        assert_allclose(ds1.t2_3darr.coords["west_east"], t2.coords["west_east"])
        assert "time" in ds1.t2_3darr.coords

        # dataarray case
        ds1.salem.transform_and_add(t2, name="NEWT2")
        assert "NEWT2" in ds1
        assert_allclose(ds1.NEWT2, ds1.T2)
        assert_allclose(ds1.t2_3darr.coords["south_north"], t2.coords["south_north"])
        assert_allclose(ds1.t2_3darr.coords["west_east"], t2.coords["west_east"])
        assert "time" in ds1.t2_3darr.coords

        # dataset case
        ds1.salem.transform_and_add(ds2[["RAINC", "RAINNC"]], name={"RAINC": "PRCPC", "RAINNC": "PRCPNC"})
        assert "PRCPC" in ds1
        assert_allclose(ds1.PRCPC, ds1.RAINC)
        assert "time" in ds1.PRCPNC.coords

        # what happens with external data?
        dse = sio.open_xr_dataset(get_demo_file("era_interim_tibet.nc"))
        out = ds1.salem.transform(dse.t2m, interp="linear")
        assert_allclose(out.coords["south_north"], t2.coords["south_north"])
        assert_allclose(out.coords["west_east"], t2.coords["west_east"])
Exemple #12
0
    def test_unstagger(self):

        wf = get_demo_file('wrf_cropped.nc')

        w = sio.open_wrf_dataset(wf).chunk()
        nc = sio.open_xr_dataset(wf).chunk()

        nc['PH_UNSTAGG'] = nc['P'] * 0.
        uns = nc['PH'].isel(bottom_top_stag=slice(0, -1)).values + \
              nc['PH'].isel(bottom_top_stag=slice(1, len(nc.bottom_top_stag))).values
        nc['PH_UNSTAGG'].values = uns * 0.5

        assert_allclose(w['PH'], nc['PH_UNSTAGG'])

        # chunk
        v = w['PH'].chunk((1, 6, 13, 13))
        assert_allclose(v.mean(), nc['PH_UNSTAGG'].mean(), atol=1e-2)

        wn = w.isel(west_east=slice(4, 8))
        ncn = nc.isel(west_east=slice(4, 8))
        assert_allclose(wn['PH'], ncn['PH_UNSTAGG'])

        wn = w.isel(south_north=slice(4, 8), time=1)
        ncn = nc.isel(south_north=slice(4, 8), Time=1)
        assert_allclose(wn['PH'], ncn['PH_UNSTAGG'])

        wn = w.isel(west_east=4)
        ncn = nc.isel(west_east=4)
        assert_allclose(wn['PH'], ncn['PH_UNSTAGG'])

        wn = w.isel(bottom_top=4)
        ncn = nc.isel(bottom_top=4)
        assert_allclose(wn['PH'], ncn['PH_UNSTAGG'])

        wn = w.isel(bottom_top=0)
        ncn = nc.isel(bottom_top=0)
        assert_allclose(wn['PH'], ncn['PH_UNSTAGG'])

        wn = w.isel(bottom_top=-1)
        ncn = nc.isel(bottom_top=-1)
        assert_allclose(wn['PH'], ncn['PH_UNSTAGG'])

        w['PH'].chunk()
Exemple #13
0
    def test_ncl_diagvars(self):

        wf = get_demo_file('wrf_cropped.nc')
        ncl_out = get_demo_file('wrf_cropped_ncl.nc')

        w = sio.open_wrf_dataset(wf)
        nc = sio.open_xr_dataset(ncl_out)

        ref = nc['TK']
        tot = w['TK']
        assert_allclose(ref, tot, rtol=1e-6)

        ref = nc['SLP']
        tot = w['SLP']
        tot = tot.values
        assert_allclose(ref, tot, rtol=1e-6)

        w = w.isel(time=1, south_north=slice(12, 16), west_east=slice(9, 16))
        nc = nc.isel(Time=1, south_north=slice(12, 16), west_east=slice(9, 16))

        ref = nc['TK']
        tot = w['TK']
        assert_allclose(ref, tot, rtol=1e-6)

        ref = nc['SLP']
        tot = w['SLP']
        tot = tot.values
        assert_allclose(ref, tot, rtol=1e-6)

        w = w.isel(bottom_top=slice(3, 5))
        nc = nc.isel(bottom_top=slice(3, 5))

        ref = nc['TK']
        tot = w['TK']
        assert_allclose(ref, tot, rtol=1e-6)

        ref = nc['SLP']
        tot = w['SLP']
        tot = tot.values
        assert_allclose(ref, tot, rtol=1e-6)
Exemple #14
0
    def test_ncl_diagvars(self):

        wf = get_demo_file("wrf_cropped.nc")
        ncl_out = get_demo_file("wrf_cropped_ncl.nc")

        w = sio.open_wrf_dataset(wf)
        nc = sio.open_xr_dataset(ncl_out)

        ref = nc["TK"]
        tot = w["TK"]
        assert_allclose(ref, tot, rtol=1e-6)

        ref = nc["SLP"]
        tot = w["SLP"]
        tot = tot.values
        assert_allclose(ref, tot, rtol=1e-6)

        w = w.isel(time=1, south_north=slice(12, 16), west_east=slice(9, 16))
        nc = nc.isel(Time=1, south_north=slice(12, 16), west_east=slice(9, 16))

        ref = nc["TK"]
        tot = w["TK"]
        assert_allclose(ref, tot, rtol=1e-6)

        ref = nc["SLP"]
        tot = w["SLP"]
        tot = tot.values
        assert_allclose(ref, tot, rtol=1e-6)

        w = w.isel(bottom_top=slice(3, 5))
        nc = nc.isel(bottom_top=slice(3, 5))

        ref = nc["TK"]
        tot = w["TK"]
        assert_allclose(ref, tot, rtol=1e-6)

        ref = nc["SLP"]
        tot = w["SLP"]
        tot = tot.values
        assert_allclose(ref, tot, rtol=1e-6)