Merge pull request #30 from NREL/boolean_masks

Enable slicing with boolean masks

rex/resource.py

@@ -134,6 +134,9 @@ def _check_slice(ds_slice):
         ds_idx = None
         if isinstance(ds_slice, (list, np.ndarray)):
             in_slice = np.array(ds_slice)
+            if np.issubdtype(in_slice.dtype, np.dtype(bool)):
+                in_slice = np.where(in_slice)[0]
+
             s = in_slice.min()
             e = in_slice.max() + 1
             ds_slice = slice(s, e, None)

rex/version.py

@@ -1,3 +1,3 @@
 """rex Version number"""
 
-__version__ = "0.2.12"
+__version__ = "0.2.13"

tests/test_resource.py

@@ -183,6 +183,12 @@ def check_dset(res_cls, ds_name):
     assert isinstance(ds, np.ndarray)
     assert ds.shape == (100,)
     assert np.allclose(arr[times, 0], ds)
+    # boolean mask
+    mask = res_cls.time_index.month == 7
+    ds = res_cls[ds_name, mask]
+    assert isinstance(ds, np.ndarray)
+    assert ds.shape == (mask.sum(), ds_shape[1])
+    assert np.allclose(arr[mask], ds)
     # time and site lists
     with pytest.raises(IndexError):
         assert res_cls[ds_name, times, sites]