Skip to content
Snippets Groups Projects
rho_sample.py 2.61 KiB
Newer Older
whooie's avatar
whooie committed
import numpy as np
import lib.pyplotdefs as pd

def pop_excited(saturation: float, detuning: float, linewidth: float) -> float:
    return (
        saturation / 2.0
        / (1.0 + saturation + (2.0 * detuning / linewidth)**2)
    )

def rho_ee(t: float, saturation: float, detuning: float, linewidth: float) -> float:
    W = np.sqrt(linewidth / 2.0 * saturation)
    w = np.sqrt(W**2 - (linewidth / 4.0)**2 + detuning**2)
    return (
        pop_excited(saturation, detuning, linewidth)
        * (1.0 - np.exp(-0.75 * linewidth * t) * np.cos(w * t))
    )

def drho_ee(t: float, saturation: float, detuning: float, linewidth: float) -> float:
    W = np.sqrt(linewidth / 2.0 * saturation)
    w = np.sqrt(W**2 - (linewidth / 4.0)**2 + detuning**2)
    return (
        pop_excited(saturation, detuning, linewidth)
        * (
            0.75 * linewidth * np.exp(-0.75 * linewidth * t) * np.cos(w * t)
            + w * np.exp(-0.75 * linewidth * t) * np.sin(w * t)
        )
    )

def rho_ee_max(saturation: float, detuning: float, linewidth: float) -> (float, float):
    W = np.sqrt(linewidth / 2.0 * saturation)
    w = np.sqrt(W**2 - (linewidth / 4.0)**2 + detuning**2)
    t0 = (
        2.0 / w
        * np.arctan(
            4.0 * w / 3.0 / linewidth
            + np.sqrt(1.0 + (4.0 * w / 3.0 / linewidth)**2)
        )
    )
    rho0 = (
        pop_excited(saturation, detuning, linewidth)
        * (1.0 - np.exp(-0.75 * linewidth * t0) * np.cos(w * t0))
    )
    return (t0, rho0)

def rho_ee_inv(saturation: float, detuning: float, linewidth: float, r: float) -> float:
    if r < 0.0 or r > 1.0:
        raise Exception("rho_ee_inv: encountered invalid probability value")
    t0, rho0 = rho_ee_max(saturation, detuning, linewidth)
    if r > rho0:
        # return t0
        return np.inf
    t = t0 / 2.0
    for _ in range(1000):
        dt  = (
            (rho_ee(t, saturation, detuning, linewidth) - r)
            / drho_ee(t, saturation, detuning, linewidth)
        )
        t -= dt
        if abs(dt) < 1e-6:
            return t
        t = max(0.0, min(t, t0))
    raise Exception("rho_ee_inv: failed to converge")

s = 5.0
d = 5.0
y = 2.0

N = 10000

r = np.random.random(size=N)
t = np.array([rho_ee_inv(s, d, y, rk) for rk in r])
t_finite = t[np.where(np.isfinite(t))]

t_mean = t_finite.mean()

t0, rho0 = rho_ee_max(s, d, y)

print(t_mean, t0, t0 / t_mean)

xplot = np.linspace(0.0, t0 * 10, 1000)
yplot = np.array([rho_ee(x, s, d, y) for x in xplot])

(pd.Plotter()
    .plot(xplot, yplot, marker="", linestyle="-", color="k")
    .plot(t, r, marker="o", linestyle="", color="C0", alpha=0.35)
    .ggrid()
    .show()
)