Skip to content

Commit ecd438b

Browse files
ChristopherMayesChristopherMayes
committed
Clean up pseudomode code
1 parent 27b9a81 commit ecd438b

File tree

2 files changed

+199
-99
lines changed

2 files changed

+199
-99
lines changed

‎docs/examples/wakefields/resistive_wall_wakefield.ipynb‎

Lines changed: 108 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,15 @@
77
"source": [
88
"# Resistive Wall Wakefield\n",
99
"\n",
10-
"The `ResistiveWallWakefield` object implements the pseudomode (damped oscillator) fits from SLAC-PUB-10707 to efficiently calculate the resistive wall wakefield from round and flat geometries, for arbitrary conductivities and relaxation times."
10+
"The `ResistiveWallWakefield` class implements the pseudomode (damped oscillator) fits from SLAC-PUB-10707 to efficiently calculate the short-range resistive wall wakefield for round and flat geometries.\n",
11+
"\n",
12+
"The wakefield is represented as a single damped sinusoid:\n",
13+
"$$W(z) = A \\, e^{d z} \\sin(k_r z + \\phi)$$\n",
14+
"where $A$, $d$, $k_r$, and $\\phi$ are derived from polynomial fits of digitized data from Figures 4, 8, and 14 of SLAC-PUB-10707.\n",
15+
"\n",
16+
"**References:**\n",
17+
"- Bane & Stupakov, [SLAC-PUB-10707](https://www.slac.stanford.edu/cgi-wrap/getdoc/slac-pub-10707.pdf) (2004)\n",
18+
"- Bane, Stupakov, Tu, [EPAC 2006 THPCH073](https://accelconf.web.cern.ch/e06/PAPERS/THPCH073.PDF)"
1119
]
1220
},
1321
{
@@ -17,125 +25,161 @@
1725
"metadata": {},
1826
"outputs": [],
1927
"source": [
20-
"from pmd_beamphysics.wakefields.resistive_wall import ResistiveWallWakefield\n",
21-
"\n",
28+
"from pmd_beamphysics.wakefields import ResistiveWallWakefield\n",
2229
"from pmd_beamphysics import ParticleGroup\n",
2330
"from pmd_beamphysics.units import epsilon_0\n",
31+
"\n",
2432
"import numpy as np\n",
2533
"import matplotlib.pyplot as plt"
2634
]
2735
},
36+
{
37+
"cell_type": "markdown",
38+
"id": "c36da92f-d4a6-46c5-a0ff-307d5d69bde7",
39+
"metadata": {},
40+
"source": [
41+
"## Basic Usage\n",
42+
"\n",
43+
"Create a wakefield by specifying the pipe geometry and material properties:"
44+
]
45+
},
2846
{
2947
"cell_type": "code",
3048
"execution_count": null,
31-
"id": "393a7b50-478f-47a0-acf8-268a34b2af45",
49+
"id": "d045a156-5b8d-439c-b3a0-0434f01657ba",
3250
"metadata": {},
3351
"outputs": [],
3452
"source": [
35-
"?ResistiveWallWakefield"
53+
"# Create a wakefield for a 2.5 mm radius copper pipe\n",
54+
"wake = ResistiveWallWakefield(\n",
55+
" radius=0.0025,\n",
56+
" conductivity=6.5e7, # σ [S/m]\n",
57+
" relaxation_time=27e-15, # τ [s]\n",
58+
" geometry=\"round\",\n",
59+
")\n",
60+
"wake"
3661
]
3762
},
3863
{
3964
"cell_type": "markdown",
40-
"id": "c36da92f-d4a6-46c5-a0ff-307d5d69bde7",
65+
"id": "aac561ff",
4166
"metadata": {},
4267
"source": [
43-
"# Basic ins"
68+
"The built-in `plot()` method shows the wakefield as a function of trailing distance:"
4469
]
4570
},
4671
{
4772
"cell_type": "code",
4873
"execution_count": null,
49-
"id": "d045a156-5b8d-439c-b3a0-0434f01657ba",
74+
"id": "0dfb3db6",
5075
"metadata": {},
5176
"outputs": [],
5277
"source": [
53-
"wake = ResistiveWallWakefield(\n",
54-
" radius=0.0025, geometry=\"round\", conductivity=6.5e7, relaxation_time=27e-15\n",
55-
")\n",
56-
"\n",
5778
"wake.plot()"
5879
]
5980
},
81+
{
82+
"cell_type": "markdown",
83+
"id": "13301e26",
84+
"metadata": {},
85+
"source": [
86+
"### Material Presets\n",
87+
"\n",
88+
"Common materials are available via `from_material()`:"
89+
]
90+
},
6091
{
6192
"cell_type": "code",
6293
"execution_count": null,
63-
"id": "47eb70ea-ff85-4b1d-8de0-f9b9f44b5f9b",
94+
"id": "4903016f",
6495
"metadata": {},
6596
"outputs": [],
6697
"source": [
67-
"print(wake.to_bmad())"
98+
"# Available material presets\n",
99+
"list(ResistiveWallWakefield.MATERIALS)"
100+
]
101+
},
102+
{
103+
"cell_type": "code",
104+
"execution_count": null,
105+
"id": "59ec3e57",
106+
"metadata": {},
107+
"outputs": [],
108+
"source": [
109+
"wake_cu = ResistiveWallWakefield.from_material(\"copper-slac-pub-10707\", radius=2.5e-3)\n",
110+
"wake_cu"
68111
]
69112
},
70113
{
71114
"cell_type": "markdown",
72-
"id": "ad1249db-19e5-42a2-a976-2ff48843a91d",
115+
"id": "c0acf32d",
73116
"metadata": {},
74117
"source": [
75-
"# Compare with SLAC-PUB-10707\n",
118+
"### Bmad Export\n",
76119
"\n",
77-
"Here we check the function against the plots in SLAC-PUB-10707."
120+
"The wakefield can be exported in Bmad format:"
78121
]
79122
},
80123
{
81124
"cell_type": "code",
82125
"execution_count": null,
83-
"id": "56a82d9d-9079-45d5-8bcd-6c8f5f2bdd32",
126+
"id": "47eb70ea-ff85-4b1d-8de0-f9b9f44b5f9b",
84127
"metadata": {},
85128
"outputs": [],
86129
"source": [
87-
"radius_ref = 2.5e-3\n",
88-
"\n",
89-
"wake = ResistiveWallWakefield.from_material(\n",
90-
" material=\"copper-slac-pub-10707\", radius=radius_ref, geometry=\"round\"\n",
91-
")\n",
92-
"\n",
93-
"wake.plot()"
130+
"print(wake_cu.to_bmad())"
94131
]
95132
},
96133
{
97-
"cell_type": "code",
98-
"execution_count": null,
99-
"id": "19cff809-238b-4cbc-90d4-1cbac8032022",
134+
"cell_type": "markdown",
135+
"id": "ad1249db-19e5-42a2-a976-2ff48843a91d",
100136
"metadata": {},
101-
"outputs": [],
102137
"source": [
103-
"raw_data = np.loadtxt(\"../data/SLAC-PUB-10707-digitized-Fig4-AC-Cu.csv\", delimiter=\",\")\n",
138+
"## Validation Against SLAC-PUB-10707\n",
104139
"\n",
105-
"# Convert to SI\n",
106-
"zref, Wref = (\n",
107-
" raw_data[:, 0] * 1e-6,\n",
108-
" raw_data[:, 1] * 4 / radius_ref**2 / (4 * np.pi * epsilon_0),\n",
109-
")\n",
110-
"plt.plot(zref, Wref)"
140+
"Compare the pseudomode fit against digitized data from Figure 4 (AC-Cu, round pipe):"
111141
]
112142
},
113143
{
114144
"cell_type": "code",
115145
"execution_count": null,
116-
"id": "eabb80d7-969c-45e6-bd56-60cd05b8d494",
146+
"id": "56a82d9d-9079-45d5-8bcd-6c8f5f2bdd32",
117147
"metadata": {},
118148
"outputs": [],
119149
"source": [
120-
"zlist = np.linspace(0, 300e-6, 200)\n",
121-
"Wz = wake.pseudomode(-zlist)\n",
150+
"# Load digitized reference data\n",
151+
"raw_data = np.loadtxt(\"../data/SLAC-PUB-10707-digitized-Fig4-AC-Cu.csv\", delimiter=\",\")\n",
152+
"radius_ref = 2.5e-3\n",
122153
"\n",
123-
"fig, ax = plt.subplots()\n",
124-
"ax.plot(zlist * 1e6, Wz * 1e-12, label=f\"ResistiveWallWakefield {wake.geometry}\")\n",
154+
"# Convert CGS units to SI\n",
155+
"zref = raw_data[:, 0] * 1e-6 # µm → m\n",
156+
"Wref = raw_data[:, 1] * 4 / radius_ref**2 / (4 * np.pi * epsilon_0) # V/pC/m\n",
125157
"\n",
126-
"ax.plot(zref * 1e6, Wref * 1e-12, \"--\", label=\"Fig. 4 AC-Cu from SLAC-PUB-10707 2004\")\n",
127-
"plt.legend()\n",
158+
"# Create wakefield with same parameters\n",
159+
"wake_ref = ResistiveWallWakefield.from_material(\n",
160+
" \"copper-slac-pub-10707\", radius=radius_ref\n",
161+
")\n",
162+
"zlist = np.linspace(0, 300e-6, 200)\n",
163+
"Wz = wake_ref(-zlist) # Evaluate at trailing positions\n",
128164
"\n",
165+
"# Plot comparison\n",
166+
"fig, ax = plt.subplots()\n",
167+
"ax.plot(zlist * 1e6, Wz * 1e-12, label=f\"ResistiveWallWakefield ({wake_ref.geometry})\")\n",
168+
"ax.plot(zref * 1e6, Wref * 1e-12, \"--\", label=\"Fig. 4 AC-Cu (digitized)\")\n",
169+
"ax.legend()\n",
129170
"ax.set_xlabel(r\"$-z$ (µm)\")\n",
130-
"ax.set_ylabel(r\"$W_z$ (V/pC/m)\")"
171+
"ax.set_ylabel(r\"$W_z$ (V/pC/m)\")\n",
172+
"ax.set_title(\"Validation: Round Copper Pipe\")"
131173
]
132174
},
133175
{
134176
"cell_type": "markdown",
135177
"id": "5d112e6f-839f-4d62-bc24-3130c91b2e9d",
136178
"metadata": {},
137179
"source": [
138-
"# Integrated wake from particles\n"
180+
"## Particle Wakefield Kicks\n",
181+
"\n",
182+
"Compute the integrated wakefield kick for each particle in a bunch:"
139183
]
140184
},
141185
{
@@ -145,27 +189,18 @@
145189
"metadata": {},
146190
"outputs": [],
147191
"source": [
192+
"# Load a particle beam\n",
148193
"P = ParticleGroup(\"../data/bmad_particles2.h5\")\n",
149-
"P.drift_to_t()"
150-
]
151-
},
152-
{
153-
"cell_type": "code",
154-
"execution_count": null,
155-
"id": "e2478b38-e658-40fb-95bb-a1574e784d50",
156-
"metadata": {},
157-
"outputs": [],
158-
"source": [
159-
"z = P.z\n",
160-
"weight = P.weight"
194+
"P.drift_to_t() # Align particles at constant time\n",
195+
"P"
161196
]
162197
},
163198
{
164199
"cell_type": "markdown",
165200
"id": "5e8c504c-d5c2-4cb4-b747-eb64dc3e1bcc",
166201
"metadata": {},
167202
"source": [
168-
"`wake.pseudomode` gives the per particle kicks"
203+
"The `particle_kicks` method computes the wakefield-induced energy change for each particle:"
169204
]
170205
},
171206
{
@@ -175,18 +210,9 @@
175210
"metadata": {},
176211
"outputs": [],
177212
"source": [
178-
"zwake = wake.pseudomode.particle_kicks(z, weight)\n",
179-
"len(zwake)"
180-
]
181-
},
182-
{
183-
"cell_type": "code",
184-
"execution_count": null,
185-
"id": "34bf0eb3-7543-4da1-90ee-a426e26e5740",
186-
"metadata": {},
187-
"outputs": [],
188-
"source": [
189-
"plt.scatter(z, zwake)"
213+
"# Compute wakefield kicks (eV/m) for each particle\n",
214+
"kicks = wake_ref.particle_kicks(P)\n",
215+
"len(kicks)"
190216
]
191217
},
192218
{
@@ -196,22 +222,21 @@
196222
"metadata": {},
197223
"outputs": [],
198224
"source": [
199-
"zscale = 1e6\n",
200225
"fig, ax = plt.subplots()\n",
201-
"ax.scatter(z * zscale, zwake, label=\"Python\", color=\"black\")\n",
202-
"ax.legend()\n",
226+
"ax.scatter(P.z * 1e6, kicks, s=1, color=\"black\")\n",
203227
"ax.set_xlabel(r\"$z$ (µm)\")\n",
204-
"ax.set_ylabel(r\"$W_z$ (eV/m)\")"
228+
"ax.set_ylabel(r\"Wakefield kick (eV/m)\")\n",
229+
"ax.set_title(\"Per-particle wakefield kicks\")"
205230
]
206231
},
207232
{
208233
"cell_type": "markdown",
209234
"id": "5d6d054a-3635-4072-82f2-203fd706c171",
210235
"metadata": {},
211236
"source": [
212-
"# ParticleGroup.wakefield_plot\n",
237+
"## Convenience: `ParticleGroup.wakefield_plot`\n",
213238
"\n",
214-
"Wakefields can be autmoatically integrated and plotted from a ParticlGroup. The function automatically determintes the horizontal axis."
239+
"The wakefield can be computed and plotted directly from a `ParticleGroup`. The plot automatically chooses the appropriate horizontal axis based on whether particles are aligned at constant $t$ or constant $z$:"
215240
]
216241
},
217242
{
@@ -221,7 +246,9 @@
221246
"metadata": {},
222247
"outputs": [],
223248
"source": [
224-
"P.wakefield_plot(wake, figsize=(12, 4))"
249+
"# Particles at constant t → use z as horizontal axis\n",
250+
"P.drift_to_t()\n",
251+
"P.wakefield_plot(wake_ref, figsize=(12, 4))"
225252
]
226253
},
227254
{
@@ -231,14 +258,15 @@
231258
"metadata": {},
232259
"outputs": [],
233260
"source": [
261+
"# Particles at constant z → use t as horizontal axis\n",
234262
"P.drift_to_z()\n",
235-
"P.wakefield_plot(wake, figsize=(12, 4))"
263+
"P.wakefield_plot(wake_ref, figsize=(12, 4))"
236264
]
237265
}
238266
],
239267
"metadata": {
240268
"kernelspec": {
241-
"display_name": "Python 3 (ipykernel)",
269+
"display_name": "pmd_beamphysics-dev",
242270
"language": "python",
243271
"name": "python3"
244272
},
@@ -252,7 +280,7 @@
252280
"name": "python",
253281
"nbconvert_exporter": "python",
254282
"pygments_lexer": "ipython3",
255-
"version": "3.13.3"
283+
"version": "3.13.9"
256284
}
257285
},
258286
"nbformat": 4,

0 commit comments

Comments
 (0)