diff --git a/docs/examples/data/genesis4.par.h5 b/docs/examples/data/genesis4.par.h5 deleted file mode 100644 index 3c0376fa..00000000 Binary files a/docs/examples/data/genesis4.par.h5 and /dev/null differ diff --git a/docs/examples/data/genesis4/end.par.h5 b/docs/examples/data/genesis4/end.par.h5 new file mode 100644 index 00000000..4307a708 Binary files /dev/null and b/docs/examples/data/genesis4/end.par.h5 differ diff --git a/docs/examples/data/genesis4/genesis4.in b/docs/examples/data/genesis4/genesis4.in new file mode 100644 index 00000000..36eeee5e --- /dev/null +++ b/docs/examples/data/genesis4/genesis4.in @@ -0,0 +1,37 @@ +&setup + rootname = genesis4 + lattice = genesis4.lat + beamline = LAT + gamma0 = 1000.0 + lambda0 = 1e-07 + delz = 0.026 + seed = 123456 + npart = 128 +&end + +&time + slen = 7.175993210500245e-05 + sample = 10 +&end + +# 100 pC +&profile_gauss + label = beamcurrent + c0 = 1000.0 + s0 = 3.5879966052501225e-05 + sig = 1.1959988684167075e-05 +&end + +&beam + gamma = 1000.0 + delgam = 1.0 + current = @beamcurrent +&end + +&track + zstop = 1.0 +&end + +&write +beam=end +&end diff --git a/docs/examples/data/genesis4/genesis4.lat b/docs/examples/data/genesis4/genesis4.lat new file mode 100644 index 00000000..5fbb73a6 --- /dev/null +++ b/docs/examples/data/genesis4/genesis4.lat @@ -0,0 +1,2 @@ +D1: drift = {l=1.0}; +LAT: LINE = {D1}; diff --git a/docs/examples/data/genesis4/genesis4.out.h5 b/docs/examples/data/genesis4/genesis4.out.h5 new file mode 100644 index 00000000..7698df75 Binary files /dev/null and b/docs/examples/data/genesis4/genesis4.out.h5 differ diff --git a/docs/examples/data/genesis4/run.sh b/docs/examples/data/genesis4/run.sh new file mode 100755 index 00000000..242431be --- /dev/null +++ b/docs/examples/data/genesis4/run.sh @@ -0,0 +1 @@ +genesis4 genesis4.in diff --git a/docs/examples/read_examples.ipynb b/docs/examples/read_examples.ipynb index e13dd05d..44af6cf6 100644 --- a/docs/examples/read_examples.ipynb +++ b/docs/examples/read_examples.ipynb @@ -11,166 +11,187 @@ "cell_type": "code", "execution_count": null, "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:18:26.404727Z", - "iopub.status.busy": "2024-10-17T23:18:26.404338Z", - "iopub.status.idle": "2024-10-17T23:18:26.431800Z", - "shell.execute_reply": "2024-10-17T23:18:26.431113Z" - }, "tags": [] }, "outputs": [], "source": [ - "# Useful for debugging\n", - "%load_ext autoreload\n", - "%autoreload 2" + "from pmd_beamphysics import ParticleGroup\n", + "from pmd_beamphysics.interfaces.elegant import elegant_h5_to_data" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# elegant hdf5 format" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This will convert to a data dict" ] }, { "cell_type": "code", "execution_count": null, "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:18:26.434439Z", - "iopub.status.busy": "2024-10-17T23:18:26.434222Z", - "iopub.status.idle": "2024-10-17T23:18:27.006225Z", - "shell.execute_reply": "2024-10-17T23:18:27.005886Z" - }, "tags": [] }, "outputs": [], "source": [ - "from pmd_beamphysics import ParticleGroup" + "data = elegant_h5_to_data(\"data/elegant_raw.h5\")\n", + "data" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "# elegant hdf5 format" + "Create `ParticleGroup` and plot:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:18:27.007895Z", - "iopub.status.busy": "2024-10-17T23:18:27.007692Z", - "iopub.status.idle": "2024-10-17T23:18:27.018400Z", - "shell.execute_reply": "2024-10-17T23:18:27.018161Z" - }, "tags": [] }, "outputs": [], "source": [ - "from pmd_beamphysics.interfaces.elegant import elegant_h5_to_data" + "P = ParticleGroup(data=data)\n", + "P.plot(\"delta_t\", \"delta_pz\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "This will convert to a data dict" + "# Genesis4 HDF5 format\n", + "\n", + "Genesis4 stores particles in `z` slices in its HDF5 file format. Each slice has a `z` length equal to the reference wavelength `lambda0`. In order to speed up the simulation, Genesis4 also has an integer parameter [`sample`](https://github.com/svenreiche/Genesis-1.3-Version4/blob/master/manual/MAIN_INPUT.md#time) that essentially skips adjacent slices, defined in the input file:\n" ] }, { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:18:27.019838Z", - "iopub.status.busy": "2024-10-17T23:18:27.019749Z", - "iopub.status.idle": "2024-10-17T23:18:27.038539Z", - "shell.execute_reply": "2024-10-17T23:18:27.038287Z" - }, - "tags": [] - }, + "metadata": {}, "outputs": [], "source": [ - "data = elegant_h5_to_data(\"data/elegant_raw.h5\")\n", - "data" + "# Input\n", + "!cat data/genesis4/genesis4.in" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Lattice\n", + "!cat data/genesis4/genesis4.lat" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "Create `ParticleGroup` and plot:" + "Running Genesis4 on this input (`genesis4 genesis4.in`) produces the raw particle file `end.par.h5` with slice particle." ] }, { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:18:27.060724Z", - "iopub.status.busy": "2024-10-17T23:18:27.060588Z", - "iopub.status.idle": "2024-10-17T23:18:27.383511Z", - "shell.execute_reply": "2024-10-17T23:18:27.383284Z" - }, - "tags": [] - }, + "metadata": {}, "outputs": [], "source": [ - "P = ParticleGroup(data=data)\n", - "P.plot(\"delta_t\", \"delta_pz\")" + "parfile = \"data/genesis4/end.par.h5\"" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "# Genesis4 HDF5 format" + "By default `ParticleGroup.from_genesis4` will load these particles relative to their slice position. Because `sample=10` we see a comb-like distribution:" ] }, { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:18:27.385069Z", - "iopub.status.busy": "2024-10-17T23:18:27.384959Z", - "iopub.status.idle": "2024-10-17T23:18:27.395330Z", - "shell.execute_reply": "2024-10-17T23:18:27.395103Z" - } - }, + "metadata": {}, "outputs": [], "source": [ - "from pmd_beamphysics.interfaces.genesis import genesis4_par_to_data" + "P_default = ParticleGroup.from_genesis4(parfile)\n", + "P_default.plot(\"z\", \"pz\", bins=200)\n", + "P_default" ] }, { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:18:27.396506Z", - "iopub.status.busy": "2024-10-17T23:18:27.396431Z", - "iopub.status.idle": "2024-10-17T23:18:27.491746Z", - "shell.execute_reply": "2024-10-17T23:18:27.491430Z" - } - }, + "metadata": {}, "outputs": [], "source": [ - "P = ParticleGroup(data=genesis4_par_to_data(\"data/genesis4.par.h5\"))\n", + "# ParticleGroup has a `smear` option that shifts particles in each slice by a random integer multiple of the slice spacing, so that the overall distrubution is smoother while preserving the bunching characteristics:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "P_smear = ParticleGroup.from_genesis4(parfile, smear=True)\n", + "P_smear.plot(\"z\", \"pz\", bins=100)\n", + "P_smear" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Looking closer:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "P = P_default\n", + "P[(P.z > 30e-6) & (P.z < 35e-6)].plot(\"z\", \"pz\", bins=100)\n", "P" ] }, { "cell_type": "code", "execution_count": null, - "metadata": { - "execution": { - "iopub.execute_input": "2024-10-17T23:18:27.493261Z", - "iopub.status.busy": "2024-10-17T23:18:27.493141Z", - "iopub.status.idle": "2024-10-17T23:18:27.692929Z", - "shell.execute_reply": "2024-10-17T23:18:27.692655Z" - } - }, + "metadata": {}, + "outputs": [], + "source": [ + "P = P_smear\n", + "P[(P.z > 30e-6) & (P.z < 35e-6)].plot(\"z\", \"pz\", bins=100)\n", + "P" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Normally Genesis4 populates the same number of particles in each slice, with each slice having its own weight (internally the `current`). In some cases it's simpler to have particles that have the same weight, so `ParticleGroup` further has an `equal_weights` option that samples particles from the slices according to the relative weights. Notice that there are fewer particles in this case, but that the projected densities are roughtly the same as above." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, "outputs": [], "source": [ - "P.plot(\"z\", \"pz\")" + "P_equal = ParticleGroup.from_genesis4(parfile, smear=True, equal_weights=True)\n", + "P_equal.plot(\"z\", \"pz\", bins=100)\n", + "P_equal" ] } ], @@ -190,7 +211,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.7" + "version": "3.13.3" } }, "nbformat": 4, diff --git a/pmd_beamphysics/interfaces/genesis.py b/pmd_beamphysics/interfaces/genesis.py index c41646b9..466be373 100644 --- a/pmd_beamphysics/interfaces/genesis.py +++ b/pmd_beamphysics/interfaces/genesis.py @@ -1,4 +1,7 @@ import os +from pathlib import Path +from collections import namedtuple +from typing import Union, Optional import numpy as np from h5py import File, Group @@ -436,133 +439,288 @@ def write_genesis4_distribution(particle_group, h5file, verbose=False): print(f"Datasets x, xp, y, yp, t, p written to: {h5file}") -def genesis4_par_to_data(h5, species="electron", smear=True): +PARFILE_SLICE_FIELDS = ("x", "px", "y", "py", "gamma", "theta", "current") +ParfileSliceData = namedtuple("RawSliceData", PARFILE_SLICE_FIELDS) + + +# Known scalars +_parfile_scalar_datasets = [ + "beamletsize", + "one4one", + "refposition", + "slicecount", + "slicelength", + "slicespacing", +] + +_parfile_skip_groups = [ + "Meta", +] + + +def genesis4_parfile_scalars(h5): + """ + Extract useful scalars and slice names from a Genesis4 .par file """ - Converts Genesis 4 data from an h5 handle or file to data for - openPMD-beamphysics. - - Genesis4 datasets in the HDF5 file are named: - 'x' - x position in meters - 'px' - = gamma * beta_x - 'y' - y position in meters - 'py' - = gamma * beta_y' - 'theta' - angle within a slice in radians - 'gamma' - relativistic gamma - 'current' - Current in a single slice (scalar) in Amps + # Allow for opening a file + if isinstance(h5, (str, Path)): + assert os.path.exists(h5), f"File does not exist: {h5}" + h5 = File(h5, "r") + + params = {} + for k in _parfile_scalar_datasets: + if k not in h5 or h5[k].shape != (1,): + raise ValueError( + f"Expected scalar dataset '{k}' with shape (1,), got shape {h5[k].shape}" + ) + params[k] = h5[k][0] + return params + + +def genesis4_parfile_slice_groups(h5): + """ + Extract useful scalars and slice names from a Genesis4 .par file + """ + # Allow for opening a file + if isinstance(h5, (str, Path)): + assert os.path.exists(h5), f"File does not exist: {h5}" + h5 = File(h5, "r") + + return sorted( + g for g in h5 if g not in _parfile_scalar_datasets + _parfile_skip_groups + ) + + +def load_parfile_slice_data(group): + """ + Returns ParfileSliceData with Genesis4's named fields in a slice group. + """ + values = [group[field][:] for field in PARFILE_SLICE_FIELDS[:-1]] + current = group["current"][:] # This is really a scalar + assert len(current) == 1 + values.append(current[0]) + return ParfileSliceData(*values) + + +def genesis4_par_to_data( + h5: Union[str, Path, File], + species: str = "electron", + smear: bool = False, + wrap: bool = False, + z0: float = 0, + slices: Optional[list[int]] = None, + equal_weights: bool = False, + cutoff: float = 1.6e-19, + rng: Optional[np.random.Generator] = None, +) -> dict: + """ + Convert Genesis 4 `.par` slice data from an HDF5 file into a dictionary + compatible with `openPMD-beamphysics` ParticleGroup input. + + Each slice in the Genesis 4 output contains sampled macroparticles with + position, momentum, and phase data. This function reconstructs full 6D + coordinates and weights, with options for smearing, resampling, and filtering. Parameters ---------- - h5 : open h5py handle or str + h5 : str, Path, or h5py.File + Path to a Genesis 4 `.par` HDF5 file, or an open h5py file handle. species : str, default="electron" + Particle species. Currently, only "electron" is supported. - smear : bool, default=True + smear : bool, default=False Genesis4 often samples the beam by skipping slices in a step called 'sample'. This will smear out the theta coordinate over these slices, preserving the modulus. + slices : list of int, optional + Specific slice indices to include. If None, all available slices are used. + + wrap: bool, default = False + If true, the z position within a slice will be modulo the slice length. + + z0 : float, default=0.0 + Initial z-position for slice index 1, in meters. + + equal_weights : bool, default=False + If True, particles are resampled within each slice so that all output + particles have equal charge weights. Useful for simulations that + require uniform macroparticles. + + cutoff : float, default=1.6e-19 + Minimum per-particle weight in Coulombs. Slices resulting in + sub-electron macroparticles are skipped to avoid numerical issues. + + rng : numpy.random.Generator or seed, optional + Random number generator or seed used for smearing and resampling. + If None, a new default RNG is created. + + Returns ------- data : dict - For ParticleGroup + Dictionary containing particle phase space and metadata in openPMD-compatible format: + { + "x", "y", "z" : position [m], + "px", "py", "pz" : momentum [eV/c], + "t" : time [s], + "status" : particle status flag (all 1), + "species" : particle species string, + "weight" : particle weight [Coulombs], + } + + Notes + ----- + - The function expects Genesis slice groups named "slice000001", etc., each containing: + - 'x': x position in meters + - 'px': gamma * beta_x + - 'y': y position in meters + - 'py': gamma * beta_y + - 'theta': longitudinal phase angle in radians + - 'gamma': relativistic gamma + - 'current': scalar slice current in Amperes + + - Scalar metadata (e.g., `slicespacing`, `slicelength`) must be stored as 1-element datasets. + + - z positions are reconstructed from theta (ponderomotive phase), with optional smearing + to account for undersampled slices, and offset by slice position and `z0`. + + - Transverse momenta px and py are scaled by mec² to convert to eV/c. + """ # Allow for opening a file - if isinstance(h5, str): + if isinstance(h5, (str, Path)): assert os.path.exists(h5), f"File does not exist: {h5}" h5 = File(h5, "r") if species != "electron": raise ValueError("Only electrons supported for Genesis4") - # Scalar arrays. - # TODO: use refposition? - scalars = [ - "beamletsize", - "one4one", - "refposition", - "slicecount", - "slicelength", - "slicespacing", - ] - - params = {} - units = {} - for k in scalars: - assert len(h5[k]) == 1 - params[k] = h5[k][0] - if "unit" in h5[k].attrs: - units[k] = h5[k].attrs["unit"].decode("utf-8") + # Extract scalar parameters + params = genesis4_parfile_scalars(h5) # Useful local variables ds_slice = params["slicelength"] # single slice length s_spacing = params["slicespacing"] # spacing between slices - sample = int(s_spacing / ds_slice) # This should be an integer - - x = [] - px = [] - y = [] - py = [] - z = [] - gamma = [] - weight = [] - - i0 = 0 - for sname in sorted(g for g in h5 if g not in scalars): + sample = round(s_spacing / ds_slice) # This should be an integer + + # + xs = [] + pxs = [] + ys = [] + pys = [] + zs = [] + gammas = [] + weights = [] + + if slices is None: + snames = genesis4_parfile_slice_groups(h5) + + else: + snames = [f"slice{ix:06}" for ix in slices] + + rng = np.random.default_rng(rng) + + for sname in snames: + ix = int(sname[5:]) # Extract slice index + + # Skip missing + if sname not in h5: + continue + g = h5[sname] if not isinstance(g, Group) or "current" not in g: # Groups like 'Meta' do not contain slice data. continue - current = g["current"][:] # I * s_spacing/c = Q - assert len(current) == 1 + pdata = load_parfile_slice_data(g) + + current = pdata.current # I * s_spacing/c = Q + n1 = len(pdata.x) + + # Convert current to weight (C) + # I * s_spacing/c = Q + # Single charge + q1 = current * s_spacing / c_light / n1 + + # Skip subphysical particles + if q1 < cutoff: + continue + # Skip zero current slices. These usually have nans in the particle data. if current == 0: - i0 += sample continue - x.append(g["x"][:]) - px.append(g["px"][:] * mec2) - y.append(g["y"][:]) - py.append(g["py"][:] * mec2) - gamma.append(g["gamma"][:]) + # Calculate z + theta = pdata.theta # Smear theta over sample slices - theta = g["theta"][:] - irel = theta / (2 * np.pi) % 1 # Relative bin position (0,1) - n1 = len(theta) + irel = theta / (2 * np.pi) + if wrap: + irel = irel % 1 # Relative bin position (0,1) + + # Random smear if smear: - z1 = ( - irel + np.random.randint(0, high=sample, size=n1) + i0 - ) * ds_slice # Random smear + z1 = (irel + rng.integers(0, sample, size=n1)) * ds_slice else: - z1 = (irel + i0) * ds_slice - z.append(z1) - - # Convert current to weight (C) - # I * s_spacing/c = Q - q1 = np.full(n1, current) * s_spacing / c_light / n1 - weight.append(q1) - - i0 += sample # skip samples - - # Collect - x = np.hstack(x) - px = np.hstack(px) - y = np.hstack(y) - py = np.hstack(py) - gamma = np.hstack(gamma) - z = np.hstack(z) - weight = np.hstack(weight) - + z1 = (irel) * ds_slice + + z1 = z1 + (ix - 1) * s_spacing + z0 # set absolute z + + # Collect arrays + xs.append(pdata.x) + pxs.append(pdata.px * mec2) + ys.append(pdata.y) + pys.append(pdata.py * mec2) + gammas.append(pdata.gamma) + zs.append(z1) + weights.append(np.full(n1, q1)) + + if equal_weights: + # resample each slice + + n_slices = len(weights) + + counts = list(set([len(w) for w in weights])) + assert ( + len(counts) == 1 + ) # All slices are supposed to have the same number of particles + n1 = counts[0] + + slice_charges = np.array([np.sum(w) for w in weights]) + max_slice_charge = np.max( + slice_charges + ) # use this as an upper limit for sampling + n_samples = (n1 * slice_charges / max_slice_charge).astype(int) + total_charge = np.sum(slice_charges) + n_samples_total = np.sum(n_samples) + weight1 = total_charge / n_samples_total # new equal weight + + # Loop over populated slices (note that some are skipped above) + for i in range(n_slices): + n_sample = n_samples[i] + samples = rng.choice(n1, n_sample, replace=False) + xs[i] = xs[i][samples] + pxs[i] = pxs[i][samples] + ys[i] = ys[i][samples] + pys[i] = pys[i][samples] + zs[i] = zs[i][samples] + gammas[i] = gammas[i][samples] + weights[i] = np.full(n_sample, weight1) + + # Stack + x = np.hstack(xs) + px = np.hstack(pxs) + y = np.hstack(ys) + py = np.hstack(pys) + gamma = np.hstack(gammas) + z = np.hstack(zs) + weight = np.hstack(weights) + + # Form final particlegroup data n = len(weight) p = np.sqrt(gamma**2 - 1) * mec2 pz = np.sqrt(p**2 - px**2 - py**2) @@ -575,7 +733,7 @@ def genesis4_par_to_data(h5, species="electron", smear=True): "px": px, "py": py, "pz": pz, - "t": np.full(n, 0), + "t": np.full(n, 0.0), "status": np.full(n, status), "species": species, "weight": weight, diff --git a/pmd_beamphysics/particles.py b/pmd_beamphysics/particles.py index f5cd70d1..5d8d0f02 100644 --- a/pmd_beamphysics/particles.py +++ b/pmd_beamphysics/particles.py @@ -2,7 +2,7 @@ import os import pathlib from copy import deepcopy -from typing import Union, Sequence +from typing import Union, Optional, Sequence import numpy as np from h5py import File @@ -13,6 +13,7 @@ from .interfaces.elegant import write_elegant from .interfaces.genesis import ( genesis2_beam_data, + genesis4_par_to_data, write_genesis2_beam_file, write_genesis4_beam, write_genesis4_distribution, @@ -907,6 +908,75 @@ def from_bmad(cls, bmad_dict): data = bmad.bmad_to_particlegroup_data(bmad_dict) return cls(data=data) + @classmethod + def from_genesis4( + cls, + h5: Union[str, pathlib.Path, File], + smear: bool = False, + wrap: bool = False, + z0: float = 0.0, + slices: Optional[list[int]] = None, + equal_weights: bool = False, + cutoff: float = 1.6e-19, + rng: Optional[np.random.Generator] = None, + ) -> "ParticleGroup": + """ + Create a ParticleGroup from a Genesis4 `.par` HDF5 file. + + This method loads slice-based macroparticle data from a Genesis 4 `.par` + file, reconstructs the 6D phase space distribution, and returns it as a + ParticleGroup. Supports optional smearing and resampling. + + Parameters + ---------- + h5 : str, pathlib.Path, or h5py.File + Path to a Genesis 4 `.par` HDF5 file, or an open h5py file handle. + + smear : bool, default=False + Genesis4 often samples the beam by skipping slices + in a step called 'sample'. + This will smear out the theta coordinate over these slices, + preserving the modulus. + + slices : list of int, optional + Specific slice indices to include. If None, all available slices are used. + + wrap: bool, default = False + If true, the z position within a slice will be modulo the slice length. + + z0 : float, default=0.0 + Initial z-position for slice index 1, in meters. + + equal_weights : bool, default=False + If True, particles are resampled within each slice so that all output + particles have equal charge weights. Useful for simulations that + require uniform macroparticles. + + cutoff : float, default=1.6e-19 + Minimum per-particle weight in Coulombs. Slices resulting in + sub-electron macroparticles are skipped to avoid numerical issues. + + rng : numpy.random.Generator or seed, optional + Random number generator or seed used for smearing and resampling. + If None, a new default RNG is created. + + Returns + ------- + ParticleGroup + A ParticleGroup instance constructed from the `.par` file. + """ + data = genesis4_par_to_data( + h5=h5, + smear=smear, + wrap=wrap, + z0=z0, + slices=slices, + equal_weights=equal_weights, + cutoff=cutoff, + rng=rng, + ) + return cls(data=data) + # ------- # Writers