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

from at.future import VariableBase, ElementVariable, RefptsVariable, CustomVariable

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}")
print(lf1.value)

lf1: ElementVariable(0.5, name='lf1')
0.5

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

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

kf1: ElementVariable(2.1, name='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, np.float64(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.9.19/x64/lib/python3.9/site-packages/at/lattice/variables.py:202, in VariableBase.set(self, value, ring)
    194 """Set the variable value
    195 
    196 Args:
   (...)
    199       may be necessary to set the variable.
    200 """
    201 if value < self.bounds[0] or value > self.bounds[1]:
--> 202     raise ValueError(f"set value must be in {self.bounds}")
    203 self._setfun(value, ring=ring)
    204 if np.isnan(self._initial):

ValueError: set value must be in (0.45, 0.55)

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

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 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

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, ring=None):
    at.DConstant.DPStep = value


def getvar1(ring=None):
    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, ring=None):
    lattice.energy = value

def getvar2(lattice, ring=None):
    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(energy_var.value)

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, *args, **kwargs):
        # Store the lattice
        self.lattice = lattice
        # Initialise the parent class
        super().__init__(*args, **kwargs)

    def _setfun(self, value, ring=None):
        self.lattice.energy = value

    def _getfun(self, ring=None):
        return self.lattice.energy

We construct the variable:

energy_var = EnergyVariable(ring)

Look at the initial state:

print(energy_var.value)

6000000000.0

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

6100000000.0

energy_var.reset()

Variables#

ElementVariable#

RefptsVariable#

Custom variables#

Using the CustomVariable#

Example 1#

Example 2#

By derivation of the VariableBase class#

Example 1#

Example 2#

`ElementVariable`#

`RefptsVariable`#

Using the `CustomVariable`#

By derivation of the `VariableBase` class#