diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 98f1d6d2..b80f9299 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -30,7 +30,7 @@ repos: - id: jupyter-notebook-clear-output name: jupyter-notebook-jupyter-clear files: \.ipynb$ - stages: [commit] + stages: [pre-commit] language: system entry: jupyter nbconvert --ClearOutputPreprocessor.enabled=True --inplace exclude: "^.*(Example3-TimeDependent|perlmutter_benchmark|fodo_scan|fodo_scan_model).ipynb$" diff --git a/docs/examples/genesis4/bmad_genesis4.ipynb b/docs/examples/genesis4/bmad_genesis4.ipynb new file mode 100644 index 00000000..618b9952 --- /dev/null +++ b/docs/examples/genesis4/bmad_genesis4.ipynb @@ -0,0 +1,257 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "4e9f7721-9286-415e-b623-dc34206c28e6", + "metadata": {}, + "source": [ + "# Bmad-Genesis4 interface\n", + "\n", + "Genesis4 input and lattice can be created from Bmad using PyTao's Tao interface." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f6b891d6-d1df-4763-8cc1-3494e5885e6f", + "metadata": {}, + "outputs": [], + "source": [ + "# %load_ext autoreload\n", + "# %autoreload 2" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9ab17586-a099-4993-b701-50026cc46d9f", + "metadata": {}, + "outputs": [], + "source": [ + "from pytao import Tao\n", + "\n", + "from genesis.version4 import Genesis4, Lattice, MainInput\n", + "\n", + "import matplotlib.pyplot as plt\n", + "\n", + "from pmd_beamphysics.units import mec2\n", + "import numpy as np" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b190eceb-f631-4454-89ea-4ca36efcd067", + "metadata": {}, + "outputs": [], + "source": [ + "bmad_lat = \"data/example1-steadystate/Example1.bmad\"" + ] + }, + { + "cell_type": "markdown", + "id": "70a43e54-0ce0-41fa-aaaf-91d21be78b12", + "metadata": {}, + "source": [ + "# Lattice from Tao" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9dcc9e69-67f6-4179-8127-4e92910e511c", + "metadata": {}, + "outputs": [], + "source": [ + "tao = Tao(lattice_file=bmad_lat, noplot=True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9cba8d49-77f8-4fc4-92ce-e5962c2ea725", + "metadata": {}, + "outputs": [], + "source": [ + "Lattice.from_tao(tao)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "040d210b-b14d-4eff-bb90-52bb3bf6840f", + "metadata": {}, + "outputs": [], + "source": [ + "Lattice.from_tao(tao).plot()" + ] + }, + { + "cell_type": "markdown", + "id": "34f5e0cc-6c3a-4588-bec5-ea6eeeb093d7", + "metadata": {}, + "source": [ + "# MainInput from Tao" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e8d5c31b-f6d3-44e6-af0a-d326924606d4", + "metadata": {}, + "outputs": [], + "source": [ + "MainInput.from_tao(tao)" + ] + }, + { + "cell_type": "markdown", + "id": "b4d3c317-0840-449d-9c6d-d3f9deadfb8a", + "metadata": {}, + "source": [ + "# Entire object" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "635a1fe7-40ce-403e-8f93-5fc574f4104f", + "metadata": {}, + "outputs": [], + "source": [ + "Genesis4.from_tao(tao)" + ] + }, + { + "cell_type": "markdown", + "id": "27a6ea47-2d78-4a89-a3cb-5ecfbd600f2a", + "metadata": {}, + "source": [ + "# Compare tracking " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8b40b0bc-0737-44f2-ae06-b3c4725fc7d9", + "metadata": {}, + "outputs": [], + "source": [ + "tao = Tao(lattice_file=bmad_lat, noplot=True)\n", + "\n", + "# Add various errors\n", + "\n", + "tao.cmd(\"set particle_start x = 50e-6\")\n", + "tao.cmd(\"set particle_start y = -40e-6\")\n", + "\n", + "tao.cmd(\"set particle_start px = 10e-6\")\n", + "tao.cmd(\"set particle_start py = 20e-6\")\n", + "tao.cmd(\"set particle_start pz = 1e-4\")\n", + "\n", + "tao.cmd(\"set ele qf x_offset = 50e-6\")\n", + "tao.cmd(\"set ele qd y_offset = -50e-6\")\n", + "\n", + "tao.cmd(\"set ele qf hkick = 10e-6\")\n", + "tao.cmd(\"set ele qd vkick = -10e-6\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a1033af2-1d2e-43b6-aa91-3c14eeaa05a9", + "metadata": {}, + "outputs": [], + "source": [ + "%%time\n", + "# Create and run Genesis4\n", + "G4 = Genesis4.from_tao(tao)\n", + "G4.run();" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9c4081de-1ff5-4305-a043-75cbd93a8011", + "metadata": {}, + "outputs": [], + "source": [ + "def comparison_plots(tao, genesis4):\n", + " s = tao.lat_list(\"*\", \"ele.s\")\n", + " z = G4.output.stat(\"zplot\")\n", + "\n", + " fig, axes = plt.subplots(4, figsize=(12, 8))\n", + "\n", + " ax = axes[0]\n", + " y0 = tao.lat_list(\"*\", \"orbit.vec.1\")\n", + " y1 = G4.output.stat(\"beam_xposition\")\n", + " ax.plot(s, y0 * 1e6, label=\"Tao\")\n", + " ax.plot(z, y1 * 1e6, \"--\", label=\"Genesis4 from Tao\")\n", + " ax.set_ylabel(r\"$\\left$ (µm)\")\n", + " ax.set_ylim(-1000, 1000)\n", + " ax.legend()\n", + "\n", + " ax = axes[1]\n", + " y0 = tao.lat_list(\"*\", \"orbit.vec.3\")\n", + " y1 = G4.output.stat(\"beam_yposition\")\n", + " ax.plot(s, y0 * 1e6, label=\"Tao\")\n", + " ax.plot(z, y1 * 1e6, \"--\", label=\"Genesis4 from Tao\")\n", + " ax.set_ylabel(r\"$\\left$ (µm)\")\n", + " ax.set_ylim(-1000, 1000)\n", + " ax.legend()\n", + "\n", + " ax = axes[2]\n", + " y0 = np.sqrt(\n", + " tao.lat_list(\"*\", \"ele.a.beta\")\n", + " * G4.output.beam.emitx[0, 0]\n", + " / (tao.lat_list(\"*\", \"ele.e_tot\") / mec2)\n", + " )\n", + " y1 = G4.output.stat(\"beam_xsize\")\n", + " ax.plot(s, y0 * 1e6, label=\"Tao\")\n", + " ax.plot(z, y1 * 1e6, \"--\", label=\"Genesis4 from Tao\")\n", + " ax.set_ylabel(r\"$\\sigma_x$ (µm)\")\n", + " ax.set_ylim(0, 30)\n", + "\n", + " ax = axes[3]\n", + " y0 = np.sqrt(\n", + " tao.lat_list(\"*\", \"ele.b.beta\")\n", + " * G4.output.beam.emity[0, 0]\n", + " / (tao.lat_list(\"*\", \"ele.e_tot\") / mec2)\n", + " )\n", + " y1 = G4.output.stat(\"beam_ysize\")\n", + " ax.plot(s, y0 * 1e6, label=\"Tao\")\n", + " ax.plot(z, y1 * 1e6, \"--\", label=\"Genesis4 from Tao\")\n", + " ax.set_ylabel(r\"$\\sigma_y$ (µm)\")\n", + " ax.set_ylim(0, 30)\n", + "\n", + " ax.legend()\n", + " ax.set_xlabel(r\"$z$ (m)\")\n", + " for ax in axes:\n", + " ax.set_xlim(0, None)\n", + "\n", + "\n", + "comparison_plots(tao, G4)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/examples/genesis4/data/example1-steadystate/Example1.bmad b/docs/examples/genesis4/data/example1-steadystate/Example1.bmad new file mode 100644 index 00000000..c4b8ee6e --- /dev/null +++ b/docs/examples/genesis4/data/example1-steadystate/Example1.bmad @@ -0,0 +1,31 @@ +no_digested +! These are the periodic Twiss +beginning[beta_a] = 8.53711320 +beginning[beta_b] = 17.38992039 +beginning[alpha_a] = -0.70330585 +beginning[alpha_b] = 1.40347554 + +beginning[e_tot] = 11357.82 * m_electron + +parameter[geometry] = open ! closed will calculate the periodic Twiss +parameter[particle] = electron + +! Genesis4 undulator parameters: +lambdau = 0.015000 +nwig = 266 +aw = 0.84853 + +D1: DRIFT, L = 0.44 +D2: DRIFT, L = 0.24 +QF: QUADRUPOLE, L = 0.080000, k1= 2.000000 +QD: QUADRUPOLE, L = 0.080000, k1= -2.000000 +UND: WIGGLER, L = lambdau*nwig, + L_period=lambdau, n_period=nwig, + field_calc = helical_model, + b_max = 2*pi*m_electron * aw / (c_light * lambdau) + +FODO: LINE= (UND,D1,QF,D2,UND,D1,QD,D2) + +FEL: LINE= (6*FODO) + +use, FEL diff --git a/environment.yml b/environment.yml index cc5022cc..d8c40e26 100644 --- a/environment.yml +++ b/environment.yml @@ -17,6 +17,7 @@ dependencies: - pydantic >2 - pydantic-settings # Developer + - pytao - ffmpeg - pytest - pytest-cov diff --git a/genesis/version4/input/core.py b/genesis/version4/input/core.py index ade59dca..e300d7b9 100644 --- a/genesis/version4/input/core.py +++ b/genesis/version4/input/core.py @@ -39,6 +39,7 @@ ) from .. import archive as _archive from ..field import FieldFile +from ..interfaces.bmad import genesis4_namelists_from_tao from ..particles import Genesis4ParticleData from ..types import ( AnyPath, @@ -377,6 +378,11 @@ def to_file( filename, ) + @classmethod + def from_tao(cls, tao): + namelists = genesis4_namelists_from_tao(tao) + return cls(namelists=namelists) + def _check_for_mistakes(self) -> None: """Check and fix simple mistakes to be friendly to the end-user.""" try: diff --git a/genesis/version4/input/lattice.py b/genesis/version4/input/lattice.py index 3231c9c4..62c318ee 100644 --- a/genesis/version4/input/lattice.py +++ b/genesis/version4/input/lattice.py @@ -40,6 +40,8 @@ ) from . import _lattice as auto_lattice from . import parsers +from ..interfaces.bmad import genesis4_elements_and_line_from_tao + try: from typing import Literal @@ -1004,6 +1006,14 @@ def from_elements( filename=pathlib.Path(filename) if filename else None, ) + @classmethod + def from_tao(cls, tao): + elements, line_labels = genesis4_elements_and_line_from_tao(tao) + + line_label = tao.branch1(1, 0)["name"] + elements[line_label] = Line(elements=line_labels) + return cls(elements=elements) + def to_file( self, filename: AnyPath, diff --git a/genesis/version4/interfaces/bmad.py b/genesis/version4/interfaces/bmad.py new file mode 100644 index 00000000..702ab563 --- /dev/null +++ b/genesis/version4/interfaces/bmad.py @@ -0,0 +1,317 @@ +from ..input._lattice import Quadrupole, Corrector, Drift, Marker, Undulator +from ..input._main import Setup, Field, Track, Beam + + +from pmd_beamphysics.units import mec2 + +from scipy.constants import c + +from math import pi, sqrt +import numpy as np + + +def label_from_bmad_name(bmad_name: str) -> str: + """ + Formats a label by standardizing case, removing backslashes, and replacing disallowed characters. + + Parameters + ---------- + bmad_name : str + The Bmad name string to be formatted. + + Returns + ------- + str + A formatted label suitable for Genesis4. + """ + name = bmad_name.upper() + label = name.split("\\")[-1] # Extracts text after backslash, if any + return label.replace(".", "_").replace("#", "_") + + +def quadrupole_and_corrector_steps(quad: Quadrupole, cx=0, cy=0, num_steps=1): + """ + Converts a Genesis4 Quadrupole into a series of quadrupole and corrector steps. + + Parameters + ---------- + quad : Quadrupole + The Quadrupole element to be divided into steps. + cx : float, optional + The horizontal corrector strength. Defaults to 0. + cy : float, optional + The vertical corrector strength. Defaults to 0. + num_steps : int, optional + The number of steps to divide the quadrupole into. Defaults to 1. + + Returns + ------- + list + A list of Genesis4 elements, alternating between Corrector and Quadrupole steps. + """ + + L1 = quad.L / num_steps + label = quad.label + L_corrector = 1e-9 # BUG: zero length does nothing, but this works. + cx1 = cx / num_steps + cy1 = cy / num_steps + eles = [] + # Interleave kicks and quad steps + for step in range(num_steps): + # First step is half strength + if step == 0: + cx = cx1 / 2 + cy = cy1 / 2 + else: + cx = cx1 + cy = cy1 + + eles.append( + Corrector(cx=cx, cy=cy, label=f"{label}_kick{step+1}", L=L_corrector) + ) + eles.append( + Quadrupole( + L=L1, + k1=quad.k1, + x_offset=quad.x_offset, + y_offset=quad.y_offset, + label=f"{label}_step{step+1}", + ) + ) + + # Final half step + eles.append( + Corrector(cx=cx1 / 2, cy=cy1 / 2, label=f"{label}_kick{step+2}", L=L_corrector) + ) + return eles + + +def genesis4_eles_from_tao_ele(tao, ele_id): + """ + Creates Genesis4 elements from a specified element in a pytao.Tao instance. + + Parameters + ---------- + tao : Tao + The Tao instance containing the element definitions. + ele_id : int or str + The identifier (name or index) of the element in the Tao instance. + + Returns + ------- + list + A list of Genesis4 elements created from the specified Tao element. + + Raises + ------ + NotImplementedError + If the element type or offset is not supported. + ValueError + If the undulator length is inconsistent with the number of periods and period length. + """ + + info = tao.ele_head(ele_id) + info.update(tao.ele_gen_attribs(ele_id)) + info.update(tao.ele_methods(ele_id)) + key = info["key"].lower() + + # Genesis4 calls this 'label': + label = label_from_bmad_name(info["name"]) + + x_offset = info.get("X_OFFSET", 0) + y_offset = info.get("Y_OFFSET", 0) + + if key == "beginning_ele": + eles = None + elif key in ("drift", "pipe"): + ele = Drift(L=info["L"], label=label) + eles = [ele] + elif key in ("marker", "monitor") and (info["L"] == 0): + ele = Marker(label=label) + eles = [ele] + eles = None # TEMP + elif key == "quadrupole": + cx = info["HKICK"] + cy = info["VKICK"] + + ele = Quadrupole( + L=info["L"], + k1=info["K1"], + label=label, + x_offset=x_offset, + y_offset=y_offset, + ) + + # Do nothing if there are no correctors + if cx == 0 and cy == 0: + eles = [ele] + else: + eles = quadrupole_and_corrector_steps( + ele, cx=cx, cy=cy, num_steps=info["NUM_STEPS"] + ) + + elif key == "wiggler": + # Check for offsets + if x_offset != 0: + raise NotImplementedError(f"x_offset not zero: {x_offset}") + if y_offset != 0: + raise NotImplementedError(f"y_offset not zero: {y_offset}") + + # aw calc + B0 = B0 = info["B_MAX"] + lambdau = info["L_PERIOD"] + K = B0 * lambdau * c / (2 * pi * mec2) + + L = info["L"] + nwig = int(info["N_PERIOD"]) + lambdau = info["L_PERIOD"] + if not np.isclose(L, nwig * lambdau): + raise ValueError( + f"Inconsistent length for undulator {label}: {L} != {nwig}*{lambdau}" + ) + + if "helical" in info["field_calc"].lower(): + helical = True + aw = K + else: + assert info["field_calc"].lower() == "planar_model" + aw = K / sqrt(2) + helical = False + + ele = Undulator( + nwig=nwig, + lambdau=lambdau, + aw=aw, + helical=helical, + label=label, + ) + eles = [ele] + else: + raise NotImplementedError(key) + + return eles + + +def genesis4_elements_and_line_from_tao(tao, match="*"): + """ + Creates a Genesis4 lattice from a pytao.Tao instance. + + Parameters + ---------- + tao : Tao + The Tao instance to extract elements from. + match : str, optional + A pattern to match elements in the Tao lattice. Defaults to "*". + + Returns + ------- + dict + A dictionary mapping element labels to Genesis4 element objects. + list of str + A list of element labels in the order of the lattice line. + + Raises + ------ + ValueError + If elements with the same label have different properties. + """ + + elements = {} + line_labels = [] + for ix_ele in tao.lat_list(match, "ele.ix_ele"): + eles = genesis4_eles_from_tao_ele(tao, ix_ele) + if eles is not None: + for ele in eles: + label = ele.label + ele.label = "" + if label in elements: + ele0 = elements[label] + if ele0 != ele: + raise ValueError( + f"Elements have the same name but different properties: {ele0}, {ele}" + ) + else: + elements[label] = ele + line_labels.append(label) + + return elements, line_labels + + +def genesis4_namelists_from_tao( + tao, ele_start: str = "beginning", branch: int = 0, universe: int = 1 +): + """ + Creates Genesis4 namelists from a Tao instance, using specified parameters to + extract relevant beamline, beam, and field configurations. + + Parameters + ---------- + tao : pytao.Tao + A running Tao instance, providing access to element attributes, Twiss parameters, + and orbit data. + ele_start : str, optional + The starting element within the specified Tao universe and branch, using the + syntax `@` for universe and `>>` for branch (e.g., `'1@0>>element_name'`). + Defaults to `'beginning'`. + branch : int, optional + The branch index within the specified Tao universe. Defaults to 0. + universe : int, optional + The universe index within the Tao object. Defaults to 1. + + Returns + ------- + list + A list of Genesis4 namelists for the setup, field, beam, and tracking configurations. + + Notes + ----- + - This function collects attributes, Twiss parameters, and orbit data from the + specified Tao element to construct Genesis4-compatible input data. + - The generated `beamline` name is based on the Tao universe and branch configuration, + with `gamma0` calculated from the total energy. + - Additional parameters for field and beam properties are hard-coded but may be generalized. + + Examples + -------- + >>> tao = pytao.Tao(...) + >>> namelists = genesis4_namelists_from_tao(tao, ele_start='beginning', branch=0, universe=1) + + """ + + # Handle Tao universe @ branch >> sytax + if ">>" not in ele_start: + ele_start = f"{universe}@{branch}>>{ele_start}" + + attrs = tao.ele_gen_attribs(ele_start) + twiss = tao.ele_twiss(ele_start) + orbit = tao.ele_orbit(ele_start) + + beamline = tao.branch1(universe, branch)["name"] + gamma0 = attrs["E_TOT"] / mec2 + + setup = Setup( + rootname=beamline, + # lattice='Example1.lat', + beamline=beamline, + gamma0=gamma0, + # nbins=8, + # shotnoise=False, + ) + + # TODO: Generalize + field = Field(power=5000.0, waist_size=3e-05, dgrid=0.0002, ngrid=255) + + beam = Beam( + alphax=twiss["alpha_a"], + alphay=twiss["alpha_b"], + betax=twiss["beta_a"], + betay=twiss["beta_b"], + xcenter=orbit["x"], + ycenter=orbit["y"], + pxcenter=orbit["px"] * gamma0, + pycenter=orbit["py"] * gamma0, + gamma=(1 + orbit["pz"]) * gamma0, + ) + namelists = [setup, field, beam, Track()] + + return namelists diff --git a/genesis/version4/run.py b/genesis/version4/run.py index 0af2b7e5..8c62bcd4 100644 --- a/genesis/version4/run.py +++ b/genesis/version4/run.py @@ -627,6 +627,19 @@ def from_archive(cls, arch: Union[AnyPath, h5py.Group]) -> Genesis4: inst.load_archive(arch) return inst + @classmethod + def from_tao(cls, tao) -> Genesis4: + """ + Create a new Genesis4 object from a pytao.Tao instance + + Parameters + ---------- + + """ + input = MainInput.from_tao(tao) + lattice = Lattice.from_tao(tao) + return cls(input=input, lattice=lattice) + @override def load_output( self, diff --git a/mkdocs.yml b/mkdocs.yml index f32eaab5..1b178fdc 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -30,6 +30,7 @@ nav: - examples/genesis4/fodo_scan_model.ipynb - examples/genesis4/genesis4_example.ipynb - examples/genesis4/perlmutter_benchmark.ipynb + - examples/genesis4/bmad-genesis4.ipynb - API: - api/genesis4.md - api/genesis4-main-input.md