PyMieDiff: A differentiable Mie scattering library

TL;DR

PyMieDiff delivers a fully differentiable, GPU‑accelerated Mie solver for core–shell spheres implemented in PyTorch, enabling backpropagation through scattering observables with respect to geometry, dispersion, and experimental setup. The framework features vectorized batched evaluation, two AD pathways, and a material interface, and it integrates with TorchGDM for end‑to‑end multi‑particle simulations. Demonstrations include gradient‑based inverse design, a Mie‑informed tandem network, and end‑to‑end metasurface optimization, all leveraging analytic Mie theory with exact derivatives. This differentiable solver bridges analytic nanophotonics and modern ML pipelines, enabling rapid gradient‑based design, sensitivity analysis, and physics‑informed learning in nano‑optics.

Abstract

Light scattering by spherical-shaped particles of sizes comparable to the wavelength is foundational in many areas of science, from chemistry to atmospheric science, photonics and nanotechnology. With the new capabilities offered by machine learning, there is a great interest in end-to-end differentiable frameworks for scattering calculations. Here we introduce PyMieDiff, a fully differentiable, GPU-compatible implementation of Mie scattering for core-shell particles in PyTorch. The library provides native, autograd-compatible spherical Bessel and Hankel functions, vectorized evaluation of Mie coefficients, and APIs for computing efficiencies, angular scattering, and near-fields. All inputs - geometry, material dispersion, wavelengths, and observation angles and positions - are represented as tensors, enabling seamless integration with gradient-based optimisation or physics-informed neural networks. The toolkit can also be combined with "TorchGDM" for end-to-end differentiable multi-particle scattering simulations. PyMieDiff is available under an open source licence at https://github.com/UoS-Integrated-Nanophotonics-group/MieDiff.

Figures (6)

• Figure 1: Overview of PyMieDiff's features. The core of the code is a differentiable Mie solver, which is used to compute various differentiable observables (efficiencies, angular scattering, near-fields). All calculations are fully implemented in PyTorch, they all support automatic differentiation and run on GPU. The shown results are from a gold-silicon core-shell particle with core radius $r_c=100\,$nm and shell radius $r_s=200\,$nm, placed in vacuum and illuminated by a linearly polarized plane wave.
• Figure 2: Benchmarks. (a) timing for increasing numbers of simultaneously evaluated wavelengths. (b) timing for increasing numbers of simultaneously evaluated particles. (c) parallelization speedup on multi-core CPUs. (d) timing comparison to other Mie toolkits (on CPU). The batched example evaluates 256 different particles simultaneously, each at 256 wavelengths. Note that "treams" is not performance-optimized for Mie evaluation since it is not a dedicated Mie toolkit (but T-matrix), in consequence the time-bar for the treams result is not fully shown.
• Figure 3: Mie forward evaluation by the example of scattering from a gold-silicon core-shell particle with core radius $r_c=20\,$nm and shell radius $r_s=100\,$nm, placed in vacuum and illuminated by a linearly polarized plane wave. Tabulated material permittivities are taken from literature.johnsonOpticalConstantsNoble1972edwardsSiliconSi1997 (a) Extinction (blue), scattering (orange), and absorption (green) efficiency spectra. (b) Scattering radiation patterns in the scattering plane at selected wavelengths for perpendicularly (blue lines) and parallel (orange dashed lines) polarized light. (c) Near-fields (electric and magnetic in top and bottom row) inside and around the particle, evaluated at the resonance at $\lambda_0=575\,$nm. The shown maps are $400\times 400\,$nm$^2$. Colorbars indicate field intensity (leftmost panels) and amplitude (second from left to right panels), relative to the incident field absolute amplitude.
• Figure 4: Optimization example. (a) Sketch of the particle optimization procedure. A large number of random core-shell particles is randomly initialized and evaluated using PyMieDiff. Their optical properties (orange lines, here scattering efficiency spectra) are compared with the design target (dotted black lines, here a Gaussian response). The derivatives of the loss function (here MSE) with respect to the geometry parameters are obtained via PyTorch autograd and used to iteratively update the geometries. (b) Sketches of the best particles of selected iteration during an optimization (top row) together with their spectral response (bottom row). Relative sizes are on the same scale, the shades of the colors indicate the real part of the refractive index (darker means higher). (c) Loss function during the optimization.
• Figure 5: "Mie-informed" tandem network. (a) Sketch of the training configuration. The design observable (e.g., scattering efficiency, or any other Mie-calculated value) is fed into the "inverse network" (here a simple MLP). The NN result is interpreted as core-shell geometry parameters and fed into PyMieDiff, which calculates analytically the response of the particle. The result of the Mie calculation is compared to the original input data using some loss function (here MSE). The inverse network is trained using backpropagation of the loss through the differentiable Mie solver and then the network itself. The inset shows the training loss for an inverse network trained on scattering efficiency spectra of dielectric particles. (b) Examples of inverse design responses (orange lines) vs. randomly sampled design targets (dashed black lines) from the training dataset .