"""Matching of lattice parameters.
The matching acts on a :py:class:`.VariableList` to statisfy the constraints
expressed in an :py:class:`.ObservableList`.
Examples of use of such classes are available in:
* :doc:`/p/notebooks/variables`
* :doc:`/p/notebooks/observables`
Examples of matching are given in :doc:`/p/notebooks/matching`.
"""
from __future__ import annotations
__all__ = ["match", "ring_match"]
import numpy as np
from scipy.optimize import least_squares
from .observablelist import ObservableList
from ..lattice import Lattice, VariableList
from ..lattice.lattice_variables import ElementVariable
[docs]
def match(
variables: VariableList,
constraints: ObservableList,
*,
method: str | None = None,
verbose: int = 2,
max_nfev: int = 1000,
optim_kw: dict | None = None,
**eval_kw,
) -> np.ndarray:
r"""Observable matching.
Minimisation is performed by the :py:func:`~scipy.optimize.least_squares`
function.
Args:
variables: Variable parameters
constraints: Constraints to fulfill
method: Minimisation algorithm (see
:py:func:`~scipy.optimize.least_squares`). If :py:obj:`None`, use
'lm' for unbounded problems, 'trf' otherwise,
verbose: Level of verbosity,
max_nfev: Maximum number of function evaluation,
optim_kw: Dictionary of optimiser keyword arguments sent to
:py:func:`~scipy.optimize.least_squares`,
Keyword Args:
\*\*eval_kw: Evaluation keywords provided to the
:py:meth:`ObservableList.evaluate` method. For instance *"ring"* (for
lattice-dependent observables), *"dp"*, *"dct"*, *"orbit"*, *"twiss_in"*,
*"r_in"*… Default values are taken from the ObservableList.
Returns:
Matching result: solution for the variable values.
Raises:
RuntimeError: If optimisation fails
"""
def fun(vals) -> np.ndarray:
"""Evaluation function for the minimiser."""
variables.set(vals, **eval_kw)
constraints.evaluate(**eval_kw)
return constraints.get_flat_weighted_deviations(err=1.0e6)
if optim_kw is None:
optim_kw = {}
vini = variables.get(initial=True, check_bounds=True, **eval_kw)
bounds = np.array([var.bounds for var in variables])
x_scale = np.array([var.delta for var in variables])
constraints.evaluate(initial=True, **eval_kw)
constraints.check()
ntargets = constraints.flat_values.size
if method is None:
if np.any(np.isfinite(bounds)) or ntargets < len(variables):
method = "trf"
else:
method = "lm"
if verbose >= 1:
print(
f"\n{ntargets} constraints, {len(variables)} variables,"
f" using method {method}\n"
)
opt = least_squares(
fun,
vini,
bounds=(bounds[:, 0], bounds[:, 1]),
x_scale=x_scale,
verbose=verbose,
max_nfev=max_nfev,
method=method,
**optim_kw,
)
if verbose >= 1:
print("\nConstraints:")
print(constraints)
print("\nVariables:")
print(variables)
return opt.x
[docs]
def ring_match(
ring: Lattice,
variables: VariableList,
constraints: ObservableList,
*,
copy: bool = False,
method: str | None = None,
verbose: int = 2,
max_nfev: int = 1000,
**kwargs,
) -> Lattice | None:
r"""Match constraints by varying variables.
Minimisation is performed by the :py:func:`~scipy.optimize.least_squares`
function.
Args:
ring: Lattice description
variables: Variable parameters
constraints: Constraints to fulfill
copy: If :py:obj:`True`, return a modified copy of *ring*, otherwise
perform the match in-line
method: Minimisation algorithm (see
:py:func:`~scipy.optimize.least_squares`). If :py:obj:`None`, use
'lm' for unbounded problems, 'trf' otherwise,
verbose: Level of verbosity,
max_nfev: Maximum number of function evaluation,
Keyword Args:
dp (float | None): Momentum deviation. Default taken from the ObservableList,
dct (float | None): Path lengthening. Default taken from the ObservableList,
df (float | None): Deviation from the nominal RF frequency. Default taken from
the ObservableList,
orbit (Orbit | None): Initial orbit. Avoids looking for the closed orbit if
it is already known. Used for :py:class:`.MatrixObservable` and
:py:class:`.LocalOpticsObservable`. Default taken from the ObservableList,
twiss_in: Initial conditions for transfer line optics
See :py:func:`.get_optics`. Used for :py:class:`.LocalOpticsObservable`.
Default taken from the ObservableList,
r_in (Orbit | None): Initial trajectory, used for
:py:class:`.TrajectoryObservable`. Default taken from the ObservableList,
\*\*kwargs: The other keyword arguments sent to,
:py:func:`~scipy.optimize.least_squares`.
Returns:
Modified Lattice if copy=True, None otherwise
Raises:
TypeError: If copy=True and ElementVariable is present
RuntimeError: If optimisation fails
Examples:
See :doc:`/p/notebooks/matching`
"""
# Separate the keywords for evaluation
eval_kw = {}
for key in ["dp", "dct", "df", "orbit", "twiss_in", "r_in"]:
v = kwargs.pop(key, None)
if v is not None:
eval_kw[key] = v
if copy:
# Check that there is no ElementVariable
for var in variables:
if isinstance(var, ElementVariable):
msg = "When 'copy' is True, no ElementVariable is accepted."
raise TypeError(msg)
newring = ring.deepcopy()
match(
variables,
constraints,
ring=newring,
method=method,
verbose=verbose,
max_nfev=max_nfev,
optim_kw=kwargs,
**eval_kw,
)
return newring
else:
match(
variables,
constraints,
ring=ring,
method=method,
verbose=verbose,
max_nfev=max_nfev,
optim_kw=kwargs,
**eval_kw,
)
return None
Lattice.match = ring_match
