Neural Networks

Use differentiable simulation for gradient-based optimization and inverse design

Overview#

WiTwin Channel is differentiable end-to-end. Gradients flow through scene geometry, material parameters, transmitter positions, and solver outputs via Dr.Jit autodiff. This enables gradient-based optimization workflows — fitting scene parameters to measured signals, optimizing antenna placement, or training neural networks on simulated channel data.

How Differentiability Works#

WiTwin Channel uses Dr.Jit as its autodiff backend. To differentiate through a simulation:

  1. Wrap your parameter in a Dr.Jit array
  2. Enable gradients with dr.enable_grad()
  3. Run the simulation
  4. Call dr.backward() or dr.forward_to() to propagate gradients

Example: Material Parameter Gradient#

The example below differentiates the path gain with respect to a material's permittivity eps_r:

import drjit as dr
import witwin.channel as wc
 
cube_eps = wc.Float(4.0)
dr.enable_grad(cube_eps)
 
scene = wc.Scene(
    structures=[
        wc.Structure(
            name="cube",
            geometry=wc.Box(
                position=(0.0, 0.0, 1.5),
                size=(1.0, 1.0, 1.0),
                device="cuda",
            ),
            material=wc.Material(eps_r=4.0, sigma_e=0.0),
        ),
    ],
    transmitters=[wc.Transmitter("tx", (0.0, -3.0, 1.5))],
    receivers=[
        wc.ReceiverGrid(
            "rm",
            axis="z",
            position=1.5,
            bounds=((-0.5, 0.5), (1.5, 2.5)),
            grid_shape=(1, 1),
        ),
    ],
    frequency=1.0e9,
    device="cuda",
)
 
scene.structure("cube").set_material_parameters(eps_r=cube_eps)
 
result = wc.deterministic.solve(
    scene=scene,
    transmitter="tx",
    receiver="rm",
    config=wc.deterministic.Config(
        num_samples=256,
        max_bounces=2,
        max_diffraction_order=1,
        edge_policy=wc.EdgePolicy(edge_selection_mode="all_edges"),
    ),
)
 
dr.backward(dr.sum(result.path_gain))
grad = float(dr.grad(cube_eps))
print(f"d(path_gain) / d(eps_r) = {grad}")

Example: Transmitter Position Recovery#

The example below uses the Dr.Jit Adam optimizer to recover a transmitter x position from a target signal measurement:

import drjit as dr
from drjit.opt import Adam
import witwin.channel as wc
 
opt = Adam(lr=0.02)
opt["tx_x"] = wc.Float(0.1)
 
for step in range(80):
    tx_x = opt["tx_x"]
 
    scene = wc.Scene(
        structures=[
            wc.Structure(
                name="wall",
                geometry=wc.Box(position=(0.0, 0.0, 1.5), size=(0.25, 4.0, 3.0), device="cuda"),
                material=wc.Material(eps_r=4.0, sigma_e=0.0),
            ),
        ],
        transmitters=[wc.Transmitter("tx", wc.Point3f(tx_x, -1.0, 0.5))],
        receivers=[
            wc.ReceiverGrid(
                "rm",
                axis="z",
                position=1.5,
                bounds=((-3.0, 3.0), (-3.0, 3.0)),
                grid_shape=(8, 8),
            ),
        ],
        frequency=1.0e9,
        device="cuda",
    )
 
    result = wc.deterministic.solve(
        scene=scene,
        transmitter="tx",
        receiver="rm",
        config=wc.deterministic.Config(num_samples=64, max_bounces=1),
    )
 
    loss = dr.mean((result.path_gain - target) ** 2)
    dr.backward(loss)
    opt.step()
 
print(f"Recovered tx_x = {float(opt['tx_x']):.4f}")

What Is Differentiable#

ParameterDifferentiable
Material eps_r✓
Material sigma_e✓
Transmitter position✓
Mesh vertex positions✓
result.path_gain✓
result.rss✓
Field components via solve_field✓

Forward and Reverse Mode#

Dr.Jit supports both forward (JVP) and reverse (VJP) mode autodiff:

# Reverse mode — gradient of scalar loss w.r.t. parameter
dr.backward(dr.sum(result.path_gain))
grad = dr.grad(cube_eps)
 
# Forward mode — directional derivative
dr.set_grad(cube_eps, 1.0)
jvp = dr.forward_to(result.path_gain)

Note

Use reverse mode when you have many parameters and a scalar loss. Use forward mode when you have few parameters and want the full Jacobian.

Next Steps#