import at
from importlib.resources import files, as_file
import numpy as np

from at.future import ElementVariable, RefptsVariable
from at import VariableBase, AttributeVariable, ItemVariable, CustomVariable, EvaluationVariable
from at import LocalOpticsObservable, ObservableList

Variables#

Variables are references to any scalar quantity. Predefined classes are available for accessing any scalar attribute of an element, or any item of an array attribute.

Any other quantity may be accessed by either subclassing the VariableBase abstract base class, or by using a CustomVariable.

`ElementVariable`#

An ElementVariable refers to a single attribute (or item of an array attribute) of one or several Element objects.

We now create a variable pointing to the length of a QF1 magnet:

qf1 = at.Quadrupole("QF1", 0.5, 2.1)
print(qf1)

Quadrupole:
       FamName: QF1
        Length: 0.5
    PassMethod: StrMPoleSymplectic4Pass
      MaxOrder: 1
   NumIntSteps: 10
      PolynomA: [0. 0.]
      PolynomB: [0.  2.1]
             K: 2.1

lf1 = ElementVariable(qf1, "Length", name="lf1")
print(f"{lf1}: {lf1.value}")

lf1: 0.5

and another variable pointing to the strength of the same magnet:

kf1 = ElementVariable(qf1, "PolynomB", index=1, name="kf1")
print(f"{kf1}: {kf1.value}")

kf1: 2.1

We can check which elements are concerned by the kf1 variable. The element container is a set, so that no element may appear twice:

kf1.elements

{Quadrupole('QF1', 0.5, 2.1)}

We can now change the strength of QF1 magnets and check again:

kf1.set(2.5)
print(qf1)

Quadrupole:
       FamName: QF1
        Length: 0.5
    PassMethod: StrMPoleSymplectic4Pass
      MaxOrder: 1
   NumIntSteps: 10
      PolynomA: [0. 0.]
      PolynomB: [0.  2.5]
             K: 2.5

We can look at the history of kf1 values

kf1.history

[np.float64(2.1), 2.5]

And revert to the initial or previous values:

kf1.set_previous()
print(qf1)

Quadrupole:
       FamName: QF1
        Length: 0.5
    PassMethod: StrMPoleSymplectic4Pass
      MaxOrder: 1
   NumIntSteps: 10
      PolynomA: [0. 0.]
      PolynomB: [0.  2.1]
             K: 2.1

An ElementVariable is linked to Elements. It will apply wherever the element appears but it will not follow any copy of the element, neither shallow nor deep. So if we make a copy of QF1:

qf2 = qf1.deepcopy()
print(f"qf1: {qf1.PolynomB[1]}")
print(f"qf2: {qf2.PolynomB[1]}")

qf1: 2.1
qf2: 2.1

and modify the kf1 variable:

kf1.set(2.6)

print(f"qf1: {qf1.PolynomB[1]}")
print(f"qf2: {qf2.PolynomB[1]}")

qf1: 2.6
qf2: 2.1

The copy of QF1 in is not affected.

One can set upper and lower bounds on a variable. Trying to set a value out of the bounds will raise a ValueError. The default is (-numpy.inf, numpy.inf).

lfbound = ElementVariable(qf1, "Length", bounds=(0.45, 0.55))

lfbound.set(0.2)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[14], line 1
----> 1 lfbound.set(0.2)

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/at/lattice/variables.py:360, in VariableBase.set(self, value, **setkw)
    348 def set(self, value: Number, **setkw) -> None:
    349     """Set the variable value.
    350 
    351     Args:
   (...)    358           They augment the keyword arguments given in the constructor.
    359     """
--> 360     self.check_bounds(value)
    361     self._setfun(value, *self.args, **(self.kwargs | setkw))
    362     if np.isnan(self._initial):

File /opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/at/lattice/variables.py:291, in VariableBase.check_bounds(self, value)
    289 if min_val is not None and value < min_val:
    290     msg = f"Value {value} must be larger or equal to {min_val}"
--> 291     raise ValueError(msg)
    292 if max_val is not None and value > max_val:
    293     msg = f"Value {value} must be smaller or equal to {max_val}"

ValueError: Value 0.2 must be larger or equal to 0.45

Variables also accept a delta keyword argument. Its value is used as the initial step in matching, and in the step_up() and step_down() methods.

`RefptsVariable`#

An RefptsVariable is similar to an ElementVariable but it is not associated with an Element itself, but with its location in a Lattice. So it will act on any lattice with the same elements.

But it needs a ring keyword in its set and get methods, to identify the selected lattice.

Let’s load a test ring and make a copy of it:

fname = "hmba.mat"
with as_file(files("machine_data") / fname) as path:
    ring = at.load_lattice(path)
newring = ring.deepcopy()

and create a RefptsVariable

kf2 = RefptsVariable("QF1[AE]", "PolynomB", index=1, name="kf2")

We can now use this variable on the two rings:

kf2.set(2.55, ring=ring)
kf2.set(2.45, ring=newring)

print(f"   ring:   {ring[5].PolynomB[1]}")
print(f"newring: {newring[5].PolynomB[1]}")

   ring:   2.55
newring: 2.45

`AttributeVariable`#

An AttributeVariable drives an attribute of an object. For example, we can drive a Lattice attribute like its energy:

energy_var = AttributeVariable(newring, "energy")

print(energy_var.value)

6000000000.0

energy_var.value = 6.1e9
print(newring.energy)

6100000000.0

We can look at the history of the variable

energy_var.history

[6000000000.0, 6100000000.0]

and go back to the initial value

energy_var.reset()

`ItemVariable`#

An ItemVariable drives an item of a directory or of a sequence.

obj = dict(key1=42.0, key2=[0.0, 1.0, 2.0, 3.0])
v1 = at.ItemVariable(obj, "key1")

v1 is associated with the item “key1” of obj (dict).

v2 = at.ItemVariable(obj, "key2", 1)

v2 is associated with the 2nd item of obj[“key2”] (list)

We can look at v1, change its value and check obj:

print(v1.value)
v1.value = 0.5
print(obj)

42.0
{'key1': 0.5, 'key2': [0.0, 1.0, 2.0, 3.0]}

We can also look at v2, change its value and check obj:

print(v2.value)
v2.value = 10.0
print(obj)

1.0
{'key1': 0.5, 'key2': [0.0, 10.0, 2.0, 3.0]}

An EvaluationVariable drives a parameter used by the ObservableList.evaluate() method. These parameters depend on the Observables involved. For example, optics observables rely on dp, orbit, twiss_in … Default values may be provided when instantiating the ObservableList. An EvaluationVariable can vary these default values before the evaluation.

Here is an ObservableList:

twiss_in = {"alpha": np.zeros(2), "beta": np.array([9.0, 2.5])}

obs = at.ObservableList(
    [
        at.LocalOpticsObservable([0], "beta", plane="x"),
        at.LocalOpticsObservable([0], "beta", plane="y"),
    ],
    ring=ring,
    dp=0.01,
    orbit=np.zeros(6),
    twiss_in=twiss_in,
)

Create a variable controlling \(\delta\):

v1 = EvaluationVariable(obs, "dp")

Create a variable controlling \(p_x\) and change its value:

v2 = EvaluationVariable(obs, "orbit", 1)
v2.value = 0.001

Create a variable controlling \(\beta_x\):

v3 = EvaluationVariable(obs, "twiss_in", "beta", 0)

We can look at their values:

print(v1.value)
print(v2.value)
print(v3.value)

01
001
0

The next evaluation of obs will use these values.

Custom variables#

Custom variables allow access to almost any quantity in AT. This can be achieved either by subclassing the VariableBase abstract base class, or by using a CustomVariable.

We will take 2 examples:

A variable accessing the DPStep AT parameter used in chromaticity computations. It does not look like a very useful variable, it’s for demonstration purpose,
A variable accessing the energy of a given lattice. This can be done more easily by using an AttributeVariable, but it’s also for demonstration.

Using the `CustomVariable`#

Using a CustomVariable makes it very easy to define simple variables: we just need to define two functions for the “get” and “set” actions, and give them to the CustomVariable constructor.

Example 1#

We define 2 functions for setting and getting the variable value:

def setvar1(value, **_):
    at.DConstant.DPStep = value


def getvar1(**_):
    return at.DConstant.DPStep

dpstep_var = CustomVariable(setvar1, getvar1, bounds=(1.0e-12, 0.1))
print(dpstep_var.value)

3e-06

dpstep_var.value = 2.0e-4
print(at.DConstant.DPStep)

0.0002

Example 2#

We can give to the CustomVariable constructor any positional or keyword argument necessary for the set and get functions. Here we will send the lattice as a positional argument:

def setvar2(value, lattice, **_):
    lattice.energy = value

def getvar2(lattice, **_):
    return lattice.energy

energy_var = CustomVariable(setvar2, getvar2, newring)

Here, the newring positional argument given to the variable constructor is available as a positional argument in both the set and get functions.

print(energy_var.value)

6000000000.0

energy_var.value = 6.1e9
print(newring.energy)

6100000000.0

We can look at the history of the variable

energy_var.history

[6000000000.0, 6100000000.0]

and go back to the initial value

energy_var.reset()

By derivation of the `VariableBase` class#

The derivation of VariableBase allows more control on the created variable by using the class constuctor and its arguments to setup the variable.

We will write a new variable class based on VariableBase abstract base class. The main task is to implement the _setfun and _getfun abstract methods.

Example 1#

class DPStepVariable(VariableBase):

    def _setfun(self, value, ring=None):
        at.DConstant.DPStep = value

    def _getfun(self, ring=None):
        return at.DConstant.DPStep

dpstep_var = DPStepVariable()
print(dpstep_var.value)

0.0002

dpstep_var.value = 3.0e-6
print(dpstep_var.value)

3e-06

Example 2#

Here we will store the lattice as an instance variable in the class constructor:

class EnergyVariable(VariableBase):
    def __init__(self, lattice):
        # Initialise the parent class
        super().__init__(lattice)

    def _setfun(self, value, lattice, **_):
        lattice.energy = value

    def _getfun(self, lattice, **_):
        return lattice.energy

We construct the variable:

energy_var = EnergyVariable(newring)

Look at the initial state:

print(energy_var.value)

6000000000.0

energy_var.value = 6.1e9
print(newring.energy)

6100000000.0

energy_var.reset()

Variables#

ElementVariable#

RefptsVariable#

AttributeVariable#

ItemVariable#

EvaluationVariable#

Custom variables#

Using the CustomVariable#

Example 1#

Example 2#

By derivation of the VariableBase class#

Example 1#

Example 2#

`ElementVariable`#

`RefptsVariable`#

`AttributeVariable`#

`ItemVariable`#

`EvaluationVariable`#

Using the `CustomVariable`#

By derivation of the `VariableBase` class#