SAX Quick Start¶
Let's go over the core functionality of SAX.
Imports¶
import jax
import jax.example_libraries.optimizers as opt
import jax.numpy as jnp
import matplotlib.pyplot as plt
from tqdm.notebook import trange
import sax
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 an 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:
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
{('in0', 'out0'): 0.7071067811865476,
('out0', 'in0'): 0.7071067811865476,
('in0', 'out1'): 0.7071067811865476j,
('out1', 'in0'): 0.7071067811865476j,
('in1', 'out0'): 0.7071067811865476j,
('out0', 'in1'): 0.7071067811865476j,
('in1', 'out1'): 0.7071067811865476,
('out1', 'in1'): 0.7071067811865476}
Tip
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:
coupler_dict = sax.reciprocal(
{
("in0", "out0"): tau,
("in0", "out1"): 1j * kappa,
("in1", "out0"): 1j * kappa,
("in1", "out1"): tau,
}
)
coupler_dict
{('in0', 'out0'): 0.7071067811865476,
('in0', 'out1'): 0.7071067811865476j,
('in1', 'out0'): 0.7071067811865476j,
('in1', 'out1'): 0.7071067811865476,
('out0', 'in0'): 0.7071067811865476,
('out1', 'in0'): 0.7071067811865476j,
('out0', 'in1'): 0.7071067811865476j,
('out1', 'in1'): 0.7071067811865476}
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:
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)
{('in0', 'out0'): 0.8366600265340756,
('in0', 'out1'): 0.5477225575051661j,
('in1', 'out0'): 0.5477225575051661j,
('in1', 'out1'): 0.8366600265340756,
('out0', 'in0'): 0.8366600265340756,
('out1', 'in0'): 0.5477225575051661j,
('out0', 'in1'): 0.5477225575051661j,
('out1', 'in1'): 0.8366600265340756}
We can define a waveguide in the same way:
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
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:
mzi, info = sax.circuit(
netlist={
"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",
},
},
models={
"coupler": coupler,
"waveguide": waveguide,
},
)
{('in0', 'in0'): Array(0.+0.j, dtype=complex128),
('in0', 'in1'): Array(0.+0.j, dtype=complex128),
('in0', 'out0'): Array(0.+0.j, dtype=complex128),
('in0', 'out1'): Array(-0.57126822+0.82076344j, dtype=complex128),
('in1', 'in0'): Array(0.+0.j, dtype=complex128),
('in1', 'in1'): Array(0.+0.j, dtype=complex128),
('in1', 'out0'): Array(-0.57126822+0.82076344j, dtype=complex128),
('in1', 'out1'): Array(0.+0.j, dtype=complex128),
('out0', 'in0'): Array(0.+0.j, dtype=complex128),
('out0', 'in1'): Array(-0.57126822+0.82076344j, dtype=complex128),
('out0', 'out0'): Array(0.+0.j, dtype=complex128),
('out0', 'out1'): Array(0.+0.j, dtype=complex128),
('out1', 'in0'): Array(-0.57126822+0.82076344j, dtype=complex128),
('out1', 'in1'): Array(0.+0.j, dtype=complex128),
('out1', 'out0'): Array(0.+0.j, dtype=complex128),
('out1', 'out1'): Array(0.+0.j, dtype=complex128)}
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.
There is also some circuit metadata information returned as a second return value of the sax.circuit call:
CircuitInfo(dag=<networkx.classes.digraph.DiGraph object at 0x7fc95aa761b0>, models={'waveguide': <function waveguide at 0x7fc95a8e2520>, 'coupler': <function coupler at 0x7fc9b82e7880>, 'top_level': <function _flat_circuit.<locals>._circuit at 0x7fc954280180>}, backend='klu')
Let's now do a simulation for the MZI we just constructed:
CPU times: user 13.9 ms, sys: 989 μs, total: 14.8 ms
Wall time: 30.3 ms
{('in0', 'in0'): Array(0.+0.j, dtype=complex128),
('in0', 'in1'): Array(0.+0.j, dtype=complex128),
('in0', 'out0'): Array(0.+0.j, dtype=complex128),
('in0', 'out1'): Array(-0.57126822+0.82076344j, dtype=complex128),
('in1', 'in0'): Array(0.+0.j, dtype=complex128),
('in1', 'in1'): Array(0.+0.j, dtype=complex128),
('in1', 'out0'): Array(-0.57126822+0.82076344j, dtype=complex128),
('in1', 'out1'): Array(0.+0.j, dtype=complex128),
('out0', 'in0'): Array(0.+0.j, dtype=complex128),
('out0', 'in1'): Array(-0.57126822+0.82076344j, dtype=complex128),
('out0', 'out0'): Array(0.+0.j, dtype=complex128),
('out0', 'out1'): Array(0.+0.j, dtype=complex128),
('out1', 'in0'): Array(-0.57126822+0.82076344j, dtype=complex128),
('out1', 'in1'): Array(0.+0.j, dtype=complex128),
('out1', 'out0'): Array(0.+0.j, dtype=complex128),
('out1', 'out1'): Array(0.+0.j, dtype=complex128)}
CPU times: user 142 ms, sys: 8.09 ms, total: 150 ms
Wall time: 127 ms
{('in0', 'in0'): Array(0.+0.j, dtype=complex128),
('in0', 'in1'): Array(0.+0.j, dtype=complex128),
('in0', 'out0'): Array(0.+0.j, dtype=complex128),
('in0', 'out1'): Array(-0.57126822+0.82076344j, dtype=complex128),
('in1', 'in0'): Array(0.+0.j, dtype=complex128),
('in1', 'in1'): Array(0.+0.j, dtype=complex128),
('in1', 'out0'): Array(-0.57126822+0.82076344j, dtype=complex128),
('in1', 'out1'): Array(0.+0.j, dtype=complex128),
('out0', 'in0'): Array(0.+0.j, dtype=complex128),
('out0', 'in1'): Array(-0.57126822+0.82076344j, dtype=complex128),
('out0', 'out0'): Array(0.+0.j, dtype=complex128),
('out0', 'out1'): Array(0.+0.j, dtype=complex128),
('out1', 'in0'): Array(-0.57126822+0.82076344j, dtype=complex128),
('out1', 'in1'): Array(0.+0.j, dtype=complex128),
('out1', 'out0'): Array(0.+0.j, dtype=complex128),
('out1', 'out1'): Array(0.+0.j, dtype=complex128)}
CPU times: user 436 μs, sys: 36 μs, total: 472 μs
Wall time: 456 μs
{('in0', 'in0'): Array(0.+0.j, dtype=complex128),
('in0', 'in1'): Array(0.+0.j, dtype=complex128),
('in0', 'out0'): Array(0.+0.j, dtype=complex128),
('in0', 'out1'): Array(-0.57126822+0.82076344j, dtype=complex128),
('in1', 'in0'): Array(0.+0.j, dtype=complex128),
('in1', 'in1'): Array(0.+0.j, dtype=complex128),
('in1', 'out0'): Array(-0.57126822+0.82076344j, dtype=complex128),
('in1', 'out1'): Array(0.+0.j, dtype=complex128),
('out0', 'in0'): Array(0.+0.j, dtype=complex128),
('out0', 'in1'): Array(-0.57126822+0.82076344j, dtype=complex128),
('out0', 'out0'): Array(0.+0.j, dtype=complex128),
('out0', 'out1'): Array(0.+0.j, dtype=complex128),
('out1', 'in0'): Array(-0.57126822+0.82076344j, dtype=complex128),
('out1', 'in1'): Array(0.+0.j, dtype=complex128),
('out1', 'out0'): Array(0.+0.j, dtype=complex128),
('out1', 'out1'): Array(0.+0.j, dtype=complex128)}
Or in the case we want an MZI with different arm lengths:
{('in0', 'in0'): Array(0.+0.j, dtype=complex128),
('in0', 'in1'): Array(0.+0.j, dtype=complex128),
('in0', 'out0'): Array(-0.28072841+0.10397039j, dtype=complex128),
('in0', 'out1'): Array(0.89474612-0.33137758j, dtype=complex128),
('in1', 'in0'): Array(0.+0.j, dtype=complex128),
('in1', 'in1'): Array(0.+0.j, dtype=complex128),
('in1', 'out0'): Array(0.89474612-0.33137758j, dtype=complex128),
('in1', 'out1'): Array(0.28072841-0.10397039j, dtype=complex128),
('out0', 'in0'): Array(-0.28072841+0.10397039j, dtype=complex128),
('out0', 'in1'): Array(0.89474612-0.33137758j, dtype=complex128),
('out0', 'out0'): Array(0.+0.j, dtype=complex128),
('out0', 'out1'): Array(0.+0.j, dtype=complex128),
('out1', 'in0'): Array(0.89474612-0.33137758j, dtype=complex128),
('out1', 'in1'): Array(0.28072841-0.10397039j, dtype=complex128),
('out1', 'out0'): Array(0.+0.j, dtype=complex128),
('out1', 'out1'): Array(0.+0.j, dtype=complex128)}
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:
Let's see what this gives:
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()
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:
@jax.jit
def loss_fn(delta_length):
S = mzi(wl=1.55, top={"length": 15.0 + delta_length}, btm={"length": 15.0})
return (abs(S["in0", "out0"]) ** 2).mean()
CPU times: user 140 ms, sys: 7.23 ms, total: 147 ms
Wall time: 135 ms
Array(0.08961828, dtype=float64)
We can use this loss function to define a grad function which works on the parameters of the loss function:
grad_fn = jax.jit(
jax.grad(
loss_fn,
argnums=0, # JAX gradient function for the first positional argument, jitted
)
)
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.
initial_delta_length = 10.0
init_fn, update_fn, params_fn = opt.adam(step_size=0.1)
state = init_fn(initial_delta_length)
Given all this, a single training step can be defined:
def step_fn(step, state):
settings = params_fn(state)
loss = loss_fn(settings)
grad = grad_fn(settings)
state = update_fn(step, grad, state)
return loss, state
And we can use this step function to start the training of the MZI:
range_ = trange(300)
for step in range_:
loss, state = step_fn(step, state)
range_.set_postfix(loss=f"{loss:.6f}")
0%| | 0/300 [00:00<?, ?it/s]
Array(9.93589744, dtype=float64)
Let's see what we've got over a range of wavelengths:
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()
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:
dc_with_arms, info = sax.circuit(
netlist={
"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",
},
},
models={
"coupler": coupler,
"waveguide": waveguide,
},
)
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:
def mzi_chain(num_mzis=1) -> sax.Model:
chain, _ = sax.circuit(
netlist={
"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": "dc0,in0",
"in1": "dc0,in1",
"out0": f"dc{num_mzis},out0",
"out1": f"dc{num_mzis},out1",
},
},
models={"dc_with_arms": dc_with_arms},
backend="klu",
return_type="SDense",
)
return sax.sdict(chain)
Let's for example create a chain with 15 MZIs. We can also update the settings dictionary as follows:
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))
We can simulate this:
%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
CPU times: user 1.02 s, sys: 45 ms, total: 1.06 s
Wall time: 1.45 s
CPU times: user 4.08 s, sys: 24.1 ms, total: 4.11 s
Wall time: 3.09 s
CPU times: user 522 μs, sys: 0 ns, total: 522 μs
Wall time: 530 μs
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 seconds (on a CPU). However, after the MZI chain has been jitted, the evaluation is in the order of about a few milliseconds!
Anyway, let's see what this gives:
%%time
S = sax.sdict(S)
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()
CPU times: user 139 ms, sys: 3.11 ms, total: 142 ms
Wall time: 237 ms