Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add Bmad interface #51

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Add Bmad interface #51

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
883d67c
Working from_tao methods
ChristopherMayes Nov 2, 2024
edd9a89
Better docstrings with help from GPT-4o
ChristopherMayes Nov 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .pre-commit-config.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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$"
257 changes: 257 additions & 0 deletions docs/examples/genesis4/bmad_genesis4.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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<x\\right>$ (µ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<y\\right>$ (µ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
}
31 changes: 31 additions & 0 deletions docs/examples/genesis4/data/example1-steadystate/Example1.bmad
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
1 change: 1 addition & 0 deletions environment.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ dependencies:
- pydantic >2
- pydantic-settings
# Developer
- pytao
- ffmpeg
- pytest
- pytest-cov
Expand Down
6 changes: 6 additions & 0 deletions genesis/version4/input/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -377,6 +378,11 @@ def to_file(
filename,
)

@classmethod
def from_tao(cls, tao):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This needs a docstring

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:
Expand Down
10 changes: 10 additions & 0 deletions genesis/version4/input/lattice.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -1004,6 +1006,14 @@ def from_elements(
filename=pathlib.Path(filename) if filename else None,
)

@classmethod
def from_tao(cls, tao):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This needs a docstring

elements, line_labels = genesis4_elements_and_line_from_tao(tao)

line_label = tao.branch1(1, 0)["name"]
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ix_uni and ix_branch should be arguments to from_tao

pytao.Tao.branch1(
    self,
    ix_uni,
    ix_branch,
    ...
 )

elements[line_label] = Line(elements=line_labels)
return cls(elements=elements)

def to_file(
self,
filename: AnyPath,
Expand Down
Loading
Loading