From ce66e0e0b081371aa1bda79ccfd226658412357a Mon Sep 17 00:00:00 2001
From: Ali Khan <alik@robarts.ca>
Date: Sun, 8 Sep 2024 01:37:46 -0400
Subject: [PATCH] fixes for prestitched data, and revert back to chunksize z=1

---
 config/config.yml                          | 14 ++++++++-----
 spimprep_run                               | 11 +++++++++++
 workflow/Snakefile                         |  2 +-
 workflow/rules/ome_zarr.smk                |  4 ++--
 workflow/scripts/tif_stacks_to_ome_zarr.py | 23 +++++++++++-----------
 5 files changed, 34 insertions(+), 20 deletions(-)
 create mode 100644 spimprep_run

diff --git a/config/config.yml b/config/config.yml
index 7fd990d..e1ec9c4 100644
--- a/config/config.yml
+++ b/config/config.yml
@@ -57,18 +57,18 @@ bigstitcher:
     downsampling: 1
     block_size_x: 256 # for storage
     block_size_y: 256
-    block_size_z: 256
+    block_size_z: 1
     block_size_factor_x: 1 #e.g. 2 will use 2*block_size for computation
     block_size_factor_y: 1
-    block_size_factor_z: 1
+    block_size_factor_z: 256
 
 ome_zarr:
   desc: stitchedflatcorr
   max_downsampling_layers: 5 # e.g. 4 levels: { 0: orig, 1: ds2, 2: ds4, 3: ds8, 4: ds16}
   rechunk_size: #z, y, x 
-    - 256 
-    - 256
-    - 256
+    - 1 
+    - 4096
+    - 4096
   scaling_method: 'local_mean' #can be nearest, gaussian, local_mean, zoom (zoom uses spline interp)
  
   omero_metadata:
@@ -76,7 +76,11 @@ ome_zarr:
       default_color: 'FFFFFF'
       color_mapping:
         autof: 'FFFFFF'
+        AutoF: 'FFFFFF'
         abeta: '00FF00'
+        Abeta: '00FF00'
+        PI: 'FFFFFF'
+        AlphaSynuclein: '00FF00'
       defaults:
         active: True
         coefficient: 1.0
diff --git a/spimprep_run b/spimprep_run
new file mode 100644
index 0000000..fe72b75
--- /dev/null
+++ b/spimprep_run
@@ -0,0 +1,11 @@
+#!/bin/bash
+
+if ! command -v singularity >/dev/null 2>&1; then
+    echo "Error: 'singularity' is not installed or not in the PATH." >&2
+    exit 1
+fi
+
+container='docker://khanlab/spimprep-deps:main'
+unset SNAKEMAKE_PROFILE
+singularity exec ${container} snakemake --config --set-resources bigstitcher:mem_mb=30000 fuse_dataset:mem_mb=30000 -pr $@
+
diff --git a/workflow/Snakefile b/workflow/Snakefile
index 155244d..14abdd6 100644
--- a/workflow/Snakefile
+++ b/workflow/Snakefile
@@ -33,7 +33,7 @@ rule all:
     input:
         get_all_targets(),
         get_bids_toplevel_targets(),
-        get_qc_targets(),
+    #        get_qc_targets(), #need to skip this if using prestitched
     localrule: True
 
 
diff --git a/workflow/rules/ome_zarr.smk b/workflow/rules/ome_zarr.smk
index 51c635e..8ce26da 100644
--- a/workflow/rules/ome_zarr.smk
+++ b/workflow/rules/ome_zarr.smk
@@ -77,10 +77,10 @@ rule tif_stacks_to_ome_zarr:
         config["containers"]["spimprep"]
     group:
         "preproc"
-    threads: 8
+    threads: config["cores_per_rule"]
     resources:
         runtime=360,
-        mem_mb=32000,
+        mem_mb=35000,
     script:
         "../scripts/tif_stacks_to_ome_zarr.py"
 
diff --git a/workflow/scripts/tif_stacks_to_ome_zarr.py b/workflow/scripts/tif_stacks_to_ome_zarr.py
index 0490ffd..3d497a7 100644
--- a/workflow/scripts/tif_stacks_to_ome_zarr.py
+++ b/workflow/scripts/tif_stacks_to_ome_zarr.py
@@ -1,17 +1,16 @@
 import json
 import zarr
 import dask.array as da
-from dask.array.image import imread as dask_imread
+from dask.array.image import imread as dask_imread 
 from ome_zarr.io import parse_url
 from ome_zarr.writer import write_image
 from ome_zarr.format import format_from_version
 from ome_zarr.scale import Scaler
-from dask.diagnostics import ProgressBar
 from upath import UPath as Path
 from lib.cloud_io import get_fsspec, is_remote
+from dask.distributed import Client, LocalCluster
 
 in_tif_glob = snakemake.params.in_tif_glob
-
 metadata_json=snakemake.input.metadata_json
 downsampling=snakemake.params.downsampling
 max_layer=snakemake.params.max_downsampling_layers #number of downsamplings by 2 to include in zarr
@@ -21,6 +20,9 @@
 scaling_method=snakemake.params.scaling_method
 uri = snakemake.params.uri
 
+cluster = LocalCluster(processes=False)
+client = Client(cluster)
+print(client.dashboard_link)
 
 # prepare metadata for ome-zarr
 with open(metadata_json) as fp:
@@ -48,6 +50,8 @@
 omero={key:val for key,val in snakemake.config['ome_zarr']['omero_metadata']['defaults'].items()}
 omero['channels']=[]
 
+
+
 darr_list=[]
 for i,stain in enumerate(stains):
     
@@ -63,7 +67,6 @@
     channel_metadata['color'] = color
     omero['channels'].append(channel_metadata)
     
-
 darr_channels = da.stack(darr_list)
 
 
@@ -72,23 +75,19 @@
     fs = get_fsspec(uri,**fs_args)
     store = zarr.storage.FSStore(Path(uri).path,fs=fs,dimension_separator='/',mode='w')
 else:
-    store = zarr.DirectoryStore(out_zarr) 
-
-
+    store = zarr.DirectoryStore(out_zarr,dimension_separator='/') 
 
 
 
 group = zarr.group(store,overwrite=True)
 scaler = Scaler(max_layer=max_layer,method=scaling_method)
 
-
-with ProgressBar():
-    write_image(image=darr_channels,
+delayed = write_image(image=darr_channels,
                             group=group,
                             scaler=scaler,
                             coordinate_transformations=coordinate_transformations,
                             axes=axes,
-                            metadata={'omero':omero}
+                            metadata={'omero':omero},
+                            compute=True
                                 )
 
-