TorchGDM: A GPU-Accelerated Python Toolkit for Multi-Scale Electromagnetic Scattering with Automatic Differentiation

TL;DR

TorchGDM presents a GPU-accelerated, PyTorch-based framework for nano-optical scattering that unifies volume discretization with effective dipole models through the Green's Dyadic Method. By enabling mixed discretization (fully discretized and GPM-based structures) and full automatic differentiation, it supports differentiable design, optimization, and inverse problems for multi-scale metasurfaces and complex environments. Key contributions include the formalism for 6N×6N coupling, a robust extraction procedure for Global Polarizability Matrix models, and a versatile set of observables (near/far fields, LDOS, Green's tensors, multipole analyses) with AD capabilities. The toolkit targets mesoscale to large-scale scattering problems, offering a path toward physics-informed learning and gradient-based metasurface design, with demonstrated capabilities in resonance searches, BICs, metalens optimization, and large mixed-structure reconstructions. The work also delineates current AD limitations and outlines a roadmap for differentiable Mie/T-Matrix components and iterative solvers to extend scalability beyond ~10^4 dipoles.

Abstract

We present "torchGDM", a numerical framework for nano-optical simulations based on the Green's Dyadic Method (GDM). This toolkit combines a hybrid approach, allowing for both fully discretized nano-structures and structures approximated by sets of effective electric and magnetic dipoles. It supports simulations in three dimensions and for infinitely long, two-dimensional structures. This capability is particularly suited for multi-scale modeling, enabling accurate near-field calculations within or around a discretized structure embedded in a complex environment of scatterers represented by effective models. Importantly, torchGDM is entirely implemented in PyTorch, a well-optimized and GPU-enabled automatic differentiation framework. This allows for the efficient calculation of exact derivatives of any simulated observable with respect to various inputs, including positions, wavelengths or permittivity, but also intermediate parameters like Green's tensor components, which can be interesting for physics informed deep learning applications. We anticipate that this toolkit will be valuable for applications merging nano-photonics and machine learning, as well as for solving nano-photonic optimization and inverse problems, such as the global design and characterization of metasurfaces, where optical interactions between structures are critical.

Figures (16)

• Figure 1: (a) Reference system and angular conventions in torchGDM. (b) Sketch of the global polarizability matrix (GPM) effective model formalism.bertrandGlobalPolarizabilityMatrix2020 A GPM describes a particle through a set of local (dark blue) and non-local polarizabilities.This means, the incident field at one specific polarizability location induces electric and magnetic dipole moments at all polarizability locations in the GPM. (c-d) available structure types for (c) 3D and (d) 2D simulations. In addition to discretized structures, effective dipole models are supported that may contain several non-local dipoles (global polarizability matrix, "GPM" (b)), or a single pair of an electric and a magnetic dipole ("dp pair"). Effective models can be extracted from discretized structures, T-Matrices or from Mie theory. Note that the Mie and T-Matrix extractions are not AD compatible.
• Figure 2: Extraction of an effective polarizability model from a discretized structure. (a) define a set of illuminations. (b) simulate the full fields for each illumination. (c) for each simulation, extract the effective electric and magnetic dipole moments of the response. (d) optimize the effective polarizability to represent as closely as possible the different dipole moments.
• Figure 3: (a) Example 3D geometry. A discretized structure $V$ with $N$ cubic volume elements $v_i$ (orange) is coupled to two other structures (green and blue), represented each by a pair of electric-electric and magnetic-magnetic effective polarizabilities $\boldsymbol{\alpha}_{\text{eff}1}$ and $\boldsymbol{\alpha}_{\text{eff}2}$. (b) 2D projection of the geometry. (c) Mixed polarizability linear coupled system, as solved in torchGDM (by inversion of the coupling matrix).
• Figure 4: TorchGDM main package structure. The entire toolkit is built around the Simulation class, which is a container for the structure, environment and illumination descriptions and manages the simulation. It also provides access to the most relevant post processing and visualization functions. Blue: Subpackages and modules. Black/gray: Classes.
• Figure 5: Gallery of the main features of torchGDM. a) Structures can be arbitrarily composed of discretized geometries (here a dielectric prism) and effective models (here small gold nano-rods and a dielectric disc). b) Electric and magnetic near fields. c) Field gradients, either calculated via automatic differentiation or by finite differences (faster). d) Partial, electric and magnetic Local density of states as well as Green's tensors. e) Rasterscan simulations with scanning illuminations. Here, a focused Gaussian (through NA 1.4, linear $X$-polarization) is raster-scanned across the structure and the scattered far-field intensity is plotted for each beam position. f) Internal fields can be obtained for all discretized structures. g) Radiation patterns (here: cut through XY plane). h) Cross section spectra. i) Exact multipole decomposition (discretized structures only). Here, the multipole decomposition (up to quadrupoles) of the contribution of the discretized Si$_3$N$_4$ triangular structure to the scattering cross section is plotted. j) Near-field enhancement spectra (here $100$ nm below the center of the triangle). All 2D plots show an area of $800\times 800$ nm$^2$. Except for the LDOS and rasterscan examples, the illumination is a plane wave in the structure plane, coming from the positive $Y$ axis, with linear polarization along $X$. Single wavelength results are calculated at $\lambda_0=550\,$nm. The host medium is vacuum.