use dask client, opts for flo and downsampling

- adds options for downsampling and reading downsampled levels
- use dask client, and provided dashboard link too

spimquant/workflow/rules/templatereg.smk

@@ -323,8 +323,9 @@ rule deform_zarr_to_template_nii:
         ref_nii=bids_tpl(root=root, template="{template}", suffix="anat.nii.gz"),
     params:
         flo_opts={"level": 2}, #downsampling level to use (TODO: set this automatically based on ref resolution?)
-        downsample_opts={'along_z': 6}, #could also be determined automatically 
-        ref_opts={"chunks": (1, 20, 20, 20)},
+        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,
@@ -335,7 +336,7 @@ rule deform_zarr_to_template_nii:
             suffix="SPIM.nii",
             **inputs["spim"].wildcards
         ),
-    threads: 4
+    threads: 32
     container: None
     script:
         "../scripts/deform_to_template_nii.py"
@@ -348,6 +349,9 @@ 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:
+        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": (

spimquant/workflow/scripts/deform_to_template_nii.py

@@ -3,7 +3,7 @@
 from  zarrnii import ZarrNii, Transform
 from dask.distributed import Client
 
-client = Client(n_workers=snakemake.threads, threads_per_worker=1,processes=False)
+client = Client(n_workers=4, threads_per_worker=2,processes=False)
 print(client.dashboard_link)
 
 #get channel index from omero metadata
@@ -17,15 +17,15 @@
 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)
 
-deform_znimg = flo_znimg.apply_transform(Transform.displacement_from_nifti(snakemake.input.warp_nii),
+if snakemake.params.do_downsample:
+    flo_ds_znimg = flo_znimg.downsample(**snakemake.params.downsample_opts)
+
+
+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)
 
 
-if 'downsample_opts' in snakemake.params:
-    print('downsampling before apply transform')
-    deform_znimg = deform_znimg.downsample(**snakemake.params.downsample_opts)
-
 deform_znimg.to_nifti(snakemake.output.nii)