made random_walk_idxs compatible with np.slice

Author: ihincks

src/qinfer/derived_models.py

@@ -620,8 +620,10 @@ class GaussianRandomWalkModel(DerivedModel):
     
     :param Model underlying_model: Model representing the likelihood with no
         random walk added.
-    :param random_walk_idxs: A list of model parameter indeces to add the 
-        random walk to.
+    :param random_walk_idxs: A list or ``np.slice`` of 
+        ``underlying_model`` model parameter indeces to add the random walk to.
+        Indeces larger than ``underlying_model.n_modelparams`` should not 
+        be touched.
     :param fixed_covariance: An ``np.ndarray`` specifying the fixed covariance 
         matrix (or diagonal thereof if ``diagonal`` is ``True``) of the 
         gaussian distribution. If set to ``None`` (default), this matrix is 
@@ -649,14 +651,18 @@ def __init__(
         ):
 
         self._diagonal = diagonal
-        self._rw_idxs = np.arange(underlying_model.n_modelparams).astype(np.int) \
+        self._rw_idxs = np.s_[:underlying_model.n_modelparams] \
             if random_walk_idxs == 'all' else random_walk_idxs
 
+        explicit_idxs = np.arange(underlying_model.n_modelparams)[self._rw_idxs]
+        if explicit_idxs.size == 0:
+            raise IndexError('At least one model parameter must take a random walk.')
+    
         self._rw_names = [
                 underlying_model.modelparam_names[idx] 
-                for idx in self._rw_idxs
+                for idx in explicit_idxs
             ]
-        self._n_rw = len(self._rw_idxs)
+        self._n_rw = len(explicit_idxs)
 
         self._srw_names = []
         if fixed_covariance is None:
@@ -700,6 +706,9 @@ def __init__(
                 )
 
         super(GaussianRandomWalkModel, self).__init__(underlying_model)
+        
+        if np.max(np.arange(self.n_modelparams)[self._rw_idxs]) > np.max(explicit_idxs):
+            raise IndexError('random_walk_idxs out of bounds; must index (a subset of ) underlying_model modelparams.')
 
         if scale_mult is None:
             self._scale_mult_fcn = (lambda expparams: 1)

src/qinfer/tests/test_concrete_models.py

@@ -407,7 +407,7 @@ def instantiate_model(self):
             m,
             fixed_covariance = cov,
             diagonal = False,
-            random_walk_idxs = [1,2,4],
+            random_walk_idxs = np.s_[:6:2],
             model_transformation = (from_simplex, to_simplex),
             scale_mult = 'n_meas'
         )
@@ -480,3 +480,23 @@ def test_est_update_covariance(self):
         cov = self.model.est_update_covariance(self.modelparams)
         eigs, v = np.linalg.eig(cov)
         assert(np.greater_equal(eigs, -1e-10).all())
+        
+class TestGaussianRandomWalkModel5(DerandomizedTestCase):
+    """
+    Tests miscillaneous properties of GaussianRandomWalkModel.
+    """
+
+    def test_indexing(self):
+        model = lambda slice: GaussianRandomWalkModel(
+                MultinomialModel(NDieModel(n=6)), 
+                random_walk_idxs = slice
+            )
+
+        assert(model('all').n_modelparams == 12)
+        assert(model(np.s_[:6]).n_modelparams == 12)
+        assert(model(np.s_[:6:2]).n_modelparams == 9)
+        assert(model([2,3,4]).n_modelparams == 9)
+        
+        self.assertRaises(IndexError, model, np.s_[:7])
+        self.assertRaises(IndexError, model, np.s_[6:])
+        self.assertRaises(IndexError, model, [1,2,8])