SAX#
S + Autograd + XLA
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 - Documentation.
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 function for your component. A SAX model is just a function that returns an ‘S-dictionary’. For example a directional coupler:
def coupler(coupling=0.5):
kappa = coupling**0.5
tau = (1-coupling)**0.5
sdict = sax.reciprocal({
("in0", "out0"): tau,
("in0", "out1"): 1j*kappa,
("in1", "out0"): 1j*kappa,
("in1", "out1"): tau,
})
return sdict
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}
Or a waveguide:
def waveguide(wl=1.55, wl0=1.55, neff=2.34, ng=3.4, length=10.0, loss=0.0):
dwl = wl - wl0
dneff_dwl = (ng - neff) / wl0
neff = neff - dwl * dneff_dwl
phase = 2 * jnp.pi * neff * length / wl
amplitude = jnp.asarray(10 ** (-loss * length / 20), dtype=complex)
transmission = amplitude * jnp.exp(1j * phase)
sdict = sax.reciprocal({("in0", "out0"): transmission})
return sdict
waveguide(length=100.0)
{('in0', 'out0'): 0.97953-0.2013j, ('out0', 'in0'): 0.97953-0.2013j}
These component models can then be combined into a circuit:
mzi, _ = sax.circuit(
netlist={
"instances": {
"lft": coupler,
"top": waveguide,
"rgt": coupler,
},
"connections": {
"lft,out0": "rgt,in0",
"lft,out1": "top,in0",
"top,out0": "rgt,in1",
},
"ports": {
"in0": "lft,in0",
"in1": "lft,in1",
"out0": "rgt,out0",
"out1": "rgt,out1",
},
}
)
type(mzi)
function
As you can see, the mzi we just created is just another component model function! To simulate it, call the mzi function with the (possibly nested) settings of its subcomponents. Global settings can be added to the ‘root’ of the circuit call and will be distributed over all subcomponents which have a parameter with the same name (e.g. ‘wl’):
wl = jnp.linspace(1.53, 1.57, 1000)
result = mzi(wl=wl, lft={'coupling': 0.3}, top={'length': 200.0}, rgt={'coupling': 0.8})
plt.plot(1e3*wl, jnp.abs(result['in0', 'out0'])**2, label="in0->out0")
plt.plot(1e3*wl, jnp.abs(result['in0', 'out1'])**2, label="in0->out1", ls="--")
plt.xlabel("λ [nm]")
plt.ylabel("T")
plt.grid(True)
plt.figlegend(ncol=2, loc="upper center")
plt.show()
Those are the basics. For more info, check out the full SAX Quick Start page or the rest of the Documentation.
Installation#
You can install SAX with pip:
pip install sax
If you want to be able to run all the example notebooks, you’ll need python>=3.10 and you should install the development version of SAX:
pip install 'sax[dev]'
License#
Copyright © 2023, Floris Laporte, Apache-2.0 License
Home
Changelog#
0.14.0#
Added#
Backend information to
CircuitInfo
.
Updated#
klujax
package.Documentation.
0.13.5#
Added#
settings
field to netlist.
Updated#
GitHub Actions
publish.yml
workflow.Dependencies.
0.13.4#
Added#
forward-only
backend (#39).Example to demonstrate usage of forward-only backend.
Bug fixes related to forward-only backend.
Fixed#
CI configuration.
Pinned minimum version of
numpy
in dependencies.
0.13.3#
Added#
Validation for the number of ports.
Fixed#
Issues in
variability
notebook.Layout-aware notebook.
0.13.2#
Added#
extra_netlist_cleanup
argument.
Updated#
Docker image.
Makefile.
0.13.1#
Updated#
Makefile.
Dockerfile.
0.13.0#
Added#
Documentation and examples.
Coercion of nets into connections.
Handle for callable instances in connections.
Tests for cleanup and refactoring.
Deprecated#
info
subdirectory in instance dictionary.sax.nn
module.
Updated#
Proper transition to
pydantic v2
.Added docstrings across functions (#34).
0.12.2#
Updated#
Pinned minimum version of
flax
.
0.12.1#
Fixed#
Minor bug fixes.
0.12.0#
Added#
Config setting to override component in
info['model']
.
0.11.4#
Updated#
Bumped
klujax
.
0.11.3#
Improved#
Flatten netlist functionality.
0.11.2#
Fixed#
Ignored ports not following the correct format in
flatten_netlist
.
0.11.1#
Updated#
Notebooks.
Bug fixes from 0.11 release.
0.11.0#
Added#
flatten_netlist
function.
Improved#
Error messages.
Analysis of dummy models to prevent dense representation during evaluation.
0.10.4#
Added#
Handling of complex function arguments (#28).
Support for omitting placements in recursive netlist (#27).
Fixed#
Issues related to
'$'
sign in netlist component values.
0.10.3#
Added#
Changelog (#23).
Fixed#
Changelog formatting.
Updated
MANIFEST.in
.
0.10.2#
Fixed#
Evaluation of ports rather than tracing in models.
0.10.1#
Added#
bump2version
as a development dependency.Test folder exposure as artifact.
Fixed#
CI workflows, including Sphinx and Binderhub configurations.
0.10.0#
Added#
Type annotations and expanded documentation.
Jupyter-book based documentation.
Updated#
Notebooks.
Improved multimode operations.
0.9.2#
Added#
Helper functions for component-to-instance mapping.
0.9.0#
Fixed#
Compatibility with both
pydantic
v1 and v2.
0.8.8#
Fixed#
Validation in
get_required_circuit_models
function.
Added#
get_required_circuit_models
function.
0.8.6#
Updated#
GitHub workflows for pypi publishing.
0.8.5#
Fixed#
Various minor cleanups in notebooks.
0.8.4#
Fixed#
Consolidation in multimode to singlemode transition.
0.8.3#
Added#
Support for circuit dependency DAGs.
0.8.2#
Fixed#
Handling of
None
settings in instances.
0.8.1#
Updated#
Notebooks and examples with new API.
Added#
Support for dict-based circuits.
0.8.0#
Added#
Full type support for SAX.
New documentation structure.
0.7.3#
Updated#
Layout-aware notebook title.
Project logo.
0.7.2#
Fixed#
circuit_from_gdsfactory
function.
0.7.1#
Updated#
Minimal version of SAX availability without JAX.
0.6.0 - 0.7.0#
Initial commits with foundational code, refactoring, and early versioning.
API Docs
- sax package
- Subpackages
- Submodules
- sax.circuit module
- sax.loss module
- sax.models module
- sax.multimode module
- sax.netlist module
- sax.saxtypes module
IntArray1D
FloatArray1D
ComplexArray1D
IntArrayND
FloatArrayND
ComplexArrayND
PortMap
PortCombination
SDict
SDense
SCoo
Settings
SType
Model
ModelFactory
is_float()
is_complex()
is_complex_float()
is_sdict()
is_scoo()
is_sdense()
is_model()
is_model_factory()
validate_model()
is_stype()
is_singlemode()
is_multimode()
is_mixedmode()
sdict()
scoo()
sdense()
modelfactory()
- sax.utils module
block_diag()
clean_string()
copy_settings()
validate_settings()
try_complex_float()
flatten_dict()
unflatten_dict()
get_ports()
get_port_combinations()
get_settings()
grouped_interp()
merge_dicts()
mode_combinations()
reciprocal()
rename_params()
rename_ports()
update_settings()
validate_not_mixedmode()
validate_multimode()
validate_sdict()
get_inputs_outputs()
hash_dict()
Normalization
normalization()
cartesian_product()
normalize()
denormalize()
- Module contents