Microring resonator add-drop filter in BeamZ

import os
import matplotlib.pyplot as plt
import numpy as np
from scipy.signal import find_peaks
from IPython.display import display

import beamz as bz
from beamz.analysis import mode_data_to_dataframe
from beamz.analysis.plotting import plot_mode_field_components

test_mode = os.environ.get("BEAMZ_DOCS_TEST") == "1"
um = bz.um

lambda0 = 1.55 * um
freq0 = bz.LIGHT_SPEED / lambda0
wavelengths = np.linspace(1.50, 1.60, 3 if test_mode else 201) * um
freqs = bz.LIGHT_SPEED / wavelengths
fwidth = 0.5 * (np.max(freqs) - np.min(freqs))

n_si, n_sio2 = 3.48, 1.444
mat_si = bz.Material(permittivity=n_si**2)
mat_sio2 = bz.Material(permittivity=n_sio2**2)

wg_height = 0.22 * um
wg_width = 0.43 * um
gap = 0.08 * um
ring_radius = 3.0 * um
buffer = 0.7 * lambda0
wg_center_y = ring_radius + wg_width + gap
print(f"Ring radius = {ring_radius / um:.2f} µm; coupling gap = {gap / um:.2f} µm")

def make_sim():
    domain_x = 2 * ring_radius + wg_width + 2 * buffer
    domain_y = domain_x + 2 * gap + 2 * wg_width
    domain_z = wg_height + 2 * buffer
    domain = (domain_x, domain_y, domain_z)

    design = bz.Design(background=mat_sio2)
    design += bz.Box(center=(0.0, wg_center_y, 0.0), size=(bz.inf, wg_width, wg_height), material=mat_si)
    design += bz.Box(center=(0.0, -wg_center_y, 0.0), size=(bz.inf, wg_width, wg_height), material=mat_si)
    # Ring.z is the lower slab coordinate for BeamZ planar geometry.
    design += bz.Ring(
        position=(0.0, 0.0, -0.5 * wg_height),
        inner_radius=ring_radius - 0.5 * wg_width,
        outer_radius=ring_radius + 0.5 * wg_width,
        depth=wg_height, points=256, material=mat_si, color="#d81b60",
    )

    steps_per_wavelength = 4 if test_mode else 15
    resolution = lambda0 / (steps_per_wavelength * n_si)
    grid_spec = (
        bz.GridSpec.uniform(resolution)
        if test_mode
        else bz.GridSpec.auto(
            wavelength=lambda0, min_steps_per_wvl=steps_per_wavelength, max_scale=1.3
        )
    )
    source_size = (0.0, 1.2 * um, 0.8 * um)
    mode_spec = bz.ModeSpec(
        num_modes=1, mode_index=0, target_neff=3.47, polarization="te",
        num_freqs=1,
    )
    source = bz.ModeSource(
        center=(-ring_radius, wg_center_y, 0.0), size=source_size, direction="+",
        source_time=bz.GaussianPulse(freq0=freq0, fwidth=fwidth, offset=0.5 if test_mode else 4.0),
        mode_spec=mode_spec,
    )
    # The access guides are single mode here, so signed Poynting flux gives the
    # through/drop power spectra without solving monitor modes at every port.
    through = bz.FluxMonitor(
        center=(ring_radius, wg_center_y, 0.0), size=source_size,
        freqs=freqs, name="through",
    )
    drop = bz.FluxMonitor(
        center=(-ring_radius, -wg_center_y, 0.0), size=source_size,
        freqs=freqs, name="drop",
    )
    field = bz.FieldMonitor(
        center=(0.0, 0.0, 0.0), size=(domain_x, domain_y, 0.0), freqs=[freq0],
        fields=("Hz",), name="field",
    )
    simulation = bz.Simulation(
        domain=domain, design=design, sources=(source,), monitors=(through, drop, field),
        boundaries=(bz.PML(thickness=0.6 * um, formulation="cpml"),),
        grid_spec=grid_spec,
        run_time=(8 / freq0) if test_mode else 30e-12,
        setup_device="cpu" if test_mode else "default",
    )
    return simulation, {"source": source, "resolution": resolution}

sim, setup = make_sim()
memory = sim.memory_estimate(num_steps=sim.num_steps)
print(f"3D grid: {sim.grid.shape} | time steps: {sim.num_steps}")
print(f"Estimated known storage: {memory['total_gib']:.2f} GiB")

fig, axes = sim.plot(
    z=0.0, y=0.0, figsize=(12, 4.5), width_ratios=(1.1, 1.0),
    source_markers=True, monitor_markers=True, show=False,
)
for axis in np.asarray(axes).flat:
    axis.grid(False)
fig.suptitle("Microring add-drop layout: xy core plane and xz centerline", y=1.03)
plt.show()

# API opportunity: a single non-compiling `Simulation.estimate()` could combine
# memory, expected runtime, and direct-CUDA availability.

source_modes = setup["source"].solve_modes(sim, freqs=[freq0])
display(mode_data_to_dataframe(source_modes))

fig, axes, neffs = plot_mode_field_components(
    source_modes, field_names=("Ey", "Ez"), mode_indices=(0,), val="abs", f=freq0,
    figsize=(8, 3.8), show=False,
)
for axis in np.asarray(axes).flat:
    axis.grid(False)
plt.show()

from time import perf_counter

execution_backend = "jax" if test_mode else "cuda_streamed"
termination = None if test_mode else bz.AutoTermination(
    max_wall_time=600.0, field_decay=1e-5, monitor_change=1e-3, chunk_steps=2000, consecutive_checks=3,
    monitor_names=("through", "drop"),
)
wall_start = perf_counter()
results = sim.run(progress=not test_mode, backend=execution_backend, termination=termination)
print(f"Stopped: {results.termination.reason if results.termination else 'fixed duration'}")
if results.termination is not None:
    print(f"Measured {results.termination.time / 1e-12:.3f} ps; converged: {results.termination.converged}")
print(f"Wall time: {perf_counter() - wall_start:.1f} s (budget: 600 s)")

results.plot_field("field", "Hz", frequency=freq0, val="real", cmap="RdBu", show_grid=False)
plt.title("Hz at the central wavelength")
plt.show()

launched_power = max(float(results.launched_power()), 1e-18)
T_through = np.maximum(np.asarray(results["through"].flux, dtype=float), 0.0) / launched_power
T_drop = np.maximum(-np.asarray(results["drop"].flux, dtype=float), 0.0) / launched_power

fig, ax = plt.subplots(figsize=(9, 4))
ax.plot(wavelengths / um, 10 * np.log10(np.maximum(T_through, 1e-12)), label="Through", color="#d81b60")
ax.plot(wavelengths / um, 10 * np.log10(np.maximum(T_drop, 1e-12)), label="Drop", color="#1f77b4")
ax.set(xlabel="Wavelength (µm)", ylabel="Transmission (dB)", title="Microring through/drop spectrum")
if results.termination is not None and not results.termination.converged:
    ax.set_title("Microring through/drop spectrum — not yet converged")
ax.legend(frameon=False)
ax.grid(False)
plt.show()

wavelength_step_um = float(np.mean(np.diff(wavelengths / um)))
minimum_peak_spacing = max(1, int(round(0.012 / wavelength_step_um)))
peak_indices, peak_properties = find_peaks(
    T_drop, prominence=0.05, distance=minimum_peak_spacing
)
if len(peak_indices) >= 2:
    resonance_wavelengths = wavelengths[peak_indices] / um
    fsr_nm = 1e3 * np.diff(resonance_wavelengths)
    print(f"Resolved resonances (µm): {resonance_wavelengths}")
    print(f"Mean free spectral range: {np.mean(fsr_nm):.3f} nm")
else:
    print("Fewer than two drop-port peaks are resolved; increase run time and/or spectral sampling.")

if len(peak_indices):
    peak = peak_indices[np.argmax(T_drop[peak_indices])]
    half_power = 0.5 * T_drop[peak]
    left = np.flatnonzero(T_drop[:peak] <= half_power)
    right = np.flatnonzero(T_drop[peak + 1:] <= half_power)
    if len(left) and len(right):
        linewidth = wavelengths[peak + 1 + right[0]] - wavelengths[left[-1]]
        print(f"Approximate loaded Q: {wavelengths[peak] / linewidth:.0f}")

# API opportunity: a public harmonic-inversion/resonance-analysis helper could
# estimate Q from a short time-domain trace rather than manual spectral peak analysis.
# Treat the printed Q as approximate: check both temporal convergence and
# spectral sampling before interpreting the linewidth quantitatively.