[ENH] add median bundle len function

AFQ/api/utils.py

@@ -125,6 +125,7 @@ def export_all_helper(api_afq_object, seg_algo, xforms, indiv, viz):
         if seg_algo == "AFQ":
             api_afq_object.export("rois")
     api_afq_object.export("sl_counts")
+    api_afq_object.export("median_bundle_lengths")
     api_afq_object.export("profiles")
 
     if viz:

AFQ/tasks/segmentation.py

@@ -214,6 +214,46 @@ def export_sl_counts(data_imap,
     return counts_df, dict(sources=bundles_files)
 
 
+@pimms.calc("median_bundle_lengths")
+@as_file('_medianBundleLengths.csv', include_track=True, include_seg=True)
+def export_bundle_lengths(data_imap,
+                          clean_bundles, bundles):
+    """
+    full path to a JSON file containing median bundle lengths
+    """
+    bundle_dict = data_imap["bundle_dict"]
+    med_len_clean_counts = []
+    med_len_counts = []
+    bundle_names = list(bundle_dict.keys())
+    if "whole_brain" not in bundle_names:
+        bundle_names.append("whole_brain")
+    bundles_files = [clean_bundles, bundles]
+    lists = [med_len_clean_counts, med_len_counts]
+
+    for bundles_file, lens in zip(bundles_files, lists):
+        seg_sft = aus.SegmentedSFT.fromfile(bundles_file)
+
+        for bundle in bundle_names:
+            if bundle == "whole_brain":
+                lens.append(np.median(
+                    seg_sft.sft._tractogram._streamlines._lengths))
+            else:
+                these_lengths = seg_sft.get_bundle(
+                    bundle)._tractogram._streamlines._lengths
+                if len(these_lengths) > 0:
+                    lens.append(np.median(
+                        these_lengths))
+                else:
+                    lens.append(0)
+
+    counts_df = pd.DataFrame(
+        data=dict(
+            median_len=med_len_counts,
+            median_len_clean=med_len_clean_counts),
+        index=bundle_names)
+    return counts_df, dict(sources=bundles_files)
+
+
 @pimms.calc("profiles")
 @as_file('_profiles.csv', include_track=True, include_seg=True)
 def tract_profiles(clean_bundles, data_imap,
@@ -341,6 +381,7 @@ def get_segmentation_plan(kwargs):
     segmentation_tasks = with_name([
         get_scalar_dict,
         export_sl_counts,
+        export_bundle_lengths,
         export_bundles,
         clean_bundles,
         segment,