Cluster scattering gradient

Which sphere radii or positions would increase the scattering cross section? Eight dielectric spheres scatter a circularly polarized plane wave at a vacuum wavelength of 700 nm. Their radii range from 45 to 54 nm; their relative permittivity is 4 + 0.02j. The objective is the scattering cross section: the scattered power divided by the incident intensity, in nm².

jax.value_and_grad computes this scalar and its derivatives with respect to all 32 parameters: one radius and three position coordinates per sphere. This is a natural reverse-mode problem: start with the one output of interest and work back to every input. Forward mode would follow 32 input directions to recover the same full gradient.

The script also changes all the radii and positions slightly, then checks that the gradient predicts the resulting change in cross section. The prediction and a central difference agree to 1e-6 relative tolerance.

# Requires: jax
# /// script
# requires-python = ">=3.12"
# dependencies = ["numpy>=2.1,<3", "jax>=0.10,<0.12", "treams-rs"]
# ///
"""Which sphere radii or positions increase the scattering cross section?"""

import jax
import jax.numpy as jnp
import numpy as np

import treams_rs as tr

# Keep double precision for the comparison with small finite differences.
jax.config.update("jax_enable_x64", True)


def make_problem(particles=8):
    """Return a scalar objective and (radius, x, y, z) parameters, all in nm."""
    # A slightly irregular stack of squares, with separated dielectric spheres.
    index = np.arange(particles)
    positions = np.column_stack(
        (180 * (index % 2), 190 * ((index // 2) % 2), 200 * (index // 4))
    ).astype(float)
    positions += np.column_stack(
        (9 * np.sin(index), 7 * np.cos(index), 11 * np.sin(2 * index))
    )
    positions -= positions.mean(axis=0)
    initial = np.column_stack((45 + 3 * (index % 4), positions))

    k0 = 2 * np.pi / 700  # vacuum wavelength: 700 nm
    incident = tr.plane_wave([0.2, -0.1, 1], "positive_helicity", k0=k0)

    def scattering(parameters):
        """Cluster scattering cross section in nm², for fixed illumination."""
        # Use the same calculation for the gradient and the difference check.
        parameters = jnp.asarray(parameters)
        spheres = [
            tr.sphere_tmatrix(
                k0=k0, lmax=1, radius=parameters[i, 0], material=4 + 0.02j
            )
            for i in range(particles)
        ]
        response = tr.Cluster(spheres, positions=parameters[:, 1:]).solve()
        return response.cross_sections(incident).scattering

    return scattering, initial


def main():
    scattering, initial = make_problem()
    value, gradient = jax.value_and_grad(scattering)(jnp.asarray(initial))
    gradient = np.asarray(gradient)

    # Perturb every radius and position together, along a unit direction.
    direction = np.sin(np.arange(initial.size) + 0.5).reshape(initial.shape)
    direction /= np.linalg.norm(direction)
    step = 1e-3  # nm
    difference = (
        scattering(initial + step * direction) - scattering(initial - step * direction)
    ) / (2 * step)
    projection = np.sum(gradient * direction)
    np.testing.assert_allclose(projection, difference, rtol=1e-6, atol=1e-9)

    print(f"Scattering cross section: {float(value):.8f} nm^2")
    print(f"One reverse pass: {initial.size} radius and position derivatives")
    print("Gradient in nm (cross section in nm^2, parameters in nm)")
    print(f"{'sphere':>7}{'radius':>13}{'x':>13}{'y':>13}{'z':>13}")
    for i, row in enumerate(gradient):
        print(f"{i:>7}" + "".join(f"{entry:>13.6f}" for entry in row))
    print(f"Joint directional derivative: {projection:.8f} nm")
    print(f"Central difference:           {float(difference):.8f} nm")


if __name__ == "__main__":
    main()

Output

Scattering cross section: 4615.81617506 nm^2
One reverse pass: 32 radius and position derivatives
Gradient in nm (cross section in nm^2, parameters in nm)
 sphere       radius            x            y            z
      0    55.551000     2.293551     2.564654     0.525056
      1    67.072027    -2.570098     3.713243     0.107474
      2    80.717968     3.700719    -2.462213     1.821015
      3    96.960034    -3.953498    -3.376162     1.053288
      4    51.931349     2.326181     2.315960    -0.445444
      5    65.804871    -2.177511     3.052389    -1.003933
      6    75.430202     3.981540    -2.905870    -0.225639
      7    88.221359    -3.600884    -2.902002    -1.831817
Joint directional derivative: -4.28130461 nm
Central difference:           -4.28130461 nm

How it works

• Rows of parameters hold (radius, x, y, z) in nm. JAX tracks the radius through sphere_tmatrix and the positions through Cluster, including multiple scattering, the incident phase, and interference in the outgoing power.
• Cluster.solve() constructs the coupled T-matrix; cross_sections turns it into one scattering cross section. Its gradient includes repeated scattering between the spheres.
• jax_enable_x64 retains double precision for the finite-difference check. The returned gradient has units of nm: cross section in nm² differentiated with respect to a length in nm.
• The material, incident wave, and multipole truncation remain fixed. The small lmax=1 dipole model keeps the example compact. Check convergence in lmax before using this calculation for quantitative design; derivative agreement only validates differentiation of the chosen model.

A positive gradient[i, 0] means that slightly increasing sphere i's radius raises the cross section, with every other parameter held fixed. The other three entries describe translations along x, y and z. For a proposed small change to all the parameters, (gradient * change).sum() predicts the change in cross section. The central-difference check tests exactly this prediction.

The field sensitivity example asks for changes across a whole map and uses forward mode. Choosing a differentiation mode compares the times to obtain the same result with either mode. See framework adapters for the differentiation interfaces and gradient checks for more validation methods.