{"id":19966068,"url":"https://github.com/bahremsd/tmmax","last_synced_at":"2025-10-15T13:35:36.979Z","repository":{"id":257036970,"uuid":"854440500","full_name":"bahremsd/tmmax","owner":"bahremsd","description":"A fast transfer matrix method written in jax for modelling optical multilayer thin films","archived":false,"fork":false,"pushed_at":"2025-10-07T14:16:03.000Z","size":32995,"stargazers_count":21,"open_issues_count":4,"forks_count":4,"subscribers_count":1,"default_branch":"master","last_synced_at":"2025-10-07T15:37:28.001Z","etag":null,"topics":["coating","jax","jax-based-optical-tools","jit","optical-coating","optical-coatings","optical-simulation","optics","photonics","thin-film","thin-film-engineering","tmm","transfer-matrix-method","vectorization","xla"],"latest_commit_sha":null,"homepage":"https://tmmax.readthedocs.io/en/latest/","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/bahremsd.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2024-09-09T07:14:00.000Z","updated_at":"2025-10-07T14:16:07.000Z","dependencies_parsed_at":null,"dependency_job_id":"31a7b64b-2b67-46d8-a608-6ef855ea3cb1","html_url":"https://github.com/bahremsd/tmmax","commit_stats":null,"previous_names":["bahremsd/tmmax"],"tags_count":3,"template":false,"template_full_name":null,"purl":"pkg:github/bahremsd/tmmax","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/bahremsd%2Ftmmax","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/bahremsd%2Ftmmax/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/bahremsd%2Ftmmax/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/bahremsd%2Ftmmax/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/bahremsd","download_url":"https://codeload.github.com/bahremsd/tmmax/tar.gz/refs/heads/master","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/bahremsd%2Ftmmax/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":279083617,"owners_count":26099674,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-10-15T02:00:07.814Z","response_time":56,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["coating","jax","jax-based-optical-tools","jit","optical-coating","optical-coatings","optical-simulation","optics","photonics","thin-film","thin-film-engineering","tmm","transfer-matrix-method","vectorization","xla"],"created_at":"2024-11-13T02:33:07.135Z","updated_at":"2025-10-15T13:35:36.973Z","avatar_url":"https://github.com/bahremsd.png","language":"Jupyter Notebook","funding_links":[],"categories":["Recently Updated","simulation","Libraries"],"sub_categories":["[Feb 16, 2025](/content/2025/02/16/README.md)","New Libraries"],"readme":"# **TMMax: High-performance modeling of multilayer thin-film structures using transfer matrix method with JAX**\n\n[![PyPI Version](https://img.shields.io/pypi/v/tmmax.svg)](https://pypi.org/project/tmmax/)\n[![License](https://img.shields.io/github/license/bahremsd/tmmax.svg)](https://github.com/serhatdann/tmmax/blob/main/LICENSE)\n[![Documentation Status](https://readthedocs.org/projects/tmmax/badge/?version=latest)](https://tmmax.readthedocs.io/en/latest/index.html)\n[![DOI](https://img.shields.io/badge/DOI-10.48550%2FarXiv.2507.11341-b31b1b.svg)](https://doi.org/10.48550/arXiv.2507.11341)\n[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.xxxxxxx.svg)](https://doi.org/10.48550/arXiv.2507.11341)\n\n\u003cdiv align=\"center\"\u003e\n  \u003ca href=\"https://pypi.org/project/tmmax/\"\u003e\n    \u003cimg src=\"https://github.com/bahremsd/tmmax/blob/master/docs/source/_static/logo_tmmax.png\" alt=\"tmmax\"\u003e\n  \u003c/a\u003e\n\u003c/div\u003e\n\n\u003c!-- TABLE OF CONTENTS --\u003e\n\u003cdetails\u003e\n  \u003csummary\u003eTable of Contents\u003c/summary\u003e\n  \u003col\u003e\n    \u003cli\u003e\u003ca href=\"#introduction\"\u003eIntroduction\u003c/a\u003e\u003c/li\u003e\n    \u003cli\u003e\u003ca href=\"#documentation\"\u003eDocumentation\u003c/a\u003e\u003c/li\u003e\n    \u003cli\u003e\u003ca href=\"#usage\"\u003eUsage\u003c/a\u003e\u003c/li\u003e\n    \u003cli\u003e\u003ca href=\"#database\"\u003eDatabase\u003c/a\u003e\u003c/li\u003e\n    \u003cli\u003e\u003ca href=\"#benchmarks\"\u003eBenchmarks\u003c/a\u003e\u003c/li\u003e\n    \u003cli\u003e\u003ca href=\"#installation\"\u003eInstallation\u003c/a\u003e\u003c/li\u003e\n    \u003cli\u003e\u003ca href=\"#license\"\u003eLicense\u003c/a\u003e\u003c/li\u003e\n    \u003cli\u003e\u003ca href=\"#credits\"\u003eCredits\u003c/a\u003e\u003c/li\u003e\n    \u003cli\u003e\u003ca href=\"#contact-and-support\"\u003eContact and Support\u003c/a\u003e\u003c/li\u003e\n  \u003c/ol\u003e\n\u003c/details\u003e\n\n## Introduction\n\n`tmmax` is a high-performance computational library designed for efficient calculation of optical properties in multilayer thin-film structures. Engineered to serve as a Swiss Army knife tool for thin-film optics research, `tmmax` integrates a comprehensive suite of advanced numerical methods. At its core, `tmmax` leverages JAX to enable just-in-time (JIT) compilation, vectorization and XLA (Accelerated Linear Algebra) operations, dramatically accelerating the evaluation of optical responses in multilayer coatings. By exploiting these capabilities, `tmmax` achieves exceptional computational speed, making it the optimal choice for modeling and analyzing complex systems.\n\nOriginally architected for CPU-based execution to ensure accessibility and scalability across diverse hardware configurations, `tmmax` seamlessly extends its computational efficiency to GPU and TPU platforms, thanks to JAX’s unified execution model. This adaptability ensures that high-performance simulations can be executed efficiently across a range of computing environments without modifying the core implementation. Moreover, `tmmax` natively supports automatic differentiation through JAX’s powerful autograd framework, allowing users to compute analytical gradients of optical properties with respect to arbitrary system parameters. This capability makes it particularly well-suited for gradient-based inverse design, machine learning-assisted optimization, and parameter estimation in photonic applications, providing a direct pathway to next-generation thin-film engineering.\n\n\n## Documentation\n\nThe complete documentation for TMMax is available in the [TMMax Read The Docs](https://tmmax.readthedocs.io/en/latest/). This repository provides an extensive set of examples demonstrating the key functionalities of tmmax, enabling users to efficiently analyze and manipulate multilayer thin-film stacks.\n\nThe examples cover fundamental and advanced use cases, including:\n\n- Material Database Management: Retrieving wavelength-dependent refractive index (n) and extinction coefficient (k) data from the built-in material database. Users can seamlessly add new materials, modify existing entries, or remove materials while maintaining database integrity.\n\n- Thin-Film Optical Properties: Computing reflection (R), transmission (T), and absorption (A) spectra for both coherent and incoherent multilayer thin-film structures.\n\n- Filter Design and Optimization: Rapid simulation of optical filters, showcasing how tmmax efficiently models various thin-film coatings, such as anti-reflective coatings, high-reflectivity mirrors, and bandpass filters.\n\n## Usage\n\nTo compute the reflection and transmission spectra of a multilayer thin-film stack using the tmmax framework, consider the following example. Suppose we have a coherent multilayer structure consisting of [Air, Y₂O₃, TiO₂, Y₂O₃, TiO₂, Y₂O₃, TiO₂, SiO₂], where the incident wavelength varies from 500 nm to 700 nm, and the angle of incidence spans from 0° to 70°. The calculation is performed as follows:\n\n```python\nimport jax.numpy as jnp\nfrom tmmax.tmm import tmm\n\n# Define your multilayer stack and simulation parameters\n\nmaterial_list = [\"Air\", \"Y2O3\", \"TiO2\", \"Y2O3\", \"TiO2\", \"Y2O3\", \"TiO2\", \"SiO2\"]\nthickness_list = jnp.array([630e-9, 200e-9, 630e-9, 200e-9, 630e-9, 200e-9])  \nwavelength_arr  = jnp.linspace(500e-9, 700e-9, 1000)\nangle_of_incidences  = jnp.linspace(0, (70*jnp.pi/180), 1000)\npolarization = 's'\n\nR_s, T_s = tmm(material_list = material_list,\n               thickness_list = thickness_list,\n               wavelength_arr = wavelength_arr,\n               angle_of_incidences = angle_of_incidences,\n               polarization = polarization)\n\npolarization = 'p'\n\nR_p, T_p = tmm(material_list = material_list,\n               thickness_list = thickness_list,\n               wavelength_arr = wavelength_arr,\n               angle_of_incidences = angle_of_incidences,\n               polarization = polarization)\n```\n\n\u003cdiv align=\"center\"\u003e\n  \u003cimg src=\"https://github.com/bahremsd/tmmax/blob/master/docs/source/_static/quickstart_example.png\" alt=\"usage_example\"\u003e\n\u003c/div\u003e\n\nFor cases where an incoherent layer is introduced within the stack, the simulation should include averaging effects of \"thick\" layers. In tmmax, incoherent layers are denoted by `1`, while coherent layers remain as `0`. The following example demonstrates the configuration of the same stack with incoherent layers:\n\n```python\nimport jax.numpy as jnp\nfrom tmmax.tmm import tmm\n\n# Define your multilayer stack and simulation parameters\n\nmaterial_list = [\"Air\", \"Y2O3\", \"TiO2\", \"Y2O3\", \"TiO2\", \"Y2O3\", \"TiO2\", \"SiO2\"]\nthickness_list = jnp.array([2000e-9, 100e-9, 2000e-9, 100e-9, 2000e-9, 100e-9])\ncoherency_list = jnp.array([1, 0, 1, 0, 1, 0])\nwavelength_arr  = jnp.linspace(500e-9, 700e-9, 1000)\nangle_of_incidences  = jnp.linspace(0, (70*jnp.pi/180), 1000)\npolarization = 's'\n\nR_s, T_s = tmm(material_list = material_list,\n               thickness_list = thickness_list,\n               wavelength_arr = wavelength_arr,\n               angle_of_incidences = angle_of_incidences,\n               coherency_list = coherency_list,\n               polarization = polarization)\n\npolarization = 'p'\n\nR_p, T_p = tmm(material_list = material_list,\n               thickness_list = thickness_list,\n               wavelength_arr = wavelength_arr,\n               angle_of_incidences = angle_of_incidences,\n               coherency_list = coherency_list,\n               polarization = polarization)\n```\n\n## Database\n\nThe `tmmax` repository features a database of 29 materials commonly used in optical thin film research. These materials is selected based on their prevalence in the field, as highlighted in the 18th chapter of _Thin-Film Optical Filters, Fifth Edition by H. Angus Macleod_ [doi](https://doi.org/10.1201/b21960). However, it’s important to clarify that these materials are not directly sourced from the book; rather, the book provides a reference for materials typically utilized in thin-film optics studies. The material data, including refractive index (n) and extinction coefficient (k) values, is available in the `tmmax/nk_data` folder in both .csv and .npy formats. While the initial data was stored in .csv format, the repository switched to .npy to leverage JAX’s faster data loading capabilities, as reading .csv files was slower in comparison.\n\nMost of the refractive index and extinction coefficient data for the materials in the database is obtained from [refractiveindex.info](https://refractiveindex.info/), which itself aggregates data from various research articles. To ensure proper attribution, we will provide references to the original sources for each material in the `docs/database_info` folder. You can access these references by reviewing the `README` file in that directory. \n\nFor example, to visualize the n and k data for SiO2, a material with widely accepted optical properties, you can use the `visualize_material_properties(material_name = \"SiO2\")` function in the `tmmax/nk_data`. This allows for a straightforward representation of the material's refractive index and extinction coefficient.\n\n\n\u003cdiv align=\"center\"\u003e\n  \u003cimg src=\"https://github.com/bahremsd/tmmax/blob/master/docs/source/_static/SiO2_nk_plot.png\" alt=\"database_example_sio2\"\u003e\n\u003c/div\u003e\n\nThe database is designed to be extensible, and we plan to include additional materials in future versions of tmmax. Contributions are welcome, and if you have a material you would like to add, please feel free to open an issue or submit a pull request. \n\n## Benchmarks\n\nIn evaluating the performance of various transfer matrix method implementations, we conducted rigorous benchmarking to compare runtime efficiency as a function of layer count, wavelength array length, and angle of incidence array length. Given that optical multilayer coatings can contain a substantial number of layers based on design constraints and performance sensitivity, scalability is a critical factor. To ensure a fair comparison, all tests were executed under identical conditions, including hardware specifications and input parameters.\n\n### Layer Size vs Run Time\n\nOne of the primary factors influencing computational complexity in tmm simulations is the number of layers in the multilayer stack. We benchmarked [`tmm`](https://github.com/sbyrnes321/tmm) and `tmmax` to assess their performance under increasing layer counts. The results indicate that as the number of layers grows, tmmax demonstrates good scalability compared to other implementations. This makes tmmax particularly well-suited for simulating highly complex multilayer structures without significant degradation in performance.\n\n\u003cdiv align=\"center\"\u003e\n  \u003cimg src=\"https://github.com/bahremsd/tmmax/blob/master/docs/source/_static/layer_benchmark_figure.png\" alt=\"layer_size_exp\"\u003e\n\u003c/div\u003e\n\n### Wavelength and Angle of Incidence Array Lengths vs Run Time\n\n\u003cdiv align=\"center\"\u003e\n  \u003cimg src=\"https://github.com/bahremsd/tmmax/blob/master/docs/source/_static/array_benchmark_figure.png\" alt=\"vmap_array_length_exp\"\u003e\n\u003c/div\u003e\n\n\n## Minimum Requirements and Installation\n\n- **Python**: ≥ 3.10  \n- **Operating System**: Linux, macOS, or Windows  \n- **Dependencies**:\n  - `jax\u003e=0.5.0`\n  - `jaxlib\u003e=0.5.0`\n  - `numpy\u003e=2.2.2`\n  - `scipy\u003e=1.15.1`\n  - `pandas\u003e=2.2.3`\n  - `matplotlib\u003e=3.10.0`\n\nThese dependencies are automatically installed when you install TMMax via pip.\n\n```bash\npip install tmmax\n```\n\n## Install from Source\n\n```bash\ngit clone https://github.com/bahremsd/tmmax.git\ncd tmmax\npip install .\n```\n\nYou may also install optional dependencies for documentation, testing, or development:\n\n```bash\npip install tmmax[test]\npip install tmmax[docs]\npip install tmmax[dev]\n```\n\n## Running TMMax on CPU, GPU, and TPU\n\nTMMax can run on multiple hardware backends via JAX — CPU by default, GPU or TPU if available.\n\n### 1. Running TMMax on CPU (Default)\n\nBy default, JAX installs CPU-only binaries.\n\n**Installation:**\n\n```bash\npip install -U jax jaxlib\n```\n\n**Verification:**\n\n```bash\npython -c \"import jax; print(jax.devices())\"\n```\n\n**Example output:**\n\n```\n[CpuDevice(id=0)]\n```\n\nIf you see `CpuDevice(id=0)`, TMMax is running on CPU.\n\n---\n\n### 2. Running TMMax on NVIDIA GPU\n\nIf your system has a CUDA-capable GPU, install JAX with GPU support.\n\n**Uninstall existing JAX (CPU-only):**\n\n```bash\npip uninstall jax jaxlib -y\n```\n\n**Install CUDA-enabled JAX (example for CUDA 12):**\n\n```bash\npip install -U \"jax[cuda12]\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n```\n\n**Verify GPU detection:**\n\n```bash\npython -c \"import jax; print(jax.devices())\"\n```\n\n**Example output:**\n\n```\n[GpuDevice(id=0, process_index=0)]\n```\n\nIf you see `GpuDevice`, TMMax will automatically use GPU acceleration.\n\n**Common Tips:**\n\n- Ensure CUDA Toolkit and cuDNN are installed correctly.  \n- No need to modify TMMax code — JAX automatically handles device dispatching.\n\n---\n\n### 3. Running TMMax on TPU (Google Cloud / Colab)\n\nTMMax also supports TPU computation through JAX on Google Cloud or Colab TPUs.\n\n**To use TPU on Colab:**\n\n1. Go to **Runtime → Change runtime type → TPU**  \n2. Run the setup cell:\n\n```python\nimport jax\nprint(jax.devices())\n```\n\n**Expected output:**\n\n```\n[TpuDevice(id=0), TpuDevice(id=1), ...]\n```\n\nJAX automatically configures TPU execution in TMMax — no code modification required.\n\nFor production or HPC TPU environments, follow the official JAX TPU installation guide: [JAX TPU Installation Guide](https://jax.readthedocs.io/en/latest/installation.html#tpu)\n\n---\n\n## Running Tests\n\nTMMax includes a robust **Pytest-based test suite** for validation and performance benchmarking.\n\n### Run All Tests\n\n```bash\npytest\n```\n\n### Run Tests with Coverage Report\n\n```bash\npytest --cov=tmmax --cov-report=term-missing\n```\n\n### Run Specific Test Files\n\n```bash\npytest tests/test_benchmark.py\npytest tests/test_integration.py\n```\n\n### Benchmark Validation\n\nTMMax includes performance benchmarks for JAX and vectorized TMM implementations.  \nTo verify benchmark correctness:\n\n```bash\npytest tests/test_benchmark.py -v\n```\n\nExpected output should include timing comparisons and correctness assertions for saved benchmark files.\n\n---\n\n## Development\n\nContributions are welcome! To set up a development environment:\n\n```bash\ngit clone https://github.com/bahremsd/tmmax.git\ncd tmmax\npip install -e .[dev]\n```\n\nBefore submitting a pull request, please:\n\n- Run code formatting:\n\n```bash\nblack tmmax\nruff check tmmax\n```\n\n- Run all tests:\n\n```bash\npytest\n```\n\n\n## Contributing\n\nWe warmly welcome contributions from the community to improve TMMax. This project thrives thanks to the collaborative effort of researchers, engineers, and enthusiasts in the field of optical multilayer thin-film simulations. Whether you are fixing a bug, suggesting a new feature, improving documentation, or enhancing performance, your input is highly valued and appreciated. For detailed instructions on how to contribute, please see the\n[Contribution Guidelines](https://tmmax.readthedocs.io/en/latest/contributing.html)\nin the documentation.\n\n\n## License\n\nThis project is licensed under the [MIT License](https://opensource.org/license/MIT), 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](https://github.com/bahremsd/tmmax/blob/master/LICENSE) file.\n\n## Credits\n\nThis repository sometimes utilizes approaches from Steven Byrnes' [`tmm`](https://github.com/sbyrnes321/tmm) library, and we thank him for its development. In the context of calculating the optical properties of incoherent multilayer thin films, we used the approach from a paper by _Charalambos C. Katsidis and Dimitrios I. Siapkas, \"General transfer-matrix method for optical multilayer systems with coherent, partially coherent, and incoherent interference\"_ (Applied Optics, 41(19):3978) 2002 [doi](https://doi.org/10.1364/AO.41.003978).\n\n\nIf you find the `tmmax` library beneficial in your work, we kindly ask that you consider citing us.\n\n```bibtex\n@misc{danis2025tmmax,\n      title={TMMax: High-performance modeling of multilayer thin-film structures using transfer matrix method with JAX}, \n      author={Bahrem Serhat Danis and Esra Zayim},\n      year={2025},\n      eprint={2507.11341},\n      archivePrefix={arXiv},\n      primaryClass={physics.comp-ph},\n      url={https://arxiv.org/abs/2507.11341}, \n}\n```\n\n## Contact and Support\n\nFor any questions, suggestions, or issues you encounter, feel free to [open an issue](https://github.com/bahremsd/tmmax/issues) on the GitHub repository. This not only ensures that your concern is shared with the community but also allows for collaborative problem-solving and creates a helpful reference for similar challenges in the future. If you would like to collaborate or contribute to the code, you can contact me via email.\n\nBahrem Serhat Danis - bdanis23@ku.edu.tr\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fbahremsd%2Ftmmax","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fbahremsd%2Ftmmax","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fbahremsd%2Ftmmax/lists"}