From 655cdcc66ef433c8c515bf558fcdbcb2873ba6e7 Mon Sep 17 00:00:00 2001 From: Ali Khan Date: Sun, 22 Dec 2024 07:48:37 -0500 Subject: [PATCH] WIP --- hippunfold/workflow/rules/native_surf.smk | 504 +++++++++++++++++- .../workflow/scripts/convert_warp_2d_to_3d.py | 28 + hippunfold/workflow/scripts/create_warps.py | 43 +- 3 files changed, 541 insertions(+), 34 deletions(-) create mode 100644 hippunfold/workflow/scripts/convert_warp_2d_to_3d.py diff --git a/hippunfold/workflow/rules/native_surf.smk b/hippunfold/workflow/rules/native_surf.smk index 63288089..8212f9f7 100644 --- a/hippunfold/workflow/rules/native_surf.smk +++ b/hippunfold/workflow/rules/native_surf.smk @@ -16,10 +16,15 @@ resample_from_density = ( ) +unfoldreg_metrics = [ 'thickness', 'curvature', 'gyrification'] + + # if we aren't doing unfolded reg, then we have these rules override correct_bad_vertices2 -ruleorder: resample_native_surf > correct_bad_vertices2 +ruleorder: resample_native_surf > cp_template_to_unfold +#ruleorder: cp_template_to_unfold > resample_native_surf +ruleorder: resample_native_surf > correct_bad_vertices2 ruleorder: warp_midthickness_to_inout > correct_bad_vertices2 - +ruleorder: unfold_spring_model > warp_unfold_native_to_unfoldreg rule prep_hipp_coords_for_meshing: input: @@ -328,7 +333,7 @@ rule resample_native_surf: root=work, datatype="surf", suffix="{surf_name}.surf.gii", - space="corobl", + space="{space}", hemi="{hemi}", label="{label}", **inputs.subj_wildcards @@ -344,7 +349,7 @@ rule resample_native_surf: root=work, datatype="surf", suffix="{surf_name}.surf.gii", - space="unfold", + space="{space}", hemi="{hemi}", label="{label}", **inputs.subj_wildcards @@ -354,7 +359,7 @@ rule resample_native_surf: root=work, datatype="surf", suffix="{surf_name,midthickness}.surf.gii", - space="corobl", + space="{space}", den="{density}", hemi="{hemi}", label="{label}", @@ -511,8 +516,8 @@ rule apply_halfsurf_warp_to_img: shell: "greedy -d 3 -rf {input.fixed} -rm {input.moving} {output.warped} -r {input.warp} " - -rule convert_warp_from_itk_to_world: +#TODO: rename the warps here according to existing custom (type=itk -> type=ras) +rule convert_inout_warp_from_itk_to_world: input: warp=bids( root=work, @@ -550,7 +555,6 @@ rule warp_midthickness_to_inout: surf_gii=bids( root=work, datatype="surf", - den="{density}", suffix="midthickness.surf.gii", space="corobl", hemi="{hemi}", @@ -578,7 +582,6 @@ rule warp_midthickness_to_inout: datatype="surf", suffix="{surfname,inner|outer}.surf.gii", space="corobl", - den="{density}", hemi="{hemi}", label="{label}", **inputs.subj_wildcards @@ -591,3 +594,486 @@ rule warp_midthickness_to_inout: "wb_command -surface-apply-warpfield {input.surf_gii} {input.warp} {output.surf_gii} && " "wb_command -set-structure {output.surf_gii} {params.structure_type} -surface-type {params.surface_type}" " -surface-secondary-type {params.secondary_type}" + + + +rule unfold_spring_model: + input: + surf_gii=bids( + root=work, + datatype="surf", + suffix="{surfname}.surf.gii", + space="unfold", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + output: + surf_gii=bids( + root=work, + datatype="surf", + suffix="{surfname}.surf.gii", + space="unfoldspring", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + container: + config["singularity"]["autotop"] + group: + "subj" + script: '../scripts/unfold_spring_model.py' + + + +rule calculate_surface_area: + input: + gii=bids( + root=work, + datatype="surf", + suffix="midthickness.surf.gii", + space="{space}", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + output: + gii=bids( + root=work, + datatype="surf", + suffix="surfarea.shape.gii", + space="{space}", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + container: + config["singularity"]["autotop"] + group: + "subj" + shell: + "wb_command -surface-vertex-areas {input} {output}" + + +rule calculate_gyrification: + """new gyrification is ratio of nativearea to unfoldarea (e.g. surface scaling or distortion factor. + this should be proportional by a constant, to the earlier gyrification on 32k surfaces.""" + input: + native_surfarea=bids( + root=work, + datatype="surf", + suffix="surfarea.shape.gii", + space="corobl", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + unfold_surfarea=bids( + root=work, + datatype="surf", + suffix="surfarea.shape.gii", + space="unfold", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + output: + gii=bids( + root=work, + datatype="surf", + suffix="gyrification.shape.gii", + space="corobl", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + container: + config["singularity"]["autotop"] + group: + "subj" + shell: + 'wb_command -metric-math "nativearea/unfoldarea" {output.gii}' + " -var nativearea {input.native_surfarea} -var unfoldarea {input.unfold_surfarea}" + + +rule calculate_curvature_from_surface: + input: + gii=bids( + root=work, + datatype="surf", + suffix="midthickness.surf.gii", + space="corobl", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + output: + gii=bids( + root=work, + datatype="surf", + suffix="curvature.shape.gii", + space="corobl", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + container: + config["singularity"]["autotop"] + group: + "subj" + shell: + "wb_command -surface-curvature {input} -mean {output}" + + +rule calculate_thickness_from_surface: + input: + inner=bids( + root=work, + datatype="surf", + suffix="inner.surf.gii", + space="corobl", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + outer=bids( + root=work, + datatype="surf", + suffix="outer.surf.gii", + space="corobl", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + output: + gii=bids( + root=work, + datatype="surf", + suffix="thickness.shape.gii", + space="corobl", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + container: + config["singularity"]["autotop"] + group: + "subj" + shell: + "wb_command -surface-to-surface-3d-distance {input.outer} {input.inner} {output}" + + + + +#for unfold_reg, need to get data into 2d image - Jordan's previous code used metric-to-volume-mapping with unfoldiso inner and outer surfaces +# we can use same approach, just calculating the metrics on native, and using the unfold inner/outer as ribbon + +rule pad_unfold_ref: + """pads the unfolded ref in XY (to improve registration by ensuring zero-displacement at boundaries) + and to deal with input vertices that are just slightly outside the original bounding box""" + input: + atlas_dir=Path(download_dir) / "atlas" / "multihist7", + params: + ref_nii=lambda wildcards, input: Path(input.atlas_dir) + / config["atlas_files"]["multihist7"]["label_nii"].format(**wildcards), + output: + ref_nii=bids( + root=work, + datatype="anat", + suffix="refvol.nii.gz", + space="unfold", + desc="padded", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + container: + config["singularity"]["autotop"] + group: + "subj" + shell: + #"c3d {params.ref_nii} -scale 0 -shift 1 -pad 16x16x0vox 16x16x0vox 0 -o {output.ref_nii}" + "c3d {params.ref_nii} -scale 0 -shift 1 -pad 64x64x0vox 64x64x0vox 0 -o {output.ref_nii}" + + +rule native_metric_to_unfold_nii: + """converts metric .gii files to .nii for use in ANTs""" + input: + metric_gii=bids( + root=work, + datatype="surf", + suffix="{metric}.shape.gii", + space="corobl", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + unfoldedsurf=bids( + root=work, + datatype="surf", + suffix="inner.surf.gii", #inner is because the multihist 2d slices are there, but ideally this should be midthickness + space="unfold", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + ref_nii=bids( + root=work, + datatype="anat", + suffix="refvol.nii.gz", + space="unfold", + desc="padded", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + params: + interp="-nearest-vertex 1", + output: + metric_nii=bids( + root=work, + datatype="anat", + suffix="{metric}.nii.gz", + space="unfold", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + container: + config["singularity"]["autotop"] + group: + "subj" + shell: + "wb_command -metric-to-volume-mapping {input.metric_gii} {input.unfoldedsurf} {input.ref_nii} {output.metric_nii} {params.interp}" + + +print(bids( + root=work, + suffix="xfm.nii.gz", + datatype="warps", + desc="SyN", + from_="native", + to_="{atlas}", + space="unfold", + type_="itk", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards, + )) + +rule unfoldreg: + """performs 2D registration in unfolded space as in [reference paper] + this is done in a shadow directory to get rid of the tmp files generated by ants.""" + input: + fixed_images=expand(bids( + root=work, + datatype="anat", + suffix="{metric}.nii.gz", + space="unfold", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + metric=unfoldreg_metrics,allow_missing=True), + atlas_dir=lambda wildcards: Path(download_dir) / "atlas" / wildcards.atlas, + params: + antsparams="-d 2 -t so -o tmp", + fixed_args=lambda wildcards, input: + ' '.join(['-f {img}'.format(img=img) for img in input.fixed_images]), + moving_args=lambda wildcards, input: + ' '.join(['-m {img}'.format(img=str( Path(input.atlas_dir) / config["atlas_files"][wildcards.atlas][metric].format(**wildcards))) + for metric in unfoldreg_metrics ] ), + cmd_copy_warps=lambda wildcards, output: + ' && '.join([f"cp tmp1Warp.nii.gz {output.warp}", + f"cp tmp1InverseWarp.nii.gz {output.invwarp}"]) + output: + warp=bids( + root=work, + suffix="xfm.nii.gz", + datatype="warps", + desc="SyN", + from_="{atlas}", + to="native", + space="unfold", + type_="itk2d", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards, + ), + invwarp=bids( + root=work, + suffix="xfm.nii.gz", + datatype="warps", + desc="SyN", + from_="native", + to="{atlas}", + space="unfold", + type_="itk2d", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards, + ), + container: + config["singularity"]["autotop"] + group: + "subj" + log: + bids( + root="logs", + suffix="unfoldreg.txt", + atlas="{atlas}", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + shadow: + "minimal" + threads: 16 + resources: + mem_mb=16000, + time=10, + shell: + #"antsRegistrationSyNQuick.sh {params.antsparams} {params.fixed_args} {params.moving_args} &> {log} && " + "antsRegistrationSyNQuick.sh {params.antsparams} {params.fixed_args} {params.moving_args} && " + "{params.cmd_copy_warps}" + + +rule extend_warp_2d_to_3d: + """ we have a 2d warp, in space of innersurf 2d slice, 254x126. we want to + 1) make it a 3d warp (create a z-displacement, as zeros) + 2) extend the 2d warp across z slices, (so each z slice is transformed identically) + 3) fill in the central part of the array, when going from 254x126 to 256x128 + + Note that this rule is hardcoded for 256x128xN ref and 254x126 2d warp; + """ + input: + warp=bids( + root=work, + suffix="xfm.nii.gz", + datatype="warps", + desc="{desc}", + from_="{from}", + to="{to}", + space="{space}", + type_="itk2d", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards, + ), + ref=bids( + root=work, + datatype="anat", + suffix="refvol.nii.gz", + space="unfold", + desc="padded", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + output: + warp=bids( + root=work, + suffix="xfm.nii.gz", + datatype="warps", + desc="{desc}", + from_="{from}", + to="{to}", + space="{space}", + type_="itk", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards, + ), + container: + config["singularity"]["autotop"] + group: + "subj" + script: '../scripts/convert_warp_2d_to_3d.py' + + +rule convert_unfoldreg_warp_from_itk_to_world: + input: + warp=bids( + root=work, + suffix="xfm.nii.gz", + datatype="warps", + desc="{desc}", + from_="{to}", #swapped around since surf needs opposite direction + to="{from}", + space="{space}", + type_="itk", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards, + ), + output: + warp=bids( + root=work, + suffix="xfm.nii.gz", + datatype="warps", + desc="{desc}", + from_="{from}", + to="{to}", + space="{space}", + type_="surface", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards, + ), + group: + "subj" + shell: + "wb_command -convert-warpfield -from-itk {input} -to-world {output}" + + + + +rule warp_unfold_native_to_unfoldreg: + input: + surf_gii=bids( + root=work, + datatype="surf", + suffix="{surfname}.surf.gii", + space="unfold", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + warp=bids( + root=work, + suffix="xfm.nii.gz", + datatype="warps", + desc="SyN", + from_="native", + to="{atlas}", + space="unfold", + type_="surface", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards, + ), + params: + structure_type=lambda wildcards: hemi_to_structure[wildcards.hemi], + secondary_type=lambda wildcards: surf_to_secondary_type[wildcards.surfname], + surface_type="FLAT", + output: + surf_gii=bids( + root=work, + datatype="surf", + suffix="{surfname}.surf.gii", + space="unfold{atlas}", + hemi="{hemi}", + label="{label}", + **inputs.subj_wildcards + ), + container: + config["singularity"]["autotop"] + group: + "subj" + shell: + "wb_command -surface-apply-warpfield {input.surf_gii} {input.warp} {output.surf_gii} && " + "wb_command -set-structure {output.surf_gii} {params.structure_type} -surface-type {params.surface_type}" + " -surface-secondary-type {params.secondary_type}" + + + diff --git a/hippunfold/workflow/scripts/convert_warp_2d_to_3d.py b/hippunfold/workflow/scripts/convert_warp_2d_to_3d.py new file mode 100644 index 00000000..56aaa5f1 --- /dev/null +++ b/hippunfold/workflow/scripts/convert_warp_2d_to_3d.py @@ -0,0 +1,28 @@ +import numpy as np +import nibabel as nib + +# Read 2D xfm nifti +xfm2d_nib = nib.load(snakemake.input.warp) +xfm2d_vol = xfm2d_nib.get_fdata() + +# Read 3d ref nifti +ref3d_nib = nib.load(snakemake.input.ref) +ref3d_vol = ref3d_nib.get_fdata() + + +# Define the new shape +Nx, Ny, Nz = ref3d_vol.shape[:3] + +if Nx != xfm2d_vol.shape[0] or Ny != xfm2d_vol.shape[1]: + print(f'ref_vol: {ref3d_vol.shape}, warp_vol: {xfm2d_vol.shape}') + raise ValueError('Ref nifti and warp nifti must have the same X and Y dimensions') + + +# Create a new array initialized with zeros +xfm3d_vol = np.zeros((Nx, Ny, Nz, 1, 3), dtype=xfm2d_vol.dtype) + +# Insert the original array, which replicates in each zslice. Leaves the z-displacement untouched (as zero) +xfm3d_vol[...,:2] = xfm2d_vol + +# Save as nifti +nib.Nifti1Image(xfm3d_vol,affine=ref3d_nib.affine,header=xfm2d_nib.header).to_filename(snakemake.output.warp) diff --git a/hippunfold/workflow/scripts/create_warps.py b/hippunfold/workflow/scripts/create_warps.py index 84664d89..a74829fe 100644 --- a/hippunfold/workflow/scripts/create_warps.py +++ b/hippunfold/workflow/scripts/create_warps.py @@ -149,31 +149,24 @@ def summary(name, array): # we have points defined by coord_flat_{ap,pd,io}, and corresponding value as native_coords_phys[:,i] # and we want to interpolate on a grid in the unfolded space -# add some noise to avoid perfectly overlapping datapoints! -points = ( - coord_flat_ap * unfold_dims[0] - + (np.random.rand(coord_flat_ap.shape[0]) - 0.5) * 1e-6, - coord_flat_pd * unfold_dims[1] - + (np.random.rand(coord_flat_ap.shape[0]) - 0.5) * 1e-6, - coord_flat_io * unfold_dims[2] - + (np.random.rand(coord_flat_ap.shape[0]) - 0.5) * 1e-6, -) - -# get unfolded grid (from 0 to 1, not world coords), using meshgrid: -# note: indexing='ij' to swap the ordering of x and y -epsilon = snakemake.params.epsilon -(unfold_gx, unfold_gy, unfold_gz) = np.meshgrid( - np.linspace( - 0 + float(epsilon[0]), unfold_dims[0] - float(epsilon[0]), unfold_dims[0] - ), - np.linspace( - 0 + float(epsilon[1]), unfold_dims[1] - float(epsilon[1]), unfold_dims[1] - ), - np.linspace( - 0 + float(epsilon[2]), unfold_dims[2] - float(epsilon[2]), unfold_dims[2] - ), - indexing="ij", -) +# unnormalize and add a bit of noise so points don't ever perfectly overlap +coord_flat_ap_unnorm = coord_flat_ap * unfold_dims[0] +coord_flat_pd_unnorm = coord_flat_pd * unfold_dims[1] +coord_flat_io_unnorm = coord_flat_io * unfold_dims[2] + +# get unfolded grid + +points = np.stack( + [ + coord_flat_ap * unfold_dims[0] + + (np.random.rand(coord_flat_ap.shape[0]) - 0.5) * 1e-6, + coord_flat_pd * unfold_dims[1] + + (np.random.rand(coord_flat_ap.shape[0]) - 0.5) * 1e-6, + coord_flat_io * unfold_dims[2] + + (np.random.rand(coord_flat_ap.shape[0]) - 0.5) * 1e-6, + ], + axis=1, + ) summary("points", points)