SAX
Project description
SAX
Autograd and XLA for S-parameters - a scatter parameter circuit simulator and optimizer for the frequency domain based on JAX.
The simulator was developed for simulating Photonic Integrated Circuits but in fact is able to perform any S-parameter based circuit simulation. The goal of SAX is to be a thin wrapper around JAX with some basic tools for S-parameter based circuit simulation and optimization. Therefore, SAX does not define any special datastructures and tries to stay as close as possible to the functional nature of JAX. This makes it very easy to get started with SAX as you only need functions and standard python dictionaries. Let's dive in...
Quick Start
Full Quick Start page - Examples - Full Docs.
Let's first import the SAX library, along with JAX and the JAX-version of numpy:
import sax
import jax
import jax.numpy as jnp
Define a model for your component. Which is a decorated function that returns an S-matrix dictionary. For example a directional coupler:
@sax.model(params={"coupling": 0.5})
def coupler(params):
kappa = params["coupling"]**0.5
tau = (1-params["coupling"])**0.5
sdict = sax.reciprocal({
("in0", "out0"): tau,
("in0", "out1"): 1j*kappa,
("in1", "out0"): 1j*kappa,
("in1", "out1"): tau,
})
return sdict
Or a waveguide:
@sax.model({"wl":1.55, "length":100.0, "neff":2.34, "ng":3.4, "wl0":1.55, "loss":0.0})
def waveguide(params):
dwl = params["wl"] - params["wl0"]
dneff_dwl = (params["ng"] - params["neff"]) / params["wl0"]
neff = params["neff"] - dwl * dneff_dwl
phase = 2 * jnp.pi * neff * params["length"] / params["wl"]
transmission = 10 ** (-params["loss"] * params["length"] / 20) * jnp.exp(1j * phase)
sdict = {
("in0", "out0"): transmission,
("out0", "in0"): transmission,
}
return sdict
These component models can then be combined into a circuit:
mzi_func, mzi_params = sax.circuit(
instances = {
"lft": coupler,
"top": waveguide,
"rgt": coupler,
},
connections={
"lft:out0": "rgt:in0",
"lft:out1": "top:in0",
"top:out0": "rgt:in1",
},
ports={
"lft:in0": "in0",
"lft:in1": "in1",
"rgt:out0": "out0",
"rgt:out1": "out1",
},
)
Simulating this is as simple as calling the mzi function with the correct parameters:
params = sax.copy_params(mzi_params)
params["top"]["length"] = 10e-5
S = mzi_func(params)
S["in0", "out0"]
DeviceArray(-0.280701+0.10398856j, dtype=complex64)
Those are the basics. For more info, check out the full SAX Quick Start page, the Examples or the Documentation.
Installation
Dependencies
- JAX & JAXLIB. Please read the JAX install
instructions here. Alternatively, you can
try running jaxinstall.sh to automatically pip-install the correct
jaxandjaxlibpackage for your python and cuda version (if that exact combination exists).
Installation
pip install sax
License
Copyright © 2021, Floris Laporte, Apache-2.0 License
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
