HACK: add diag_indices helper for numpy < 1.4.0

matthew-brett · matthew-brett · commit 4e913137a2d8 · 2012-07-20T16:27:34.000-07:00
Use of diag_indices causing error for numpy &lt; 1.4.0; replicate in
private function.  Sometime soon after the release of Debian Wheezy as
stable, we should drop numpy 1.3.0 compatibility and delete this helper.
diff --git a/nipy/algorithms/segmentation/segmentation.py b/nipy/algorithms/segmentation/segmentation.py
@@ -106,7 +106,7 @@ def set_markov_prior(self, beta, U=None):
             self.U = np.asarray(U).copy()
         else:  # Potts model
             U = np.ones((self.nclasses, self.nclasses))
-            U[np.diag_indices(self.nclasses)] = 0
+            U[_diag_indices(self.nclasses)] = 0
             self.U = U
         self.beta = float(beta)
 
@@ -206,6 +206,13 @@ def free_energy(self, ppm=None):
         return f1 + f2
 
 
+def _diag_indices(n, ndim=2):
+    # diag_indices function present in numpy 1.4 and later.  This for
+    # compatibility with numpy < 1.4
+    idx = np.arange(n)
+    return (idx,) * ndim
+
+
 def moment_matching(dat, mu, sigma, glob_mu, glob_sigma):
     """
     Moment matching strategy for parameter initialization to feed a
