spheres.py

import numpy as np
import tifffile


def make_cylinder_profile(x, radius):
    valid = np.where(x ** 2 <= radius ** 2)
    profile = np.zeros(m)
    profile[valid] = 2 * np.sqrt(radius ** 2 - x[valid] ** 2)
    
    return profile


def make_sphere(x, y, radius):
    y, x = np.mgrid[-radius:radius, -radius:radius]
    valid = np.where(x ** 2 + y ** 2 <= radius ** 2)
    sphere = np.zeros((2 * radius, 2 * radius))
    sphere[valid] = 2 * np.sqrt(radius ** 2 - x[valid] ** 2 - y[valid] ** 2)
    
    return sphere


n = 512
m = 1024
x = np.arange(m) -  m // 2

# Cylinders
outer_cylinder_radius = m / 3
inner_cylinder_radius = outer_cylinder_radius - 50
outer_cylinder = make_cylinder_profile(x, outer_cylinder_radius)
inner_cylinder = make_cylinder_profile(x, inner_cylinder_radius)
capillary = np.tile(outer_cylinder - inner_cylinder, [n, 1])

# Spheres
r_min = 10
r_max = 50
num_spheres = 100
x = np.random.randint(
    m // 2 - inner_cylinder_radius + r_max,
    high=m // 2 + inner_cylinder_radius - r_max,
    size=num_spheres
)
y = np.random.randint(
    r_max,
    high=n - r_max,
    size=num_spheres
)
radii = np.random.randint(r_min, high=r_max, size=num_spheres)

for i, radius in enumerate(radii):
    sphere = make_sphere(x[i], y[i], radius)
    capillary[y[i] - radius:y[i] + radius, x[i] - radius:x[i] + radius] += sphere