small python testing scripts

test-scripts/dipole_rad.py

+import numpy as np
+import lib.pyplotdefs as pd
+import sys
+
+trials = 100000
+p = 1.0
+s = 1 - p
+
+pdf = lambda p, s, th: (
+    p * 3 / 8 / np.pi * np.sin(th)**2
+    + s * 3 / 16 / np.pi * (1 + np.cos(th)**2)
+)
+
+cdf = lambda p, s, th: (
+    0.5
+    - (2 * p + s) * 3 / 8 * np.cos(th)
+    + (2 * p - s) / 8 * np.cos(th)**3
+)
+
+xplot = np.linspace(0.0, np.pi, 1000)
+pdfplot = pdf(p, s, xplot)
+cdfplot = cdf(p, s, xplot)
+
+# (pd.Plotter()
+#     .plot(xplot, pdfplot)
+#     .plot(xplot, cdfplot)
+#     .ggrid()
+#     .show()
+#     .close()
+# )
+# sys.exit(0)
+
+def newton_raphson(p: float, s: float, r: float) -> float:
+    th = r * np.pi
+    th = np.pi / 2
+    dth = np.inf
+    for _ in range(1000):
+        dth = (cdf(p, s, th) - r) / (pdf(p, s, th) * 2 * np.pi * np.sin(th))
+        th -= dth
+        if abs(dth) < 1e-6:
+            return th
+        # th = min(max(th, 1e-6), np.pi * (1 - 1e-6))
+    (pd.Plotter()
+        .plot(xplot, pdfplot)
+        .plot(xplot, cdfplot)
+        .plot(th, r, marker="o", linestyle="", color="r")
+        .show()
+        .close()
+    )
+    raise Exception(f"didn't converge: {p=:.3f}; {s=:.3f}; {r=:.3f}")
+
+data = np.zeros((trials, 2), dtype=np.float64)
+for k in range(trials):
+    print(f"\r  {k = :5.0f} ", end="", flush=True)
+    r = np.random.random()
+    th = newton_raphson(p, s, r)
+    data[k, :] = (th, r)
+print()
+
+(pd.Plotter()
+    .plot(
+        xplot, cdfplot,
+        marker="", linestyle="-", color="k"
+    )
+    .plot(
+        data[:, 0], data[:, 1],
+        marker="o", linestyle="", color="C0", alpha=0.35
+    )
+    .ggrid()
+    .show()
+    .close()
+)
+

test-scripts/lib

+../lib
\ No newline at end of file

test-scripts/mean_excite_time.py

+import numpy as np
+import lib.pyplotdefs as pd
+
+saturation = 10.0
+detuning = 1.0
+linewidth = 5.0
+
+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)
+
+W = np.sqrt(linewidth / 2.0 * saturation)
+w = np.sqrt(W**2 + (linewidth / 4.0)**2 + detuning**2)
+t0, rho0 = rho_ee_max(saturation, detuning, linewidth)
+K1 = np.exp(0.75 * linewidth * t0) - np.cos(w * t0)
+K2 = 9.0 * linewidth**2 + 16.0 * w**2
+tbar_anl = (
+    12.0 * linewidth * K1
+    - K2 * t0 * np.cos(w * t0)
+     + 16.0 * w * np.sin(w * t0)
+) / (K1 * K2)
+
+t = np.linspace(0.0, t0, 10000)
+tbar_num = np.trapz(
+    (
+        t * drho_ee(t, saturation, detuning, linewidth) / rho0
+    ),
+    dx=t[1] - t[0]
+)
+
+print(f"{{Y: {linewidth}, w: {w}, t1: {t0}}}")
+print(tbar_anl)
+print(tbar_num)
+
+(pd.Plotter()
+    .plot(t, rho_ee(t, saturation, detuning, linewidth) / rho0, color="k")
+    .plot(t, drho_ee(t, saturation, detuning, linewidth) / rho0, color="0.5")
+    .axvline(tbar_anl, color="b")
+    .axvline(tbar_num, color="r")
+    .ggrid()
+    .show()
+    .close()
+)
+

test-scripts/rho_sample.py

+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()
+)
+

test-scripts/scatter_inv.py

+import numpy as np
+import lib.pyplotdefs as pd
+
+# def cdf_inv(r: float, rho: float, sigma: float) -> float:
+#     Mp = 2 * rho + sigma
+#     Mm = 2 * rho - sigma
+#     p = -3 * Mp / Mm
+#     q = -8 * (r - 0.5) / Mm
+#     if 4 * p**3 + 27 * q**2 > 0 and p < 0:
+#         branch = True
+#         t = (
+#             -2 * (q / abs(q)) * np.sqrt(-p / 3)
+#             * np.cosh(
+#                 np.acosh((-3 * abs(q)) / (2 * p) * np.sqrt(-3 / p)) / 3
+#             )
+#         )
+#     else:
+#         branch = False
+#         p = complex(p, 0.0)
+#         q = complex(q, 0.0)
+#         x0 = 2 * np.sqrt(-p / 3)
+#         x1 = np.arccos((3 * q) / (2 * p) * np.sqrt(-3 / p)) / 3
+#         for k in range(3):
+#             tk = x0 * np.cos(x1 - 2 * np.pi * k / 3)
+#             if tk.imag < 1e-6:
+#                 break
+#         t = tk.real
+#         if tk.imag < 1e-6:
+#             if tk.real >= -1 and tk.real <= 1:
+#                 t = tk.real
+#             elif abs(tk.real + 1.0) < 1e-6:
+#                 t = -1.0
+#             elif abs(tk.real - 1.0) < 1e-6:
+#                 t = 1.0
+#             else:
+#                 raise Exception(
+#                     f"cubic root is out of bounds\n"
+#                     f"{r=} ; {rho=} ; {sigma=}\n"
+#                     f"{tk=}"
+#                 )
+#         else:
+#             raise Exception(
+#                 f"cubic root is out of bounds\n"
+#                 f"{r=} ; {rho=} ; {sigma=}\n"
+#                 f"{tk=}"
+#             )
+#     return np.arccos(t), branch
+
+def cdf_inv(r: float, rho: float, sigma: float) -> float:
+    Mp = 2 * rho + sigma
+    Mm = 2 * rho - sigma
+
+    a = Mm / 8
+    b = 0
+    c = -3 * Mp / 8
+    d = 0.5 - r
+
+    D0 = complex(b**2 - 3 * a * c, 0.0)
+    D1 = complex(2 * b**3 - 9 * a * b * c + 27 * a**2 * d, 0.0)
+    Cp = pow((D1 + np.sqrt(D1**2 - 4 * D0**3)) / 2, 1 / 3)
+    Cm = pow((D1 - np.sqrt(D1**2 - 4 * D0**3)) / 2, 1 / 3)
+    if abs(Cp) > 1e-6:
+        C = Cp
+    elif abs(Cm) > 1e-6:
+        C = Cm
+    else:
+        return -b / 3 / a
+    xi = (-1 + np.sqrt(3) * 1j) / 2
+    for k in [0, 1, 2]:
+        x = -(b + xi**k * C + D0 / (xi**k * C)) / (3 * a)
+        if abs(x.imag) < 1e-6:
+            break
+    if abs(x.imag) < 1e-6:
+        if abs(x.real) <= 1.0:
+            return np.arccos(x.real), k
+        elif abs(x.real + 1.0) < 1e-6:
+            return np.arccos(-1.0), k
+        elif abs(x.real - 1.0) < 1e-6:
+            return np.arccos(+1.0), k
+        else:
+            raise Exception(
+                "cubic root is out of bounds\n"
+                f"{r=} ; {rho=} ; {sigma=}\n"
+                f"{x=}"
+            )
+    else:
+        raise Exception(
+            "cubic root is out of bounds\n"
+            f"{r=} ; {rho=} ; {sigma=}\n"
+            f"{x=}"
+        )
+
+def testpq(r: float, rho: float, sigma: float) -> float:
+    Mp = 2 * rho + sigma
+    Mm = 2 * rho - sigma
+    p = -3 * Mp / Mm
+    q = -8 * (r - 0.5) / Mm
+    return 4 * p**3 + 27 * q**2, p
+
+r_test = np.linspace(0.0, 1.0, 51)
+rho_test = np.linspace(0.0, 1.0, 51)
+
+# B = np.zeros((51, 51), dtype=np.float64)
+# p = np.zeros((51, 51), dtype=np.float64)
+# for i, r_t in enumerate(r_test):
+#     for j, rho_t in enumerate(rho_test):
+#         B[i, j], p[i, j] = testpq(r_t, rho_t, 1 - rho_t)
+# pd.Plotter().imshow(B > 0).colorbar()
+# pd.Plotter().imshow(p < 0).colorbar()
+
+theta = np.zeros((51, 51), dtype=np.float64)
+branch = np.zeros((51, 51), dtype=np.int8)
+for i, r_t in enumerate(r_test):
+    for j, rho_t in enumerate(rho_test):
+        theta[i, j], branch[i, j] = cdf_inv(r_t, rho_t, 1 - rho_t)
+pd.Plotter().imshow(theta).colorbar()
+pd.Plotter().imshow(branch, cmap="gray").colorbar()
+
+pd.show()
+