Example-10: Dipole (element)

# Comparison of dipole element with MADX-PTC and other features

# Note, cylindrical multipole is included upto octupole order
# Potential is not truncated in paraxial case, only sqrt is expanded
# In exact case effects of multipoles are not accounted in wedges (cases with e1 or e2 not equal to zero)

from pathlib import Path
from os import system

import torch
from model.library.drift import Drift
from model.library.quadrupole import Quadrupole
from model.library.sextupole import Sextupole
from model.library.dipole import Dipole

# Tracking (paraxial)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = False
align = False

length = 2.0
angle = 0.05
e1 = 0.0
e2 = 0.0
kn = 0.0
ks = 0.0
ms = 0.0
mo = 0.0
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()

[5.160257777686003e-05, -0.004956273150572646, -0.0030019980019983517, 0.001]
[5.1602577776445975e-05, -0.004956273150572915, -0.0030019980019981656, 0.0010000000000000434]
[4.1405680967221414e-16, 2.688821387764051e-16, -1.8604909279851256e-16, -4.336808689942018e-17]

# Tracking (paraxial, e1 & e2)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = False
align = False

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
kn = 0.0
ks = 0.0
ms = 0.0
mo = 0.0
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()

[6.408749411377232e-05, -0.00494998958983017, -0.002995752944647053, 0.0010049983869644134]
[6.408749411335537e-05, -0.004949989589830473, -0.0029957529446465857, 0.0010049983869644566]
[4.169434979564568e-16, 3.0270924655795284e-16, -4.670742959067553e-16, -4.3151246464923076e-17]

# Tracking (paraxial, e1 & e2, alignment)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = False
align = True

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
kn = 0.0
ks = 0.0
ms = 0.0
mo = 0.0
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()

[0.0027250502286323983, -0.004697460778410744, -0.007781317874336121, -0.004015687563474919]
[0.0027250502286311207, -0.004697460778410097, -0.007781317874338802, -0.004015687563476177]
[1.2776238400569184e-15, -6.47051856539349e-16, 2.6810151321221554e-15, 1.2576745200831851e-15]

# Tracking (paraxial, kn & ks)

# Note, in paraxial case, MADX seems to truncate cylindrical potential

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = False
align = False

length = 2.0
angle = 0.05
e1 = 0.0
e2 = 0.0
kn = 0.1
ks = 0.1
ms = 0.0
mo = 0.0
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()

[-0.002043638956596164, -0.006566087436123452, -0.00259447159458907, 0.0010855376916213359]
[-0.0020440852924812135, -0.006566346702412346, -0.002594440402484201, 0.001085525276668085]
[4.46335885049675e-07, 2.592662888935629e-07, -3.119210486906762e-08, 1.2414953250768773e-08]

# Tracking (paraxial, kn & ks, alignment)

# Note, in paraxial case, MADX seems to truncate cylindrical potential

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = False
align = True

length = 2.0
angle = 0.05
e1 = 0.0
e2 = 0.0
kn = 0.1
ks = 0.1
ms = 0.0
mo = 0.0
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()

[0.013167193439103305, 0.005511119368028012, -0.010781508233295324, -0.007329921037793256]
[0.013155910631910853, 0.0055012372192967465, -0.010779780613239438, -0.00732819304969001]
[1.128280719245138e-05, 9.882148731265271e-06, -1.727620055886822e-06, -1.7279881032459046e-06]

# Tracking (paraxial, kn & ks, ms & mo, alignment)

# Note, in paraxial case, MADX seems to truncate cylindrical potential

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = False
align = True

length = 2.0
angle = 0.05
e1 = 0.0
e2 = 0.0
kn = 0.1
ks = 0.1
ms = 0.1
mo = 0.1
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()

[0.013072849472264409, 0.005416718324594963, -0.010908836090706359, -0.007452351251381879]
[0.013061618858281962, 0.0054068762860350235, -0.010906993838412838, -0.007450515116390934]
[1.1230613982447443e-05, 9.842038559939453e-06, -1.842252293521307e-06, -1.8361349909453567e-06]

# Tracking (paraxial, e1 & e2, kn & ks, ms & mo, alignment)

# Note, in paraxial case, MADX seems to truncate cylindrical potential
# Note, model dipole doesn't account for multipoles in wedges

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = False
align = True

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
kn = 0.1
ks = 0.1
ms = 0.1
mo = 0.1
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()

[0.013030972087034147, 0.005379072261796265, -0.010940636831296861, -0.007479587493733525]
[0.013019736470996354, 0.005369216085022485, -0.0109387887264908, -0.007477745021700929]
[1.1235616037793758e-05, 9.85617677377957e-06, -1.8481048060618732e-06, -1.8424720325954658e-06]

# Tracking (exact, alignment)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = True
align = True

length = 2.0
angle = 0.05
e1 = 0.0
e2 = 0.0
kn = 0.0
ks = 0.0
ms = 0.0
mo = 0.0
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()

[0.002770798305588665, -0.004651371270674763, -0.007745680488823617, -0.0039921009984161026]
[0.0027707983054956944, -0.004651371270673946, -0.007745680488833436, -0.003992100998417344]
[9.2970769971501e-14, -8.161873954470877e-16, 9.819402235766717e-15, 1.2411946470614055e-15]

# Tracking (exact, kn & ks, ms & mo, alignment)

# Note, in model dipole cylindrical potential is truncated at octupole

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = True
align = True

length = 2.0
angle = 0.05
e1 = 0.0
e2 = 0.0
kn = 0.1
ks = 0.1
ms = 0.1
mo = 0.1
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()

[0.013055723507830729, 0.005405610520344896, -0.01089583005089035, -0.007447552679263773]
[0.013055723646617522, 0.005405610850033292, -0.010895828250285611, -0.007447551055280752]
[-1.387867928220876e-10, -3.2968839637492753e-10, -1.8006047382279622e-09, -1.6239830209763273e-09]

# Tracking (exact, e1 & e2, kn & ks, ms & mo, alignment)

# Note, in model dipole cylindrical potential is truncated at octupole
# Note, model dipole doesn't account for multipoles in wedges

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = True
align = True

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
kn = 0.1
ks = 0.1
ms = 0.1
mo = 0.1
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()

[0.013018589129527029, 0.005372361621197678, -0.010919187705983761, -0.007468383469439608]
[0.013013881710648583, 0.0053679558070361885, -0.010927596957624465, -0.007474796553899775]
[4.707418878445446e-06, 4.4058141614898225e-06, 8.409251640703955e-06, 6.413084460166543e-06]

# Deviation/error variables

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
kn = 0.1
ks = 0.1
ms = 0.1
mo = 0.1
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp)

# Each element has two variant of a call method
# In the first case only state is passed, it is transformed using parameters specified on initializaton

print(D(state))
print()

# Deviation errors can be also passed to call method
# These variables are added to corresponding parameters specified on initializaton
# For example, element lenght can changed

print(D(state, data={**D.data(), **{'dl': -0.5*D.length}}))
print()

# In the above D.data() creates default deviation dictionary (with zero values for each deviaton)
# {**D.data(), **{'dl': -D.length}} replaces the 'dl' key value

# Additionaly, alignment errors are passed as deivation variables
# They are used if alignment flag is raised

print(D(state, data={**D.data(), **error}, alignment=True))
print()

# The following elements can be made equivalent using deviation variables

DA = Dipole('DA', length, angle, e1, e2, kn, ks, ms, mo, dp)
DB = Dipole('DB', length - 0.1, angle, e1, e2, kn, ks, ms, mo, dp)

print(DA(state) - DB(state, data={**DB.data(), **{'dl': torch.tensor(+0.1, dtype=DB.dtype)}}))

# Note, while in some cases float values can be passed as values to deviation variables
# The correct behaviour in guaranteed only for tensors

tensor([-0.0020, -0.0066, -0.0026,  0.0011], dtype=torch.float64)

tensor([ 0.0044, -0.0061, -0.0038,  0.0013], dtype=torch.float64)

tensor([ 0.0130,  0.0054, -0.0109, -0.0075], dtype=torch.float64)

tensor([0., 0., 0., 0.], dtype=torch.float64)

# Insertion element

import torch
from model.library.dipole import Dipole

# In this mode elements are treated as thin insertions (at the center)
# Using parameters specified on initialization, transport two matrices are computed
# These matrices are used to insert the element
# Input state is transformed from the element center to its entrance
# Next, transformation from the entrance frame to the exit frame is performed
# This transformation can contain errors
# The final step is to transform state from the exit frame back to the element center
# Without errors, this results in identity transformation for linear elements

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
kn = 0.0
ks = 0.0
ms = 0.0
mo = 0.0
dp = 0.0
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp, exact=False, insertion=True)

# For dipole insertion is an identity transformation in the following is true
# dp = 0 (chomatic orbit change)
# kn & ks and ms & mo are all equal to zero (nonlinearity of cylindrical potential)

print(D(state) - state)

# Represents effect of an error (any nonzero value of strengh or a change in other parameter)

print(D(state, data={**D.data(), **{'dp': 0.001}}) - state)

# Exact tracking corresponds to inclusion of kinematic term as errors

D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp, exact=True, insertion=True)

print(D(state) - state)

tensor([ 1.0408e-17, -4.3368e-18, -4.3368e-18,  8.6736e-19],
       dtype=torch.float64)
tensor([ 9.9380e-06,  4.9995e-05, -1.9980e-06,  6.5052e-19],
       dtype=torch.float64)
tensor([-2.3859e-06, -5.8600e-06, -7.2738e-07, -2.4947e-07],
       dtype=torch.float64)

# Mapping over a set of initial conditions

# Call method can be used to map over a set of initial conditions
# Note, device can be set to cpu or gpu via base element classvariables

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
kn = 0.5
ks = 0.5
ms = 1.0
mo = 1.0
dp = 0.001

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp)

state = 1.0E-3*torch.randn((512, 4), dtype=D.dtype, device=D.device)

print(torch.vmap(D)(state).shape)

# To map over deviations parameters a wrapper function (or a lambda expression) can be used

def wrapper(state, dp):
    return D(state, data={**D.data(), **{'dp': dp}})

dp = 1.0E-3*torch.randn(512, dtype=D.dtype, device=D.device)

print(torch.vmap(wrapper)(state, dp).shape)

torch.Size([512, 4])
torch.Size([512, 4])

# Differentiability

# Both call methods are differentiable
# Derivative with respect to state can be computed directly
# For deviation variables, wrapping is required

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
kn = 0.5
ks = 0.5
ms = 1.0
mo = 1.0
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp)

# Compute derivative with respect to state

print(torch.func.jacrev(D)(state))
print()

# Compute derivative with respect to a deviation variable

e1 = torch.tensor(0.0, dtype=torch.float64)

def wrapper(state, e1):
    return D(state, data={**D.data(), **{'e1': e1}})

print(torch.func.jacrev(wrapper, 1)(state, e1))
print()

tensor([[ 0.3083,  1.4545,  1.0125,  0.6652],
        [-0.3993,  0.3083,  1.0602,  1.0168],
        [ 1.0168,  0.6652,  2.3568,  2.8065],
        [ 1.0602,  1.0125,  1.7383,  2.3568]], dtype=torch.float64)

tensor([0.0004, 0.0002, 0.0005, 0.0005], dtype=torch.float64)

# Output at each step

# It is possible to collect output of state or tangent matrix at each integration step
# Number of integratin steps is controlled by ns parameter on initialization
# Alternatively, desired integration step length can be passed
# Number of integration steps is computed as ceil(length/ds)

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
kn = 0.5
ks = 0.5
ms = 1.0
mo = 1.0
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp, exact=False, ns=10, output=True, matrix=True)

# Final state is still returned

print(D(state))

# Data is added to special attributes (state and tangent matrix)

print(D.container_output.shape)
print(D.container_matrix.shape)

# Number of integration steps can be changed

D.ns = 100

D(state)
print(D.container_output.shape)
print(D.container_matrix.shape)

tensor([-0.0086, -0.0098, -0.0022, -0.0008], dtype=torch.float64)
torch.Size([10, 4])
torch.Size([10, 4, 4])
torch.Size([100, 4])
torch.Size([100, 4, 4])

# Integration order is set on initialization (default value is zero)
# This order is related to difference order as 2n + 2
# Thus, zero corresponds to second order difference method

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
kn = 0.5
ks = 0.5
ms = 1.0
mo = 1.0
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp, exact=True, ns=10)

# For dipole integration is always performed
# In exact case, kinematic term error is added

D.ns = 10
ref = D(state)

D.ns = 100
res = D(state)

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())
print()

# Integrator parameters are stored in data attribute (if integration is actually performed)

maps, weights = D._data
print(maps)
print(weights)

[-0.00858751580977462, -0.009819760091907982, -0.0021662845096666727, -0.0008078063877268448]
[-0.008587660168877631, -0.00982000432362944, -0.0021667806526843485, -0.0008081645802833406]
[1.443591030117486e-07, 2.442317214579326e-07, 4.961430176758305e-07, 3.5819255649581234e-07]

[0, 1, 2, 3, 4, 3, 2, 1, 0]
[0.5, 0.5, 0.5, 0.5, 1.0, 0.5, 0.5, 0.5, 0.5]

# Derivatives of twiss parameters (chromaticity)

# pip install git+https://github.com/i-a-morozov/twiss.git@main
# pip install git+https://github.com/i-a-morozov/ndmap.git@main

from twiss import twiss

from ndmap.pfp import parametric_fixed_point
from ndmap.evaluate import evaluate

QF = Quadrupole('QF', 0.5, +0.21)
QD = Quadrupole('QD', 0.5, -0.19)
SF = Sextupole('SF', 0.25)
SD = Sextupole('SD', 0.25)
DR = Drift('DR', 0.25)
BM = Dipole('DR', 3.50, 0.15)

def fodo(state, dp, ms):
    dp, *_ = dp
    msf, msd, *_ = ms
    state = QF(state, data={**QF.data(), **{'dp': dp}})
    state = DR(state, data={**DR.data(), **{'dp': dp}})
    state = SF(state, data={**SF.data(), **{'dp': dp, 'ms': msf}})
    state = DR(state, data={**DR.data(), **{'dp': dp}})
    state = BM(state, data={**BM.data(), **{'dp': dp}})
    state = DR(state, data={**DR.data(), **{'dp': dp}})
    state = SD(state, data={**SD.data(), **{'dp': dp, 'ms': msd}})
    state = DR(state, data={**DR.data(), **{'dp': dp}})
    state = QD(state, data={**QD.data(), **{'dp': dp}})
    state = QD(state, data={**QD.data(), **{'dp': dp}})
    state = DR(state, data={**DR.data(), **{'dp': dp}})
    state = SD(state, data={**SD.data(), **{'dp': dp, 'ms': msd}})
    state = DR(state, data={**DR.data(), **{'dp': dp}})
    state = BM(state, data={**BM.data(), **{'dp': dp}})
    state = DR(state, data={**DR.data(), **{'dp': dp}})
    state = SF(state, data={**SF.data(), **{'dp': dp, 'ms': msf}})
    state = DR(state, data={**DR.data(), **{'dp': dp}})
    state = QF(state, data={**QF.data(), **{'dp': dp}})
    return state

# Set deviation parameters

msf = torch.tensor(0.0, dtype=torch.float64)
msd = torch.tensor(0.0, dtype=torch.float64)
ms = torch.stack([msf, msd])
dp = torch.tensor([0.0], dtype=torch.float64)

# Set fixed point

fp = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)


# Compute parametrix fixed point (first order in momentum deviation)
# Note, all parameters must be vectors

pfp, *_ = parametric_fixed_point((1, ), fp, [dp], fodo, ms)

# Define transformation around fixed point

def pfp_fodo(state, dp, ms):
    return fodo(state + evaluate(pfp, [dp]), dp, ms) - evaluate(pfp, [dp])

# Tune

def tune(dp, ms):
    matrix = torch.func.jacrev(pfp_fodo)(fp, dp, ms)
    tune, *_ = twiss(matrix)
    return tune

# Chromaticity

def chromaticity(ms):
    return torch.func.jacrev(tune)(dp, ms)

# Compute tunes

tunes = tune(dp, ms)
print(tunes)


# Compute chromaticity

chromaticities = chromaticity(ms).squeeze()
print(chromaticities)

# Compute derivative of chromaticities

jacobian = torch.func.jacrev(chromaticity)(ms).squeeze()
print(jacobian)

# Correct chomaticity

print((chromaticity(ms - torch.linalg.pinv(jacobian) @ chromaticities)).squeeze())

tensor([0.2210, 0.1703], dtype=torch.float64)
tensor([-0.2310, -0.2107], dtype=torch.float64)
tensor([[ 1.6013,  0.3613],
        [-0.7593, -1.2452]], dtype=torch.float64)
tensor([-1.1102e-16,  1.3878e-16], dtype=torch.float64)