examples.py

"""
Example calculations using the energyloss package.
Run with:  python examples.py
"""

import math
import numpy as np
import matplotlib.pyplot as plt
import energyloss as el

SEP = "-" * 60


# ---------------------------------------------------------------------------
# 1. Bethe-Bloch dE/dx for common particles in liquid argon
# ---------------------------------------------------------------------------

print(SEP)
print("1. Bethe-Bloch dE/dx in liquid argon  [MeV/(g/cm²)]")
print(SEP)
print(f"{'Particle':>10}  {'P (GeV)':>8}  {'Bethe':>8}  {'Restricted':>10}  {'with δ':>8}")

lar = el.get_material("argon", "liquid")
particles = [
    ("muon",   el.MUON,   [0.1, 0.3, 1.0, 3.0, 10.0]),
    ("pion",   el.PION,   [0.3, 1.0, 3.0]),
    ("proton", el.PROTON, [0.5, 1.0, 3.0]),
]

for name, p, momenta in particles:
    for P in momenta:
        bethe  = el.dedx(P, p, lar)
        rest   = el.dedx(P, p, lar, tcut=5.0)
        with_d = el.dedx(P, p, lar, density_correction=True)
        print(f"{name:>10}  {P:>8.2f}  {bethe:>8.4f}  {rest:>10.4f}  {with_d:>8.4f}")


# ---------------------------------------------------------------------------
# 2. z² scaling: proton vs alpha vs carbon ion in silicon
# ---------------------------------------------------------------------------

print()
print(SEP)
print("2. Charge scaling (z²) at same β — silicon")
print(SEP)

si = el.get_material("silicon", "solid")
carbon = el.Particle(mass=11.177, charge=6)  # ¹²C

# Fix β ≈ 0.7 (roughly minimum ionising for proton)
import math
beta_target = 0.70
particles_z = [
    ("proton", el.PROTON),
    ("alpha",  el.ALPHA),
    ("carbon", carbon),
]

print(f"β = {beta_target}  (same velocity)")
print(f"{'Particle':>10}  {'z':>3}  {'P (GeV)':>9}  {'dE/dx':>9}  {'ratio/z²':>10}")
ref_dedx = None
for name, p in particles_z:
    bg = beta_target / math.sqrt(1 - beta_target**2)
    P = p.mass * bg
    d = el.dedx(P, p, si)
    if ref_dedx is None:
        ref_dedx = d
    ratio = d / (p.charge**2 * ref_dedx)
    print(f"{name:>10}  {p.charge:>3}  {P:>9.4f}  {d:>9.4f}  {ratio:>10.4f}")


# ---------------------------------------------------------------------------
# 3. CSDA range vs momentum for proton in several materials
# ---------------------------------------------------------------------------

print()
print(SEP)
print("3. Proton CSDA range (cm) vs momentum")
print(SEP)

materials = {
    "LAr":    el.get_material("argon",   "liquid"),
    "LXe":    el.get_material("xenon",   "liquid"),
    "Si":     el.get_material("silicon", "solid"),
    "H₂O":   el.get_material("water",   "liquid"),
    "Fe":     el.get_material("iron",    "solid"),
}

momenta = [0.3, 0.5, 1.0, 2.0, 5.0]
header = f"{'P (GeV)':>8}" + "".join(f"  {name:>8}" for name in materials)
print(header)

for P in momenta:
    row = f"{P:>8.2f}"
    for mat in materials.values():
        r = el.range_dedx(P, el.PROTON, mat)
        row += f"  {r['range']:>8.2f}"
    print(row)


# ---------------------------------------------------------------------------
# 4. Landau most-probable loss vs Bethe mean — 300 µm silicon
# ---------------------------------------------------------------------------

print()
print(SEP)
print("4. Landau most-probable loss vs Bethe mean — 300 µm silicon")
print("   (dx = 0.0699 g/cm²,  particle = pion)")
print(SEP)

si = el.get_material("silicon", "solid")
dx = 0.0699  # 300 µm × 2.33 g/cc

print(f"{'P (GeV)':>8}  {'Bethe×dx':>10}  {'MP loss':>9}  {'σ (keV)':>9}  {'MP/Bethe':>10}")
for P in [0.3, 0.5, 1.0, 3.0, 10.0]:
    mean  = el.dedx(P, el.PION, si) * dx
    mp    = el.most_probable_loss(P, el.PION, si, dx)
    sig   = el.sigma(P, el.PION, si, dx)
    print(f"{P:>8.2f}  {mean:>10.4f}  {mp:>9.4f}  {sig*1000:>9.2f}  {mp/mean:>10.4f}")


# ---------------------------------------------------------------------------
# 5. Radiation length
# ---------------------------------------------------------------------------

print()
print(SEP)
print("5. Radiation length X₀ (g/cm²)")
print(SEP)

rad_materials = [
    ("H₂O",     1,   1.008),
    ("Al",      13,  26.982),
    ("Si",      14,  28.086),
    ("Ar",      18,  39.948),
    ("Fe",      26,  55.845),
    ("Cu",      29,  63.546),
    ("Pb",      82, 207.200),
]

print(f"{'Material':>10}  {'Z':>4}  {'A':>7}  {'1/X₀ (cm²/g)':>14}  {'X₀ (g/cm²)':>12}")
for name, Z, A in rad_materials:
    inv = el.inv_rad_length(Z, A)
    print(f"{name:>10}  {Z:>4}  {A:>7.3f}  {inv:>14.6f}  {1/inv:>12.3f}")


# ---------------------------------------------------------------------------
# 6. Plot: muon dE/dx vs momentum (0.1 – 50 GeV)
# ---------------------------------------------------------------------------

print()
print(SEP)
print("6. Plotting muon dE/dx vs momentum → dedx_muon.png")
print(SEP)

momenta = np.logspace(np.log10(0.1), np.log10(50), 300)

plot_materials = {
    "Liquid Ar":          el.get_material("argon",       "liquid"),
    "Liquid Xe":          el.get_material("xenon",       "liquid"),
    "Silicon":            el.get_material("silicon",     "solid"),
    "Water":              el.get_material("water",       "liquid"),
    "Iron":               el.get_material("iron",        "solid"),
}

fig, ax = plt.subplots(figsize=(8, 5))

for label, mat in plot_materials.items():
    y = [el.dedx(P, el.MUON, mat, density_correction=True) for P in momenta]
    ax.plot(momenta, y, label=label)

ax.set_xscale("log")
ax.set_xlabel("Muon momentum (GeV/c)")
ax.set_ylabel(r"$\langle dE/dx \rangle$  (MeV g$^{-1}$ cm$^{2}$)")
ax.set_title("Muon energy loss (Bethe-Bloch with density correction)")
ax.legend()
ax.grid(True, which="both", linestyle="--", alpha=0.4)
ax.set_xlim(0.1, 50)

plt.tight_layout()
plt.savefig("dedx_muon.png", dpi=150)
print("Saved dedx_muon.png")
plt.show()