Sparse COO

The most efficient S-representation is the COO representation. It’s however also the least user friendly:

def my_coo(*, wl=1.5):
    num_input_ports = 4
    num_output_ports = 1
    pm = {
        **{f"in{i}": i for i in range(num_input_ports)},
        **{f"out{i}": i + num_input_ports for i in range(num_output_ports)},
    }
    thru = jnp.ones(num_input_ports)
    i = jnp.arange(0, num_input_ports, 1)
    j = jnp.zeros_like(i) + num_input_ports

    # make reciprocal
    i, j = jnp.concatenate([i, j]), jnp.concatenate([j, i])
    thru = jnp.concatenate([thru, thru], 0)
    return (i, j, thru, pm)

circuit, _ = sax.circuit(
    netlist={
        "instances": {
            "coo": "coo",
        },
        "connections": {},
        "ports": {
            "in0": "coo,in0",
            "out0": "coo,out0",
        },
    },
    models={
        "coo": my_coo,
    },
    backend="klu",
    return_type="scoo",
)

circuit()

(Array([0, 0, 1, 1], dtype=int64),
 Array([0, 1, 0, 1], dtype=int64),
 Array([0.+0.j, 1.+0.j, 1.+0.j, 0.+0.j], dtype=complex128),
 {'in0': 0, 'out0': 1})

sax.sdict(circuit())

{('in0', 'in0'): Array(0.+0.j, dtype=complex128),
 ('in0', 'out0'): Array(1.+0.j, dtype=complex128),
 ('out0', 'in0'): Array(1.+0.j, dtype=complex128),
 ('out0', 'out0'): Array(0.+0.j, dtype=complex128)}