Response matrices

docs/p/api/at.latticetools.response_matrix.rst

@@ -0,0 +1,15 @@
+at.latticetools.response\_matrix
+================================
+
+.. automodule:: at.latticetools.response_matrix
+   :inherited-members:
+
+
+   .. rubric:: Classes
+
+   .. autosummary::
+
+      ResponseMatrix
+      OrbitResponseMatrix
+      TrajectoryResponseMatrix
+

docs/p/index.rst

@@ -36,7 +36,8 @@ Sub-packages
    howto/multiprocessing
    howto/CavityControl
    howto/Collective
-   Working with MAD-X files <api/at.load.madx>
+   Work with MAD-X files <api/at.load.madx>
+   Use response matrices <notebooks/response_matrices>
 
 .. autosummary::
    :toctree: api

docs/p/notebooks/observables.ipynb

@@ -936,7 +936,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.20"
+   "version": "3.11.11"
   }
  },
  "nbformat": 4,

pyat/at/lattice/axisdef.py

@@ -1,23 +1,23 @@
 """Helper functions for axis and plane descriptions"""
 
 from __future__ import annotations
-from typing import Optional, Union
 
-# For sys.version_info.minor < 9:
+# Necessary for type aliases in python <= 3.8 :
 from typing import Tuple
+from typing import Union
 
 AxisCode = Union[str, int, slice, None, type(Ellipsis)]
 AxisDef = Union[AxisCode, Tuple[AxisCode, AxisCode]]
 
-_axis_def = dict(
-    x=dict(index=0, label="x", unit=" [m]"),
-    px=dict(index=1, label=r"$p_x$", unit=" [rad]"),
-    y=dict(index=2, label="y", unit=" [m]"),
-    py=dict(index=3, label=r"$p_y$", unit=" [rad]"),
-    dp=dict(index=4, label=r"$\delta$", unit=""),
-    ct=dict(index=5, label=r"$\beta c \tau$", unit=" [m]"),
-)
-for xk, xv in [it for it in _axis_def.items()]:
+_axis_def = {
+    "x": {"index": 0, "label": "x", "unit": " [m]"},
+    "px": {"index": 1, "label": r"$p_x$", "unit": " [rad]"},
+    "y": {"index": 2, "label": "y", "unit": " [m]"},
+    "py": {"index": 3, "label": r"$p_y$", "unit": " [rad]"},
+    "dp": {"index": 4, "label": r"$\delta$", "unit": ""},
+    "ct": {"index": 5, "label": r"$\beta c \tau$", "unit": " [m]"},
+}
+for xk, xv in list(_axis_def.items()):
     xv["code"] = xk
     _axis_def[xv["index"]] = xv
     _axis_def[xk.upper()] = xv
@@ -26,41 +26,43 @@
 _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=""),
-)
-for xk, xv in [it for it in _plane_def.items()]:
+_axis_def[None] = {"index": None, "label": "", "unit": "", "code": ":"}
+_axis_def[Ellipsis] = {"index": Ellipsis, "label": "", "unit": "", "code": "..."}
+
+_plane_def = {
+    "x": {"index": 0, "label": "x", "unit": " [m]"},
+    "y": {"index": 1, "label": "y", "unit": " [m]"},
+    "z": {"index": 2, "label": "z", "unit": ""},
+}
+for xk, xv in list(_plane_def.items()):
     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=None, label="", unit="", code=":")
-_plane_def[Ellipsis] = dict(index=Ellipsis, label="", unit="", code="...")
+_plane_def[None] = {"index": None, "label": "", "unit": "", "code": ":"}
+_plane_def[Ellipsis] = {"index": Ellipsis, "label": "", "unit": "", "code": "..."}
 
 
-def _descr(dd: dict, arg: AxisDef, key: Optional[str] = None):
-    if isinstance(arg, tuple):
-        return tuple(_descr(dd, a, key=key) for a in arg)
-    else:
-        try:
-            descr = dd[arg]
-        except (TypeError, KeyError):
-            descr = dict(index=arg, code=arg, label="", unit="")
-        if key is None:
-            return descr
+def _descr(dd: dict, *args: AxisDef, key: str | None = None):
+    for arg in args:
+        if isinstance(arg, tuple):
+            for a in arg:
+                yield from _descr(dd, a, key=key)
         else:
-            return descr[key]
+            if isinstance(arg, slice):
+                descr = {"index": arg, "code": arg, "label": "", "unit": ""}
+            else:
+                descr = dd[arg]
+            if key is None:
+                yield descr
+            else:
+                yield descr[key]
 
 
-def axis_(axis: AxisDef, key: Optional[str] = None):
+def axis_(*axis: AxisDef, key: str | None = None):
     r"""Return axis descriptions
 
     Parameters:
@@ -100,28 +102,32 @@ def axis_(axis: AxisDef, key: Optional[str] = None):
 
     Examples:
 
-        >>> axis_(('x','dp'), key='index')
+        >>> axis_("x", "dp", key="index")
         (0, 4)
 
         returns the indices in the standard coordinate vector
 
-        >>> dplabel = axis_('dp', key='label')
+        >>> dplabel = axis_("dp", key="label")
         >>> print(dplabel)
         $\delta$
 
         returns the coordinate label for plot annotation
 
-        >>> axis_((0,'dp'))
+        >>> axis_(0, "dp")
         ({'plane': 0, 'label': 'x', 'unit': ' [m]', 'code': 'x'},
          {'plane': 4, 'label': '$\\delta$', 'unit': '', 'code': 'dp'})
 
         returns the entire description directories
 
     """
-    return _descr(_axis_def, axis, key=key)
+    ret = tuple(_descr(_axis_def, *axis, key=key))
+    if len(ret) > 1:
+        return ret
+    else:
+        return ret[0]
 
 
-def plane_(plane: AxisDef, key: Optional[str] = None):
+def plane_(*plane: AxisDef, key: str | None = None):
     r"""Return plane descriptions
 
     Parameters:
@@ -154,16 +160,20 @@ def plane_(plane: AxisDef, key: Optional[str] = None):
 
     Examples:
 
-        >>> plane_('v', key='index')
+        >>> plane_("v", key="index")
         1
 
         returns the indices in the standard coordinate vector
 
-        >>> plane_(('x','y'))
-        ({'plane': 0, 'label': 'h', 'unit': ' [m]', 'code': 'h'},
-         {'plane': 1, 'label': 'v', 'unit': ' [m]', 'code': 'v'})
+        >>> plane_("x", "y")
+        ({'plane': 0, 'label': 'x', 'unit': ' [m]', 'code': 'h'},
+         {'plane': 1, 'label': 'y', 'unit': ' [m]', 'code': 'v'})
 
         returns the entire description directories
 
     """
-    return _descr(_plane_def, plane, key=key)
+    ret = tuple(_descr(_plane_def, *plane, key=key))
+    if len(ret) > 1:
+        return ret
+    else:
+        return ret[0]