import numpy as np
from pydantic import BaseModel, ConfigDict
from ..common.element import __pyaml_repr__
from ..common.exception import PyAMLException
from ..configuration.curve import Curve
from ..configuration.inline_curve import ConfigModel as InlineCurveModel
from ..configuration.inline_curve import InlineCurve
from ..configuration.matrix import Matrix
from ..control.deviceaccess import DeviceAccess
from .linear_model import ConfigModel as LinearConfigModel
from .linear_model import LinearMagnetModel
from .model import MagnetModel
# Define the main class name for this module
PYAMLCLASS = "LinearSerializedMagnetModel"
[docs]
class ConfigModel(BaseModel):
"""
Configuration model for linear serialized magnet model
Parameters
----------
curves : Curve or list[Curve]
Excitation curves, 1 curve for all or 1 curve per magnet
calibration_factors : float or list[float], optional
Correction factor applied to curves, 1 factor for all or 1 factor per magnet.
Default: ones
calibration_offsets : float or list[float], optional
Correction offset applied to curves, 1 offset for all or 1 offset per magnet.
Default: zeros
crosstalk : float or list[float], optional
Crosstalk factors. Default: 1.0
powerconverter : DeviceAccess
The hardware can be a single power supply or a list of power supplies.
If a list is provided, the same value will be affected to all of them
unit : str
Strength unit: rad, m-1, m-2
"""
model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
curves: Curve | list[Curve]
calibration_factors: float | list[float] = None
calibration_offsets: float | list[float] = None
crosstalk: float | list[float] = 1.0
powerconverter: DeviceAccess
unit: str
def _get_length(elem) -> int:
if elem is None:
return 0
if isinstance(elem, list):
return len(elem)
else:
return 1
def _get_max_length(*args, **kwargs) -> int:
max_args = max([_get_length(elem) for elem in args]) if args else 0
max_kwargs = max([_get_length(elem) for elem in kwargs.values()]) if kwargs else 0
return max(max_args, max_kwargs)
def _to_list_of_length(elem, length: int) -> list:
if isinstance(elem, list):
return elem
else:
return [elem] * length
def _check_len(obj, name, expected_length):
length = len(obj)
if length != expected_length:
raise PyAMLException(
f"{name} does not have the expected number of items ({expected_length} items expected but got {length})"
)
[docs]
class LinearSerializedMagnetModel(MagnetModel):
"""
Class providing a simple linear model for combined function magnets. A matrix
can handle separation of multipoles. A pseudo current is a linear combination
of power supply currents associated to a single function.
"""
def __init__(self, cfg: ConfigModel):
self._cfg = cfg
self.__brho = np.nan
# Check config
self.__nbMagnets: int = _get_max_length(cfg.curves, cfg.calibration_factors, cfg.calibration_offsets, cfg.crosstalk)
self.__calibration_factors = np.ones(self.__nbMagnets)
self.__calibration_offsets = np.ones(self.__nbMagnets)
self.__crosstalk = np.ones(self.__nbMagnets)
self.__curves = _to_list_of_length(self._cfg.curves, self.__nbMagnets)
self.__sub_models: list[LinearMagnetModel] = []
def __initialize(self):
if self._cfg.calibration_factors is None:
self.__calibration_factors = np.ones(self.__nbMagnets)
else:
self.__calibration_factors = _to_list_of_length(self._cfg.calibration_factors, self.__nbMagnets)
if self._cfg.calibration_offsets is None:
self.__calibration_offsets = np.zeros(self.__nbMagnets)
else:
self.__calibration_offsets = _to_list_of_length(self._cfg.calibration_offsets, self.__nbMagnets)
if self._cfg.crosstalk is None:
self.__crosstalk = np.zeros(self.__nbMagnets)
else:
self.__crosstalk = _to_list_of_length(self._cfg.crosstalk, self.__nbMagnets)
self.__curves = _to_list_of_length(self._cfg.curves, self.__nbMagnets)
if isinstance(self._cfg.curves, list):
self.__curves = self._cfg.curves
else:
self.__curves: list[Curve] = []
for _ in range(self.__nbMagnets):
curve = InlineCurve(InlineCurveModel(mat=self._cfg.curves.get_curve()))
self.__curves.append(curve)
_check_len(self.__calibration_factors, "calibration_factors", self.__nbMagnets)
_check_len(self.__calibration_offsets, "calibration_offsets", self.__nbMagnets)
_check_len(self.__crosstalk, "crosstalk", self.__nbMagnets)
_check_len(self.__curves, "curves", self.__nbMagnets)
self.__sub_models: list[LinearMagnetModel] = []
for magnet_idx in range(self.__nbMagnets):
sub_model = LinearConfigModel(
curve=self.__curves[magnet_idx],
calibration_factor=self.__calibration_factors[magnet_idx],
calibration_offset=self.__calibration_offsets[magnet_idx],
crosstalk=self.__crosstalk[magnet_idx],
powerconverter=self._cfg.powerconverter,
unit=self._cfg.unit,
)
self.__sub_models.append(LinearMagnetModel(sub_model))
[docs]
def set_number_of_magnets(self, nb_magnets: int):
self.__nbMagnets = nb_magnets
self.__initialize()
[docs]
def get_sub_model(self, index: int) -> LinearMagnetModel:
return self.__sub_models[index]
[docs]
def compute_hardware_values(self, strengths: np.array) -> np.array:
currents = [model.compute_hardware_values([s])[0] for s, model in zip(strengths, self.__sub_models, strict=True)]
return np.array([np.mean(currents)])
[docs]
def compute_strengths(self, currents: np.array) -> np.array:
current = currents[0]
return np.array([model.compute_strengths([current])[0] for model in self.__sub_models])
[docs]
def get_strength_units(self) -> list[str]:
return [self._cfg.unit] * self.__nbMagnets
[docs]
def get_hardware_units(self) -> list[str]:
return [self.__sub_models[0].get_hardware_units()[0]]
[docs]
def get_devices(self) -> list[DeviceAccess]:
return [self._cfg.powerconverter]
[docs]
def set_magnet_rigidity(self, brho: np.double):
self.__brho = brho
[model.set_magnet_rigidity(brho) for model in self.__sub_models]
[docs]
def get_magnet_rigidity(self) -> np.double:
return self.__brho
def __repr__(self):
return __pyaml_repr__(self)
