motion-correction for b0s when not using topup

snakedwi/workflow/rules/prepdwi.smk

@@ -119,6 +119,60 @@ rule mrdegibbs:
         "cp {input[3]} {output[3]}"
 
 
+rule moco_bzeros:
+    """ run motion-correction (rigid reg to init volume) on the bzeros """
+    input:
+        nii_4d=bids(
+            root=work,
+            suffix="b0s.nii.gz",
+            datatype="dwi",
+            desc="degibbs",
+            **subj_wildcards
+        ),
+    output:
+        affine_dir=directory(
+            bids(
+                root=work,
+                suffix="transforms",
+                desc="moco",
+                datatype="dwi",
+                **subj_wildcards
+            )
+        ),
+        nii_4d=bids(
+            root=work,
+            suffix="b0s.nii.gz",
+            desc="moco",
+            datatype="dwi",
+            **subj_wildcards
+        ),
+        nii_avg3d=bids(
+            root=work,
+            suffix="b0.nii.gz",
+            desc="moco",
+            datatype="dwi",
+            **subj_wildcards
+        ),
+    threads: 8  #doesn't need to be more than the number of bzeros 
+    resources:
+        mem_mb=32000,
+    shadow:
+        "minimal"
+    container:
+        config["singularity"]["prepdwi"]  #-- this rule needs niftyreg, c3d and mrtrix
+    group:
+        "subj"
+    shell:
+        "c4d {input.nii_4d} -slice w 0:-1 -oo dwi_%03d.nii && "
+        "parallel --eta --jobs {threads} "
+        "reg_aladin -flo dwi_{{1}}.nii  -ref dwi_000.nii -res warped_{{1}}.nii -aff affine_xfm_ras_{{1}}.txt "
+        " ::: `ls dwi_???.nii | tail -n +2 | grep -Po '(?<=dwi_)[0-9]+'` && "
+        " mkdir -p {output.affine_dir} && cp affine_xfm_ras_*.txt {output.affine_dir} && "
+        " echo -e '1 0 0 0\n0 1 0 0\n0 0 1 0\n0 0 0 1' > {output.affine_dir}/affine_xfm_ras_000.txt && "
+        " mrcat dwi_000.nii warped_*.nii {output.nii_4d} && "
+        " mrmath {output.nii_4d} mean {output.nii_avg3d} -axis 3"
+
+
 # now have nii with just the b0's, want to create the topup phase-encoding text files for each one:
 rule get_phase_encode_txt:
     input:
@@ -688,6 +742,23 @@ rule get_shell_avg:
         "../scripts/get_shell_avg.py"
 
 
+# this gets vols from a a particular shell (can use to get b0s)
+rule get_shell_vols:
+    input:
+        dwi="{dwi_prefix}_dwi.nii.gz",
+        shells="{dwi_prefix}_dwi.shells.json",
+    params:
+        bval="{shell}",
+    output:
+        avgshell="{dwi_prefix}_b{shell}s.nii.gz",
+    group:
+        "subj"
+    container:
+        config["singularity"]["python"]
+    script:
+        "../scripts/get_shell_vols.py"
+
+
 def get_mask_for_eddy():
     if config["masking_method"] == "b0_BET":
         method = "bet_from-b0"
@@ -806,24 +877,16 @@ else:
 
 def get_dwi_ref(wildcards):
 
-    # this gets the number of DWI scans for this subject(session)
-    filtered = filter_list(input_zip_lists["dwi"], wildcards)
-    num_scans = len(filtered["subject"])
-
-    if num_scans > 1 and not config["no_topup"]:
+    if config["no_topup"]:
         return bids(
-            root=work,
-            suffix="b0.nii.gz",
-            desc="topup",
-            method="jac",
-            datatype="dwi",
-            **subj_wildcards
+            root=work, suffix="b0.nii.gz", datatype="dwi", desc="moco", **subj_wildcards
         )
     else:
         return bids(
             root=work,
             suffix="b0.nii.gz",
-            desc="degibbs",
+            desc="topup",
+            method="jac",
             datatype="dwi",
             **subj_wildcards
         )

snakedwi/workflow/scripts/get_shell_vols.py

@@ -0,0 +1,38 @@
+import json
+import nibabel as nib
+import numpy as np
+import sys
+
+with open(snakemake.input.shells) as f:
+    shells_dict = json.load(f)
+
+# get bval parameter:
+bval = snakemake.params.bval
+
+indices = shells_dict["shell_to_vol"][bval]
+
+# input dwi
+dwi_nib = nib.load(snakemake.input.dwi)
+
+# create output shape, 4th dim length of indices (# of volumes at this shell)
+newshape = np.array(
+    [dwi_nib.shape[0], dwi_nib.shape[1], dwi_nib.shape[2], len(indices)]
+)
+
+shell_vols = np.zeros(newshape.astype(int))
+
+
+if len(dwi_nib.shape) == 3 and len(indices) == 1 and indices[0] == 0:
+    # we have 3d vol (e.g. b0 only), so just grab it..
+    shell_vols = dwi_nib.get_fdata()[:, :, :]
+elif len(dwi_nib.shape) == 4:
+    # otherwise, pick out indices and average
+    shell_vols = dwi_nib.get_fdata()[:, :, :, indices]
+else:
+    # if not either of these cases, then something weird with indices and volumes
+    print("unable to get map indices to get shell volumes")
+    sys.exit()
+
+# now save as image
+shell_vols_nii = nib.Nifti1Image(shell_vols, affine=dwi_nib.affine)
+shell_vols_nii.to_filename(snakemake.output[0])