Hierarchical circuits

Let's discuss hierarchical circuits

import jax.numpy as jnp
import matplotlib.pyplot as plt

import sax

Models

create a dictionary of models to be used

models = {
    "coupler": sax.models.coupler_ideal,
    "waveguide": sax.models.straight,
}

Flat Circuit

Probably best to start from a reference circuit. Let's build a flat MZI netlist (introduced in the SAX Quick Start):

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",
    },
}

we can easily simulate this netlist as we did before:

# created the circuit function
mzi, _ = sax.circuit(netlist=netlist, models=models)

# simulate
wl = jnp.linspace(1.5, 1.6)
result = mzi(wl=wl, top={"length": 20})

# plot
plt.plot(wl, abs(result["in0", "out0"]) ** 2)
plt.show()

Hierarchical Circuit

We can quite easily convert this into a hierarchical netlist:

hierarchical_netlist = {
    "top_level": {
        "instances": {
            "top_lft": "top_lft",
            "btm_rgt": "btm_rgt",
        },
        "connections": {
            "top_lft,out0": "btm_rgt,in0",
            "top_lft,out1": "btm_rgt,in1",
        },
        "ports": {
            "in0": "top_lft,in0",
            "in1": "top_lft,in1",
            "out0": "btm_rgt,out0",
            "out1": "btm_rgt,out1",
        },
    },
    "top_lft": {
        "instances": {
            "lft": "coupler",
            "top": "waveguide",
        },
        "connections": {
            "lft,out1": "top,in0",
        },
        "ports": {
            "in0": "lft,in0",
            "in1": "lft,in1",
            "out0": "lft,out0",
            "out1": "top,out0",
        },
    },
    "btm_rgt": {
        "instances": {
            "btm": "waveguide",
            "rgt": "coupler",
        },
        "connections": {
            "btm,out0": "rgt,in0",
        },
        "ports": {
            "in0": "btm,in0",
            "in1": "rgt,in1",
            "out0": "rgt,out0",
            "out1": "rgt,out1",
        },
    },
}

and simulate it just as before

# created the circuit function
mzi, _ = sax.circuit(netlist=hierarchical_netlist, models=models)

# simulate
wl = jnp.linspace(1.5, 1.6)
result = mzi(wl=wl, top_lft={"top": {"length": 20}})

# plot
plt.plot(wl, abs(result["in0", "out0"]) ** 2)
plt.show()