Optimize with gradients

Fit wall material parameters to measured path gains with a PyTorch optimizer, and avoid the common pitfalls of gradient-based inverse problems in ray tracing.

Channel results are differentiable, so you can fit scene parameters to measurements with any PyTorch optimizer. This guide calibrates the permittivity and conductivity of a wall from ten path-gain measurements. The transmitter localization tutorial does the same for a position.

Fit a wall material#

import torch
 
from witwin.core import AntennaState, Mesh, PhysicalMaterial, Scene, Structure
from witwin.channel import deterministic
 
FREQUENCY_HZ = 3.5e9
vertices = torch.tensor([[2.5, -3.0, -1.0], [2.5, 3.0, -1.0], [2.5, -3.0, 2.0], [2.5, 3.0, 2.0]])
faces = torch.tensor([[0, 1, 2], [1, 3, 2]])
tx = AntennaState(1, "tx", torch.tensor([0.0, 0.0, 0.5]), power_w=1.0)
# Five receivers behind the wall and five in front of it.
receivers = [AntennaState(10 + i, "rx", torch.tensor([5.0, -2.0 + i, 0.5])) for i in range(5)]
receivers += [AntennaState(20 + i, "rx", torch.tensor([0.0, 1.0 + i, 0.5])) for i in range(5)]
config = deterministic.Config(max_depth=1, components={"reflection", "transmission"}, ad_mode="vjp")
 
 
def solve(eps_r, sigma_e):
    wall = Structure(
        Mesh(vertices, faces, recenter=False, fill_mode="surface"),
        PhysicalMaterial(eps_r=eps_r, sigma_e=sigma_e, thickness_m=0.0125),
    )
    scene = Scene(structures=[wall], endpoints=[tx, *receivers])
    return deterministic.solve(scene, config, reference_frequency_hz=FREQUENCY_HZ).path_gain
 
 
# "Measurements" from a wall with unknown properties.
with torch.no_grad():
    measured_db = 10 * torch.log10(solve(2.73, 0.0085))
 
# Start from a guess and fit both parameters to the measured gains.
eps_r = torch.tensor(2.0, device="cuda", requires_grad=True)
log_sigma = torch.tensor(-3.0, device="cuda", requires_grad=True)   # sigma_e = exp(log_sigma) stays positive
optimizer = torch.optim.Adam([eps_r, log_sigma], lr=0.05)
for step in range(501):
    optimizer.zero_grad()
    predicted_db = 10 * torch.log10(solve(eps_r, log_sigma.exp()))
    loss = ((predicted_db - measured_db) ** 2).mean()
    loss.backward()
    optimizer.step()
    if step % 100 == 0:
        print(f"step {step:3d}  loss {loss.item():.2e} dB^2  eps_r {eps_r.item():.3f}  sigma_e {log_sigma.exp().item():.4f}")
print(f"fitted eps_r {eps_r.item():.3f}, sigma_e {log_sigma.exp().item():.4f} S/m")

Output (about 16 s on an RTX 5080):

step   0  loss 5.60e+00 dB^2  eps_r 2.050  sigma_e 0.0474
step 100  loss 3.01e-03 dB^2  eps_r 2.742  sigma_e 0.0140
step 200  loss 5.30e-04 dB^2  eps_r 2.736  sigma_e 0.0108
step 300  loss 1.10e-04 dB^2  eps_r 2.733  sigma_e 0.0095
step 400  loss 2.22e-05 dB^2  eps_r 2.731  sigma_e 0.0090
step 500  loss 4.01e-06 dB^2  eps_r 2.731  sigma_e 0.0087

The fit recovers eps_r = 2.73 and approaches sigma_e = 0.0085. Conductivity converges more slowly because a 12.5 mm sheet absorbs little, so the loss is less sensitive to it.

The pattern#

  1. Make the parameter a tensor leaf. Pass a 0-d tensor with requires_grad=True wherever Channel accepts a number: PhysicalMaterial fields, AntennaState(position=...), Mesh(vertices=...) or reference_frequency_hz. See Differentiability for the full list.
  2. Set ad_mode="vjp" in the solver Config. With the default "none", results have no gradient.
  3. Rebuild the scene inside the loop. Scene objects are immutable; build a new Structure or AntennaState from the current tensor each step. Channel caches what it can.
  4. Reparametrize constrained values. sigma_e must be non-negative, so the example optimizes its logarithm. Use the same trick for thicknesses.
  5. Compare in dB. Path gains span many orders of magnitude; a loss on 10 * log10(gain) weights all receivers evenly.

Pitfalls#

Non-convex losses. Wave physics produces oscillating losses. A thick slab is a resonator: with the same setup and a 20 cm concrete wall (eps_r = 5.24), the loss oscillates along eps_r, and Adam started from eps_r = 3.0 and 4.0 stopped at 2.84 and 4.36 instead of 5.24. Use measurements that constrain the parameter from several angles, start near a plausible value, or scan a coarse grid first and refine with gradients.

Topology changes. Gradients follow the paths found in the current solve. When a step moves a reflection point off a wall or a receiver into a shadow, a path appears or disappears and the loss jumps; the gradient did not predict that. Keep steps small relative to the scene geometry, and watch for jumps in the loss.

Unsupported combinations raise. Coupled paths with mesh-vertex gradients and frequency gradients with dispersive materials are refused. With "scattering" the solve runs but backward() raises NotImplementedError. Remove the component or the leaf.

Monte Carlo noise. montecarlo.basic supports gradients of the power map, but its values are noisy estimates. Fix the seed so successive steps sample the same rays, and use enough samples.

Forward mode#

For the derivative along one direction (a Jacobian-vector product), use forward mode instead of backward():

import torch.autograd.forward_ad as fwAD
 
with fwAD.dual_level():
    position = fwAD.make_dual(torch.tensor([0.0, -1.0, 0.5], device="cuda"), torch.tensor([1.0, 0.0, 0.0], device="cuda"))
    scene = Scene(structures=[wall], endpoints=[AntennaState(1, "tx", position, power_w=1.0), rx])
    result = deterministic.solve(scene, deterministic.Config(max_depth=1, components={"los", "reflection"}, ad_mode="jvp"), reference_frequency_hz=3.0e9)
    print(fwAD.unpack_dual(result.path_gain).tangent)   # d(path_gain) / d(x of the transmitter)