Fig2b_3.py

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import root
from tqdm import tqdm
import warnings
import os

# Suppress runtime warnings
warnings.filterwarnings("ignore")

# --- 1. MFT Equations (Returning Derivatives) ---
def mft_derivatives(y, N, eta, g, kappa, gamma, chi):
    # y = [n, m_re, m_im, Jz, S_re]
    n, m_re, m_im, Jz, S_re = y
    m = m_re + 1j * m_im
    
    m_minus_m_conj = 2j * m_im 
    
    # Eq 5: Photon Number
    dn = -kappa * n + g * N * m_im
    
    # Eq 6: Atom-Field Coherence
    term1 = -( (eta + gamma + kappa)/2.0 + chi ) * m
    term2 = 1j * g / 2.0 * ( Jz * n + (Jz + 1)/2.0 + (N - 1) * S_re )
    dm = term1 + term2
    
    # Eq 7: Inversion
    dJz = 1j * g * m_minus_m_conj - gamma * (1 + Jz) + eta * (1 - Jz)
    
    # Eq 8: Spin-Spin Correlation
    dS = -(gamma + eta + 2*chi) * S_re - (1j * g / 2.0 * Jz * m_minus_m_conj).real
    
    return [dn, dm.real, dm.imag, dJz.real, dS.real]

# --- 2. Parameters ---
g = 2 * np.pi * 1e3 
# Convert g from rad/s to Hz and format as an integer
g_hz = round(g / (2 * np.pi))     
kappa = 2 * np.pi * 100e6  
gamma = 2 * np.pi * 100e3 
chi = 2 * np.pi * 10e6    
N_crit = (gamma * kappa) / (g**2)
print(f"Critical Atom Number N_c = {N_crit:.2e}")


# --- 3. Grid Setup ---
res = 5
eta_range = np.logspace(4.55, 11, res) 
N_range = np.logspace(8.5, 12, res)  
X, Y = np.meshgrid(eta_range, N_range)

Z_linewidth = np.zeros_like(X)

print(f"Solving Steady-State Roots for {res}x{res} grid...")

# --- 4. Main Solver Loop ---
for i in tqdm(range(res)):
    for j in range(res):
        N_val = Y[i, j]
        eta_hz = X[i, j]
        eta_rad = eta_hz * (2 * np.pi)
        
        # Initial Guess
        guess = [1e4, 1e2, 1e2, -0.5, 0] 
        
        try:
            sol = root(mft_derivatives, guess, 
                       args=(N_val, eta_rad, g, kappa, gamma, chi),
                       method='lm')
            
            if sol.success:
                n_ss = sol.x[0]
            else:
                n_ss = 0 
            
            # --- Linewidth Calculation ---
            if n_ss > 1.0: 
                lw = (gamma + eta_rad) / (4 * np.pi * n_ss)
            else:
                lw = kappa / (2 * np.pi)
                
            Z_linewidth[i, j] = lw
            
        except Exception:
            Z_linewidth[i, j] = kappa / (2 * np.pi)

# ---------------------------------------------------------
# --- 5. SAVE THE COMPUTED DATA ---
# ---------------------------------------------------------
save_dir = "Final version"
os.makedirs(save_dir, exist_ok=True) # Ensure the folder exists

data_filepath = os.path.join(save_dir, f"linewidth_data_g{g_hz}Hz.npz")
np.savez_compressed(data_filepath, X=X, Y=Y, Z_linewidth=Z_linewidth)
print(f"Data successfully saved to {data_filepath}")
# ---------------------------------------------------------

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

# Linewidth Heatmap (Magma Reverse)
im = plt.pcolormesh(X, Y, np.log10(Z_linewidth), shading='auto', cmap='magma_r', vmin=-1, vmax=8)

# Vertical Markers
plt.axvline(x=gamma/(2*np.pi), color='lime', linestyle=':', linewidth=2, label=r'$\gamma$')
plt.axvline(x=kappa/(2*np.pi), color='cyan', linestyle='-', linewidth=2, label=r'$\kappa$')
plt.axvline(x=chi/(2*np.pi), color='yellow', linestyle='-.', linewidth=2, label=r'$\chi$')

# --- NEW WHITE DASHED LINE: N = (N_c / \gamma) * \eta --------
eta_plot_hz = np.logspace(4.55, 11, 100)  # X-axis values in Hz
eta_plot_rad = eta_plot_hz * 2 * np.pi    # Converted to rad/s for math equation
N_line = (N_crit / gamma) * eta_plot_rad  # Y-axis values for N

plt.plot(eta_plot_hz, N_line, color='white', linestyle='--', linewidth=3, label=r'$N = (N_c/\gamma)\eta$')
# -------------------------------------------------------------

plt.xscale('log')
plt.yscale('log')
plt.xlim(5e4, 1e11)
plt.ylim(5e8, 1e12)

# --- LARGER FONT SIZES ---
plt.xlabel(r'Pumping rate $\eta/2\pi$ (Hz)', fontsize=16)
plt.ylabel(r'Number of atoms $N$', fontsize=16)
plt.title(r'Laser Linewidth (MFT)', fontsize=18)

# Increase tick label size
plt.tick_params(axis='both', which='major', labelsize=14)

cbar = plt.colorbar(im)
cbar.set_label(r'$\log_{10} \Delta \nu$ (Hz)', rotation=270, labelpad=25, fontsize=16)
cbar.ax.tick_params(labelsize=14)  # Colorbar tick size

plt.legend(loc='upper right', fontsize=14) 
plt.tight_layout()

# Save the figure dynamically based on g
plt.savefig(os.path.join(save_dir, f"Fig2b_g={g_hz}Hz.png"), dpi=300)
plt.show()