FigS3.py

import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt

# --- System Parameters ---
# Converted to angular frequencies (rad/s)
twopi = 2 * np.pi
gamma = 100e3 * twopi     # Atomic decay rate (1 kHz)
chi = 10e6 * twopi     # Inhomogeneous broadening (100 kHz)
kappa = 100e6 * twopi     # Cavity linewidth (5 MHz)
g = 1.5e3 * twopi       # Coupling constant (1.4 kHz)
N = 1e11                # Number of atoms


# Incoherent pumping rates
eta_values = [251e3 * twopi, 500e3 * twopi, 20e3 * twopi, 63e3 * twopi, 39e3 * twopi]
eta_labels = ['251 kHz', '158 kHz', '100 kHz', '63 kHz', '39 kHz']
colors = ['blue', 'red', 'black', 'green', 'magenta']

def mft_equations(t, y, eta):
    """
    Mean Field Theory (MFT) ODEs derived from Eqs. 5-8.
    Assuming detuning delta = 0, the atom-field coherence is purely imaginary.
    
    Variables:
    y[0] = n   : Intracavity photon number <a^dag a>
    y[1] = c_I : Imaginary part of atom-field coherence Im(<a^dag sigma_1^->)
    y[2] = z   : Atomic population inversion <sigma^z>
    y[3] = s   : Spin-spin correlation <sigma_1^+ sigma_2^->
    """
    n, c_I, z, s = y
    
    # Eq. 5: Photon number dynamics
    dn = -kappa * n + g * N * c_I
    
    # Eq. 6: Atom-field coherence dynamics (Imaginary part)
    Gamma = (eta + gamma + kappa) / 2.0 + chi
    dc_I = -Gamma * c_I + (g / 2.0) * (z * n + (z + 1.0) / 2.0 + (N - 1.0) * s)
    
    # Eq. 7: Population inversion dynamics
    dz = -2.0 * g * c_I - gamma * (1.0 + z) + eta * (1.0 - z)
    
    # Eq. 8: Spin-spin correlation dynamics
    ds = -(gamma + eta + 2.0 * chi) * s + g * z * c_I
    
    return [dn, dc_I, dz, ds]

# --- Simulation Setup ---
t_span = (0, 1e-5)
# Use a dense time grid for smooth oscillatory plots
t_eval = np.linspace(t_span[0], t_span[1], 1000)

# Initial conditions (Based on Fig S3 starting at <sigma_z> = 1)
# [n(0), c_I(0), z(0), s(0)]
y0 = [0.0, 0.0, 1.0, 0.0]

# --- Plotting ---
plt.figure(figsize=(10, 6))

for eta, label, color in zip(eta_values, eta_labels, colors):
    # 'Radau' is an implicit Runge-Kutta method ideal for stiff ODEs
    sol = solve_ivp(mft_equations, t_span, y0, args=(eta,), 
                    t_eval=t_eval, method='Radau')
    
    # sol.y[2] corresponds to the z variable <sigma_z>
    plt.plot(sol.t, sol.y[2], label=rf'$\eta = {label}$', color=color, linewidth=1.2)

# --- Plot Formatting (Matching Fig S3) ---
plt.xlabel('Time(s)', fontsize=14)
plt.ylabel(r'$< \hat{\sigma}_z >$', fontsize=14)
plt.ylim(-1, 1)
plt.xlim(0, 5e-6)

# Formatting scientific notation on the x-axis
plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)

# Replicating the legend style
legend = plt.legend(loc='upper right', fontsize=12, framealpha=1)
legend.get_frame().set_edgecolor('black')

plt.tight_layout()
plt.show()