mlresearchatosram / tmm_fast

tmm_fast is a lightweight package to speed up optical planar multilayer thin-film device computation. Developed by Alexander Luce (@Nerrror) in cooperation with Heribert Wankerl (@HarryTheBird).

tmm_fast

tmm_fast or transfer-matrix-method_fast is a lightweight package to speed up optical planar multilayer thin-film device computation. Typically, one is interested in computing the Reflection and/or Transmission of light through a multilayer thin-film depending on the wavelength and angle of incidence of the incoming light. The package is build on the original tmm package from sjbyrnes but quite a lot faster. Depending on the number of layers, wavelength range and angular range speed-ups of ~100x are possible. The physics behind the transfer matrix method can be studied in any textbook on optical devices or in Multilayer optical calculations from Steven J. Byrnes. More detailed information about the package and its applications can be found in TMM-Fast: A Transfer Matrix Computation Package for Multilayer Thin-Film Optimization.

Figure 1: Overview about what you can do with this package.

Getting Started

To get a local copy up and running follow these steps:

Installation

tmm-fast package requires pytorch, therefore it is better to install it first following official instructions.

Stable version from pipy.org:

pip install tmm-fast

Latest version from github.com:

pip install git+https://github.com/MLResearchAtOSRAM/tmm_fast

Unified functionality of tmm_fast: Sponge PyTorch functionality for free

Parallelized computation of reflection and transmission for coherent light spectra that traverse a bunch of multilayer thin-films with dispersive materials. This package is essentially build on Pytorch and its related functionalities such as GPU accelerations and Autograd. It naturally allows for:

• GPU accelerated computations
• To compute gradients regarding the multilayer thin-film (i.e. N, T) thanks to Pytorch Autograd
• Vectorized computations, i.e. using Einstein summation convention

In general, tmm_fast is a lightweight package to speed up the computations of reflection and transmission of optical planar multilayer thin-films by vectorization regarding

• a set of multilayer thin-films consisting of various layers of particular material and thickness,
• a set of wavelengths, and
• a set of incident angles.

For old guards: Numpy is fully supported

All of the inputs can also be a numpy array format instead of torch tensors. However, all internal computations are processed via PyTorch and thus the output data is converted back to numpy arrays again. Hence, the use of numpy input may increase computation time due to data type conversions.

Benefits and conducted sanity checks, backgrounds

Depending on the number of thin films an their number of layers as well as the considered wavelengths that irradiate the thin film under particular angles of incident, the computation time can be decreased by 2-3 orders of magnitude. This claim is supported by several cross-checks (https://arxiv.org/abs/2111.13667), conducted with the code provided by Steven J. Byrnes (https://arxiv.org/abs/1603.02720). Of course, the checks covered both, computational time and physical outputs.

The physics behind the transfer matrix method can be studied in any textbook on optical devices or related papers, e.g.

• Chapter on thin films in Microphotonics (http://www.photonics.intec.ugent.be/download/ocs129.pdf) by Dries Van Thourhout, Roel Baets (UGent) and Heidi Ottevaere (VUB)
• Multilayer optical computations (https://arxiv.org/abs/1603.02720) by Steven J. Byrnes
• The Fresnel Coefficient of Thin Film Multilayer Using Transfer Matrix Method (https://iopscience.iop.org/article/10.1088/1757-899X/518/3/032026) by Zahraa Hummam Mohammed

gym-multilayerthinfilm

The proposed OpenAI/Farama-Foundation gymnasium environment utilizes the parallelized transfer-matrix method (TMM-Fast) to implement the optimization of multi-layer thin films as parameterized Markov decision processes. A very intuitive example is provided in example.py. Whereas the contained physical methods are well-studied and known since decades, the contribution of this code lies the transfer to an OpenAI/Farama-Foundation gymnasium environment. The intention is to enable AI researchers without optical expertise to solve the corresponding parameterized Markov decision processes. Due to their structure, the solution of such problems is still an active field of research in the AI community.
The publication Parameterized Reinforcement learning for Optical System Optimization used a related environment.

Getting started

To get started you can do the example py-files for tmm (example_tmm.py) or the gymnasium environment (example_gym.py)!

Multi-layer thin films meet parameterized reinforcement learning

Reinforcement learning is an area of machine learning concerned with how intelligent agents ought to take actions in an environment in order to maximize the notion of reward. The code to be published implements such an environment for the optimization of multi-layer thin films. In principle, the proposed code allows to execute actions taken by an agent. These actions determine which material and with which thickness to stack next, thereby consecutively forming a multi-layer thin film as illustrated in Figure 1. Such a multilayer thin-film exhibits optical characteristics. By comparison between the actual and user-defined desired characteristics, a notion of numeric reward is computed based on which the agent learns to distinguish between good and bad design choices. Due to its physical and mathematical structure, the optimization of multi-layer thin film remains a challenging and thus still active field of research in the scientific community. As such it gained recent attention in many publications. Therefore, naturally the need for a standardized environment arises to make the corresponding research more trustful, comparable and consistent.

Figure 2: Principal idea of an OpenAI/Farama-Foundation gymnasium environment. The agent takes an action that specifies the material and thickness of the layer to stack next. The environment implements the multi-layer thin film generation as consecutive conduction of actions and assigns a reward to a proposed multi-layer thin film based on how close the actual (solid orange line) fulfils a desired (dashed orange line) characteristic. The made experience is used to adapt the taken actions made in order to increase the reward and thus generate more and more sophisticated multi-layer thin films.

Description of key features

The environment can include
• cladding of the multi-layer thin film (e.g. substrate and ambient materials),
• dispersive and dissipative materials,
• spectral and angular optical behavior of multi-layer thin films (See figure 3),
• … and many more.

The environment class allows to
• conduct so-called parameterized actions (See publication) that define a multi-layer thin film,
• evaluate the generated thin film given a desired optical response, and
• render the results (See figure 3).

In general, the comprehensive optimization of multi-layer thin films in regards of optical response encompasses
• the number of layers (integer),
• the thickness of each layer (float),
• the material of each layer (categorical, integer).

Figure 3: Rendered output of the environment. Reflectivity (left) over angle of incidence and spectrum of a multi-layer thin film (right). Here, the stack features four layers and each layer’s material was chosen from a set of eight alternatives. The reward is computed based on a desired reflectivity, which is one for each angle and wavelength, but not displayed in this figure.

Citing

If you use the code from this repository for your projects, please cite
TMM-Fast: A Transfer Matrix Computation Package for Multilayer Thin-Film Optimization or/and
Parameterized reinforcement learning for optical system optimization in your publications.

tmm_fast's Issues

It works only with GPU

Hi, if I try to run the demo script example_tmm.py in a CPU-only environment, the following error appears:
AssertionError: It's not clear which beam is incoming vs outgoing. Weird index maybe?

Is there any suggested fix or workaround? Thanks!

Calculation result different from tmm package

Hi,

I would like to use tmm_fast for reflection calculation an optimisation expecially in a context like example 6 in tmm examples.

I took the example from the original tmm package and tried to reimplement the same calculation in tmm_fast. The outputs do not match.

import numpy as np
import matplotlib.pyplot as plt
import torch
## tmm

import tmm as tmm
"""
An example reflection plot with a surface plasmon resonance (SPR) dip.
Compare with http://doi.org/10.2320/matertrans.M2010003 ("Spectral and
Angular Responses of Surface Plasmon Resonance Based on the Kretschmann
Prism Configuration") Fig 6a
"""
# list of layer thicknesses in nm
d_list = [np.inf, 5, 30, np.inf]
# list of refractive indices
n_list = [1.517, 3.719+4.362j, 0.130+3.162j, 1]
# wavelength in nm
lam_vac = 633
# list of angles to plot
theta_list = np.linspace(30*np.pi/180, 60*np.pi/180, num=300)
# initialize lists of y-values to plot
Rp = []
for theta in theta_list:
    Rp.append(tmm.coh_tmm('p', n_list, d_list, theta, lam_vac)['R'])
    
## tmm_fast
from tmm_fast import coh_tmm

pol="p"
n_list= torch.tensor(n_list, dtype=torch.complex64)
d_list= np.array(d_list)

N = torch.zeros(1,4,1, dtype=torch.complex128)
N[0,:,0]= n_list

T=torch.Tensor(1,4)
T[0,:] = torch.Tensor(d_list)

theta = np.linspace(30*np.pi/180, 60*np.pi/180, num=300)
lambda_vacuum = torch.Tensor([633])
theta_incidence = torch.Tensor(theta)

tm = coh_tmm(pol,N,T,theta_incidence,lambda_vacuum)    

fig,ax = plt.subplots(1,2,figsize=(8,4))
ax[0].plot(theta_list/np.pi*180, Rp, 'blue')
ax[1].plot(theta_list/np.pi*180, np.array(tm['R']).reshape(-1,1), 'blue')
for a in ax:
    a.set_xlabel('theta (degree)')
    a.set_ylabel('Fraction reflected')
    a.set_xlim(30, 60)
    a.set_ylim(0, 1)
plt.suptitle('Reflection of p-polarized light with Surface Plasmon Resonance\n'
          'Compare with http://doi.org/10.2320/matertrans.M2010003 Fig 6a')
plt.show()

Which produces the attached output (left original tmm; right tmm_fast; the original tmm output is the expected outcome).

Am I using the package incorrectly or is there some other issue? tmm_fast is in version 0.2.1. tmm in 0.1.8

missing tmm_fast_core

tmm_fast_torch.py seams broken without tmm_fast_core

missing example.ipynb and quickstarter.py referenced in Readme

missing example.ipynb and quickstarter.py referenced in Readme
would be helpful if included

Nan will sometimes appear when calculating reflection coefficients

I used the package for genetic optimization and found that nan would appear when calculating reflectance in a large-scale search

Error when dealing with absorbing material

When the substrate is an absorbing material, the calculation is correct. However, when it is placed at upper layer of the photonic structure with thickness large enough, for instance a structure like W(e. g. 1000 nm)/SiO2/W(substrate), the upper layer of W (Although this structure is meaningless) would lead to 'nan' in calculated reflectance.

Simplify dependencies

Since the tmm-fast is published on pipy.org I think it is better to keep dependencies as minimal as possible.

For instance, it would be better to remove dask from dependencies, as it is not directly used anywhere in the project.

The main functionality in vectorized_tmm_dispersive_multistack.py and tmm_fast_torch.py does not require gymnasium, matplotlib and seaborn, therefore we can declare them as optional ones.

Probably, it could be even better to split the TMM and AI functionality into two separate projects...

input datatype for inc_vec_tmm_disp_lstack

inc_vec_tmm_disp_lstack must take torch data type as the input. numpy input result in the error message below:

AttributeError: 'numpy.ndarray' object has no attribute 'requires_grad'

Move from gym to gymnasium

Since the package gym is renamed to gymnasium, we have to also update the dependencies.
I have prepared a pull request for it.

Use `api-array-compat` as a backend

What do you think about the api-array-compat project?

It should be not so difficult to change the backend from PyTorch to api-array-compat, and it will give many advantages:

• Simplify array type selection/verification logic
• Add cupy support for those who want to calculate on GPU, but don't need backpropagation or don't want to install PyTorch
• Make PyTorch dependence optional

AttributeError: can't set attribute

When running example example_gym.py, failing at line

https://github.com/MLResearchAtOSRAM/tmm_fast/blob/main/tmm_fast/gym_multilayerthinfilm/gym_class.py#L133

with:

AttributeError: can't set attribute

object:
self.action_space.spaces Out [1]: (Discrete(5), Box(0.0, 1.0, (1,), float32), Box(0.0, 1.0, (1,), float32), Box(0.0, 1.0, (1,), float32), Box(0.0, 1.0, (1,), float32), Box(0.0, 1.0, (1,), float32))

Dataset generation example

It would be helpful to add an example how to do dataset generation via tmm_fast and Dask

Implementation of an equivalent to tmm.tmm_core.inc_tmm()

Hi again,
I was wondering if there are any plans to implement a vectorized version of the "incoherent, or partly-incoherent-partly-coherent, transfer matrix method" as implemented in the original tmm package by S. Byrnes?

For one of my projects, I reproduced the original functionality for a particular case (fitting to FTIR data of a DBR on a thick bulk slab, with vacuum as a sub- and superstrate).

I imagine that the partly-coherent method might benefit people in similar use cases, both in measurement and simulation.

My crude implementation already resulted in a considerable speed improvement compared to the original tmm code. If you're interested, I could take a deeper dive to generalize and optimize the code further for a possible contribution to tmm_fast.

I'm open to discussing this further should you be interested.

Best, superluminescent

Provide tmm_fast via conda-forge

Dear devs,
thanks for providing this extremely useful module and the accompanying tutorial J. Opt. Soc. Am. A.
I used an earlier version for a recent publication and made sure to cite your work.

However, with Anaconda/conda being popular entities in the scientific/R&D community, I have a general request: Would it be possible to provide tmm_fast via conda-forge for the current and future releases via @MLResearchAtOSRAM ?

As I'm only a conda user, I cannot estimate the workload of this request (see conda docs here). Still, it would be nice to use the full capabilities of conda with your package (environments, dependencies etc.) as I found it to be much more resilient than pip features such as venv.

Thanks for considering my request!

Best,
superluminescent (Lukas)

BUG: Unexpected squeezing of Theta and lambda_vacuum inputs in coh_vec_tmm_disp_mstack

I have tried to compare the output of the coh_vec_tmm_disp_mstack function with the original tmm package.
For this I have reduced the Theta and lambda_vacuum vectors to a single element as following:

import numpy as np
from tmm_fast.vectorized_tmm_dispersive_multistack import coh_vec_tmm_disp_mstack as tmm

# wl = np.asarray([400., 500.,]) * 1e-9 # This works well
wl = np.asarray([400.,]) * 1e-9 # AssertionError: N and T are not of same shape, as they are of dimensions 3 and 2
# theta = np.asarray([0., 45.,]) # This also works well
theta = np.asarray([0.,]) # IndexError: tuple index out of range
mode = 'T'
num_layers = 4
num_stacks = 128

refractive_index = np.ones([num_stacks, num_layers, wl.shape[0]])
thickness = np.ones([num_stacks, num_layers]) * 100

tmm(
    pol="s",
    N=refractive_index,
    T=thickness,
    Theta=theta,
    lambda_vacuum=wl,
)

However, I have faced the following error in the case of monochromatic frequency:

AssertionError: N and T are not of same shape, as they are of dimensions 3 and 2

And in the case of a single angle:

IndexError: tuple index out of range

I think that both of these problems could be fixed by removing unnecessary squeeze() from the converter function in vectorized_tmm_dispersive_multistack module:

def converter(data, device):
    if type(data) is not torch.Tensor:
        if type(data) is np.ndarray:
            data = torch.from_numpy(data.copy())
        else:
            raise ValueError('At least one of the inputs (i.e. N, Theta, ...) is not of type numpy.array or torch.Tensor!')
    data = data.type(torch.cfloat).to(device)
    return data.squeeze()

Request: Provide tmm_fast package via conda-forge

Dear devs,
thanks for providing this extremely useful module and the accompanying tutorial J. Opt. Soc. Am. A.
I used an earlier version for a recent publication and made sure to cite your work.

However, with Anaconda/conda being popular entities in the scientific/R&D community, I have a general request: Would it be possible to provide tmm_fast via conda-forge for the current and future releases via @MLResearchAtOSRAM ?

As I'm only a conda user, I cannot estimate the workload of this request (see conda docs here). Still, it would be nice to use the full capabilities of conda with your package (environments, dependencies etc.) as I found it to be much more resilient than pip features such as venv.

Thanks for considering my request!

Best,
superluminescent (Lukas)