merged from response_matrices

Author: lfarv

pyat/at/latticetools/observables.py

@@ -66,6 +66,7 @@
 from typing import Tuple
 
 import numpy as np
+import numpy.typing as npt
 
 from ..lattice import AtError, AxisDef, axis_, plane_
 from ..lattice import Lattice, Refpts, End
@@ -204,8 +205,8 @@ def __init__(
         fun: Callable,
         *args,
         name: str | None = None,
-        target=None,
-        weight=1.0,
+        target: npt.ArrayLike | None = None,
+        weight: npt.ArrayLike = 1.0,
         bounds=(0.0, 0.0),
         needs: Set[Need] | None = None,
         **kwargs,
@@ -256,12 +257,12 @@ def __init__(
         self.fun: Callable = fun  #: Evaluation function
         self.needs: Set[Need] = needs or set()  #: Set of requirements
         self.name: str = name  #: Observable name
-        self.target = target  #: Target value
-        self.w = weight
+        self.target: npt.ArrayLike | None = target  #: Target value
+        self.w: npt.ArrayLike = weight
         self.lbound, self.ubound = bounds
-        self.initial = None
-        self._value = None
-        self._shape = None
+        self.initial: npt.NDArray[float] | None = None
+        self._value: npt.NDArray[float] | Exception | None = None
+        self._shape: tuple[int, ...] | None = None
         self.args = args
         self.kwargs = kwargs
 
@@ -299,23 +300,25 @@ def _all_lines(self):
         vnow = self._value
         if vnow is None or isinstance(vnow, Exception):
             deviation = None
-        else:
-            deviation = self.deviation
-        if self.target is None:
             vmin = None
             vmax = None
         else:
-            target = np.broadcast_to(self.target, vnow.shape)
-            vmin = target + self.lbound
-            vmax = target + self.ubound
+            deviation = self.deviation
+            if self.target is None:
+                vmin = None
+                vmax = None
+            else:
+                target = np.broadcast_to(self.target, vnow.shape)  # type: ignore
+                vmin = target + self.lbound
+                vmax = target + self.ubound
         values = self._line("", self.initial, vnow, vmin, vmax, deviation)
         return "\n".join((self.name, values))
 
     def _setup(self, ring: Lattice):
         """Setup function called when the observable is added to a list."""
         pass
 
-    def evaluate(self, *data, initial: bool = False):
+    def evaluate(self, *data, initial: bool = False) -> npt.NDArray[float] | Exception:
         """Compute and store the value of the observable.
 
         The direct evaluation of a single :py:class:`Observable` is normally
@@ -327,6 +330,9 @@ def evaluate(self, *data, initial: bool = False):
               sent to the evaluation function
             initial:    It :py:obj:`None`, store the result as the initial
               value
+
+        Returns:
+            value:      The value of the observable or the error in evaluation
         """
         for d in data:
             if isinstance(d, Exception):
@@ -354,30 +360,34 @@ def check(self) -> bool:
         return self.value is not None
 
     @staticmethod
-    def check_value(value):
+    def check_value(value: npt.NDArray[float] | Exception) -> npt.NDArray[float]:
         if isinstance(value, Exception):
             raise type(value)(value.args[0]) from value
         return value
 
     @property
-    def value(self):
+    def value(self) -> npt.NDArray[float]:
         """Value of the observable."""
         return self.check_value(self._value)
 
     @property
-    def weight(self):
+    def weight(self) -> npt.NDArray[float]:
         """Observable weight."""
-        return np.broadcast_to(self.w, self._value.shape)
+        return np.broadcast_to(self.w, self._value.shape)  # type: ignore
+
+    @weight.setter
+    def weight(self, w: npt.ArrayLike):
+        self.w = w
 
     @property
-    def weighted_value(self):
+    def weighted_value(self) -> npt.NDArray[float]:
         """Weighted value of the Observable, computed as
         :pycode:`weighted_value = value/weight`.
         """
         return self.value / self.w
 
     @property
-    def deviation(self):
+    def deviation(self) -> npt.NDArray[float]:
         """Deviation from target value, computed as
         :pycode:`deviation = value-target`.
         """
@@ -395,12 +405,12 @@ def deviation(self):
         return deviation
 
     @property
-    def weighted_deviation(self):
+    def weighted_deviation(self) -> npt.NDArray[float]:
         """:pycode:`weighted_deviation = (value-target)/weight`."""
         return self.deviation / self.w
 
     @property
-    def residual(self):
+    def residual(self) -> npt.NDArray[float]:
         """residual, computed as :pycode:`residual = ((value-target)/weight)**2`."""
         return (self.deviation / self.w) ** 2
 
@@ -527,7 +537,7 @@ def __init__(
         self._excluded = None
         self._locations = [""]
 
-    def check(self):
+    def check(self) -> bool:
         ok = super().check()
         shp = self._shape
         if ok and shp and shp[0] <= 0:
@@ -552,7 +562,7 @@ def _all_lines(self):
                     vmin = repeat(None)
                     vmax = repeat(None)
                 else:
-                    target = np.broadcast_to(self.target, vnow.shape)
+                    target = np.broadcast_to(self.target, vnow.shape)  # type: ignore
                     vmin = target + self.lbound
                     vmax = target + self.ubound
             vini = self.initial
@@ -655,8 +665,8 @@ def __init__(
 
             Observe the horizontal closed orbit at monitor locations
         """
-        name = self._set_name(name, "orbit", axis_(axis, "code"))
-        fun = _ArrayAccess(axis_(axis, "index"))
+        name = self._set_name(name, "orbit", axis_(axis, key="code"))
+        fun = _ArrayAccess(axis_(axis, key="index"))
         needs = {Need.ORBIT}
         super().__init__(fun, refpts, needs=needs, name=name, **kwargs)
 
@@ -705,8 +715,8 @@ def __init__(
             Observe the transfer matrix from origin to monitor locations and
             extract T[0,1]
         """
-        name = self._set_name(name, "matrix", axis_(axis, "code"))
-        fun = _ArrayAccess(axis_(axis, "index"))
+        name = self._set_name(name, "matrix", axis_(axis, key="code"))
+        fun = _ArrayAccess(axis_(axis, key="index"))
         needs = {Need.MATRIX}
         super().__init__(fun, refpts, needs=needs, name=name, **kwargs)
 
@@ -739,12 +749,12 @@ def __init__(
         shape of *value*.
         """
         needs = {Need.GLOBALOPTICS}
-        name = self._set_name(name, param, plane_(plane, "code"))
+        name = self._set_name(name, param, plane_(plane, key="code"))
         if callable(param):
             fun = param
             needs.add(Need.CHROMATICITY)
         else:
-            fun = partial(_record_access, param, plane_(plane, "index"))
+            fun = partial(_record_access, param, plane_(plane, key="index"))
             if param == "chromaticity":
                 needs.add(Need.CHROMATICITY)
         super().__init__(fun, needs=needs, name=name, **kwargs)
@@ -881,8 +891,8 @@ def __init__(
             ax_ = plane_
 
         needs = {Need.LOCALOPTICS}
-        name = self._set_name(name, param, ax_(plane, "code"))
-        index = _all_rows(ax_(plane, "index"))
+        name = self._set_name(name, param, ax_(plane, key="code"))
+        index = _all_rows(ax_(plane, key="index"))
         if callable(param):
             fun = param
         else:
@@ -922,7 +932,9 @@ def __init__(
 
         Example:
 
-            >>> obs = LatticeObservable(at.Sextupole, "KickAngle", index=0, statfun=np.sum)
+            >>> obs = LatticeObservable(
+            ...     at.Sextupole, "KickAngle", index=0, statfun=np.sum
+            ... )
 
             Observe the sum of horizontal kicks in Sextupoles
         """
@@ -967,8 +979,8 @@ def __init__(
         The *target*, *weight* and *bounds* inputs must be broadcastable to the
         shape of *value*.
         """
-        name = self._set_name(name, "trajectory", axis_(axis, "code"))
-        fun = _ArrayAccess(axis_(axis, "index"))
+        name = self._set_name(name, "trajectory", axis_(axis, key="code"))
+        fun = _ArrayAccess(axis_(axis, key="index"))
         needs = {Need.TRAJECTORY}
         super().__init__(fun, refpts, needs=needs, name=name, **kwargs)
 
@@ -1020,11 +1032,11 @@ def __init__(
 
             Observe the horizontal emittance
         """
-        name = self._set_name(name, param, plane_(plane, "code"))
+        name = self._set_name(name, param, plane_(plane, key="code"))
         if callable(param):
             fun = param
         else:
-            fun = partial(_record_access, param, plane_(plane, "index"))
+            fun = partial(_record_access, param, plane_(plane, key="index"))
         needs = {Need.EMITTANCE}
         super().__init__(fun, needs=needs, name=name, **kwargs)
 
@@ -1087,10 +1099,10 @@ def GlobalOpticsObservable(
     """
     if param == "tune" and use_integer:
         # noinspection PyProtectedMember
-        name = ElementObservable._set_name(name, param, plane_(plane, "code"))
+        name = ElementObservable._set_name(name, param, plane_(plane, key="code"))
         return LocalOpticsObservable(
             End,
-            _Tune(plane_(plane, "index")),
+            _Tune(plane_(plane, key="index")),
             name=name,
             summary=True,
             all_points=True,