Simple Richardson-Lucy 3D on the CPU for testing purposes.

rl.py

#!/usr/bin/env python3
import numpy as np

def centered(arr, newshape):
    # Return the center newshape portion of the array.
    newshape = np.asarray(newshape)
    currshape = np.array(arr.shape)
    startind = (currshape - newshape) // 2
    endind = startind + newshape
    myslice = [slice(startind[k], endind[k]) for k in range(len(endind))]
    return arr[tuple(myslice)]

def rl_cpu(image, psfs, obj_xy, n_iter):
    from scipy.signal import convolve
    from joblib import Parallel, delayed
    Pz, Ix, Iy = psfs.shape
    Ox, Oy = obj_xy
    eps = 0.0001 * image.max()          # epsilon used for regularization and avoiding x/0
    obj = 0.5 * np.ones((Pz, Ox, Oy))   # initialize object (can be anything nonzero)
    psfs = psfs / psfs.sum()            # normalize to prevent overflow in convolutions

    with Parallel(n_jobs=-2, backend='threading') as pool:
        for _ in range(n_iter):
            # Forward projection done plane-by-plane, then summed to get estimate of sensor image
            est = sum(pool(delayed(convolve)(psfs[p,:,:], obj[p,:,:], mode='same', method='fft') for p in range(Pz)))
            
            # Relative blur (unexplained image intensity)
            rb = image / (est + eps)
            print("Relative blur magnitude {:0.02f}".format(np.sum(np.square(rb))))
            
            # Back projection (every plane of PSF with relative blur image to create volume
            bp = np.stack(pool(delayed(convolve)(psfs[p,::-1,::-1], rb, mode='same', method='fft') for p in range(Pz)))
            
            # Update estimated of the 3D object, crop the back projection to object size.
            obj *= centered(bp, (Pz, Ox, Oy))

    return obj

test_utils.py

def make_blobs(shape, blob_shape=None, Nprototypes=5, replicates=3):
    """Synthesize a volume to reconstruct.
    
    Args:
        shape: array shape tuple for the volume.
        blob_shape: size of bounding box for blobs in the volume.
        Nprototypes: number of unique blob classes.
        replicates: number of copies of each unique class.
    """
    from scipy.ndimage.filters import gaussian_filter
    from skimage.morphology import disk
    
    Z, Xo, Yo = shape
    if blob_shape is None:
        blob_shape = (3, 7, 7)
        
    obj = np.zeros((Z, Xo, Yo), dtype=np.float32)

    for i in range(Nprototypes):
        prototype = np.random.gamma(shape=2.5, scale=1, size=blob_shape) * disk(blob_shape[-1] // 2)
        prototype = gaussian_filter(prototype, sigma=1)
        positions = np.random.uniform(low=(0,0,0), 
                                      high=[x - y for (x, y) in zip(obj.shape, blob_shape)],
                                      size=(replicates, len(blob_shape)))
        for j in range(replicates):
            p = np.round(positions[j,:]).astype(np.int)
            m = np.array(blob_shape)
            obj[p[0]:p[0]+m[0],p[1]:p[1]+m[1],p[2]:p[2]+m[2]] += prototype * np.random.rand()
    
    return obj