Data

General

USE_CV2 = True module-attribute

RandomDiskGenerator

Source code in src/chromatix/data/data.py

class RandomDiskGenerator:  # TODO avoid overlapping disks
    def __init__(
        self,
        N: int,
        num_points: int,
        radius: int,
        shape: tuple[int, int, int],
        z_range: tuple[int, int],
    ):
        """
        Create a dataset of random 3D coordinates and their associated image.
        Each generated sample consists of an array of x y z coordinates with
        shape: (n_points 3), accompanied by a 3D image with shape `shape`.
        The last dimension of `shape` represents the z axis, if it exists. The
        number of planes will be inferred from the `shape` argument. This is
        meant for TensorFlow and PyTorch data loaders that support generators.

        Args:
            N: Number of samples in the dataset. Avoid large N as the coordinates
                on each epoch are pre-stored for speed. On each new epoch the
                samples are randomized again (new random coordinates are generated)
                therefore you can easily use small `N` to avoid memory issues.
            num_points: Number of points in each sample. For 3D samples these samples
                will be randomly split between the planes.
            radius: Radius of the disks to be drawn on each plane.
            shape: Shape of the output image. Dimensions are [h w n_z] where n_z is
                the number of planes in 3D. For 2D samples use z=1.
            z_range: list
                Minimum and Maximum values for z values [min, max]. The returned
                coordinates are [x y z] and this parameter determines the range of
                the z coordinates.
        """

        assert len(shape) == 3, (
            "Shape must specify three dimensions, shape parameter is: {}".format(
                len(shape)
            )
        )
        self.N = N
        self.radius = radius
        self.shape = shape
        self.z_range = z_range
        self.num_points = num_points
        self.num_planes = shape[-1]
        self.reset()

    def reset(self):
        """
        Generate all the random coordinates. This is called when generator
        is instantiated or the last sample in the generator is reached.
        """
        self.y = np.random.randint(
            low=self.radius,
            high=self.shape[0] - self.radius,
            size=(self.N, self.num_points),
        )
        self.x = np.random.randint(
            low=self.radius,
            high=self.shape[1] - self.radius,
            size=(self.N, self.num_points),
        )

        if self.num_planes > 1:
            self.z_indices = np.random.randint(
                low=0, high=self.num_planes, size=(self.N, self.num_points)
            )
            self.z_values = np.random.rand(self.N, self.num_planes) * (
                self.z_range[1] - self.z_range[0]
            )
            self.z_values += self.z_range[0]
            self.z_values.sort(axis=1)
            self.z = np.zeros_like(self.x).astype(np.float32)

    def __len__(self) -> int:
        return self.N

    def __getitem__(self, idx: int) -> np.ndarray:
        """
        Get a new sample.

        Args:
            idx: Index of the current sample that needs to be generated.

        Returns:
            numpy.ndarray
                A (num_points 2) or (num_points 3) array containing the
                    coordinates.
            numpy.ndarray
                2D or 3D Image corresponding to the coordinates.
        """
        coords = np.array([self.x[idx], self.y[idx]]).T

        if self.num_planes > 1:
            canvas = np.zeros(self.shape)
            for i in range(self.num_planes):
                canvas[:, :, i] = draw_disks(
                    self.shape[:-1],
                    coords[self.z_indices[idx] == i, :],
                    self.radius,
                    color=255,
                )  # TODO add weight
                self.z[idx, self.z_indices[idx] == i] = self.z_values[idx, i]
            return (
                np.array([self.x[idx], self.y[idx], self.z[idx]]).T,
                canvas,
            )  # TODO add weight

        else:
            image = draw_disks(self.shape[:-1], coords, self.radius, color=255)[
                ..., None
            ]  # TODO add weight
            return coords, image

    def __call__(self) -> tuple[np.ndarray, np.ndarray]:
        """
        Get a new sample. Automatically iterates through samples of coordinates
        with every call. Will cause the random coordinates to be regenerated
        when the last sample is reached.

        Returns:
            numpy.ndarray
                A (num_points 2) or (num_points 3) array containing the
                    coordinates.
            numpy.ndarray
                2D or 3D Image corresponding to the coordinates.
        """
        for i in range(self.__len__()):
            yield self.__getitem__(i)
            if i == self.__len__() - 1:
                self.reset()

N = N instance-attribute

num_planes = shape[-1] instance-attribute

num_points = num_points instance-attribute

radius = radius instance-attribute

shape = shape instance-attribute

z_range = z_range instance-attribute

__call__()

Get a new sample. Automatically iterates through samples of coordinates with every call. Will cause the random coordinates to be regenerated when the last sample is reached.

Returns:

Type	Description
ndarray	numpy.ndarray A (num_points 2) or (num_points 3) array containing the coordinates.
ndarray	numpy.ndarray 2D or 3D Image corresponding to the coordinates.

Source code in src/chromatix/data/data.py

def __call__(self) -> tuple[np.ndarray, np.ndarray]:
    """
    Get a new sample. Automatically iterates through samples of coordinates
    with every call. Will cause the random coordinates to be regenerated
    when the last sample is reached.

    Returns:
        numpy.ndarray
            A (num_points 2) or (num_points 3) array containing the
                coordinates.
        numpy.ndarray
            2D or 3D Image corresponding to the coordinates.
    """
    for i in range(self.__len__()):
        yield self.__getitem__(i)
        if i == self.__len__() - 1:
            self.reset()

__getitem__(idx)

Get a new sample.

Parameters:

Name	Type	Description	Default
idx	int	Index of the current sample that needs to be generated.	required

Returns:

Type	Description
ndarray	numpy.ndarray A (num_points 2) or (num_points 3) array containing the coordinates.
ndarray	numpy.ndarray 2D or 3D Image corresponding to the coordinates.

Source code in src/chromatix/data/data.py

def __getitem__(self, idx: int) -> np.ndarray:
    """
    Get a new sample.

    Args:
        idx: Index of the current sample that needs to be generated.

    Returns:
        numpy.ndarray
            A (num_points 2) or (num_points 3) array containing the
                coordinates.
        numpy.ndarray
            2D or 3D Image corresponding to the coordinates.
    """
    coords = np.array([self.x[idx], self.y[idx]]).T

    if self.num_planes > 1:
        canvas = np.zeros(self.shape)
        for i in range(self.num_planes):
            canvas[:, :, i] = draw_disks(
                self.shape[:-1],
                coords[self.z_indices[idx] == i, :],
                self.radius,
                color=255,
            )  # TODO add weight
            self.z[idx, self.z_indices[idx] == i] = self.z_values[idx, i]
        return (
            np.array([self.x[idx], self.y[idx], self.z[idx]]).T,
            canvas,
        )  # TODO add weight

    else:
        image = draw_disks(self.shape[:-1], coords, self.radius, color=255)[
            ..., None
        ]  # TODO add weight
        return coords, image

__init__(N, num_points, radius, shape, z_range)

Create a dataset of random 3D coordinates and their associated image. Each generated sample consists of an array of x y z coordinates with shape: (n_points 3), accompanied by a 3D image with shape shape. The last dimension of shape represents the z axis, if it exists. The number of planes will be inferred from the shape argument. This is meant for TensorFlow and PyTorch data loaders that support generators.

Parameters:

Name	Type	Description	Default
N	int	Number of samples in the dataset. Avoid large N as the coordinates on each epoch are pre-stored for speed. On each new epoch the samples are randomized again (new random coordinates are generated) therefore you can easily use small N to avoid memory issues.	required
num_points	int	Number of points in each sample. For 3D samples these samples will be randomly split between the planes.	required
radius	int	Radius of the disks to be drawn on each plane.	required
shape	tuple[int, int, int]	Shape of the output image. Dimensions are [h w n_z] where n_z is the number of planes in 3D. For 2D samples use z=1.	required
z_range	tuple[int, int]	list Minimum and Maximum values for z values [min, max]. The returned coordinates are [x y z] and this parameter determines the range of the z coordinates.	required

Source code in src/chromatix/data/data.py

def __init__(
    self,
    N: int,
    num_points: int,
    radius: int,
    shape: tuple[int, int, int],
    z_range: tuple[int, int],
):
    """
    Create a dataset of random 3D coordinates and their associated image.
    Each generated sample consists of an array of x y z coordinates with
    shape: (n_points 3), accompanied by a 3D image with shape `shape`.
    The last dimension of `shape` represents the z axis, if it exists. The
    number of planes will be inferred from the `shape` argument. This is
    meant for TensorFlow and PyTorch data loaders that support generators.

    Args:
        N: Number of samples in the dataset. Avoid large N as the coordinates
            on each epoch are pre-stored for speed. On each new epoch the
            samples are randomized again (new random coordinates are generated)
            therefore you can easily use small `N` to avoid memory issues.
        num_points: Number of points in each sample. For 3D samples these samples
            will be randomly split between the planes.
        radius: Radius of the disks to be drawn on each plane.
        shape: Shape of the output image. Dimensions are [h w n_z] where n_z is
            the number of planes in 3D. For 2D samples use z=1.
        z_range: list
            Minimum and Maximum values for z values [min, max]. The returned
            coordinates are [x y z] and this parameter determines the range of
            the z coordinates.
    """

    assert len(shape) == 3, (
        "Shape must specify three dimensions, shape parameter is: {}".format(
            len(shape)
        )
    )
    self.N = N
    self.radius = radius
    self.shape = shape
    self.z_range = z_range
    self.num_points = num_points
    self.num_planes = shape[-1]
    self.reset()

__len__()

Source code in src/chromatix/data/data.py

def __len__(self) -> int:
    return self.N

reset()

Generate all the random coordinates. This is called when generator is instantiated or the last sample in the generator is reached.

Source code in src/chromatix/data/data.py

def reset(self):
    """
    Generate all the random coordinates. This is called when generator
    is instantiated or the last sample in the generator is reached.
    """
    self.y = np.random.randint(
        low=self.radius,
        high=self.shape[0] - self.radius,
        size=(self.N, self.num_points),
    )
    self.x = np.random.randint(
        low=self.radius,
        high=self.shape[1] - self.radius,
        size=(self.N, self.num_points),
    )

    if self.num_planes > 1:
        self.z_indices = np.random.randint(
            low=0, high=self.num_planes, size=(self.N, self.num_points)
        )
        self.z_values = np.random.rand(self.N, self.num_planes) * (
            self.z_range[1] - self.z_range[0]
        )
        self.z_values += self.z_range[0]
        self.z_values.sort(axis=1)
        self.z = np.zeros_like(self.x).astype(np.float32)

draw_disks(shape, coordinates, radius, color=255)

Create a grayscale image with disks drawn at each provided coordinate.

Parameters:

Name	Type	Description	Default
image_size		The desired image size as (height, width).	required
coordinates	ndarray	A list of (x, y) coordinates where disks should be drawn.	required
radius	int	The radius of the disks.	required
color	int	An optional intensity for the disks (0-255).	255

Returns:

Type	Description
ndarray	numpy.ndarray: The resulting grayscale image with disks drawn at the specified coordinates.

Source code in src/chromatix/data/data.py

def draw_disks(
    shape: tuple[int, int], coordinates: np.ndarray, radius: int, color: int = 255
) -> np.ndarray:
    """
    Create a grayscale image with disks drawn at each provided coordinate.

    Args:
        image_size: The desired image size as (height, width).
        coordinates: A list of (x, y) coordinates where disks should be
            drawn.
        radius: The radius of the disks.
        color: An optional intensity for the disks (0-255).

    Returns:
        numpy.ndarray: The resulting grayscale image with disks drawn at
            the specified coordinates.
    """
    image = np.zeros([s + radius * 2 for s in shape], dtype=np.uint8)
    _samples = np.linspace(-radius, radius, num=radius * 2, dtype=np.float32)
    circle = color * np.uint8(
        np.sum(np.array(np.meshgrid(_samples, _samples)) ** 2, axis=0) <= radius**2
    )
    for c in coordinates:
        slices = (slice(c[0], c[0] + radius * 2), slice(c[1], c[1] + radius * 2))
        image[slices] = image[slices] | circle
    image = image[radius : radius + shape[0], radius : radius + shape[1]]
    return image

draw_line(arr, start, stop, thickness=0.3, intensity=1.0)

Draw a line in a 3D object with a given thickness and intensity.

Parameters:

Name	Type	Description	Default
arr	ndarray	The object to draw the line in.	required
start	ndarray	The start of the line.	required
end		The end of the line.	required
thickness	float	The thickness of the line.	0.3
intensity	float	The intensity of the line.	1.0

Source code in src/chromatix/data/data.py

def draw_line(
    arr: np.ndarray,
    start: np.ndarray,
    stop: np.ndarray,
    thickness: float = 0.3,
    intensity: float = 1.0,
) -> np.ndarray:
    """
    Draw a line in a 3D object with a given thickness and intensity.

    Args:
        arr: The object to draw the line in.
        start: The start of the line.
        end: The end of the line.
        thickness: The thickness of the line.
        intensity: The intensity of the line.
    """
    direction = np.subtract(stop, start)
    line_length = np.sqrt(np.sum(np.square(direction)))
    n = direction / line_length

    sigma2 = 2 * thickness**2

    z, y, x = np.meshgrid(
        np.arange(arr.shape[0]),
        np.arange(arr.shape[1]),
        np.arange(arr.shape[2]),
        indexing="ij",
    )
    d2, t = sqr_dist_to_line(z, y, x, start, n)

    line_weight = (
        (t > 0) * (t < line_length)
        + (t <= 0) * np.exp(-(t**2) / sigma2)
        + (t >= line_length) * np.exp(-((t - line_length) ** 2) / sigma2)
    )
    return arr + intensity * np.exp(-d2 / sigma2) * line_weight

filaments_3d(sz, intensity=1.0, radius=0.8, rand_offset=0.05, rel_theta=1.0, num_filaments=50, thickness=0.3, seed=True)

Create a 3D representation of filaments.

Parameters:

Name	Type	Description	Default
sz	tuple[int, int, int]	A 3D shape tuple representing the size of the object.	required
radius	float | tuple[float, float, float]	A tuple of real numbers (or a single real number) representing the relative radius of the volume in which the filaments will be created. Default is 0.8. If a tuple is used, the filamets will be created in a corresponding elliptical volume. Note that the radius is only enforced in the version filaments_3d which creates the array rather than adding.	0.8
rand_offset	float	A tuple of real numbers representing the random offsets of the filaments in relation to the size. Default is 0.05.	0.05
rel_theta	float	A real number representing the relative theta range of the filaments. Default is 1.0.	1.0
num_filaments	int	An integer representing the number of filaments to be created. Default is 50.	50
thickness	float	A real number representing the thickness of the filaments in pixels. Default is 0.8.	0.3
seed	int | None	An optional integer used to seed the random number generator for randomly generating filaments. This can be used to repeatedly generate the same sample. Default is None in which case no seed is set and the default Numpy seed is used.	True

This code is based on the SyntheticObjects.jl package by Hossein Zarei Oshtolagh and Rainer Heintzmann.

Source code in src/chromatix/data/data.py

def filaments_3d(
    sz: tuple[int, int, int],
    intensity: float = 1.0,
    radius: float | tuple[float, float, float] = 0.8,
    rand_offset: float = 0.05,
    rel_theta: float = 1.0,
    num_filaments: int = 50,
    thickness: float = 0.3,
    seed: int | None = True,
) -> np.ndarray:
    """
    Create a 3D representation of filaments.

    Args:
        sz: A 3D shape tuple representing the size of the object.
        radius: A tuple of real numbers (or a single real number) representing
            the relative radius of the volume in which the filaments will be
            created. Default is 0.8. If a tuple is used, the filamets will be
            created in a corresponding elliptical volume. Note that the radius
            is only enforced in the version `filaments_3d` which creates the
            array rather than adding.
        rand_offset: A tuple of real numbers representing the random offsets of
            the filaments in relation to the size. Default is 0.05.
        rel_theta: A real number representing the relative theta range of the
            filaments. Default is 1.0.
        num_filaments: An integer representing the number of filaments to be
            created. Default is 50.
        thickness: A real number representing the thickness of the filaments in
            pixels. Default is 0.8.
        seed: An optional integer used to seed the random number generator
            for randomly generating filaments. This can be used to repeatedly
            generate the same sample. Default is ``None`` in which case no seed
            is set and the default Numpy seed is used.

    This code is based on the SyntheticObjects.jl package by Hossein Zarei
    Oshtolagh and Rainer Heintzmann.
    """

    sz = np.array(sz)
    radius = np.array(radius)
    rng = np.random.default_rng(seed)
    obj = np.zeros(sz, dtype=np.float32)
    mid = sz // 2
    # Draw random lines equally distributed over the 3D sphere
    for n in range(num_filaments):
        phi = 2 * np.pi * rng.random()
        # Theta should be scaled such that the distribution over the unit sphere is uniform
        theta = np.arccos(rel_theta * (1 - 2 * rng.random()))
        pos = (sz * radius / 2) * np.array(
            [
                np.sin(theta) * np.cos(phi),
                np.sin(theta) * np.sin(phi),
                np.cos(theta),
            ]
        )
        pos_offset = np.array(rand_offset * sz * (rng.random((3,)) - 0.5))
        # Draw line
        obj = draw_line(
            obj,
            pos + pos_offset + mid,
            mid + pos_offset - pos,
            thickness=thickness,
            intensity=intensity,
        )
    return obj

pollen_3d(sz, intensity=1.0, radius=0.8, dphi=0.0, dtheta=0.0, thickness=0.8, filled=False, filled_rel_intensity=0.1)

Create a 3D representation of a pollen grain.

Parameters:

Name	Type	Description	Default
sz	tuple[int, int, int]	A tuple of three integers representing the size of the volume in which the pollen grain will be created. Default is (128, 128, 128).	required
radius	float	Roughly the relative radius of the pollen grain.	0.8
dphi	float	A float representing the phi angle offset in radians. Default is 0.0.	0.0
dtheta	float	A float representing the theta angle offset in radians. Default is 0.0.	0.0
thickness	float	A float representing the thickness of the pollen grain. Default is 0.8.	0.8
filled	bool	A boolean representing whether the pollen grain should be filled. Default is False.	False
filled_rel_intensity		A float representing the relative intensity of the filled part of the pollen grain. Default is 0.1.	0.1

Returns: ret: A 3D array representing the pollen grain.

This code is based on the SyntheticObjects.jl package by Hossein Zarei Oshtolagh and Rainer Heintzmann and the original code by Kai Wicker.

Source code in src/chromatix/data/data.py

def pollen_3d(
    sz: tuple[int, int, int],
    intensity: float = 1.0,
    radius: float = 0.8,
    dphi: float = 0.0,
    dtheta: float = 0.0,
    thickness: float = 0.8,
    filled: bool = False,
    filled_rel_intensity=0.1,
) -> np.ndarray:
    """
    Create a 3D representation of a pollen grain.

    Args:
        sz: A tuple of three integers representing the size of the volume in
            which the pollen grain will be created. Default is ``(128, 128,
            128)``.
        radius: Roughly the relative radius of the pollen grain.
        dphi: A float representing the phi angle offset in radians. Default
            is 0.0.
        dtheta: A float representing the theta angle offset in radians. Default
            is 0.0.
        thickness: A float representing the thickness of the pollen grain.
            Default is 0.8.
        filled: A boolean representing whether the pollen grain should be
            filled. Default is ``False``.
        filled_rel_intensity: A float representing the relative intensity of the
            filled part of the pollen grain. Default is 0.1.
    Returns:
        ret: A 3D array representing the pollen grain.

    This code is based on the SyntheticObjects.jl package by Hossein Zarei
    Oshtolagh and Rainer Heintzmann and the original code by Kai Wicker.
    """

    sz = np.array(sz)
    z, y, x = np.meshgrid(
        np.linspace(-radius, radius, sz[0]),
        np.linspace(-radius, radius, sz[1]),
        np.linspace(-radius, radius, sz[2]),
        indexing="ij",
    )
    thickness = thickness / sz[0]

    r = x**2 + y**2 + z**2

    phi = np.atan2(y, x)
    theta = np.asin(z / (np.sqrt(x**2 + y**2 + z**2) + 1e-2)) + dtheta

    a = np.abs(np.cos(theta * 20))
    b = np.abs(
        np.sin(
            (phi + dphi) * np.sqrt(np.maximum(0, np.cos(theta) * (20.0**2)))
            - theta
            + np.pi / 2
        )
    )

    # calculate the relative distance to the surface of the pollen grain
    dc = ((0.4 + 1 / 20.0 * (a * b) ** 5) + np.cos(phi + dphi) * 1 / 20) - r
    # return dc

    sigma2 = 2 * (thickness**2)
    res = (
        intensity * np.exp(-(dc**2) / sigma2)
        + filled * (dc > 0) * intensity * filled_rel_intensity
    )

    return res

siemens_star(num_pixels=512, num_spokes=32, radius=None)

Generates a 2D Siemens star image of shape num_pixels. A single input is interpreted as a square-shaped array. radius is the radius of the star in pixels. If not provided, it will be half of the image size along each dimension.

Number of spokes in the star can be controlled with num_spokes. Spokes will alternate between black and white (0.0 and 1.0).

Source code in src/chromatix/data/data.py

def siemens_star(
    num_pixels: int = 512, num_spokes: int = 32, radius: int | None = None
) -> np.ndarray:
    """
    Generates a 2D Siemens star image of shape ``num_pixels``. A single input
    is interpreted as a square-shaped array. ``radius`` is the radius of the
    star in pixels. If not provided, it will be half of the image size along
    each dimension.

    Number of spokes in the star can be controlled with ``num_spokes``. Spokes
    will alternate between black and white (0.0 and 1.0).
    """

    num_pixels = np.atleast_1d(num_pixels)
    if num_pixels.size == 1:
        num_pixels = np.repeat(num_pixels, 2)
    if radius is None:
        radius = num_pixels / 2
    radius = np.atleast_1d(radius)
    if radius.size == 1:
        radius = np.repeat(radius, 2)
    ctr = num_pixels // 2
    X, Y = np.mgrid[
        -ctr[0] : num_pixels[0] - ctr[0], num_pixels[1] - ctr[1] : -ctr[1] : -1
    ]
    R = np.sqrt((X / radius[1]) ** 2 + (Y / radius[0]) ** 2)
    theta = np.arctan2(X, Y) + np.pi
    S = np.zeros_like(R)
    for spoke in range(num_spokes):
        in_spoke = (theta >= ((spoke) * 2 * np.pi / num_spokes)) & (
            theta <= ((spoke + 1) * 2 * np.pi / num_spokes)
        )
        if not spoke % 2:
            S[in_spoke] = 1.0
    S *= R < 1.0
    return S

sqr_dist_to_line(z, y, x, start, n)

Returns an array with each pixel being assigned to the square distance to that line and an array with the distance along the line.

Source code in src/chromatix/data/data.py

def sqr_dist_to_line(
    z: ArrayLike, y: ArrayLike, x: ArrayLike, start: ArrayLike, n: ArrayLike
) -> ArrayLike:
    """
    Returns an array with each pixel being assigned to the square distance to
    that line and an array with the distance along the line.
    """
    dx = x - start[2]
    dy = y - start[1]
    dz = z - start[0]
    dot_dn = dx * n[2] + dy * n[1] + dz * n[0]
    return (dx - dot_dn * n[2]) ** 2 + (dy - dot_dn * n[1]) ** 2 + (
        dz - dot_dn * n[0]
    ) ** 2, dot_dn

Objects

create_radial_pattern(shape)

Create a basic radial pattern image.

Parameters:

Name	Type	Description	Default
shape	tuple	Shape of the image (height, width).	required

Returns:

Type	Description
	jnp.ndarray: Radial pattern image.

Source code in src/chromatix/data/objects.py

def create_radial_pattern(shape):
    """
    Create a basic radial pattern image.

    Args:
        shape (tuple): Shape of the image (height, width).

    Returns:
        jnp.ndarray: Radial pattern image.
    """
    # Create a grid of coordinates
    y, x = jnp.indices(shape)

    # Calculate the center of the image
    center_y, center_x = shape[0] // 2, shape[1] // 2

    # Compute the distances from the center
    distances = jnp.sqrt((x - center_x) ** 2 + (y - center_y) ** 2)

    # Normalize distances to range [0, 2*pi] for phase pattern
    max_distance = jnp.sqrt(center_x**2 + center_y**2)
    phase_pattern = (distances / max_distance) * 2 * jnp.pi

    return phase_pattern

normalize_grayscale_image(input_path, output_path)

Source code in src/chromatix/data/objects.py

def normalize_grayscale_image(input_path, output_path):
    # Read the image
    img = imageio.imread(input_path)

    # Normalize the grayscale image
    normalized_img = img / img.max()

    # Convert the normalized image to 8-bit unsigned integer format
    normalized_img_ubyte = img_as_ubyte(normalized_img)

    # Save the normalized grayscale image as a PNG
    imageio.imsave(output_path, normalized_img_ubyte)

save_phase_pattern()

Source code in src/chromatix/data/objects.py

def save_phase_pattern():
    # Create the radial pattern
    shape = (512, 512)
    radial_pattern = create_radial_pattern(shape)

    # Save the pattern as a PNG file
    plt.imshow(radial_pattern, cmap="hsv")
    plt.colorbar()
    plt.title("Radial Phase Pattern")
    plt.axis("off")  # Hide the axis
    plt.savefig("data/radial_pattern.png", bbox_inches="tight", pad_inches=0)
    plt.show()

Permittivity Tensors

calc_scattering_potential(epsilon_r, refractive_index, wavelength)

Create the scattering potential from the permittivity tensor.

Parameters:

Name	Type	Description	Default
epsilon_r	ndarray	The permittivity tensor of the material.	required
refractive_index	float	The refractive index of the background medium.	required
wavelength	float	The wavelength of the light (microns).	required

Returns:

Type	Description
Array	The scattering potential.

Source code in src/chromatix/data/permittivity_tensors.py

def calc_scattering_potential(
    epsilon_r: jax.Array, refractive_index: float, wavelength: float
) -> jax.Array:
    """
    Create the scattering potential from the permittivity tensor.

    Args:
        epsilon_r (jnp.ndarray): The permittivity tensor of the material.
        refractive_index (float): The refractive index of the background medium.
        wavelength (float): The wavelength of the light (microns).

    Returns:
        The scattering potential.
    """
    k_0 = 2 * jnp.pi / wavelength
    vol_shape = epsilon_r.shape[:3]
    epsilon_m = jnp.tile(jnp.eye(3) * refractive_index**2, (*vol_shape, 1, 1))
    scattering_potential = k_0**2 * (epsilon_m - epsilon_r)
    return scattering_potential

create_calcite_crystal(shape, extraordinary_axis='z')

Create a calcite crystal phantom.

Parameters:

Name	Type	Description	Default
shape	Tuple[int, int, int]	Shape of the phantom (z, y, x)	required
extraordinary_axis	Optional[str]	Axis which is extraordinary ('x', 'y', or 'z')	'z'

Returns:

Type	Description
Array	4D array representing the phantom with the permittivity tensor at each voxel.

Source code in src/chromatix/data/permittivity_tensors.py

def create_calcite_crystal(
    shape: Tuple[int, int, int], extraordinary_axis: Optional[str] = "z"
) -> jax.Array:
    """
    Create a calcite crystal phantom.

    Args:
        shape: Shape of the phantom (z, y, x)
        extraordinary_axis: Axis which is extraordinary ('x', 'y', or 'z')

    Returns:
        4D array representing the phantom with the permittivity tensor at each voxel.
    """
    n_o = 1.658
    n_e = 1.486
    return create_homogeneous_phantom(shape, n_o, n_e, extraordinary_axis)

create_homogeneous_phantom(shape, n_o, n_e, extraordinary_axis='x')

Create a homogeneous uniaxial anisotropic phantom.

Parameters:

Name	Type	Description	Default
shape	Tuple[int, int, int]	Shape of the phantom (z, y, x)	required
n_o	float	Ordinary refractive index	required
n_e	float	Extraordinary refractive index	required
extraordinary_axis	Optional[str]	Axis which is extraordinary ('x', 'y', or 'z')	'x'

Returns:

Type	Description
Array	4D array representing the phantom with the permittivity tensor at each voxel.

Source code in src/chromatix/data/permittivity_tensors.py

def create_homogeneous_phantom(
    shape: Tuple[int, int, int],
    n_o: float,
    n_e: float,
    extraordinary_axis: Optional[str] = "x",
) -> jax.Array:
    """
    Create a homogeneous uniaxial anisotropic phantom.

    Args:
        shape: Shape of the phantom (z, y, x)
        n_o: Ordinary refractive index
        n_e: Extraordinary refractive index
        extraordinary_axis: Axis which is extraordinary ('x', 'y', or 'z')

    Returns:
        4D array representing the phantom with the permittivity tensor at each voxel.
    """
    epsilon_tensor = generate_permittivity_tensor(n_o, n_e, extraordinary_axis)
    phantom = jnp.tile(epsilon_tensor, (*shape, 1, 1))
    return phantom

create_homogeneous_scattering_potential(shape, n_o, n_e, background_permittivity)

Create a homogeneous uniaxial anisotropic scattering potential.

Parameters:

Name	Type	Description	Default
shape	tuple	Shape of the phantom (z, y, x)	required
n_o	float	Ordinary refractive index	required
n_e	float	Extraordinary refractive index	required
background_permittivity	float	Background permittivity	required

Returns:

Type	Description
Array	4D array representing the scattering potential.

Source code in src/chromatix/data/permittivity_tensors.py

def create_homogeneous_scattering_potential(
    shape: Tuple[int, int, int], n_o: float, n_e: float, background_permittivity: float
) -> jax.Array:
    """
    Create a homogeneous uniaxial anisotropic scattering potential.

    Args:
        shape (tuple): Shape of the phantom (z, y, x)
        n_o (float): Ordinary refractive index
        n_e (float): Extraordinary refractive index
        background_permittivity (float): Background permittivity

    Returns:
        4D array representing the scattering potential.
    """
    permittivity_tensor = create_homogeneous_phantom(shape, n_o, n_e)
    scattering_potential = create_scattering_potential(
        permittivity_tensor, background_permittivity
    )
    return scattering_potential

create_scattering_potential(permittivity_tensor, background_permittivity)

Create the scattering potential from the permittivity tensor.

Parameters:

Name	Type	Description	Default
permittivity_tensor	ndarray	The permittivity tensor of the material.	required
background_permittivity	float	The permittivity of the background medium.	required

Returns:

Type	Description
	The scattering potential.

Source code in src/chromatix/data/permittivity_tensors.py

def create_scattering_potential(permittivity_tensor, background_permittivity):
    """
    Create the scattering potential from the permittivity tensor.

    Args:
        permittivity_tensor (jnp.ndarray): The permittivity tensor of the material.
        background_permittivity (float): The permittivity of the background medium.

    Returns:
        The scattering potential.
    """
    # Calculate the permittivity contrast
    contrast = permittivity_tensor - background_permittivity

    # Scattering potential is proportional to the permittivity contrast
    scattering_potential = contrast / background_permittivity

    return scattering_potential

expand_potential_dims(tensor)

Source code in src/chromatix/data/permittivity_tensors.py

def expand_potential_dims(tensor: jax.Array) -> jax.Array:
    potential = jnp.expand_dims(tensor, axis=(1, 4))
    return potential

generate_dummy_potential(vol_shape)

Source code in src/chromatix/data/permittivity_tensors.py

def generate_dummy_potential(vol_shape: Tuple[int, int, int]) -> jax.Array:
    potential = expand_potential_dims(jnp.ones((*vol_shape, 3, 3)))
    return potential

generate_permittivity_tensor(n_o, n_e, extraordinary_axis='x')

Generate the permittivity tensor for a uniaxial anisotropic material.

Parameters:

Name	Type	Description	Default
n_o	float	Ordinary refractive index	required
n_e	float	Extraordinary refractive index	required
extraordinary_axis	str	Axis which is extraordinary ('x', 'y', or 'z')	'x'

Returns:

Type	Description
Array	Permittivity tensor with the order of axes as zyx.

Source code in src/chromatix/data/permittivity_tensors.py

def generate_permittivity_tensor(
    n_o: float, n_e: float, extraordinary_axis: Optional[str] = "x"
) -> jax.Array:
    """
    Generate the permittivity tensor for a uniaxial anisotropic material.

    Args:
        n_o (float): Ordinary refractive index
        n_e (float): Extraordinary refractive index
        extraordinary_axis (str): Axis which is extraordinary ('x', 'y', or 'z')

    Returns:
        Permittivity tensor with the order of axes as zyx.
    """
    epsilon_o = n_o**2
    epsilon_e = n_e**2
    if extraordinary_axis == "z":
        epsilon_tensor = jnp.array(
            [[epsilon_e, 0, 0], [0, epsilon_o, 0], [0, 0, epsilon_o]]
        )
    elif extraordinary_axis == "y":
        epsilon_tensor = jnp.array(
            [[epsilon_o, 0, 0], [0, epsilon_e, 0], [0, 0, epsilon_o]]
        )
    elif extraordinary_axis == "x":
        epsilon_tensor = jnp.array(
            [[epsilon_o, 0, 0], [0, epsilon_o, 0], [0, 0, epsilon_e]]
        )
    else:
        raise ValueError("extraordinary_axis must be one of 'x', 'y', or 'z'")
    return epsilon_tensor

permittivity_tensor_from_pixel(pixel_value, n_o_base=1.55, n_e_base=1.55, scale=0.5)

Source code in src/chromatix/data/permittivity_tensors.py

def permittivity_tensor_from_pixel(
    pixel_value, n_o_base=1.55, n_e_base=1.55, scale=0.5
):
    # The difference between n_o and n_e increases with the pixel value
    n_o = n_o_base + scale * pixel_value
    n_e = n_e_base - scale * pixel_value
    return generate_permittivity_tensor(n_o, n_e)

process_image_to_epsilon_r(input_path, n_o=1.658, n_e=1.486)

Source code in src/chromatix/data/permittivity_tensors.py

def process_image_to_epsilon_r(
    input_path: jax.Array, n_o: float = 1.658, n_e: float = 1.486
) -> jax.Array:
    img = imageio.imread(input_path)
    img = img / img.max()
    jax_img = jnp.array(img)

    n_avg = (n_o + n_e) / 2
    scale = (n_o - n_e) / 2
    epsilon_img = vectorized_permittivity_tensor_from_pixel(
        jax_img, n_avg, n_avg, scale
    )

    # Tile the epsilon tensor
    epsilon_r = jnp.tile(epsilon_img, (10, 1, 1, 1, 1))

    return epsilon_r

vectorized_permittivity_tensor_from_pixel(img, n_o_base=1.55, n_e_base=1.55, scale=0.5)

Source code in src/chromatix/data/permittivity_tensors.py

def vectorized_permittivity_tensor_from_pixel(
    img, n_o_base=1.55, n_e_base=1.55, scale=0.5
):
    vmap_func = jax.vmap(
        lambda pixel: permittivity_tensor_from_pixel(pixel, n_o_base, n_e_base, scale)
    )
    return jax.vmap(vmap_func)(img)