diff --git a/pyaml/lattice/abstract_impl.py b/pyaml/lattice/abstract_impl.py index 615d7ea0..a734a587 100644 --- a/pyaml/lattice/abstract_impl.py +++ b/pyaml/lattice/abstract_impl.py @@ -708,13 +708,14 @@ class BPMScalarAggregator(ScalarAggregator): Return the value unit. """ - def __init__(self, ring: at.Lattice): + def __init__(self, ring: at.Lattice, twiss_in: dict): """ Initialize the BPMScalarAggregator. """ self._lattice = ring self._refpts = [] self._matrices = [] + self._twiss_in = twiss_in def add_elem(self, elem: at.Element): """ @@ -764,9 +765,15 @@ def unit(self) -> str: def _transform(self) -> np.array: """Transform the configured value for the lattice interface.""" - _, orbit = at.find_orbit(self._lattice, refpts=self._refpts) + if self._twiss_in is None: + _, orbit = at.find_orbit(self._lattice, refpts=self._refpts) + pts = orbit[:, [0, 2]] # Extract x,y + else: + _, _, elemdata = at.get_optics(self._lattice, refpts=self._refpts, twiss_in=self._twiss_in) + pts = [e.closed_orbit[[0, 2]] for e in elemdata] + ones = np.ones(len(self._refpts)) - pts = orbit[:, [0, 2]] # Extract x,y + # Transfom BPM coordinates (offset and tilt) # Batch matrices multiplication (homogeneous coordinates) return np.matmul(self._matrices, np.column_stack([pts, ones])[:, :, None]).squeeze(-1) @@ -856,19 +863,25 @@ class RBpmArray(abstract.ReadFloatArray): Return the value unit. """ - def __init__(self, element: at.Element, lattice: at.Lattice): + def __init__(self, element: at.Element, lattice: at.Lattice, twiss_in: dict): """ Initialize the RBpmArray. """ self._element = element self._lattice = lattice + self._twiss_in = twiss_in # Gets the value def get(self) -> np.array: """Return the current value.""" index = self._lattice.index(self._element) - _, orbit = at.find_orbit(self._lattice, refpts=index) - pts = orbit[0, [0, 2]] + if self._twiss_in is None: + _, orbit = at.find_orbit(self._lattice, refpts=index) + pts = orbit[0, [0, 2]] + else: + _, _, elemdata = at.get_optics(self._lattice, refpts=index, twiss_in=self._twiss_in) + pts = elemdata[0].closed_orbit[[0, 2]] + # Transfom BPM coordinates (offset and tilt) return np.dot(self._element._transform, np.hstack([pts, 1])) # Use homogeneous coordinates # Gets the unit of the value diff --git a/pyaml/lattice/simulator.py b/pyaml/lattice/simulator.py index 738ddaf2..e0922ae3 100644 --- a/pyaml/lattice/simulator.py +++ b/pyaml/lattice/simulator.py @@ -48,6 +48,7 @@ from ..tuning_tools.tuning_tool import TuningTool from ..validation import DynamicValidation, register_schema from .lattice_elements_linker import LatticeElementsLinker +from .twiss_in import TwissIn if TYPE_CHECKING: from ..configuration.unbound_element import UnboundElement @@ -79,6 +80,8 @@ class Simulator(ElementHolder, DynamicValidation): mat_key : str, optional Variable name of the lattice when loading a MATLAB ``.mat`` lattice file. + twiss_in : TwissIn, optional + Describe input beam parameters of a tranfer line, when not specified the lattice is considered as a ring linker : LatticeElementsLinker, optional Custom linker used to associate PyAML elements with PyAT lattice elements. If omitted, elements are matched by name. @@ -123,6 +126,7 @@ def __init__( name: str, lattice: str, mat_key: str | None = None, + twiss_in: TwissIn | None = None, linker: LatticeElementsLinker | None = None, description: str | None = None, ): @@ -135,6 +139,11 @@ def __init__( self._lattice = lattice self._mat_key = mat_key self.description = description + self.twiss_in = twiss_in + if twiss_in is None: + self._at_twiss_in = None + else: + self._at_twiss_in = self.twiss_in._to_at() path: Path = ROOT.get() / self._lattice @@ -226,9 +235,9 @@ def create_bpm_aggregators(self, bpms: list[BPM]) -> list[ScalarAggregator]: list[ScalarAggregator] Aggregators for combined, horizontal, and vertical BPM positions. """ - agg = BPMScalarAggregator(self.get_lattice()) - aggh = BPMHScalarAggregator(self.get_lattice()) - aggv = BPMVScalarAggregator(self.get_lattice()) + agg = BPMScalarAggregator(self.get_lattice(), self._at_twiss_in) + aggh = BPMHScalarAggregator(self.get_lattice(), self._at_twiss_in) + aggv = BPMVScalarAggregator(self.get_lattice(), self._at_twiss_in) for b in bpms: e = self.get_at_elems(b)[0] agg.add_elem(e) @@ -293,7 +302,9 @@ def _fill_bpm(self, bpm: BPM) -> None: raise PyAMLException(f"BPM {bpm.get_name()} offset must be a 2-element array.") update_bpm_transform_matrix(bpm_elt) self.diagnostic.bpm.add( - bpm.attach(self, RBpmArray(bpm_elt, self.ring), RWBpmOffsetArray(bpm_elt), RWBpmTiltScalar(bpm_elt)) + bpm.attach( + self, RBpmArray(bpm_elt, self.ring, self._at_twiss_in), RWBpmOffsetArray(bpm_elt), RWBpmTiltScalar(bpm_elt) + ) ) def _fill_rf_plant(self, rf_plant: RFPlant) -> None: diff --git a/pyaml/lattice/twiss_in.py b/pyaml/lattice/twiss_in.py new file mode 100644 index 00000000..c42ee6c7 --- /dev/null +++ b/pyaml/lattice/twiss_in.py @@ -0,0 +1,52 @@ +""" +Initial input beam parameters for a transfer line +""" + +import numpy + +from ..validation import DynamicValidation, register_schema + +# Define the main class name for this module +PYAMLCLASS = "TwissIn" + + +@register_schema +class TwissIn(DynamicValidation): + """ + Class that describes input beam parameters of a transfer line + + Attributes + ---------- + alpha : list[float] + Initial alpha parameters (x,y) + beta : list[float] | None + Initial beta parameters (x,y), default [0,0] + position : list[float] | None + Initial position (6D), default [0,0,0,0,0,0] + dispersion : list[float] | None + Initial dispersion (4D), default [0,0,0,0] + """ + + def __init__( + self, + alpha: list[float], + beta: list[float] | None = None, + position: list[float] | None = None, + dispersion: list[float] | None = None, + ): + """ + Create TwissIn object + """ + super().__init__() + self.alpha = alpha + self.beta = beta if beta is not None else [0.0, 0.0] + self.position = position if position is not None else [0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + self.dispersion = dispersion if dispersion is not None else [0.0, 0.0, 0.0, 0.0] + + def _to_at(self) -> dict: + return { + "beta": numpy.array(self.alpha, dtype=float), + "alpha": numpy.array(self.beta, dtype=float), + "closed_orbit": numpy.array(self.position, dtype=float), + "dispersion": numpy.array(self.dispersion, dtype=float), + } diff --git a/tests/config/tl2/tl2.mat b/tests/config/tl2/tl2.mat new file mode 100644 index 00000000..b871250f Binary files /dev/null and b/tests/config/tl2/tl2.mat differ diff --git a/tests/config/tl2/tl2.yml b/tests/config/tl2/tl2.yml new file mode 100644 index 00000000..8509e34d --- /dev/null +++ b/tests/config/tl2/tl2.yml @@ -0,0 +1,215 @@ +type: pyaml.accelerator +facility: ESRF +machine: tl2 +energy: 6e9 + +simulators: + - class: pyaml.lattice.simulator.Simulator + lattice: tl2.mat + name: design + twiss_in: + class: pyaml.lattice.twiss_in.TwissIn + alpha: [-1.71, 1.47] + beta: [5.9 , 5.65] + dispersion: [-0.145, -0.092, 0, 0] + +data_folder: /data/store + +arrays: + + - type: pyaml.arrays.magnet + name: HCORR + elements: + - CH* + - type: pyaml.arrays.magnet + name: VCORR + elements: + - CV* + - type: pyaml.arrays.bpm + name: BPMS + elements: + - BPM* + +devices: + + - class: pyaml.tuning_tools.orbit.Orbit + bpm_array_name: BPMS + hcorr_array_name: HCORR + vcorr_array_name: VCORR + name: DEFAULT_ORBIT_CORRECTION + singular_values: 5 + response_matrix: ideal_tl2_orm.json + + - class: pyaml.tuning_tools.orbit_response_matrix.OrbitResponseMatrix + bpm_array_name: BPMS + hcorr_array_name: HCORR + vcorr_array_name: VCORR + corrector_delta: 1e-6 + name: DEFAULT_ORBIT_RESPONSE_MATRIX + + - class: pyaml.diagnostics.bpm.bpm.BPM + name: BPM_CV1 + x_pos: tl2/d-bpm/cv1/X_Position + y_pos: tl2/d-bpm/cv1/Z_Position + - class: pyaml.diagnostics.bpm.bpm.BPM + name: BPM_CH3 + x_pos: tl2/d-bpm/d2/X_Position + y_pos: tl2/d-bpm/d2/Z_Position + - class: pyaml.diagnostics.bpm.bpm.BPM + name: BPM_QD5 + x_pos: tl2/d-bpm/q5/X_Position + y_pos: tl2/d-bpm/q5/Z_Position + - class: pyaml.diagnostics.bpm.bpm.BPM + name: BPM_QF6 + x_pos: tl2/d-bpm/q6/X_Position + y_pos: tl2/d-bpm/q6/Z_Position + - class: pyaml.diagnostics.bpm.bpm.BPM + name: BPM_SX + x_pos: tl2/d-bpm/sx/X_Position + y_pos: tl2/d-bpm/sx/Z_Position + - class: pyaml.diagnostics.bpm.bpm.BPM + name: BPM_QD8 + x_pos: tl2/d-bpm/q8/X_Position + y_pos: tl2/d-bpm/q8/Z_Position + - class: pyaml.diagnostics.bpm.bpm.BPM + name: BPM_QD12 + x_pos: tl2/d-bpm/q12/X_Position + y_pos: tl2/d-bpm/q12/Z_Position + - class: pyaml.diagnostics.bpm.bpm.BPM + name: BPM_QF14 + x_pos: tl2/d-bpm/q14/X_Position + y_pos: tl2/d-bpm/q14/Z_Position + - class: pyaml.diagnostics.bpm.bpm.BPM + name: BPM_S1PM + x_pos: tl2/d-bpm/s1pm/X_Position + y_pos: tl2/d-bpm/s1pm/Z_Position + + - class: pyaml.magnet.hcorrector.HCorrector + name: CH1 + model: + class: pyaml.magnet.linear_model.LinearMagnetModel + calibration_factor: 161e-6 + unit: 1/m + hardware_unit: A + powerconverter: tl2/ps-c1/ch1/Current + - class: pyaml.magnet.hcorrector.HCorrector + name: CH2 + model: + class: pyaml.magnet.linear_model.LinearMagnetModel + calibration_factor: 161e-6 + unit: 1/m + hardware_unit: A + powerconverter: tl2/ps-c1/ch2/Current + - class: pyaml.magnet.hcorrector.HCorrector + name: CH3 + model: + class: pyaml.magnet.linear_model.LinearMagnetModel + calibration_factor: 161e-6 + unit: 1/m + hardware_unit: A + powerconverter: tl2/ps-c1/ch3/Current + - class: pyaml.magnet.hcorrector.HCorrector + name: CH4 + model: + class: pyaml.magnet.linear_model.LinearMagnetModel + calibration_factor: 161e-6 + unit: 1/m + hardware_unit: A + powerconverter: tl2/ps-c1/ch4/Current + - class: pyaml.magnet.hcorrector.HCorrector + name: CH5 + model: + class: pyaml.magnet.linear_model.LinearMagnetModel + calibration_factor: 161e-6 + unit: 1/m + hardware_unit: A + powerconverter: tl2/ps-c1/ch5/Current + - class: pyaml.magnet.hcorrector.HCorrector + name: CH6 + model: + class: pyaml.magnet.linear_model.LinearMagnetModel + calibration_factor: 161e-6 + unit: 1/m + hardware_unit: A + powerconverter: tl2/ps-c1/ch6/Current + - class: pyaml.magnet.hcorrector.HCorrector + name: CH7 + model: + class: pyaml.magnet.linear_model.LinearMagnetModel + calibration_factor: 161e-6 + unit: 1/m + hardware_unit: A + powerconverter: tl2/ps-c1/ch7/Current + - class: pyaml.magnet.hcorrector.HCorrector + name: CH8 + model: + class: pyaml.magnet.linear_model.LinearMagnetModel + calibration_factor: 161e-6 + unit: 1/m + hardware_unit: A + powerconverter: tl2/ps-c1/ch8/Current + + - class: pyaml.magnet.vcorrector.VCorrector + name: CV0 + model: + class: pyaml.magnet.linear_model.LinearMagnetModel + calibration_factor: 161e-6 + unit: 1/m + hardware_unit: A + powerconverter: tl2/ps-c1/cv0/Current + - class: pyaml.magnet.vcorrector.VCorrector + name: CV1 + model: + class: pyaml.magnet.linear_model.LinearMagnetModel + calibration_factor: 161e-6 + unit: 1/m + hardware_unit: A + powerconverter: tl2/ps-c1/cv1/Current + - class: pyaml.magnet.vcorrector.VCorrector + name: CV3 + model: + class: pyaml.magnet.linear_model.LinearMagnetModel + calibration_factor: 161e-6 + unit: 1/m + hardware_unit: A + powerconverter: tl2/ps-c1/cv3/Current + - class: pyaml.magnet.vcorrector.VCorrector + name: CV4 + model: + class: pyaml.magnet.linear_model.LinearMagnetModel + calibration_factor: 161e-6 + unit: 1/m + hardware_unit: A + powerconverter: tl2/ps-c1/cv4/Current + - class: pyaml.magnet.vcorrector.VCorrector + name: CV5 + model: + class: pyaml.magnet.linear_model.LinearMagnetModel + calibration_factor: 161e-6 + unit: 1/m + hardware_unit: A + powerconverter: tl2/ps-c1/cv5/Current + - class: pyaml.magnet.vcorrector.VCorrector + name: CV6 + model: + class: pyaml.magnet.linear_model.LinearMagnetModel + calibration_factor: 161e-6 + unit: 1/m + hardware_unit: A + powerconverter: tl2/ps-c1/cv6/Current + - class: pyaml.magnet.vcorrector.VCorrector + name: CV7 + model: + class: pyaml.magnet.linear_model.LinearMagnetModel + calibration_factor: 161e-6 + unit: 1/m + hardware_unit: A + powerconverter: tl2/ps-c1/cv7/Current + - class: pyaml.magnet.vcorrector.VCorrector + name: CV8 + model: + class: pyaml.magnet.linear_model.LinearMagnetModel + calibration_factor: 161e-6 + unit: 1/m + hardware_unit: A + powerconverter: tl2/ps-c1/cv8/Current diff --git a/tests/tuning_tools/test_transfer_line.py b/tests/tuning_tools/test_transfer_line.py new file mode 100644 index 00000000..cd4dfe3e --- /dev/null +++ b/tests/tuning_tools/test_transfer_line.py @@ -0,0 +1,47 @@ +import numpy +import pytest + +from pyaml.accelerator import Accelerator + + +def test_transfer_line(): + tl2 = Accelerator.load("tests/config/tl2/tl2.yml") + tl2.design.tool.orm.measure() + orm_data = tl2.design.tool.orm.get() + mat = numpy.array(orm_data["matrix"]) + + orm_hv = mat[:9, 8:] + orm_vh = mat[9:, :8] + assert numpy.allclose(orm_hv, 0) + assert numpy.allclose(orm_vh, 0) + + orm_h = mat[:9, :8] + orm_v = mat[9:, 8:] + + assert numpy.allclose( + orm_h[8, :], [1.49746082, 3.11487458, 10.60522905, -0.84014865, -19.27774502, -14.92314963, 1.29287711, 1.60223207] + ) + assert numpy.allclose( + orm_v[8, :], [10.95115436, 10.56436269, 16.43975007, 32.48697172, 12.7218589, 10.31009581, 10.46821268, 2.60568252] + ) + + tl2.design.tool.orm.save("tl2.json") + tl2.design.tool.orbit.load("tl2.json") + + # Mangle the orbit + hcorr = tl2.design.magnets.get("HCORR") + vcorr = tl2.design.magnets.get("VCORR") + numpy.random.seed(1) + std_kick = 1e-6 # rad + hcorr.strengths.set(hcorr.strengths.get() + std_kick * numpy.random.normal(size=len(hcorr))) + vcorr.strengths.set(vcorr.strengths.get() + std_kick * numpy.random.normal(size=len(vcorr))) + + orbit = tl2.design.diagnostic.bpms.get("BPMS").positions + assert numpy.allclose(numpy.std(orbit.get(), axis=0), [1.78601145e-05, 3.57671504e-05]) + + tl2.design.tool.orbit.correct() + tl2.design.tool.orbit.correct() + + assert numpy.allclose(numpy.std(orbit.get(), axis=0), [7.50786048e-07, 7.26005222e-07]) + + assert numpy.allclose(orbit.get()[2], tl2.design.diagnostic.bpm.get("BPM_QD5").positions.get())