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

[![DOI](https://joss.theoj.org/papers/10.21105/joss.08572/status.svg)](https://doi.org/10.21105/joss.08572)
[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.12602789.svg)](https://doi.org/10.5281/zenodo.12602789)
[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/jaxlayerlumos)](https://pypi.org/project/jaxlayerlumos/)
![GitHub Release](https://img.shields.io/github/v/release/JaxLayerLumos/jaxlayerlumos)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
[![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
[![Documentation Status](https://readthedocs.org/projects/jaxlayerlumos/badge/?version=main)](https://jaxlayerlumos.readthedocs.io)

<p align="center">
<img src="https://raw.githubusercontent.com/JaxLayerLumos/JaxLayerLumos/main/assets/jaxlayerlumos.jpg" width="400" />
</p>

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

<p align="center">
<img src="https://raw.githubusercontent.com/JaxLayerLumos/JaxLayerLumos/main/assets/TMM.png" width="700" />
</p>

## 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](https://pypi.org/project/jaxlayerlumos/).

```bash
pip install jaxlayerlumos
```

Alternatively, JaxLayerLumos can be installed from source.

```bash
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

<p align="center">
<img src="https://raw.githubusercontent.com/JaxLayerLumos/JaxLayerLumos/main/assets/applications.png" width="800" />
</p>

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](examples/angle-variation.ipynb)
2. [Color Conversion](examples/color-conversion.ipynb)
3. [Color Exploration with Thin-Film Structures](examples/color-exploration.ipynb)
4. [Gradient Computation](examples/gradient-computation.ipynb)
5. [Visualization of Light Sources](examples/light-source-visualization.ipynb)
6. [Plotting of Optical Constants](examples/n-k-extrapolation.ipynb)
7. [Thin-Film Structure Optimization with Bayesian Optimization](examples/optimization-bayeso.ipynb)
8. [Thin-Film Structure Optimization with DoG Optimizer](examples/optimization-dog.ipynb)
9. [Reflection Spectra over Frequencies for Radar Design](examples/radar-design.ipynb)
10. [Analysis of Solar Cells](examples/solar-cell-analysis.ipynb)
11. [Transmission Spectra over Wavelengths Varying Thicknesses](examples/thickness-variation.ipynb)
12. [Triple Junction Solar Cells](examples/triple-junction-solar-cells.ipynb)

## Comparison of TMM Packages

We compare [Ansys Optics](https://www.ansys.com/products/optics), [TMM-Fast](https://github.com/MLResearchAtOSRAM/tmm_fast), and [tmm](https://github.com/sbyrnes321/tmm) to our software.

| Feature | Ansys Optics (stackrt) | TMM-Fast (PyTorch/NumPy) | tmm (Pure Python) | JaxLayerLumos (JAX) |
|-----|-----|-----|-----|-----|
| **Lightweight** | ❌ Bulky | ✅ | ✅ | ✅ |
| **Speed** | ❌ Slow | ✅ Fast | ❌ Slow | 🟨 Moderate |
| **Gradient Support** | ❌ | ✅ | ❌ | ✅ |
| **GPU Support** | ❌ | ✅ | ❌ | ✅ |
| **TPU Support**$^1$ | ❌ | ❌ | ❌ | ✅ |
| **Position-Dependent Absorption** | ❌ | ❌ | ✅ | ✅ |
| **Optical Simulations** | ✅ | ✅ | ✅ | ✅ |
| **Infrared Simulations** | 🟨 Limited | 🟨 Limited | ❌ | ✅ |
| **Radio Wave Simulations** | 🟨 Limited | ❌ | ❌ | ✅ Handles magnetic materials |
| **Open Source** | ❌ Commercial | ✅ MIT | ✅ BSD-3-Clause | ✅ MIT |

$^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](markdowns/COMPARISONS.md).
These comparisons include the results of [Ansys Optics](https://www.ansys.com/products/optics), [TMM-Fast](https://github.com/MLResearchAtOSRAM/tmm_fast), and [tmm](https://github.com/sbyrnes321/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.

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

For details, you can refer to [the official instruction of PyTorch](https://pytorch.org).
Then, the required packages can be installed by the following command.

```bash
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](markdowns/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](markdowns/CONTRIBUTING.md) for our guidelines on issues, enhancements, and pull requests. Follow the outlined standards to keep the project consistent and collaborative.

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

## License

JaxLayerLumos is released under the [MIT License](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)](https://mds-rely.org), which is the [Industry-University Cooperative Research Center (IUCRC)](https://iucrc.nsf.gov) of [National Science Foundation](https://www.nsf.gov).
- The University of Pittsburgh, Case Western Reserve University, and Carnegie Mellon University are participating institutions in MDS-Rely.
