refine numerical corrections in guassion weighting; fix other bundles…

… code to consider affine

AFQ/_fixes.py

@@ -2,6 +2,7 @@
 import logging
 
 from scipy.special import lpmv, gammaln
+from scipy.linalg import pinvh
 
 from tqdm import tqdm
 
@@ -215,17 +216,20 @@ def gaussian_weights(bundle, n_points=100, return_mahalnobis=False,
     for i in range(n_nodes):
         # This should come back as a 3D covariance matrix with the spatial
         # variance covariance of this node across the different streamlines,
-        # reorganized as an upper diagonal matrix for expected Mahalanobis
+        # converted to a positive semi-definite matrix if necessary
         cov = np.cov(sls[:, i, :].T, ddof=0)
-        while np.any(np.linalg.eigvals(cov) < 0):
-            cov += np.eye(cov.shape[0]) * 1e-12
+        if np.any(np.linalg.eigvals(cov) < 0):
+            eigenvalues, eigenvectors = np.linalg.eigh((cov + cov.T) / 2)
+            eigenvalues[eigenvalues < 0] = 0
+            cov = eigenvectors @ np.diag(eigenvalues) @ eigenvectors.T
 
         # calculate Mahalanobis for node in every fiber
         if np.any(cov > 0):
-            ci = np.linalg.inv(cov)
-
-            dist = (diff[:, i, :] @ ci) * diff[:, i, :]
-            weights[:, i] = np.sqrt(np.sum(dist, axis=1))
+            weights[:, i] = np.sqrt(np.einsum(
+                'ij,jk,ik->i',
+                diff[:, i, :],
+                pinvh(cov),
+                diff[:, i, :]))
 
         # In the special case where all the streamlines have the exact same
         # coordinate in this node, the covariance matrix is all zeros, so

AFQ/recognition/criteria.py

@@ -304,7 +304,8 @@ def clean_by_other_bundle(b_sls, bundle_def,
         cleaned_idx_core = abo.clean_relative_to_other_core(
             bundle_def[other_bundle_name]['core'].lower(),
             preproc_imap["fgarray"][b_sls.selected_fiber_idxs],
-            np.array(abu.resample_tg(other_bundle_sls, 20)))
+            np.array(abu.resample_tg(other_bundle_sls, 20)),
+            img.affine)
         cleaned_idx = np.logical_and(cleaned_idx, cleaned_idx_core)
 
     b_sls.select(cleaned_idx, other_bundle_name)

AFQ/recognition/other_bundles.py

@@ -24,7 +24,7 @@ def clean_by_other_density_map(this_bundle_sls,
     return cleaned_idx
 
 
-def clean_relative_to_other_core(core, this_fgarray, other_fgarray):
+def clean_relative_to_other_core(core, this_fgarray, other_fgarray, affine):
     """
     Remove any fibers that are on the wrong side of the core
     """
@@ -47,6 +47,10 @@ def clean_relative_to_other_core(core, this_fgarray, other_fgarray):
         core_axis = 0
         core_direc = 1
 
+    for ii in range(3):
+        if affine[ii, ii] < 0:
+            core_direc = -core_direc
+
     core_bundle = np.median(other_fgarray, axis=0)
     cleaned_idx_core = np.zeros(this_fgarray.shape[0], dtype=np.bool8)
     for ii, sl in enumerate(this_fgarray):

AFQ/recognition/tests/test_other_bundles.py

@@ -31,7 +31,8 @@ def test_clean_relative_to_other_core():
         cleaned_idx_core = abo.clean_relative_to_other_core(
             core,
             this_bundle_sls_sample,
-            other_bundle_sls_sample
+            other_bundle_sls_sample,
+            np.eye(4)
         )
 
         assert isinstance(cleaned_idx_core, np.ndarray)