Example-16: Custom (element)

# In this example steps to build element with custom transformatiion are shown
# This element can be used to wrap a custom transformation and used it with other elements
# Construction requires transformation to be passed as a parameter
# This transformation is expected to have certain signature
# The first parameter is state (qx, px, qy, py) followed by other tensor knobs
# These knobs are passed simular to deviation parameters in standart elements

import torch

from model.library.transformations import quadrupole
from model.library.quadrupole import Quadrupole
from model.library.custom import Custom

# Build custom element using a transformation

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

kn = torch.tensor(5.0, dtype=torch.float64)
ks = torch.tensor(1.0, dtype=torch.float64)
dp = torch.tensor(0.001, dtype=torch.float64)
length = torch.tensor(1.0, dtype=torch.float64)

print(quadrupole(state, kn, ks, dp, length))

C = Custom('C', quadrupole, ['kn', 'ks', 'dp', 'length'], length=1.0, angle=0.0, dp=dp.item())

print(C(state, data={**C.data(), **{'kn': kn, 'ks': ks, 'dp': dp, 'length': length}}, alignment=False))

tensor([-0.0101, -0.0189, -0.0173, -0.0385], dtype=torch.float64)
tensor([-0.0101, -0.0189, -0.0173, -0.0385], dtype=torch.float64)

# Use regular element step

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

kn = torch.tensor(5.0, dtype=torch.float64)
ks = torch.tensor(1.0, dtype=torch.float64)
dp = torch.tensor(0.001, dtype=torch.float64)
dl = torch.tensor(0.0, dtype=torch.float64)

Q = Quadrupole('Q', 1.0, 5.0, 1.0, dp=0.001, dx=-0.001)
print(Q(state, data={**Q.data(), **{'kn': kn, 'ks': ks, 'dp': dp, 'dl': dl}}, alignment=False))

print(Q._step(state, kn, ks, dp, dl))

C = Custom('C', Q._step, ['kn', 'ks', 'dp', 'dl'], length=1.0, angle=0.0, dp=0.001, dx=-0.001)
print(C(state, data={**C.data(), **{'kn': kn, 'ks': ks, 'dp': dp, 'dl': dl}}, alignment=False))

tensor([-0.0149, -0.0080, -0.0454, -0.1485], dtype=torch.float64)
tensor([-0.0149, -0.0080, -0.0454, -0.1485], dtype=torch.float64)
tensor([-0.0149, -0.0080, -0.0454, -0.1485], dtype=torch.float64)

# Alignment support

print(Q(state, data={**Q.data(), **{'kn': kn, 'ks': ks, 'dp': dp, 'dl': dl}}, alignment=True))
print(C(state, data={**C.data(), **{'kn': kn, 'ks': ks, 'dp': dp, 'dl': dl}}, alignment=True))

tensor([-0.0168, -0.0074, -0.0441, -0.1447], dtype=torch.float64)
tensor([-0.0168, -0.0074, -0.0441, -0.1447], dtype=torch.float64)