Inverse design

Use differentiable simulation and MaterialRegion for gradient-based optimization

Overview#

WiTwin Maxwell supports gradient-based inverse design through PyTorch autodiff. Gradients flow through the FDTD simulation via the adjoint backward path. This tutorial shows two approaches:

  1. Trainable geometry parameter — differentiate with respect to a structure position
  2. MaterialRegion — density-based topology optimization

Approach 1 — Trainable Geometry Parameter#

Differentiate the probe field with respect to a cube's x position:

import torch
import witwin.maxwell as mw
 
# Trainable parameter
box_x = torch.tensor(0.10, device="cuda", requires_grad=True)
box_position = torch.stack((box_x, box_x.new_tensor(0.0), box_x.new_tensor(0.06)))
 
scene = mw.Scene(
    domain=mw.Domain(bounds=((-0.24, 0.24), (-0.24, 0.24), (-0.24, 0.24))),
    grid=mw.GridSpec.uniform(0.12),
    boundary=mw.BoundarySpec.pml(num_layers=2, strength=1.0),
    device="cuda",
    subpixel_samples=5,
)
 
scene.add_structure(
    mw.Structure(
        name="cube",
        geometry=mw.Box(position=box_position, size=(0.18, 0.18, 0.18)),
        material=mw.Material(eps_r=20.0),
    )
)
scene.add_source(
    mw.PointDipole(
        position=(0.0, 0.0, -0.06),
        polarization="Ez",
        width=0.04,
        source_time=mw.GaussianPulse(frequency=1.0e9, fwidth=0.25e9, amplitude=50.0),
    )
)
scene.add_monitor(mw.PointMonitor("probe", (0.0, 0.0, 0.06), fields=("Ez",)))
 
sim = mw.Simulation.fdtd(
    scene,
    frequencies=[1.0e9],
    run_time=mw.TimeConfig(time_steps=32),
    spectral_sampler=mw.SpectralSampler(window="none"),
    full_field_dft=True,
)
 
result = sim.run()
probe = result.monitor("probe")["data"]
loss = torch.abs(probe) ** 2
loss.backward()
 
print("d(loss)/d(box_x) =", box_x.grad)

Approach 2 — MaterialRegion Topology Optimization#

Use a trainable density tensor to optimize a material distribution:

import torch
import witwin.maxwell as mw
 
# Trainable density field
density = torch.ones(8, 8, 8, device="cuda", requires_grad=True)
 
class MyDesign(mw.SceneModule):
    def __init__(self):
        super().__init__()
        self.density = torch.nn.Parameter(torch.ones(8, 8, 8))
 
    def to_scene(self):
        scene = mw.Scene(
            domain=mw.Domain(bounds=((-0.24, 0.24), (-0.24, 0.24), (-0.24, 0.24))),
            grid=mw.GridSpec.uniform(0.12),
            boundary=mw.BoundarySpec.pml(num_layers=2, strength=1.0),
            device="cuda",
        )
        scene.add_material_region(
            mw.MaterialRegion(
                name="design",
                geometry=mw.Box(position=(0.0, 0.0, 0.0), size=(0.18, 0.18, 0.18)),
                density=self.density,
                eps_bounds=(1.0, 12.0),
                filter_radius=2.0,
                projection_beta=10.0,
            )
        )
        scene.add_source(
            mw.PointDipole(
                position=(0.0, 0.0, -0.06),
                polarization="Ez",
                width=0.04,
                source_time=mw.GaussianPulse(frequency=1.0e9, fwidth=0.25e9, amplitude=50.0),
            )
        )
        scene.add_monitor(mw.PointMonitor("probe", (0.0, 0.0, 0.06), fields=("Ez",)))
        return scene
 
model = MyDesign()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
 
for step in range(10):
    optimizer.zero_grad()
    sim = mw.Simulation.fdtd(
        model,
        frequencies=[1.0e9],
        run_time=mw.TimeConfig(time_steps=32),
        spectral_sampler=mw.SpectralSampler(window="none"),
        full_field_dft=True,
    )
    result = sim.run()
    probe = result.monitor("probe")["data"]
    loss = -torch.abs(probe) ** 2
    loss.backward()
    optimizer.step()
    print(f"step {step}: loss={float(loss):.4f}")

What the Adjoint Does#

When scene contains trainable parameters or is a SceneModule, sim.run() automatically routes through the adjoint backward path. The forward FDTD solve is checkpointed and the backward pass computes gradients with respect to compiled material tensors and supported trainable geometry parameters.

Note

The current public backward path targets trainable inputs that flow into compiled material tensors — MaterialRegion.density and supported trainable geometry parameters. Source placement and other non-material branches are not yet part of the public adjoint path.

Next Steps#