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

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 differnt 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:

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. For example, Ag, Air, Al2O3, Al, aSi, Au, BK7, Cr, cSi, Cu, Fe, FusedSilica, GaAs, GaInP, GaP, Ge, InP, ITO, Mg, Mn, Ni, Pb, Pd, Pt, Sapphire, Si3N4, SiO2, TiN, TiO2, Ti, W, ZnO, Zn are included in JaxLayerLumos.

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, 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.

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.

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}
      }