Four-channel WDM inverse design in BeamZ

from pathlib import Path
import time, json, gc
import numpy as np
import jax
import jax.numpy as jnp
from jax.scipy.special import logsumexp
import matplotlib.pyplot as plt
from IPython.display import display, Markdown
from beamz import (LIGHT_SPEED, Design, Material, Rectangle, Simulation, ModeSource,
                   ModeMonitor, FieldMonitor, FluxMonitor, ModeSpec, GaussianPulse, PML)
import beamz.optimization as opt

um = 1e-6
wavelengths_design = np.array([1.270, 1.290, 1.310, 1.330])*um
frequencies_design = LIGHT_SPEED / wavelengths_design
frequency = frequencies_design.mean()
fwidth = np.ptp(frequencies_design)
channel_width = abs(np.diff(frequencies_design).mean()) / 2
channel_frequencies = np.array([np.linspace(f-channel_width/2, f+channel_width/2, 5)
                                for f in frequencies_design])
training_frequencies = channel_frequencies.ravel()
output_y = (1.5 + np.linspace(.5625, 3.9375, 4))*um


root=next(p for p in [Path.cwd(),*Path.cwd().parents] if (p/"pyproject.toml").exists())
artifacts=root/"results/topology_examples/topology_wdm_4channel"
artifacts.mkdir(parents=True,exist_ok=True)
plt.rcParams.update({"figure.dpi":110,"font.size":10})
print("JAX:",jax.__version__,"backend:",jax.default_backend())

def make_wdm(resolution=50e-9, duration=6e-12, backend="adjoint", grouping="frequency"):
    air, silicon = Material(1), Material(3.49**2)
    geometry = Design(width=7.5*um, height=7.5*um, material=air)
    geometry += Rectangle(position=(0,3.6*um), width=1.5*um, height=.3*um, material=silicon)
    for y in output_y:
        geometry += Rectangle(position=(6*um,y-.15*um), width=1.5*um, height=.3*um, material=silicon)
    mode = ModeSpec(num_modes=1, polarization="tm")
    source = ModeSource(center=(.75*um,3.75*um,0),size=(0,1.125*um,um),direction="+",
        mode_spec=mode, source_time=GaussianPulse(freq0=frequency,fwidth=fwidth,offset=5/(2*np.pi)))
    outputs = [ModeMonitor(center=(6.75*um,y,0),size=(0,1.0125*um,um),freqs=training_frequencies,
                mode_spec=mode,name=f"mode_{i}") for i,y in enumerate(output_y)]
    incoming = ModeMonitor(center=(1.125*um,3.75*um,0),size=(0,1.125*um,um),
                           freqs=training_frequencies,mode_spec=mode,name="input")
    simulation = Simulation(design=geometry,resolution=resolution,run_time=duration,polarization="tm",
        sources=[source],monitors=[*outputs,incoming],boundaries=[PML(thickness=.45*um,formulation="cpml")])
    region = opt.TopologyDesignRegion(center=(3.75*um,3.75*um,0),size=(4.5*um,4.5*um,np.inf),
        pixel_size=resolution,eps_bounds=(1,3.49**2),transformations=[opt.FilterProject(radius=.1*um,beta=1)],
        penalties=[opt.ErosionDilationPenalty(length_scale=.1*um,beta=10)],penalty_input="parameters")
    return opt.InverseDesign(simulation,region,output_monitor_names=[f"mode_{i}" for i in range(4)],
        gradient_backend=backend,adjoint_source_grouping=grouping,checkpoint_interval=128)


def channel_powers(data):
    # Rows are output ports; columns are wavelength channels, five frequencies each.
    return jnp.stack([jnp.abs(data[f"mode_{i}"].amps.sel(direction="+",mode_index=0).values.reshape(4,5))**2
                      for i in range(4)]).mean(axis=-1)


def post_process(data, leak_weight=0.0):
    power = channel_powers(data)
    desired = jnp.diag(power)
    leakage = (power.sum(axis=1)-desired)/3
    return -logsumexp(-(desired-leak_weight*leakage))

schedule = opt.OptimizationSchedule(beta=opt.LinearSchedule(1,50,50),
    objective_kwargs={"leak_weight":opt.StepSchedule(0,1,17)})

design=make_wdm()
params0=design.design_region.initial_parameters
print("Design parameters:",params0.shape,params0.size)
print("Training frequencies:",len(training_frequencies),"timesteps:",len(design.simulation.time))
fig,ax=design.initial_simulation.plot(figsize=(7,6))
ax.set_title("Four-channel WDM · initial geometry")
plt.show()

check=make_wdm(resolution=75e-9,backend="autodiff")
probe=check.design_region.initial_parameters+.015*np.random.default_rng(7).normal(size=check.design_region.params_shape)
settings=schedule.values(20)
values,gradients,timings,diagnostics={},{},{},{}
for label,backend,grouping in [("autodiff","autodiff","frequency"),("frequency","adjoint","frequency"),("auto","adjoint","auto")]:
    candidate=check if label=="autodiff" else check.updated_copy(gradient_backend=backend,adjoint_source_grouping=grouping)
    fn=candidate.make_objective_fn(post_process,**settings)
    started=time.perf_counter()
    value,gradient=jax.value_and_grad(fn)(probe)
    gradients[label]=np.asarray(gradient)
    values[label]=float(value)
    timings[label]=time.perf_counter()-started
    if backend=="adjoint":diagnostics[label]=candidate.gradient_diagnostics
    print(label,"score",values[label],"seconds",timings[label],diagnostics.get(label,{}),flush=True)
gradient_errors={name:float(np.linalg.norm(g-gradients["autodiff"])/np.linalg.norm(gradients["autodiff"])) for name,g in gradients.items()}
direction=gradients["autodiff"]/np.linalg.norm(gradients["autodiff"])
fn=check.make_objective_fn(post_process,**settings)
h=.01
finite_difference=float((fn(probe+h*direction)-fn(probe-h*direction))/(2*h))
np.testing.assert_allclose(finite_difference,np.sum(gradients["autodiff"]*direction),rtol=.01,atol=1e-6)
assert max(gradient_errors.values())<.005
print("Relative gradient errors:",gradient_errors,"directional finite difference:",finite_difference)
del check,candidate,fn,gradients
jax.clear_caches();gc.collect()

chosen=min(timings,key=timings.get)
backend="autodiff" if chosen=="autodiff" else "adjoint"
grouping="auto" if chosen=="auto" else "frequency"
design=design.updated_copy(gradient_backend=backend,adjoint_source_grouping=grouping)
fn=design.make_objective_fn(post_process,**schedule.values(0))
elapsed=[]
for repeat in range(2):
    start=time.perf_counter()
    value,gradient=jax.value_and_grad(fn)(params0)
    np.asarray(gradient)
    elapsed.append(time.perf_counter()-start)
profile={"coarse_seconds":timings,"coarse_adjoint_diagnostics":diagnostics,
         "training_first_seconds":elapsed[0],"training_warm_seconds":elapsed[1]}
print("Selected backend:",chosen,"training profile:",profile,flush=True)

optimizer = opt.AdamOptimizer(design=design, learning_rate=.1, schedule=schedule,
    objective_id="reference-wdm-four-channel-v1")
checkpoint = artifacts / "history.npz"
started = time.perf_counter()
progress_rows = []
def report(result):
    h = result.history[-1]
    row = dict(step=h.step, score=h.objective_before, post=h.post_process_val,
        penalty=h.penalty, elapsed=time.perf_counter()-started,
        **result.settings(), **(design.gradient_diagnostics or {}))
    progress_rows.append(row)
    (artifacts/"progress.json").write_text(json.dumps(progress_rows, indent=2))
    print(f"Step {row['step']:2d}: score={row['score']:.5f}, beta={row['beta']:.1f}, "
          f"leak={row['post_process_kwargs']['leak_weight']:.0f}, elapsed={row['elapsed']:.1f}s", flush=True)
partial = optimizer.run(post_process, steps=25, checkpoint=checkpoint,
                        checkpoint_every=5, callback=report)
restored = optimizer.load_result(checkpoint, post_process)
np.testing.assert_array_equal(partial.final_params, restored.final_params)
assert restored.settings()["beta"] == 25
result = optimizer.run(post_process, steps=50, resume=checkpoint,
                       checkpoint=checkpoint, checkpoint_every=5, callback=report)
optimization_seconds = time.perf_counter()-started
assert result.completed_steps == 50 and result.settings()["beta"] == 50
print("Completed all 50 updates with checkpoint continuation.")

fig,axes=plt.subplots(1,2,figsize=(12,4))
result.plot_optimization(ax=axes[0])
axes[0].axvline(17,color="black",ls=":",label="Leakage penalty enabled")
axes[0].legend()
axes[1].plot([s["beta"] for s in result.schedule_history])
axes[1].set(xlabel="Iteration",ylabel="Projection beta",title="Continuation schedule")
fig.tight_layout();plt.show()
final_simulation=result.to_simulation()
fig,ax=final_simulation.plot_eps(z=0,figsize=(7,6))
ax.set_title("Optimized four-channel WDM · beta=50")
plt.show()

frequency_spacing=abs(np.diff(frequencies_design).mean())
freqs_measure=np.linspace(frequencies_design.min()-frequency_spacing,frequencies_design.max()+frequency_spacing,151)
mode_monitors=[m.updated_copy(freqs=freqs_measure) for m in final_simulation.monitors]
flux_monitors=[FluxMonitor(center=m.center,size=m.size,freqs=freqs_measure,name=f"flux_{i}") for i,m in enumerate(mode_monitors[:4])]
field_monitor=FieldMonitor(center=(3.75*um,3.75*um,0),size=(7.5*um,7.5*um,0),freqs=frequencies_design,fields=("Ez",),name="field")
verified=final_simulation.updated_copy(monitors=[*mode_monitors,*flux_monitors,field_monitor]).run(backend="jax",performance=False)
modal_power=np.stack([np.abs(verified.mode(f"mode_{i}").amps.sel(direction="+",mode_index=0).values).reshape(-1)**2 for i in range(4)])
incident_power=np.abs(verified.mode("input").amps.sel(direction="+",mode_index=0).values).reshape(-1)**2
flux_power=np.stack([np.asarray(verified[f"flux_{i}"].flux).reshape(-1) for i in range(4)])
wavelengths_nm=LIGHT_SPEED/freqs_measure/1e-9
fig,axes=plt.subplots(1,2,figsize=(13,4))
for ax,powers,title in [(axes[0],modal_power,"Fundamental-mode power"),(axes[1],flux_power,"Total flux")]:
    for i in range(4):
        ax.plot(wavelengths_nm,10*np.log10(np.maximum(powers[i],1e-15)),label=f"Port {i}: {wavelengths_design[i]/1e-9:.0f} nm")
        bounds=LIGHT_SPEED/channel_frequencies[i,[0,-1]]/1e-9
        ax.axvspan(min(bounds),max(bounds),alpha=.08)
    ax.set(xlabel="Wavelength (nm)",ylabel="Source-normalized power (dB)",title=title,ylim=(-45,3))
    ax.legend();ax.grid(alpha=.25)
fig.tight_layout();plt.show()
fig,axes=plt.subplots(1,4,figsize=(16,4))
for f,ax,w in zip(frequencies_design,axes,wavelengths_design):
    verified.plot_field("field",field_name="Ez",val="abs^2",frequency=f,ax=ax)
    ax.set_title(f"{w/1e-9:.0f} nm")
fig.tight_layout();plt.show()

binary=result.export_design()
verification_rows=[]
for resolution,duration in [(50e-9,6e-12),(30e-9,6e-12),(30e-9,12e-12)]:
    sim=design.simulation.updated_copy(design=binary,resolution=resolution,run_time=duration)
    data=sim.run(backend="jax",performance=False)
    output=np.asarray(channel_powers({f"mode_{i}":data.mode(f"mode_{i}") for i in range(4)}))
    input_samples=np.abs(data.mode("input").amps.sel(direction="+",mode_index=0).values.reshape(4,5))**2
    normalized=np.stack([np.abs(data.mode(f"mode_{i}").amps.sel(direction="+",mode_index=0).values.reshape(4,5))**2/input_samples for i in range(4)]).mean(axis=-1)
    desired=np.diag(normalized)
    unwanted=normalized.copy();np.fill_diagonal(unwanted,0)
    isolation=10*np.log10(desired/np.maximum(unwanted.max(axis=0),1e-15))
    row=dict(mesh_nm=resolution/1e-9,run_time_ps=duration/1e-12,desired=desired.tolist(),max_unwanted=unwanted.max(axis=0).tolist(),isolation_db=isolation.tolist())
    verification_rows.append(row)
    print(row,flush=True)
fig,ax=sim.plot(figsize=(7,6));ax.set_title("Exported binary WDM");plt.show()
metrics=dict(gradient_errors=gradient_errors,profile=profile,gradient_backend=design.gradient_backend,adjoint_source_grouping=design.adjoint_source_grouping,optimization_seconds=optimization_seconds,verification=verification_rows)
(artifacts/"metrics.json").write_text(json.dumps(metrics,indent=2))
np.savez_compressed(artifacts/"spectra.npz",frequencies=freqs_measure,modal_power=modal_power,incident_power=incident_power,flux_power=flux_power)
display(Markdown("**Finer-grid binary verification:** desired channel transmissions = "+", ".join(f"{v:.1%}" for v in verification_rows[-1]["desired"])+". See the full per-channel leakage and mesh/duration comparisons above."))

import os
from uuid import uuid4
from beamz.design import export_gds

# Unique cell names allow re-running this cell in the same gdsfactory session.

gds_root = next(p for p in [Path.cwd(), *Path.cwd().parents] if (p / "pyproject.toml").exists())
gds_dir = Path(os.environ.get("BEAMZ_GDS_DIR", gds_root / "results/topology_examples/gds")).expanduser()
gds_dir.mkdir(parents=True, exist_ok=True)

binary_layout = result.export_design(threshold=0.5)
# Clearing polygons represent cladding, not material to pattern.
core_layout = Design(
    width=binary_layout.width, height=binary_layout.height, depth=binary_layout.depth,
    background=binary_layout.background,
    structures=[s for s in binary_layout.structures
                if s.material.permittivity > binary_layout.background.permittivity],
)
gds_path = export_gds(core_layout, gds_dir / "topology_wdm_4channel.gds",
                      cell_name=f"wdm_4channel_{uuid4().hex[:8]}")
print(f"GDS core mask (layer 0/0): {gds_path.resolve()}")