Source code for pyaml.magnet.serialized_magnet

import numpy as np
from scipy.constants import speed_of_light

from .. import PyAMLException
from ..common import abstract
from ..common.element import Element, ElementConfigModel, __pyaml_repr__
from ..configuration import Factory
from ..control.deviceaccess import DeviceAccess
from .function_mapping import function_map
from .magnet import Magnet, MagnetConfigModel
from .model import MagnetModel

# Define the main class name for this module
PYAMLCLASS = "SerializedMagnets"




class ConfigModel(ElementConfigModel):
    function: str
    """List of magnets"""
    elements: list[str] | str
    """List of magnets"""
    model: MagnetModel | None = None
    """Object in charge of converting magnet strengths to currents"""





class ReadWriteSerializedStrengths(abstract.ReadWriteFloatScalar):
    def __init__(self, cfg: ConfigModel, elements: list[abstract.ReadWriteFloatScalar]):
        self.elements = elements
        self._cfg = cfg



    def get(self) -> float:
        return sum([elem.get() for elem in self.elements])




    def set(self, value: float):
        self.elements[0].set(value)




    def set_and_wait(self, value: float):
        raise NotImplementedError("Not implemented yet.")




    def unit(self) -> str:
        return self._cfg.model.get_strength_units()[0]




    def get_model(self) -> MagnetModel:
        return self._cfg.model




    def get_elements(self):
        return self.elements




    def set_magnet_rigidity(self, brho: np.double):
        [element.set_magnet_rigidity(brho) for element in self.elements]






class ReadWriteSerializedHardwares(ReadWriteSerializedStrengths):
    def __init__(self, cfg: ConfigModel, elements: list[abstract.ReadWriteFloatScalar]):
        super().__init__(cfg, elements)



    def unit(self) -> str:
        return self._cfg.model.get_hardware_units()[0]




    def set_magnet_rigidity(self, brho: np.double):
        [element.set_magnet_rigidity(brho) for element in self.elements]






class SerializedMagnets(Element):
    """
    Class managing serialized magnets: a set of magnet with the same set point.
    The set point is usually managed by only one power supply but it can be covered by several ones.
    If several power supplies


    Parameters
    ----------
    cfg : ConfigModel
        Configuration object TODO: to describe

    Raises
    ------
    pyaml.PyAMLException
        In case of wrong initialization
    """

    def __init__(self, cfg: ConfigModel, peer=None):
        super().__init__(cfg.name)
        self._cfg = cfg
        self.model = cfg.model
        self.polynom = None
        self.__strengths = None
        self.__hardwares = None
        self.__virtuals: list[Magnet] = []
        self.__elements = cfg.elements if isinstance(cfg.elements, list) else [cfg.elements]
        self.model.set_number_of_magnets(len(self.__elements))
        if peer is None:
            # Configuration part
            self.polynom = function_map[self._cfg.function].polynom
            if self._cfg.function not in function_map:
                raise PyAMLException(self._cfg.function + " not implemented for serialized magnet")
            for element in self.__elements:
                # Check mapping validity
                # Create the virtual magnet for the corresponding magnet
                vm = self.__create_virtual_magnet(element)
                self.__virtuals.append(vm)
                # Register the virtual element in the factory to have a coherent factory and improve error reporting
                Factory.register_element(vm)
        else:
            # Attach
            self._peer = peer

    def __create_virtual_magnet(self, name: str) -> Magnet:
        args = {"name": name, "model": self.model}
        virtual: Magnet = function_map[self._cfg.function](MagnetConfigModel(**args))
        virtual.set_model_name(self.get_name())
        return virtual



    def get_nb_magnets(self) -> int:
        return len(self.__elements)




    def get_magnets(self) -> list[Magnet]:
        return self.__virtuals




    def attach(
        self,
        peer,
        strengths: list[abstract.ReadWriteFloatScalar],
        hardwares: list[abstract.ReadWriteFloatScalar],
    ) -> list[Magnet]:
        l = []
        n_ser_mag = SerializedMagnets(self._cfg, peer)
        n_ser_mag.__strengths = ReadWriteSerializedStrengths(self._cfg, strengths)
        n_ser_mag.__hardwares = ReadWriteSerializedHardwares(self._cfg, hardwares)
        l.append(n_ser_mag)
        # Construct a single magnet for each magnet.
        sub_magnets: list[Magnet] = []
        for idx, _ in enumerate(self.__elements):
            strength = strengths[idx]
            hardware = hardwares[idx] if self.model.has_hardware() else None
            sub_magnets.append(self.__virtuals[idx].attach(peer, strength, hardware))
        n_ser_mag.__virtuals.extend(sub_magnets)
        l.extend(sub_magnets)
        return l


    @property
    def strength(self) -> abstract.ReadWriteFloatScalar:
        """
        Gives access to the strengths of those magnets in physics unit
        """
        self.check_peer()
        if self.__strengths is None:
            raise PyAMLException(f"{str(self)} has no model that supports physics units")
        return self.__strengths

    @property
    def hardware(self) -> abstract.ReadWriteFloatScalar:
        """
        Gives access to the strengths of this those magnets in hardware unit when possible
        """
        self.check_peer()
        if self.__hardwares is None:
            raise PyAMLException(f"{str(self)} has no model that supports hardware units")
        return self.__hardwares



    def set_energy(self, energy: float):
        brho = energy / speed_of_light
        if self.model is not None:
            self.model.set_magnet_rigidity(brho)
        if self.__hardwares is not None:
            self.__hardwares.set_magnet_rigidity(brho)
        if self.__strengths is not None:
            self.__strengths.set_magnet_rigidity(brho)


    def __repr__(self):
        return __pyaml_repr__(self)



    def get_devices(self) -> list[DeviceAccess]:
        return self.model.get_devices()