Pycudadecon + dask

[Eddymorphling]

Hi Talley,
I am trying to use pycudadecon on large 3D datasets generated with our gaussian ligthsheet. As you may be aware, the datasets are large and loading the entire volume during decon leads to memory issues.

I have acquired PSFs but wanted to reach out about your suggestions for doing this using dask-arrays for efficient GPU/CPU memory usage when running decon jobs using GPU on my cluster. Any notebooks/scripts that you may have that uses pycudadecon+dask to help me get started? Much appreciated. Thank you.

[tlambert03]

i don't have one directly using pycudadecon, but in general, you'll want to use dask map_blocks or map_overlap. As @VolkerH mentioned in #13 (comment), he wrote a short tutorial on breaking up large deconvolution problems into tiles. (there he used flowdec, but the split/combine principle is the same)

https://github.com/hammerlab/flowdec/blob/master/python/examples/notebooks/Tile-by-tile%20deconvolution%20using%20dask.ipynb

[Eddymorphling]

Great, you are the best! Thanks.

[Eddymorphling]

Hi again @tlambert03 , The CLI help command does not seems to work. Is it cudaDeconv --help ?

[tlambert03]

no v, just cudaDecon ... see the C++ library for the cli: https://github.com/scopetools/cudaDecon

[tlambert03]

oh! my docs are wrong aren't they! Sorry about that... will update

[Eddymorphling]

Yeah, I was following the docs. No worries. Thank you.

[Eddymorphling]

Sorry to bother you again @tlambert03. Testing a set of images using the CLI command line and it seems to be stuck at the step without any progress. Output included below. Any suggestions why? Pytorch sees my GPU.

# of tiles: 3. # of overlaps: 2. Tile size: 1024 Y pix. Overlap: 213 Y pix. 

Tile: 1 out of 3.   
raw_image size             : 2478 x 1024 x 1

[Eddymorphling]

Here is my CLI

radialft /home/psf_4x.tif /home/psf_4x_OTF.tif --fixorigin 10 --wavelength 525 --na 0.35 --nimm 1.56 --xyres 3 --zres 3 --nocleanup 
cudaDecon /data/z-slices fused --otf-file /home/psf_4x_OTF.tif --wavelength 525 --tile 1024 

[tlambert03]

not immediately sure, but let's first determine whether it has to do with your data specifically. Try the test data available here: https://github.com/tlambert03/pycudadecon/tree/main/tests/test_data

use something like

radialft.exe .\test_data\psf.tif .\test_data\otf.tif --nocleanup --fixorigin 10
cudaDecon.exe -z 0.3 -i 10 -D 31.5 .\test_data\im_raw .\test_data\otf.tif

[Eddymorphling]

Yes that worked with your test dataset. Not sure what's not compatible with my input datasets. Here is some background.

1. My PSF is not experimental, it is computationally generated (xyz = 256x256x256).
2. The OTF file seems good.
3. My input files are single tiff slices in a folder. These are gaussian LS tiff slices. They are 4339x5942 (xy). uneven pixel count as they are stitched and fused.

I kept it simple this time and here is what I ran. But gets stuck in the step I mentioned above.

radialft /home/scratch/pycudadecon/psf_4x.tif /home/scratch/pycudadecon/psf_4x_OTF.tif --fixorigin 10 --wavelength 525 --na 0.35 --nimm 1.56 --xyres 3 --zres 3 --nocleanup 
cudaDecon -i 10 /home/scratch/pycudadecon/z-slices fused --otf-file /home/scratch/pycudadecon/psf_4x_OTF.tif

[Eddymorphling]

Ah I figured it out. When the documentation said the input folder can be tiff files, what I assumed is they can be 2D tiff files generated from a 3D stack, not multiple 3D stacks. My bad. Running it on a 3D stack works now. Thanks Talley.

[tlambert03]

great!

[Eddymorphling]

Hi again @tlambert03 I have been trying to setup a workflow that uses dask and pycudadecon for just 3D gaussian light sheet datasets but have been facing trouble. I adapted this example which I believe was written by you? Not sure.. I removed the deskew function and just told it to crop and deconvolve as they are 3D (xyz) gaussian LS datasets.

I adapted what I found and tried to run dask+pycudadecon+GPU on our cluster. Here is what I came up with

import napari
import pycudadecon
from functools import partial
from skimage import io
from dask_image.imread import imread
import dask.array as da
import tifffile
from dask.distributed import Client

def main():
 
    client = Client()  # This will use all available cores by default
    print(client)

    # Load stacks with dask_image, and PSF with skimage
    # Each file in the folder is a single Z slice from the 3D volume
    stack = imread("/data/z-slices/*.tif")
    psf = io.imread("/home/psf_4x.tif")

    # Print the chunk sizes of the stack
    print("Chunk sizes:", stack.chunks)

    # Prepare some functions that accept a numpy array and return a processed array
    def crop(array):
        # Simple cropping function
        return array[:, 2:, 10:-20, :500]

    # https://docs.python.org/3.8/library/functools.html#functools.partial
    deconv = partial(pycudadecon.decon, psf=psf, background=10)

    # Map and chain those functions across all dask blocks
    deconvolved = stack.map_blocks(deconv, dtype="float16")
    cropped = deconvolved.map_blocks(crop, dtype="float16")

    output_path = "/home/scratch/decon.tif"
    tifffile.imwrite(output_path, deconvolved.compute())

    # Put the resulting dask array into napari.
    # (Don't forget the contrast limits and multiscale==False !)
    v = napari.view_image(
        cropped,
        contrast_limits=[90, 1500],
        multiscale=False,
        ndisplay=3,
        scale=(3, 1, 1),
    )

    napari.run()

if __name__ == '__main__':
    main()

My input folder is tiff files series where each file is a single slice, while my PSF is a 3D volume. When I run this, it does not save any deconvolved image nor open the image in napari. Below is what I see in the terminal, Would you happen to know what's going on? Thank you.

<Client: 'tcp://127.0.0.1:45707' processes=8 threads=64, memory=503.62 GiB>
Chunk sizes: ((1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1), (5942,), (4339,))
New chunk sizes: ((5, 5, 3), (1000, 1000, 1000, 1000, 1000, 942), (1000, 1000, 1000, 1000, 339))
callednx=74, ny=74, nz=74
Center of mass is (37.500, 37.500, 38.100)
Background is 0.000
In radialft()

[tlambert03]

looking, quick question: where is this getting printed?

New chunk sizes: ((5, 5, 3), (1000, 1000, 1000, 1000, 1000, 942), (1000, 1000, 1000, 1000, 339))

[tlambert03]

as for the failure, it looks like something is going wrong at the moment of OTF generation (in radialft)...
I would simplify this for now, leaving out napari, and simply calling cropped.compute(), also, try putting prints and/or breakpoints in the various functions to see exactly what dimensions/data are getting passed.

Try also leaving out dask and all that complexity too. If you're trying to build up to something that does block-wise deconvolution and stitches them all back together, then first just try to manually deconvolve one of those blocks (i.e. just literally crop your data down to a smaller subvolume that would match the chunk size you're trying to create). Then just deconvolve that single chunk and see how it goes.

If you're failing on the radialft part, then it suggests there might be a problem with your psf file... so I would suggest manually going through what the decon function does, and see where it fails. for example maybe first try to make a TemporaryOTF from your psf like this line in the decon function does. In general, break it down to much smaller steps, before building back up to the full map-blocks-chunked decon in napari