JaxLayerLumos: A JAX-based Differentiable Optical and Radio Frequency Simulator for Multilayer Structures

Overview

JaxLayerLumos is open-source transfer-matrix method (TMM) software designed for scientists, engineers, and researchers in optics and photonics. It provides a powerful yet intuitive interface for calculating the reflection and transmission (RT) of light through multi-layer optical structures. By inputting the refractive index, thickness of each layer, and the frequency vector, users can analyze how light interacts with layered materials, including the option to adjust for incidence angles. Our mission is to offer a lightweight, flexible, and fast alternative to commercial software, enabling users to perform complex optical simulations with ease. JaxLayerLumos is built with performance and usability in mind, facilitating the exploration of optical phenomena in research and development settings.

Features

• Lightweight and Efficient: Optimized for performance, JaxLayerLumos ensures rapid calculations without the overhead of large-scale commercial software.
• Gradient Calculation: Calculates the gradients over any variables involved in RT, powered by JAX.
• Flexibility: Accommodates a wide range of materials and structures by allowing users to specify complex refractive indices, layer thicknesses, and frequency vectors.
• Angle of Incidence Support: Expands simulation capabilities to include angled light incidence, providing more detailed analysis for advanced optical designs.
• Open Source and Community-Driven: Encourages contributions and feedback from the community, ensuring continuous improvement and innovation.
• Comprehensive Material Database: Includes a growing database of materials with their optical properties, streamlining the simulation setup process.

Installation

JaxLayerLumos can be easily installed by the following command using the PyPI repository.

pip install jaxlayerlumos

Alternatively, JaxLayerLumos can be installed from source.

pip install .

In addition, we support three installation modes, dev, benchmarking, and examples, where dev is defined for installing the packages required for development and software testing, benchmarking is for installing the packages required for benchmarking against different TMM software programs, and examples is needed for running the examples included in the examples directory. One of these modes can be used by commanding pip install .[dev], pip install .[benchmarking], or pip install .[examples].

Examples

A collection of examples in the examples directory exhibits various use cases and capabilities of our software. We provide the following examples:

1. Reflection Spectra over Wavelengths Varying Incidence Angles
2. Color Conversion
3. Color Exploration with Thin-Film Structures
4. Gradient Computation
5. Effective Medium Approximation Across a Partially Graded Interface
6. Graded Layers with Selected Interfaces
7. Visualization of Light Sources
8. Plotting of Optical Constants
9. Thin-Film Structure Optimization with Bayesian Optimization
10. Thin-Film Structure Optimization with DoG Optimizer
11. Reflection Spectra over Frequencies for Radar Design
12. Analysis of Solar Cells
13. Transmission Spectra over Wavelengths Varying Thicknesses
14. Triple Junction Solar Cells

Comparison of TMM Packages

We compare Ansys Optics, TMM-Fast, and tmm to our software.

$^1$ Because TPUs are optimized for low-precision computation, their simulation results may show reduced numerical precision.

Benchmarking against Other Software

We benchmark JaxLayerLumos against other software. Detailed benchmarking results can be found in COMPARISONS.md. These comparisons include the results of Ansys Optics, TMM-Fast, and tmm.

To obtain these results, you should install additional required packages. Before installing the packages, you should install PyTorch first. In particular, if you need the CPU version of PyTorch, you can install it using the following command.

pip install torch --index-url https://download.pytorch.org/whl/cpu

For details, you can refer to the official PyTorch instruction. Then, the required packages can be installed by the following command.

pip install .[benchmarking]

Finally, you can run the benchmarking code compare_methods.py in the benchmarking directory.

Software Testing and Test Automation

We provide a variety of test files in the tests directory. Before running the test files, the required packages should be installed by using pip install .[dev]. They can be run by commanding pytest tests/. Moreover, these test files are automatically tested via GitHub Actions, of which the configuration is defined in .github/workflows/pytest.yml.

Supported Materials

Materials supported by our software are described in MATERIALS.md.

JaxLayerLumos includes a growing library of materials, which are specified using either complex refractive indices or complex permittivities and permeabilities, which can be sourced from the literature or specified by users based on experimental data. When only complex refractive indices are provided, magnetic effects are assumed to be negligible, and the relative permeability is set to unity ($\mu_{r,j} = 1$), an assumption typically valid at optical frequencies. In the RF and microwave regimes, the electromagnetic properties of metals are derived from their electrical conductivity and magnetic susceptibility, while dielectrics are generally modeled with constant permittivity and negligible loss.

Contributing Guidelines

To contribute, please read CONTRIBUTING.md for our guidelines on issues, enhancements, and pull requests. Follow the outlined standards to keep the project consistent and collaborative.

Declaration of Large Language Model Usage

This project was initiated before large language models were incorporated into its development workflow. They are now actively used to support code revision, feature development, software maintenance, and text editing in this repository. In particular, we use Codex as an AI coding tool.

Citation

@article{LiM2025joss,
    author={Li, Mingxuan and Kim, Jungtaek and Leu, Paul W.},
    title={{JaxLayerLumos}: A {JAX}-based Differentiable Optical and Radio Frequency Simulator for Multilayer Structures},
    journal={Journal of Open Source Software},
    volume={10},
    number={114},
    pages={8572},
    year={2025}
}

Contributors

The Journal of Open Source Software paper author list is fixed at publication. This section recognizes software contributions to the project.

For the full commit history, see the GitHub contributors graph.

License

JaxLayerLumos is released under the MIT License, promoting open and unrestricted access to software for academic and commercial use.

Acknowledgments

• We sincerely thank all contributors and users for your support and feedback.
• This work is supported by the Center for Materials Data Science for Reliability and Degradation (MDS-Rely), which is the Industry-University Cooperative Research Center (IUCRC) of National Science Foundation.
• The University of Pittsburgh, Case Western Reserve University, and Carnegie Mellon University are participating institutions in MDS-Rely.