import numpy as np
from pydantic import BaseModel, ConfigDict
from ..common.element import __pyaml_repr__
from ..common.exception import PyAMLException
from ..control.deviceaccess import DeviceAccess
from .curve import Curve
from .matrix import Matrix
from .model import MagnetModel
# Define the main class name for this module
PYAMLCLASS = "LinearCFMagnetModel"
[docs]
class ConfigModel(BaseModel):
"""
Configuration model for linear combined function magnet model
Parameters
----------
multipoles : list[str]
List of supported functions: A0, B0, A1, B1, etc (i.e. [B0, A1, B2])
curves : list[Curve]
Excitation curves, 1 curve per function
calibration_factors : list[float], optional
Correction factor applied to curves, 1 factor per function.
Default: ones
calibration_offsets : list[float], optional
Correction offset applied to curves, 1 offset per function.
Default: zeros
pseudo_factors : list[float], optional
Factors applied to 'pseudo currents', 1 factor per function.
Default: ones
pseudo_offsets : list[float], optional
Offsets applied to 'pseudo currents', 1 offset per function.
Default: zeros
powerconverters : list[DeviceAccess]
List of power converter devices to apply currents (can be different
from number of functions)
matrix : Matrix, optional
n x m matrix (n rows for n functions, m columns for m currents)
to handle multipoles separation. Default: Identity
units : list[str]
List of strength units (i.e. ['rad', 'm-1', 'm-2'])
"""
model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
multipoles: list[str]
curves: list[Curve]
calibration_factors: list[float] = None
calibration_offsets: list[float] = None
pseudo_factors: list[float] = None
pseudo_offsets: list[float] = None
powerconverters: list[DeviceAccess | None]
matrix: Matrix = None
units: list[str]
[docs]
class LinearCFMagnetModel(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.__nbFunction: int = len(cfg.multipoles)
self.__nbPS: int = len(cfg.powerconverters)
if cfg.calibration_factors is None:
self.__calibration_factors = np.ones(self.__nbFunction)
else:
self.__calibration_factors = cfg.calibration_factors
if cfg.calibration_offsets is None:
self.__calibration_offsets = np.zeros(self.__nbFunction)
else:
self.__calibration_offsets = cfg.calibration_offsets
if cfg.pseudo_factors is None:
self.__pf = np.ones(self.__nbFunction)
else:
self.__pf = cfg.pseudo_factors
if cfg.pseudo_offsets is None:
self.__po = np.zeros(self.__nbFunction)
else:
self.__po = cfg.pseudo_factors
self.__check_len(self.__calibration_factors, "calibration_factors", self.__nbFunction)
self.__check_len(self.__calibration_offsets, "calibration_offsets", self.__nbFunction)
self.__check_len(self.__pf, "pseudo_factors", self.__nbFunction)
self.__check_len(self.__po, "pseudo_offsets", self.__nbFunction)
self.__check_len(cfg.units, "units", self.__nbFunction)
self.__check_len(cfg.curves, "curves", self.__nbFunction)
if cfg.matrix is None:
self.__matrix = np.identity(self.__nbFunction)
else:
self.__matrix = cfg.matrix.get_matrix()
_s = np.shape(self.__matrix)
if len(_s) != 2 or _s[0] != self.__nbFunction or _s[1] != self.__nbPS:
raise PyAMLException(
f"matrix wrong dimension ({self.__nbFunction}x{self.__nbPS} expected but got {_s[0]}x{_s[1]})"
)
self.__curves = []
self.__rcurves = []
# Apply factor and offset
for idx, c in enumerate(cfg.curves):
self.__curves.append(c.get_curve())
self.__curves[idx][:, 1] *= self.__calibration_factors[idx]
self.__curves[idx][:, 1] += self.__calibration_offsets[idx]
self.__rcurves.append(Curve.inverse(self.__curves[idx]))
# Compute pseudo inverse
self.__inv = np.linalg.pinv(self.__matrix)
def __check_len(self, obj, name, expected_len):
lgth = len(obj)
if lgth != expected_len:
raise PyAMLException(
f"{name} does not have the expected number of items ({expected_len} items expected but got {lgth})"
)
[docs]
def compute_hardware_values(self, strengths: np.array) -> np.array:
_pI = np.zeros(self.__nbFunction)
for idx, c in enumerate(self.__rcurves):
_pI[idx] = self.__pf[idx] * np.interp(strengths[idx] * self._brho, c[:, 0], c[:, 1]) + self.__po[idx]
_currents = np.matmul(self.__inv, _pI)
return _currents
[docs]
def compute_strengths(self, currents: np.array) -> np.array:
_strength = np.zeros(self.__nbFunction)
_pI = np.matmul(self.__matrix, currents)
for idx, c in enumerate(self.__curves):
_strength[idx] = np.interp((_pI[idx] - self.__po[idx]) / self.__pf[idx], c[:, 0], c[:, 1]) / self._brho
return _strength
[docs]
def get_strength_units(self) -> list[str]:
return self._cfg.units
[docs]
def get_hardware_units(self) -> list[str]:
return np.array([p.unit() for p in self._cfg.powerconverters])
[docs]
def get_devices(self) -> list[DeviceAccess]:
return self._cfg.powerconverters
[docs]
def set_magnet_rigidity(self, brho: np.double):
self._brho = brho
[docs]
def has_hardware(self) -> bool:
return (self.__nbPS == self.__nbFunction) and np.allclose(self.__matrix, np.eye(self.__nbFunction))
def __repr__(self):
return __pyaml_repr__(self)