gpbeam 0.0.1__py3-none-any.whl

gpbeam/DumpCoordinates.py

@@ -0,0 +1,139 @@
+import os
+import numpy as np
+from scipy.stats import truncnorm
+from .Settings import AcceleratorConfig
+
+try:
+    import xtrack
+    xtrack_loaded = True
+except ImportError:
+    xtrack_loaded = False
+
+try:
+    import RF_Track
+    rf_track_loaded = True
+except ImportError:
+    rf_track_loaded = False
+
+
+def dump_coordinates(cur_workdir, coords: np.ndarray | xtrack.particles.particles.Particles, positron: bool):
+    if type(coords) == np.ndarray:
+        if not len(coords.shape) == 2 or coords.shape[1] != 6:
+            raise ValueError("coordinate array must have shape (N,6)")
+    elif xtrack_loaded and isinstance(coords, xtrack.particles.particles.Particles):
+        particles = coords
+        coords = np.empty((6,len(particles.x)))
+
+        coords[:,0] = particles.energy / 1e9  # [GeV]
+        coords[:,1] = particles.x * 1e6  # [um]
+        coords[:,2] = particles.y * 1e6  # [um]
+        coords[:,3] = particles.zeta * 1e6  # [um]
+        coords[:,4] = particles.px * 1e6
+        coords[:,5] = particles.py * 1e6
+    elif rf_track_loaded and isinstance(coords, RF_Track.Bunch6d):
+        particles = coords
+        coords = np.empty((6,len(particles.x)))
+
+        coords[:,0] = particles.get_phase_space("%K") / 1e3  # [GeV]
+        coords[:,1] = particles.get_phase_space("%x") * 1e3  # [um]
+        coords[:,2] = particles.get_phase_space("%y") * 1e3  # [um]
+        coords[:,3] = particles.get_phase_space("%t") * 1e3 * RF_Track.clight  # [um]
+        coords[:,4] = particles.get_phase_space("%xp") * 1e3
+        coords[:,5] = particles.get_phase_space("%yp") * 1e3
+    else:
+        raise ValueError("coords must be of either np.ndarray or xtrack.particles.particles.Particles type")
+    
+    if positron:
+        fname = "positron.ini"
+    else:
+        fname = "electron.ini"
+
+    np.savetxt(os.path.join(cur_workdir,fname), coords)
+
+    return
+
+
+def sampleGaussianWaist(accelerator: AcceleratorConfig, cut: int = 3, N: int = 100000,):
+    """
+    Sample Gaussian waist particle distributions from an AcceleratorConfig.
+    """
+    betx_1 = accelerator.beta_x[0] * 1e3  # mm → µm
+    betx_2 = accelerator.beta_x[1] * 1e3
+
+    bety_1 = accelerator.beta_y[0] * 1e3
+    bety_2 = accelerator.beta_y[1] * 1e3
+
+    sigma_z_1 = accelerator.sigma_z[0]
+    sigma_z_2 = accelerator.sigma_z[1]
+
+    energy_1 = accelerator.energy[0]
+    energy_2 = accelerator.energy[1]
+
+    espread_1 = accelerator.espread[0]
+    espread_2 = accelerator.espread[1]
+
+    lorentz_1 = (energy_1 * 1e9 - 511e3) / 511e3
+    lorentz_2 = (energy_2 * 1e9 - 511e3) / 511e3
+
+    ex_1 = accelerator.emitt_x[0] / lorentz_1
+    ex_2 = accelerator.emitt_x[1] / lorentz_2
+
+    ey_1 = accelerator.emitt_y[0] / lorentz_1
+    ey_2 = accelerator.emitt_y[1] / lorentz_2
+
+    x0_1 = accelerator.offset_x[0]
+    x0_2 = accelerator.offset_x[1]
+
+    y0_1 = accelerator.offset_y[0]
+    y0_2 = accelerator.offset_y[1]
+
+    which_espread_1 = accelerator.which_espread[0]
+    which_espread_2 = accelerator.which_espread[1]
+
+    # sample initial conditions
+    beam_1 = [
+        truncnorm.rvs(-cut, cut, loc=0, scale=np.sqrt(ex_1 * betx_1), size=N),
+        truncnorm.rvs(-cut, cut, loc=0, scale=np.sqrt(ey_1 * bety_1), size=N),
+        truncnorm.rvs(-cut, cut, loc=0, scale=sigma_z_1, size=N),
+        truncnorm.rvs(-cut, cut, loc=0, scale=np.sqrt(ex_1 * 1/betx_1), size=N),
+        truncnorm.rvs(-cut, cut, loc=0, scale=np.sqrt(ey_1 * 1/bety_1), size=N),
+    ]
+
+    beam_2 = [
+        truncnorm.rvs(-cut, cut, loc=0, scale=np.sqrt(ex_2 * betx_2), size=N),
+        truncnorm.rvs(-cut, cut, loc=0, scale=np.sqrt(ey_2 * bety_2), size=N),
+        truncnorm.rvs(-cut, cut, loc=0, scale=sigma_z_2, size=N),
+        truncnorm.rvs(-cut, cut, loc=0, scale=np.sqrt(ex_2 * 1/betx_2), size=N),
+        truncnorm.rvs(-cut, cut, loc=0, scale=np.sqrt(ey_2 * 1/bety_2), size=N),
+    ]
+
+    if which_espread_1 == 0:
+        beam_1.insert(0, np.ones(N) * energy_1)
+    elif which_espread_1 == 1:
+        beam_1.insert(0, ((np.random.random(N) - 1/2) * 2 * espread_1 + 1) * energy_1 )
+    elif which_espread_1 == 3:
+        beam_1.insert(0, truncnorm.rvs(-cut, cut, loc=energy_1, scale=espread_1, size=N))
+    else:
+        raise NotImplementedError("Beam_1: unknown setting for which_espread")
+
+    if which_espread_2 == 0:
+        beam_2.insert(0, np.ones(N) * energy_2)
+    elif which_espread_2 == 1:
+        beam_2.insert(0, ((np.random.random(N) - 1/2) * 2 * espread_2 + 1) * energy_2 )
+    elif which_espread_2 == 3:
+        beam_2.insert(0, truncnorm.rvs(-cut, cut, loc=energy_2, scale=espread_2, size=N))
+    else:
+        raise NotImplementedError("Beam_2: unknown setting for which_espread")
+
+    # add offset
+    beam_1, beam_2 = np.array(beam_1).T, np.array(beam_2).T
+
+    beam_1[:,1:] -= beam_1[:,1:].mean(axis=0)
+    beam_2[:,1:] -= beam_2[:,1:].mean(axis=0)
+
+    beam_1[:,1] += x0_1
+    beam_2[:,1] += x0_2
+    beam_1[:,2] += y0_1
+    beam_2[:,2] += y0_2
+
+    return beam_1, beam_2

gpbeam/GuineaPig.py

@@ -0,0 +1,320 @@
+import os
+import time
+import re
+import subprocess
+import shutil
+
+from dataclasses import dataclass, field, asdict
+from typing import Dict, Optional
+from pathlib import Path
+
+from .Settings import AcceleratorConfig, SimulationConfig
+from .DumpCoordinates import dump_coordinates, sampleGaussianWaist
+from .Result import parse_guinea_output
+
+
+##########################
+### general config
+##########################
+from .config import WORKDIR, EXECUTABLE_GUINEA_PLUSPLUS
+
+
+def _promote_scalars(cls, raw_dict):
+    """
+    Promote scalar values to (value, value) if dataclass field expects tuple.
+    """
+
+    from typing import get_origin
+
+    promoted = {}
+
+    for field in cls.__dataclass_fields__.values():
+        name = field.name
+        expected_type = field.type
+
+        if name not in raw_dict:
+            continue
+
+        value = raw_dict[name]
+
+        # If expected type is Tuple[...] but value is scalar → promote
+        if (
+            get_origin(expected_type) is tuple
+            and not isinstance(value, tuple)
+        ):
+            promoted[name] = (value, value)
+        else:
+            promoted[name] = value
+
+    return promoted
+
+
+@dataclass
+class GuineaPig:
+    """
+    Container for a full GuineaPig configuration project.
+
+    Holds multiple accelerator definitions and multiple simulation
+    parameter sets.
+    """
+
+    accelerators: Dict[str, "AcceleratorConfig"] = field(default_factory=dict)
+    simulations: Dict[str, "SimulationConfig"] = field(default_factory=dict)
+
+    active_accelerator: Optional[str] = None
+    active_simulation: Optional[str] = None
+
+    # ============================================================
+    # Basic management
+    # ============================================================
+
+    def add_accelerator(self, name: str, config: "AcceleratorConfig"):
+        self.accelerators[name] = config
+
+    def add_simulation(self, name: str, config: "SimulationConfig"):
+        self.simulations[name] = config
+
+    def set_active(self, accelerator: str, simulation: str):
+        if accelerator not in self.accelerators:
+            raise KeyError(f"Unknown accelerator '{accelerator}'")
+        if simulation not in self.simulations:
+            raise KeyError(f"Unknown simulation '{simulation}'")
+
+        self.active_accelerator = accelerator
+        self.active_simulation = simulation
+
+    # ============================================================
+    # Accessors
+    # ============================================================
+    @property
+    def accelerator(self) -> "AcceleratorConfig":
+        if self.active_accelerator is None:
+            # implicitly activate only choice
+            if len(self.accelerators) == 1:
+                self.active_accelerator = tuple(self.accelerators.keys())[0]
+                return self.accelerators[self.active_accelerator]
+            else:
+                raise RuntimeError("No active accelerator selected")
+        return self.accelerators[self.active_accelerator]
+
+    @property
+    def simulation(self) -> "SimulationConfig":
+        if self.active_simulation is None:
+            # implicitly activate only choice
+            if len(self.simulations) == 1:
+                self.active_simulation = tuple(self.simulations.keys())[0]
+                return self.simulations[self.active_simulation]
+            else:
+                raise RuntimeError("No active simulation selected")
+        return self.simulations[self.active_simulation]
+
+    # ============================================================
+    # File I/O
+    # ============================================================
+    @classmethod
+    def load(cls, filepath: str | Path) -> "GuineaPig":
+        with open(filepath, "r") as f:
+            text = f.read()
+
+        section_pattern = re.compile(
+            r"\$(\w+)::\s*(\w+)\s*\{([^}]*)\}",
+            re.DOTALL
+        )
+
+        project = cls()
+
+        for section, name, body in section_pattern.findall(text):
+
+            raw = {}
+            tuple_accumulator = {}
+
+            for entry in body.split(";"):
+                entry = entry.strip()
+                if not entry:
+                    continue
+
+                if "=" not in entry:
+                    continue
+
+                key, value = entry.split("=", 1)
+                key = key.strip()
+                value = value.strip()
+
+                # DO NOT evaluate expressions like 3.0*sigma_z.1
+                try:
+                    value_eval = eval(value, {"__builtins__": None}, {})
+                except Exception:
+                    value_eval = value  # keep expression string
+
+                # convert GP naming → python naming
+                key = key.replace(" ", "_")
+
+                # handle .1 / .2
+                if "." in key:
+                    base, idx = key.split(".")
+                    idx = int(idx) - 1
+
+                    tuple_accumulator.setdefault(base, [None, None])
+                    tuple_accumulator[base][idx] = value_eval
+                else:
+                    raw[key] = value_eval
+
+            # finalize tuples
+            for k, v in tuple_accumulator.items():
+                if v[0] is None or v[1] is None:
+                    raise ValueError(f"Incomplete tuple for {k}")
+                raw[k] = tuple(v)
+
+            # Promote scalar to tuple if dataclass expects tuple
+            if section.upper() == "ACCELERATOR":
+                obj = AcceleratorConfig(**_promote_scalars(AcceleratorConfig, raw))
+                project.accelerators[name] = obj
+
+            elif section.upper() == "PARAMETERS":
+                obj = SimulationConfig(**_promote_scalars(SimulationConfig, raw))
+                project.simulations[name] = obj
+
+        return project
+
+    def save(self, filepath: str | Path):
+        lines = []
+
+        def write_section(section_name, objects):
+            for name, obj in objects.items():
+
+                lines.append(f"${section_name}:: {name}")
+                lines.append("{")
+
+                data = asdict(obj)
+
+                for key, value in data.items():
+                    # gp_key = key.replace("_", " ")
+                    gp_key = key
+
+                    if isinstance(value, tuple) and len(value) == 2:
+                        lines.append(f"    {gp_key}.1={value[0]};")
+                        lines.append(f"    {gp_key}.2={value[1]};")
+                    else:
+                        lines.append(f"    {gp_key}={value};")
+
+                lines.append("}")
+
+        write_section("ACCELERATOR", self.accelerators)
+        write_section("PARAMETERS", self.simulations)
+
+        with open(filepath, "w") as f:
+            f.write("\n".join(lines) + "\n")    
+            
+    
+    # ============================================================
+    # Convenience utilities
+    # ============================================================
+    def remove_accelerator(self, name: str):
+        if name in self.accelerators:
+            del self.accelerators[name]
+            if self.active_accelerator == name:
+                self.active_accelerator = None
+
+    def remove_simulation(self, name: str):
+        if name in self.simulations:
+            del self.simulations[name]
+            if self.active_simulation == name:
+                self.active_simulation = None
+
+    def list_accelerators(self):
+        return list(self.accelerators.keys())
+
+    def list_simulations(self):
+        return list(self.simulations.keys())
+
+    def to_dict(self):
+        """
+        Export project as nested dictionary (dataclasses preserved).
+        """
+        return {
+            "ACCELERATOR": self.accelerators,
+            "PARAMETERS": self.simulations,
+        }
+
+    # ============================================================
+    # Validation
+    # ============================================================
+    def validate(self):
+        """
+        Re-run dataclass validation explicitly.
+        Useful if configs were modified after creation.
+        """
+        for acc in self.accelerators.values():
+            acc.__post_init__()
+
+        for sim in self.simulations.values():
+            sim.__post_init__()
+
+    # ============================================================
+    # Representation
+    # ============================================================
+    def __repr__(self):
+        return (
+            f"GuineaPigProject("
+            f"{len(self.accelerators)} accelerators, "
+            f"{len(self.simulations)} simulations, "
+            f"active=({self.active_accelerator}, {self.active_simulation})"
+            f")"
+        )
+
+    # ============================================================
+    # Run Guinea-Pig
+    # ============================================================
+    def simulate(self, particles_1 = None, particles_2 = None, retain_files: bool = False):
+        unique_id = abs(hash((os.getpid(), time.time())))
+        cur_workdir = os.path.normpath(WORKDIR +"/workdir_guinea_" + str(unique_id))
+        os.makedirs(cur_workdir, exist_ok=False)
+
+        # ensure proper configuration
+        _ = self.accelerator
+        _ = self.simulation
+
+        self.save(os.path.join(cur_workdir, "acc.dat"))
+                  
+        # initial conditions
+        if particles_1 is not None and particles_2 is not None:
+            dump_coordinates(cur_workdir, particles_1,positron=False)
+            dump_coordinates(cur_workdir, particles_2,positron=True)
+        elif particles_1 is not None and particles_2 is None:
+            dump_coordinates(cur_workdir, particles_1,positron=False)
+            dump_coordinates(cur_workdir, particles_1,positron=True)
+        elif particles_1 is None and particles_2 is None:
+            pass
+            particles_1, particles_2 = sampleGaussianWaist(self.accelerator)
+
+            dump_coordinates(cur_workdir, particles_1,positron=False)
+            dump_coordinates(cur_workdir, particles_2,positron=True)
+        else:
+            raise ValueError("particles_2 must be None if particles_1 is not provided")
+        
+        # call guinea <- guinea seems to read from stdin always as interactive program <- feed it with fake stdin
+        out_file = os.path.join(cur_workdir, "g++_result.out")
+        with open(os.devnull, "r") as nullin:
+            result = subprocess.run(
+                [EXECUTABLE_GUINEA_PLUSPLUS(), self.active_accelerator, self.active_simulation, out_file],
+                stdin=nullin,
+                stdout=subprocess.PIPE,    # Capture stdout
+                stderr=subprocess.STDOUT,  # Capture stderr too
+                text=True,                 # Decode bytes -> str
+                cwd=cur_workdir,
+                check=True
+            )
+        
+        simulation_result = parse_guinea_output(out_file)
+
+        if simulation_result.extra["miss"] > 0.3:
+            print("Warning: particles outside of grid, miss rate = ", simulation_result.extra["miss"])
+
+        if retain_files:
+            with open(os.path.join(cur_workdir, "guinea.log"), "w") as logfile:
+                logfile.write(result.stdout)
+        else:
+            shutil.rmtree(cur_workdir, ignore_errors=True)
+
+        return simulation_result
+

gpbeam/Result.py

@@ -0,0 +1,132 @@
+import re
+from dataclasses import dataclass, field
+from typing import Dict
+
+
+@dataclass
+class BeamResult:
+    tracked_macroparticles: int
+    energy_loss: float
+    avg_photon_energy: float
+    photons_per_particle: float
+
+
+@dataclass
+class GeneralResults:
+    lumi_fine: float
+    lumi_ee: float
+    lumi_ee_high: float
+    lumi_pp: float
+    lumi_eg: float
+    lumi_ge: float
+    lumi_gg: float
+    lumi_gg_high: float
+    upsmax: float
+    E_cm: float
+    E_cm_var: float
+
+
+@dataclass
+class GuineaPigOutput:
+    beam1: BeamResult
+    beam2: BeamResult
+    general: GeneralResults
+
+    # additional scalar key=value style results
+    extra: Dict[str, float] = field(default_factory=dict)
+
+
+# -------------------------------------------------
+# Parse GuineaPig output file
+# -------------------------------------------------
+def parse_guinea_output(file_path: str) -> GuineaPigOutput:
+    with open(file_path, "r") as f:
+        content = f.read()
+
+    # Remove everything after step table
+    content = content.split("step   lumi total   lumi peak")[0]
+
+    # -------------------------------------------------
+    # Extract beam results (compact key=value section)
+    # -------------------------------------------------
+    kv_pairs = dict(
+        # re.findall(r"([A-Za-z0-9_.]+)\s*=\s*([-+0-9.eE]+)", content)
+        re.findall(r"([A-Za-z0-9_.-]+)\s*=\s*([-+0-9.eE]+)", content)
+    )
+
+    def get_float(key, default=0.0):
+        return float(kv_pairs.get(key, default))
+
+    def get_int_from_line(pattern):
+        m = re.search(pattern, content)
+        return int(m.group(1)) if m else 0
+
+    def get_int_multiple_line(pattern):
+        m = re.search(pattern, content, re.DOTALL)
+        return int(m.group(1)) if m else 0
+
+    # -------------------------------------------------
+    # Beam 1
+    # -------------------------------------------------
+    beam1 = BeamResult(
+        tracked_macroparticles=get_int_multiple_line(
+            r"beam1.*?number of tracked macroparticles\s*:\s*(\d+)"
+        ),
+        energy_loss=get_float("de1"),
+        avg_photon_energy=get_float("phot-e1"),
+        photons_per_particle=get_float("n_phot1"),
+    )
+
+    # -------------------------------------------------
+    # Beam 2
+    # -------------------------------------------------
+    beam2 = BeamResult(
+        tracked_macroparticles=get_int_multiple_line(
+            r"beam2.*?number of tracked macroparticles\s*:\s*(\d+)"
+        ),
+        energy_loss=get_float("de2"),
+        avg_photon_energy=get_float("phot-e2"),
+        photons_per_particle=get_float("n_phot2"),
+    )
+
+    # -------------------------------------------------
+    # General results
+    # -------------------------------------------------
+    def extract_float(pattern):
+        m = re.search(pattern, content)
+        return float(m.group(1)) if m else 0.0
+
+    general = GeneralResults(
+        lumi_fine=extract_float(r"lumi_fine\s*=\s*([-+0-9.eE]+)"),
+        lumi_ee=extract_float(r"lumi_ee\s*=\s*([-+0-9.eE]+)"),
+        lumi_ee_high=extract_float(r"lumi_ee_high\s*=\s*([-+0-9.eE]+)"),
+        lumi_pp=extract_float(r"lumi_pp\s*=\s*([-+0-9.eE]+)"),
+        lumi_eg=extract_float(r"lumi_eg\s*=\s*([-+0-9.eE]+)"),
+        lumi_ge=extract_float(r"lumi_ge\s*=\s*([-+0-9.eE]+)"),
+        lumi_gg=extract_float(r"lumi_gg\s*=\s*([-+0-9.eE]+)"),
+        lumi_gg_high=extract_float(r"lumi_gg_high\s*=\s*([-+0-9.eE]+)"),
+        upsmax=extract_float(r"upsmax=\s*([-+0-9.eE]+)"),
+        E_cm=extract_float(r"E_cm\s*=\s*([-+0-9.eE]+)"),
+        E_cm_var=extract_float(r"E_cm_var\s*=\s*([-+0-9.eE]+)"),
+    )
+
+    # -------------------------------------------------
+    # Extra numeric key=value pairs (angles, hadrons, etc.)
+    # -------------------------------------------------
+    extra = {
+        k: float(v)
+        for k, v in kv_pairs.items()
+        if k not in {
+            "phot-e1", "phot-e2",
+            "n_phot1", "n_phot2",
+            "de1", "de2",
+            "lumi_ee", "lumi_ee_high"
+        }
+    }
+
+    return GuineaPigOutput(
+        beam1=beam1,
+        beam2=beam2,
+        general=general,
+        extra=extra,
+    )