How to use the thewalrus.symplectic.expand function in thewalrus

"""

    (n, m) = S.shape
    if n != m:
        raise ValueError("The matrix S is not square")
    if n % 2 != 0:
        raise ValueError("The matrix S is not of even size")

    nmodes = n // 2

    ch = np.cosh(r) * np.identity(nmodes)
    sh = np.sinh(r) * np.identity(nmodes)
    zh = np.zeros([nmodes, nmodes])
    Schoi = np.block([[ch, sh, zh, zh], [sh, ch, zh, zh], [zh, zh, ch, -sh], [zh, zh, -sh, ch]])

    return expand(S, list(range(nmodes)), n) @ Schoi

if not is_symplectic(S):
            raise ValueError("The matrix S is not symplectic")

    # And that S and alpha have compatible dimensions
    m, _ = S.shape
    if m // 2 != len(alpha):
        raise ValueError("The matrix S and the vector alpha do not have compatible dimensions")

    # Construct the covariance matrix of l two-mode squeezed vacua pairing modes i and i+l
    l = m // 2
    ch = np.cosh(choi_r) * np.identity(l)
    sh = np.sinh(choi_r) * np.identity(l)
    zh = np.zeros([l, l])
    Schoi = np.block([[ch, sh, zh, zh], [sh, ch, zh, zh], [zh, zh, ch, -sh], [zh, zh, -sh, ch]])
    # And then its Choi expanded symplectic
    S_exp = expand(S, list(range(l)), 2 * l) @ Schoi
    # And this is the corresponding covariance matrix
    cov = S_exp @ S_exp.T
    alphat = np.array(list(alpha) + ([0] * l))
    x = 2 * alphat.real
    p = 2 * alphat.imag
    mu = np.concatenate([x, p])

    tensor = state_vector(
        mu, cov, normalize=False, cutoff=cutoff, hbar=2, check_purity=False, choi_r=choi_r
    )

    if sf_order:
        sf_indexing = tuple(chain.from_iterable([[i, i + l] for i in range(l)]))
        return tensor.transpose(sf_indexing)

    return tensor

elif name == "S2gate":
                    S = expand(
                        two_mode_squeezing(params[0], params[1]),
                        [dict_indices[modes[0]], dict_indices[modes[1]]],
                        nmodes,
                    )
                elif name == "Interferometer":
                    S = expand(
                        interferometer(params[0]), [dict_indices[mode] for mode in modes], nmodes
                    )
                elif name == "GaussianTransform":
                    S = expand(
                        params[0], [dict_indices[mode] for mode in modes], nmodes
                    )
                elif name == "BSgate":
                    S = expand(beam_splitter(params[0], params[1]), [dict_indices[modes[0]], dict_indices[modes[1]]], nmodes)
                Snet = S @ Snet
                rnet = S @ rnet

        # Having obtained the net displacement we simply convert it into complex notation
        alphas = 0.5 * (rnet[0:nmodes] + 1j * rnet[nmodes : 2 * nmodes])
        # And now we just pass the net transformation as a big Symplectic operation plus displacements
        ord_reg = [r for r in list(registers) if r.ind in used_modes]
        ord_reg = sorted(list(ord_reg), key=lambda x: x.ind)
        if np.allclose(Snet, np.identity(2 * nmodes)):
            A = []
        else:
            A = [Command(ops.GaussianTransform(Snet), ord_reg)]
        B = [
            Command(ops.Dgate(alphas[i]), ord_reg[i])
            for i in range(len(ord_reg))
            if not np.allclose(alphas[i], 0.0)

# Now we will go through each operation in the sequence `seq` and apply it in quadrature space
        # We will keep track of the net transforation in the Symplectic matrix `Snet` and the quadrature
        # vector `rnet`.
        for operations in seq:
            name = operations.op.__class__.__name__
            params = par_evaluate(operations.op.p)
            modes = [modes_label.ind for modes_label in operations.reg]
            if name == "Dgate":
                rnet = rnet + expand_vector(
                    params[0] * (np.exp(1j * params[1])), dict_indices[modes[0]], nmodes
                )
            else:
                if name == "Rgate":
                    S = expand(rotation(params[0]), dict_indices[modes[0]], nmodes)
                elif name == "Sgate":
                    S = expand(squeezing(params[0], params[1]), dict_indices[modes[0]], nmodes)
                elif name == "S2gate":
                    S = expand(
                        two_mode_squeezing(params[0], params[1]),
                        [dict_indices[modes[0]], dict_indices[modes[1]]],
                        nmodes,
                    )
                elif name == "Interferometer":
                    S = expand(
                        interferometer(params[0]), [dict_indices[mode] for mode in modes], nmodes
                    )
                elif name == "GaussianTransform":
                    S = expand(
                        params[0], [dict_indices[mode] for mode in modes], nmodes
                    )
                elif name == "BSgate":
                    S = expand(beam_splitter(params[0], params[1]), [dict_indices[modes[0]], dict_indices[modes[1]]], nmodes)

rnet = np.zeros(2 * nmodes)

        # Now we will go through each operation in the sequence `seq` and apply it in quadrature space
        # We will keep track of the net transforation in the Symplectic matrix `Snet` and the quadrature
        # vector `rnet`.
        for operations in seq:
            name = operations.op.__class__.__name__
            params = par_evaluate(operations.op.p)
            modes = [modes_label.ind for modes_label in operations.reg]
            if name == "Dgate":
                rnet = rnet + expand_vector(
                    params[0] * (np.exp(1j * params[1])), dict_indices[modes[0]], nmodes
                )
            else:
                if name == "Rgate":
                    S = expand(rotation(params[0]), dict_indices[modes[0]], nmodes)
                elif name == "Sgate":
                    S = expand(squeezing(params[0], params[1]), dict_indices[modes[0]], nmodes)
                elif name == "S2gate":
                    S = expand(
                        two_mode_squeezing(params[0], params[1]),
                        [dict_indices[modes[0]], dict_indices[modes[1]]],
                        nmodes,
                    )
                elif name == "Interferometer":
                    S = expand(
                        interferometer(params[0]), [dict_indices[mode] for mode in modes], nmodes
                    )
                elif name == "GaussianTransform":
                    S = expand(
                        params[0], [dict_indices[mode] for mode in modes], nmodes
                    )

if name == "Rgate":
                    S = expand(rotation(params[0]), dict_indices[modes[0]], nmodes)
                elif name == "Sgate":
                    S = expand(squeezing(params[0], params[1]), dict_indices[modes[0]], nmodes)
                elif name == "S2gate":
                    S = expand(
                        two_mode_squeezing(params[0], params[1]),
                        [dict_indices[modes[0]], dict_indices[modes[1]]],
                        nmodes,
                    )
                elif name == "Interferometer":
                    S = expand(
                        interferometer(params[0]), [dict_indices[mode] for mode in modes], nmodes
                    )
                elif name == "GaussianTransform":
                    S = expand(
                        params[0], [dict_indices[mode] for mode in modes], nmodes
                    )
                elif name == "BSgate":
                    S = expand(beam_splitter(params[0], params[1]), [dict_indices[modes[0]], dict_indices[modes[1]]], nmodes)
                Snet = S @ Snet
                rnet = S @ rnet

        # Having obtained the net displacement we simply convert it into complex notation
        alphas = 0.5 * (rnet[0:nmodes] + 1j * rnet[nmodes : 2 * nmodes])
        # And now we just pass the net transformation as a big Symplectic operation plus displacements
        ord_reg = [r for r in list(registers) if r.ind in used_modes]
        ord_reg = sorted(list(ord_reg), key=lambda x: x.ind)
        if np.allclose(Snet, np.identity(2 * nmodes)):
            A = []
        else:
            A = [Command(ops.GaussianTransform(Snet), ord_reg)]

# vector `rnet`.
        for operations in seq:
            name = operations.op.__class__.__name__
            params = par_evaluate(operations.op.p)
            modes = [modes_label.ind for modes_label in operations.reg]
            if name == "Dgate":
                rnet = rnet + expand_vector(
                    params[0] * (np.exp(1j * params[1])), dict_indices[modes[0]], nmodes
                )
            else:
                if name == "Rgate":
                    S = expand(rotation(params[0]), dict_indices[modes[0]], nmodes)
                elif name == "Sgate":
                    S = expand(squeezing(params[0], params[1]), dict_indices[modes[0]], nmodes)
                elif name == "S2gate":
                    S = expand(
                        two_mode_squeezing(params[0], params[1]),
                        [dict_indices[modes[0]], dict_indices[modes[1]]],
                        nmodes,
                    )
                elif name == "Interferometer":
                    S = expand(
                        interferometer(params[0]), [dict_indices[mode] for mode in modes], nmodes
                    )
                elif name == "GaussianTransform":
                    S = expand(
                        params[0], [dict_indices[mode] for mode in modes], nmodes
                    )
                elif name == "BSgate":
                    S = expand(beam_splitter(params[0], params[1]), [dict_indices[modes[0]], dict_indices[modes[1]]], nmodes)
                Snet = S @ Snet
                rnet = S @ rnet

rnet = rnet + expand_vector(
                    params[0] * (np.exp(1j * params[1])), dict_indices[modes[0]], nmodes
                )
            else:
                if name == "Rgate":
                    S = expand(rotation(params[0]), dict_indices[modes[0]], nmodes)
                elif name == "Sgate":
                    S = expand(squeezing(params[0], params[1]), dict_indices[modes[0]], nmodes)
                elif name == "S2gate":
                    S = expand(
                        two_mode_squeezing(params[0], params[1]),
                        [dict_indices[modes[0]], dict_indices[modes[1]]],
                        nmodes,
                    )
                elif name == "Interferometer":
                    S = expand(
                        interferometer(params[0]), [dict_indices[mode] for mode in modes], nmodes
                    )
                elif name == "GaussianTransform":
                    S = expand(
                        params[0], [dict_indices[mode] for mode in modes], nmodes
                    )
                elif name == "BSgate":
                    S = expand(beam_splitter(params[0], params[1]), [dict_indices[modes[0]], dict_indices[modes[1]]], nmodes)
                Snet = S @ Snet
                rnet = S @ rnet

        # Having obtained the net displacement we simply convert it into complex notation
        alphas = 0.5 * (rnet[0:nmodes] + 1j * rnet[nmodes : 2 * nmodes])
        # And now we just pass the net transformation as a big Symplectic operation plus displacements
        ord_reg = [r for r in list(registers) if r.ind in used_modes]
        ord_reg = sorted(list(ord_reg), key=lambda x: x.ind)