Merge pull request #6 from AlexVCaron/versaflow_update

Versaflow update

modules/processes/segment.nf

@@ -2,8 +2,8 @@
 
 nextflow.enable.dsl=2
 
-params.segmentation_classes = ["csf", "gm", "dgm", "wm", "bstem"]
-params.atropos_n4_tissues = ["wm", "bstem"]
+params.segmentation_classes = ["csf", "wm", "gm", "dgm", "pdgm", "blood"]
+params.atropos_n4_tissues = ["wm", "dgm", "pdgm"]
 params.atropos_prior_weight = 0.2
 params.atropos_mrf_weight = 0.3
 params.atropos_mrf_neighborhood = 1

modules/processes/tracking.nf

@@ -2,6 +2,8 @@
 
 nextflow.enable.dsl=2
 
+params.local_tracking_gpu_batch_size = 20000
+
 params.streamline_compression_factor = 0.2
 params.pve_threshold = 0.05
 
@@ -142,15 +144,18 @@ process Local_prob_tracking_opencl {
         export ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS=1
         export OMP_NUM_THREADS=1
         export OPENBLAS_NUM_THREADS=1
-        scil_compute_local_tracking_gpu.py $fodf $seeding_mask $tracking_mask \
+
+        scil_compute_local_tracking.py $fodf $seeding_mask $tracking_mask \
             tmp.trk \
             --step $step_length \
             --theta $theta \
             --min_length $params.local_min_len \
             --max_length $params.local_max_len \
             --sub_sphere $sphere_sub \
             --${seeding_strategy} $n_seeds \$compress \
-            --rng_seed $seed
+            --seed $seed \
+            --use_gpu \
+            --batch_size $params.local_tracking_gpu_batch_size
         scil_remove_invalid_streamlines.py tmp.trk \
             ${sid}_local_gpu_prob_in_${tracking_mask_type}_seed_${seed}_${seeding_strategy}${n_seeds}_in_${seeding_mask_type}_step_${step_length}_theta_${theta}_tracking.trk \
             --remove_single_point

modules/processes/utils.nf

@@ -13,8 +13,8 @@ params.safe_csf_mask_dilation = 1
 params.safe_gm_mask_dilation = 1
 params.duplicates_merge_method = "mean"
 params.validate_bvecs_fa_thr = 0.2
-params.segmentation_classes = ["csf", "gm", "dgm", "wm", "bstem"]
-params.tissue_masks_mapping = ["wm": ["wm", "bstem"], "gm": ["gm", "dgm"], "csf": ["csf"]]
+params.segmentation_classes = ["csf", "wm", "gm", "dgm", "pdgm", "blood"]
+params.tissue_masks_mapping = ["wm": ["wm", "dgm", "pdgm"], "gm": ["gm"], "csf": ["csf"]]
 
 include { remove_alg_suffixes; add_suffix } from '../functions.nf'
 

modules/workflows/segment.nf

@@ -49,8 +49,7 @@ params.register_d99 = true
 params.register_charm = true
 params.register_sarm = true
 params.register_inia19 = true
-
-params.segmentation_classes = ["csf", "gm", "dgm", "wm", "bstem"]
+params.segmentation_classes = ["csf", "wm", "gm", "dgm", "pdgm", "blood"]
 
 params.tissue_segmentation_root = "${get_data_path()}/maccaca_mulatta/tissue_segmentation"
 params.wm_segmentation_root = "${get_data_path()}/maccaca_mulatta/wm_segmentation"

nextflow.config

@@ -170,8 +170,8 @@ params.quick_denoised_t1_registration = false
 params.t1_registration_in_subject_space = false
 
 // Segmentation
-params.segmentation_classes = ["csf", "gm", "dgm", "wm", "bstem"]
-params.atropos_n4_tissues = ["wm", "bstem"]
+params.segmentation_classes = ["csf", "wm", "gm", "dgm", "pdgm", "blood"]
+params.atropos_n4_tissues = ["wm", "dgm", "pdgm"]
 params.atropos_prior_weight = 0.2
 params.atropos_mrf_weight = 0.3
 params.atropos_mrf_neighborhood = 1

tracking/main.nf

@@ -88,7 +88,7 @@ workflow load_tracking_data {
 
         wm_mask_channel = Channel.fromFilePairs("$root/**/segmentation/*_wm_mask.nii.gz", size: 1, flat: true)
             { get_id(it.parent.parent, root) }
-        tissue_pvf_channel = Channel.fromFilePairs("$root/**/segmentation/*pvf_3t_{csf,gm,wm}.nii.gz", size: 3, flat: true)
+        tissue_pvf_channel = Channel.fromFilePairs("$root/**/segmentation/*_3t_{csf,gm,wm}_pvf.nii.gz", size: 3, flat: true)
             { get_id(it.parent.parent, root) }
 
         enforce_sid_convention_dwi(dwi_channel.map{ [it[0], it[1..-1], ["dwi"] * (it.size() - 1)] })
@@ -104,7 +104,7 @@ workflow load_tracking_data {
             tissue_pvf_channel.map{[
                 it[0],
                 it[1..-1],
-                it[1..-1].collect{ i -> i.simpleName.tokenize("_")[-1] + "_pvf"}
+                it[1..-1].collect{ i -> i.simpleName.tokenize("_")[-2] + "_pvf"}
             ]}
         )
 

tracking/nextflow.config

@@ -2,12 +2,14 @@
 // Output management
 params.publish_all = false
 params.publish_mode = "link"
+params.publish_all_mode = "link"
 params.output_root = "outputs"
 
 // Tractography techniques
 params.pft_tracking = true
 params.local_tracking = true
 params.run_commit = true
+params.concatenate_tractograms = true
 
 params.use_opencl_tracking = false
 params.use_cuda = false
@@ -45,6 +47,7 @@ params.local_sub_sphere = 0
 
 params.b0_threshold = 40
 params.streamline_compression_factor = 0.1
+params.local_tracking_gpu_batch_size = 20000
 
 params.local_min_len = 5
 params.local_max_len = 80
@@ -86,9 +89,11 @@ if ( params.use_cuda ) {
     docker.runOptions = "--gpus all"
 }
 
-process {
-    withName: '.*(opencl|gpu).*' {
-        maxForks = params.cuda_max_parallel
+if (params.cuda_max_parallel > 0) {
+    process {
+        withName: '.*(opencl|gpu).*' {
+            maxForks = params.cuda_max_parallel
+        }
     }
 }
 
@@ -100,4 +105,4 @@ params.on_hcp = false
 includeConfig './profiles.config'
 
 // Include tower config, which can be empty if no tower connection needs to be made
-includeConfig '../.base/tower.config'
+includeConfig '../.base/tower.config'

workflows/io.nf

@@ -68,10 +68,7 @@ workflow load_dataset {
         enforce_sid_convention_metadata(dwi_json_channel.map{ it + ["dwi"] })
         enforce_sid_convention_rev_metadata(rev_json_channel.map{ it + ["rev"] })
 
-        // Unpack base images (T1 + DWI)
-        anat_channel = enforce_sid_convention_anat.out.image
-        dwi_channel = enforce_sid_convention_dwi.out.image
-            .map{ [it[0], it[1][2], it[1][0], it[1][1]] }
+        // Filter subjects to get only those specified by the user
         ref_id_channel = enforce_sid_convention_anat.out.image
             .map{ [it[0]] }
 
@@ -82,27 +79,40 @@ workflow load_dataset {
             ref_id_channel = ref_id_channel.filter{ !params.exclude_sid.contains(it[0]) }
         }
 
+        // Unpack base images (T1 + DWI)
+        anat_channel = ref_id_channel
+            .join(enforce_sid_convention_anat.out.image)
+
+        dwi_channel = ref_id_channel
+            .join(enforce_sid_convention_dwi.out.image)
+            .map{ [it[0], it[1][2], it[1][0], it[1][1]] }
+
         // Unpack reverse phase images
         rev_bval_bvec_channel = fill_missing_datapoints(
-            enforce_sid_convention_rev_bval_bvec.out.image.map{ it.flatten() },
+            ref_id_channel
+                .join(enforce_sid_convention_rev_bval_bvec.out.image)
+                .map{ it.flatten() },
             ref_id_channel,
             1, ["", ""]
         )
 
         rev_channel = fill_missing_datapoints(
-            enforce_sid_convention_rev.out.image,
+            ref_id_channel
+                .join(enforce_sid_convention_rev.out.image),
             ref_id_channel,
             1, [""]
         ).join(rev_bval_bvec_channel)
 
         // Unpack image metadata
         dwi_json_channel = fill_missing_datapoints(
-            enforce_sid_convention_metadata.out.image,
+            ref_id_channel
+                .join(enforce_sid_convention_metadata.out.image),
             ref_id_channel,
             1, [""]
         )
         rev_json_channel = fill_missing_datapoints(
-            enforce_sid_convention_rev_metadata.out.image,
+            ref_id_channel
+                .join(enforce_sid_convention_rev_metadata.out.image),
             ref_id_channel,
             1, [""]
         )
@@ -128,20 +138,23 @@ workflow load_dataset {
 
         // Unpack tissues PVF images
         pvf_channel = fill_missing_datapoints(
-            enforce_sid_convention_pvf.out.image
+            ref_id_channel
+                .join(enforce_sid_convention_pvf.out.image)
                 .map{ [it[0], it[1].reverse()] },
             ref_id_channel,
             1, [[]]
         )
 
         // Load available masks (T1 and/or DWI)
         dwi_mask_channel = fill_missing_datapoints(
-            enforce_sid_convention_dwi_mask.out.image,
+            ref_id_channel
+                .join(enforce_sid_convention_dwi_mask.out.image),
             ref_id_channel,
             1, [""]
         )
         anat_mask_channel = fill_missing_datapoints(
-            enforce_sid_convention_anat_mask.out.image,
+            ref_id_channel
+                .join(enforce_sid_convention_anat_mask.out.image),
             ref_id_channel,
             1, [""]
         )

workflows/preprocess.nf

@@ -15,7 +15,8 @@ include {
 include {
     extract_b0 as dwi_b0;
     extract_b0 as extract_epi_corrected_b0;
-    extract_b0 as extract_b0_preprocessed
+    extract_b0 as extract_b0_preprocessed;
+    extract_b0 as extract_b0_reference
 } from '../modules/processes/preprocess.nf'
 include {
     scil_compute_dti_fa
@@ -39,7 +40,8 @@ include {
     ants_transform as ants_transform_safe_wm_mask;
     ants_transform as ants_transform_raw_t1_mask;
     ants_transform as apply_transform_epi_dwi;
-    ants_transform as apply_transform_epi_rev
+    ants_transform as apply_transform_epi_rev;
+    ants_transform as apply_transform_epi_field
 } from '../modules/processes/register.nf'
 include {
     convert_float_to_integer as convert_wm_segmentation;
@@ -394,10 +396,15 @@ workflow preprocess_wkf {
                     .map{ [it[0], it[1].find{m -> m.simpleName.contains("_rev")}, "rev__ec_input_rev_metadata"] }
             ).map{ it.flatten() }
 
+            extract_b0_reference(
+                ec_input_dwi_channel.map{ it[0..2] + [""] },
+                "preprocess", "false", params.extract_mean_b0_base_config
+            )
 
+            b0_reference_for_registration = extract_b0_reference.out.b0
             apply_transform_epi_rev(
                 ec_input_rev_channel.map{ it[0..1] }
-                    .join(epi_correction_wkf.out.transform_reference)
+                    .join(b0_reference_for_registration)
                     .join(epi_correction_wkf.out.forward_transform)
                     .join(epi_correction_wkf.out.reverse_transform)
                     .map{ it[0..-3] + [it[-2] + it[-1]] }
@@ -408,6 +415,7 @@ workflow preprocess_wkf {
                 "",
                 params.ants_transform_base_config
             )
+
             rev_channel = replace_dwi_file(rev_channel, apply_transform_epi_rev.out.image)
 
             // Applied estimated susceptibility correction to DWI
@@ -443,10 +451,23 @@ workflow preprocess_wkf {
                 epi_displacement_field_channel = epi_correction_wkf.out.field
                 epi_fieldmap_channel = epi_correction_wkf.out.fieldmap
 
+                apply_transform_epi_field(
+                    epi_displacement_field_channel
+                        .join(b0_reference_for_registration)
+                        .join(epi_correction_wkf.out.forward_transform)
+                        .map{ it[0..-2] + [it[-1]] }
+                        .map{ it + [["true"], "", ""] },
+                    "preprocess",
+                    "",
+                    "false",
+                    "",
+                    params.ants_transform_base_config
+                )
+
                 apply_epi_field_wkf(
                     ec_input_dwi_channel.map{ it[0..1] },
                     apply_transform_epi_rev.out.image,
-                    epi_displacement_field_channel,
+                    apply_transform_epi_field.out.image,
                     ec_input_dwi_meta_channel
                         .map{ [it[0], it[1]] },
                     ""

workflows/tracking.nf

@@ -17,6 +17,7 @@ def asArray ( object ) {
 params.pft_tracking = true
 params.local_tracking = true
 params.run_commit = true
+params.concatenate_tractograms = true
 
 params.use_opencl_tracking = false
 
@@ -99,7 +100,7 @@ workflow tracking_wkf {
 
         PFT_maps(wm_vf.join(gm_vf).join(csf_vf), "tracking")
         PFT_tracking(
-            fodfs.join(PFT_maps.out.maps).join(PFT_maps.out.wm_gm_interface),
+            fodfs.join(PFT_maps.out.maps).join(PFT_maps.out.wm_gm_interface).map{ it + ["interface"] },
             "tracking",
             pft_random_seed,
             pft_tracking_algorithm,
@@ -278,9 +279,11 @@ workflow ensemble_tracking_wkf {
 
             out_tractograms = Commit.out.filtered_tractogram
         }
-
-        Ensemble_Tractograms(out_tractograms.groupTuple().join(fa_channel), "tracking")
-        out_ensemble_tractogram = Ensemble_Tractograms.out.tractogram
+
+        if (params.concatenate_tractograms) {
+            Ensemble_Tractograms(out_tractograms.groupTuple().join(fa_channel), "tracking")
+            out_ensemble_tractogram = Ensemble_Tractograms.out.tractogram
+        }
 
     emit:
         tractograms = out_tractograms