Beam Class Regression Tests¶

---

In [1]:

%load_ext autoreload
%autoreload 2

%load_ext autoreload %autoreload 2

In [2]:

import numpy as np
from distgen import Generator
from distgen.physical_constants import PHYSICAL_CONSTANTS
from distgen.tools import check_abs_and_rel_tols
from glob import glob

import numpy as np from distgen import Generator from distgen.physical_constants import PHYSICAL_CONSTANTS from distgen.tools import check_abs_and_rel_tols from glob import glob

In [3]:

yaml_files = list(glob("../examples/data/*.yaml"))

coordinates = {
    "x",
    "y",
    "z",
    "px",
    "py",
    "pz",
    "r",
    "theta",
    "pr",
    "ptheta",
    "xp",
    "yp",
    "thetax",
    "thetay",
    "gamma",
    "energy",
    "kinetic_energy",
    "beta_x",
    "beta_y",
    "beta_z",
}


def run_test_on_input_file(input_file, test):
    G = Generator(input_file, verbose=0)
    G["n_particle"] = 10_000
    test(G.beam())


def run_test_on_input_files(input_files, test, verbose=False):
    for input_file in input_files:
        if verbose:
            print("Testing:", input_file)

        run_test_on_input_file(input_file, test)

yaml_files = list(glob("../examples/data/*.yaml")) coordinates = { "x", "y", "z", "px", "py", "pz", "r", "theta", "pr", "ptheta", "xp", "yp", "thetax", "thetay", "gamma", "energy", "kinetic_energy", "beta_x", "beta_y", "beta_z", } def run_test_on_input_file(input_file, test): G = Generator(input_file, verbose=0) G["n_particle"] = 10_000 test(G.beam()) def run_test_on_input_files(input_files, test, verbose=False): for input_file in input_files: if verbose: print("Testing:", input_file) run_test_on_input_file(input_file, test)

Statistical Tests¶

$\sum_i w_i = 1$

In [4]:

def test_weight_normalization(beam):
    check_abs_and_rel_tols(
        "macroparticle weights", np.sum(beam["w"]), 1.0, abs_tol=1e-12, rel_tol=1e-15
    )


run_test_on_input_files(yaml_files, test_weight_normalization)

def test_weight_normalization(beam): check_abs_and_rel_tols( "macroparticle weights", np.sum(beam["w"]), 1.0, abs_tol=1e-12, rel_tol=1e-15 ) run_test_on_input_files(yaml_files, test_weight_normalization)

$\langle \mathcal{O}\rangle = \sum_i w_i \mathcal{O}_i$

In [5]:

def test_avg(beam):
    for var in coordinates:
        avg_beam, avg_numpy = beam.avg(var), np.sum(beam["w"] * beam[var])
        check_abs_and_rel_tols(
            "beam.avg", avg_beam, avg_numpy, abs_tol=1e-12, rel_tol=1e-15
        )


run_test_on_input_files(yaml_files, test_avg)

def test_avg(beam): for var in coordinates: avg_beam, avg_numpy = beam.avg(var), np.sum(beam["w"] * beam[var]) check_abs_and_rel_tols( "beam.avg", avg_beam, avg_numpy, abs_tol=1e-12, rel_tol=1e-15 ) run_test_on_input_files(yaml_files, test_avg)

$\sigma_{\mathcal{O}}^2 = \sum_i w_i (\mathcal{O}_i-\langle \mathcal{O}\rangle)^2 $

In [6]:

def test_std(beam):
    for var in coordinates:
        sigma2_beam = beam.std(var) ** 2
        sigma2_numpy = np.sum(beam["w"] * (beam[var] - beam.avg(var)) ** 2)
        check_abs_and_rel_tols(
            "beam.std", sigma2_beam, sigma2_numpy, abs_tol=1e-8, rel_tol=1e-15
        )


run_test_on_input_files(yaml_files, test_std)

def test_std(beam): for var in coordinates: sigma2_beam = beam.std(var) ** 2 sigma2_numpy = np.sum(beam["w"] * (beam[var] - beam.avg(var)) ** 2) check_abs_and_rel_tols( "beam.std", sigma2_beam, sigma2_numpy, abs_tol=1e-8, rel_tol=1e-15 ) run_test_on_input_files(yaml_files, test_std)

Cylindrical Coordinates¶

---

Getting¶

$r=\sqrt{ x^2 + y^2 }$

In [7]:

def test_r(beam):
    check_abs_and_rel_tols(
        "beam.r",
        beam["r"],
        np.sqrt(beam["x"] ** 2 + beam["y"] ** 2),
        abs_tol=1e-15,
        rel_tol=1e-15,
    )


run_test_on_input_files(yaml_files, test_r)

def test_r(beam): check_abs_and_rel_tols( "beam.r", beam["r"], np.sqrt(beam["x"] ** 2 + beam["y"] ** 2), abs_tol=1e-15, rel_tol=1e-15, ) run_test_on_input_files(yaml_files, test_r)

$x=r\cos\theta$

In [8]:

def test_x(beam):
    check_abs_and_rel_tols(
        "beam.x = r cos(theta)",
        beam["x"],
        beam["r"] * np.cos(beam["theta"]),
        abs_tol=1e-12,
        rel_tol=1e-11,
    )


run_test_on_input_files(yaml_files, test_x)

def test_x(beam): check_abs_and_rel_tols( "beam.x = r cos(theta)", beam["x"], beam["r"] * np.cos(beam["theta"]), abs_tol=1e-12, rel_tol=1e-11, ) run_test_on_input_files(yaml_files, test_x)

$y = r\sin\theta$

In [9]:

def test_y(beam):
    check_abs_and_rel_tols(
        "beam.y = r sin(theta)",
        beam["y"],
        beam["r"] * np.sin(beam["theta"]),
        abs_tol=1e-12,
        rel_tol=1e-11,
    )


run_test_on_input_files(yaml_files, test_y)

def test_y(beam): check_abs_and_rel_tols( "beam.y = r sin(theta)", beam["y"], beam["r"] * np.sin(beam["theta"]), abs_tol=1e-12, rel_tol=1e-11, ) run_test_on_input_files(yaml_files, test_y)

$p_r = p_x\cos\theta + p_y\sin\theta$

In [10]:

def test_pr(beam):
    check_abs_and_rel_tols(
        "beam.pr = px cos(theta) + py sin(theta)",
        beam["pr"],
        beam["px"] * np.cos(beam["theta"]) + beam["py"] * np.sin(beam["theta"]),
        abs_tol=1e-12,
        rel_tol=1e-15,
    )


run_test_on_input_files(yaml_files, test_pr)

def test_pr(beam): check_abs_and_rel_tols( "beam.pr = px cos(theta) + py sin(theta)", beam["pr"], beam["px"] * np.cos(beam["theta"]) + beam["py"] * np.sin(beam["theta"]), abs_tol=1e-12, rel_tol=1e-15, ) run_test_on_input_files(yaml_files, test_pr)

$p_{\theta} = -p_x\sin\theta + p_y\cos\theta$

In [11]:

def test_ptheta(beam):
    check_abs_and_rel_tols(
        "beam.ptheta = -px sin(theta) + py cos(theta)",
        beam["ptheta"],
        -beam["px"] * np.sin(beam["theta"]) + beam["py"] * np.cos(beam["theta"]),
        abs_tol=1e-12,
        rel_tol=1e-15,
    )


run_test_on_input_files(yaml_files, test_ptheta)

def test_ptheta(beam): check_abs_and_rel_tols( "beam.ptheta = -px sin(theta) + py cos(theta)", beam["ptheta"], -beam["px"] * np.sin(beam["theta"]) + beam["py"] * np.cos(beam["theta"]), abs_tol=1e-12, rel_tol=1e-15, ) run_test_on_input_files(yaml_files, test_ptheta)

Transverse Derivatives and Angles¶

---

Getting¶

$x^{\prime} = p_x/p_z$

In [12]:

def test_xp(beam):
    check_abs_and_rel_tols(
        "beam.xp = px/pz",
        beam["xp"],
        beam["px"].to(beam["pz"].units) / beam["pz"],
        abs_tol=1e-14,
        rel_tol=1e-15,
    )


run_test_on_input_files(yaml_files, test_xp)

def test_xp(beam): check_abs_and_rel_tols( "beam.xp = px/pz", beam["xp"], beam["px"].to(beam["pz"].units) / beam["pz"], abs_tol=1e-14, rel_tol=1e-15, ) run_test_on_input_files(yaml_files, test_xp)

$y^{\prime} = p_y/p_z$

In [13]:

def test_yp(beam):
    check_abs_and_rel_tols(
        "beam.yp = py/pz",
        beam["yp"],
        beam["py"].to(beam["pz"].units) / beam["pz"],
        abs_tol=1e-14,
        rel_tol=1e-15,
    )


run_test_on_input_files(yaml_files, test_yp)

def test_yp(beam): check_abs_and_rel_tols( "beam.yp = py/pz", beam["yp"], beam["py"].to(beam["pz"].units) / beam["pz"], abs_tol=1e-14, rel_tol=1e-15, ) run_test_on_input_files(yaml_files, test_yp)

$\theta_x = \arctan(p_x/p_z)$

In [14]:

def test_thetax(beam):
    check_abs_and_rel_tols(
        "beam.thetax = arctan(px/pz)",
        beam["thetax"],
        np.arctan2(beam["px"].to(beam["pz"].units), beam["pz"]),
        abs_tol=1e-15,
        rel_tol=1e-15,
    )


run_test_on_input_files(yaml_files, test_thetax)

def test_thetax(beam): check_abs_and_rel_tols( "beam.thetax = arctan(px/pz)", beam["thetax"], np.arctan2(beam["px"].to(beam["pz"].units), beam["pz"]), abs_tol=1e-15, rel_tol=1e-15, ) run_test_on_input_files(yaml_files, test_thetax)

$\theta_y = \arctan(p_y/p_z)$

In [15]:

def test_thetay(beam):
    check_abs_and_rel_tols(
        "beam.thetay = arctan(py/pz)",
        beam["thetay"],
        np.arctan2(beam["py"].to(beam["pz"].units), beam["pz"]),
        abs_tol=1e-15,
        rel_tol=1e-15,
    )


run_test_on_input_files(yaml_files, test_thetay)

def test_thetay(beam): check_abs_and_rel_tols( "beam.thetay = arctan(py/pz)", beam["thetay"], np.arctan2(beam["py"].to(beam["pz"].units), beam["pz"]), abs_tol=1e-15, rel_tol=1e-15, ) run_test_on_input_files(yaml_files, test_thetay)

Relativistic Quantities¶

---

$p=\sqrt{p_x^2 + p_y^2 + p_z^2}$

In [16]:

def test_p(beam):
    check_abs_and_rel_tols(
        "beam.p = sqrt(px^2 + py^2 + pz^2)",
        beam["p"],
        np.sqrt(beam["px"] ** 2 + beam["py"] ** 2 + beam["pz"] ** 2),
        abs_tol=1e-15,
        rel_tol=1e-15,
    )


run_test_on_input_files(yaml_files, test_p)

def test_p(beam): check_abs_and_rel_tols( "beam.p = sqrt(px^2 + py^2 + pz^2)", beam["p"], np.sqrt(beam["px"] ** 2 + beam["py"] ** 2 + beam["pz"] ** 2), abs_tol=1e-15, rel_tol=1e-15, ) run_test_on_input_files(yaml_files, test_p)

$E = \sqrt{c^2|\vec{p}|^2 + (mc^2)^2}$

In [17]:

def test_energy(beam):
    c = PHYSICAL_CONSTANTS["speed of light in vacuum"]

    check_abs_and_rel_tols(
        "beam.energy = sqrt(c^2p^2 + (mc^2)^2)",
        beam["energy"],
        np.sqrt(c**2 * beam["p"] ** 2 + beam.mc2**2),
        abs_tol=1e-9,
        rel_tol=1e-15,
    )


run_test_on_input_files(yaml_files, test_energy)

def test_energy(beam): c = PHYSICAL_CONSTANTS["speed of light in vacuum"] check_abs_and_rel_tols( "beam.energy = sqrt(c^2p^2 + (mc^2)^2)", beam["energy"], np.sqrt(c**2 * beam["p"] ** 2 + beam.mc2**2), abs_tol=1e-9, rel_tol=1e-15, ) run_test_on_input_files(yaml_files, test_energy)

$\gamma = \sqrt{1+\left(\frac{p}{mc}\right)^2}$, $E/mc^2$

In [18]:

def test_gamma(beam):
    mc = beam.species_mass * PHYSICAL_CONSTANTS["speed of light in vacuum"]

    check_abs_and_rel_tols(
        "beam.gamma = sqrt( 1 + (p/mc)^2 )",
        beam["gamma"],
        np.sqrt(1 + (beam["p"] / mc).to_reduced_units() ** 2),
        abs_tol=1e-10,
        rel_tol=1e-10,
    )

    check_abs_and_rel_tols(
        "beam.gamma = E/mc^2",
        beam["gamma"],
        beam["energy"] / beam.mc2,
        abs_tol=1e-15,
        rel_tol=1e-15,
    )


run_test_on_input_files(yaml_files, test_gamma)

def test_gamma(beam): mc = beam.species_mass * PHYSICAL_CONSTANTS["speed of light in vacuum"] check_abs_and_rel_tols( "beam.gamma = sqrt( 1 + (p/mc)^2 )", beam["gamma"], np.sqrt(1 + (beam["p"] / mc).to_reduced_units() ** 2), abs_tol=1e-10, rel_tol=1e-10, ) check_abs_and_rel_tols( "beam.gamma = E/mc^2", beam["gamma"], beam["energy"] / beam.mc2, abs_tol=1e-15, rel_tol=1e-15, ) run_test_on_input_files(yaml_files, test_gamma)

$\beta = \frac{c|\vec{p}|}{E}$

In [19]:

def test_beta(beam):
    check_abs_and_rel_tols(
        "beam.beta = c|p|/E",
        beam["beta"],
        PHYSICAL_CONSTANTS["speed of light in vacuum"] * beam.p / beam.energy,
        abs_tol=1e-11,
        rel_tol=1e-14,
    )

    assert max(beam["beta"]) < 1, "max(beta) > 1, faster than light particle!"


run_test_on_input_files(yaml_files, test_beta)

def test_beta(beam): check_abs_and_rel_tols( "beam.beta = c|p|/E", beam["beta"], PHYSICAL_CONSTANTS["speed of light in vacuum"] * beam.p / beam.energy, abs_tol=1e-11, rel_tol=1e-14, ) assert max(beam["beta"]) < 1, "max(beta) > 1, faster than light particle!" run_test_on_input_files(yaml_files, test_beta)

$\beta_{x_i} = \frac{cp_{x_i}}{E}$, $\beta_x = x^{\prime}\beta_z$, $\beta_y = y^{\prime}\beta_z$, $\beta_z = \frac{\beta}{\sqrt{1+(x^{\prime})^2 +(y^{\prime})^2}}$

In [20]:

def test_beta_xi(beam):
    for var in ["x", "y", "z"]:
        check_abs_and_rel_tols(
            "beam.beta_xi = c pxi/E )",
            beam[f"beta_{var}"],
            (
                PHYSICAL_CONSTANTS["speed of light in vacuum"]
                * beam[f"p{var}"]
                / beam["energy"]
            ).to_reduced_units(),
            abs_tol=1e-15,
            rel_tol=1e-15,
        )

        check_abs_and_rel_tols(
            "beam.beta_z = sign(pz)*beta/sqrt( 1 + x'^2 + y'^2 )",
            beam["beta_z"],
            np.sign(beam["pz"])
            * beam["beta"]
            / np.sqrt(1 + beam["xp"] ** 2 + beam["yp"] ** 2),
            abs_tol=1e-11,
            rel_tol=5e-9,
        )


run_test_on_input_files(yaml_files, test_beta_xi)

def test_beta_xi(beam): for var in ["x", "y", "z"]: check_abs_and_rel_tols( "beam.beta_xi = c pxi/E )", beam[f"beta_{var}"], ( PHYSICAL_CONSTANTS["speed of light in vacuum"] * beam[f"p{var}"] / beam["energy"] ).to_reduced_units(), abs_tol=1e-15, rel_tol=1e-15, ) check_abs_and_rel_tols( "beam.beta_z = sign(pz)*beta/sqrt( 1 + x'^2 + y'^2 )", beam["beta_z"], np.sign(beam["pz"]) * beam["beta"] / np.sqrt(1 + beam["xp"] ** 2 + beam["yp"] ** 2), abs_tol=1e-11, rel_tol=5e-9, ) run_test_on_input_files(yaml_files, test_beta_xi)

KE = $mc^2(\gamma-1)$, $E-mc^2$

In [21]:

def test_kinetic_energy(beam):
    if PHYSICAL_CONSTANTS.species(beam["species"])["mass"] > 0:
        check_abs_and_rel_tols(
            "beam.kinetic_energy = mc2*(gamma-1)",
            beam["kinetic_energy"],
            beam.mc2 * (beam["gamma"] - 1),
            abs_tol=1e-9,
            rel_tol=1e-05,
        )

    check_abs_and_rel_tols(
        "beam.kinetic_energy = E - mc2",
        beam["kinetic_energy"],
        beam["energy"] - beam.mc2,
        abs_tol=1e-15,
        rel_tol=1e-15,
    )


run_test_on_input_files(yaml_files, test_kinetic_energy)

def test_kinetic_energy(beam): if PHYSICAL_CONSTANTS.species(beam["species"])["mass"] > 0: check_abs_and_rel_tols( "beam.kinetic_energy = mc2*(gamma-1)", beam["kinetic_energy"], beam.mc2 * (beam["gamma"] - 1), abs_tol=1e-9, rel_tol=1e-05, ) check_abs_and_rel_tols( "beam.kinetic_energy = E - mc2", beam["kinetic_energy"], beam["energy"] - beam.mc2, abs_tol=1e-15, rel_tol=1e-15, ) run_test_on_input_files(yaml_files, test_kinetic_energy)

Twiss Parameters¶

---

Getting¶

$\epsilon_{n,x_i} = \frac{1}{mc}\sqrt{\sigma_{x_i}^2\sigma_{p_{x_i}}^2 - \langle \left(x_i-\langle x_i\rangle\right)\left(p_{x_i}-\langle p_{x_i}\rangle\right)\rangle^2 }$

In [22]:

def test_emitt_normalized(beam):
    for var in ["x", "y"]:
        mc = beam.species_mass * PHYSICAL_CONSTANTS["speed of light in vacuum"]

        stdx = beam.std(var)
        stdp = (beam.std(f"p{var}") / mc).to_reduced_units()
        dx = beam[var] - beam.avg(var)
        dp = ((beam[f"p{var}"] - beam.avg(f"p{var}")) / mc).to_reduced_units()

        check_abs_and_rel_tols(
            "beam.emitt (normalized)",
            beam.emitt(var),
            np.sqrt(stdx**2 * stdp**2 - (np.sum(beam["w"] * dx * dp)) ** 2),
            abs_tol=1e-11,
            rel_tol=1e-11,
        )


run_test_on_input_files(yaml_files, test_emitt_normalized)

def test_emitt_normalized(beam): for var in ["x", "y"]: mc = beam.species_mass * PHYSICAL_CONSTANTS["speed of light in vacuum"] stdx = beam.std(var) stdp = (beam.std(f"p{var}") / mc).to_reduced_units() dx = beam[var] - beam.avg(var) dp = ((beam[f"p{var}"] - beam.avg(f"p{var}")) / mc).to_reduced_units() check_abs_and_rel_tols( "beam.emitt (normalized)", beam.emitt(var), np.sqrt(stdx**2 * stdp**2 - (np.sum(beam["w"] * dx * dp)) ** 2), abs_tol=1e-11, rel_tol=1e-11, ) run_test_on_input_files(yaml_files, test_emitt_normalized)

$\epsilon_{x} = \sqrt{\sigma_x^2\sigma_{x^{\prime}}^2 - \langle \left(x-\langle x\rangle\right)\left(x^{\prime}-\langle x^{\prime}\rangle\right)\rangle^2 }$

In [23]:

def test_emitt_geometric(beam):
    for var in ["x", "y"]:
        stdx = beam.std(var)
        stdp = beam.std(f"{var}p")
        dx = beam[var] - beam.avg(var)
        dp = beam[f"{var}p"] - beam.avg(f"{var}p")

        check_abs_and_rel_tols(
            "beam.emitt (geometric)",
            beam.emitt(var, "geometric"),
            np.sqrt(stdx**2 * stdp**2 - (np.sum(beam["w"] * dx * dp)) ** 2),
            abs_tol=1e-14,
            rel_tol=1e-15,
        )


run_test_on_input_files(yaml_files, test_emitt_geometric)

def test_emitt_geometric(beam): for var in ["x", "y"]: stdx = beam.std(var) stdp = beam.std(f"{var}p") dx = beam[var] - beam.avg(var) dp = beam[f"{var}p"] - beam.avg(f"{var}p") check_abs_and_rel_tols( "beam.emitt (geometric)", beam.emitt(var, "geometric"), np.sqrt(stdx**2 * stdp**2 - (np.sum(beam["w"] * dx * dp)) ** 2), abs_tol=1e-14, rel_tol=1e-15, ) run_test_on_input_files(yaml_files, test_emitt_geometric)

Twiss $\beta_{x_i} = \frac{\sigma_x^2}{\epsilon_x}$

In [24]:

def twiss_beta_xi(beam):
    for var in ["x", "y"]:
        stdx = beam.std(var)
        epsx = beam.emitt(var, "geometric")

        if epsx > 0:
            check_abs_and_rel_tols(
                "beam.Beta_xi (Twiss)",
                beam.Beta(var),
                stdx**2 / epsx,
                abs_tol=1e-14,
                rel_tol=1e-15,
            )


run_test_on_input_files(yaml_files, test_beta_xi)

def twiss_beta_xi(beam): for var in ["x", "y"]: stdx = beam.std(var) epsx = beam.emitt(var, "geometric") if epsx > 0: check_abs_and_rel_tols( "beam.Beta_xi (Twiss)", beam.Beta(var), stdx**2 / epsx, abs_tol=1e-14, rel_tol=1e-15, ) run_test_on_input_files(yaml_files, test_beta_xi)

Twiss $\alpha_{x_i} = -\frac{\langle(x-\langle x\rangle)(x^{\prime}-\langle x^{\prime}\rangle)\rangle}{\epsilon_x}$

In [25]:

def test_alpha_xi(beam):
    for var in ["x", "y"]:
        dx = beam[var] - beam.avg(var)
        dp = beam[f"{var}p"] - beam.avg(f"{var}p")
        epsx = beam.emitt(var, "geometric")

        if epsx > 0:
            check_abs_and_rel_tols(
                "beam.Alpha_xi (Twiss)",
                beam.Alpha(var),
                -sum(beam["w"] * dx * dp) / epsx,
                abs_tol=1e-14,
                rel_tol=1e-14,
            )


run_test_on_input_files(yaml_files, test_alpha_xi)

def test_alpha_xi(beam): for var in ["x", "y"]: dx = beam[var] - beam.avg(var) dp = beam[f"{var}p"] - beam.avg(f"{var}p") epsx = beam.emitt(var, "geometric") if epsx > 0: check_abs_and_rel_tols( "beam.Alpha_xi (Twiss)", beam.Alpha(var), -sum(beam["w"] * dx * dp) / epsx, abs_tol=1e-14, rel_tol=1e-14, ) run_test_on_input_files(yaml_files, test_alpha_xi)

Twiss $\gamma_{x_i} = \frac{\sigma_{x^{\prime}}^2}{\epsilon_x}$

In [26]:

def test_gamma_xi(beam):
    for var in ["x", "y"]:
        stdp = beam.std(f"{var}p")
        epsx = beam.emitt(var, "geometric")

        if epsx > 0:
            check_abs_and_rel_tols(
                "beam.Gamma_xi (Twiss)",
                beam.Gamma(var),
                stdp**2 / epsx,
                abs_tol=1e-14,
                rel_tol=1e-14,
            )


run_test_on_input_files(yaml_files, test_gamma_xi)

def test_gamma_xi(beam): for var in ["x", "y"]: stdp = beam.std(f"{var}p") epsx = beam.emitt(var, "geometric") if epsx > 0: check_abs_and_rel_tols( "beam.Gamma_xi (Twiss)", beam.Gamma(var), stdp**2 / epsx, abs_tol=1e-14, rel_tol=1e-14, ) run_test_on_input_files(yaml_files, test_gamma_xi)

In [ ]:

In [ ]: