<jupyter_start><jupyter_text>SAX Quick Start Let's go over the core functionality of SAX.  Environment variables SAX is based on JAX... here are some useful environment variables for working with JAX:<jupyter_code># select float32 or float64 as default dtype
%env JAX_ENABLE_X64=0

# select cpu or gpu
%env JAX_PLATFORM_NAME=cpu

# set custom CUDA location for gpu:
%env XLA_FLAGS=--xla_gpu_cuda_data_dir=/usr/lib/cuda

# Using GPU?
from jax.lib import xla_bridge
print(xla_bridge.get_backend().platform)<jupyter_output><empty_output><jupyter_text>Imports<jupyter_code>import jax
import jax.example_libraries.optimizers as opt
import jax.numpy as jnp
import matplotlib.pyplot as plt  # plotting
import sax
import tqdm  # progress bars<jupyter_output><empty_output><jupyter_text>Scatter *dictionaries*The core datastructure for specifying scatter parameters in SAX is a dictionary... more specifically a dictionary which maps a port combination (2-tuple) to a scatter parameter (or an array of scatter parameters when considering multiple wavelengths for example). Such a specific dictionary mapping is called ann `SDict` in SAX (`SDict ≈ Dict[Tuple[str,str], float]`).Dictionaries are in fact much better suited for characterizing S-parameters than, say, (jax-)numpy arrays due to the inherent sparse nature of scatter parameters. Moreover, dictonaries allow for string indexing, which makes them much more pleasant to use in this context. Let’s for example create an `SDict` for a 50/50 coupler: ```in1          out1   \        /    ========   /        \in0          out0```<jupyter_code>coupling = 0.5
kappa = coupling ** 0.5
tau = (1 - coupling) ** 0.5
coupler_dict = {
    ("in0", "out0"): tau,
    ("out0", "in0"): tau,
    ("in0", "out1"): 1j * kappa,
    ("out1", "in0"): 1j * kappa,
    ("in1", "out0"): 1j * kappa,
    ("out0", "in1"): 1j * kappa,
    ("in1", "out1"): tau,
    ("out1", "in1"): tau,
}
coupler_dict<jupyter_output><empty_output><jupyter_text>Only the non-zero port combinations need to be specified. Any non-existent port-combination (for example `("in0", "in1")`) is considered to be zero by SAX.Obviously, it can still be tedious to specify every port in the circuit manually. SAX therefore offers the `reciprocal` function, which auto-fills the reverse connection if the forward connection exist. For example:<jupyter_code>coupler_dict = sax.reciprocal(
    {
        ("in0", "out0"): tau,
        ("in0", "out1"): 1j * kappa,
        ("in1", "out0"): 1j * kappa,
        ("in1", "out1"): tau,
    }
)

coupler_dict<jupyter_output><empty_output><jupyter_text>Parametrized Models Constructing such an `SDict` is easy, however, usually we're more interested in having parametrized models for our components. To parametrize the coupler `SDict`, just wrap it in a function to obtain a SAX `Model`, which is a keyword-only function mapping to an `SDict`:<jupyter_code>def coupler(coupling=0.5) -> sax.SDict:
    kappa = coupling ** 0.5
    tau = (1 - coupling) ** 0.5
    coupler_dict = sax.reciprocal(
        {
            ("in0", "out0"): tau,
            ("in0", "out1"): 1j * kappa,
            ("in1", "out0"): 1j * kappa,
            ("in1", "out1"): tau,
        }
    )
    return coupler_dict


coupler(coupling=0.3)<jupyter_output><empty_output><jupyter_text>We can define a waveguide in the same way:<jupyter_code>def waveguide(wl=1.55, wl0=1.55, neff=2.34, ng=3.4, length=10.0, loss=0.0) -> sax.SDict:
    dwl = wl - wl0
    dneff_dwl = (ng - neff) / wl0
    neff = neff - dwl * dneff_dwl
    phase = 2 * jnp.pi * neff * length / wl
    transmission = 10 ** (-loss * length / 20) * jnp.exp(1j * phase)
    sdict = sax.reciprocal(
        {
            ("in0", "out0"): transmission,
        }
    )
    return sdict<jupyter_output><empty_output><jupyter_text>That's pretty straightforward. Let's now move on to parametrized circuits:  Circuit Models Existing models can now be combined into a circuit using `sax.circuit`, which basically creates a new `Model` function:<jupyter_code>mzi = sax.circuit(
    instances={
        "lft": coupler,
        "top": waveguide,
        "btm": waveguide,
        "rgt": coupler,
    },
    connections={
        "lft,out0": "btm,in0",
        "btm,out0": "rgt,in0",
        "lft,out1": "top,in0",
        "top,out0": "rgt,in1",
    },
    ports={
        "in0": "lft,in0",
        "in1": "lft,in1",
        "out0": "rgt,out0",
        "out1": "rgt,out1",
    },
)
mzi?<jupyter_output><empty_output><jupyter_text>The `circuit` function just creates a similar function as we created for the waveguide and the coupler, but in stead of taking parameters directly it takes parameter *dictionaries* for each of the instances in the circuit. The keys in these parameter dictionaries should correspond to the keyword arguments of each individual subcomponent.  Let's now do a simulation for the MZI we just constructed:<jupyter_code>%time mzi()
mzi2 = jax.jit(mzi)
%time mzi2()<jupyter_output><empty_output><jupyter_text>Or in the case we want an MZI with different arm lengths:<jupyter_code>mzi(top={"length": 25.0}, btm={"length": 15.0})<jupyter_output><empty_output><jupyter_text>Simulating the parametrized MZI We can simulate the above mzi for multiple wavelengths as well by specifying the wavelength at the top level of the circuit call. Each setting specified at the top level of the circuit call will be propagated to all subcomponents of the circuit which have that setting:<jupyter_code>wl = jnp.linspace(1.51, 1.59, 1000)
%time S = mzi(wl=wl, top={"length": 25.0}, btm={"length": 15.0})<jupyter_output><empty_output><jupyter_text>Let's see what this gives:<jupyter_code>plt.plot(wl * 1e3, abs(S["in0", "out0"]) ** 2)
plt.ylim(-0.05, 1.05)
plt.xlabel("λ [nm]")
plt.ylabel("T")
plt.ylim(-0.05, 1.05)
plt.show()<jupyter_output><empty_output><jupyter_text>Optimization We'd like to optimize an MZI such that one of the minima is at 1550nm. To do this, we need to define a loss function for the circuit at 1550nm. This function should take the parameters that you want to optimize as positional arguments:<jupyter_code>@jax.jit
def loss(delta_length):
    S = mzi(wl=1.55, top={"length": 15.0 + delta_length}, btm={"length": 15.0})
    return (abs(S["in0", "out0"]) ** 2).mean()
%time loss(10.0)<jupyter_output><empty_output><jupyter_text>We can use this loss function to define a grad function which works on the parameters of the loss function:<jupyter_code>grad = jax.jit(
    jax.grad(
        loss,
        argnums=0,  # JAX gradient function for the first positional argument, jitted
    )
)<jupyter_output><empty_output><jupyter_text>Next, we need to define a JAX optimizer, which on its own is nothing more than three more functions:  an initialization function with which to initialize the optimizer state, an update function which will update the optimizer state (and with it the model parameters). The third function that's being returned will give the model parameters given the optimizer state.<jupyter_code>initial_delta_length = 10.0
optim_init, optim_update, optim_params = opt.adam(step_size=0.1)
optim_state = optim_init(initial_delta_length)<jupyter_output><empty_output><jupyter_text>Given all this, a single training step can be defined:<jupyter_code>def train_step(step, optim_state):
    settings = optim_params(optim_state)
    lossvalue = loss(settings)
    gradvalue = grad(settings)
    optim_state = optim_update(step, gradvalue, optim_state)
    return lossvalue, optim_state<jupyter_output><empty_output><jupyter_text>And we can use this step function to start the training of the MZI:<jupyter_code>range_ = tqdm.trange(300)
for step in range_:
    lossvalue, optim_state = train_step(step, optim_state)
    range_.set_postfix(loss=f"{lossvalue:.6f}")
delta_length = optim_params(optim_state)
delta_length<jupyter_output><empty_output><jupyter_text>Let's see what we've got over a range of wavelengths:<jupyter_code>S = mzi(wl=wl, top={"length": 15.0 + delta_length}, btm={"length": 15.0})
plt.plot(wl * 1e3, abs(S["in1", "out1"]) ** 2)
plt.xlabel("λ [nm]")
plt.ylabel("T")
plt.ylim(-0.05, 1.05)
plt.plot([1550, 1550], [0, 1])
plt.show()<jupyter_output><empty_output><jupyter_text>The minimum of the MZI is perfectly located at 1550nm.  MZI Chain Let's now create a chain of MZIs. For this, we first create a subcomponent: a directional coupler with arms:```                                 top                            in0 ----- out0 -> out1    in1 <- in1        out1                               \  dc  /                                 ======                                /      \                    in0 <- in0        out0       btm                                in0 ----- out0 -> out0```<jupyter_code>dc_with_arms = sax.circuit(
    instances={
        "lft": coupler,
        "top": waveguide,
        "btm": waveguide,
    },
    connections={
        "lft,out0": "btm,in0",
        "lft,out1": "top,in0",
    },
    ports={
        "in0": "lft,in0",
        "in1": "lft,in1",
        "out0": "btm,out0",
        "out1": "top,out0",
    },
)<jupyter_output><empty_output><jupyter_text>An MZI chain can now be created by cascading these directional couplers with arms:```      _    _    _    _             _    _      \/   \/   \/   \/     ...    \/   \/       /\_  /\_  /\_  /\_           /\_  /\_  ```Let's create a *model factory* (`ModelFactory`) for this. In SAX, a *model factory* is any keyword-only function that generates a `Model`:<jupyter_code>def mzi_chain(num_mzis=1) -> sax.Model:
    chain = sax.circuit(
        instances={f"dc{i}": dc_with_arms for i in range(num_mzis + 1)},
        connections={
            **{f"dc{i},out0": f"dc{i+1},in0" for i in range(num_mzis)},
            **{f"dc{i},out1": f"dc{i+1},in1" for i in range(num_mzis)},
        },
        ports={
            "in0": f"dc0,in0",
            "in1": f"dc0,in1",
            "out0": f"dc{num_mzis},out0",
            "out1": f"dc{num_mzis},out1",
        },
    )
    return chain<jupyter_output><empty_output><jupyter_text>Let's for example create a chain with 15 MZIs. We can also update the settings dictionary as follows:<jupyter_code>chain = mzi_chain(num_mzis=15)
settings = sax.get_settings(chain)
for dc in settings:
    settings[dc]["top"]["length"] = 25.0
    settings[dc]["btm"]["length"] = 15.0
settings = sax.update_settings(settings, wl=jnp.linspace(1.5, 1.6, 1000))<jupyter_output><empty_output><jupyter_text>We can simulate this:<jupyter_code>%time S = chain(**settings) # time to evaluate the MZI
func = jax.jit(chain)
%time S = func(**settings) # time to jit the MZI
%time S = func(**settings) # time to evaluate the MZI after jitting<jupyter_output><empty_output><jupyter_text>Where we see that the unjitted evaluation of the MZI chain takes about a second, while the jitting of the MZI chain takes about two minutes (on a CPU). However, after the MZI chain has been jitted, the evaluation is in the order of about a millisecond!Anyway, let's see what this gives:<jupyter_code>plt.plot(1e3 * settings["dc0"]["top"]["wl"], jnp.abs(S["in0", "out0"]) ** 2)
plt.ylim(-0.05, 1.05)
plt.xlabel("λ [nm]")
plt.ylabel("T")
plt.ylim(-0.05, 1.05)
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12613900303734144, 'fraction_numerical': 0.026800071466857246, 'mean_word_length': 3.5619396903015486, 'pattern_counts': {'":': 36, '<': 91, '<?xml version=': 0, '>': 94, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '22', 'dir': 'jupyter-structured-clean-dedup', 'id': '123878', 'license': 'Apache-2.0', 'n_tokens_mistral': 4046, 'n_tokens_neox': 3807, 'n_words': 1124, 'path': 'examples/01_quick_start.ipynb', 'repo_name': 'joamatab/sax', 'size': '11194'}