Sphere radius optimization

A lossless sphere of refractive index 2 sits in vacuum, lit at a wavelength of 600 nm. Its scattering efficiency Q, the scattering cross section divided by π r², depends on the radius r. Starting at r = 190 nm, five steps of gradient ascent, r → r + 50 nm² · dQ/dr, reach the maximum Q = 5.759 at r = 201.75 nm. The derivative dQ/dr comes from Advect through the analytic derivatives of treams-rs; before the first step, the script checks it against central differences. treams computes no derivatives, so this example has no treams version.

# Requires: advect
"""Gradient ascent on the radius of a dielectric sphere to maximize its scattering."""

import advect
import numpy as np

import treams_rs as tr
from treams_rs.testing import check_gradient

STEPS = 5

# Lengths are in nm: light of 600 nm wavelength meets a sphere of refractive index 2.
k0 = 2 * np.pi / 600
material = 4  # permittivity 4, refractive index 2
lmax = 6
start = 190.0
# The gradient has units of 1/nm, so the step size has units of nm^2.
step_size = 50.0


def efficiency(radius):
    """Scattering cross section divided by the geometric cross section."""
    sphere = tr.sphere_tmatrix(k0=k0, lmax=lmax, radius=radius, material=material)
    # A sphere scatters every plane wave alike, so one wave gives the average.
    incident = tr.plane_wave([0, 0, 1], "positive_helicity", k0=k0)
    return sphere.cross_sections(incident).scattering / (np.pi * radius**2)


# Compare the gradient with central differences before using it.
check_gradient(
    efficiency, advect.grad(efficiency), np.asarray(start), step=1e-3, rtol=1e-6, atol=0
)

# Rows hold (radius, efficiency, gradient) before each step and after the last.
history = np.zeros((STEPS + 1, 3))
radius = start
for i in range(STEPS + 1):
    value, gradient = advect.value_and_grad(efficiency)(np.asarray(radius))
    history[i] = radius, value, gradient
    radius = radius + step_size * gradient

if __name__ == "__main__":
    print("Gradient ascent on the scattering efficiency Q")
    print(f"{'step':>6}{'radius/nm':>14}{'Q':>14}{'dQ/dr in 1/nm':>16}")
    for i, (r, q, dq) in enumerate(history):
        print(f"{i:>6}{r:>14.8g}{q:>14.8g}{dq:>16.6g}")

Output

Gradient ascent on the scattering efficiency Q
  step     radius/nm             Q   dQ/dr in 1/nm
     0           190     4.6472917        0.121572
     1     196.07859     5.4235962        0.112069
     2     201.68203     5.7591286      0.00144098
     3     201.75408     5.7591807     6.83677e-06
     4     201.75442     5.7591807      5.7162e-08
     5     201.75442     5.7591807     4.78922e-10

How it works

• treams_rs selects Advect from the traced radius, so advect.value_and_grad(efficiency) returns Q and dQ/dr. The constant incident wave composes with the differentiated sphere automatically. Framework adapters lists the JAX, PyTorch and HIPS Autograd equivalents.
• Rust computes each derivative analytically: a pullback maps the gradient with respect to an output, here the cross section, back to the inputs, here the radius. Differentiation describes these pullbacks.
• check_gradient compares dQ/dr with central differences of step 1e-3 nm in a random direction and raises an error when they differ by more than 1e-6 relative; here they differ by about 1e-9. Gradient checks lists its options.