Example-15: BPM (element)

[1]:

from pathlib import Path
from os import system

import torch
from model.library.drift import Drift
from model.library.bpm import BPM

[2]:

# BPM acts as identity transformation with calibration error

state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

B = BPM('B', direction='forward')
B(state)

[2]:

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

[3]:

# Calibration errors can be passed as deviation variables

# qx -> (1 + xx) qx + xy qy
# qy -> yx qx + (1 + yy) qy

B.data()

[3]:

{'xx': tensor(0., dtype=torch.float64),
 'xy': tensor(0., dtype=torch.float64),
 'yx': tensor(0., dtype=torch.float64),
 'yy': tensor(0., dtype=torch.float64),
 'dp': tensor(0., dtype=torch.float64),
 'dx': tensor(0., dtype=torch.float64),
 'dy': tensor(0., dtype=torch.float64),
 'dz': tensor(0., dtype=torch.float64),
 'wx': tensor(0., dtype=torch.float64),
 'wy': tensor(0., dtype=torch.float64),
 'wz': tensor(0., dtype=torch.float64)}

[4]:

# Transform to BPM frame and back to beam frame

xx = torch.tensor(+0.05, dtype=torch.float64)
xy = torch.tensor(+0.01, dtype=torch.float64)
yx = torch.tensor(+0.05, dtype=torch.float64)
yy = torch.tensor(-0.06, dtype=torch.float64)

B = BPM('B', direction='forward')

state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
print(state)

state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})
print(state)

B.direction = 'inverse'

state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})
print(state)

tensor([ 0.0100, -0.0050, -0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0105, -0.0048, -0.0042,  0.0011], dtype=torch.float64)
tensor([ 0.0100, -0.0050, -0.0050,  0.0010], dtype=torch.float64)

[5]:

# Transform to BPM frame and back to beam frame using a pair of BPMS

xx = torch.tensor(+0.05, dtype=torch.float64)
xy = torch.tensor(+0.01, dtype=torch.float64)
yx = torch.tensor(+0.05, dtype=torch.float64)
yy = torch.tensor(-0.06, dtype=torch.float64)

BA = BPM('B', direction='forward')
BB = BPM('B', direction='inverse')

state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
print(state)

state = BA(state, data={**BA.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})
print(state)

state = BB(state, data={**BB.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})
print(state)

tensor([ 0.0100, -0.0050, -0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0105, -0.0048, -0.0042,  0.0011], dtype=torch.float64)
tensor([ 0.0100, -0.0050, -0.0050,  0.0010], dtype=torch.float64)

[6]:

# Differentiability

B = BPM('B', direction='forward')
D = Drift('D', 1.0)

xx = torch.tensor(+0.05, dtype=torch.float64)
xy = torch.tensor(+0.01, dtype=torch.float64)
yx = torch.tensor(+0.05, dtype=torch.float64)
yy = torch.tensor(-0.06, dtype=torch.float64)

error = torch.stack([xx, xy, yx, yy])

state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

def line(state, error):
    xx, xy, yx, yy = error
    state = D(state)
    state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})
    return state

print(torch.func.jacrev(line, 1)(state, error))

def line(state, error):
    xx, xy, yx, yy = error
    state = D(state)
    B.direction = 'forward'
    state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})
    B.direction = 'inverse'
    state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})
    return state

print(torch.func.jacrev(line, 1)(state, error))

tensor([[ 5.0000e-03, -4.0000e-03,  0.0000e+00,  0.0000e+00],
        [ 4.5880e-03, -2.4404e-04, -1.0625e-03,  5.6516e-05],
        [ 0.0000e+00,  0.0000e+00,  5.0000e-03, -4.0000e-03],
        [-4.8809e-05,  5.1249e-03,  1.1303e-05, -1.1868e-03]],
       dtype=torch.float64)
tensor([[ 0.0000e+00,  8.6736e-19,  0.0000e+00, -5.4888e-19],
        [ 0.0000e+00,  0.0000e+00,  3.3881e-20,  0.0000e+00],
        [-2.7105e-19,  0.0000e+00,  8.6736e-19,  0.0000e+00],
        [ 2.1684e-19,  3.9302e-19, -3.3881e-21,  2.1684e-19]],
       dtype=torch.float64)

[7]:

# Alignment support

xx = torch.tensor(+0.05, dtype=torch.float64)
xy = torch.tensor(+0.01, dtype=torch.float64)
yx = torch.tensor(+0.05, dtype=torch.float64)
yy = torch.tensor(-0.06, 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}

B = BPM('B', direction='forward')

state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}, **error})
print(state)

state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}, **error}, alignment=True)
print(state)

tensor([ 0.0105, -0.0048, -0.0042,  0.0011], dtype=torch.float64)
tensor([ 0.0086, -0.0048, -0.0082,  0.0008], dtype=torch.float64)