Attachment of orbit correction to simulator/control system and generic ResponseMatrix pyAML object

[kparasch]

No description provided.

[kparasch]

@JeanLucPons do you think it is ok to pass the element_holder like this?

[JeanLucPons]

I did like you for the attachement process.
I put here my orbit.py which is working.
I let you check that I didn't make wrong assumption concerning hv corrector order.
It also use the fileloader patch for the json file.

Usage:

from pyaml.accelerator import Accelerator
sr = Accelerator.load("tests/config/EBSOrbit.yaml",use_fast_loader=False)
orbit = sr.live.get_orbit_tuning("ORBIT")
orbit.correct()

I added to EBSOrbit.yaml:

- type: pyaml.tuning_tools.orbit
  name: ORBIT
  bpm_array_name: BPM
  hcorr_array_name: HCorr
  vcorr_array_name: VCorr
  singular_values: 162
  response_matrix_file: ideal_orm_disp.json

import logging
from pathlib import Path
from typing import Literal, Optional
from typing import TYPE_CHECKING

from ..common.element import Element,ElementConfigModel
from ..external.pySC.pySC import ResponseMatrix
from ..external.pySC.pySC.apps import orbit_correction
from ..external.pySC_interface import pySCInterface
from ..configuration import get_root_folder
from ..arrays.magnet_array import MagnetArray

if TYPE_CHECKING:
    from ..common.element_holder import ElementHolder

try:
    from typing import Self  # Python 3.11+
except ImportError:
    from typing_extensions import Self  # Python 3.10 and earlier

logger = logging.getLogger(__name__)
logging.getLogger("pyaml.external.pySC").setLevel(logging.WARNING)

PYAMLCLASS = "Orbit"


class ConfigModel(ElementConfigModel):
    """
    Orbit tuning tool configuration

    Parameters
    ----------

    bpm_array_name: str
        Name of BPM array used to get the orbit
    hcorr_array_name: str
        Name of horizontal steerer array used to correct the orbit
    vcorr_array_name: str
        Name of vertical steerer array used to correct the orbit
    singular_values: int
        Number of singular value used for orbit correction
    response_matrix_file: str
        Orbit response matrix file to load    
    """

    bpm_array_name: str
    hcorr_array_name: str
    vcorr_array_name: str
    singular_values: int
    response_matrix_file: str


class Orbit(Element):

    def __init__(self, cfg: ConfigModel):

        super().__init__(cfg.name)
        self._cfg = cfg
        rm_file_name: Path = get_root_folder() / cfg.response_matrix_file
        self.response_matrix = ResponseMatrix.from_json(rm_file_name)
        self._hcorr:"MagnetArray" = None
        self._vcorr:"MagnetArray" = None
        self._hvcorr:"MagnetArray" = None

    def correct(
        self,
        reference=None,
        gain: float = 1.0,
        plane: Optional[Literal["H", "V"]] = None,
    ):        
        self.check_peer()

        # Lazilly attach magnet arrays
        if self._hcorr is None:
            self._hcorr = self._peer.get_magnets(self._cfg.hcorr_array_name)
            self._vcorr = self._peer.get_magnets(self._cfg.vcorr_array_name)
            hvElts = []
            hvElts.extend(self._hcorr)
            hvElts.extend(self._vcorr)
            self._hvcorr = MagnetArray("HVCorr",hvElts)


        interface = pySCInterface(
            element_holder=self._peer,
            bpm_array_name=self._cfg.bpm_array_name,
            hcorr_array_name=self._cfg.hcorr_array_name,
            vcorr_array_name=self._cfg.vcorr_array_name,
        )

        if plane is None or plane == "H":
            trims_h = orbit_correction(
                interface=interface,
                response_matrix=self.response_matrix,
                method="svd_values",
                parameter=self._cfg.singular_values,
                zerosum=True,
                apply=False,
                plane="H",
                reference=reference,
            )

        if plane is None or plane == "V":
            trims_v = orbit_correction(
                interface=interface,
                response_matrix=self.response_matrix,
                method="svd_values",
                parameter=self._cfg.singular_values,
                zerosum=False,
                apply=False,
                plane="V",
                reference=reference,
            )

        if plane is None:
            data_to_send_hv = self._hvcorr.strengths.get()
            for idx,name in enumerate(trims_h.keys()):
                data_to_send_hv[idx] += trims_h[name] * gain
            l = len(trims_h.keys())
            for idx,name in enumerate(trims_v.keys()):
                data_to_send_hv[idx+l] += trims_v[name] * gain
            self._hvcorr.strengths.set(data_to_send_hv)
        elif plane == "H":
            data_to_send_h = self._hcorr.strengths.get()
            for idx,name in enumerate(trims_h.keys()):
                data_to_send_h[idx] += trims_h[name] * gain
            self._hcorr.strengths.set(data_to_send_h)
        elif plane == "V":
            data_to_send_v = self._vcorr.strengths.get()
            for idx,name in enumerate(trims_v.keys()):
                data_to_send_v[idx] += trims_v[name] * gain
            self._vcorr.strengths.set(data_to_send_v)
    
    def attach(
        self,
        peer: "ElementHolder",
    ) -> Self:
        """
        Create a new reference to attach this tune object to a simulator
        or a control system.
        """
        obj = self.__class__(self._cfg)
        obj._peer = peer
        obj.response_matrix = self.response_matrix # Copy only ref
        return obj

[kparasch]

@JeanLucPons
The assumption on the ordering will "usually" work depending how the response matrix was built, because the keys in trims comes from response_matrix.input_names.

A more robust way would be to iterate over names of the array:

# join h and v
trims = {}
for key in trims_h.keys():
    trims[key] = trims_h[key]
for key in trims_h.keys():
    trims[key] = trims_v[key]


arr_names = self._hvcorr.names()
data_to_send_hv = self._hvcorr.strengths.get()
for idx,name in enumerate(arr_names):
     data_to_send_hv[idx] += trims[name] * gain

and we should probably also make a check that all correctors from trims will be set in the arrays.

for name in trims.keys():
    assert name in arr_names

self._hvcorr.strengths.set(data_to_send_hv)

Do you want to merge the branches and one of us continues?

[JeanLucPons]

I rebased and merged my mods.
I let you optimize the join of h and v.
The sample below is working:

from pyaml.accelerator import Accelerator
sr = Accelerator.load("tests/config/EBSOrbit.yaml",use_fast_loader=False)
sr.live.orbit.correct()

[kparasch]

@JeanLucPons This is basically finished except one thing:

    def post_init(self):
        self._hcorr = self._peer.get_magnets(self._cfg.hcorr_array_name)
        self._vcorr = self._peer.get_magnets(self._cfg.vcorr_array_name)
        hvElts = []
        hvElts.extend(self._hcorr)
        hvElts.extend(self._vcorr)
        self._hvcorr = MagnetArray("HVCorr", hvElts)

Here we create a magnet family with a hardcoded name "HVCorr". I think we should not do this so that it doesn't break if someone decides to define a magnet family called "HVCorr".
I don't have any good solution.

[JeanLucPons]

@JeanLucPons This is basically finished except one thing:

    def post_init(self):
        self._hcorr = self._peer.get_magnets(self._cfg.hcorr_array_name)
        self._vcorr = self._peer.get_magnets(self._cfg.vcorr_array_name)
        hvElts = []
        hvElts.extend(self._hcorr)
        hvElts.extend(self._vcorr)
        self._hvcorr = MagnetArray("HVCorr", hvElts)

Here we create a magnet family with a hardcoded name "HVCorr". I think we should not do this so that it doesn't break if someone decides to define a magnet family called "HVCorr". I don't have any good solution.

In fact this name does not conflict as this array is not added in the holder. This is as for dynamic array which have no name for the moment. You can also use "" if you prefer.

[kparasch]

Ok great! I have then used "" since anyway the name is not supposed to be seen by anyone I think.

The PR is ready to be merged.

We can now do:

sr.live.orbit.correct()
sr.design.orbit.correct()

[JeanLucPons]

Thanks Kostas.