BF: fix code in format_matrix, update tests

Author: CindeeM

brainx/tests/test_util.py

@@ -92,10 +92,19 @@ def test_cost2thresh(self):
         npt.assert_almost_equal(thr, real_thr, decimal=7)
 
     def test_format_matrix(self):
-        thresh_matrix = util.format_matrix2(self.data_5d, 0,0,0,
+        bool_matrix = util.format_matrix2(self.data_5d, 0,0,0,
                 self.lookup, self.costs[100])
-        print thresh_matrix.sum()
-        npt.assert_equal(thresh_matrix.sum(), 100 -1)
+        npt.assert_equal(bool_matrix.sum(), 100 -1)
+        thresh_matrix = util.format_matrix2(self.data_5d, 0,0,0,
+                self.lookup, self.costs[100],asbool = False)
+        npt.assert_equal(bool_matrix.sum()== thresh_matrix.sum(), False)
+        npt.assert_almost_equal(thresh_matrix.sum(), 
+                94.183321784530804, decimal=7)
+        ## test format_matrix call on format_matrix2
+        bool_matrix_sm = util.format_matrix(self.data_5d[0].squeeze(),
+                0,0, self.lookup[0].squeeze(), self.costs[100])
+        npt.assert_equal(bool_matrix.sum(), bool_matrix_sm.sum())
+
 
     def test_threshold_adjacency_matrix(self):
         adj_matrix = self.data_5d[0,0,0].squeeze()

brainx/util.py

@@ -37,7 +37,8 @@ def slice_data(data, sub, block, subcond=None):
     return data[subcond, block, sub]
 
 
-def format_matrix(data,s,b,lk,co,idc = [],costlist=[],nouptri = False):
+def format_matrix(data, s, b, lk, co, idc=[], costlist=[],
+        nouptri=False, asbool=True):
     """ Function which thresholds the adjacency matrix for a particular 
     subject and particular block, using lookuptable to find thresholds, 
     cost value to find threshold, costlist
@@ -46,24 +47,37 @@ def format_matrix(data,s,b,lk,co,idc = [],costlist=[],nouptri = False):
     Parameters
     -----------
     data : full data array 4D (block, sub, node, node)
-    s = subject
-    b = block
-    lk = lookup table for study
-    co = cost value to threshold at
-    idc = index of ideal cost
-    costlist = list (size num_edges) with ordered values used to find 
+    s : int
+        subject
+    b : int
+        block
+    lk : numpy array
+        lookup table for study
+    co : int
+        cost value to threshold at
+    idc : int
+        index of ideal cost
+    costlist : list
+        list (size num_edges) with ordered values used to find 
         threshold to control number of edges
-    nouptri = if False only keeps upper tri, True yields symmetric matrix
+    nouptri : bool
+        if False only keeps upper tri, True yields symmetric matrix
+    asbool : bool
+        if True return boolean mask, otherwise returns thesholded
+        weight matrix
     """
-    cmat = data[b,s]
+    cmat = slice_data(data, b,s)
     th = cost2thresh(co,s,b,lk,idc,costlist) #get the right threshold
     cmat = thresholded_arr(cmat,th,fill_val=0)
     if not nouptri:
         cmat = np.triu(cmat,1)
+    if asbool:
+        return ~(cmat == 0)
     return cmat
 
 
-def format_matrix2(data,s,sc,c,lk,co,idc = [],costlist=[],nouptri = False):
+def format_matrix2(data, s, sc, c, lk, co, idc=[],
+        costlist=[], nouptri=False, asbool=True):
     """ Function which formats matrix for a particular subject and 
     particular block (thresholds, upper-tris it) so that we can 
     make a graph object out of it
@@ -88,14 +102,19 @@ def format_matrix2(data,s,sc,c,lk,co,idc = [],costlist=[],nouptri = False):
         list of possible costs
     nouptri : bool
         False zeros out diag and below, True returns symmetric matrix
+    asbool : bool
+        If true returns boolean mask, otherwise returns thresholded w
+        weighted matrix
     """
     cmat = slice_data(data, s, c, sc) 
     th = cost2thresh2(co,s,sc,c,lk,[],idc,costlist) #get the right threshold
     cmat = thresholded_arr(cmat,th,fill_val=0)
     if not nouptri:
         cmat = np.triu(cmat,1)
-    # return boolean mask
-    return ~(cmat == 0)
+    if asbool:
+        # return boolean mask
+        return ~(cmat == 0)
+    return cmat
 
 def threshold_adjacency_matrix(adj_matrix, cost):
     """docstring for threshold_adjacency_matrix(adj_matrix, cost"""