Geometric-phase metalens: focal distance and spot size

This notebook reconstructs the official Tidy3D metalens example with a compact 7×7 array of rotated TiO₂ pillars. It builds and voxelizes the actual FDTDX scene, launches a circularly polarized pulse, records an x-z field phasor, and extracts focal distance, transverse FWHM, and focal-core contrast.

Imports and parameters

import os
os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
from time import perf_counter
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
import numpy as np
import pandas as pd
import jax
import fdtdx
from IPython import get_ipython

from benchmarks.cases.device_metalens.run import (
    DL, DOMAIN, LENS_TOP_Z, PML_LAYERS, RUN_TIME, UM, WAVELENGTH,
    _build, _centers, _fwhm,
)
from benchmarks.device_geometry import metalens_pillars_um

get_ipython().run_line_magic("matplotlib", "inline")
polygons_um, centers_um, focal_length_design_um = metalens_pillars_um()
pd.Series({
    "wavelength_um": WAVELENGTH / UM,
    "grid_pitch_um": DL / UM,
    "domain_um": tuple(value / UM for value in DOMAIN),
    "array": "7 × 7 pillars",
    "pillar_height_um": 0.600,
    "design_focal_length_um": focal_length_design_um,
    "maximum_time_ps": RUN_TIME * 1e12,
    "PML_cells": PML_LAYERS,
}, name="value").to_frame()

	value
wavelength_um	0.6
grid_pitch_um	0.025
domain_um	(3.625, 3.625, 6.35)
array	7 × 7 pillars
pillar_height_um	0.6
design_focal_length_um	2.606736
maximum_time_ps	0.3
PML_cells	12

1. Set and plot the pillar geometry

fig, ax = plt.subplots(figsize=(6.5, 6.5))
for vertices, center in zip(polygons_um, centers_um, strict=True):
    ax.add_patch(Polygon(vertices + center, facecolor="#55c2b5", edgecolor="#087d83", lw=0.6))
ax.set(xlabel="x (µm)", ylabel="y (µm)", title="Exact rotated-pillar array", aspect="equal")
limit = 1.75
ax.set(xlim=(-limit, limit), ylim=(-limit, limit))
plt.show()

2. Build and voxelize the FDTDX scene

object_list, constraints, config, focal_length_design_um = _build()
key = jax.random.PRNGKey(0)
objects, arrays, design_params, config, _ = fdtdx.place_objects(
    object_list, config, constraints, key
)
arrays, objects, _ = fdtdx.apply_params(arrays, objects, design_params, key)
print(f"Yee grid: {objects.volume.grid_shape}")
print(f"Time steps: {config.time_steps_total:,}")

Yee grid: (145, 145, 254)
Time steps: 6,294

fig, axes = plt.subplots(1, 2, figsize=(11, 4.4))
fdtdx.plot_material_from_side(config, arrays, "z", position=-1.875 * UM, ax=axes[0], plot_legend=True)
fdtdx.plot_material_from_side(config, arrays, "y", position=0.0, ax=axes[1], plot_legend=True)
axes[0].set_title("Voxelized pillar plane")
axes[1].set_title("Voxelized vertical section")
plt.tight_layout(); plt.show()

3. Run the circularly polarized simulation

started = perf_counter()
_, result_arrays = fdtdx.run_fdtd(arrays=arrays, objects=objects, config=config, key=key, show_progress=False)
print(f"Local FDTDX runtime: {perf_counter() - started:.1f} s")

Local FDTDX runtime: 13.4 s

4. Plot the simulated focal field

detector = objects["xz"]
phasor = np.asarray(result_arrays.detector_states["xz"]["phasor"])[0, 0]
intensity = np.sum(np.abs(phasor) ** 2, axis=0).squeeze()
x_um = _centers(config, 0, detector.grid_slice_tuple[0]) / UM
z_from_lens_um = (_centers(config, 2, detector.grid_slice_tuple[2]) - LENS_TOP_Z) / UM

fig, ax = plt.subplots(figsize=(9, 5.5))
image = ax.pcolormesh(x_um, z_from_lens_um, intensity.T / intensity.max(), cmap="magma", shading="auto")
ax.axhline(focal_length_design_um, color="cyan", ls="--", label="design focus")
ax.set(xlabel="x (µm)", ylabel="distance from lens (µm)", title="Executed x-z intensity")
ax.legend(); fig.colorbar(image, ax=ax, label="normalized |E|²"); plt.show()

5. Extract focus, FWHM, and contrast

axis_intensity = intensity[np.argmin(np.abs(x_um)), :]
search = np.flatnonzero((z_from_lens_um > 1.0) & (z_from_lens_um < 4.0))
focus_index = int(search[np.argmax(axis_intensity[search])])
focal_distance_um = float(z_from_lens_um[focus_index])
fwhm_um = _fwhm(x_um, intensity[:, focus_index])
core = np.abs(x_um) <= 0.5
contrast = float(axis_intensity[focus_index] / np.mean(intensity[core, focus_index]))

fig, ax = plt.subplots(figsize=(8, 4.2))
profile = intensity[:, focus_index] / intensity[:, focus_index].max()
ax.plot(x_um, profile, color="#087d83")
ax.axhline(0.5, color="#8221a8", ls="--")
ax.set(xlabel="x (µm)", ylabel="normalized intensity", title=f"Focal plane · FWHM = {fwhm_um:.3f} µm")
ax.grid(alpha=0.25); plt.show()

golden = np.load("benchmarks/goldens/device_metalens.npz", allow_pickle=True)
pd.DataFrame({
    "FDTDX": [focal_distance_um, fwhm_um, contrast],
    "Tidy3D": [float(golden["focal_distance_um"][0]), float(golden["fwhm_um"][0]), float(golden["focal_core_contrast"][0])],
}, index=["focal distance (µm)", "FWHM (µm)", "focal-core contrast"])

	FDTDX	Tidy3D
focal distance (µm)	1.962500	1.850000
FWHM (µm)	0.502048	0.497904
focal-core contrast	1.924352	1.876280

Reproduce

For a production lens, increase the aperture and perform a grid/aperture convergence study.

print("uv run fdtdx-bench run --case device_metalens")

uv run fdtdx-bench run --case device_metalens