Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

TMMax: High-performance modeling of multilayer thin-film structures using transfer matrix method with JAX

Bahrem Serhat Danis, Esra Zayim

Abstract

Optical multilayer thin-films are fundamental components that enable the precise control of reflectance, transmittance, and phase shift in the design of photonic systems. Rapid and accessible simulation of these structures holds critical importance for designing and analyzing complex coatings. While researchers commonly use the traditional transfer matrix method for designing these structures, its scalar approach to wavelength and angle of incidence causes redundant recalculations and inefficiencies in large-scale simulations. Furthermore, traditional method implementations do not support automatic differentiation, which limits their applicability in gradient-based inverse design approaches. Here, we present TMMax, a Python library that fully vectorizes and accelerates transfer matrix method using the high-performance machine learning library JAX. TMMax supports CPU, GPU, and TPU hardware, and includes a publicly available material database. Our approach, demonstrated through benchmarking, allows us to model thin-film stacks with hundreds of layers within seconds. This illustrates that our method achieves a simulation speedup of x100s over a baseline NumPy implementation, enabling optical engineers and thin-film researchers in optics and photonics to efficiently design complex dielectric multilayer structures through rapid and scalable simulations.

TMMax: High-performance modeling of multilayer thin-film structures using transfer matrix method with JAX

Abstract

Optical multilayer thin-films are fundamental components that enable the precise control of reflectance, transmittance, and phase shift in the design of photonic systems. Rapid and accessible simulation of these structures holds critical importance for designing and analyzing complex coatings. While researchers commonly use the traditional transfer matrix method for designing these structures, its scalar approach to wavelength and angle of incidence causes redundant recalculations and inefficiencies in large-scale simulations. Furthermore, traditional method implementations do not support automatic differentiation, which limits their applicability in gradient-based inverse design approaches. Here, we present TMMax, a Python library that fully vectorizes and accelerates transfer matrix method using the high-performance machine learning library JAX. TMMax supports CPU, GPU, and TPU hardware, and includes a publicly available material database. Our approach, demonstrated through benchmarking, allows us to model thin-film stacks with hundreds of layers within seconds. This illustrates that our method achieves a simulation speedup of x100s over a baseline NumPy implementation, enabling optical engineers and thin-film researchers in optics and photonics to efficiently design complex dielectric multilayer structures through rapid and scalable simulations.

Paper Structure

This paper contains 4 sections, 1 equation, 3 figures.

Table of Contents

  1. Statemement of need
  2. Benchmarks
  3. Installation
  4. Acknowledgements

Figures (3)

  • Figure 1: Schematic of two strategies for calculating transmission, reflection, and absorption in multilayer thin-film simulations. The system (a) is modeled either by sequentially multiplying 2×2 transfer matrices for each wavelength and incidence angle (b) or by vectorizing these operations across both axes (c).
  • Figure 2: Run time vs. layer count comparing tmm (orange) and TMMax (blue).
  • Figure 3: The colormaps show the runtime performance of tmm and TMMax across varying simulation grid sizes, comparing 8- and 80-layer stacks in (a) and (b), respectively.