Handy function to check 3D particle tracking results

check_tracking_example.py

import numpy as np
import matplotlib.pyplot as plt
from scipy import ndimage
from matplotlib.pyplot import plot, scatter, imshow


def get_fake(size, number, radius, blur):
    """
    Generate a fake image for ideal gas (random positions).
    
    Args:
        size (tuple): the shape of the image, could be in any dimension
        number (int): the number of random gas paritcles in this image
        radius (int): the radius of the particles, assuming paticles are cubes.
        blur (float): sigma value for extra gaussian blur
    
    Return:
        tuple: the n-dimensional image and the positions of the idea gas
    """
    img = np.zeros(size)
    view = img.ravel()
    positions_1d = np.random.randint(0, img.size, number)
    positions_nd = np.array(np.unravel_index(positions_1d, img.shape)).T
    view[positions_1d] = 1
    img = ndimage.maximum_filter(img, radius)
    img = ndimage.gaussian_filter(img, blur)
    return img, positions_nd


def see_slice(image, positions, s, radius, axis=2, sizes=(10, 8)):
    """
    Check the 3D tracking result for a given slice
    
    Args:
        image (np.ndarray): the 3D volumetirc image, shape (x, y, z)
        positions (np.ndarray): the xyz positions of the particles, shape (n, 3)
        s (int): the index of slice to inspect. s = 5 means check the 5th slice
        radius (float or np.ndarray): the radius of the particles.
            if it is a float, then all particles were assumed to have the same radius
            if it is a numpy array, then different particles can have different radii
        axis (int): the axis index along which to take the slices.
            Setting axis=2 means slice along z-axis
        sizes (tuple): the figure size of the final plot
        
    Return:
        None: a figure would be plotted
    """
    fig = plt.figure()
    ax = fig.add_subplot(111)
    to_show = np.moveaxis(image, axis, 0)
    plt.imshow(to_show[s].T)
    if axis == -1:
        axis = 2
    shown_axes = np.array([i for i in range(3) if i != axis])
    if isinstance(radius, np.ndarray):
        radius_all_particles = radius
    else:
        radius_all_particles = [radius] * len(positions)
    for p, r in zip(positions, radius_all_particles):
        x, y = p[shown_axes]
        z = p[axis]
        dz = abs(z - s)
        if z > s:
            color='k'
        else:
            color='w'
        if dz < r:
            r_slice = np.sqrt(r**2 - dz**2)
            circle = plt.Circle([x, y], radius=r_slice, color=color, fill=None, lw=2)
            ax.add_patch(circle)
    fig.set_size_inches(sizes[0], sizes[1])
    plt.axis('off')
    plt.show()


if __name__ == "__main__":
    fake, pos = get_fake(size=(100, 100, 100), number=100, radius=5, blur=2)

    # Handy function to check your particle tracking result
    see_slice(fake, pos, 20, radius=6, axis=2)

This is the output