TMMax: High-performance modeling of multilayer thin-film structures using transfer matrix method with JAX¶
Optical multilayer thin films are essential building blocks in modern photonic systems, enabling precise control over reflectance, transmittance, and phase response. Fast and reliable simulation of these structures is critical for the design and analysis of advanced coatings such as distributed Bragg reflectors, anti-reflection layers, and spectral filters.
Traditional implementations of the transfer matrix method remain widely used for simulating these structures. However, their scalar treatment of wavelength and angle of incidence often leads to redundant recalculations, resulting in inefficiencies for large-scale simulations. In addition, conventional approaches lack native support for automatic differentiation, limiting their utility in gradient-based inverse design. To address these challenges, we introduce TMMax, a Python library that fully vectorizes and accelerates the transfer matrix method using the high-performance machine learning framework JAX. Designed as a versatile and extensible tool for thin-film optics research, TMMax integrates a suite of advanced numerical techniques and leverages just-in-time (JIT) compilation for optimal performance.
Originally developed with CPU-based execution to ensure broad accessibility, TMMax seamlessly scales to GPU and TPU platforms through JAX’s unified execution model. This ensures that users can benefit from both flexibility and computational efficiency, regardless of their available hardware.
TMMax also provides a curated material database with over 30 commonly used thin-film materials. Most refractive index and extinction coefficient datasets are sourced from refractiveindex.info, which compiles values from peer-reviewed literature. The database is fully extensible, and contributions from the community are encouraged. Researchers can easily add new materials by submitting issues or pull requests, supporting collaborative growth of the resource.
Beyond forward simulation, TMMax natively supports automatic differentiation via JAX’s autograd functionality. This enables analytical gradient calculations of optical properties with respect to arbitrary system parameters. Such capabilities open the door to gradient-based inverse design, optimization workflows assisted by machine learning, and parameter estimation tasks in photonics, thereby establishing TMMax as a powerful enabler for next-generation thin-film engineering.
Benchmarking demonstrates that TMMax is capable of simulating thin-film stacks containing hundreds of layers within seconds. Compared to baseline NumPy implementations, TMMax achieves speedups of several hundred times, providing a substantial advantage in computational throughput. This scalability empowers researchers to efficiently design and optimize large and complex multilayer structures, significantly accelerating the research and development cycle.
Contents
Citing TMMax¶
If you find the TMMax library useful in your work, we kindly ask that you consider citing it.
The current version is available as an arXiv preprint. A submission to the Journal of Open Source Software (JOSS) is currently under review, and we will update the reference once it is published.
Here is the recommended citation in BibTeX format:
@misc{danis2025tmmaxhighperformancemodelingmultilayer,
title={TMMax: High-performance modeling of multilayer thin-film structures using transfer matrix method with JAX},
author={Bahrem Serhat Danis and Esra Zayim},
year={2025},
eprint={2507.11341},
archivePrefix={arXiv},
primaryClass={physics.comp-ph},
url={https://arxiv.org/abs/2507.11341},
}
License¶
This project is licensed under the MIT License, which permits free use, modification, and distribution of the software, provided that the original copyright notice and license terms are included in all copies or substantial portions of the software. For a detailed explanation of the terms and conditions, please refer to the LICENSE file.
Acknowledgements¶
This work was supported by the Scientific and Technological Research Council of Türkiye (TUBITAK) under the 2209-A Research Project Support Programme for Undergraduate Students, 2022 First-Term Call.