Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
241 changes: 241 additions & 0 deletions examples/ad-examples.ipynb
Original file line number Diff line number Diff line change
@@ -0,0 +1,241 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"id": "1ff2c3b4-e83b-4348-8c02-8504740c02b2",
"metadata": {},
"outputs": [],
"source": [
"using SciBmad\n",
"using DifferentiationInterface\n",
"import DifferentiationInterface as DI\n",
"# Different differentiation packages:\n",
"using ForwardDiff, GTPSA, ReverseDiff, FiniteDiff, FiniteDifferences"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "2f720d76-f35b-4575-9967-97de69c67152",
"metadata": {},
"outputs": [],
"source": [
"K1 = 0.36\n",
"K2 = 0.1\n",
"\n",
"qf = Quadrupole(L=0.5, Kn1=K1)\n",
"sf = Sextupole(L=0.2, Kn2=K2)\n",
"d = Drift(L=0.6)\n",
"b = SBend(L=6.0, angle=π/132)\n",
"qd = Quadrupole(L=0.5, Kn1=-K1)\n",
"sd = Sextupole(L=0.2, Kn2=-K2)\n",
"fodo = Beamline([qf, d, sf, b, sd, d, qd], species_ref=Species(\"electron\"), E_ref=18e9);"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "3de3d2f8-abe3-4559-9ce8-1741afb8b9ae",
"metadata": {},
"outputs": [],
"source": [
"# Returns exit coordinates for 1 particle\n",
"function track_a_particle(v0, beamline)\n",
" v0_ = similar(v0)\n",
" v0_ .= v0\n",
" b0 = Bunch(v0_; species=beamline.species_ref, R_ref=beamline.R_ref)\n",
" track!(b0, beamline)\n",
" return b0.coords.v\n",
"end\n",
"x0 = [1e-2, 0.0, 1e-2, 0.0, 0.0, 0.0];"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "b481def7-4306-47bd-ae1b-802fcb05a41b",
"metadata": {
"scrolled": true
},
"outputs": [
{
"data": {
"text/plain": [
"6×6 Matrix{Float64}:\n",
" -0.501301 8.58617 0.000765453 -0.00240033 0.0 0.122029\n",
" -0.257394 2.41377 -0.000629133 -0.00359376 0.0 0.0439702\n",
" 0.00154156 -0.00171283 2.41678 8.59153 0.0 -0.013781\n",
" 4.81099e-5 -0.00281752 -0.256971 -0.499744 0.0 0.00253851\n",
" -0.00937178 -0.0829428 -0.0027041 -0.0152556 1.0 -0.000483029\n",
" 0.0 0.0 0.0 0.0 0.0 1.0"
]
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# Forward AD using ForwardDiff.jl:\n",
"prep_FD = DI.prepare_jacobian(track_a_particle, AutoForwardDiff(), x0, Constant(fodo))\n",
"J_FD = DI.jacobian(track_a_particle, prep_FD, AutoForwardDiff(), x0, Constant(fodo))"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "fa4ffadf-e727-4951-9163-75a760a0b4ff",
"metadata": {
"scrolled": true
},
"outputs": [
{
"data": {
"text/plain": [
"6×6 Matrix{Float64}:\n",
" -0.501301 8.58617 0.000765453 -0.00240033 0.0 0.122029\n",
" -0.257394 2.41377 -0.000629133 -0.00359376 0.0 0.0439702\n",
" 0.00154156 -0.00171283 2.41678 8.59153 0.0 -0.013781\n",
" 4.81099e-5 -0.00281752 -0.256971 -0.499744 0.0 0.00253851\n",
" -0.00937178 -0.0829428 -0.0027041 -0.0152556 1.0 -0.000483029\n",
" 0.0 0.0 0.0 0.0 0.0 1.0"
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# Taylor AD using GTPSA.jl:\n",
"prep_GTPSA = DI.prepare_jacobian(track_a_particle, AutoGTPSA(), x0, Constant(fodo))\n",
"J_GTPSA = DI.jacobian(track_a_particle, prep_GTPSA, AutoGTPSA(), x0, Constant(fodo))"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "351aa79a-7bbd-423c-8613-a0a9ceb0b0bd",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"6×6 Matrix{Float64}:\n",
" -0.501301 8.58617 0.000765453 -0.00240033 0.0 0.122029\n",
" -0.257394 2.41377 -0.000629133 -0.00359376 0.0 0.0439702\n",
" 0.00154156 -0.00171283 2.41678 8.59153 0.0 -0.013781\n",
" 4.81099e-5 -0.00281752 -0.256971 -0.499744 0.0 0.00253851\n",
" -0.00937178 -0.0829428 -0.0027041 -0.0152556 1.0 -0.000483029\n",
" 0.0 0.0 0.0 0.0 0.0 1.0"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# Reverse AD using ReverseDiff.jl:\n",
"prep_RD = DI.prepare_jacobian(track_a_particle, AutoReverseDiff(;compile=true), x0, Constant(fodo))\n",
"J_RD = DI.jacobian(track_a_particle, prep_RD, AutoReverseDiff(;compile=true), x0, Constant(fodo))"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "d835d6a6-7ebb-4971-850e-35995f9df4c4",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"6×6 Matrix{Float64}:\n",
" -0.501301 8.58617 0.000765455 -0.00240034 0.0 0.122029\n",
" -0.257394 2.41377 -0.000629133 -0.00359377 0.0 0.0439702\n",
" 0.00154156 -0.00171282 2.41678 8.59153 0.0 -0.013781\n",
" 4.81099e-5 -0.00281752 -0.256971 -0.499744 0.0 0.00253851\n",
" -0.00937178 -0.0829429 -0.00270414 -0.0152558 1.0 -0.000483002\n",
" 0.0 0.0 0.0 0.0 0.0 1.0"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# Finite differencing using FiniteDiff.jl:\n",
"prep_fdif1 = DI.prepare_jacobian(track_a_particle, AutoFiniteDiff(), x0, Constant(fodo))\n",
"J_fdif1 = DI.jacobian(track_a_particle, prep_fdif1, AutoFiniteDiff(), x0, Constant(fodo))"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "c88033cb-f9ab-4a71-a3af-f2eb9e162c18",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"6×6 Matrix{Float64}:\n",
" -0.501301 8.58617 0.000765453 -0.00240033 0.0 0.122029\n",
" -0.257394 2.41377 -0.000629133 -0.00359376 0.0 0.0439702\n",
" 0.00154156 -0.00171283 2.41678 8.59153 0.0 -0.013781\n",
" 4.81099e-5 -0.00281752 -0.256971 -0.499744 0.0 0.00253851\n",
" -0.00937178 -0.0829428 -0.0027041 -0.0152556 1.0 -0.000483029\n",
" 0.0 0.0 0.0 0.0 0.0 1.0"
]
},
"execution_count": 8,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# Finite differencing using FiniteDifferences.jl:\n",
"prep_fdif2 = DI.prepare_jacobian(track_a_particle, AutoFiniteDifferences(;fdm=central_fdm(3,1)), x0, Constant(fodo))\n",
"J_fdif2 = DI.jacobian(track_a_particle, prep_fdif2, AutoFiniteDifferences(;fdm=central_fdm(3,1)), x0, Constant(fodo))"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "3cd2cc3d-8016-4aa6-bea5-77591bd20397",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"true"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# Show they're all equal:\n",
"J_FD ≈ J_GTPSA ≈ J_RD ≈ J_fdif1 ≈ J_fdif2"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Julia Global Project (8 threads) 1.10.10",
"language": "julia",
"name": "julia-global-project-_8-threads_-1.10"
},
"language_info": {
"file_extension": ".jl",
"mimetype": "application/julia",
"name": "julia",
"version": "1.10.10"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
4 changes: 1 addition & 3 deletions lattices/toy_ags.jl
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,6 @@ Lb = 2 * asin(2.66666666666666652E+000 * 1.96327112309048618E-002 / 2)/1.9632711
csnk = Marker()
wsnk = Marker()
end1 = Marker()
rfbc = RFCavity(L = 1.00000000000000000E+000, harmon = 1,
voltage = 3.20000000000000000E+5)
end


Expand All @@ -37,6 +35,6 @@ ring = Beamline([q1h, dr, b, dr, q2, dr, b, dr, q1h, q1h, dr, b, dr, q2, dr, b,
dr, b, dr, q1h, q1h, dr, b, dr, q2, dr, b, dr, q1h, q1h, dr, b, dr, q2, dr, b, dr, q1h, q1h, dr, b, dr,
q2, dr, b, dr, q1h, q1h, dr, b, dr, q2, dr, b, dr, q1h, q1h, dr, b, dr, q2, dr, b, dr, q1h, q1h, dr, b,
dr, q2, dr, b, dr, q1h, end1],
E_ref =2.40490243299096680E+010, # 3.5587247443204529E+10 + 2*1.9018517828440544e10*Time(),
E_ref = 3.5587247443204529E+10 + 2*1.9018517828440544e10*Time(), #2.40490243299096680E+010
species_ref = Species("proton"))

3 changes: 2 additions & 1 deletion src/SciBmad.jl
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ end

const BTBL = Base.get_extension(BeamTracking, :BeamTrackingBeamlinesExt)

export twiss, find_closed_orbit, track!, track
export twiss, find_closed_orbit, track!, track, Time

function fast_coast_check(bl; use_KA=false, use_explicit_SIMD=false)
# Just check if dpz/dz == 0
Expand Down Expand Up @@ -288,6 +288,7 @@ include("newton.jl")


@setup_workload begin

@compile_workload begin
# We want to compile drift-kick-drift, matrix-kick-matrix
# and solenoid kick for different numbers of multipoles
Expand Down
2 changes: 1 addition & 1 deletion src/newton.jl
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ Finds roots of f!(y, x) using Newton's method.

# Keyword arguments
- `abstol`: Convergence absolute tolerance (default: 1e-13)
- `rektol`: Convergence relative tolerance (default: 1e-13)
- `reltol`: Convergence relative tolerance (default: 1e-13)
- `max_iter`: Maximum number of iterations (default: 100)

Returns `NamedTuple` containing newton search results.
Expand Down
Loading