__all__ = ["LocalEmittanceObservable"]
from collections.abc import Callable
from functools import partial
from typing import ClassVar
import numpy as np
from ..lattice import AxisDef, plane_, axis_
from ..lattice import Refpts
from .observables import ElementObservable, Need
# noinspection PyProtectedMember
from .observables import _all_rows, _record_access, _fun_access, _subscript
def _sigma4(_, index, data, **kwargs):
"""Betatron beam size: square root of diagonal terms of r44."""
idx = index[1]
return np.sqrt(data.r44[_all_rows((idx, idx))])
def _sigma6(_, index, data, **kwargs):
"""Beam size: square root of diagonal terms of r66."""
idx = index[1]
return np.sqrt(data.r66[_all_rows((idx, idx))])
_opdata = {"sigma4": _sigma4, "sigma6": _sigma6}
[docs]
class LocalEmittanceObservable(ElementObservable):
"""Observe emittance-related parameters."""
# Class attributes
_pinfo: ClassVar[dict] = {
"r66": (r"$R66_{{{plane}}}$", "R", _subscript),
"r44": (r"$R44_{{{plane}}}$", "R", _subscript),
"emitXY": (
r"$\epsilon_{{{plane}}}$",
"Emittance [m]",
partial(plane_, key="label"),
),
"emitXYZ": (
r"$\epsilon_{{{plane}}}$",
"Emittance [m]",
partial(plane_, key="label"),
),
"m66": (r"$\mathrm{{T}}_{{{plane}}}$", "T [m]", _subscript),
"orbit6": (r"${{{plane}}}_{{co}}$", "closed orbit", partial(axis_, key="code")),
"sigma4": (r"$\sigma_{{{plane}}}$", "sigma", partial(axis_, key="label")),
"sigma6": (r"$\sigma_{{{plane}}}$", "sigma", partial(axis_, key="label")),
}
def __init__(
self,
refpts: Refpts,
param: str | Callable,
*,
plane: AxisDef = Ellipsis,
axis: AxisDef = Ellipsis,
name: str | None = None,
summary: bool = False,
label: str | None = None,
**eval_kw,
):
# noinspection PyUnresolvedReferences
r"""Args:
refpts: Observation points.
See ":ref:`Selecting elements in a lattice <refpts>`"
param: :ref:`Emittance parameter name <localemit_param>`
or :ref:`user-defined evaluation function <localemit_eval>`
axis: Index in the parameter array. If :py:obj:`Ellipsis`, the
whole array is specified
plane: Index in the parameter array for *emitXY* and *emitXYZ*
parameters. If :py:obj:`Ellipsis`, the whole array is specified
name: Observable name. If :py:obj:`None`, an explicit
name will be generated.
Keyword Args:
summary: Set to :py:obj:`True` if the user-defined
evaluation function returns a single item (see below)
postfun: Post-processing function. It can be any numpy ufunc or a
function name in {"real", "imag", "abs", "angle", "log", "exp", "sqrt"}.
statfun: Statistics post-processing function. it can be a numpy
function or a function name in {"mean", "std", "var", "min", "max"}.
Example: :pycode:`statfun=numpy.mean`.
target: Target value for a constraint. If :py:obj:`None`
(default), the residual will always be zero.
weight: Weight factor: the residual is
:pycode:`((value-target)/weight)**2`
bounds: Tuple of lower and upper bounds. The parameter
is constrained in the interval
[*target*\ +\ *low_bound* *target*\ +\ *up_bound*]
.. rubric:: Evaluation keywords
These values must be provided to the :py:meth:`~.ObservableList.evaluate`
method. Default values may be given at instantiation.
* **ring**: Lattice description,
* **dp**: Momentum deviation. Defaults to :py:obj:`None`,
* **dct**: Path lengthening. Defaults to :py:obj:`None`,
* **df**: Deviation from the nominal RF frequency.
Defaults to :py:obj:`None`,
* **orbit**: Initial orbit. Avoids looking for the closed orbit if it is
already known,
.. rubric:: Shape of the value
If the requested attribute has shape :pycode:`shp`, then
*value* has shape :pycode:`(nrefs,) + shp` with :pycode:`nrefs` number of
reference points: a :pycode:`nrefs` dimension is prepended to the shape of
the attribute, **even if nrefs == 1**. The *target*, *weight* and *bounds*
inputs must be broadcastable to the shape of *value*. For instance, a *target*
with shape :pycode:`shp` will automatically broadcast and apply to all
reference points.
.. _localemit_param:
.. rubric:: Emittance parameter names
In addition to :py:func:`.ohmi_envelope` parameter names,
LocalEmittanceObservable adds the *sigma* parameter:
=========== ===================================================
**r66** (6, 6) equilibrium envelope matrix :math:`\Sigma`
**r44** (4, 4) betatron emittance matrix (dp = 0)
**m66** (6, 6) transfer matrix from the start of the ring
**orbit6** (6,) closed orbit
**emitXY** (2,) betatron emittance projected on xxp and yyp
**emitXYZ** (3,) 6x6 emittance projected on xxp, yyp, ldp
**sigma6** (6,) standard deviation of the beam (square root
of the diagonal terms of *r66*)
**sigma4** (6,) standard deviation of the monochromatic beam (square root
of the diagonal terms of *r44*)
=========== ===================================================
.. _localemit_eval:
.. rubric:: User-defined evaluation function
The observable value is computed as:
:pycode:`value = fun(emit, **eval_kw)`
- *emit* is the output of :py:func:`.ohmi_envelope`, evaluated at the
*refpts* of the observable,
- *eval_kw* are the keyword arguments provided to the observable constructor,
to the constructor of the enclosing :py:class:`.ObservableList` and to the
:py:meth:`~.ObservableList.evaluate` method,
- *value* is the value of the Observable and must have one line per
refpoint. Alternatively, it may be a single line, but then the
*summary* keyword must be set to :py:obj:`True`,
Examples:
>>> obs = LocalEmittanceObservable(at.Monitor, "sigma6")
Observe the beam standard deviation in all 6 axes at Monitor locations
>>> obs = LocalEmittanceObservable(
... at.Quadrupole, "sigma6", axis="x", statfun=np.max
... )
Observe the maximum horizontal beam size in Quadrupoles
"""
needs = {Need.LOCALEMIT}
descr = plane_(plane) if param in {"emitXY", "emitXYZ"} else axis_(axis)
name = self._set_name(name, param, descr["code"])
index = _all_rows(descr["index"])
if callable(param):
if summary:
fun = partial(_fun_access, param, descr["index"])
else:
fun = partial(_fun_access, param, index)
else:
fun = partial(_opdata.get(param, _record_access), param, index)
super().__init__(
fun,
refpts,
needs=needs,
name=name,
summary=summary,
label=self._pl_lab(param, descr["index"]),
axis_label=self._ax_lab(param, descr["index"]),
**eval_kw,
)
if label:
self.label = label
