combine.py

import dask.array as da
import os
from zarr_tools.convert import to_zarr
import fire

def combine(BF_zarr_2D_path:str, TRITC_zarr_3D_path:str, out_zarr_path:str):
    bd2d = da.from_zarr(os.path.join(BF_zarr_2D_path,'0/'))
    print('Opened BF:', bd2d)
    fd3d = da.from_zarr(os.path.join(TRITC_zarr_3D_path,'0/'))
    print('Opened TRITC:', fd3d)
    fd2d = fd3d.max(axis=1)
    bd2d = da.stack([bd2d, fd2d], axis=1)
    print('Resulting stack:', bd2d)
    to_zarr(
        bd2d, 
        path=out_zarr_path, 
        steps=4, 
        name=['BF','TRITC'], 
        colormap=['gray','orange'],
        lut=((1000,30000),(440, 600)),
    )

# combine(snakemake.input[0], snakemake.input[1], snakemake.output[0])
if __name__ == "__main__":
    fire.Fire(combine)