Example-14: Linear (element)

# Can be used to model elements as linear transformations
# Constant offset can be used to model first order dispersion

import torch
from model.library.quadrupole import Quadrupole
from model.library.dipole import Dipole
from model.library.linear import Linear

# Linear quadrupole

length = 1.0
kn = - 2.0
ks = + 1.5
dp = 0.0
Q = Quadrupole('Q', length, kn, ks, dp)

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

matrix = torch.func.jacrev(lambda state, dp: Q(state, data={**Q.data(), **{'dp': dp}}), 0)(state, Q.dp)
vector = torch.func.jacrev(lambda state, dp: Q(state, data={**Q.data(), **{'dp': dp}}), 1)(state, Q.dp)

length = 1.0
kn = - 2.0
ks = + 1.5
dp = 0.0001
Q = Quadrupole('Q', length, kn, ks, dp)

L = Linear('L', (dp*vector).tolist(), matrix.tolist())

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

print(Q(state).tolist())
print(L(state).tolist())
print((Q(state) - L(state)).tolist())

[0.012284014325104016, 0.013015328226371463, 0.005866573649483532, 0.01748719571870364]
[0.012284330668424102, 0.013015814342968814, 0.005867420763768485, 0.01748647747810468]
[-3.163433200861071e-07, -4.861165973507608e-07, -8.471142849530294e-07, 7.182405989576701e-07]

# Linear dipole

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
dp = 0.0
D = Dipole('D', length, angle, e1, e2, dp)

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

matrix = torch.func.jacrev(lambda state, dp: D(state, data={**D.data(), **{'dp': dp}}), 0)(state, D.dp)
vector = torch.func.jacrev(lambda state, dp: D(state, data={**D.data(), **{'dp': dp}}), 1)(state, D.dp)

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
dp = 0.0001
D = Dipole('D', length, angle, e1, e2, dp)

L = Linear('L', (dp*vector).tolist(), matrix.tolist())

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

print(D(state).tolist())
print(L(state).tolist())
print((D(state) - L(state)).tolist())

[2.8326139356223135e-06, -0.005001000818821848, -0.002994615078158236, 0.001004198140800915]
[9.165104284458664e-06, -0.004994998958072864, -0.002993748697591074, 0.001004997134048961]
[-6.33249034883635e-06, -6.001860748984511e-06, -8.663805671619597e-07, -7.989932480460815e-07]