Example-12: Gradient (element)

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

from pathlib import Path
from os import system

import torch
from model.library.gradient import Gradient

# Tracking

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

exact = True
align = False

kn = +5.0E-3
ks = -2.5E-3

dp = 0.005

state = torch.tensor([0.001, 0.0, -0.005, 0.0], 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: multipole,knl={{0.0,{kn}}},ksl={{0.0,{ks}}};
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=10,sector_nmul=10 ;
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)

G = Gradient('G', kn=kn, ks=ks, dp=dp)
res = G(state, alignment=align, data={**G.data(), **error})

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

ptc.unlink()
obs.unlink()

[0.001, 7.5e-06, -0.005, -2.75e-05]
[0.001, 7.5e-06, -0.005, -2.75e-05]
[0.0, 0.0, 0.0, 0.0]

# Tracking (alignment)

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

exact = True
align = True

kn = +5.0E-3
ks = -2.5E-3

dp = 0.005

state = torch.tensor([0.001, 0.0, -0.005, 0.0], 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: multipole,knl={{0.0,{kn}}},ksl={{0.0,{ks}}};
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=10,sector_nmul=10 ;
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)

G = Gradient('G', kn=kn, ks=ks, dp=dp)
res = G(state, alignment=align, data={**G.data(), **error})

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

ptc.unlink()
obs.unlink()

[0.000991887042015914, 0.00016412084104346932, -0.005011606734840507, 0.00023479809553215934]
[0.0009918870420159223, 0.00016412084104346948, -0.005011606734840506, 0.00023479809553215904]
[-8.239936510889834e-18, -1.6263032587282567e-19, -8.673617379884035e-19, 2.981555974335137e-19]

# Deviation/error variables

kn = +5.0E-3
ks = -2.5E-3
dp = 0.005
state = torch.tensor([0.001, 0.0, -0.005, 0.0], dtype=torch.float64)

dx = torch.tensor(+0.01, dtype=torch.float64)
dy = torch.tensor(-0.01, dtype=torch.float64)
dz = torch.tensor(0.0, dtype=torch.float64)

wx = torch.tensor(0.0, dtype=torch.float64)
wy = torch.tensor(0.0, dtype=torch.float64)
wz = torch.tensor(torch.pi, dtype=torch.float64)

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

G = Gradient('G', kn, ks, 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(G(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(G(state, data={**G.data(), **{'kn': -kn, 'ks': -ks}}))
print()

# In the above G.data() creates default deviation dictionary (with zero values for each deviaton)
# {**G.data(), **{'kn': -kn, 'ks': -ks}} replaces the 'kn' and 'ks' key values

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

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


# The following elements can be made equivalent using deviation variables

GA = Gradient('GA', kn, ks, dp)
GB = Gradient('GB', kn - 0.001, ks, dp)

print(GA(state) - GB(state, data={**GB.data(), **{'kn': + 0.001}}))

tensor([ 1.0000e-03,  7.5000e-06, -5.0000e-03, -2.7500e-05],
       dtype=torch.float64)

tensor([ 0.0010,  0.0000, -0.0050,  0.0000], dtype=torch.float64)

tensor([ 1.0000e-03,  3.2500e-05, -5.0000e-03,  4.7500e-05],
       dtype=torch.float64)

tensor([0., 0., 0., 0.], 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

kn = +5.0E-3
ks = -2.5E-3
dp = 0.005

dx = torch.tensor(+0.01, dtype=torch.float64)
dy = torch.tensor(-0.01, dtype=torch.float64)
dz = torch.tensor(0.0, dtype=torch.float64)

wx = torch.tensor(0.0, dtype=torch.float64)
wy = torch.tensor(0.0, dtype=torch.float64)
wz = torch.tensor(torch.pi, dtype=torch.float64)

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

G = Gradient('G', kn, ks, dp)

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

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

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

def wrapper(state, kn, ks):
    return G(state, data={**G.data(), **{'kn': kn, 'ks': ks}})

kn = 1.0E-3*torch.randn(512, dtype=G.dtype, device=G.device)
ks = 1.0E-3*torch.randn(512, dtype=G.dtype, device=G.device)

print(torch.vmap(wrapper)(state, kn, ks).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

kn = +5.0E-3
ks = -2.5E-3
dp = 0.005
state = torch.tensor([0.001, 0.0, -0.005, 0.0], dtype=torch.float64)

dx = torch.tensor(+0.01, dtype=torch.float64)
dy = torch.tensor(-0.01, dtype=torch.float64)
dz = torch.tensor(0.0, dtype=torch.float64)

wx = torch.tensor(0.0, dtype=torch.float64)
wy = torch.tensor(0.0, dtype=torch.float64)
wz = torch.tensor(torch.pi, dtype=torch.float64)
error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

G = Gradient('G', kn, ks, dp)

# Compute derivative with respect to state

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

# Compute derivative with respect to a deviation variable

dkn = torch.tensor(0.0, dtype=torch.float64)
dks = torch.tensor(0.0, dtype=torch.float64)
dk = torch.stack([dkn, dks])

def wrapper(state, dk):
    dkn, dks = dk
    return G(state, data={**G.data(), **{'kn': dkn, 'ks': dks}})

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

tensor([[ 1.0000,  0.0000,  0.0000,  0.0000],
        [-0.0050,  1.0000, -0.0025,  0.0000],
        [ 0.0000,  0.0000,  1.0000,  0.0000],
        [-0.0025,  0.0000,  0.0050,  1.0000]], dtype=torch.float64)

tensor([[-0.0000,  0.0000],
        [-0.0010, -0.0050],
        [-0.0000,  0.0000],
        [-0.0050,  0.0010]], dtype=torch.float64)