Example code doesn't work

[harshvardhan-pandey]

While working with the example sax.ipynb, the following piece of code throws error

y = siepic.y_branch(wl=wl)
plt.plot(1e3*wl, jnp.abs(y['o0', 'o0'])**2, label="in0->out0")
plt.plot(1e3*wl, jnp.abs(y['o0', 'o1'])**2, label="in0->out0")
plt.plot(1e3*wl, jnp.abs(y['o0', 'o2'])**2, label="in0->out0")

KeyError                                  Traceback (most recent call last)
Cell In[22], [line 2]
      1 y = siepic.y_branch(wl=wl)
----> plt.plot(1e3*wl, jnp.abs(y['o0', 'o0'])**2, label="in0->out0")
      3 plt.plot(1e3*wl, jnp.abs(y['o0', 'o1'])**2, label="in0->out0")
      4 plt.plot(1e3*wl, jnp.abs(y['o0', 'o2'])**2, label="in0->out0")

KeyError: ('o0', 'o0')

I changed names 'o0', 'o1', and 'o2' to 'port_0', 'port_1', and 'port_2' respectively, which works fine for this cell, but later in

mzi, info = sax.circuit(
    netlist={
        "instances": {
            "ingc": "gc",
            "lft": "ybranch",
            "top": "wg_long",
            "btm": "wg_short",
            "rgt": "ybranch",
            "outgc": "gc",
        },
        "connections": {
            "ingc,o1": "lft,port_1",
            "lft,o1": "btm,port_1",
            "btm,o1": "rgt,port_2",
            "lft,o2": "top,port_1",
            "top,o1": "rgt,port_3",
            "rgt,o0": "outgc,port_2",
        },
        "ports": {
            "in": "ingc,port_1",
            "out": "outgc,port_1",
        },
    },
    models={
        "ybranch": Y_fixed,
        # "wg_long": lambda wl: WG_fixed(wl=wl, length=150),
        # "wg_short": lambda wl: WG_fixed(wl=wl, length=50),
        # "wg_long": wg_150,
        # "wg_short": wg_50,
        "wg_long": partial(siepic.waveguide, length=150),
        "wg_short": partial(siepic.waveguide, length=50),
        "gc": GC_fixed,
    }
)

wl = jnp.linspace(1.51, 1.59, 2000)

jitted_mzi = jit(mzi)
%timeit S = jitted_mzi(wl=wl)
%timeit S = jitted_mzi(wl=wl)
S = jitted_mzi(wl=wl)

plt.plot(wl * 1e3, abs(S["in", "out"]) ** 2)
plt.xlabel("λ [nm]")
plt.ylabel("T")
# plt.ylim(-0.05, 1.05)
plt.show()

this still shows

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
File ~/lightml-hawai/.venv/lib/python3.12/site-packages/sax/circuit.py:324, in _get_multimode_connections(connections, inst_port_mode, ignore_missing_ports)
    323 try:
--> 324     modes1 = inst_port_mode[inst1][port1]
    325 except KeyError as err:

KeyError: 'o1'

The above exception was the direct cause of the following exception:

RuntimeError                              Traceback (most recent call last)
Cell In[21], line 1
----> 1 mzi, info = sax.circuit(
      2     netlist={
      3         "instances": {
      4             "ingc": "gc",
      5             "lft": "ybranch",
      6             "top": "wg_long",
      7             "btm": "wg_short",
      8             "rgt": "ybranch",
      9             "outgc": "gc",
     10         },
     11         "connections": {
     12             "ingc,o1": "lft,port_1",
     13             "lft,o1": "btm,port_1",
     14             "btm,o1": "rgt,port_2",
     15             "lft,o2": "top,port_1",
     16             "top,o1": "rgt,port_3",
     17             "rgt,o0": "outgc,port_2",
     18         },
     19         "ports": {
     20             "in": "ingc,port_1",
     21             "out": "outgc,port_1",
     22         },
     23     },
     24     models={
     25         "ybranch": Y_fixed,
     26         # "wg_long": lambda wl: WG_fixed(wl=wl, length=150),
     27         # "wg_short": lambda wl: WG_fixed(wl=wl, length=50),
     28         # "wg_long": wg_150,
     29         # "wg_short": wg_50,
     30         "wg_long": partial(siepic.waveguide, length=150),
     31         "wg_short": partial(siepic.waveguide, length=50),
     32         "gc": GC_fixed,
     33     }
     34 )
     36 wl = jnp.linspace(1.51, 1.59, 2000)
     38 jitted_mzi = jit(mzi)

File ~/lightml-hawai/.venv/lib/python3.12/site-packages/sax/circuit.py:84, in circuit(netlist, models, backend, return_type, ignore_missing_ports)
     81     current_models |= new_models
     82     new_models = {}
---> 84     current_models[model_name] = circuit = _flat_circuit(
     85         flatnet.instances,
     86         flatnet.connections,
     87         flatnet.ports,
     88         current_models,
     89         backend,
     90         ignore_missing_ports=ignore_missing_ports,
     91     )
     92 if circuit is None:
     93     msg = "No circuit was created"

File ~/lightml-hawai/.venv/lib/python3.12/site-packages/sax/circuit.py:190, in _flat_circuit(instances, connections, ports, models, backend, ignore_missing_ports)
    186 dummy_instances = analyze_insts_fn(instances, models)
    187 inst_port_mode = {
    188     k: _port_modes_dict(get_ports(s)) for k, s in dummy_instances.items()
    189 }
--> 190 connections = _get_multimode_connections(
    191     connections, inst_port_mode, ignore_missing_ports=ignore_missing_ports
    192 )
    193 ports = _get_multimode_ports(
    194     ports, inst_port_mode, ignore_missing_ports=ignore_missing_ports
    195 )
    197 inst2model = {}

File ~/lightml-hawai/.venv/lib/python3.12/site-packages/sax/circuit.py:332, in _get_multimode_connections(connections, inst_port_mode, ignore_missing_ports)
    327         continue
    328     msg = (
    329         f"Instance {inst1} does not contain port {port1}. "
    330         f"Available ports: {list(inst_port_mode[inst1])}."
    331     )
--> 332     raise RuntimeError(msg) from err
    333 try:
    334     modes2 = inst_port_mode[inst2][port2]

RuntimeError: Instance lft does not contain port o1. Available ports: ['port_1', 'port_2', 'port_3'].

I am using sax version 0.14.7