sax

What it is and what it does

SAX wraps JAX to provide a functional, dictionary-based framework for simulating and optimizing circuits described by S-parameters (scattering matrices). It was built for photonic integrated circuits but works with any frequency-domain circuit model. Rather than defining custom data structures, SAX stays close to JAX's functional paradigm: you write component models as functions returning S-dictionaries, compose them into circuits via netlists, and optimize them using JAX's autodiff and XLA compilation.

The package handles the boilerplate of circuit composition and parameter passing, allowing you to define a directional coupler, waveguide, or other component once and then combine them into larger systems like Mach-Zehnder interferometers. You can then sweep parameters (wavelength, coupling strength, length) across ranges and plot transmission spectra or use JAX's optimization tools to tune component parameters for a target response.

Use it for:

• Design and simulate photonic integrated circuits by composing waveguides, couplers, and other optical components into larger systems.
• Optimize photonic circuit parameters (coupling ratios, waveguide lengths, phase shifts) to achieve target transmission or reflection spectra.
• Perform frequency-domain S-parameter analysis on arbitrary circuits without writing custom matrix algebra.
• Leverage JAX autodiff to compute gradients for circuit optimization or sensitivity analysis across wavelength ranges.
• Combine circuit simulation with machine learning workflows via JAX's ecosystem (Optax, Flax, etc.).

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

SAX is a JAX-based simulator and optimizer for S-parameter circuits, primarily designed for photonic integrated circuits but applicable to any frequency-domain circuit simulation using scattering parameters.

Yes, if you work with S-parameter circuit simulation or photonic design. SAX is actively maintained, has no known vulnerabilities, and offers a clean functional interface built on JAX's autodiff and XLA compilation. The dependency footprint is large but manageable. Install with caution if your environment has strict dependency constraints; otherwise, it is a solid choice for frequency-domain circuit work.

Install

sax on PyPI

pip install sax

uv add sax

poetry add sax

Installing sax

Before you install

Low install friction with a pure-wheel distribution. Active maintenance (66 days since last release) and support for current Python versions (3.11–3.14). Requires 19 runtime dependencies including JAX, NumPy, and scientific libraries; installation is straightforward but the dependency stack is substantial.

License in practice

Apache-2.0 permissive license allows commercial and private use with minimal restrictions; you must include a copy of the license and state significant changes, but there are no copyleft obligations.

Quickstart

pip install sax

import sax
import jax.numpy as jnp

def coupler(coupling=0.5):
    kappa = coupling**0.5
    tau = (1-coupling)**0.5
    return sax.reciprocal({
        ("in0", "out0"): tau,
        ("in0", "out1"): 1j*kappa,
        ("in1", "out0"): 1j*kappa,
        ("in1", "out1"): tau,
    })

result = coupler(coupling=0.3)

Requires Python >=3.11.0; JAX installation may require additional system dependencies depending on your platform (CPU vs. GPU/TPU support).

Verify before relying

• Whether the package's 19 dependencies can be installed together without version conflicts in typical environments.
• Performance characteristics and scalability limits for large circuit topologies or high-dimensional parameter sweeps.

Package facts

Evidence: sax-0.18.2-py3-none-any.whl

Keywords: simulation, optimization, autograd, simulation-framework, circuit, physics-simulation, photonics, s-parameters, jax, xla, photonic-circuit, photonic-optimization

Tags

More Scientific/Engineering packages