<jupyter_start><jupyter_text>Thin film optimization and wavelength-dependent parametersIn this notebook, we apply SAX to thin-film optimization and show how it can be used for wavelength-dependent parameter optimization.The language of transfer/scatter matrices is commonly used to calculate optical properties of thin-films. Many specialized methods exist for their optimization. However, SAX can be useful to cut down on developer time by circumventing the need to manually take gradients of complicated or often-changed objective functions, and by generating efficient code from simple syntax.   Imports<jupyter_code>import tqdm
import matplotlib.pyplot as plt

# GPU setup
#%env XLA_FLAGS=--xla_gpu_cuda_data_dir=/path/to/cuda

import jax
import jax.numpy as jnp
import jax.experimental.optimizers as opt

# sax circuit simulator
import sax
# Using GPU?
from jax.lib import xla_bridge 
print(xla_bridge.get_backend().platform)<jupyter_output><empty_output><jupyter_text>Dielectric mirror Fabry-PérotConsider a stack composed of only two materials, $n_A$ and $n_B$. Two types of transfer matrices characterize wave propagation in the system : interfaces described by Fresnel's equations, and propagation.<jupyter_code># Reflection at i-->j interface
def r_fresnel_ij(params):
    """
    Normal incidence amplitude reflection from Fresnel's equations
    ni : refractive index of the initial medium
    nf : refractive index of the final
    """
    return (params["ni"] - params["nj"]) / (params["ni"] + params["nj"])

# Transmission at i-->j interface
def t_fresnel_ij(params):
    """
    Normal incidence amplitude transmission from Fresnel's equations
    """
    return 2 * params["ni"] / (params["ni"] + params["nj"])

# Propagation through medium A
def prop_i(params):
    """
    Phase shift acquired as a wave propagates through medium A
    wl : wavelength (arb. units)
    ni : refractive index of medium (at wavelength wl)
    di : thickness of layer (same arb. unit as wl)
    """
    return jnp.exp(1j * 2*jnp.pi * params["ni"] / params["wl"] * params["di"])<jupyter_output><empty_output><jupyter_text>For the two-material stack, this leads to 4 scatter matrices coefficients. Through reciprocity they can be constructed out of two independent ones :<jupyter_code># Arbitrary default parameters
fresnel_mirror_ij = {
    ("in", "in"): r_fresnel_ij,
    ("in", "out"): t_fresnel_ij,
    ("out", "in"): lambda params: (1 - r_fresnel_ij(params)**2)/t_fresnel_ij(params), # t_ji,
    ("out", "out"): lambda params: -1*r_fresnel_ij(params), # r_ji,
    "default_params": {
        "ni": 1.,
        "nj": 1.,
        "wl": 532.,
    }
}

propagation_i = {
    ("in", "out"): prop_i,
    ("out", "in"): prop_i,
    "default_params": {
        "ni": 1.,
        "di": 500.,
        "wl": 532.,
    }
}<jupyter_output><empty_output><jupyter_text>A resonant cavity can be formed when a high index region is surrounded by low-index region :<jupyter_code>dielectric_fabry_perot = sax.circuit(
    models = {
        "air-B": fresnel_mirror_ij,
        "B": propagation_i,
        "B-air": fresnel_mirror_ij,
    },
    connections={
        "air-B:out": "B:in",
        "B:out": "B-air:in",
    },
    ports={
        "air-B:in": "in",
        "B-air:out": "out",
    },
)

params = dielectric_fabry_perot["default_params"]
params<jupyter_output><empty_output><jupyter_text>Let's choose $n_A = 1$, $n_B = 2$, $d_B = 1000$ nm, and compute over the visible spectrum :<jupyter_code>params = sax.copy_params(params)
params["air-B"]["nj"] = 2.
params["B"]["ni"] = 2.
params["B-air"]["ni"] = 2.

wls = jnp.linspace(380, 750, 200)
params = sax.set_global_params(params, wl=wls)<jupyter_output><empty_output><jupyter_text>Compute transmission and reflection, and compare to another package's results (https://github.com/sbyrnes321/tmm) :<jupyter_code>fabry_perot_transmitted = dielectric_fabry_perot["in","out"]
fabry_perot_transmitted = jax.jit(dielectric_fabry_perot["in", "out"])
transmitted = fabry_perot_transmitted(params)

fabry_perot_reflected = dielectric_fabry_perot["in","in"]
fabry_perot_reflected = jax.jit(dielectric_fabry_perot["in", "in"])
reflected = fabry_perot_reflected(params)
# import sys
# !$sys.executable -m pip install tmm
from tmm import coh_tmm

# tmm syntax (https://github.com/sbyrnes321/tmm)
d_list = [jnp.inf,500,jnp.inf]
n_list = [1,2,1]
# initialize lists of y-values to plot
rnorm=[]
tnorm = []
Tnorm = []
Rnorm = []
for l in wls:
    rnorm.append(coh_tmm('s',n_list, d_list, 0, l)['r'])
    tnorm.append(coh_tmm('s',n_list, d_list, 0, l)['t'])
    Tnorm.append(coh_tmm('s',n_list, d_list, 0, l)['T'])
    Rnorm.append(coh_tmm('s',n_list, d_list, 0, l)['R'])
plt.scatter(wls, jnp.real(transmitted), label='t SAX')
plt.plot(wls, jnp.real(jnp.array(tnorm)), 'k', label='t tmm')
plt.scatter(wls, jnp.real(reflected), label='r SAX')
plt.plot(wls, jnp.real(jnp.array(rnorm)), 'k--', label='r tmm')
plt.xlabel("λ [nm]")
plt.ylabel("Transmitted and reflected amplitude")
plt.legend(loc="upper right")
plt.title("Real part")
plt.show()

plt.scatter(wls, jnp.imag(transmitted), label='t SAX')
plt.plot(wls, jnp.imag(jnp.array(tnorm)), 'k', label='t tmm')
plt.scatter(wls, jnp.imag(reflected), label='r SAX')
plt.plot(wls, jnp.imag(jnp.array(rnorm)), 'k--', label='r tmm')
plt.xlabel("λ [nm]")
plt.ylabel("Transmitted and reflected amplitude")
plt.legend(loc="upper right")
plt.title("Imaginary part")
plt.show()<jupyter_output><empty_output><jupyter_text>In terms of powers, we get the following. Due to the reflections at the interfaces, resonant behaviour is observed, with evenly-spaced maxima/minima in wavevector space :<jupyter_code>plt.scatter(2*jnp.pi/wls, jnp.abs(transmitted)**2, label='T SAX')
plt.plot(2*jnp.pi/wls, Tnorm, 'k', label='T tmm')
plt.scatter(2*jnp.pi/wls, jnp.abs(reflected)**2, label='R SAX')
plt.plot(2*jnp.pi/wls, Rnorm, 'k--', label='R tmm')
plt.vlines(jnp.arange(3,6)*jnp.pi/(2*500), ymin=0, ymax=1, color='k', linestyle='--', label='m$\pi$/nd')
plt.xlabel("k = 2$\pi$/λ [nm]")
plt.ylabel("Transmitted and reflected intensities")
plt.legend(loc="upper right")
plt.show()<jupyter_output><empty_output><jupyter_text>Optimization testLet's attempt to minimize transmission at 500 nm by varying thickness.<jupyter_code>@jax.jit
def loss(thickness):
    params = sax.set_global_params(dielectric_fabry_perot["default_params"], wl=500.)
    params["B"]["di"] = thickness
    params["air-B"]["nj"] = 2.
    params["B"]["ni"] = 2.
    params["B-air"]["ni"] = 2.
    detected = dielectric_fabry_perot["in", "out"](params)
    return abs(detected)**2
%time loss(500.)
grad = jax.jit(jax.grad(loss))
%time grad(500.)
initial_thickness = 500.
optim_init, optim_update, optim_params = opt.adam(step_size=10)
optim_state = optim_init(initial_thickness)

@jax.jit
def train_step(step, optim_state):
    params = optim_params(optim_state)
    lossvalue = loss(params)
    gradvalue = grad(params)
    optim_state = optim_update(step, gradvalue, optim_state)
    return lossvalue, optim_state
range_ = tqdm.trange(1000)
for step in range_:
    lossvalue, optim_state = train_step(step, optim_state)
    range_.set_postfix(loss=f"{lossvalue:.6f}")
thickness = optim_params(optim_state)
thickness
params = sax.set_global_params(dielectric_fabry_perot["default_params"], wl=wls)
params["B"]["di"] = thickness
params["air-B"]["nj"] = 2.
params["B"]["ni"] = 2.
params["B-air"]["ni"] = 2.
detected = dielectric_fabry_perot["in", "out"](params)

plt.plot(wls, jnp.abs(transmitted)**2, label='Before (500 nm)')
plt.plot(wls, jnp.abs(detected)**2, label="After ({} nm)".format(thickness))
plt.vlines(500, 0.6, 1, 'k', linestyle='--')
plt.xlabel("λ [nm]")
plt.ylabel("Transmitted intensity")
plt.legend(loc="lower right")
plt.title("Thickness optimization")
plt.show()<jupyter_output><empty_output><jupyter_text>General Fabry-Pérot étalonWe reuse the propagation matrix above, and instead of simple interface matrices, model Fabry-Pérot mirrors as general lossless reciprocal scatter matrices :$$ \left(\begin{array}{c} E_t \\E_r\end{array}\right) = E_{out} = SE_{in} = \left(\begin{array}{cc} t & r \\r & t\end{array}\right) \left(\begin{array}{c} E_0 \\0\end{array}\right) $$For lossless reciprocal systems, we further have the requirements$$ |t|^2 + |r|^2 = 1 $$and$$ \angle t - \angle r = \pm \pi/2 $$The general Fabry-Pérot cavity is analytically described by :<jupyter_code># General Fabry-Pérot transfer function (Airy formulas)
def airy_t13(t12, t23, r21, r23, wl, d=1000., n=1.):
    '''
    Assumptions 
    Each mirror lossless, reciprocal : tij = tji, rij = rji
    
    Inputs
    t12 and r12 : S-parameters of the first mirror
    t23 and r23 : S-parameters of the second mirror
    wl : wavelength
    d : gap between the two mirrors (in units of wavelength)
    n : index of the gap between the two mirrors
    
    Returns
    t13 : complex transmission amplitude of the mirror-gap-mirror system
    '''
    # Assume each mirror lossless, reciprocal : tij = tji, rij = rji
    phi = n*2*jnp.pi/wl*d
    return t12*t23*jnp.exp(-1j*phi)/( 1 - r21*r23*jnp.exp(-2j*phi) )

def airy_r13(t12, t23, r21, r23, wl, d=1000., n=1.):
    '''
    Assumptions, inputs : see airy_t13
    
    Returns
    r13 : complex reflection amplitude of the mirror-gap-mirror system
    '''
    phi = n*2*jnp.pi/wl*d
    return r21 + t12*t12*r23*jnp.exp(-2j*phi)/( 1 - r21*r23*jnp.exp(-2j*phi) )<jupyter_output><empty_output><jupyter_text>We need to implement the relationship between $t$ and $r$ for lossless reciprocal mirrors. The design parameter will be the amplitude and phase of the tranmission coefficient. The reflection coefficient is then fully determined :<jupyter_code>def t_complex(params):
    # Transmission coefficient (design parameter)
    return params['t_amp']*jnp.exp(-1j*params['t_ang'])

def r_complex(params):
    # Reflection coefficient, derived from transmission coefficient
    # Magnitude from |t|^2 + |r|^2 = 1
    # Phase from phase(t) - phase(r) = pi/2
    r_amp = jnp.sqrt( ( 1. - params['t_amp']**2 ) )
    r_ang = params['t_ang'] - jnp.pi/2
    return r_amp*jnp.exp(-1j*r_ang)<jupyter_output><empty_output><jupyter_text>Let's see the expected result for half-mirrors :<jupyter_code>t_initial = jnp.sqrt(0.5)
d_gap = 2000.
n_gap = 1.
params_analytical_test = {"t_amp": t_initial, "t_ang": 0.0}
r_initial = r_complex(params_analytical_test)

wls = jnp.linspace(380, 780, 500)

T_analytical_initial = jnp.abs(airy_t13(t_initial, t_initial, r_initial, r_initial, wls, d=d_gap, n=n_gap))**2
R_analytical_initial = jnp.abs(airy_r13(t_initial, t_initial, r_initial, r_initial, wls, d=d_gap, n=n_gap))**2

plt.title(f't={t_initial:1.3f}, d={d_gap} nm, n={n_gap}')
plt.plot(2*jnp.pi/wls, T_analytical_initial, label='T')
plt.plot(2*jnp.pi/wls, R_analytical_initial, label='R')
plt.vlines(jnp.arange(6,11)*jnp.pi/2000, ymin=0, ymax=1, color='k', linestyle='--', label='m$\pi$/nd')
plt.xlabel('k = 2$\pi$/$\lambda$ (/nm)')
plt.ylabel('Power (units of input)')
plt.legend()
plt.show()

plt.title(f't={t_initial:1.3f}, d={d_gap} nm, n={n_gap}')
plt.plot(wls, T_analytical_initial, label='T')
plt.plot(wls, R_analytical_initial, label='R')
plt.xlabel('$\lambda$ (nm)')
plt.ylabel('Power (units of input)')
plt.legend()
# Is power conserved? (to within 0.1%)
assert jnp.isclose(R_analytical_initial + T_analytical_initial, 1, 0.001).all()<jupyter_output><empty_output><jupyter_text>Now let's do the same with SAX by defining new elements :<jupyter_code>mirror = {
    ("in", "in"): r_complex,
    ("in", "out"): t_complex,
    ("out", "in"): t_complex, # lambda params: (1 - r_complex(params)**2)/t_complex(params), # t_ji,
    ("out", "out"): r_complex, # lambda params: -1*r_complex(params), # r_ji,
    "default_params": {
        "t_amp": jnp.sqrt(0.5),
        "t_ang": 0.0,
    }
}

fabry_perot_tunable = sax.circuit(
    models = {
        "mirror1": mirror,
        "gap": propagation_i,
        "mirror2": mirror,
    },
    connections={
        "mirror1:out": "gap:in",
        "gap:out": "mirror2:in",
    },
    ports={
        "mirror1:in": "in",
        "mirror2:out": "out",
    },
)

params = fabry_perot_tunable["default_params"]
params = sax.copy_params(params)
params
fabry_perot_tunable = sax.circuit(
    models = {
        "mirror1": mirror,
        "gap": propagation_i,
        "mirror2": mirror,
    },
    connections={
        "mirror1:out": "gap:in",
        "gap:out": "mirror2:in",
    },
    ports={
        "mirror1:in": "in",
        "mirror2:out": "out",
    },
)

params = fabry_perot_tunable["default_params"]
params = sax.copy_params(params)
params
N = 100
wls = jnp.linspace(380, 780, N)
params = sax.copy_params(fabry_perot_tunable["default_params"])
params = sax.set_global_params(params, wl=wls)
params = sax.set_global_params(params, t_amp=jnp.sqrt(0.5))
params = sax.set_global_params(params, t_ang=0.0)
params["gap"]["ni"] = 1.
params["gap"]["di"] = 2000.
transmitted_initial = fabry_perot_tunable["in","out"](params)
reflected_initial = fabry_perot_tunable["out","out"](params)
T_analytical_initial = jnp.abs(airy_t13(t_initial, t_initial, r_initial, r_initial, wls, d=d_gap, n=n_gap))**2
R_analytical_initial = jnp.abs(airy_r13(t_initial, t_initial, r_initial, r_initial, wls, d=d_gap, n=n_gap))**2

plt.title(f't={t_initial:1.3f}, d={d_gap} nm, n={n_gap}')
plt.plot(wls, T_analytical_initial, label='T theory')
plt.scatter(wls, jnp.abs(transmitted_initial)**2, label='T SAX')
plt.plot(wls, R_analytical_initial, label='R theory')
plt.scatter(wls, jnp.abs(reflected_initial)**2, label='R SAX')
#plt.vlines(jnp.arange(6,11)*jnp.pi/2000, ymin=0, ymax=1, color='k', linestyle='--', label='m$\pi$/nd')
plt.xlabel('k = 2$\pi$/$\lambda$ (/nm)')
plt.ylabel('Power (units of input)')
plt.legend()
plt.show()<jupyter_output><empty_output><jupyter_text>Wavelength-dependent Fabry-Pérot étalonLet's repeat with a model where parameters can be wavelength-dependent. To comply with the optimizer object, we will stack all design parameters in a single array :<jupyter_code>ts_initial = jnp.zeros(2*N)
ts_initial = jax.ops.index_update(ts_initial, jax.ops.index[0:N], jnp.sqrt(0.5))<jupyter_output><empty_output><jupyter_text>We will simply loop over all wavelengths, and use different $t$ parameters at each wavelength.<jupyter_code>wls = jnp.linspace(380, 780, N)
transmitted = jnp.zeros_like(wls)
reflected = jnp.zeros_like(wls)

for i in range(N):
    # Update parameters
    params = sax.copy_params(fabry_perot_tunable["default_params"])
    params = sax.set_global_params(params, wl=wls[i])
    params = sax.set_global_params(params, t_amp=ts_initial[i])
    params = sax.set_global_params(params, t_ang=ts_initial[N+i])
    params["gap"]["ni"] = 1.
    params["gap"]["di"] = 2000.
    # Perform computation
    transmission_i = fabry_perot_tunable["in","out"](params)
    transmitted = jax.ops.index_update(transmitted, jax.ops.index[i], jnp.abs(transmission_i)**2)
    reflected_i = fabry_perot_tunable["in","in"](params)
    reflected = jax.ops.index_update(reflected, jax.ops.index[i], jnp.abs(reflected_i)**2)
plt.plot(wls, T_analytical_initial, label='T theory')
plt.scatter(wls, transmitted, label='T SAX')
plt.plot(wls, R_analytical_initial, label='R theory')
plt.scatter(wls, reflected, label='R SAX')
plt.xlabel("λ [nm]")
plt.ylabel("Transmitted and reflected intensities")
plt.legend(loc="upper right")
plt.title(f't={t_initial:1.3f}, d={d_gap} nm, n={n_gap}')
plt.show()<jupyter_output><empty_output><jupyter_text>Since it seems to work, let's add a target and optimize some harmonics away :<jupyter_code>def lorentzian(l0, dl, wl, A):
    return A/((wl - l0)**2 + (0.5*dl)**2)

target = lorentzian(533, 20, wls, 100)

plt.scatter(wls, transmitted, label='T SAX')
plt.scatter(wls, reflected, label='R SAX')
plt.plot(wls, target, 'r', linewidth=2, label='target')
plt.xlabel("λ [nm]")
plt.ylabel("Transmitted and reflected intensities")
plt.legend(loc="upper right")
plt.title(f't={t_initial:1.3f}, d={d_gap} nm, n={n_gap}')
plt.show()<jupyter_output><empty_output><jupyter_text>Defining the loss as the mean squared error between transmission and target transmission. Note that we can use JAX's looping functions to great speedups after compilation here :<jupyter_code>@jax.jit
def loss(ts):

    N = len(ts[::2])
    wls = jnp.linspace(380, 780, N)
    transmitted = jnp.zeros_like(wls)
    target = lorentzian(533, 20, wls, 100)
    
    def inner_loop(transmitted, i):
        # Update parameters
        params = sax.copy_params(fabry_perot_tunable["default_params"])
        params = sax.set_global_params(params, wl=wls[i])
        params = sax.set_global_params(params, t_amp=ts[i])
        params = sax.set_global_params(params, t_ang=ts[N+i])
        params["gap"]["ni"] = 1.
        params["gap"]["di"] = 2000.
        # Perform computation
        transmission_i = fabry_perot_tunable["in","out"](params)
        transmitted = jax.ops.index_update(transmitted, jax.ops.index[i], jnp.abs(transmission_i)**2)
        return transmitted, i

    transmitted, _ = jax.lax.scan(inner_loop, transmitted, jnp.arange(N, dtype=jnp.int32))
        
    return (jnp.abs(transmitted - target)**2).mean()
grad = jax.jit(jax.grad(loss))
optim_init, optim_update, optim_params = opt.adam(step_size=0.001)

@jax.jit
def train_step(step, optim_state):
    params = optim_params(optim_state)
    lossvalue = loss(params)
    gradvalue = grad(params)
    optim_state = optim_update(step, gradvalue, optim_state)
    return lossvalue, optim_state
range_ = tqdm.trange(2000)

optim_state = optim_init(ts_initial)
for step in range_:
    lossvalue, optim_state = train_step(step, optim_state)
    range_.set_postfix(loss=f"{lossvalue:.6f}")<jupyter_output><empty_output><jupyter_text>The optimized parameters are now wavelength-dependent :<jupyter_code>ts_optimal = optim_params(optim_state)

plt.scatter(wls, ts_initial[:N], label='t initial')
plt.scatter(wls, ts_optimal[:N], label='t optimal')
plt.xlabel("λ [nm]")
plt.ylabel("|t| $(\lambda)$")
plt.legend(loc="best")
plt.title(f'd={d_gap} nm, n={n_gap}')
plt.show()

plt.scatter(wls, ts_initial[N:], label='t initial')
plt.scatter(wls, ts_optimal[N:], label='t optimal')
plt.xlabel("λ [nm]")
plt.ylabel("angle $t (\lambda)$ (rad)")
plt.legend(loc="best")
plt.title(f'd={d_gap} nm, n={n_gap}')
plt.show()<jupyter_output><empty_output><jupyter_text>Visualizing the result :<jupyter_code>wls = jnp.linspace(380, 780, N)
transmitted_optimal = jnp.zeros_like(wls)
reflected_optimal = jnp.zeros_like(wls)

for i in range(N):
    # Update parameters
    params = sax.copy_params(fabry_perot_tunable["default_params"])
    params = sax.set_global_params(params, wl=wls[i])
    params = sax.set_global_params(params, t_amp=ts_optimal[i])
    params = sax.set_global_params(params, t_ang=ts_optimal[N+i])
    params["gap"]["ni"] = 1.
    params["gap"]["di"] = 2000.
    # Perform computation
    transmission_i = fabry_perot_tunable["in","out"](params)
    transmitted_optimal = jax.ops.index_update(transmitted_optimal, jax.ops.index[i], jnp.abs(transmission_i)**2)
    reflected_i = fabry_perot_tunable["in","in"](params)
    reflected_optimal = jax.ops.index_update(reflected_optimal, jax.ops.index[i], jnp.abs(reflected_i)**2)
plt.scatter(wls, transmitted_optimal, label='T')
plt.scatter(wls, reflected_optimal, label='R')
plt.plot(wls, lorentzian(533, 20, wls, 100), 'r', label='target')
plt.xlabel("λ [nm]")
plt.ylabel("Transmitted and reflected intensities")
plt.legend(loc="upper right")
plt.title(f'Optimized t($\lambda$), d={d_gap} nm, n={n_gap}')
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15890327183773315, 'fraction_numerical': 0.019824686576291918, 'mean_word_length': 5.049013563501849, 'pattern_counts': {'":': 34, '<': 73, '<?xml version=': 0, '>': 75, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '18', 'dir': 'jupyter-structured-clean-dedup', 'id': '422386', 'license': 'Apache-2.0', 'n_tokens_mistral': 7886, 'n_tokens_neox': 7344, 'n_words': 1565, 'path': 'examples/02_thinfilm.ipynb', 'repo_name': 'simbilod/sax', 'size': '19622'}