# \tref{ap:printed_code}

This appendix contains the 'core' model code. Full code for generating the figures in this thesis, and the data used in those figures, can be found in the GitHub repository for this thesis https://github.com/Quantum-Integration-Laboratory/MariaNicolaeHonoursThesis.

## Three-Level Transduction Replication

These codes are my own implementations of the prior three-level transduction models of \ref{ch:prior_transduction}. They were written quite early in the project, and so follow the notation of the original sources rather than the notation in this thesis.

### Single Cavity

This Python code replicates the single-cavity model in Reference \cite{fernandez-gonzalvo_2019}. When run as a script, it replicates Figure 4(e) in that paper.

    \
    \verb|import numpy as np
    from scipy import integrate, optimize, stats
    import matplotlib.pyplot as plt
    import sympy

    def getL():
        # symbols
        s12 = sympy.Matrix([[0, 1, 0], [0, 0, 0], [0, 0, 0]])
        s13 = sympy.Matrix([[0, 0, 1], [0, 0, 0], [0, 0, 0]])
        s23 = sympy.Matrix([[0, 0, 0], [0, 0, 1], [0, 0, 0]])
        s21 = s12.H
        s31 = s13.H
        s32 = s23.H
        s11 = s12*s21
        s22 = s21*s12
        s33 = s31*s13

        delta_mu = sympy.symbols('delta_mu', real=True)
        delta_s = sympy.symbols('delta_s', real=True)
        gamma_mu = sympy.symbols('gamma_mu', real=True)
        gamma31, gamma32 = sympy.symbols('gamma3(1:3)', real=True)
        gamma2d, gamma3d = sympy.symbols('gamma(2:4)d', real=True)
        nbath = sympy.symbols('nbath', real=True)
        Omega_mu = sympy.symbols('Omega_mu', complex=True)
        Omega_o = sympy.symbols('Omega_o', complex=True)
        A = sympy.symbols('A', complex=True)

        def master_equation_rhs(rho):
            H = Omega_o*s32 + Omega_mu*s21 + A*s31
            H = H + H.H
            H = H + delta_mu*s22 + delta_s*s33

            L21 = gamma_mu/2 * (nbath+1) * (2*s12*rho*s21 - s22*rho - rho*s22)
            L12 = gamma_mu/2 * nbath * (2*s21*rho*s12 - s11*rho - rho*s11)
            L32 = gamma32/2 * (2*s23*rho*s32 - s33*rho - rho*s33)
            L31 = gamma31/2 * (2*s13*rho*s31 - s33*rho - rho*s33)
            L22 = gamma2d/2 * (2*s22*rho*s22 - s22*s22*rho - rho*s22*s22)
            L33 = gamma3d/2 * (2*s33*rho*s33 - s33*s33*rho - rho*s33*s33)
            loss = L21 + L12 + L32 + L31 + L22 + L33

            return -sympy.I*(H*rho - rho*H) + loss

        # obtain matrix representation of differential operator L
        L = sympy.zeros(9)
        for i in range(3):
            for j in range(3):
                rho = sympy.zeros(3)
                rho[i,j] = 1
                Lcol = master_equation_rhs(rho)
                col = 3*i+j
                for ip in range(3):
                    for jp in range(3):
                        row = 3*ip+jp
                        L[row,col] = Lcol[ip,jp]

        # replace first row with row computing the trace of rho
        L[0,:] = sympy.Matrix([1, 0, 0, 0, 1, 0, 0, 0, 1]).T

        # lambdify
        args = (delta_mu, delta_s, gamma_mu, gamma31, gamma32, gamma2d, gamma3d, nbath, Omega_mu, Omega_o, A)
        Lfunc = sympy.lambdify(args, L, 'numpy')

        return L, Lfunc

    L, Lfunc = getL()

    # fundamental constants
    hbar = 1.05457e-34
    kB = 1.380649e-23
    c = 299792458
    eps0 = 8.854187817e-12
    mu0 = 4*np.pi*1e-7
    muB = 9.274009994e-24

    def rho_steady_state_many_args(delta_mu, delta_s, gamma_mu, gamma31, gamma32, gamma2d, gamma3d, nbath, Omega_mu, Omega_o, A):
        Lmatrix = Lfunc(
            delta_mu=delta_mu,
            delta_s=delta_s,
            gamma_mu=gamma_mu,
            gamma31=gamma31,
            gamma32=gamma32,
            gamma2d=gamma2d,
            gamma3d=gamma3d,
            nbath=nbath,
            Omega_mu=Omega_mu,
            Omega_o=Omega_o,
            A=A
        )
        b = np.array([1, 0, 0, 0, 0, 0, 0, 0, 0])
        x = np.linalg.solve(Lmatrix, b)
        rho = np.array([[x[0], x[1], x[2]],
                        [x[3], x[4], x[5]],
                        [x[6], x[7], x[8]]])
        return rho

    def Omega_mu_from_Pin(Pin, omega_mu):
        V_mu_cavity = Vsample / fill_factor
        Pmu = 1e-3 * 10**(Pin/10)
        Q = omega_mu / (kappa_mi + 2*kappa_mc)
        S21 = 4 * kappa_mc**2 / (kappa_mi + 2*kappa_mc)**2
        energy_mu_cavity = 2*Pmu*Q*np.sqrt(S21) / omega_mu
        Bmu = np.sqrt(mu0*energy_mu_cavity / (2*V_mu_cavity))
        return -mu12*Bmu / hbar

    def Omega_o_from_Pin(Pin, omega_o):
        Po = 1e-3 * 10**(Pin/10)
        pflux = Po / (hbar*omega_o)
        n_in = 4*pflux*kappa_oc / (kappa_oc+kappa_oi)**2
        Sspot = np.pi * Woc**2
        V_o_cavity = (Sspot*Loc + Sspot*Lsample*nYSO**3) / 2
        Eo = np.sqrt(n_in*hbar*omega_o / (2*eps0*V_o_cavity))
        return -d23*Eo / hbar

    def rho_steady_state(Pin_mu, Pin_o, delta_mu, delta_o, a):
        omega_mu = omega_12 - delta_mu
        nbath = 1 / (np.exp(hbar*(omega_mu)/(kB*T))-1)
        rho = rho_steady_state_many_args(
            delta_mu = delta_mu,
            delta_s = delta_o - delta_mu,
            gamma_mu = 1/tau2 * 1/(nbath+1),
            gamma31 = 1/tau3 * d13**2 / (d13**2 + d23**2),
            gamma32 = 1/tau3 * d23**2 / (d13**2 + d23**2),
            gamma2d = 1e6,
            gamma3d = 1e6,
            nbath = nbath,
            Omega_mu = Omega_mu_from_Pin(Pin_mu, omega_12 - delta_mu),
            Omega_o = Omega_o_from_Pin(Pin_o, omega_23 - delta_o),
            A = g*a
        )
        return rho

    def rho13_steady_state_ensemble(Pin_mu, Pin_o, mean_delta_mu, mean_delta_s, a):
        standard_norm = lambda z: np.exp(-z**2/2) / np.sqrt(2*np.pi)
        zrange = 3

        def ensemble_integrand(z_mu, z_s):
            delta_mu = mean_delta_mu + w12*z_mu
            delta_s = mean_delta_mu + w13*z_s
            envelope = standard_norm(z_mu)*standard_norm(z_s) / (w12*w13)
            rho13 = rho_steady_state(Pin_mu, Pin_o, delta_mu, delta_s, a)[0,2]
            jacobian = w12 * w13
            return envelope * rho13 * jacobian

        real_integrand = lambda z_mu, z_s: np.real(ensemble_integrand(z_mu, z_s))
        imag_integrand = lambda z_mu, z_s: np.imag(ensemble_integrand(z_mu, z_s))

        y_re, abserr_re = integrate.dblquad(real_integrand, -zrange, zrange, -zrange, zrange)
        y_im, abserr_im = integrate.dblquad(imag_integrand, -zrange, zrange, -zrange, zrange)
        y = y_re + 1j*y_im
        return y

    def steady_a(Pin_mu, Pin_o, delta_oc, rescaling=1):
        def S13(a):
            return N*g*rho13_steady_state_ensemble(Pin_mu, Pin_o, 0, 0, a)

        def ffunc(a):
            return -1j*delta_oc*a - 1j*S13(a) - (kappa_oi+kappa_oc)*a/2

        def ffunc_R2toR2(a):
            [a_re, a_im] = a
            a = a_re + 1j*a_im
            fa = ffunc(a)
            fa *= rescaling
            return [np.real(fa), np.imag(fa)]

        result = optimize.root(ffunc_R2toR2, [0, 0])
        [a_re, a_im] = result.x
        return a_re + 1j*a_im, result

    if __name__ == '__main__':
        # recreate Figure 4c from Fernandez-Gonzalvo et. al. 2019
        # (Phys. Rev. A 100, 033807)

        nYSO = 1.76 # refractive index of YSO
        T = 4.6 # experiment temperature
        N = 1.28e15 # erbium number density
        g = 51.9 # s13 to optical coupling

        omega_12 = 2*np.pi*5.186e9
        omega_23 = 2*np.pi*195113.36e9 # 1536.504 nm
        w12 = 2*np.pi*25e6
        w13 = 2*np.pi*170e6
        mu12 = 4.3803*muB
        d13 = 1.63e-32
        d23 = 1.15e-32
        tau2 = 1e-3
        tau3 = 11e-3

        kappa_mi = 2*np.pi*650e3
        kappa_mc = 2*np.pi*70e3
        kappa_oi = 2*np.pi*1.7e6
        kappa_oc = 2*np.pi*7.95e6

        Woc = 0.6e-3
        Loc = 49.5e-3
        fill_factor = 0.8 # fraction of microwave cavity filled by sample

        dsample = 5e-3
        Lsample = 12e-3
        Vsample = np.pi * dsample**2 * Lsample / 4

        Pin_mu = 0
        Pin_o = 10.4135

        a, _ = steady_a(-20, Pin_o, 0, rescaling=1e-6)

        delta_mu_v = 30e6 * np.linspace(-2*np.pi, 2*np.pi, 101) / (2*np.pi)
        delta_o_v = 50e6 * np.linspace(-2*np.pi, 2*np.pi, 101) / (2*np.pi)
        delta_mu, delta_o = np.meshgrid(delta_mu_v, delta_o_v)

        def steady_pop(Pin_mu, Pin_o, delta_mu, delta_o, a):
            rho = rho_steady_state(Pin_mu, Pin_o, delta_mu, delta_o, a)
            return np.real(np.diag(rho))

        sig = '(),(),(),(),()->(n)'
        pop = np.vectorize(steady_pop, signature=sig)(-20, Pin_o, delta_mu, delta_o, a)

        p11 = pop[:,:,0]
        p33 = pop[:,:,2]
        mhz = 1e6

        fig, ax = plt.subplots(1, 1)
        ax.pcolormesh(delta_o/mhz, delta_mu/mhz, p11-p33, vmin=0, vmax=0.05)
        ax.invert_yaxis()
        ax.set_aspect('equal')
        ax.set_title('rho11 - rho33 (Mine)')
        ax.set_xlabel('delta_o (rad MHz)')
        ax.set_ylabel('delta_mu (rad MHz)')
        plt.show()
    |

### Double Cavity

This Python code replicates the double-cavity model in Reference \cite{barnett_longdell_2020}. It is 'library' code that is not a script in its own right. The GitHub repository contains scripts that `import` this code.

    \
    \verb|import numpy as np
    import matplotlib.pyplot as plt
    from scipy import integrate, optimize, stats
    import sympy
    import pickle
    import itertools

    gauss_lobatto_n = 20

    legendre_coeffs = (0,)*(gauss_lobatto_n-1) + (1,)
    roots = np.polynomial.legendre.Legendre(legendre_coeffs).deriv().roots()
    gauss_lobatto_points = np.concatenate([[-1], roots, [1]])
    gauss_lobatto_points = (gauss_lobatto_points+1) / 2

    k = np.arange(gauss_lobatto_n)
    i, j = np.indices((gauss_lobatto_n, gauss_lobatto_n))
    M = gauss_lobatto_points[j]**i
    b = 1/(k+1)
    gauss_lobatto_weights = np.linalg.solve(M, b)

    # flattening and unflattening arrays
    unravel = np.array([
        [0, 3, 6],
        [1, 4, 7],
        [2, 5, 8]
    ])

    ravel_i = np.zeros(9, dtype=int)
    ravel_j = np.zeros(9, dtype=int)
    for i in range(3):
        for j in range(3):
            k = unravel[i,j]
            ravel_i[k] = i
            ravel_j[k] = j
    ravel = (ravel_i, ravel_j)

    def get_symbolic():
        # parameters
        g_o = sympy.symbols('g_o', real=True)
        g_mu = sympy.symbols('g_mu', real=True)
        Omega = sympy.symbols('Omega')
        alpha = sympy.symbols('alpha')
        beta = sympy.symbols('beta')
        delta_mu = sympy.symbols('delta_mu', real=True)
        delta_o = sympy.symbols('delta_o', real=True)
        delta_amu = sympy.symbols('delta_amu', real=True)
        delta_ao = sympy.symbols('delta_ao', real=True)
        gamma_12 = sympy.symbols('gamma_12', real=True)
        gamma_13 = sympy.symbols('gamma_13', real=True)
        gamma_23 = sympy.symbols('gamma_23', real=True)
        gamma_2d = sympy.symbols('gamma_2d', real=True)
        gamma_3d = sympy.symbols('gamma_3d', real=True)
        n_bath = sympy.symbols('n_bath', real=True)
        symbols = {
            'g_o': g_o,
            'g_mu': g_mu,
            'Omega': Omega,
            'alpha': alpha,
            'beta': beta,
            'delta_mu': delta_mu,
            'delta_o': delta_o,
            'delta_amu': delta_amu,
            'delta_ao': delta_ao,
            'gamma_12': gamma_12,
            'gamma_13': gamma_13,
            'gamma_23': gamma_23,
            'gamma_2d': gamma_2d,
            'gamma_3d': gamma_3d,
            'n_bath': n_bath
        }

        # matrix symbols
        s12 = sympy.Matrix([[0, 1, 0], [0, 0, 0], [0, 0, 0]])
        s13 = sympy.Matrix([[0, 0, 1], [0, 0, 0], [0, 0, 0]])
        s23 = sympy.Matrix([[0, 0, 0], [0, 0, 1], [0, 0, 0]])
        s21 = s12.H
        s31 = s13.H
        s32 = s23.H
        s11 = s12*s21
        s22 = s21*s12
        s33 = s31*s13

        # Hamiltonian
        H_0 = (delta_amu-delta_mu)*s22 + (delta_ao-delta_o)*s33 + Omega*s32 + Omega.conjugate()*s23
        H_alpha = g_o*s31
        H_alpha_c = g_o*s13
        H_beta = g_mu*s21
        H_beta_c = g_mu*s12
        H = H_0 + alpha*H_alpha + alpha.conjugate()*H_alpha_c + beta*H_beta + beta.conjugate()*H_beta_c

        # loss superoperator
        def loss_superoperator(rho):
            L21 = gamma_12/2 * (n_bath+1) * (2*s12*rho*s21 - s22*rho - rho*s22)
            L12 = gamma_12/2 * n_bath * (2*s21*rho*s12 - s11*rho - rho*s11)
            L32 = gamma_23/2 * (2*s23*rho*s32 - s33*rho - rho*s33)
            L31 = gamma_13/2 * (2*s13*rho*s31 - s33*rho - rho*s33)
            L22 = gamma_2d/2 * (2*s22*rho*s22 - s22*rho - rho*s22)
            L33 = gamma_3d/2 * (2*s33*rho*s33 - s33*rho - rho*s33)
            loss = L21 + L12 + L32 + L31 + L22 + L33
            return loss

        # get liouvillan superoperator matrix
        def liouvillan_matrix(H, loss=None):
            def master_equation(rho):
                master_operator = -sympy.I * (H*rho-rho*H)
                if loss is not None:
                    master_operator += loss(rho)
                return master_operator

            L = sympy.zeros(9)
            for column, (icol, jcol) in enumerate(zip(*ravel)):
                # basis matrix
                rho = sympy.zeros(3)
                rho[icol,jcol] = 1

                # get column of supermatrix
                L_column = master_equation(rho)
                for row, (irow, jrow) in enumerate(zip(*ravel)):
                    L[row,column] = L_column[irow,jrow]

            return L

        L_0 = liouvillan_matrix(H_0, loss=loss_superoperator)
        L_alpha = liouvillan_matrix(H_alpha)
        L_alpha_c = liouvillan_matrix(H_alpha_c)
        L_beta = liouvillan_matrix(H_beta)
        L_beta_c = liouvillan_matrix(H_beta_c)
        L = L_0 + alpha*L_alpha + alpha.conjugate()*L_alpha_c + beta*L_beta + beta.conjugate()*L_beta_c

        return {
            'symbols': symbols,
            'H': H,
            'L': L,
            'L_0': L_0,
            'L_alpha': L_alpha,
            'L_alpha_c': L_alpha_c,
            'L_beta': L_beta,
            'L_beta_c': L_beta_c
        }

    symbolic = get_symbolic()

    symbolic_args_full = (
        symbolic['symbols']['g_o'],
        symbolic['symbols']['g_mu'],
        symbolic['symbols']['Omega'],
        symbolic['symbols']['gamma_12'],
        symbolic['symbols']['gamma_13'],
        symbolic['symbols']['gamma_23'],
        symbolic['symbols']['gamma_2d'],
        symbolic['symbols']['gamma_3d'],
        symbolic['symbols']['n_bath'],
        symbolic['symbols']['delta_mu'],
        symbolic['symbols']['delta_o'],
        symbolic['symbols']['delta_amu'],
        symbolic['symbols']['delta_ao'],
        symbolic['symbols']['alpha'],
        symbolic['symbols']['beta']
    )

    symbolic_args_0 = (
        symbolic['symbols']['g_o'],
        symbolic['symbols']['g_mu'],
        symbolic['symbols']['Omega'],
        symbolic['symbols']['gamma_12'],
        symbolic['symbols']['gamma_13'],
        symbolic['symbols']['gamma_23'],
        symbolic['symbols']['gamma_2d'],
        symbolic['symbols']['gamma_3d'],
        symbolic['symbols']['n_bath'],
        symbolic['symbols']['delta_mu'],
        symbolic['symbols']['delta_o'],
        symbolic['symbols']['delta_amu'],
        symbolic['symbols']['delta_ao']
    )

    symbolic_args_H = (
        symbolic['symbols']['g_o'],
        symbolic['symbols']['g_mu'],
        symbolic['symbols']['Omega'],
        symbolic['symbols']['delta_mu'],
        symbolic['symbols']['delta_o'],
        symbolic['symbols']['delta_amu'],
        symbolic['symbols']['delta_ao'],
        symbolic['symbols']['alpha'],
        symbolic['symbols']['beta']
    )

    H = sympy.lambdify(symbolic_args_H, symbolic['H'], 'numpy')
    L = sympy.lambdify(symbolic_args_full, symbolic['L'], 'numpy')
    L_0 = sympy.lambdify(symbolic_args_0, symbolic['L_0'], 'numpy')
    L_alpha = sympy.lambdify(symbolic['symbols']['g_o'], symbolic['L_alpha'], 'numpy')
    L_alpha_c = sympy.lambdify(symbolic['symbols']['g_o'], symbolic['L_alpha_c'], 'numpy')
    L_beta = sympy.lambdify(symbolic['symbols']['g_mu'], symbolic['L_beta'], 'numpy')
    L_beta_c = sympy.lambdify(symbolic['symbols']['g_mu'], symbolic['L_beta_c'], 'numpy')

    def rho_steady_state(g_o, g_mu, Omega, gamma_12, gamma_13,
                         gamma_23, gamma_2d, gamma_3d, n_bath,
                         delta_mu, delta_o, delta_amu, delta_ao, alpha, beta):
        L_matrix = L(
            g_o=g_o,
            g_mu=g_mu,
            Omega=Omega,
            gamma_12=gamma_12,
            gamma_13=gamma_13,
            gamma_23=gamma_23,
            gamma_2d=gamma_2d,
            gamma_3d=gamma_3d,
            n_bath=n_bath,
            delta_mu=delta_mu,
            delta_o=delta_o,
            delta_amu=delta_amu,
            delta_ao=delta_ao,
            alpha=alpha,
            beta=beta
        )
        L_matrix[0,:] = np.array([1, 0, 0, 0, 1, 0, 0, 0, 1])
        b = np.array([1, 0, 0, 0, 0, 0, 0, 0, 0])
        rho = np.linalg.solve(L_matrix, b)
        return rho[unravel]

    def rho_steady_state_ensemble(g_o, g_mu, Omega, gamma_12, gamma_13,
                                  gamma_23, gamma_2d, gamma_3d, n_bath,
                                  delta_mu, delta_o, delta_amu, delta_ao,
                                  sigma_mu, sigma_o, alpha, beta):
        # degenerate dressed state detuning
        def minor_det(M, i, j):
                return M[i,i]*M[j,j] - M[i,j]*M[j,i]

        def H_disc(delta_amu, delta_ao):
            H_matrix = H(
                g_mu=g_mu,
                g_o=g_o,
                Omega=Omega,
                alpha=alpha,
                beta=beta,
                delta_mu=delta_mu,
                delta_o=delta_o,
                delta_amu=delta_amu,
                delta_ao=delta_ao
            )
            a = -1
            b = np.trace(H_matrix)
            c = -(minor_det(H_matrix, 0, 1)
                + minor_det(H_matrix, 0, 2)
                + minor_det(H_matrix, 1, 2))
            d = np.linalg.det(H_matrix)
            Delta = 18*a*b*c*d - 4*b**3*d + b*b*c*c - 4*a*c**3 - 27*a*a*d*d
            return np.abs(Delta)

        def delta_amu_degenerate(delta_ao):
            # get close guess
            if delta_ao == delta_o:
                delta_amu0 = delta_mu
            elif np.abs(beta) < np.abs(Omega):
                delta_amu0 = -np.abs(Omega)**2/(delta_ao-delta_o) + delta_ao - delta_o + delta_mu
            else:
                delta_amu0 = np.abs(g_mu*beta)**2/(delta_ao-delta_o) + delta_mu

            f = lambda d_amu: H_disc(d_amu, delta_ao)
            result = optimize.minimize_scalar(f, bracket=(delta_amu0-sigma_mu, delta_amu0+sigma_mu))
            return result.x

        # gauss-lobatto integration
        def gauss_lobatto_nodes_weights(a, b):
            nodes = a + gauss_lobatto_points*(b-a)
            weights = gauss_lobatto_weights*(b-a)
            return nodes, weights

        def composite_gauss_lobatto_nodes_weights(bounds):
            bounds = np.sort(bounds)
            nodes = np.array([], dtype=float)
            weights = np.array([], dtype=float)
            for a, b in zip(bounds[:-1], bounds[1:]):
                new_nodes, new_weights = gauss_lobatto_nodes_weights(a, b)
                nodes = np.concatenate([nodes, new_nodes])
                weights = np.concatenate([weights, new_weights])
            return nodes, weights

        # rho helper
        def rho_fewer_args(d_ao, d_amu):
            return rho_steady_state(
                g_o=g_o,
                g_mu=g_mu,
                Omega=Omega,
                gamma_12=gamma_12,
                gamma_13=gamma_13,
                gamma_23=gamma_23,
                gamma_2d=gamma_2d,
                gamma_3d=gamma_3d,
                n_bath=n_bath,
                delta_mu=delta_mu,
                delta_o=delta_o,
                delta_amu=d_amu,
                delta_ao=d_ao,
                alpha=alpha,
                beta=beta
            )

        # generate outer integral nodes and weights
        bounds = [
            delta_ao,
            delta_ao - sigma_o,
            delta_ao + sigma_o,
            delta_ao - 3*sigma_o,
            delta_ao + 3*sigma_o,
            delta_ao - 10*sigma_o,
            delta_ao + 10*sigma_o,
            delta_o,
            delta_o - gamma_3d,
            delta_o + gamma_3d,
            delta_o - 5*gamma_3d,
            delta_o + 5*gamma_3d
        ]
        d_ao_nodes, d_ao_weights = composite_gauss_lobatto_nodes_weights(bounds)

        # outer integral envelope function
        z_ao_nodes = (d_ao_nodes-delta_ao) / sigma_o
        G_ao_nodes = stats.norm.pdf(z_ao_nodes) / sigma_o

        # perform integrals
        rho_integral = np.zeros((3, 3), dtype=complex)
        for d_ao, d_ao_w, G_ao in zip(d_ao_nodes, d_ao_weights, G_ao_nodes):
            # generate inner integral nodes and weights
            bounds = [
                delta_amu,
                delta_amu - sigma_mu,
                delta_amu + sigma_mu,
                delta_amu - 3*sigma_mu,
                delta_amu + 3*sigma_mu,
                delta_amu - 10*sigma_mu,
                delta_amu + 10*sigma_mu,
                delta_mu,
                delta_mu - gamma_2d,
                delta_mu + gamma_2d,
                delta_mu - 5*gamma_2d,
                delta_mu + 5*gamma_2d,
                delta_amu_degenerate(d_ao)
            ]
            d_amu_nodes, d_amu_weights = composite_gauss_lobatto_nodes_weights(bounds)

            # inner integral envelope function
            z_amu_nodes = (d_amu_nodes-delta_amu) / sigma_mu
            G_amu_nodes = stats.norm.pdf(z_amu_nodes) / sigma_mu

            # perform inner integral
            for d_amu, d_amu_w, G_amu in zip(d_amu_nodes, d_amu_weights, G_amu_nodes):
                rho = rho_fewer_args(d_ao, d_amu)
                rho_integral += G_ao * G_amu * rho * d_ao_w * d_amu_w

        return rho_integral

    def alpha_beta_langevin_differential(N_o, N_mu, g_o, g_mu, Omega, gamma_12, gamma_13,
                                         gamma_23, gamma_2d, gamma_3d, gamma_muc, gamma_mui,
                                         gamma_oc, gamma_oi, n_bath, delta_mu, delta_o, delta_amu,
                                         delta_ao, sigma_mu, sigma_o, alpha, beta, alpha_in, beta_in,
                                         use_rho_ji=True):
        rho_steady = rho_steady_state_ensemble(
            g_o=g_o,
            g_mu=g_mu,
            Omega=Omega,
            gamma_12=gamma_12,
            gamma_13=gamma_13,
            gamma_23=gamma_23,
            gamma_2d=gamma_2d,
            gamma_3d=gamma_3d,
            n_bath=n_bath,
            delta_mu=delta_mu,
            delta_o=delta_o,
            delta_amu=delta_amu,
            delta_ao=delta_ao,
            sigma_mu=sigma_mu,
            sigma_o=sigma_o,
            alpha=alpha,
            beta=beta
        )
        S12 = N_mu * g_mu * (rho_steady[1,0] if use_rho_ji else rho_steady[0,1])
        S13 = N_o * g_o * (rho_steady[2,0] if use_rho_ji else rho_steady[0,2])
        alpha_diff = 1j*delta_o*alpha - 1j*S13 - (gamma_oc+gamma_oi)*alpha/2 + np.sqrt(gamma_oc)*alpha_in
        beta_diff = 1j*delta_mu*beta - 1j*S12 - (gamma_muc+gamma_mui)*beta/2 + np.sqrt(gamma_muc)*beta_in
        return alpha_diff, beta_diff

    def alpha_beta_steady_state(N_o, N_mu, g_o, g_mu, Omega, gamma_12, gamma_13,
                                gamma_23, gamma_2d, gamma_3d, gamma_muc, gamma_mui,
                                gamma_oc, gamma_oi, n_bath, delta_mu, delta_o, delta_amu,
                                delta_ao, sigma_mu, sigma_o, alpha_in, beta_in, use_rho_ji=True):
        def root_function(alpha_beta_vec):
            alpha_r = alpha_beta_vec[0]
            alpha_i = alpha_beta_vec[1]
            beta_r = alpha_beta_vec[2]
            beta_i = alpha_beta_vec[3]
            alpha = alpha_r + 1j*alpha_i
            beta = beta_r + 1j*beta_i
            alpha_res, beta_res = alpha_beta_langevin_differential(
                N_o=N_o,
                N_mu=N_mu,
                g_o=g_o,
                g_mu=g_mu,
                Omega=Omega,
                gamma_12=gamma_12,
                gamma_13=gamma_13,
                gamma_23=gamma_23,
                gamma_2d=gamma_2d,
                gamma_3d=gamma_3d,
                gamma_muc=gamma_muc,
                gamma_mui=gamma_mui,
                gamma_oc=gamma_oc,
                gamma_oi=gamma_oi,
                n_bath=n_bath,
                delta_mu=delta_mu,
                delta_o=delta_o,
                delta_amu=delta_amu,
                delta_ao=delta_ao,
                sigma_mu=sigma_mu,
                sigma_o=sigma_o,
                alpha=alpha,
                beta=beta,
                alpha_in=alpha_in,
                beta_in=beta_in,
                use_rho_ji=use_rho_ji
            )
            alpha_beta_res_vec = np.zeros(4, dtype=float)
            alpha_beta_res_vec[0] = np.real(alpha_res)
            alpha_beta_res_vec[1] = np.imag(alpha_res)
            alpha_beta_res_vec[2] = np.real(beta_res)
            alpha_beta_res_vec[3] = np.imag(beta_res)
            return alpha_beta_res_vec

        # initial guess for root-finding; root for zero atomic interaction
        alpha_0 = alpha_in*np.sqrt(gamma_oc) / ((gamma_oi+gamma_oc)/2 - 1j*delta_o)
        beta_0 = beta_in*np.sqrt(gamma_muc) / ((gamma_mui+gamma_muc)/2 - 1j*delta_mu)

        # perform root finding
        x0 = np.zeros(4, dtype=float)
        x0[0] = np.real(alpha_0)
        x0[1] = np.imag(alpha_0)
        x0[2] = np.real(beta_0)
        x0[3] = np.imag(beta_0)
        result = optimize.root(root_function, x0=x0, tol=1e-12)

        # restore result to complex numbers
        alpha_r = result.x[0]
        alpha_i = result.x[1]
        beta_r = result.x[2]
        beta_i = result.x[3]
        alpha = alpha_r + 1j*alpha_i
        beta = beta_r + 1j*beta_i

        return (alpha, beta), result
    |

## Four-Level Transduction

This Python code implements the four-level transduction model in \ref{ch:four_level_transduction}. When imported into a Python script, it expects to be able to save and load a file into a directory called `result-cache`. When run as a script, it does the aforementioned saving and loading and nothing else. The notation in this code differs from that in this thesis, mainly in that the \m{23}, \m{13}, \m{24}, and \m{14} transitions are labelled \m{A}, \m{B}, \m{D}, and \m{E} respectively. Additionally, \m{\delta_p} is instead called \m{\delta_B}.

    \
    \verb|import numpy as np
    from scipy import signal
    import sympy
    import pickle
    import os

    # fundamental constants

    hbar = 1.054571817e-34
    kB = 1.380649e-23

    # flattening and unflattening of density matrices

    unflatten = np.array([
        [ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15]
    ])

    flatten_i = np.zeros(16, dtype=int)
    flatten_j = np.zeros(16, dtype=int)
    for i in range(4):
        for j in range(4):
            k = unflatten[i,j]
            flatten_i[k] = i
            flatten_j[k] = j
    flatten = (flatten_i, flatten_j)

    # computer algebra for Liouvillan matrix

    s11 = sympy.Matrix([[1,0,0,0],[0,0,0,0],[0,0,0,0],[0,0,0,0]])
    s12 = sympy.Matrix([[0,1,0,0],[0,0,0,0],[0,0,0,0],[0,0,0,0]])
    s13 = sympy.Matrix([[0,0,1,0],[0,0,0,0],[0,0,0,0],[0,0,0,0]])
    s14 = sympy.Matrix([[0,0,0,1],[0,0,0,0],[0,0,0,0],[0,0,0,0]])
    s21 = s12.T
    s31 = s13.T
    s41 = s14.T
    s22 = s21*s12
    s23 = s21*s13
    s24 = s21*s14
    s32 = s23.T
    s42 = s24.T
    s33 = s31*s13
    s34 = s31*s14
    s43 = s34.T
    s44 = s43*s34

    omega12 = sympy.symbols('omega_12', real=True)
    delta_B = sympy.symbols('delta_B', real=True)
    delta_mu = sympy.symbols('delta_mu', real=True)
    Omega_mu = sympy.symbols('Omega_mu')
    Omega_A = sympy.symbols('Omega_A')
    Omega_B = sympy.symbols('Omega_B')
    Omega_D = sympy.symbols('Omega_D')
    Omega_E = sympy.symbols('Omega_E')

    gamma12, gamma13, gamma14 = sympy.symbols('gamma_1(2:5)', real=True)
    gamma23, gamma24 = sympy.symbols('gamma_2(3:5)', real=True)
    gamma34 = sympy.symbols('gamma_34', real=True)
    gamma2d, gamma3d, gamma4d = sympy.symbols('gamma_(2:5)d', real=True)
    nbath_12 = sympy.symbols('n_12', real=True)
    nbath_34 = sympy.symbols('n_34', real=True)

    H = Omega_A*s32 + Omega_B*s31 + Omega_D*s42 + Omega_E*s41 + Omega_mu*s43
    H += H.H
    H += omega12*s22 + delta_B*s33 + (delta_B+delta_mu)*s44

    def loss_superoperator(rho):
        L12 = gamma12*(nbath_12+1)/2 * (2*s12*rho*s21 - s22*rho - rho*s22)
        L21 = gamma12*nbath_12/2 * (2*s21*rho*s12 - s11*rho - rho*s11)
        L13 = gamma13/2 * (2*s13*rho*s31 - s33*rho - rho*s33)
        L14 = gamma14/2 * (2*s14*rho*s41 - s44*rho - rho*s44)
        L23 = gamma23/2 * (2*s23*rho*s32 - s33*rho - rho*s33)
        L24 = gamma24/2 * (2*s24*rho*s42 - s44*rho - rho*s44)
        L34 = gamma34*(nbath_34+1)/2 * (2*s34*rho*s43 - s44*rho - rho*s44)
        L43 = gamma34*nbath_34/2 * (2*s43*rho*s34 - s33*rho - rho*s33)
        L2d = gamma2d/2 * (2*s22*rho*s22 - s22*rho - rho*s22)
        L3d = gamma3d/2 * (2*s33*rho*s33 - s33*rho - rho*s33)
        L4d = gamma4d/2 * (2*s44*rho*s44 - s44*rho - rho*s44)
        return L12 + L21 + L13 + L14 + L23 + L24 + L34 + L43 + L2d + L3d + L4d

    def liouvillan_superoperator(rho):
        return -sympy.I*(H*rho-rho*H) + loss_superoperator(rho)

    # construct Liouvillan matrix
    L = sympy.zeros(16)
    for column, (icol, jcol) in enumerate(zip(*flatten)):
        # basis matrix
        rho = sympy.zeros(4)
        rho[icol,jcol] = 1

        # get column of supermatrix
        L_column = liouvillan_superoperator(rho)
        for row, (irow, jrow) in enumerate(zip(*flatten)):
            L[row,column] = L_column[irow,jrow]

    # matrix to convert complex Hermitian to real nonsymmetric
    C = sympy.zeros(16)
    for k in range(4):
        ik = unflatten[k,k]
        C[ik,ik] = 1
    for j in range(3):
        for k in range(j+1, 4):
            ij = unflatten[j,k]
            ik = unflatten[k,j]
            C[ij,ij] = sympy.Rational(1, 2)
            C[ij,ik] = sympy.Rational(1, 2)
            C[ik,ij] = -sympy.I/2
            C[ik,ik] = sympy.I/2

    # Liouvillan matrix converted to real
    C_inv = C.inv()
    Lreal = sympy.re(C*L*C.inv())

    # lambdified functions
    def short_lambdify(params, expr):
        return sympy.lambdify(params, expr, 'numpy', cse=True, docstring_limit=0)

    Hsymbols = (omega12, delta_B, delta_mu, Omega_mu,
                Omega_A, Omega_B, Omega_D, Omega_E)
    Lsymbols = Hsymbols + (gamma12, gamma13, gamma14, gamma23, gamma24, gamma34,
                            gamma2d, gamma3d, gamma4d, nbath_12, nbath_34)
    Hfunc = short_lambdify(Hsymbols, H)
    Lrealfunc = short_lambdify(Lsymbols, Lreal)

    # gross hack to get numerical matrices from sympy
    x = sympy.symbols('x')
    complex_to_real = short_lambdify(x, C)(None)
    real_to_complex = short_lambdify(x, C_inv)(None)

    # discriminant symbolic and lambdified functions for numerical methods

    def expr_to_poly_coeffs(expr, x):
        poly = sympy.poly(expr, x)
        n = poly.degree(x)
        coeffs = sympy.zeros(n+1, 1)
        for (k,), coeff in zip(poly.monoms(), poly.coeffs()):
            coeffs[k] = coeff
        return coeffs

    def diff_poly_coeffs(coeffs):
        n = coeffs.shape[0]-1
        diff_coeffs = sympy.zeros(n, 1)
        for k in range(1, n+1):
            diff_coeffs[k-1,0] = sympy.Integer(k)*coeffs[k,0]
        return diff_coeffs

    # restore the symbolic expressions from a cache, if it exists
    fname = 'result-cache/Delta_expr.pkl'
    if os.path.isfile(fname):
        with open(fname, 'rb') as f:
            Delta_expr = pickle.load(f)
            Delta = Delta_expr['Delta']
            Delta_coeffs_delta_B = Delta_expr['Delta_coeffs_delta_B']
            Delta_coeffs_delta_mu = Delta_expr['Delta_coeffs_delta_mu']
            dmu_Delta = Delta_expr['dmu_Delta']
            dmu_Delta_coeffs_delta_B = Delta_expr['dmu_Delta_coeffs_delta_B']
            dmu_Delta_coeffs_delta_mu = Delta_expr['dmu_Delta_coeffs_delta_mu']
            dB_Delta = Delta_expr['dB_Delta']
            dB_Delta_coeffs_delta_B = Delta_expr['dB_Delta_coeffs_delta_B']
            dB_Delta_coeffs_delta_mu = Delta_expr['dB_Delta_coeffs_delta_mu']
            d2mu_Delta = Delta_expr['d2mu_Delta']
            d2B_Delta = Delta_expr['d2B_Delta']
    else:
        # compute the discriminant Delta of the characteristic polynomial
        # of the Hamiltonian; this slightly convoluted method of using
        # a generic quartic is faster than the obvious way
        print('Computing the symbolic polynomial coefficients of the discriminant')
        print('This takes about 20 minutes on my machine, so get ready to wait')
        print(f'The result will be saved as {fname} to be re-used in future builds')
        a = [a0, a1, a2, a3, a4] = sympy.symbols('a_(0:5)', real=True)
        x = sympy.symbols('x', real=True)
        poly = sympy.poly(a4*x**4 + a3*x**3 + a2*x**2 + a1*x + a0, x)
        Delta = poly.discriminant()
        charpoly = H.charpoly()
        for (k,), coeff in zip(charpoly.monoms(), charpoly.coeffs()):
            Delta = Delta.subs(a[k], coeff)

        # compute various polynomial coefficients and derivatives of Delta
        Delta_coeffs_delta_B = expr_to_poly_coeffs(Delta, delta_B)
        Delta_coeffs_delta_mu = expr_to_poly_coeffs(Delta, delta_mu)
        dmu_Delta = sympy.diff(Delta, delta_mu)
        dmu_Delta_coeffs_delta_B = sympy.diff(Delta_coeffs_delta_B, delta_mu)
        dmu_Delta_coeffs_delta_mu = diff_poly_coeffs(Delta_coeffs_delta_mu)
        dB_Delta = sympy.diff(Delta, delta_B)
        dB_Delta_coeffs_delta_B = diff_poly_coeffs(Delta_coeffs_delta_B)
        dB_Delta_coeffs_delta_mu = sympy.diff(Delta_coeffs_delta_mu, delta_B)
        d2mu_Delta = sympy.diff(Delta, delta_mu, 2)
        d2B_Delta = sympy.diff(Delta, delta_B, 2)
        Delta_expr = {
            'Delta': Delta,
            'Delta_coeffs_delta_B': Delta_coeffs_delta_B,
            'Delta_coeffs_delta_mu': Delta_coeffs_delta_mu,
            'dmu_Delta': dmu_Delta,
            'dmu_Delta_coeffs_delta_B': dmu_Delta_coeffs_delta_B,
            'dmu_Delta_coeffs_delta_mu': dmu_Delta_coeffs_delta_mu,
            'dB_Delta': dB_Delta,
            'dB_Delta_coeffs_delta_B': dB_Delta_coeffs_delta_B,
            'dB_Delta_coeffs_delta_mu': dB_Delta_coeffs_delta_mu,
            'd2mu_Delta': d2mu_Delta,
            'd2B_Delta': d2B_Delta
        }
        with open(fname, 'wb') as f:
            pickle.dump(Delta_expr, f)

    H_symbols_common = (omega12, Omega_A, Omega_B, Omega_D, Omega_E, Omega_mu)
    H_symbols_no_dB = H_symbols_common + (delta_mu,)
    H_symbols_no_dmu = H_symbols_common + (delta_B,)
    H_symbols_all = Hsymbols

    # lambdify these expressions; this can't be pickled
    Delta_func = short_lambdify(H_symbols_all, Delta)
    Delta_coeffs_delta_B_func = short_lambdify(H_symbols_no_dB,
                                               Delta_coeffs_delta_B)
    Delta_coeffs_delta_mu_func = sympy.lambdify(H_symbols_no_dmu,
                                                Delta_coeffs_delta_mu)
    dmu_Delta_func = sympy.lambdify(H_symbols_all, dmu_Delta)
    dmu_Delta_coeffs_delta_B_func = sympy.lambdify(H_symbols_no_dB,
                                                   dmu_Delta_coeffs_delta_B)
    dmu_Delta_coeffs_delta_mu_func = sympy.lambdify(H_symbols_no_dmu,
                                                    dmu_Delta_coeffs_delta_mu)
    dB_Delta_func = sympy.lambdify(H_symbols_all, dB_Delta)
    dB_Delta_coeffs_delta_B_func = sympy.lambdify(H_symbols_no_dB,
                                                  dB_Delta_coeffs_delta_B)
    dB_Delta_coeffs_delta_mu_func = sympy.lambdify(H_symbols_no_dmu,
                                                   dB_Delta_coeffs_delta_mu)
    d2mu_Delta_func = sympy.lambdify(H_symbols_all, d2mu_Delta)
    d2B_Delta_func = sympy.lambdify(H_symbols_all, d2B_Delta)

    # helper function for multivariate normal distribution; scipy's has a bad API

    def multivariate_normal_pdf(Sigma, *args):
        d = len(args)
        x = np.array(args)
        coef = 1 / np.sqrt((2*np.pi)**d * np.linalg.det(Sigma))

        # permutation sending the first dimension to the second-last
        axes = list(range(x.ndim))
        axes[0] = -2
        axes[-2] = 0

        xDivSigma = np.linalg.solve(Sigma, np.transpose(x, axes=axes))
        xDivSigma = np.transpose(xDivSigma, axes=axes)

        return coef * np.exp(-np.einsum('i...,i...', x, xDivSigma)/2)

    # computes a grid of atom output signal, without multisampling

    def rho_steady_state(omega_12, omega_34, delta_B, delta_mu, Omega_A,
                         Omega_B, Omega_D, Omega_E, Omega_mu, tau_12,
                         tau_34, gamma_13, gamma_14, gamma_23, gamma_24,
                         gamma_2d, gamma_3d, gamma_4d, T):
        n_12 = 1 / np.expm1(hbar*omega_12/(kB*T))
        n_34 = 1 / np.expm1(hbar*omega_34/(kB*T))
        L = Lrealfunc(
            omega_12=omega_12,
            delta_B=delta_B,
            delta_mu=delta_mu,
            Omega_A=Omega_A,
            Omega_B=Omega_B,
            Omega_D=Omega_D,
            Omega_E=Omega_E,
            Omega_mu=Omega_mu,
            gamma_12=1/(tau_12*(n_12+1)),
            gamma_13=gamma_13,
            gamma_14=gamma_14,
            gamma_23=gamma_23,
            gamma_24=gamma_24,
            gamma_34=1/(tau_34*(n_34+1)),
            gamma_2d=gamma_2d,
            gamma_3d=gamma_3d,
            gamma_4d=gamma_4d,
            n_12=n_12,
            n_34=n_34
        )
        L[0,:] = np.identity(4)[flatten]

        b = np.zeros(16)
        b[0] = 1

        rho_real = np.linalg.solve(L, b)
        rho = real_to_complex @ rho_real
        return rho[unflatten]

    def atom_signal(omega_12, omega_34, delta_B, delta_mu, Omega_A,
                    Omega_B, Omega_D, Omega_E, Omega_mu, tau_12, tau_34,
                    gamma_13, gamma_14, gamma_23, gamma_24, C_14,
                    C_24, gamma_2d, gamma_3d, gamma_4d, T):
        rho = rho_steady_state(
            omega_12=omega_12,
            omega_34=omega_34,
            delta_B=delta_B,
            delta_mu=delta_mu,
            Omega_A=Omega_A,
            Omega_B=Omega_B,
            Omega_D=Omega_D,
            Omega_E=Omega_E,
            Omega_mu=Omega_mu,
            tau_12=tau_12,
            tau_34=tau_34,
            gamma_13=gamma_13,
            gamma_14=gamma_14,
            gamma_23=gamma_23,
            gamma_24=gamma_24,
            gamma_2d=gamma_2d,
            gamma_3d=gamma_3d,
            gamma_4d=gamma_4d,
            T=T
        )

        signal_D = C_24 * rho[3,1]
        signal_E = C_14 * rho[3,0]
        photon_rate_out = np.abs(signal_D + signal_E)**2
        return photon_rate_out

    def atom_scan(delta_mu_min, delta_mu_max, delta_mu_points,
                  delta_B_min, delta_B_max, delta_B_points,
                  omega_12, omega_34, Omega_A, Omega_B, Omega_D, Omega_E,
                  Omega_mu, tau_12, tau_34, gamma_13, gamma_14, C_14, C_24,
                  gamma_23, gamma_24, gamma_2d, gamma_3d, gamma_4d, T):
        delta_mu = np.linspace(delta_mu_min, delta_mu_max, delta_mu_points)
        delta_B = np.linspace(delta_B_min, delta_B_max, delta_B_points)
        delta_mu, delta_B = np.meshgrid(delta_mu, delta_B, indexing='ij')

        scan = np.zeros_like(delta_mu)
        for i in range(scan.shape[0]):
            for j in range(scan.shape[1]):
                scan[i,j] = atom_signal(
                    omega_12=omega_12,
                    omega_34=omega_34,
                    delta_B=delta_B[i,j],
                    delta_mu=delta_mu[i,j],
                    Omega_A=Omega_A,
                    Omega_B=Omega_B,
                    Omega_D=Omega_D,
                    Omega_E=Omega_E,
                    Omega_mu=Omega_mu,
                    tau_12=tau_12,
                    tau_34=tau_34,
                    gamma_13=gamma_13,
                    gamma_14=gamma_14,
                    gamma_23=gamma_23,
                    gamma_24=gamma_24,
                    gamma_2d=gamma_2d,
                    gamma_3d=gamma_3d,
                    gamma_4d=gamma_4d,
                    C_14=C_14,
                    C_24=C_24,
                    T=T
                )

        return delta_mu, delta_B, scan

    # multisampling functions

    def poly_real_roots(poly_coeffs):
        poly = np.polynomial.Polynomial(poly_coeffs)
        roots = poly.roots()
        real_roots = np.array([np.real(x) for x in roots if np.imag(x)==0])
        return real_roots

    def poly_positive_minima(poly_func, diff_coeffs):
        critical_x = poly_real_roots(diff_coeffs)
        critical_y = poly_func(critical_x)

        # base cases
        if len(critical_x) == 0:
            return []
        if len(critical_x) <= 1:
            return critical_x

        minima = []

        # check if leftmost critical point is minimum
        if critical_y[0] < critical_y[1]:
            minima.append(critical_x[0])

        # check for mimima between critical points
        for i in range(len(critical_x)-2):
            y1 = critical_y[i]
            x2 = critical_x[i+1]
            y2 = critical_y[i+1]
            y3 = critical_y[i+2]
            if y1 > y2 and y3 > y2:
                minima.append(x2)

        # check if rightmost critical point is mimimum
        if critical_y[-1] < critical_y[-2]:
            minima.append(critical_x[-1])

        return minima

    def H_get_minimal_delta_B(omega_12, delta_mu, Omega_A,
                              Omega_B, Omega_D, Omega_E, Omega_mu):
        # minima along delta_B
        dB_Delta_coeffs = dB_Delta_coeffs_delta_B_func(
            omega_12=omega_12,
            delta_mu=delta_mu,
            Omega_A=Omega_A,
            Omega_B=Omega_B,
            Omega_D=Omega_D,
            Omega_E=Omega_E,
            Omega_mu=Omega_mu
        )[:,0]
        Delta_short_func = lambda delta_B: Delta_func(
            omega_12=omega_12,
            delta_mu=delta_mu,
            delta_B=delta_B,
            Omega_A=Omega_A,
            Omega_B=Omega_B,
            Omega_D=Omega_D,
            Omega_E=Omega_E,
            Omega_mu=Omega_mu
        )
        minima = list(poly_positive_minima(Delta_short_func, dB_Delta_coeffs))
        is_parallel = [True]*len(minima)

        # minima along delta_mu
        dmu_Delta_coeffs = dmu_Delta_coeffs_delta_B_func(
            omega_12=omega_12,
            delta_mu=delta_mu,
            Omega_A=Omega_A,
            Omega_B=Omega_B,
            Omega_D=Omega_D,
            Omega_E=Omega_E,
            Omega_mu=Omega_mu
        )[:,0]
        roots = poly_real_roots(dmu_Delta_coeffs)
        for root in roots:
            curvature = d2mu_Delta_func(
                omega_12=omega_12,
                delta_mu=delta_mu,
                delta_B=root,
                Omega_A=Omega_A,
                Omega_B=Omega_B,
                Omega_D=Omega_D,
                Omega_E=Omega_E,
                Omega_mu=Omega_mu
            )
            if curvature >= 0:
                minima.append(root)
                is_parallel.append(False)

        return minima, is_parallel

    def H_get_minimal_delta_mu(omega_12, delta_B, Omega_A,
                               Omega_B, Omega_D, Omega_E, Omega_mu):
        # minima along delta_mu
        dmu_Delta_coeffs = dmu_Delta_coeffs_delta_mu_func(
            omega_12=omega_12,
            delta_B=delta_B,
            Omega_A=Omega_A,
            Omega_B=Omega_B,
            Omega_D=Omega_D,
            Omega_E=Omega_E,
            Omega_mu=Omega_mu
        )[:,0]
        Delta_short_func = lambda delta_mu: Delta_func(
            omega_12=omega_12,
            delta_mu=delta_mu,
            delta_B=delta_B,
            Omega_A=Omega_A,
            Omega_B=Omega_B,
            Omega_D=Omega_D,
            Omega_E=Omega_E,
            Omega_mu=Omega_mu
        )
        minima = list(poly_positive_minima(Delta_short_func, dmu_Delta_coeffs))
        is_parallel = [True]*len(minima)

        # minima along delta_mu
        dB_Delta_coeffs = dB_Delta_coeffs_delta_mu_func(
            omega_12=omega_12,
            delta_B=delta_B,
            Omega_A=Omega_A,
            Omega_B=Omega_B,
            Omega_D=Omega_D,
            Omega_E=Omega_E,
            Omega_mu=Omega_mu
        )[:,0]
        roots = poly_real_roots(dB_Delta_coeffs)
        for root in roots:
            curvature = d2B_Delta_func(
                omega_12=omega_12,
                delta_mu=root,
                delta_B=delta_B,
                Omega_A=Omega_A,
                Omega_B=Omega_B,
                Omega_D=Omega_D,
                Omega_E=Omega_E,
                Omega_mu=Omega_mu
            )
            if curvature >= 0:
                minima.append(root)
                is_parallel.append(False)

        return minima, is_parallel

    def find_curve_pixel_intersections(delta_mu_min, delta_mu_max, delta_mu_points,
                                       delta_B_min, delta_B_max, delta_B_points,
                                       omega_12, Omega_A, Omega_B,
                                       Omega_D, Omega_E, Omega_mu):
        def linspace_edges(start, stop, count):
            dx = (stop-start) / (count-1)
            edges = np.linspace(start-dx/2, stop+dx/2, count+1)
            return edges

        delta_mu_edges = linspace_edges(delta_mu_min, delta_mu_max, delta_mu_points)
        delta_B_edges = linspace_edges(delta_B_min, delta_B_max, delta_B_points)
        delta_mu_edges_min, delta_mu_edges_max = delta_mu_edges[0], delta_mu_edges[-1]
        delta_B_edges_min, delta_B_edges_max = delta_B_edges[0], delta_B_edges[-1]

        intersection_points = []
        intersecting_pixels = dict()
        is_parallel = []

        for i, dmu in enumerate(delta_mu_edges):
            dBs, paras = H_get_minimal_delta_B(
                omega_12=omega_12,
                delta_mu=dmu,
                Omega_A=Omega_A,
                Omega_B=Omega_B,
                Omega_D=Omega_D,
                Omega_E=Omega_E,
                Omega_mu=Omega_mu
            )
            for dB, para in zip(dBs, paras):
                if not delta_B_edges_min <= dB <= delta_B_edges_max:
                    continue
                point = (dmu, dB)
                intersection_points.append((dmu, dB))
                is_parallel.append(para)

                j = int(delta_B_points * (dB-delta_B_edges_min) / (delta_B_edges_max-delta_B_edges_min))
                if i > 0:
                    if (i-1,j) not in intersecting_pixels:
                        intersecting_pixels[i-1,j] = []
                    intersecting_pixels[i-1,j].append(point)
                if i < delta_mu_points:
                    if (i,j) not in intersecting_pixels:
                        intersecting_pixels[i,j] = []
                    intersecting_pixels[i,j].append(point)

        for j, dB in enumerate(delta_B_edges):
            dmus, paras = H_get_minimal_delta_mu(
                omega_12=omega_12,
                delta_B=dB,
                Omega_A=Omega_A,
                Omega_B=Omega_B,
                Omega_D=Omega_D,
                Omega_E=Omega_E,
                Omega_mu=Omega_mu
            )
            for dmu, para in zip(dmus, paras):
                if not delta_mu_edges_min <= dmu <= delta_mu_edges_max:
                    continue
                point = (dmu, dB)
                intersection_points.append((dmu, dB))
                is_parallel.append(para)

                i = int(delta_mu_points * (dmu-delta_mu_edges_min) / (delta_mu_edges_max-delta_mu_edges_min))
                if j > 0:
                    if (i,j-1) not in intersecting_pixels:
                        intersecting_pixels[i,j-1] = []
                    intersecting_pixels[i,j-1].append(point)
                if j < delta_B_points:
                    if (i,j) not in intersecting_pixels:
                        intersecting_pixels[i,j] = []
                    intersecting_pixels[i,j].append(point)

        return delta_mu_edges, delta_B_edges, intersection_points, intersecting_pixels, is_parallel

    # precompute Gauss-Lobatto quadrature nodes and weights
    gauss_lobatto_nodes_weights = dict()
    for n in range(2, 100):
        legendre_coeffs = (0,)*(n-1) + (1,)
        legendre_poly = np.polynomial.Legendre(legendre_coeffs)
        nodes = legendre_poly.deriv().roots()
        nodes = np.concatenate([[-1.0], nodes, [1.0]])
        weights = 2 / (n*(n-1)*legendre_poly(nodes)**2)

        # convert from [-1, 1] to [0, 1]
        nodes = (nodes+1)/2
        weights = weights/2

        gauss_lobatto_nodes_weights[n] = (nodes, weights)

    def composite_gauss_lobatto_nodes_weights(n, a, b, intermediates=[]):
        if n not in gauss_lobatto_nodes_weights:
            raise ValueError
        if a >= b:
            raise ValueError

        intermediates = [x for x in intermediates if a < x < b]
        points = sorted([a, b] + intermediates)
        num_points = len(points)
        base_nodes, base_weights = gauss_lobatto_nodes_weights[n]
        if num_points == 2:
            L = points[1] - points[0]
            nodes = points[0] + base_nodes*L
            weights = base_weights*L
            return nodes, weights

        count = (num_points-1)*(n-1) + 1
        nodes = np.zeros(count)
        weights = np.zeros(count)

        # first and last node
        nodes[0] = points[0]
        weights[0] = base_weights[0] * (points[1]-points[0])
        nodes[1] = points[-1]
        weights[1] = base_weights[-1] * (points[-1]-points[-2])
        filled = 2

        # nodes at intermediate points
        for i in range(1,num_points-1):
            nodes[filled] = points[i]
            weights[filled] = base_weights[-1]*(points[i]-points[i-1])
            weights[filled] += base_weights[0]*(points[i+1]-points[i])
            filled += 1

        # nodes between points
        if n == 2:
            return nodes, weights
        for i in range(num_points-1):
            L = points[i+1]-points[i]
            nodes[filled:filled+n-2] = points[i] + base_nodes[1:-1]*L
            weights[filled:filled+n-2] = base_weights[1:-1]*L
            filled += n-2

        return nodes, weights

    def atom_signal_integrated(delta_mu_min, delta_mu_max, delta_B_min, delta_B_max,
                               edge_points, omega_12, omega_34, Omega_A, Omega_B, Omega_D,
                               Omega_E, Omega_mu, tau_12, tau_34, gamma_13, gamma_14,
                               gamma_23, gamma_24, gamma_2d, gamma_3d, gamma_4d, C_14, C_24, T):
        outer_order = 5
        inner_order = 10

        delta_B_linewidth = gamma_3d
        delta_mu_linewidth = gamma_3d + gamma_4d

        area = (delta_mu_max-delta_mu_min) * (delta_B_max-delta_B_min)

        def lerp(a, b, t):
            return a + t*(b-a)
        def ilerp(a, b, x):
            return (x-a)/(b-a)

        def normalise_delta_mu(delta_mu):
            return ilerp(delta_mu_min, delta_mu_max, delta_mu)
        def unnormalise_delta_mu(t):
            return lerp(delta_mu_min, delta_mu_max, t)
        def normalise_delta_B(delta_B):
            return ilerp(delta_B_min, delta_B_max, delta_B)
        def unnormalise_delta_B(t):
            return lerp(delta_B_min, delta_B_max, t)

        def bounds_range(array):
            if len(array) == 0:
                return [], 0
            if len(array) == 1:
                return list(array), 0

            a = np.min(array)
            b = np.max(array)
            return [a, b], b-a

        def atom_signal_short(delta_mu, delta_B):
            return atom_signal(
                omega_12=omega_12,
                omega_34=omega_34,
                delta_mu=delta_mu,
                delta_B=delta_B,
                Omega_A=Omega_A,
                Omega_B=Omega_B,
                Omega_D=Omega_D,
                Omega_E=Omega_E,
                Omega_mu=Omega_mu,
                tau_12=tau_12,
                tau_34=tau_34,
                gamma_13=gamma_13,
                gamma_14=gamma_14,
                gamma_23=gamma_23,
                gamma_24=gamma_24,
                gamma_2d=gamma_2d,
                gamma_3d=gamma_3d,
                gamma_4d=gamma_4d,
                C_14=C_14,
                C_24=C_24,
                T=T,
            )

        edge_mu_points = [x[0] for x in edge_points]
        edge_B_points = [x[1] for x in edge_points]

        # get direction for outer integral
        if len(edge_points) == 0:
            outer = 'delta_mu' # arbitrary
        elif len(edge_points) == 1:
            edge_mu = edge_mu_points[0]
            edge_B = edge_B_points[0]
            if normalise_delta_mu(edge_mu) > normalise_delta_B(edge_B):
                outer = 'delta_mu'
            else:
                outer = 'delta_B'
        else:
            edge_mu_min = np.min(edge_mu_points)
            edge_mu_max = np.max(edge_mu_points)
            edge_B_min = np.min(edge_B_points)
            edge_B_max = np.max(edge_B_points)
            width_mu = (edge_mu_max-edge_mu_min) / (delta_mu_max-delta_mu_min)
            width_B = (edge_B_max-edge_B_min) / (delta_B_max-delta_B_min)
            if width_mu > width_B:
                outer = 'delta_mu'
            else:
                outer = 'delta_B'

        # orientation-neutral definitions
        def atom_signal_short_neutral(delta_outer, delta_inner):
            if outer == 'delta_mu':
                delta_mu = delta_outer
                delta_B = delta_inner
            else:
                delta_mu = delta_inner
                delta_B = delta_outer
            return atom_signal_short(delta_mu, delta_B)

        def H_get_minimal_delta_inner(delta_outer):
            if outer == 'delta_mu':
                return H_get_minimal_delta_B(
                    omega_12=omega_12,
                    delta_mu=delta_outer,
                    Omega_A=Omega_A,
                    Omega_B=Omega_B,
                    Omega_D=Omega_D,
                    Omega_E=Omega_E,
                    Omega_mu=Omega_mu
                )
            else:
                return H_get_minimal_delta_mu(
                    omega_12=omega_12,
                    delta_B=delta_outer,
                    Omega_A=Omega_A,
                    Omega_B=Omega_B,
                    Omega_D=Omega_D,
                    Omega_E=Omega_E,
                    Omega_mu=Omega_mu
                )

        if outer == 'delta_mu':
            delta_outer_min = delta_mu_min
            delta_outer_max = delta_mu_max
            delta_inner_min = delta_B_min
            delta_inner_max = delta_B_max
            delta_inner_linewidth = delta_B_linewidth
            edge_outer_points = edge_mu_points
            edge_inner_points = edge_B_points
        else:
            delta_outer_min = delta_B_min
            delta_outer_max = delta_B_max
            delta_inner_min = delta_mu_min
            delta_inner_max = delta_mu_max
            delta_inner_linewidth = delta_mu_linewidth
            edge_outer_points = edge_B_points
            edge_inner_points = edge_mu_points

        edge_outer_bounds, _ = bounds_range(edge_outer_points)

        # set up diagnostic logging
        log = {
            'delta_mu_min': delta_mu_min,
            'delta_mu_max': delta_mu_max,
            'delta_B_min': delta_B_min,
            'delta_B_max': delta_B_max,
            'edge_points': edge_points,
            'omega_12': omega_12,
            'omega_34': omega_34,
            'Omega_A': Omega_A,
            'Omega_B': Omega_B,
            'Omega_D': Omega_D,
            'Omega_E': Omega_E,
            'Omega_mu': Omega_mu,
            'tau_12': tau_12,
            'tau_34': tau_34,
            'gamma_13': gamma_13,
            'gamma_14': gamma_14,
            'gamma_23': gamma_23,
            'gamma_24': gamma_24,
            'gamma_2d': gamma_2d,
            'gamma_3d': gamma_3d,
            'gamma_4d': gamma_4d,
            'C_14': C_14,
            'C_24': C_24,
            'T': T,
            'outer': outer
        }

        # perform integral
        integral = 0
        outer_nodes, outer_weights = composite_gauss_lobatto_nodes_weights(
            outer_order, delta_outer_min, delta_outer_max, intermediates=edge_outer_bounds)
        log['outer_integral'] = {
            'intermediates': edge_outer_bounds,
            'nodes': outer_nodes,
            'weights': outer_weights,
            'inner_integrals': []
        }
        for outer_node, outer_weight in zip(outer_nodes, outer_weights):
            inner_intersections, _ = H_get_minimal_delta_inner(outer_node)
            inner_intermediates = [z*delta_inner_linewidth+x
                                   for x in inner_intersections
                                   for z in (-10, -3, -1, 0, 1, 3, 10)]
            log_inner_integral = {
                'intersections': inner_intersections,
                'intermediates': inner_intermediates
            }
            inner_nodes, inner_weights = composite_gauss_lobatto_nodes_weights(
                inner_order, delta_inner_min, delta_inner_max, intermediates=inner_intermediates)
            log_inner_integral['nodes'] = inner_nodes
            log_inner_integral['weights'] = inner_weights
            log_inner_integral['values'] = []
            for inner_node, inner_weight in zip(inner_nodes, inner_weights):
                weight = outer_weight * inner_weight
                sample = atom_signal_short_neutral(outer_node, inner_node) / area
                log_inner_integral['values'].append(sample)
                integral += weight * sample
            log['outer_integral']['inner_integrals'].append(log_inner_integral)

        log['integral'] = integral
        return integral, log

    def atom_scan_integrated(delta_mu_min, delta_mu_max, delta_mu_points,
                             delta_B_min, delta_B_max, delta_B_points,
                             omega_12, omega_34, Omega_A, Omega_B, Omega_D, Omega_E,
                             Omega_mu, tau_12, tau_34, gamma_13, gamma_14, gamma_23,
                             gamma_24, gamma_2d, gamma_3d, gamma_4d, C_14, C_24, T):

        (delta_mu_edges,
         delta_B_edges,
         intersection_points,
         intersecting_pixels, _) = find_curve_pixel_intersections(
            delta_mu_min=delta_mu_min,
            delta_mu_max=delta_mu_max,
            delta_mu_points=delta_mu_points,
            delta_B_min=delta_B_min,
            delta_B_max=delta_B_max,
            delta_B_points=delta_B_points,
            omega_12=omega_12,
            Omega_A=Omega_A,
            Omega_B=Omega_B,
            Omega_D=Omega_D,
            Omega_E=Omega_E,
            Omega_mu=Omega_mu
        )

        scan = np.full((delta_mu_points, delta_B_points), np.nan, dtype=float)
        logs = dict()
        for (i,j), edge_points in intersecting_pixels.items():
            scan[i,j], logs[i,j] = atom_signal_integrated(
                delta_mu_min=delta_mu_edges[i],
                delta_mu_max=delta_mu_edges[i+1],
                delta_B_min=delta_B_edges[j],
                delta_B_max=delta_B_edges[j+1],
                edge_points=edge_points,
                omega_12=omega_12,
                omega_34=omega_34,
                Omega_A=Omega_A,
                Omega_B=Omega_B,
                Omega_D=Omega_D,
                Omega_E=Omega_E,
                Omega_mu=Omega_mu,
                tau_12=tau_12,
                tau_34=tau_34,
                gamma_13=gamma_13,
                gamma_14=gamma_14,
                gamma_23=gamma_23,
                gamma_24=gamma_24,
                gamma_2d=gamma_2d,
                gamma_3d=gamma_3d,
                gamma_4d=gamma_4d,
                C_14=C_14,
                C_24=C_24,
                T=T
            )

        delta_mu = np.linspace(delta_mu_min, delta_mu_max, delta_mu_points)
        delta_B = np.linspace(delta_B_min, delta_B_max, delta_B_points)
        delta_mu, delta_B = np.meshgrid(delta_mu, delta_B, indexing='ij')
        for i in range(scan.shape[0]):
            for j in range(scan.shape[1]):
                if not np.isnan(scan[i,j]):
                    continue
                scan[i,j] = atom_signal(
                    omega_12=omega_12,
                    omega_34=omega_34,
                    delta_B=delta_B[i,j],
                    delta_mu=delta_mu[i,j],
                    Omega_A=Omega_A,
                    Omega_B=Omega_B,
                    Omega_D=Omega_D,
                    Omega_E=Omega_E,
                    Omega_mu=Omega_mu,
                    tau_12=tau_12,
                    tau_34=tau_34,
                    gamma_13=gamma_13,
                    gamma_14=gamma_14,
                    gamma_23=gamma_23,
                    gamma_24=gamma_24,
                    gamma_2d=gamma_2d,
                    gamma_3d=gamma_3d,
                    gamma_4d=gamma_4d,
                    C_14=C_14,
                    C_24=C_24,
                    T=T
                )

        return (delta_mu, delta_B, scan), logs

    def signal_scan(delta_mu_min, delta_mu_max, delta_mu_points,
                    delta_B_min, delta_B_max, delta_B_points,
                    omega_12, omega_34, Omega_A, Omega_B, Omega_D, Omega_E,
                    Omega_mu, tau_12, tau_34, gamma_13, gamma_14,
                    gamma_23, gamma_24, gamma_2d, gamma_3d, gamma_4d, T,
                    C_14, C_24, Sigma, z_max=5, integrated=True):
        delta_mu_res = (delta_mu_max-delta_mu_min) / (delta_mu_points-1)
        delta_B_res = (delta_B_max-delta_B_min) / (delta_B_points-1)

        # produce Gaussian filter kernel
        sigma_13 = np.sqrt(Sigma[0,0])
        sigma_34 = np.sqrt(Sigma[1,1])
        sigma_13_px = sigma_13 / delta_B_res
        sigma_34_px = sigma_34 / delta_mu_res
        radius_13_px = int(sigma_13_px*z_max)
        radius_34_px = int(sigma_34_px*z_max)
        delta_13_points = 2*radius_13_px + 1
        delta_34_points = 2*radius_34_px + 1
        delta_13 = np.linspace(-sigma_13*z_max, sigma_13*z_max, delta_13_points)
        delta_34 = np.linspace(-sigma_34*z_max, sigma_34*z_max, delta_34_points)
        delta_34, delta_13 = np.meshgrid(delta_34, delta_13, indexing='ij')
        kernel = multivariate_normal_pdf(Sigma, delta_13, delta_34)
        kernel /= np.sum(kernel)

        # collect atomic emission data
        deltap_B_min = delta_B_min - delta_B_res*radius_13_px
        deltap_B_max = delta_B_max + delta_B_res*radius_13_px
        deltap_B_points = delta_B_points + delta_13_points - 1
        deltap_mu_min = delta_mu_min - delta_mu_res*radius_34_px
        deltap_mu_max = delta_mu_max + delta_mu_res*radius_34_px
        deltap_mu_points = delta_mu_points + delta_34_points - 1

        if integrated:
            (_, _, atom_scan), _ = atom_scan_integrated(
                delta_mu_min=deltap_mu_min,
                delta_mu_max=deltap_mu_max,
                delta_mu_points=deltap_mu_points,
                delta_B_min=deltap_B_min,
                delta_B_max=deltap_B_max,
                delta_B_points=deltap_B_points,
                omega_12=omega_12,
                omega_34=omega_34,
                Omega_A=Omega_A,
                Omega_B=Omega_B,
                Omega_D=Omega_D,
                Omega_E=Omega_E,
                Omega_mu=Omega_mu,
                tau_12=tau_12,
                tau_34=tau_34,
                gamma_13=gamma_13,
                gamma_14=gamma_14,
                gamma_23=gamma_23,
                gamma_24=gamma_24,
                gamma_2d=gamma_2d,
                gamma_3d=gamma_3d,
                gamma_4d=gamma_4d,
                C_14=C_14,
                C_24=C_24,
                T=T
            )
        else:
            _, _, atom_scan = atom_scan(
                delta_mu_min=deltap_mu_min,
                delta_mu_max=deltap_mu_max,
                delta_mu_points=deltap_mu_points,
                delta_B_min=deltap_B_min,
                delta_B_max=deltap_B_max,
                delta_B_points=deltap_B_points,
                omega_12=omega_12,
                omega_34=omega_34,
                Omega_A=Omega_A,
                Omega_B=Omega_B,
                Omega_D=Omega_D,
                Omega_E=Omega_E,
                Omega_mu=Omega_mu,
                tau_12=tau_12,
                tau_34=tau_34,
                gamma_13=gamma_13,
                gamma_14=gamma_14,
                gamma_23=gamma_23,
                gamma_24=gamma_24,
                gamma_2d=gamma_2d,
                gamma_3d=gamma_3d,
                gamma_4d=gamma_4d,
                C_14=C_14,
                C_24=C_24,
                T=T
            )

        # convolve
        ensemble_signal = signal.fftconvolve(atom_scan, kernel, mode='valid')

        # return results
        delta_mu = np.linspace(delta_mu_min, delta_mu_max, delta_mu_points)
        delta_B = np.linspace(delta_B_min, delta_B_max, delta_B_points)
        delta_mu, delta_B = np.meshgrid(delta_mu, delta_B, indexing='ij')
        return delta_mu, delta_B, ensemble_signal
    |

## Biphoton Generation

These codes implement the three-level biphoton generation models in \ref{ch:biphoton_generation}.

### Steady State

This Python code implements the steady-state model. When run as a script, it performs root finding for cavity steady-states, and prints the (non-convergent) results of this root finding, for two sets of detuning parameters.

    \
    \verb|import numpy as np
    import matplotlib.pyplot as plt
    from scipy import optimize, stats
    import sympy

    hbar = 1.054571817e-34
    kB = 1.380649e-23

    gauss_lobatto_nodes_weights = dict()

    for n in range(2, 100):
        legendre_coeffs = (0,)*(n-1) + (1,)
        legendre_poly = np.polynomial.Legendre(legendre_coeffs)
        nodes = legendre_poly.deriv().roots()
        nodes = np.concatenate([[-1.0], nodes, [1.0]])
        weights = 2 / (n*(n-1)*legendre_poly(nodes)**2)

        # convert from [-1, 1] to [0, 1]
        nodes = (nodes+1)/2
        weights = weights/2

        gauss_lobatto_nodes_weights[n] = (nodes, weights)

    def composite_gauss_lobatto_nodes_weights(n, points):
        if n not in gauss_lobatto_nodes_weights:
            raise ValueError

        num_points = len(points)
        if num_points < 2:
            raise ValueError

        base_nodes, base_weights = gauss_lobatto_nodes_weights[n]
        if num_points == 2:
            L = points[1] - points[0]
            nodes = points[0] + base_nodes*L
            weights = base_weights*L
            return nodes, weights

        count = (num_points-1)*(n-1) + 1
        nodes = np.zeros(count)
        weights = np.zeros(count)

        # first and last node
        nodes[0] = points[0]
        weights[0] = base_weights[0] * (points[1]-points[0])
        nodes[1] = points[-1]
        weights[1] = base_weights[-1] * (points[-1]-points[-2])
        filled = 2

        # nodes at intermediate points
        for i in range(1,num_points-1):
            nodes[filled] = points[i]
            weights[filled] = base_weights[-1]*(points[i]-points[i-1])
            weights[filled] += base_weights[0]*(points[i+1]-points[i])
            filled += 1

        # nodes between points
        if n == 2:
            return nodes, weights
        for i in range(num_points-1):
            L = points[i+1]-points[i]
            nodes[filled:filled+n-2] = points[i] + base_nodes[1:-1]*L
            weights[filled:filled+n-2] = base_weights[1:-1]*L
            filled += n-2

        return nodes, weights

    # unit matrices
    s12 = sympy.Matrix([[0, 1, 0], [0, 0, 0], [0, 0, 0]])
    s13 = sympy.Matrix([[0, 0, 1], [0, 0, 0], [0, 0, 0]])
    s21 = s12.H
    s31 = s13.H
    s11 = s12*s21
    s22 = s21*s12
    s23 = s21*s13
    s32 = s23.H
    s33 = s31*s13

    # Hamiltonians for Lambda-system and V-system

    Omega_p = sympy.symbols('Omega_p')
    Omega_o = sympy.symbols('Omega_o')
    Omega_mu = sympy.symbols('Omega_mu')
    delta_o = sympy.symbols('delta_o', real=True)
    delta_mu = sympy.symbols('delta_mu', real=True)

    HLambda = s21*Omega_mu + s32*Omega_o + s31*Omega_p
    HLambda += HLambda.H
    HLambda += s22*delta_mu + s33*(delta_mu+delta_o)

    HV = s21*Omega_o + s32*Omega_mu + s31*Omega_p
    HV += HV.H
    HV += s22*delta_o + s33*(delta_mu+delta_o)

    systems = ('Lambda', 'V')
    H = {'Lambda': HLambda, 'V': HV}

    # decomposition of Hamiltonians into linear components
    Omega_or, Omega_oi = sympy.symbols('Omega_or Omega_oi', real=True)
    Omega_mur, Omega_mui = sympy.symbols('Omega_\\mu\\ r Omega_\\mu\\ i', real=True)
    subs = [
        (Omega_o, Omega_or+sympy.I*Omega_oi),
        (Omega_mu, Omega_mur+sympy.I*Omega_mui)
    ]

    H_subs = dict()
    for sys in systems:
        H_subs[sys] = H[sys].subs(subs)

    H0 = dict()
    Hor = dict()
    Hoi = dict()
    Hmur = dict()
    Hmui = dict()
    for sys in systems:
        H0[sys] = H[sys].subs([(Omega_o, 0), (Omega_mu, 0)])
        Hor[sys] = sympy.diff(H_subs[sys], Omega_or)
        Hoi[sys] = sympy.diff(H_subs[sys], Omega_oi)
        Hmur[sys] = sympy.diff(H_subs[sys], Omega_mur)
        Hmui[sys] = sympy.diff(H_subs[sys], Omega_mui)

    # discriminant of characteristic polynomial of Hamiltonian

    Delta_poly_coeffs = dict()
    for sys in systems:
        Delta = H[sys].charpoly().discriminant()
        Delta_poly = sympy.poly(Delta, delta_mu, delta_o)
        n_delta_mu = Delta_poly.degree(delta_mu)
        n_delta_o = Delta_poly.degree(delta_o)
        Delta_poly_coeffs[sys] = sympy.zeros(n_delta_mu+1, n_delta_o+1)
        for (i, j), coeff in zip(Delta_poly.monoms(), Delta_poly.coeffs()):
            Delta_poly_coeffs[sys][i,j] = coeff

    # liouvillan superoperators with loss

    gamma_12, gamma_13 = sympy.symbols('gamma_1(2:4)', real=True, negative=False)
    gamma_23 = sympy.symbols('gamma_23', real=True, negative=False)
    gamma_2d, gamma_3d = sympy.symbols('gamma_(2:4)d', real=True, negative=False)
    n_b = sympy.symbols('n_b', real=True, negative=False)

    def loss_operator_common(rho):
        L13 = gamma_13/2 * (2*s13*rho*s31 - rho*s33 - s33*rho)
        L2d = gamma_2d/2 * (2*s22*rho*s22 - rho*s22 - s22*rho)
        L3d = gamma_3d/2 * (2*s33*rho*s33 - rho*s33 - s33*rho)
        return L13 + L2d + L3d

    def loss_operator_Lambda(rho):
        L12 = gamma_12*(n_b+1)/2 * (2*s12*rho*s21 - rho*s22 - s22*rho)
        L21 = gamma_12*n_b/2 * (2*s21*rho*s12 - rho*s11 - s11*rho)
        L23 = gamma_23/2 * (2*s23*rho*s32 - rho*s33 - s33*rho)
        return L12 + L21 + L23 + loss_operator_common(rho)

    def loss_operator_V(rho):
        L12 = gamma_12/2 * (2*s12*rho*s21 - rho*s22 - s22*rho)
        L23 = gamma_23*(n_b+1)/2 * (2*s23*rho*s32 - rho*s33 - s33*rho)
        L32 = gamma_23*n_b/2 * (2*s32*rho*s23 - rho*s22 - s22*rho)
        return L12 + L23 + L32 + loss_operator_common(rho)

    loss_operator = {'Lambda': loss_operator_Lambda, 'V': loss_operator_V}

    def liouvillan_superoperator(H, rho, loss=None):
        Lrho = -sympy.I*(H*rho - rho*H)
        if loss is not None:
            Lrho += loss(rho)
        return Lrho

    # liouvillan matrices
    def flattening_indices(order):
        n = order.shape[0]
        unflatten = order

        flatten_i = np.zeros(n**2, dtype=int)
        flatten_j = np.zeros(n**2, dtype=int)
        for i in range(n):
            for j in range(n):
                k = order[i,j]
                flatten_i[k] = i
                flatten_j[k] = j
        flatten = (flatten_i, flatten_j)

        return unflatten, flatten

    def liouvillan_matrix(H, order, loss=None):
        n = order.shape[0]
        mtx = sympy.zeros(n**2)
        for icol in range(n):
            for jcol in range(n):
                col = order[icol,jcol]
                rho = sympy.zeros(n)
                rho[icol,jcol] = 1
                Lrho = liouvillan_superoperator(H, rho, loss=loss)

                for irow in range(n):
                    for jrow in range(n):
                        row = order[irow,jrow]
                        mtx[row,col] = Lrho[irow,jrow]
        return mtx

    order = np.array([
        [0, 1, 2],
        [3, 4, 5],
        [6, 7, 8]
    ])
    unflatten, flatten = flattening_indices(order)

    L0 = dict()
    Lor = dict()
    Loi = dict()
    Lmur = dict()
    Lmui = dict()
    L = dict()
    for sys in systems:
        L0[sys] = liouvillan_matrix(H0[sys], order, loss=loss_operator[sys])
        Lor[sys] = liouvillan_matrix(Hor[sys], order)
        Loi[sys] = liouvillan_matrix(Hoi[sys], order)
        Lmur[sys] = liouvillan_matrix(Hmur[sys], order)
        Lmui[sys] = liouvillan_matrix(Hmui[sys], order)
        L[sys] = liouvillan_matrix(H[sys], order, loss=loss_operator[sys])

    # real matrices

    def hermitian_complex_to_real(order):
        n = order.shape[0]
        C = sympy.zeros(n**2)

        # diagonals are already real, so keep them as-is
        for k in range(n):
            ik = order[k,k]
            C[ik,ik] = 1

        # transform off-diagonal pairs from (z, z*) to (Re z, Im z)
        for j in range(n-1):
            for k in range(j+1, n):
                i_upper = order[j,k]
                i_lower = order[k,j]
                C[i_upper,i_upper] = sympy.Rational(1, 2)
                C[i_upper,i_lower] = sympy.Rational(1, 2)
                C[i_lower,i_upper] = -sympy.I/2
                C[i_lower,i_lower] = sympy.I/2

        return C

    CtoR = hermitian_complex_to_real(order)
    RtoC = CtoR.inv()

    L0_real = dict()
    Lor_real = dict()
    Loi_real = dict()
    Lmur_real = dict()
    Lmui_real = dict()
    L_real = dict()
    for sys in systems:
        L0_real[sys] = sympy.re(CtoR * L0[sys] * RtoC)
        Lor_real[sys] = sympy.re(CtoR * Lor[sys] * RtoC)
        Loi_real[sys] = sympy.re(CtoR * Loi[sys] * RtoC)
        Lmur_real[sys] = sympy.re(CtoR * Lmur[sys] * RtoC)
        Lmui_real[sys] = sympy.re(CtoR * Lmui[sys] * RtoC)
        L_real[sys] = sympy.re(CtoR * L[sys] * RtoC)

    # numerical matrices

    def lambdify_wrapper(args, expr):
        return sympy.lambdify(args, expr, 'numpy', cse=True, docstring_limit=0)

    def numerify_zero_args_expr(expr):
        # evil hack
        x = sympy.symbols('x')
        return lambdify_wrapper(x, expr)(None)

    CtoR_num = numerify_zero_args_expr(CtoR)
    RtoC_num = numerify_zero_args_expr(RtoC)

    args_d = (Omega_p, Omega_mu, Omega_o)
    args_h0 = (delta_mu, delta_o, Omega_p)
    args_in = (Omega_mu, Omega_o)
    args_decay = (gamma_12, gamma_13, gamma_23, gamma_2d, gamma_3d, n_b)

    args_h = args_h0 + args_in
    args_l0 = args_decay + args_h0
    args_l = args_l0 + args_in

    H0_func = dict()
    Hor_num = dict()
    Hoi_num = dict()
    Hmur_num = dict()
    Hmui_num = dict()
    H_func = dict()
    for sys in systems:
        H0_func[sys] = lambdify_wrapper(args_h0, H0[sys])
        Hor_num[sys] = numerify_zero_args_expr(Hor[sys])
        Hoi_num[sys] = numerify_zero_args_expr(Hoi[sys])
        Hmur_num[sys] = numerify_zero_args_expr(Hmur[sys])
        Hmui_num[sys] = numerify_zero_args_expr(Hmui[sys])
        H_func[sys] = lambdify_wrapper(args_h, H[sys])

    Delta_poly_coeffs_func = dict()
    for sys in systems:
        Delta_poly_coeffs_func[sys] = lambdify_wrapper(args_d, Delta_poly_coeffs[sys])

    L0_func = dict()
    Lor_num = dict()
    Loi_num = dict()
    Lmur_num = dict()
    Lmui_num = dict()
    L_func = dict()
    for sys in systems:
        L0_func[sys] = lambdify_wrapper(args_l0, L0[sys])
        Lor_num[sys] = numerify_zero_args_expr(Lor[sys])
        Loi_num[sys] = numerify_zero_args_expr(Loi[sys])
        Lmur_num[sys] = numerify_zero_args_expr(Lmur[sys])
        Lmui_num[sys] = numerify_zero_args_expr(Lmui[sys])
        L_func[sys] = lambdify_wrapper(args_l, L[sys])

    L0_real_func = dict()
    Lor_real_num = dict()
    Loi_real_num = dict()
    Lmur_real_num = dict()
    Lmui_real_num = dict()
    L_real_func = dict()
    for sys in systems:
        L0_real_func[sys] = lambdify_wrapper(args_l0, L0_real[sys])
        Lor_real_num[sys] = numerify_zero_args_expr(Lor_real[sys])
        Loi_real_num[sys] = numerify_zero_args_expr(Loi_real[sys])
        Lmur_real_num[sys] = numerify_zero_args_expr(Lmur_real[sys])
        Lmui_real_num[sys] = numerify_zero_args_expr(Lmui_real[sys])
        L_real_func[sys] = lambdify_wrapper(args_l, L_real[sys])

    def poly_bivariate_coeffs_diff_left(coeffs):
        nx = coeffs.shape[0]-1
        ny = coeffs.shape[1]-1
        i, j = np.indices((nx, ny+1))
        return (i+1) * coeffs[1:,:]

    def poly_bivariate_coeffs_diff_right(coeffs):
        nx = coeffs.shape[0]-1
        ny = coeffs.shape[1]-1
        i, j = np.indices((nx+1, ny))
        return (j+1) * coeffs[:,1:]

    def poly_coeffs_diff(coeffs):
        n = len(coeffs)-1
        k = np.arange(1, n+1)
        return k * coeffs[1:]

    def poly_bivariate_coeffs_evaluate_left(coeffs, x):
        n = coeffs.shape[0]-1
        k = np.arange(n+1)
        return (x**k) @ coeffs

    def poly_bivariate_coeffs_evaluate_right(coeffs, y):
        n = coeffs.shape[1]-1
        k = np.arange(n+1)
        return coeffs @ (y**k)

    def poly_bivariate_coeffs_evaluate(coeffs, x, y):
        i, j = np.indices(coeffs.shape)
        return np.sum(coeffs * x**i * y**j)

    def poly_coeffs_evaluate(coeffs, x):
        n = len(coeffs)-1
        k = np.arange(n+1)
        return np.sum(coeffs * x**k)

    def poly_coeffs_roots(coeffs):
        roots = np.polynomial.polynomial.polyroots(coeffs)
        roots = np.unique(roots)
        is_real = (np.imag(roots)==0)
        return np.real(roots[is_real])

    def rho_steady_state(gamma_12, gamma_13, gamma_23, gamma_2d, gamma_3d, n_b,
                         delta_mu, delta_o, Omega_p, Omega_mu, Omega_o, sys):
        L_mtx = L_real_func[sys](
            gamma_12=gamma_12,
            gamma_13=gamma_13,
            gamma_23=gamma_23,
            gamma_2d=gamma_2d,
            gamma_3d=gamma_3d,
            n_b=n_b,
            delta_mu=delta_mu,
            delta_o=delta_o,
            Omega_p=Omega_p,
            Omega_mu=Omega_mu,
            Omega_o=Omega_o
        )
        L_mtx[0,:] = np.identity(3)[flatten]
        b = np.zeros(9)
        b[0] = 1
        rho_real = np.linalg.solve(L_mtx, b)
        rho = RtoC_num @ rho_real
        rho = rho[unflatten]
        return rho

    def rho_linear_steady_state_components(gamma_12, gamma_13, gamma_23, gamma_2d, gamma_3d,
                                           n_b, delta_mu, delta_o, Omega_p, sys):
        L0_mtx = L0_real_func[sys](
            gamma_12=gamma_12,
            gamma_13=gamma_13,
            gamma_23=gamma_23,
            gamma_2d=gamma_2d,
            gamma_3d=gamma_3d,
            n_b=n_b,
            delta_mu=delta_mu,
            delta_o=delta_o,
            Omega_p=Omega_p
        )
        L0_mtx[0,:] = np.identity(3)[flatten]

        b = np.zeros(9)
        b[0] = 1
        rho_0_real = np.linalg.solve(L0_mtx, b)
        rho_0 = RtoC_num @ rho_0_real

        b = -Lor_real_num[sys] @ rho_0_real
        b[0] = 0
        rho_or_real = np.linalg.solve(L0_mtx, b)
        rho_or = RtoC_num @ rho_or_real

        b = -Loi_real_num[sys] @ rho_0_real
        b[0] = 0
        rho_oi_real = np.linalg.solve(L0_mtx, b)
        rho_oi = RtoC_num @ rho_oi_real

        b = -Lmur_real_num[sys] @ rho_0_real
        b[0] = 0
        rho_mur_real = np.linalg.solve(L0_mtx, b)
        rho_mur = RtoC_num @ rho_mur_real

        b = -Lmui_real_num[sys] @ rho_0_real
        b[0] = 0
        rho_mui_real = np.linalg.solve(L0_mtx, b)
        rho_mui = RtoC_num @ rho_mui_real

        rho_0 = rho_0[unflatten]
        rho_or = rho_or[unflatten]
        rho_oi = rho_oi[unflatten]
        rho_mur = rho_mur[unflatten]
        rho_mui = rho_mui[unflatten]
        return rho_0, rho_or, rho_oi, rho_mur, rho_mui

    def rho_steady_state_ensemble(gamma_12, gamma_13, gamma_23, gamma_2d, gamma_3d,
                                  n_b, delta_mu, delta_o, Omega_p, Omega_mu,
                                  Omega_o, sigma_mu, sigma_o, sys, logging=False):
        gauss_lobatto_order = 20
        gamma_oh = gamma_2d + (gamma_3d if sys=='Lambda' else 0)
        gamma_muh = gamma_2d + (gamma_3d if sys=='V' else 0)

        # distribution envelope
        G_mu = stats.norm(loc=delta_mu, scale=sigma_mu).pdf
        G_o = stats.norm(loc=delta_o, scale=sigma_o).pdf
        G = lambda dp_mu, dp_o: G_mu(dp_mu) * G_o(dp_o)

        # set up for curve finding
        Delta_coeffs_2 = Delta_poly_coeffs_func[sys](
            Omega_p=Omega_p,
            Omega_mu=Omega_mu,
            Omega_o=Omega_o
        )
        dmu_Delta_coeffs_2 = poly_bivariate_coeffs_diff_left(Delta_coeffs_2)
        do_Delta_coeffs_2 = poly_bivariate_coeffs_diff_right(Delta_coeffs_2)

        # get integral points

        node_deltap_o = np.array([], dtype=float)
        node_deltap_mu = np.array([], dtype=float)
        node_weight = np.array([], dtype=float)

        # set up outer integral
        deltap_o_intervals = [delta_o + z*sigma_o
                              for z in (-10, -3, -1, 0, 1, 3, 10)]
        deltap_o_intervals += [z*gamma_oh for z in (-5, -1, 0, 1, 5)]
        deltap_o_intervals = sorted(deltap_o_intervals)
        deltap_o_nodes, deltap_o_weights = composite_gauss_lobatto_nodes_weights(
            gauss_lobatto_order, deltap_o_intervals)

        # set up inner integral
        for dp_o, w_o in zip(deltap_o_nodes, deltap_o_weights):
            dmu_Delta_coeffs = poly_bivariate_coeffs_evaluate_right(dmu_Delta_coeffs_2, dp_o)
            do_Delta_coeffs = poly_bivariate_coeffs_evaluate_right(do_Delta_coeffs_2, dp_o)
            dmu_critical = np.concatenate([
                poly_coeffs_roots(dmu_Delta_coeffs),
                poly_coeffs_roots(do_Delta_coeffs)
            ])

            deltap_mu_intervals = [delta_mu + z*sigma_mu
                                   for z in (-10, -3, -1, 1, 3, 10)]
            deltap_mu_intervals += [z*gamma_muh for z in (-5, -1, 0, 1, 5)]
            deltap_mu_intervals += list(dmu_critical)

            deltap_mu_intervals = sorted(deltap_mu_intervals)
            deltap_mu_nodes, deltap_mu_weights = composite_gauss_lobatto_nodes_weights(
                gauss_lobatto_order, deltap_mu_intervals)

            node_deltap_o = np.concatenate([node_deltap_o, np.full_like(deltap_mu_nodes, dp_o)])
            node_deltap_mu = np.concatenate([node_deltap_mu, deltap_mu_nodes])
            node_weight = np.concatenate([node_weight, w_o*deltap_mu_weights])

        node_G = G(node_deltap_mu, node_deltap_o)

        # perform integral

        integral = np.zeros((3,3), dtype=complex)
        for dp_o, dp_mu, w, g in zip(node_deltap_o, node_deltap_mu, node_weight, node_G):
            integral += w*g * rho_steady_state(
                gamma_12=gamma_12,
                gamma_13=gamma_13,
                gamma_23=gamma_23,
                gamma_2d=gamma_2d,
                gamma_3d=gamma_3d,
                n_b=n_b,
                delta_mu=dp_mu,
                delta_o=dp_o,
                Omega_p=Omega_p,
                Omega_mu=Omega_mu,
                Omega_o=Omega_o,
                sys=sys
            )

        if logging:
            return (node_deltap_mu, node_deltap_o, node_weight, node_G), integral
        else:
            return integral

    def Omega_cavity(g, alpha):
        return g * np.exp(1j*np.angle(alpha)) * np.sqrt(np.abs(alpha)**2 + 1)

    def cavity_langevin_diff(gamma_12, gamma_13, gamma_23, gamma_2d, gamma_3d,
                             n_b, gamma_oi, gamma_oc, gamma_mui, gamma_muc,
                             N_o, N_mu, g_o, g_mu, Omega_p, alpha, beta,
                             alpha_in, beta_in, delta_mu, delta_o, delta_co,
                             delta_cmu, sigma_mu, sigma_o, sys, return_S=False):
        S = rho_steady_state_ensemble(
            gamma_12=gamma_12,
            gamma_13=gamma_13,
            gamma_23=gamma_23,
            gamma_2d=gamma_2d,
            gamma_3d=gamma_3d,
            n_b=n_b,
            delta_mu=delta_mu,
            delta_o=delta_o,
            Omega_p=Omega_p,
            Omega_o=Omega_cavity(g_o, alpha),
            Omega_mu=Omega_cavity(g_mu, beta),
            sigma_o=sigma_o,
            sigma_mu=sigma_mu,
            sys=sys
        )
        S_alpha = N_o*np.conj(g_o) * (S[1,0] if sys=='V' else S[2,1])
        S_beta = N_mu*np.conj(g_mu) * (S[2,1] if sys=='V' else S[1,0])

        d_alpha_S = -1j*S_alpha
        d_beta_S = -1j*S_beta
        d_alpha_not_S = -1j*delta_co*alpha - (gamma_oi+gamma_oc)*alpha/2 + np.sqrt(gamma_oc)*alpha_in
        d_beta_not_S = -1j*delta_cmu*beta - (gamma_mui+gamma_muc)*beta/2 + np.sqrt(gamma_muc)*beta_in
        d_alpha = d_alpha_S + d_alpha_not_S
        d_beta = d_beta_S + d_beta_not_S

        if return_S:
            return d_alpha_S, d_beta_S, d_alpha_not_S, d_beta_not_S
        else:
            return d_alpha, d_beta

    def cavity_steady_state(gamma_12, gamma_13, gamma_23, gamma_2d, gamma_3d,
                            n_b, gamma_oi, gamma_oc, gamma_mui, gamma_muc,
                            N_o, N_mu, g_o, g_mu, Omega_p, alpha_in, beta_in,
                            delta_mu, delta_o, delta_co, delta_cmu, sigma_mu,
                            sigma_o, sys, alpha_0=1, beta_0=1):
        def pack_vector(alpha, beta):
            vec = np.zeros(4)
            vec[0] = np.real(alpha)
            vec[1] = np.imag(alpha)
            vec[2] = np.real(beta)
            vec[3] = np.imag(beta)
            return vec

        def unpack_vector(vec):
            alpha_r = vec[0]
            alpha_i = vec[1]
            beta_r = vec[2]
            beta_i = vec[3]
            alpha = alpha_r + 1j*alpha_i
            beta = beta_r + 1j*beta_i
            return alpha, beta

        def diff_func(vec):
            alpha, beta = unpack_vector(vec)
            d_alpha, d_beta = cavity_langevin_diff(
                gamma_12=gamma_12,
                gamma_13=gamma_13,
                gamma_23=gamma_23,
                gamma_2d=gamma_2d,
                gamma_3d=gamma_3d,
                n_b=n_b,
                gamma_oi=gamma_oi,
                gamma_oc=gamma_oc,
                gamma_mui=gamma_mui,
                gamma_muc=gamma_muc,
                N_o=N_o,
                N_mu=N_mu,
                g_o=g_o,
                g_mu=g_mu,
                Omega_p=Omega_p,
                alpha=alpha,
                beta=beta,
                alpha_in=alpha_in,
                beta_in=beta_in,
                delta_mu=delta_mu,
                delta_o=delta_o,
                delta_co=delta_co,
                delta_cmu=delta_cmu,
                sigma_mu=sigma_mu,
                sigma_o=sigma_o,
                sys=sys
            )

            d_vec = pack_vector(d_alpha, d_beta)
            return d_vec

        v0 = pack_vector(alpha_0, beta_0)
        result = optimize.root(diff_func, v0)#, method='broyden1', options={'ftol':1e-12})
        return result

    def planck_excitation(T, omega):
        return 1 / np.expm1(hbar*omega/(kB*T))

    if __name__ == '__main__':
        sys = 'Lambda'
        omega_12 = 2*np.pi*5.186e9
        d13 = 1.63e-32
        d23 = 1.15e-32
        tau_12 = 11
        tau_3 = 0.011
        gamma_2d = 1e6
        gamma_3d = 1e6
        sigma_o = 2*np.pi*419e6
        sigma_mu = 2*np.pi*5e6
        N = 1e16
        gamma_oi = 2*np.pi*7.95e6
        gamma_oc = 2*np.pi*1.7e6
        gamma_mui = 2*np.pi*650e3
        gamma_muc = 2*np.pi*1.5e6
        g_o = 51.9
        g_mu = 1.04

        T = 4.6
        n_b = planck_excitation(T, omega_12)
        tau_13 = tau_3 * d13**2 / (d13**2 + d23**2)
        tau_23 = tau_3 * d23**2 / (d13**2 + d23**2)
        gamma_12 = 1 / (tau_12*(n_b+1))
        gamma_13 = 1 / tau_13
        gamma_23 = 1 / tau_23
        N_o = N
        N_mu = N

        Omega_p = 35000.0

        delta_o_small = -100e3
        delta_mu_small = 1e6
        delta_o_large = -6.5*sigma_o
        delta_mu_large = 8*sigma_mu

        iterator = [(delta_o_small, delta_mu_small, 'Small detuning'),
                    (delta_o_large, delta_mu_large, 'Large detuning')]
        for delta_o, delta_mu, regime in iterator:
            print(f'\n{regime} regime')
            result = cavity_steady_state(
                gamma_12=gamma_12,
                gamma_13=gamma_13,
                gamma_23=gamma_23,
                gamma_2d=gamma_2d,
                gamma_3d=gamma_3d,
                n_b=n_b,
                gamma_oi=gamma_oi,
                gamma_oc=gamma_oc,
                gamma_mui=gamma_mui,
                gamma_muc=gamma_muc,
                N_o=N_o,
                N_mu=N_mu,
                g_o=g_o,
                g_mu=g_mu,
                Omega_p=Omega_p,
                alpha_in=0.0,
                beta_in=0.0,
                delta_mu=delta_mu,
                delta_o=delta_o,
                delta_co=0.0,
                delta_cmu=0.0,
                sigma_mu=sigma_mu,
                sigma_o=sigma_o,
                sys=sys,
                alpha_0=1,
                beta_0=1
            )
            print(result)
    |

### Super-Atom Dynamics

This CUDA code implements the super-atom dynamical model. When compiled and ran, it performs a simulation and saves the results as binary files in the working directory.

    \
    \verb|// compile: nvcc biphoton_super_atom.cu -o sim -O3 -rdc=true -lm -arch=sm_60
    #include
    #include
    #include
    #include
    #include
    #include

    #define PI 3.1415926535897932384626433832795
    const double HBAR = 1.054571817e-34;
    const double K_B = 1.380649e-23;

    #define NUM_THREAD_BLOCKS 256
    #define NUM_THREADS_IN_BLOCK 512

    enum SystemType
    {
        LAMBDA_SYSTEM,
        V_SYSTEM
    };
    typedef enum SystemType SystemType;

    struct DensityMatrix
    {
        double r11;
        double r22;
        double r33;
        double r12;
        double i12;
        double r13;
        double i13;
        double r23;
        double i23;
    };
    typedef struct DensityMatrix DensityMatrix;

    const DensityMatrix GROUND_STATE_MATRIX = {
        1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
    };

    struct AtomSample
    {
        double weight;
        double g_or;
        double g_oi;
        double g_mur;
        double g_mui;
        double delta_12;
        double delta_23;
    };
    typedef struct AtomSample AtomSample;

    struct SystemState
    {
        double alpha_r;
        double alpha_i;
        double beta_r;
        double beta_i;
        DensityMatrix rho[];
    };
    typedef struct SystemState SystemState;

    size_t SizeofStateStruct(size_t nAtomSamples)
    {
        return sizeof(SystemState) + nAtomSamples * sizeof(DensityMatrix);
    }

    bool HostAllocateStateStruct(SystemState **ptr, size_t nAtomSamples)
    {
        *ptr = (SystemState *) malloc(SizeofStateStruct(nAtomSamples));
        return (*ptr != NULL);
    }

    bool DeviceAllocateStateStruct(SystemState **ptr, size_t nAtomSamples)
    {
        cudaError_t result = cudaMalloc(ptr, SizeofStateStruct(nAtomSamples));
        return (result == cudaSuccess);
    }

    void CopyStateStruct(SystemState *dst, const SystemState *src,
                         size_t nAtomSamples, cudaMemcpyKind kind)
    {
        cudaMemcpy(dst, src, SizeofStateStruct(nAtomSamples), kind);
    }

    bool HostAllocateAtomSamples(AtomSample **ptr, size_t nAtomSamples)
    {
        size_t size = nAtomSamples * sizeof(AtomSample);
        *ptr = (AtomSample *) malloc(size);
        return (*ptr != NULL);
    }

    bool DeviceAllocateAtomSamples(AtomSample **ptr, size_t nAtomSamples)
    {
        size_t size = nAtomSamples * sizeof(AtomSample);
        cudaError_t result = cudaMalloc(ptr, size);
        return (result == cudaSuccess);
    }

    void CopyAtomSamples(AtomSample *dst, const AtomSample *src,
                         size_t nAtomSamples, cudaMemcpyKind kind)
    {
        cudaMemcpy(dst, src, nAtomSamples * sizeof(AtomSample), kind);
    }

    __device__
    void MasterDerivative(DensityMatrix *diff, SystemType sys,
        const DensityMatrix *rho, double Omega_mur, double Omega_mui,
        double Omega_or, double Omega_oi, double Omega_pr, double Omega_pi,
        double deltap_mu, double deltap_o, double n_b, double gamma_12,
        double gamma_13, double gamma_23, double gamma_2d, double gamma_3d)
    {
        DensityMatrix tmp = *rho;

        // code generated using computer algebra
        if (sys == LAMBDA_SYSTEM)
        {
            diff->r11 = -2*Omega_mui*tmp.r12 - 2*Omega_mur*tmp.i12 - 2*Omega_pi*tmp.r13 - 2*Omega_pr*tmp.i13 - gamma_12*n_b*tmp.r11 + gamma_12*tmp.r22*(n_b + 1) + gamma_13*tmp.r33;
            diff->r22 = 2*Omega_mui*tmp.r12 + 2*Omega_mur*tmp.i12 - 2*Omega_oi*tmp.r23 - 2*Omega_or*tmp.i23 + gamma_12*n_b*tmp.r11 - gamma_12*tmp.r22*(n_b + 1) + gamma_23*tmp.r33;
            diff->r33 = 2*Omega_oi*tmp.r23 + 2*Omega_or*tmp.i23 + 2*Omega_pi*tmp.r13 + 2*Omega_pr*tmp.i13 - gamma_13*tmp.r33 - gamma_23*tmp.r33;
            diff->r12 = Omega_mui*tmp.r11 - Omega_mui*tmp.r22 - Omega_oi*tmp.r13 - Omega_or*tmp.i13 - Omega_pi*tmp.r23 - Omega_pr*tmp.i23 - deltap_mu*tmp.i12 - 1.0/2.0*gamma_12*n_b*tmp.r12 - 1.0/2.0*gamma_12*tmp.r12*(n_b + 1) - 1.0/2.0*gamma_2d*tmp.r12;
            diff->i12 = Omega_mur*tmp.r11 - Omega_mur*tmp.r22 - Omega_oi*tmp.i13 + Omega_or*tmp.r13 + Omega_pi*tmp.i23 - Omega_pr*tmp.r23 + deltap_mu*tmp.r12 - 1.0/2.0*gamma_12*n_b*tmp.i12 - 1.0/2.0*gamma_12*tmp.i12*(n_b + 1) - 1.0/2.0*gamma_2d*tmp.i12;
            diff->r13 = -Omega_mui*tmp.r23 + Omega_mur*tmp.i23 + Omega_oi*tmp.r12 - Omega_or*tmp.i12 + Omega_pi*tmp.r11 - Omega_pi*tmp.r33 - 1.0/2.0*gamma_12*n_b*tmp.r13 - 1.0/2.0*gamma_13*tmp.r13 - 1.0/2.0*gamma_23*tmp.r13 - 1.0/2.0*gamma_3d*tmp.r13 + tmp.i13*(-deltap_mu - deltap_o);
            diff->i13 = -Omega_mui*tmp.i23 - Omega_mur*tmp.r23 + Omega_oi*tmp.i12 + Omega_or*tmp.r12 + Omega_pr*tmp.r11 - Omega_pr*tmp.r33 - 1.0/2.0*gamma_12*n_b*tmp.i13 - 1.0/2.0*gamma_13*tmp.i13 - 1.0/2.0*gamma_23*tmp.i13 - 1.0/2.0*gamma_3d*tmp.i13 - tmp.r13*(-deltap_mu - deltap_o);
            diff->r23 = Omega_mui*tmp.r13 + Omega_mur*tmp.i13 + Omega_oi*tmp.r22 - Omega_oi*tmp.r33 + Omega_pi*tmp.r12 + Omega_pr*tmp.i12 + deltap_mu*tmp.i23 - 1.0/2.0*gamma_12*tmp.r23*(n_b + 1) - 1.0/2.0*gamma_13*tmp.r23 - 1.0/2.0*gamma_23*tmp.r23 - 1.0/2.0*gamma_2d*tmp.r23 - 1.0/2.0*gamma_3d*tmp.r23 + tmp.i23*(-deltap_mu - deltap_o);
            diff->i23 = Omega_mui*tmp.i13 - Omega_mur*tmp.r13 + Omega_or*tmp.r22 - Omega_or*tmp.r33 - Omega_pi*tmp.i12 + Omega_pr*tmp.r12 - deltap_mu*tmp.r23 - 1.0/2.0*gamma_12*tmp.i23*(n_b + 1) - 1.0/2.0*gamma_13*tmp.i23 - 1.0/2.0*gamma_23*tmp.i23 - 1.0/2.0*gamma_2d*tmp.i23 - 1.0/2.0*gamma_3d*tmp.i23 - tmp.r23*(-deltap_mu - deltap_o);
        }
        else // sys == V_SYSTEM
        {
            diff->r11 = -2*Omega_oi*tmp.r12 - 2*Omega_or*tmp.i12 - 2*Omega_pi*tmp.r13 - 2*Omega_pr*tmp.i13 + gamma_12*tmp.r22 + gamma_13*tmp.r33;
            diff->r22 = -2*Omega_mui*tmp.r23 - 2*Omega_mur*tmp.i23 + 2*Omega_oi*tmp.r12 + 2*Omega_or*tmp.i12 - gamma_12*tmp.r22 - gamma_23*n_b*tmp.r22 + gamma_23*tmp.r33*(n_b + 1);
            diff->r33 = 2*Omega_mui*tmp.r23 + 2*Omega_mur*tmp.i23 + 2*Omega_pi*tmp.r13 + 2*Omega_pr*tmp.i13 - gamma_13*tmp.r33 + gamma_23*n_b*tmp.r22 - gamma_23*tmp.r33*(n_b + 1);
            diff->r12 = -Omega_mui*tmp.r13 - Omega_mur*tmp.i13 + Omega_oi*tmp.r11 - Omega_oi*tmp.r22 - Omega_pi*tmp.r23 - Omega_pr*tmp.i23 - deltap_o*tmp.i12 - 1.0/2.0*gamma_12*tmp.r12 - 1.0/2.0*gamma_23*n_b*tmp.r12 - 1.0/2.0*gamma_2d*tmp.r12;
            diff->i12 = -Omega_mui*tmp.i13 + Omega_mur*tmp.r13 + Omega_or*tmp.r11 - Omega_or*tmp.r22 + Omega_pi*tmp.i23 - Omega_pr*tmp.r23 + deltap_o*tmp.r12 - 1.0/2.0*gamma_12*tmp.i12 - 1.0/2.0*gamma_23*n_b*tmp.i12 - 1.0/2.0*gamma_2d*tmp.i12;
            diff->r13 = Omega_mui*tmp.r12 - Omega_mur*tmp.i12 - Omega_oi*tmp.r23 + Omega_or*tmp.i23 + Omega_pi*tmp.r11 - Omega_pi*tmp.r33 - 1.0/2.0*gamma_13*tmp.r13 - 1.0/2.0*gamma_23*tmp.r13*(n_b + 1) - 1.0/2.0*gamma_3d*tmp.r13 + tmp.i13*(-deltap_mu - deltap_o);
            diff->i13 = Omega_mui*tmp.i12 + Omega_mur*tmp.r12 - Omega_oi*tmp.i23 - Omega_or*tmp.r23 + Omega_pr*tmp.r11 - Omega_pr*tmp.r33 - 1.0/2.0*gamma_13*tmp.i13 - 1.0/2.0*gamma_23*tmp.i13*(n_b + 1) - 1.0/2.0*gamma_3d*tmp.i13 - tmp.r13*(-deltap_mu - deltap_o);
            diff->r23 = Omega_mui*tmp.r22 - Omega_mui*tmp.r33 + Omega_oi*tmp.r13 + Omega_or*tmp.i13 + Omega_pi*tmp.r12 + Omega_pr*tmp.i12 + deltap_o*tmp.i23 - 1.0/2.0*gamma_12*tmp.r23 - 1.0/2.0*gamma_13*tmp.r23 - 1.0/2.0*gamma_23*n_b*tmp.r23 - 1.0/2.0*gamma_23*tmp.r23*(n_b + 1) - 1.0/2.0*gamma_2d*tmp.r23 - 1.0/2.0*gamma_3d*tmp.r23 + tmp.i23*(-deltap_mu - deltap_o);
            diff->i23 = Omega_mur*tmp.r22 - Omega_mur*tmp.r33 + Omega_oi*tmp.i13 - Omega_or*tmp.r13 - Omega_pi*tmp.i12 + Omega_pr*tmp.r12 - deltap_o*tmp.r23 - 1.0/2.0*gamma_12*tmp.i23 - 1.0/2.0*gamma_13*tmp.i23 - 1.0/2.0*gamma_23*n_b*tmp.i23 - 1.0/2.0*gamma_23*tmp.i23*(n_b + 1) - 1.0/2.0*gamma_2d*tmp.i23 - 1.0/2.0*gamma_3d*tmp.i23 - tmp.r23*(-deltap_mu - deltap_o);
        }
    }

    __device__
    void Omega(double *Omega_r, double *Omega_i,
               double g_r, double g_i,
               double alpha_r, double alpha_i)
    {

        double abs_alpha_sqr = alpha_r*alpha_r + alpha_i*alpha_i;
        if (abs_alpha_sqr == 0.0)
        {
            *Omega_r = g_r;
            *Omega_i = g_i;
            return;
        }

        double alpha_rescale = sqrt(1.0 + 1.0/abs_alpha_sqr);
        alpha_r *= alpha_rescale;
        alpha_i *= alpha_rescale;

        *Omega_r = g_r*alpha_r - g_i*alpha_i;
        *Omega_i = g_r*alpha_i + g_i*alpha_r;
    }

    __global__
    void EnsembleDerivativeAndSum(size_t nAtomSamples, SystemState *diff,
        const SystemState *state, const AtomSample *atomSamples, SystemType sys,
        double alpha_r, double alpha_i, double beta_r, double beta_i,
        double Omega_pr, double Omega_pi, double delta_mu, double delta_o,
        double n_b, double gamma_12, double gamma_13,
        double gamma_23, double gamma_2d, double gamma_3d)
    {
        size_t threadId = blockIdx.x*blockDim.x + threadIdx.x;
        size_t threadCount = gridDim.x*blockDim.x;
        for (size_t k = threadId; k < nAtomSamples; k += threadCount)
        {
            // get sample-specific variables
            double g_or = atomSamples[k].g_or;
            double g_oi = atomSamples[k].g_oi;
            double g_mur = atomSamples[k].g_mur;
            double g_mui = atomSamples[k].g_mui;
            double delta_12 = atomSamples[k].delta_12;
            double delta_23 = atomSamples[k].delta_23;
            double w = atomSamples[k].weight;
            const DensityMatrix *rho = &(state->rho[k]);
            DensityMatrix *rhoDiff = &(diff->rho[k]);

            // contribution to ensemble terms
            double rho_or = ((sys == V_SYSTEM) ? rho->r12 : rho->r23);
            double rho_oi = ((sys == V_SYSTEM) ? -rho->i12 : -rho->i23);
            double rho_mur = ((sys == V_SYSTEM) ? rho->r23 : rho->r12);
            double rho_mui = ((sys == V_SYSTEM) ? -rho->i23 : -rho->i12);
            double S_alpha_r = w * (g_or*rho_or + g_oi*rho_oi);
            double S_alpha_i = w * (g_oi*rho_or - g_or*rho_oi);
            double S_beta_r = w * (g_mur*rho_mur + g_mui*rho_mui);
            double S_beta_i = w * (g_mui*rho_mur - g_mur*rho_mui);
            atomicAdd(&(diff->alpha_r), S_alpha_i);
            atomicAdd(&(diff->alpha_i), -S_alpha_r);
            atomicAdd(&(diff->beta_r), S_beta_i);
            atomicAdd(&(diff->beta_i), -S_beta_r);

            // drive Rabi frequencies
            double Omega_mur;
            double Omega_mui;
            double Omega_or;
            double Omega_oi;
            Omega(&Omega_mur, &Omega_mui, g_mur, g_mui, beta_r, beta_i);
            Omega(&Omega_or, &Omega_oi, g_or, g_oi, alpha_r, alpha_i);

            // inhomogeneous shift
            double deltap_mu = delta_mu + (sys == V_SYSTEM ? delta_23 : delta_12);
            double deltap_o = delta_o + (sys == V_SYSTEM ? delta_12 : delta_23);

            MasterDerivative(
                rhoDiff,
                sys,
                rho,
                Omega_mur,
                Omega_mui,
                Omega_or,
                Omega_oi,
                Omega_pr,
                Omega_pi,
                deltap_mu,
                deltap_o,
                n_b,
                gamma_12,
                gamma_13,
                gamma_23,
                gamma_2d,
                gamma_3d);
        }
    }

    void SystemDerivative(size_t nAtomSamples,
        SystemState *diff, const SystemState *state,
        const AtomSample *atomSamples, SystemType sys,
        double Omega_pr, double Omega_pi, double delta_mu, double delta_o,
        double n_b, double gamma_12, double gamma_13, double gamma_23,
        double gamma_2d, double gamma_3d, double gamma_o, double gamma_mu)
    {
        double alpha_r;
        double alpha_i;
        double beta_r;
        double beta_i;
        cudaMemcpy(&alpha_r, &(state->alpha_r),
            sizeof(double), cudaMemcpyDeviceToHost);
        cudaMemcpy(&alpha_i, &(state->alpha_i),
            sizeof(double), cudaMemcpyDeviceToHost);
        cudaMemcpy(&beta_r, &(state->beta_r),
            sizeof(double), cudaMemcpyDeviceToHost);
        cudaMemcpy(&beta_i, &(state->beta_i),
            sizeof(double), cudaMemcpyDeviceToHost);

        double d_alpha_r = -alpha_r * gamma_o/2.0;
        double d_alpha_i = -alpha_i * gamma_o/2.0;
        double d_beta_r = -beta_r * gamma_mu/2.0;
        double d_beta_i = -beta_i * gamma_mu/2.0;
        cudaMemcpy(&(diff->alpha_r), &d_alpha_r,
            sizeof(double), cudaMemcpyHostToDevice);
        cudaMemcpy(&(diff->alpha_i), &d_alpha_i,
            sizeof(double), cudaMemcpyHostToDevice);
        cudaMemcpy(&(diff->beta_r), &d_beta_r,
            sizeof(double), cudaMemcpyHostToDevice);
        cudaMemcpy(&(diff->beta_i), &d_beta_i,
            sizeof(double), cudaMemcpyHostToDevice);

        EnsembleDerivativeAndSum<<>>(
            nAtomSamples,
            diff,
            state,
            atomSamples,
            sys,
            alpha_r,
            alpha_i,
            beta_r,
            beta_i,
            Omega_pr,
            Omega_pi,
            delta_mu,
            delta_o,
            n_b,
            gamma_12,
            gamma_13,
            gamma_23,
            gamma_2d,
            gamma_3d);
        cudaDeviceSynchronize();
    }

    __global__
    void GlobalDensityMatricesWeightedSum(size_t nAtomSamples,
        DensityMatrix *dst, const DensityMatrix *a, double aw,
        const DensityMatrix *b, double bw)
    {
        size_t threadId = blockIdx.x*blockDim.x + threadIdx.x;
        size_t threadCount = gridDim.x*blockDim.x;
        for (size_t k = threadId; k < nAtomSamples; k += threadCount)
        {
            dst[k].r11 = aw*a[k].r11 + bw*b[k].r11;
            dst[k].r22 = aw*a[k].r22 + bw*b[k].r22;
            dst[k].r33 = aw*a[k].r33 + bw*b[k].r33;
            dst[k].r12 = aw*a[k].r12 + bw*b[k].r12;
            dst[k].i12 = aw*a[k].i12 + bw*b[k].i12;
            dst[k].r13 = aw*a[k].r13 + bw*b[k].r13;
            dst[k].i13 = aw*a[k].i13 + bw*b[k].i13;
            dst[k].r23 = aw*a[k].r23 + bw*b[k].r23;
            dst[k].i23 = aw*a[k].i23 + bw*b[k].i23;
        }
    }

    __global__
    void GlobalCavityStateWeightedSum(size_t nAtomSamples,
        SystemState *dst, const SystemState *a, double aw,
        const SystemState *b, double bw)
    {
        dst->alpha_r = aw*a->alpha_r + bw*b->alpha_r;
        dst->alpha_i = aw*a->alpha_i + bw*b->alpha_i;
        dst->beta_r = aw*a->beta_r + bw*b->beta_r;
        dst->beta_i = aw*a->beta_i + bw*b->beta_i;
    }

    void SystemStateWeightedSum(size_t nAtomSamples, SystemState *dst,
        const SystemState *a, double aw, const SystemState *b, double bw)
    {
        GlobalCavityStateWeightedSum<<<1,1>>>(nAtomSamples, dst, a, aw, b, bw);
        cudaDeviceSynchronize();
        GlobalDensityMatricesWeightedSum<<>>(nAtomSamples,
            dst->rho, a->rho, aw, b->rho, bw);
        cudaDeviceSynchronize();
    }

    void SystemRK4Step(size_t nAtomSamples, SystemState *__restrict__ state,
        const AtomSample *__restrict__ atomSamples,
        SystemState *__restrict__ tmp1, SystemState *__restrict__ tmp2,
        SystemType sys, double dt, double Omega_pr, double Omega_pi,
        double delta_mu, double delta_o, double n_b, double gamma_12,
        double gamma_13, double gamma_23, double gamma_2d,
        double gamma_3d, double gamma_o, double gamma_mu)
    {
        SystemDerivative( // k1
            nAtomSamples,
            tmp1,
            state,
            atomSamples,
            sys,
            Omega_pr,
            Omega_pi,
            delta_mu,
            delta_o,
            n_b,
            gamma_12,
            gamma_13,
            gamma_23,
            gamma_2d,
            gamma_3d,
            gamma_o,
            gamma_mu);
        SystemStateWeightedSum(nAtomSamples, tmp2, state, 1.0, tmp1, dt/6.0);
        SystemStateWeightedSum(nAtomSamples, tmp1, state, 1.0, tmp1, dt/2.0);
        SystemDerivative( // k2
            nAtomSamples,
            tmp1,
            tmp1,
            atomSamples,
            sys,
            Omega_pr,
            Omega_pi,
            delta_mu,
            delta_o,
            n_b,
            gamma_12,
            gamma_13,
            gamma_23,
            gamma_2d,
            gamma_3d,
            gamma_o,
            gamma_mu);
        SystemStateWeightedSum(nAtomSamples, tmp2, tmp1, dt/3.0, tmp2, 1.0);
        SystemStateWeightedSum(nAtomSamples, tmp1, state, 1.0, tmp1, dt/2.0);
        SystemDerivative( // k3
            nAtomSamples,
            tmp1,
            tmp1,
            atomSamples,
            sys,
            Omega_pr,
            Omega_pi,
            delta_mu,
            delta_o,
            n_b,
            gamma_12,
            gamma_13,
            gamma_23,
            gamma_2d,
            gamma_3d,
            gamma_o,
            gamma_mu);
        SystemStateWeightedSum(nAtomSamples, tmp2, tmp1, dt/3.0, tmp2, 1.0);
        SystemStateWeightedSum(nAtomSamples, tmp1, state, 1.0, tmp1, dt);
        SystemDerivative( // k4
            nAtomSamples,
            tmp1,
            tmp1,
            atomSamples,
            sys,
            Omega_pr,
            Omega_pi,
            delta_mu,
            delta_o,
            n_b,
            gamma_12,
            gamma_13,
            gamma_23,
            gamma_2d,
            gamma_3d,
            gamma_o,
            gamma_mu);
        SystemStateWeightedSum(nAtomSamples, state, tmp1, dt/6.0, tmp2, 1.0);
    }

    double PlanckExcitation(double T, double omega)
    {
        return 1.0 / expm1(HBAR*omega / (K_B*T));
    }

    double RandomUniform(void)
    {
        return ((double) rand()) / ((double) RAND_MAX);
    }

    double RandomGaussian(void)
    {
        double U1 = RandomUniform();
        double U2 = RandomUniform();

        // U1 == 0.0 results in an error when taking its log
        while (U1 == 0.0)
        {
            U1 = RandomUniform();
        }

        double R = sqrt(-2.0 * log(U1));
        double Theta = 2.0*PI * U2;
        return R * cos(Theta);
    }

    int main(void)
    {
        // system constants
        const SystemType sys = LAMBDA_SYSTEM;
        const double omega_12 = 2.0*PI*5.186e9;
        const double d13 = 1.63e-32;
        const double d23 = 1.15e-32;
        const double tau_12 = 11.0;
        const double tau_3 = 0.011;
        const double gamma_2d = 1e6;
        const double gamma_3d = 1e6;
        const double sigma_o = 2.0*PI*419e6;
        const double sigma_mu = 2.0*PI*5e6;
        const double N = 1e16;
        const double gamma_oi = 2.0*PI*7.95e6;
        const double gamma_oc = 2.0*PI*1.7e6;
        const double gamma_mui = 2.0*PI*650e3;
        const double gamma_muc = 2.0*PI*1.5e6;
        const double g_or = 51.9;
        const double g_oi = 0.0;
        const double g_mur = 1.04;
        const double g_mui = 0.0;

        const double T = 4.6;
        const double n_b = PlanckExcitation(T, omega_12);
        const double tau_13 = tau_3 * d13*d13 / (d13*d13 + d23*d23);
        const double tau_23 = tau_3 * d23*d23 / (d13*d13 + d23*d23);
        const double gamma_12 = 1.0 / (tau_12*(n_b+1.0));
        const double gamma_13 = 1.0 / tau_13;
        const double gamma_23 = 1.0 / tau_23;

        const double sigma_12 = ((sys == V_SYSTEM) ? sigma_o : sigma_mu);
        const double sigma_23 = ((sys == V_SYSTEM) ? sigma_mu : sigma_o);

        // parameters
        size_t nAtomSamples = 1'000'000;
        const double weight = N / ((double) nAtomSamples);
        const double delta_mu = 2.0*sigma_mu;
        const double delta_o = 2.0*sigma_o;
        const double Omega_pr = 35000.0;
        const double Omega_pi = 0.0;
        const double alpha_0r = 1.0;
        const double alpha_0i = 0.0;
        const double beta_0r = 1.0;
        const double beta_0i = 0.0;

        // allocate arrays
        AtomSample *atomSamples;
        if (!(HostAllocateAtomSamples(&atomSamples, nAtomSamples)))
        {
            printf("Host memory allocation failure\n");
            return -1;
        }

        AtomSample *deviceAtomSamples;
        if (!(DeviceAllocateAtomSamples(&deviceAtomSamples, nAtomSamples)))
        {
            free(atomSamples);
            printf("Device memory allocation failure\n");
            return -1;
        }

        SystemState *state;
        if (!(HostAllocateStateStruct(&state, nAtomSamples)))
        {
            free(atomSamples);
            cudaFree(deviceAtomSamples);
            printf("Host memory allocation failure\n");
            return -1;
        }

        SystemState *deviceState;
        if (!(DeviceAllocateStateStruct(&deviceState, nAtomSamples)))
        {
            free(atomSamples);
            cudaFree(deviceAtomSamples);
            free(state);
            printf("Device memory allocation failure\n");
            return -1;
        }

        SystemState *tmp1;
        if (!(DeviceAllocateStateStruct(&tmp1, nAtomSamples)))
        {
            free(atomSamples);
            cudaFree(deviceAtomSamples);
            free(state);
            cudaFree(deviceState);
            printf("Device memory allocation failure\n");
            return -1;
        }

        SystemState *tmp2;
        if (!(DeviceAllocateStateStruct(&tmp2, nAtomSamples)))
        {
            free(atomSamples);
            cudaFree(deviceAtomSamples);
            free(state);
            cudaFree(deviceState);
            cudaFree(tmp1);
            printf("Device memory allocation failure\n");
            return -1;
        }

        // populate arrays and copy to device memory
        state->alpha_r = alpha_0r;
        state->alpha_i = alpha_0i;
        state->beta_r = beta_0r;
        state->beta_i = beta_0i;

        for (size_t k = 0; k < nAtomSamples; k++)
        {
            atomSamples[k].g_or = g_or;
            atomSamples[k].g_oi = g_oi;
            atomSamples[k].g_mur = g_mur;
            atomSamples[k].g_mui = g_mui;
            atomSamples[k].weight = weight;
            atomSamples[k].delta_12 = sigma_12 * RandomGaussian();
            atomSamples[k].delta_23 = sigma_23 * RandomGaussian();

            state->rho[k] = GROUND_STATE_MATRIX;
        }

        CopyAtomSamples(deviceAtomSamples, atomSamples, nAtomSamples,
                        cudaMemcpyHostToDevice);
        CopyStateStruct(deviceState, state, nAtomSamples, cudaMemcpyHostToDevice);

        // run simulation and save results to binary files
        const char *dirname = ".";
        char atomSamplesFpath[256];
        sprintf(atomSamplesFpath, "%s/atom_samples", dirname);
        FILE *atomSamplesFp = fopen(atomSamplesFpath, "wb");
        fwrite(atomSamples, sizeof(AtomSample), nAtomSamples, atomSamplesFp);
        printf("Saved atom sample data as %s\n", atomSamplesFpath);
        fclose(atomSamplesFp);

        const double dt = 10e-12;
        size_t numSteps = 0;

        size_t numPrint = 0;
        while (true)
        {
            size_t toPrint;
            if (numPrint < 100)
            {
                toPrint = numPrint;
            }
            else
            {
                size_t n = (numPrint-100) / 900;
                toPrint = (numPrint-100) % 900 + 100;
                for (size_t i = 0; i < n; i++)
                {
                    toPrint *= 10;
                }
            }

            while (numSteps < toPrint)
            {
                SystemRK4Step(
                    nAtomSamples,
                    deviceState,
                    deviceAtomSamples,
                    tmp1,
                    tmp2,
                    sys,
                    dt,
                    Omega_pr,
                    Omega_pi,
                    delta_mu,
                    delta_o,
                    n_b,
                    gamma_12,
                    gamma_13,
                    gamma_23,
                    gamma_2d,
                    gamma_3d,
                    gamma_oi+gamma_oc,
                    gamma_mui+gamma_muc);
                numSteps++;
            }

            char fname[256];
            sprintf(fname, "state_step_%zd_dt_%zd_ps",
                numSteps, (size_t) round(dt*1e12));
            char fpath[256];
            sprintf(fpath, "%s/%s", dirname, fname);

            FILE *fp = fopen(fpath, "wb");
            CopyStateStruct(state, deviceState,
                nAtomSamples, cudaMemcpyDeviceToHost);
            fwrite(state, SizeofStateStruct(nAtomSamples), 1, fp);
            fclose(fp);

            printf("\nSaved step %zd as %s\n", numSteps, fpath);
            printf("alpha_r %f\n", state->alpha_r);
            printf("alpha_i %f\n", state->alpha_i);
            printf("beta_r %f\n", state->beta_r);
            printf("beta_i %f\n", state->beta_i);
            printf("rho[0].r11 %f\n", state->rho[0].r11);
            printf("rho[0].r22 %f\n", state->rho[0].r22);
            printf("rho[0].r33 %f\n", state->rho[0].r33);
            printf("rho[0].r12 %f\n", state->rho[0].r12);
            printf("rho[0].i12 %f\n", state->rho[0].i12);
            printf("rho[0].r13 %f\n", state->rho[0].r13);
            printf("rho[0].i13 %f\n", state->rho[0].i13);
            printf("rho[0].r23 %f\n", state->rho[0].r23);
            printf("rho[0].i23 %f\n", state->rho[0].i23);

            numPrint++;
        }

        free(atomSamples);
        cudaFree(deviceAtomSamples);
        free(state);
        cudaFree(deviceState);
        cudaFree(tmp1);
        cudaFree(tmp2);
        return 0;
    }
    |