From db6b42c7d61b7ed6febd9abbacde6c0bdf8e9c2e Mon Sep 17 00:00:00 2001 From: oscarxblanco Date: Fri, 31 Jul 2026 15:29:27 +0200 Subject: [PATCH 1/2] set default amplitudes, freq, phase for default mode SINE --- pyat/at/lattice/elements/variable_elements.py | 75 +++++++++++-------- 1 file changed, 44 insertions(+), 31 deletions(-) diff --git a/pyat/at/lattice/elements/variable_elements.py b/pyat/at/lattice/elements/variable_elements.py index a8dd94a67..f450f3da0 100644 --- a/pyat/at/lattice/elements/variable_elements.py +++ b/pyat/at/lattice/elements/variable_elements.py @@ -7,7 +7,6 @@ import numpy as np -from ..exceptions import AtError from .conversions import _array from .element_object import Element @@ -100,37 +99,52 @@ def __init__( * For ``mode=at.ACMode.ARBITRARY`` the ``Func(A,B)`` corresponding to the ``Amplitude(A,B)`` has to be provided """ - if len(kwargs) > 0: - self.Mode = mode.value - self.ModeName = mode.name - kwargs.setdefault("PassMethod", "VariableThinMPolePass") - self.MaxOrder = kwargs.pop("MaxOrder", 0) - self.Periodic = kwargs.pop("Periodic", True) - if AmplitudeA is None and AmplitudeB is None: - msg = "Please provide at least one amplitude for A or B" - raise AtError(msg) - AmplitudeB = self._set_params(AmplitudeB, "B", **kwargs) - AmplitudeA = self._set_params(AmplitudeA, "A", **kwargs) - self._setmaxorder(AmplitudeA, AmplitudeB) - if self.Mode == ACMode.WHITENOISE: - self.Seed = kwargs.pop("Seed", datetime.now().timestamp()) - self.PolynomA = np.zeros(self.MaxOrder + 1) - self.PolynomB = np.zeros(self.MaxOrder + 1) - ramps = kwargs.pop("Ramps", None) - if ramps is not None: - assert len(ramps) == 4, "Ramps has to be a vector with 4 elements" - self.Ramps = ramps + + def _default_amplitudes(ampa, ampb): + if ampa is None and ampb is None: + ampb = np.array([0]) + if np.isscalar(ampa): + ampa = np.array([ampa]) + if np.isscalar(ampb): + ampb = np.array([ampb]) + return ampa, ampb + + def _getmaxorder(ampa, ampb): + mxa, mxb = 0, 0 + if ampa is not None: + mxa = np.max(np.append(np.nonzero(ampa), 0)) + if ampb is not None: + mxb = np.max(np.append(np.nonzero(ampb), 0)) + return max(mxa, mxb) + + self.Mode = mode.value + self.ModeName = mode.name + kwargs.setdefault("PassMethod", "VariableThinMPolePass") + + AmplitudeA, AmplitudeB = _default_amplitudes(AmplitudeA, AmplitudeB) + + # MaxOrder is set finally by the user if given + max_order_ampab = _getmaxorder(AmplitudeA, AmplitudeB) + self.MaxOrder = kwargs.get("MaxOrder", max_order_ampab) + # after the definition of MaxOrder we can create Amplitudes + self._set_amplitudes(AmplitudeA, AmplitudeB) + + self.Periodic = kwargs.pop("Periodic", True) + AmplitudeB = self._set_params(AmplitudeB, "B", **kwargs) + AmplitudeA = self._set_params(AmplitudeA, "A", **kwargs) + if self.Mode == ACMode.WHITENOISE: + self.Seed = kwargs.pop("Seed", datetime.now().timestamp()) + self.PolynomA = np.zeros(self.MaxOrder + 1) + self.PolynomB = np.zeros(self.MaxOrder + 1) + ramps = kwargs.pop("Ramps", None) + if ramps is not None: + assert len(ramps) == 4, "Ramps has to be a vector with 4 elements" + self.Ramps = ramps super().__init__(family_name, **kwargs) - def _setmaxorder(self, ampa, ampb): - mxa, mxb = 0, 0 - if ampa is not None: - mxa = np.max(np.nonzero(ampa)) - if ampb is not None: - mxb = np.max(np.nonzero(ampb)) - self.MaxOrder = max(mxa, mxb) + def _set_amplitudes(self, ampa, ampb): if ampa is not None: - delta = self.MaxOrder - len(ampa) + delta = self.MaxOrder + 1 - len(ampa) if delta > 0: ampa = np.pad(ampa, (0, delta)) self.AmplitudeA = ampa @@ -152,9 +166,8 @@ def _set_params(self, amplitude, ab, **kwargs): return amplitude def _set_sine(self, ab, **kwargs): - frequency = kwargs.pop("Frequency" + ab, None) + frequency = kwargs.pop("Frequency" + ab, 0) phase = kwargs.pop("Phase" + ab, 0) - assert frequency is not None, "Please provide a value for Frequency" + ab setattr(self, "Frequency" + ab, frequency) setattr(self, "Phase" + ab, phase) From 1859fb226171f2170fc8ab50041b107c9e49dd16 Mon Sep 17 00:00:00 2001 From: oscarxblanco Date: Mon, 3 Aug 2026 15:28:07 +0200 Subject: [PATCH 2/2] setwhitenoise --- atmat/lattice/element_creation/atvariablethinmultipole.m | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/atmat/lattice/element_creation/atvariablethinmultipole.m b/atmat/lattice/element_creation/atvariablethinmultipole.m index 947198dff..8e6748c5a 100644 --- a/atmat/lattice/element_creation/atvariablethinmultipole.m +++ b/atmat/lattice/element_creation/atvariablethinmultipole.m @@ -95,6 +95,10 @@ end end + function rsrc = setwhitenoise(rsrc, ~) + % it will later implement a buffer + end + function rsrc = setarb(rsrc, ab) funcarg=strcat('Func',ab); if ~isfield(rsrc,funcarg)