Variables (#700)

* rebased on master

* Element array attributes are always ParamArrays

* Documentation of Variables

* Documentation of Parameters

* PEP8

* [0.0, k] instead of [0, k]

* Removed parameters

* merged from master

* lattice_variables

* Refactored Variable to VariableBase

* Refactored Variable to VariableBase

* "black" formatting

* merged elements.py from master

* str and repr

* str and repr

* small fixes

* utils.py

* changed Number

* minimised differences with master

* Optimised imports

* Documentation

* PEP8 checks

* variable notebook

* axisdef

* rebased on master

* merged from master

* merged from master

* code style

* Changed the examples of custom variables in the "variables" notebook

Author: lfarv

atintegrators/atelem.c

@@ -193,16 +193,18 @@ static long atGetLong(const PyObject *element, const char *name)
 {
     const PyObject *attr = PyObject_GetAttrString((PyObject *)element, name);
     if (!attr) return 0L;
+    long l = PyLong_AsLong((PyObject *)attr);
     Py_DECREF(attr);
-    return PyLong_AsLong((PyObject *)attr);
+    return l;
 }
 
 static double atGetDouble(const PyObject *element, const char *name)
 {
     const PyObject *attr = PyObject_GetAttrString((PyObject *)element, name);
     if (!attr) return 0.0;
+    double d = PyFloat_AsDouble((PyObject *)attr);
     Py_DECREF(attr);
-    return PyFloat_AsDouble((PyObject *)attr);
+    return d;
 }
 
 static long atGetOptionalLong(const PyObject *element, const char *name, long default_value)

docs/conf.py

@@ -40,8 +40,9 @@
               'sphinx.ext.intersphinx',
               'sphinx.ext.githubpages',
               'sphinx.ext.viewcode',
-              'myst_parser',
+              'myst_nb',
               'sphinx_copybutton',
+              'sphinx_design',
               ]
 
 intersphinx_mapping = {'python': ('https://docs.python.org/3', None),
@@ -56,7 +57,7 @@
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
 # This pattern also affects html_static_path and html_extra_path.
-exclude_patterns = ["README.rst", "**/*.so"]
+exclude_patterns = ["README.rst", "**/*.so", "_build/*"]
 rst_prolog = """
 .. role:: pycode(code)
    :language: python
@@ -92,6 +93,8 @@
     "deflist"
 ]
 myst_heading_anchors = 3
+nb_execution_mode = "auto"
+nb_execution_allow_errors = True
 
 # -- Options for HTML output -------------------------------------------------
 

docs/p/index.rst

@@ -26,6 +26,7 @@ Sub-packages
 
    howto/Installation
    howto/Primer
+   notebooks/variables
 
 .. toctree::
    :maxdepth: 2

docs/p/notebooks/variables.ipynb

@@ -0,0 +1,2 @@
+from .lattice.variables import *
+from .lattice.lattice_variables import *

pyat/at/future.py

@@ -9,11 +9,14 @@
 from .axisdef import *
 from .options import DConstant, random
 from .particle_object import Particle
+# from .variables import *
+from .variables import VariableList
 from .elements import *
 from .rectangular_bend import *
 from .idtable_element import InsertionDeviceKickMap
 from .utils import *
 from .lattice_object import *
+# from .lattice_variables import *
 from .cavity_access import *
 from .variable_elements import *
 from .deprecated import *

pyat/at/lattice/__init__.py

@@ -1,6 +1,8 @@
 """Helper functions for axis and plane descriptions"""
+
 from __future__ import annotations
 from typing import Optional, Union
+
 # For sys.version_info.minor < 9:
 from typing import Tuple
 
@@ -16,31 +18,31 @@
     ct=dict(index=5, label=r"$\beta c \tau$", unit=" [m]"),
 )
 for xk, xv in [it for it in _axis_def.items()]:
-    xv['code'] = xk
-    _axis_def[xv['index']] = xv
+    xv["code"] = xk
+    _axis_def[xv["index"]] = xv
     _axis_def[xk.upper()] = xv
-_axis_def['delta'] = _axis_def['dp']
-_axis_def['xp'] = _axis_def['px']   # For backward compatibility
-_axis_def['yp'] = _axis_def['py']   # For backward compatibility
-_axis_def['s'] = _axis_def['ct']
-_axis_def['S'] = _axis_def['ct']
-_axis_def[None] = dict(index=slice(None), label="", unit="", code=":")
+_axis_def["delta"] = _axis_def["dp"]
+_axis_def["xp"] = _axis_def["px"]  # For backward compatibility
+_axis_def["yp"] = _axis_def["py"]  # For backward compatibility
+_axis_def["s"] = _axis_def["ct"]
+_axis_def["S"] = _axis_def["ct"]
+_axis_def[None] = dict(index=None, label="", unit="", code=":")
 _axis_def[Ellipsis] = dict(index=Ellipsis, label="", unit="", code="...")
 
 _plane_def = dict(
     x=dict(index=0, label="x", unit=" [m]"),
     y=dict(index=1, label="y", unit=" [m]"),
-    z=dict(index=2, label="z", unit="")
+    z=dict(index=2, label="z", unit=""),
 )
 for xk, xv in [it for it in _plane_def.items()]:
-    xv['code'] = xk
-    _plane_def[xv['index']] = xv
+    xv["code"] = xk
+    _plane_def[xv["index"]] = xv
     _plane_def[xk.upper()] = xv
-_plane_def['h'] = _plane_def['x']
-_plane_def['v'] = _plane_def['y']
-_plane_def['H'] = _plane_def['x']
-_plane_def['V'] = _plane_def['y']
-_plane_def[None] = dict(index=slice(None), label="", unit="", code=":")
+_plane_def["h"] = _plane_def["x"]
+_plane_def["v"] = _plane_def["y"]
+_plane_def["H"] = _plane_def["x"]
+_plane_def["V"] = _plane_def["y"]
+_plane_def[None] = dict(index=None, label="", unit="", code=":")
 _plane_def[Ellipsis] = dict(index=Ellipsis, label="", unit="", code="...")
 
 

pyat/at/lattice/axisdef.py

@@ -33,7 +33,7 @@
 from .elements import Element
 from .particle_object import Particle
 from .utils import AtError, AtWarning, Refpts
-from .utils import get_s_pos, get_elements,get_value_refpts, set_value_refpts
+from .utils import get_s_pos, get_elements, get_value_refpts, set_value_refpts
 # noinspection PyProtectedMember
 from .utils import get_uint32_index, get_bool_index, _refcount, Uint32Refpts
 from .utils import refpts_iterator, checktype, set_shift, set_tilt, get_geometry
@@ -296,11 +296,11 @@ def _addition_filter(self, elems: Iterable[Element], copy_elements=False):
         if cavities and not hasattr(self, '_cell_harmnumber'):
             cavities.sort(key=lambda el: el.Frequency)
             try:
-                self._cell_harmnumber = getattr(cavities[0], 'HarmNumber')
+                self._cell_harmnumber = cavities[0].HarmNumber
             except AttributeError:
                 length += self.get_s_pos(len(self))[0]
                 rev = self.beta * clight / length
-                frequency = getattr(cavities[0], 'Frequency')
+                frequency = cavities[0].Frequency
                 self._cell_harmnumber = int(round(frequency / rev))
         self._radiation |= params.pop('_radiation')
 
@@ -314,13 +314,13 @@ def insert(self, idx: SupportsIndex, elem: Element, copy_elements=False):
                                  If :py:obj:`True` a deep copy of elem
                                  is used.
         """
-        # noinspection PyUnusedLocal
         # scan the new element to update it
-        elist = list(self._addition_filter([elem],
+        elist = list(self._addition_filter([elem],  # noqa: F841
                      copy_elements=copy_elements))
         super().insert(idx, elem)
 
     def extend(self, elems: Iterable[Element], copy_elements=False):
+        # noinspection PyUnresolvedReferences
         r"""This method adds all the elements of `elems` to the end of the
         lattice. The behavior is the same as for a :py:obj:`list`
 
@@ -343,6 +343,7 @@ def extend(self, elems: Iterable[Element], copy_elements=False):
         super().extend(elems)
 
     def append(self, elem: Element, copy_elements=False):
+        # noinspection PyUnresolvedReferences
         r"""This method overwrites the inherited method :py:meth:`list.append()`,
         its behavior is changed, it accepts only AT lattice elements
         :py:obj:`Element` as input argument.
@@ -361,6 +362,7 @@ def append(self, elem: Element, copy_elements=False):
         self.extend([elem], copy_elements=copy_elements)
 
     def repeat(self, n: int, copy_elements: bool = True):
+        # noinspection SpellCheckingInspection,PyUnresolvedReferences,PyRedeclaration
         r"""This method allows to repeat the lattice `n` times.
         If `n` does not divide `ring.periodicity`, the new ring
         periodicity is set to 1, otherwise  it is set to
@@ -405,6 +407,7 @@ def copy_fun(elem, copy):
 
     def concatenate(self, *lattices: Iterable[Element],
                     copy_elements=False, copy=False):
+        # noinspection PyUnresolvedReferences,SpellCheckingInspection,PyRedeclaration
         """Concatenate several `Iterable[Element]` with the lattice
 
         Equivalent syntaxes:
@@ -439,6 +442,7 @@ def concatenate(self, *lattices: Iterable[Element],
         return lattice if copy else None
 
     def reverse(self, copy=False):
+        # noinspection PyUnresolvedReferences
         r"""Reverse the order of the lattice and swapt the faces
         of elements. Alignment errors are not swapped
 
@@ -516,7 +520,7 @@ def copy(self) -> Lattice:
     def deepcopy(self) -> Lattice:
         """Returns a deep copy of the lattice"""
         return copy.deepcopy(self)
-        
+
     def slice_elements(self, refpts: Refpts, slices: int = 1) -> Lattice:
         """Create a new lattice by slicing the elements at refpts
 
@@ -538,7 +542,7 @@ def slice_generator(_):
                 else:
                     yield el
 
-        return Lattice(slice_generator, iterator=self.attrs_filter)        
+        return Lattice(slice_generator, iterator=self.attrs_filter)
 
     def slice(self, size: Optional[float] = None, slices: Optional[int] = 1) \
             -> Lattice:
@@ -635,8 +639,8 @@ def energy(self) -> float:
     def energy(self, energy: float):
         # Set the Energy attribute of radiating elements
         for elem in self:
-            if (isinstance(elem, (elt.RFCavity, elt.Wiggler)) or
-                    elem.PassMethod.endswith('RadPass')):
+            if (isinstance(elem, (elt.RFCavity, elt.Wiggler))
+                    or elem.PassMethod.endswith('RadPass')):
                 elem.Energy = energy
         # Set the energy attribute of the Lattice
         # Use a numpy scalar to allow division by zero
@@ -1471,7 +1475,7 @@ def params_filter(params, elem_filter: Filter, *args) -> Generator[Element, None
     cavities = []
     cell_length = 0
 
-    for idx, elem in enumerate(elem_filter(params, *args)):
+    for elem in elem_filter(params, *args):
         if isinstance(elem, elt.RFCavity):
             cavities.append(elem)
         elif hasattr(elem, 'Energy'):

pyat/at/lattice/elements.py

@@ -0,0 +1,137 @@
+"""Variables are **references** to scalar attributes of lattice elements. There are 2
+kinds of element variables:
+
+- an :py:class:`ElementVariable` is associated to an element object, and acts on all
+  occurences of this object. But it will not affect any copy, neither shallow nor deep,
+  of the original object,
+- a :py:class:`RefptsVariable` is not associated to an element object, but to an element
+  location in a :py:class:`.Lattice`. It acts on any copy of the initial lattice. A
+  *ring* argument must be provided to the *set* and *get* methods to identify the
+  lattice, which may be a possibly modified copy of the original lattice
+"""
+
+from __future__ import annotations
+
+__all__ = ["RefptsVariable", "ElementVariable"]
+
+from collections.abc import Sequence
+from typing import Union, Optional
+
+import numpy as np
+
+from .elements import Element
+from .lattice_object import Lattice
+from .utils import Refpts, getval, setval
+from .variables import VariableBase
+
+
+class RefptsVariable(VariableBase):
+    r"""A reference to a scalar attribute of :py:class:`.Lattice` elements.
+
+    It can refer to:
+
+    * a scalar attribute or
+    * an element of an array attribute
+
+    of one or several :py:class:`.Element`\ s of a lattice.
+
+    A :py:class:`RefptsVariable` is not associated to element objets themselves, but
+    to the location of these elements in a lattice. So a :py:class:`RefptsVariable`
+    will act equally on any copy of the initial ring.
+    As a consequence, a *ring* keyword argument (:py:class:`.Lattice` object) must be
+    supplied for getting or setting the variable.
+    """
+
+    def __init__(
+        self, refpts: Refpts, attrname: str, index: Optional[int] = None, **kwargs
+    ):
+        r"""
+        Parameters:
+            refpts:     Location of variable :py:class:`.Element`\ s
+            attrname:   Attribute name
+            index:      Index in the attribute array. Use :py:obj:`None` for
+              scalar attributes
+
+        Keyword Args:
+            name (str):     Name of the Variable. Default: ``''``
+            bounds (tuple[float, float]):   Lower and upper bounds of the
+              variable value. Default: (-inf, inf)
+            delta (float):  Step. Default: 1.0
+            ring (Lattice): If specified, it is used to get and store the initial
+              value of the variable. Otherwise, the initial value is set to None
+        """
+        self._getf = getval(attrname, index=index)
+        self._setf = setval(attrname, index=index)
+        self.refpts = refpts
+        super().__init__(**kwargs)
+
+    def _setfun(self, value: float, ring: Lattice = None):
+        if ring is None:
+            raise ValueError("Can't set values if ring is None")
+        for elem in ring.select(self.refpts):
+            self._setf(elem, value)
+
+    def _getfun(self, ring: Lattice = None) -> float:
+        if ring is None:
+            raise ValueError("Can't get values if ring is None")
+        values = np.array([self._getf(elem) for elem in ring.select(self.refpts)])
+        return np.average(values)
+
+
+class ElementVariable(VariableBase):
+    r"""A reference to a scalar attribute of :py:class:`.Lattice` elements.
+
+    It can refer to:
+
+    * a scalar attribute or
+    * an element of an array attribute
+
+    of one or several :py:class:`.Element`\ s of a lattice.
+
+    An :py:class:`ElementVariable` is associated to an element object, and acts on all
+    occurrences of this object. But it will not affect any copy, neither shallow nor
+    deep, of the original object.
+    """
+
+    def __init__(
+        self,
+        elements: Union[Element, Sequence[Element]],
+        attrname: str,
+        index: Optional[int] = None,
+        **kwargs,
+    ):
+        r"""
+        Parameters:
+            elements:   :py:class:`.Element` or Sequence[Element] whose
+              attribute is varied
+            attrname:   Attribute name
+            index:      Index in the attribute array. Use :py:obj:`None` for
+              scalar attributes
+
+        Keyword Args:
+            name (str):     Name of the Variable. Default: ``''``
+            bounds (tuple[float, float]):   Lower and upper bounds of the
+              variable value. Default: (-inf, inf)
+            delta (float):  Step. Default: 1.0
+        """
+        # Ensure the uniqueness of elements
+        if isinstance(elements, Element):
+            self._elements = {elements}
+        else:
+            self._elements = set(elements)
+        self._getf = getval(attrname, index=index)
+        self._setf = setval(attrname, index=index)
+        super().__init__(**kwargs)
+
+    def _setfun(self, value: float, **kwargs):
+        for elem in self._elements:
+            self._setf(elem, value)
+
+    def _getfun(self, **kwargs) -> float:
+        values = np.array([self._getf(elem) for elem in self._elements])
+        return np.average(values)
+
+    @property
+    def elements(self):
+        """Return the set of elements acted upon by the variable"""
+        return self._elements

pyat/at/lattice/lattice_object.py

@@ -37,6 +37,7 @@
 from typing import Union, Tuple, List, Type
 from enum import Enum
 from itertools import compress
+from operator import attrgetter
 from fnmatch import fnmatch
 from .elements import Element, Dipole
 
@@ -58,7 +59,7 @@
            'set_shift', 'set_tilt', 'set_rotation',
            'tilt_elem', 'shift_elem', 'rotate_elem',
            'get_value_refpts', 'set_value_refpts', 'Refpts',
-           'get_geometry']
+           'get_geometry', 'setval', 'getval']
 
 _axis_def = dict(
     x=dict(index=0, label="x", unit=" [m]"),
@@ -113,6 +114,72 @@ def _type_error(refpts, types):
         "Invalid refpts type {0}. Allowed types: {1}".format(tp, types))
 
 
+# setval and getval return pickleable functions: no inner, nested function
+# are allowed. So nested functions are replaced be module-level callable
+# class instances
+class _AttrItemGetter(object):
+    __slots__ = ["attrname", "index"]
+
+    def __init__(self, attrname: str, index: int):
+        self.attrname = attrname
+        self.index = index
+
+    def __call__(self, elem):
+        return getattr(elem, self.attrname)[self.index]
+
+
+def getval(attrname: str, index: Optional[int] = None) -> Callable:
+    """Return a callable object which fetches item *index* of
+    attribute *attrname* of its operand. Examples:
+
+    - After ``f = getval('Length')``, ``f(elem)`` returns ``elem.Length``
+    - After ``f = getval('PolynomB, index=1)``, ``f(elem)`` returns
+      ``elem.PolynomB[1]``
+
+    """
+    if index is None:
+        return attrgetter(attrname)
+    else:
+        return _AttrItemGetter(attrname, index)
+
+
+class _AttrSetter(object):
+    __slots__ = ["attrname"]
+
+    def __init__(self, attrname: str):
+        self.attrname = attrname
+
+    def __call__(self, elem, value):
+        setattr(elem, self.attrname, value)
+
+
+class _AttrItemSetter(object):
+    __slots__ = ["attrname", "index"]
+
+    def __init__(self, attrname: str, index: int):
+        self.attrname = attrname
+        self.index = index
+
+    def __call__(self, elem, value):
+        getattr(elem, self.attrname)[self.index] = value
+
+
+def setval(attrname: str, index: Optional[int] = None) -> Callable:
+    """Return a callable object which sets the value of  item *index* of
+    attribute *attrname* of its 1st argument to it 2nd orgument.
+
+    - After ``f = setval('Length')``, ``f(elem, value)`` is equivalent to
+      ``elem.Length = value``
+    - After ``f = setval('PolynomB, index=1)``, ``f(elem, value)`` is
+      equivalent to ``elem.PolynomB[1] = value``
+
+    """
+    if index is None:
+        return _AttrSetter(attrname)
+    else:
+        return _AttrItemSetter(attrname, index)
+
+
 # noinspection PyIncorrectDocstring
 def axis_descr(*args, key=None) -> Tuple:
     r"""axis_descr(axis [ ,axis], key=None)
@@ -779,13 +846,7 @@ def get_value_refpts(ring: Sequence[Element], refpts: Refpts,
     Returns:
         attrvalues: numpy Array of attribute values.
     """
-    if index is None:
-        def getf(elem):
-            return getattr(elem, attrname)
-    else:
-        def getf(elem):
-            return getattr(elem, attrname)[index]
-
+    getf = getval(attrname, index=index)
     return numpy.array([getf(elem) for elem in refpts_iterator(ring, refpts,
                                                                regex=regex)])
 
@@ -822,13 +883,7 @@ def set_value_refpts(ring: Sequence[Element], refpts: Refpts,
              elements are shared with the original lattice.
              Any further modification will affect both lattices.
     """
-    if index is None:
-        def setf(elem, value):
-            setattr(elem, attrname, value)
-    else:
-        def setf(elem, value):
-            getattr(elem, attrname)[index] = value
-
+    setf = setval(attrname, index=index)
     if increment:
         attrvalues += get_value_refpts(ring, refpts,
                                        attrname, index=index,
@@ -841,8 +896,7 @@ def setf(elem, value):
     # noinspection PyShadowingNames
     @make_copy(copy)
     def apply(ring, refpts, values, regex):
-        for elm, val in zip(refpts_iterator(ring, refpts,
-                                            regex=regex), values):
+        for elm, val in zip(refpts_iterator(ring, refpts, regex=regex), values):
             setf(elm, val)
 
     return apply(ring, refpts, attrvalues, regex)