Skip to content

Changes across a field map

How does the scattered field above three spheres change when all three grow slightly, or when the illumination wavelength changes? This example computes both changes at every point of a 5 × 5 grid. It keeps the real and imaginary parts of all three electric-field components, so changes in phase remain available as well as changes in amplitude.

The spheres have radius 80 nm and permittivity 3.2 + 0.05j. A circularly polarized plane wave of wavelength 650 nm illuminates them along +z, and the grid lies 350 nm above their centers. Only two inputs vary: the common radius and the vacuum wavenumber k0 = 2*pi/wavelength. The script converts the latter derivative into a change per nanometer of wavelength for the printed table.

jax.jacfwd(field_map) asks for every output derivative with respect to both inputs. This is forward mode: follow an input change through the calculation to every field value. With two inputs and many outputs, only two input directions are needed. The returned array has shape (5, 5, 3, 2): grid positions, electric-field component, and input parameter (k0, radius).

# Requires: jax
# /// script
# requires-python = ">=3.12"
# dependencies = ["numpy>=2.1,<3", "jax>=0.10,<0.12", "treams-rs"]
# ///
"""How does the scattered field change with wavelength and sphere radius?"""

import jax
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)

GRID_POINTS = 5


def make_problem(grid_points=GRID_POINTS):
    """Return the field calculation and starting [k0 (1/nm), radius (nm)]."""
    positions = np.array([[-160, 0, 0], [160, 0, 0], [0, 250, 0]], dtype=float)
    x, y = np.meshgrid(
        np.linspace(-320, 320, grid_points),
        np.linspace(-200, 400, grid_points),
        indexing="ij",
    )
    # Every observation point is outside the spheres, on a plane 350 nm above them.
    points = np.stack((x, y, np.full_like(x, 350)), axis=-1)
    initial = np.array([2 * np.pi / 650, 80.0])

    def field_map(parameters):
        k0, radius = parameters
        sphere = tr.sphere_tmatrix(k0=k0, lmax=2, radius=radius, material=3.2 + 0.05j)
        cluster = tr.Cluster([sphere, sphere, sphere], positions=positions)
        incident = tr.plane_wave([0, 0, 1], "positive_helicity", k0=k0)
        return cluster.scatter(incident).efield(points)

    return field_map, initial


def main():
    field_map, parameters = make_problem()
    field = np.asarray(field_map(parameters))
    # The last axis labels k0 and radius. Every grid point and electric-field
    # component gets both derivatives in this one call.
    derivatives = np.asarray(jax.jacfwd(field_map)(parameters))

    # Changing wavelength changes k0 = 2*pi/wavelength.
    wavelength, radius = 2 * np.pi / parameters[0], parameters[1]
    wavelength_derivative = derivatives[..., 0] * (-2 * np.pi / wavelength**2)
    radius_derivative = derivatives[..., 1]

    # Check both entire derivative maps, using steps in their own units.
    for parameter, step in enumerate((1e-7, 1e-3)):
        shift = np.eye(2)[parameter] * step
        difference = (
            np.asarray(field_map(parameters + shift))
            - np.asarray(field_map(parameters - shift))
        ) / (2 * step)
        np.testing.assert_allclose(
            derivatives[..., parameter], difference, rtol=2e-6, atol=1e-9
        )

    center = GRID_POINTS // 2
    print(f"Three spheres: radius {radius:g} nm, wavelength {wavelength:g} nm")
    print(f"Field map: {GRID_POINTS} x {GRID_POINTS} points, three complex components")
    print("At the middle grid point, scattered electric field and changes per nm")
    print(
        f"{'E':>3} {'field (real, imag)':>23} {'d/d wavelength':>23} {'d/d radius':>23}"
    )
    for component, label in enumerate("xyz"):
        values = (
            field[center, center, component],
            wavelength_derivative[center, center, component],
            radius_derivative[center, center, component],
        )
        print(
            f"{label:>3} "
            + " ".join(f"{value.real:>11.4g}{value.imag:>11.4g}j" for value in values)
        )
    print("Both complete derivative maps agree with central differences.")


if __name__ == "__main__":
    main()

Output

Three spheres: radius 80 nm, wavelength 650 nm
Field map: 5 x 5 points, three complex components
At the middle grid point, scattered electric field and changes per nm
  E      field (real, imag)          d/d wavelength              d/d radius
  x     0.05511    0.09236j   0.0003481 -0.0006258j    0.001661   0.003903j
  y    -0.09403    0.05506j   0.0006362  0.0003554j   -0.003981   0.001681j
  z    0.002894   0.002279j   1.165e-06 -2.637e-05j   0.0001102  9.908e-05j
Both complete derivative maps agree with central differences.

Using the result

derivatives[..., 1] tells you how every complex field value changes per nanometer added to all three sphere radii. For a small radius change dr, field + derivatives[..., 1] * dr is its first-order prediction. The example checks both entire derivative maps against central differences before printing a few entries.

The incident wave and the spheres receive the same changing k0, so the wavelength derivative includes both the excitation and the scattering. The permittivity stays fixed; this example does not include material dispersion. Lengths are in nanometers, the grid stays fixed, and the displayed field is the scattered field. Add incident.efield(points) inside field_map if you want the total field instead.

The small grid and lmax=2 keep the example quick. Increase the grid density and check convergence in lmax for a calculation you intend to use. The derivative checks establish agreement for this chosen discretization.

When many particle parameters control a single quantity, the cluster gradient example uses reverse mode. Choosing a differentiation mode compares both modes on these same physical questions, including a runnable timing comparison.