From 79e61a214888e57531e3f3b7a9d6d907f6fe1d94 Mon Sep 17 00:00:00 2001
From: Ali Khan <alik@robarts.ca>
Date: Tue, 14 Jan 2025 13:51:47 -0500
Subject: [PATCH] refactor to copy channel in separate rule

better for parallelization
---
 workflow/rules/imaris.smk                  | 46 ++++++++++++++++++++++
 workflow/scripts/imaris_channel_to_zarr.py | 10 +++++
 workflow/scripts/imaris_to_ome_zarr.py     | 17 ++------
 3 files changed, 59 insertions(+), 14 deletions(-)
 create mode 100644 workflow/scripts/imaris_channel_to_zarr.py

diff --git a/workflow/rules/imaris.smk b/workflow/rules/imaris.smk
index eaaf4bd..239d17d 100644
--- a/workflow/rules/imaris.smk
+++ b/workflow/rules/imaris.smk
@@ -35,10 +35,56 @@ rule imaris_to_metadata:
     script:
         "../scripts/imaris_to_metadata.py"
 
+rule imaris_channel_to_zarr:
+    input:
+        ims=get_input_sample,
+    params:
+        channel=lambda wildcards: get_stains(wildcards).index(wildcards.stain),
+    output:
+        zarr=directory(bids(
+            root=root,
+            subject="{subject}",
+            datatype="micr",
+            sample="{sample}",
+            acq="{acq}",
+            stain="{stain}",
+            suffix="imaris.zarr",
+        )),
+    log:
+        bids(
+            root="logs",
+            subject="{subject}",
+            datatype="imaris_channel_to_zarr",
+            sample="{sample}",
+            acq="{acq}",
+            stain="{stain}",
+            suffix="log.txt",
+        ),
+    container:
+        config["containers"]["spimprep"]
+    group:
+        "preproc"
+    threads: 1
+    resources:
+        runtime=360,
+        mem_mb=1000,
+    shadow: 'minimal'
+    script:
+        "../scripts/imaris_channel_to_zarr.py"
+
 
 rule imaris_to_ome_zarr:
     input:
         ims=get_input_sample,
+        zarr=lambda wildcards: expand(bids(
+            root=root,
+            subject="{subject}",
+            datatype="micr",
+            sample="{sample}",
+            acq="{acq}",
+            stain="{stain}",
+            suffix="imaris.zarr",
+        ),stain=get_stains(wildcards),allow_missing=True),
         metadata_json=rules.prestitched_to_metadata.output.metadata_json,
     params:
         max_downsampling_layers=config["ome_zarr"]["max_downsampling_layers"],
diff --git a/workflow/scripts/imaris_channel_to_zarr.py b/workflow/scripts/imaris_channel_to_zarr.py
new file mode 100644
index 0000000..8d33077
--- /dev/null
+++ b/workflow/scripts/imaris_channel_to_zarr.py
@@ -0,0 +1,10 @@
+import h5py
+import hdf5plugin
+import zarr
+from sys import stdout #change this to log file later..
+
+source = h5py.File(snakemake.input.ims, mode='r')
+dest = zarr.open_group(snakemake.output.zarr, mode='w')
+zarr.copy(source['DataSet/ResolutionLevel 0/TimePoint 0/Channel {chan}/Data'.format(chan=snakemake.params.channel)], dest,log=stdout,compressor=None)
+source.close()
+
diff --git a/workflow/scripts/imaris_to_ome_zarr.py b/workflow/scripts/imaris_to_ome_zarr.py
index 23bb9a4..d92a76a 100644
--- a/workflow/scripts/imaris_to_ome_zarr.py
+++ b/workflow/scripts/imaris_to_ome_zarr.py
@@ -14,17 +14,6 @@
 
 
 stains=snakemake.params.stains
-
-
-
-source = h5py.File(snakemake.input.ims, mode='r')
-dest = zarr.open_group('copy_hdf5.zarr', mode='w')
-from sys import stdout
-for chan in range(len(stains)):
-    zarr.copy(source[f'DataSet/ResolutionLevel 0/TimePoint 0/Channel {chan}/Data'], dest, name=f'channel_{chan}',log=stdout)
-source.close()
-
-in_zarr='copy_hdf5.zarr'
 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
@@ -73,15 +62,15 @@
 
 
 darr_list=[]
-for zarr_i,stain in enumerate(stains):
+for zarr_i,in_zarr in enumerate(snakemake.input.zarr):
     #open zarr to get group name
     zi = zarr.open(in_zarr)
-    darr_list.append(da.from_zarr(in_zarr,component=f'channel_{zarr_i}').rechunk(rechunk_size))
+    darr_list.append(da.from_zarr(in_zarr,component='Data').rechunk(rechunk_size))
 
 
     #append to omero metadata
     channel_metadata={key:val for key,val in snakemake.config['ome_zarr']['omero_metadata']['channels']['defaults'].items()}
-    channel_name=stain
+    channel_name=stains[zarr_i]
     channel_metadata['label'] = channel_name
     default_color=snakemake.config['ome_zarr']['omero_metadata']['channels']['default_color']
     color=snakemake.config['ome_zarr']['omero_metadata']['channels']['color_mapping'].get(channel_name,default_color)