Small bugfixes in exotic WRF diag variables indexing (#89)

docs/whats-new.rst

@@ -32,6 +32,7 @@ Bug fixes
   (:issue:`74`)
 - small bugfixes in the projections and warning handling
 - PRESSURE variable was given in Pa, not hPa. This is corrected now.
+- Small bugfixes in exotic WRF diag variables indexing
 
 
 v0.2.1 (07 February 2017)

salem/sio.py

@@ -1108,8 +1108,7 @@ def open_mf_wrf_dataset(paths, chunks=None,  compat='no_conflicts', lock=None,
     if preprocess is not None:
         datasets = [preprocess(ds) for ds in datasets]
 
-    # TODO: add compat=compat when xarray 9.0 is out
-    combined = xr.auto_combine(datasets, concat_dim='time')
+    combined = xr.auto_combine(datasets, concat_dim='time', compat=compat)
     combined._file_obj = _MultiFileCloser(file_objs)
     combined.attrs = datasets[0].attrs
 

salem/tests/test_misc.py

@@ -504,7 +504,7 @@ def tearDown(self):
     @requires_xarray
     def test_era(self):
 
-        ds = sio.open_xr_dataset(get_demo_file('era_interim_tibet.nc'))
+        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')
 
@@ -525,7 +525,7 @@ def test_era(self):
     def test_basic_wrf(self):
         import xarray as xr
 
-        ds = sio.open_xr_dataset(get_demo_file('wrf_tip_d1.nc'))
+        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
@@ -554,7 +554,7 @@ def test_geo_em(self):
 
         for i in [1, 2, 3]:
             fg = get_demo_file('geo_em_d0{}_lambert.nc'.format(i))
-            ds = sio.open_wrf_dataset(fg)
+            ds = sio.open_wrf_dataset(fg).chunk()
             self.assertFalse('Time' in ds.dims)
             self.assertTrue('time' in ds.dims)
             self.assertTrue('south_north' in ds.dims)
@@ -565,7 +565,7 @@ def test_geo_em(self):
     def test_wrf(self):
         import xarray as xr
 
-        ds = sio.open_wrf_dataset(get_demo_file('wrf_tip_d1.nc'))
+        ds = sio.open_wrf_dataset(get_demo_file('wrf_tip_d1.nc')).chunk()
 
         # this is because read_dataset changes some stuff, let's see if
         # georef still ok
@@ -597,7 +597,7 @@ def test_ncl_diagvars(self):
         ncl_out = get_demo_file('wrf_cropped_ncl.nc')
 
         with warnings.catch_warnings(record=True) as war:
-            w = sio.open_wrf_dataset(wf)
+            w = sio.open_wrf_dataset(wf).chunk()
             self.assertEqual(len(war), 0)
         nc = xr.open_dataset(ncl_out)
 
@@ -639,8 +639,8 @@ def test_unstagger(self):
 
         wf = get_demo_file('wrf_cropped.nc')
 
-        w = sio.open_wrf_dataset(wf)
-        nc = sio.open_xr_dataset(wf)
+        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 + \
@@ -683,19 +683,7 @@ def test_unstagger(self):
     def test_diagvars(self):
 
         wf = get_demo_file('wrf_d01_allvars_cropped.nc')
-        w = sio.open_wrf_dataset(wf)
-
-        # ws
-        w['ws_ref'] = np.sqrt(w['U']**2 + w['V']**2)
-        assert_allclose(w['ws_ref'], w['WS'])
-        wcrop = w.isel(west_east=slice(4, 8), bottom_top=4)
-        assert_allclose(wcrop['ws_ref'], wcrop['WS'])
-
-    @requires_xarray
-    def test_diagvars(self):
-
-        wf = get_demo_file('wrf_d01_allvars_cropped.nc')
-        w = sio.open_wrf_dataset(wf)
+        w = sio.open_wrf_dataset(wf).chunk()
 
         # ws
         w['ws_ref'] = np.sqrt(w['U']**2 + w['V']**2)
@@ -708,7 +696,7 @@ def test_prcp(self):
 
         wf = get_demo_file('wrfout_d01.nc')
 
-        w = sio.open_wrf_dataset(wf)
+        w = sio.open_wrf_dataset(wf).chunk()
         nc = sio.open_xr_dataset(wf)
 
         nc['REF_PRCP_NC'] = nc['RAINNC']*0.
@@ -764,8 +752,8 @@ def test_prcp(self):
     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'))
+        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)
@@ -820,8 +808,8 @@ def test_transform_logic(self):
     @requires_geopandas
     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'))
+        dsw = sio.open_wrf_dataset(get_demo_file('wrfout_d01.nc')).chunk()
+        dse = sio.open_xr_dataset(get_demo_file('era_interim_tibet.nc')).chunk()
         out = dse.salem.lookup_transform(dsw.T2C.isel(time=0), method=len)
         # qualitative tests (quantitative testing done elsewhere)
         assert out[0, 0] == 0
@@ -843,7 +831,7 @@ def test_full_wrf_wfile(self):
 
         # TODO: these tests are qualitative and should be compared against ncl
         f = get_demo_file('wrf_d01_allvars_cropped.nc')
-        ds = sio.open_wrf_dataset(f)
+        ds = sio.open_wrf_dataset(f).chunk()
 
         # making a repr was causing trouble because of the small chunks
         _ = ds.__repr__()
@@ -871,7 +859,7 @@ def test_full_wrf_wfile(self):
     def test_3d_interp(self):
 
         f = get_demo_file('wrf_d01_allvars_cropped.nc')
-        ds = sio.open_wrf_dataset(f)
+        ds = sio.open_wrf_dataset(f).chunk()
 
         out = ds.salem.wrf_zlevel('Z', levels=6000.)
         ref_2d = out*0. + 6000.

salem/wrftools.py

@@ -112,7 +112,7 @@ def __getitem__(self, item):
         else:
             item[self.ds] = slice(start, stop-1)
             itemr[self.ds] = slice(start+1, stop)
-        return 0.5*(self.ncvar[item] + self.ncvar[itemr])
+        return 0.5*(self.ncvar[tuple(item)] + self.ncvar[tuple(itemr)])
 
 
 class FakeVariable(object):
@@ -212,7 +212,9 @@ def __getitem__(self, item):
 
         # time is always going to be first dim I hope
         sl = item[0]
+        was_scalar = False
         if np.isscalar(sl) and not isinstance(sl, slice):
+            was_scalar = True
             sl = slice(sl, sl+1)
 
         # Ok, get the indexes right
@@ -243,6 +245,8 @@ def __getitem__(self, item):
         else:
             out = var[itemr]
             out -= var[item]
+        if was_scalar:
+            out = out[0, ...]
         return out * self._factor