Computed Generated Holography using a Digital Micromirror Device¶

In [ ]:

# If in Colab, install Chromatix. Don't forget to select a GPU!
!pip install --upgrade pip
!pip install git+https://github.com/chromatix-team/chromatix.git

# If in Colab, install Chromatix. Don't forget to select a GPU! !pip install --upgrade pip !pip install git+https://github.com/chromatix-team/chromatix.git

In this example, we'll demonstrate a version of computer generated holography (CGH) using a digital micromirror device (DMD). We can think of this as optimizing a binary amplitude mask instead of a phase mask, as in the first CGH example. Because we will be binarizing the phase, we'll be using a surrogate gradient. All of this is handled in Chromatix automatically. This style of CGH is inspired by the system presented in [1], though we demonstrate a much more simplified version of that system.

[1]: Conference presentation: M. Hossein Eybposh, Aram Moossavi, Vincent R. Curtis, and Nicolas C. Pegard "Optimization of time-multiplexed computer-generated holograms with surrogate gradients", Proc. SPIE PC12014, Emerging Digital Micromirror Device Based Systems and Applications XIV, PC1201406 (9 March 2022); https://doi.org/10.1117/12.2607781.

In [1]:

from typing import Tuple, Union

import equinox as eqx
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import optax
from jax import random
from jaxtyping import Array
from skimage.data import cat

from chromatix import Field
from chromatix.functional import (
    amplitude_change,
    asm_propagate,
    plane_wave,
)
from chromatix.ops import binarize

key = random.PRNGKey(4)

from typing import Tuple, Union import equinox as eqx import jax import jax.numpy as jnp import matplotlib.pyplot as plt import numpy as np import optax from jax import random from jaxtyping import Array from skimage.data import cat from chromatix import Field from chromatix.functional import ( amplitude_change, asm_propagate, plane_wave, ) from chromatix.ops import binarize key = random.PRNGKey(4)

Creating a 2D natural image target hologram¶

This is the hologram we'll be trying to recreate using our DMD. This example shows the optimization of a 2D hologram of a natural image. For an example of holography with a (very simple) 3D sample, see the other CGH example.

In [2]:

im = cat().mean(2)
im = im[:, 100:400]
data = jnp.array(im.astype(np.float32))

plt.figure(dpi=150)
plt.imshow(im, cmap="gray")
plt.axis("off")
plt.show()

im = cat().mean(2) im = im[:, 100:400] data = jnp.array(im.astype(np.float32)) plt.figure(dpi=150) plt.imshow(im, cmap="gray") plt.axis("off") plt.show()

Constructing the CGH model with a DMD¶

We've seen in other examples that defining optical systems is straightforward in Chromatix. Here, we show a CGH system that uses a new element, the AmplitudeMask. The AmplitudeMask can optionally be binarized, which means that the amplitudes will be squashed to either a 1 or a 0, and gradients through the binarization will be computed using the surrogate gradient method (as if no transformation has been applied to the amplitude mask).

In this version of the CGH model, we'll also demonstrate how to build a model that can propagate to different distances as desired by taking the z value as an argument.

In [3]:

class CGH(eqx.Module):
    amplitude: (
        Array  # This is the phase mask we want to optimize, and is not marked static!
    )
    shape: Tuple[int, int] = eqx.field(static=True, default=(300, 300))
    spacing: float = eqx.field(static=True, default=7.56)
    spectrum: float = eqx.field(static=True, default=0.66)
    n: float = eqx.field(static=True, default=1.0)
    pad_width: int = eqx.field(static=True, default=0)

    def __call__(self, z: Union[float, Array]) -> Field:
        # Chromatix does the work of simulating the propagation of the plane wave into
        # the desired hologram at multiple depths. This lets us define a CGH simulation
        # in just a few lines.
        field = plane_wave(self.shape, self.spacing, self.spectrum)
        field = amplitude_change(
            field, binarize(self.amplitude)
        )  # Modeling a DMD as a binary amplitude mask
        return asm_propagate(field, z, self.n, self.pad_width, mode="same")


z = 13e4
model = CGH(jax.nn.initializers.uniform(1.0)(key, shape=(300, 300), dtype=jnp.float32))

class CGH(eqx.Module): amplitude: ( Array # This is the phase mask we want to optimize, and is not marked static! ) shape: Tuple[int, int] = eqx.field(static=True, default=(300, 300)) spacing: float = eqx.field(static=True, default=7.56) spectrum: float = eqx.field(static=True, default=0.66) n: float = eqx.field(static=True, default=1.0) pad_width: int = eqx.field(static=True, default=0) def __call__(self, z: Union[float, Array]) -> Field: # Chromatix does the work of simulating the propagation of the plane wave into # the desired hologram at multiple depths. This lets us define a CGH simulation # in just a few lines. field = plane_wave(self.shape, self.spacing, self.spectrum) field = amplitude_change( field, binarize(self.amplitude) ) # Modeling a DMD as a binary amplitude mask return asm_propagate(field, z, self.n, self.pad_width, mode="same") z = 13e4 model = CGH(jax.nn.initializers.uniform(1.0)(key, shape=(300, 300), dtype=jnp.float32))

Optimizing the CGH system¶

We've discussed how to optimize a system in Chromatix using Flax and Optax in the first CGH example, so we won't detail those things here. The training setup and loop is basically the same, so we'll just go through the whole training process here:

In [4]:

# Defining the loss function
def loss_fn(model, data, z):
    eps = 1e-7
    approx = model(z).intensity.squeeze()
    loss = optax.cosine_distance(
        predictions=approx.reshape(-1), targets=data.reshape(-1), epsilon=eps
    ).mean()
    correlation = jnp.sum(approx * data) / (
        jnp.sqrt(jnp.sum(approx**2) * jnp.sum(data**2)) + eps
    )
    return loss, {"loss": loss, "correlation": correlation}


# Defining the optimizer
optimizer = optax.adam(learning_rate=2.0)
opt_state = optimizer.init(model)


# Defining the function which updates the parameters
@jax.jit
def update(model, opt_state, data, z):
    grads, metrics = jax.grad(loss_fn, has_aux=True)(model, data, z)
    updates, opt_state = optimizer.update(grads, opt_state, model)
    model = optax.apply_updates(model, updates)
    return model, opt_state, metrics

# Defining the loss function def loss_fn(model, data, z): eps = 1e-7 approx = model(z).intensity.squeeze() loss = optax.cosine_distance( predictions=approx.reshape(-1), targets=data.reshape(-1), epsilon=eps ).mean() correlation = jnp.sum(approx * data) / ( jnp.sqrt(jnp.sum(approx**2) * jnp.sum(data**2)) + eps ) return loss, {"loss": loss, "correlation": correlation} # Defining the optimizer optimizer = optax.adam(learning_rate=2.0) opt_state = optimizer.init(model) # Defining the function which updates the parameters @jax.jit def update(model, opt_state, data, z): grads, metrics = jax.grad(loss_fn, has_aux=True)(model, data, z) updates, opt_state = optimizer.update(grads, opt_state, model) model = optax.apply_updates(model, updates) return model, opt_state, metrics

In [5]:

%%time
# Simple training loop
max_iterations = 400
history = {
    "loss": np.zeros((max_iterations)),
    "correlation": np.zeros((max_iterations)),
}
for iteration in range(max_iterations):
    model, opt_state, metrics = update(model, opt_state, data, z)
    for m in metrics:
        history[m][iteration] = metrics[m]
    if iteration % 40 == 0 or iteration == (max_iterations - 1):
        print(iteration, metrics)
        approx1 = model(z).intensity.squeeze()

%%time # Simple training loop max_iterations = 400 history = { "loss": np.zeros((max_iterations)), "correlation": np.zeros((max_iterations)), } for iteration in range(max_iterations): model, opt_state, metrics = update(model, opt_state, data, z) for m in metrics: history[m][iteration] = metrics[m] if iteration % 40 == 0 or iteration == (max_iterations - 1): print(iteration, metrics) approx1 = model(z).intensity.squeeze()

0 {'correlation': Array(0.73058045, dtype=float32), 'loss': Array(0.91078347, dtype=float32)}
40 {'correlation': Array(0.86155796, dtype=float32), 'loss': Array(0.848808, dtype=float32)}
80 {'correlation': Array(0.8622059, dtype=float32), 'loss': Array(0.8484709, dtype=float32)}
120 {'correlation': Array(0.8627496, dtype=float32), 'loss': Array(0.8482904, dtype=float32)}
160 {'correlation': Array(0.8632451, dtype=float32), 'loss': Array(0.8481377, dtype=float32)}
200 {'correlation': Array(0.86358774, dtype=float32), 'loss': Array(0.8480373, dtype=float32)}
240 {'correlation': Array(0.86377096, dtype=float32), 'loss': Array(0.8479716, dtype=float32)}
280 {'correlation': Array(0.8639819, dtype=float32), 'loss': Array(0.8478785, dtype=float32)}
320 {'correlation': Array(0.8642313, dtype=float32), 'loss': Array(0.8477973, dtype=float32)}
360 {'correlation': Array(0.8643794, dtype=float32), 'loss': Array(0.84772563, dtype=float32)}
399 {'correlation': Array(0.86453927, dtype=float32), 'loss': Array(0.8476653, dtype=float32)}
CPU times: user 10.5 s, sys: 481 ms, total: 11 s
Wall time: 11.7 s

In [6]:

fig, ax1 = plt.subplots(dpi=150)

ax1.plot(np.array(history["loss"]), color="black")
ax1.set_ylabel("loss")
ax1.set_xlabel("iterations")
ax2 = ax1.twinx()

ax2.plot(np.array(history["correlation"]), color="red")
ax2.set_ylabel("pearson correlation", color="red")
plt.title("optimization progress")
plt.show()

fig, ax1 = plt.subplots(dpi=150) ax1.plot(np.array(history["loss"]), color="black") ax1.set_ylabel("loss") ax1.set_xlabel("iterations") ax2 = ax1.twinx() ax2.plot(np.array(history["correlation"]), color="red") ax2.set_ylabel("pearson correlation", color="red") plt.title("optimization progress") plt.show()

Evaluation¶

Let's have a look at how we did! We'll look at what happens when we apply our binarized amplitude mask to a plane wave and then propagate to the desired distance. We can also take a look at the optimized amplitude mask and its binarized form, which is the pattern we would see on the DMD.

In [7]:

plt.figure(dpi=200)
plt.subplot(1, 4, 1)
plt.axis("off")
plt.imshow(data, cmap="gray")
plt.title("Original")

plt.subplot(1, 4, 2)
plt.imshow(approx1, vmax=np.percentile(approx1, 95), cmap="gray")
plt.title("Generated")
plt.axis("off")

plt.subplot(1, 4, 3)
amp = model.amplitude
plt.imshow(amp, cmap="gray")
plt.title("Amplitude")
plt.axis("off")

plt.subplot(1, 4, 4)
plt.imshow(np.float32(amp > 0.5), cmap="gray")
plt.title("Binarized")
plt.axis("off")
plt.show()

plt.figure(dpi=200) plt.subplot(1, 4, 1) plt.axis("off") plt.imshow(data, cmap="gray") plt.title("Original") plt.subplot(1, 4, 2) plt.imshow(approx1, vmax=np.percentile(approx1, 95), cmap="gray") plt.title("Generated") plt.axis("off") plt.subplot(1, 4, 3) amp = model.amplitude plt.imshow(amp, cmap="gray") plt.title("Amplitude") plt.axis("off") plt.subplot(1, 4, 4) plt.imshow(np.float32(amp > 0.5), cmap="gray") plt.title("Binarized") plt.axis("off") plt.show()

In [8]:

z_list = np.linspace(0.0, z, 20)
images = model(z_list).intensity.squeeze()

plt.figure(figsize=(20, 10), dpi=150)
plt.suptitle("Varying Propagation Distance")

columns = 5
for i, image in enumerate(images):
    plt.subplot(len(images) // columns + 1, columns, i + 1)
    plt.imshow(image, vmax=np.percentile(image, 95), cmap="gray")
    plt.axis("off")
    plt.title(f"z = {z_list[i].squeeze() // 1}")
    plt.tight_layout()

z_list = np.linspace(0.0, z, 20) images = model(z_list).intensity.squeeze() plt.figure(figsize=(20, 10), dpi=150) plt.suptitle("Varying Propagation Distance") columns = 5 for i, image in enumerate(images): plt.subplot(len(images) // columns + 1, columns, i + 1) plt.imshow(image, vmax=np.percentile(image, 95), cmap="gray") plt.axis("off") plt.title(f"z = {z_list[i].squeeze() // 1}") plt.tight_layout()

We can also look at what happens along the propagation path by propagating to intermediate distances. This reveals the image of the cat forming as we get closer and closer to the target propagation distance. We can see that in the end we produced a hologram pretty close to the target image, though we have a lot of noise in the hologram! We'll leave fixing that as another exercise to the reader.