A PyTorch-based differentiable thin film solver for multi-layer optical coatings. Supports both isotropic and anisotropic materials with full autograd for inverse design.
| Feature | NumPy TMM (sbyrnes321/tmm) | DiffTMM |
|---|---|---|
| Differentiable | No | Yes (PyTorch autograd) |
| GPU acceleration | No (CPU only) | Yes (CUDA) |
| Batch processing | No (sequential) | Yes (vectorized) |
| Anisotropic materials | No (isotropic only) | Yes (4x4 transfer matrix) |
| Speed (batch=16) | 1x baseline | ~190x (isotropic 2x2), ~134x (anisotropic 4x4) |
git clone https://github.com/singer-yang/DiffTMM.git
cd DiffTMM
pip install torch numpy matplotlib scipyInitialize a film stack with known refractive indices and thicknesses, then compute Fresnel coefficients at arbitrary wavelengths and angles.
import torch
from difftmm import IsotropicFilmSolver
# Define film stack: Glass | Ta2O5 | SiO2 | Ta2O5 | Glass
solver = IsotropicFilmSolver(
mat_in=1.5, # incident medium
mat_out=1.5, # exit medium
mat_ls=[2.10, 1.46, 2.10], # interior layer indices
thickness_ls=[0.080, 0.120, 0.080], # thicknesses in um
device=torch.device("cuda"),
)
# Compute Fresnel coefficients: ts, tp, rs, rp
angles = torch.linspace(0, 1.2, 100, device=solver.device)
ts, tp, rs, rp = solver.simulate(theta=angles, wvln=[0.45, 0.55, 0.65])
# Output shape: (n_mirrors, n_wvlns, n_angles)Given target Fresnel coefficients, recover unknown film thicknesses using gradient-based optimization.
import torch
from difftmm import create_jones_matrix_isotropic
# Film stack with unknown thicknesses
n_list = torch.tensor([2.10, 1.46, 2.10, 1.46, 2.10], device="cuda")
d_param = torch.nn.Parameter(torch.randn(5, device="cuda") * 0.5)
def param_to_thickness(p):
return torch.sigmoid(p) * 0.19 + 0.01 # map to [0.01, 0.20] um
# Optimization loop
optimizer = torch.optim.Adam([d_param], lr=0.02)
for step in range(3000):
optimizer.zero_grad()
d = param_to_thickness(d_param)
pred = forward_tmm(n_list, d, n_in=1.0, n_out=1.52, inp=inp)
loss = ((pred - target).real ** 2 + (pred - target).imag ** 2).mean()
loss.backward()
optimizer.step()Result: Layer thicknesses recovered from random initialization:
Layer GT (nm) Recovered (nm) Error (nm)
1 60.00 60.00 0.00
2 130.00 130.00 0.00
3 85.00 85.00 0.00
4 110.00 110.00 0.00
5 70.00 70.00 0.00
difftmm.IsotropicFilmSolver— Fast 2x2 transfer matrix method for isotropic materials (~190x faster than NumPy TMM)difftmm.FilmSolver(alsoAnisotropicFilmSolver) — General 4x4 transfer matrix method for both isotropic and anisotropic materials (~134x faster than NumPy TMM)
Both solvers share the same API:
solver = Solver(
mat_in=1.0, # incident medium refractive index
mat_out=1.52, # exit medium refractive index
mat_ls=[2.1, 1.46], # interior layer refractive indices
thickness_ls=[0.08, 0.12], # thicknesses in um (optional, random if None)
device=torch.device("cuda"),
)
ts, tp, rs, rp = solver.simulate(theta, wvln)DiffTMM ships with wavelength-dependent refractive index support via the
Material class. Look up materials by name (case-insensitive):
from difftmm import IsotropicFilmSolver, Material, list_materials
# Bundled catalogs: CDGM/SCHOTT/MISC AGF glasses + thin-film n+k tables
print(len(list_materials()), "materials available")
# Pass material names directly to a solver — they're auto-wrapped in Material()
solver = IsotropicFilmSolver(
mat_in="air",
mat_out="N-BK7", # Sellmeier (AGF)
mat_ls=["TiO2", "SiO2"], # n+k tables for thin-film materials
thickness_ls=[0.06, 0.10],
)
ts, tp, rs, rp = solver.simulate(theta=angles, wvln=[0.45, 0.55, 0.65])Scalars (float/complex) and strings can be mixed freely in mat_ls.
For the 4×4 FilmSolver, anisotropic per-axis dispersion is expressed as
a (mat_x, mat_y, mat_z) tuple per layer.
Dispersion models supported in v1: Sellmeier (analytical, real n) and linear interpolation (lookup tables, complex n + ik).
For stacks containing layers thicker than the source's coherence length
(typically anything thicker than ~10 μm for broadband illumination), the
fully-coherent calculation produces dense Fabry-Perot ripples that do not
appear in real measurements. The IncoherentIsotropicFilmSolver lets you
mark individual layers as incoherent ('i') while keeping thin films
coherent ('c'):
import torch
from difftmm import IncoherentIsotropicFilmSolver
# Stack: air | 100 nm TiO2 | 1 mm glass | air
solver = IncoherentIsotropicFilmSolver(
mat_in=1.0,
mat_out=1.0,
mat_ls=[2.40, 1.52],
c_list=["c", "i"], # TiO2 coherent, glass incoherent
thickness_ls=[0.100, 1000.0],
device=torch.device("cuda"),
)
Rs, Rp, Ts, Tp = solver.simulate(
theta=torch.tensor([0.0]),
wvln=[0.55],
)
# Returns real power coefficients in [0, 1].c_list is per-interior-layer; the two semi-infinite media are always
treated as incoherent. The coherent path (IsotropicFilmSolver) and the
incoherent path (IncoherentIsotropicFilmSolver) share the same forward
mathematics for coherent stacks, so an all-'c' c_list (with
incoherent semi-infinite endpoints) produces results consistent with the
coherent solver up to the loss of complex phase.
Only the 2x2 isotropic solver supports incoherent layers today. Anisotropic incoherent TMM is tracked as future work.
Validated against the reference NumPy TMM library (sbyrnes321/tmm) on surface plasmon resonance (SPR) calculations:
The anisotropic 4x4 solver is validated for energy conservation, isotropic limit accuracy, cross-polarization coupling, and reciprocity:
| Layers | TMM NumPy (s) | Anisotropic 4x4 (s) | Isotropic 2x2 (s) | Speedup (4x4) | Speedup (2x2) |
|---|---|---|---|---|---|
| 3 | 0.281 | 0.003 | 0.001 | 84.1x | 233.0x |
| 11 | 0.577 | 0.005 | 0.003 | 128.4x | 201.1x |
| 25 | 1.076 | 0.008 | 0.006 | 133.7x | 186.4x |
| 39 | 1.574 | 0.010 | 0.009 | 165.1x | 182.2x |
The isotropic 2x2 solver uses ~23x less GPU memory than the anisotropic 4x4 solver. NumPy TMM is CPU-only (0 GPU memory).
├── difftmm/ # Importable package
│ ├── __init__.py # Public API
│ ├── film_solver_isotropic.py # 2x2 isotropic solver (fast)
│ ├── film_solver_anisotropic.py # 4x4 anisotropic solver (general)
│ ├── film_solver_incoherent.py # 2x2 isotropic solver with incoherent layers
│ └── material/ # Wavelength-dependent materials
│ ├── __init__.py
│ ├── materials.py # Material class and catalog loaders
│ └── catalogs/ # Bundled glass + thin-film data
├── 1_forward_simu.ipynb # Example: forward simulation
├── 2_inverse_design.ipynb # Example: differentiable inverse design
├── 3_real_materials.ipynb # Example: real materials
├── 4_incoherent_film.ipynb # Example: incoherent / thick-substrate solver
├── tmm_numpy/ # Reference NumPy TMM library
├── benchmarks/ # Accuracy and performance benchmarks
├── tests/ # Pytest suite
├── pyproject.toml # Packaging metadata
├── CITATION.cff # Citation metadata
└── README.md
- 2x2 transfer matrix method: Standard formulation for isotropic multi-layer films
- 4x4 transfer matrix method: General formulation for anisotropic (birefringent) media
- Snell's law, Fresnel equations, evanescent wave handling beyond critical angle
- Bidirectional propagation (forward and reverse through the film stack)
- Complete polarization handling via Jones calculus
- S. J. Byrnes, "Multilayer optical calculations," arXiv:1603.02720
- Steven Byrnes' TMM library: github.com/sbyrnes321/tmm
- Yang, X., Liu, Z., Nie, Z., Fan, Q., Shi, Z., Bonar, J., & Heidrich, W. (2026). "End-to-end differentiable design of geometric waveguide displays." arXiv preprint arXiv:2601.04370
DiffTMM is licensed under the Apache License 2.0.
The bundled NumPy TMM reference library (tmm_numpy/) is by Steven Byrnes and is licensed under the MIT License.