Merge pull request #5 from khanlab/optim-dask

Reduce runtime for warping to template, by initial downsampling

pyproject.toml

@@ -31,6 +31,7 @@ ome-zarr = "^0.9.0"
 pybids = "^0.16.5"
 sparse = "^0.15.1"
 zarrnii = "^0.1.1a1"
+bokeh = "^3.4.1"
 
 [tool.poetry.scripts]
 spimquant = "spimquant.run:main"

spimquant/workflow/rules/templatereg.smk

@@ -275,8 +275,7 @@ rule affine_zarr_to_template_nii:
         xfm_ras=rules.affine_reg.output.xfm_ras,
         ref_nii=bids_tpl(root=root, template="{template}", suffix="anat.nii.gz"),
     params:
-        chunks=(50, 50, 50),
-        zooms=None,
+        ref_opts={"chunks": (1, 50, 50, 50)},
     output:
         nii=bids(
             root=root,
@@ -298,8 +297,7 @@ rule affine_zarr_to_template_ome_zarr:
         xfm_ras=rules.affine_reg.output.xfm_ras,
         ref_nii=bids_tpl(root=root, template="{template}", suffix="anat.nii.gz"),
     params:
-        chunks=(50, 50, 50),
-        zooms=None,
+        ref_opts={"chunks": (1, 50, 50, 50)},
     output:
         ome_zarr=directory(
             bids(
@@ -324,8 +322,10 @@ rule deform_zarr_to_template_nii:
         warp_nii=rules.deform_reg.output.warp,
         ref_nii=bids_tpl(root=root, template="{template}", suffix="anat.nii.gz"),
     params:
-        chunks=(50, 50, 50),
-        zooms=None,
+        flo_opts={"level": 2}, #downsampling level to use (TODO: set this automatically based on ref resolution?)
+        do_downsample=True, #whether to perform further downsampling before transforming
+        downsample_opts={'along_z': 4}, #could also be determined automatically 
+        ref_opts={"chunks": (1, 100, 100, 100)},
     output:
         nii=bids(
             root=root,
@@ -337,6 +337,7 @@ rule deform_zarr_to_template_nii:
             **inputs["spim"].wildcards
         ),
     threads: 32
+    container: None
     script:
         "../scripts/deform_to_template_nii.py"
 
@@ -348,13 +349,17 @@ rule deform_to_template_nii_zoomed:
         warp_nii=rules.deform_reg.output.warp,
         ref_nii=bids_tpl(root=root, template="{template}", suffix="anat.nii.gz"),
     params:
-        chunks=(50, 50, 50),
-        zooms=lambda wildcards: (
-            float(wildcards.res) / 1000,
-            float(wildcards.res) / 1000,
-            float(wildcards.res) / 1000,
-        ),
-        #None #same resolution as template if NOne
+        flo_opts={"level": 1}, #downsampling level to use (TODO: set this automatically based on ref resolution?)
+        do_downsample=False, #whether to perform further downsampling before transforming
+        downsample_opts={'along_z': 4}, #could also be determined automatically 
+        ref_opts=lambda wildcards: {
+            "chunks": (1, 50, 50, 50),
+            "zooms": (
+                float(wildcards.res) / 1000,
+                float(wildcards.res) / 1000,
+                float(wildcards.res) / 1000,
+            ),
+        },
     output:
         nii=bids(
             root=root,

spimquant/workflow/scripts/affine_to_template_nii.py

@@ -11,7 +11,7 @@
 
 #member function of floating image
 flo_znimg = ZarrNii.from_path(snakemake.input.ome_zarr, channels=[channel_index])
-ref_znimg = ZarrNii.from_path_as_ref(snakemake.input.ref_nii, channels=[channel_index],chunks=snakemake.params.chunks,zooms=snakemake.params.zooms)
+ref_znimg = ZarrNii.from_path_as_ref(snakemake.input.ref_nii, channels=[channel_index],**snakemake.params.ref_opts)
 
 out_znimg = flo_znimg.apply_transform(Transform.affine_ras_from_txt(snakemake.input.xfm_ras),ref_znimg=ref_znimg)
 

spimquant/workflow/scripts/affine_to_template_ome_zarr.py

@@ -12,7 +12,7 @@
 
 
 flo_znimg = ZarrNii.from_path(snakemake.input.ome_zarr, channels=[channel_index])
-ref_znimg = ZarrNii.from_path_as_ref(snakemake.input.ref_nii, channels=[channel_index],chunks=snakemake.params.chunks,zooms=snakemake.params.zooms)
+ref_znimg = ZarrNii.from_path_as_ref(snakemake.input.ref_nii, channels=[channel_index],**snakemake.params.ref_opts)
 
 out_znimg = flo_znimg.apply_transform(Transform.affine_ras_from_txt(snakemake.input.xfm_ras),ref_znimg=ref_znimg)
 

spimquant/workflow/scripts/deform_to_template_nii.py

@@ -1,7 +1,10 @@
 import zarr
 import nibabel as nib
 from  zarrnii import ZarrNii, Transform
-from dask.diagnostics import ProgressBar
+from dask.distributed import Client
+
+client = Client(n_workers=4, threads_per_worker=2,processes=False)
+print(client.dashboard_link)
 
 #get channel index from omero metadata
 zi = zarr.open(snakemake.input.ome_zarr)
@@ -11,17 +14,20 @@
 
 
 #member function of floting image
-flo_znimg = ZarrNii.from_path(snakemake.input.ome_zarr, channels=[channel_index])
-ref_znimg = ZarrNii.from_path_as_ref(snakemake.input.ref_nii, channels=[channel_index],chunks=snakemake.params.chunks,zooms=snakemake.params.zooms)
+flo_znimg = ZarrNii.from_path(snakemake.input.ome_zarr, channels=[channel_index], **snakemake.params.flo_opts)
+ref_znimg = ZarrNii.from_path_as_ref(snakemake.input.ref_nii, channels=[channel_index],**snakemake.params.ref_opts)
+
+if snakemake.params.do_downsample:
+    flo_ds_znimg = flo_znimg.downsample(**snakemake.params.downsample_opts)
 
-deform_znimg = flo_znimg.apply_transform(Transform.displacement_from_nifti(snakemake.input.warp_nii),
+
+deform_znimg = flo_ds_znimg.apply_transform(Transform.displacement_from_nifti(snakemake.input.warp_nii),
                     Transform.affine_ras_from_txt(snakemake.input.xfm_ras),
                     ref_znimg=ref_znimg)
 
-with ProgressBar():
-
-    deform_znimg.to_nifti(snakemake.output.nii, scheduler='single-threaded')
 
+deform_znimg.to_nifti(snakemake.output.nii)
 
 
 
+client.close()

spimquant/workflow/scripts/deform_to_template_ome_zarr.py

@@ -3,7 +3,7 @@
 from dask.diagnostics import ProgressBar
 
 flo_znimg = ZarrNii.from_path(snakemake.input.ome_zarr, channels=[snakemake.params.channel_index])
-ref_znimg = ZarrNii.from_path_as_ref(snakemake.input.ref_nii, channels=[snakemake.params.channel_index],chunks=snakemake.params.chunks,zooms=snakemake.params.zooms)
+ref_znimg = ZarrNii.from_path_as_ref(snakemake.input.ref_nii, channels=[snakemake.params.channel_index],**snakemake.params.ref_opts)
 
 
 out_znimg = flo_znimg.apply_transform(Transform.displacement_from_nifti(snakemake.input.warp_nii),