From 3d4ab91f8d53e704e0ce7ff476d52f920594149e Mon Sep 17 00:00:00 2001 From: oscarxblanco <108396048+oscarxblanco@users.noreply.github.com> Date: Thu, 29 Aug 2024 15:54:27 +0200 Subject: [PATCH] Add feeddown polba (#815) * black nonlinear.py * add feed down functions * isort; add scipy special comb * change feed_down to feeddown * solving flake8 * solving flake8 * fix typo * fix docstring * fix docstring * fix DocString * restore nonlinear.py from master * add magnet_tools * removes unused code; use polb and pola Sequence[float] * removes debugprint * import magnet_tools * remove debugprint; change DocString to Magnet Tools * removing raise ValueError * code reordered to avoid padding zero-length polynoms * fix bug in array length * black * flake8 * add function feeddown_pol_from_element * add verbose * implements feedback_pol_from_element * improve verbose * fix verbose * check if component is zero --- pyat/at/physics/__init__.py | 1 + pyat/at/physics/magnet_tools.py | 190 ++++++++++++++++++++++++++++++++ 2 files changed, 191 insertions(+) create mode 100644 pyat/at/physics/magnet_tools.py diff --git a/pyat/at/physics/__init__.py b/pyat/at/physics/__init__.py index 6bfc5ca2f..9be5dfd11 100644 --- a/pyat/at/physics/__init__.py +++ b/pyat/at/physics/__init__.py @@ -17,3 +17,4 @@ from .nonlinear import * from .fastring import * from .frequency_maps import fmap_parallel_track +from .magnet_tools import * diff --git a/pyat/at/physics/magnet_tools.py b/pyat/at/physics/magnet_tools.py new file mode 100644 index 000000000..f4556ed6a --- /dev/null +++ b/pyat/at/physics/magnet_tools.py @@ -0,0 +1,190 @@ +"""Magnet tools.""" + +from __future__ import annotations + +from typing import Sequence + +import numpy +from scipy.special import comb + +from ..lattice.elements import Element + +__all__ = [ + "feeddown_polynomba", + "feeddown_from_nth_order", + "feeddown_pol_from_element", +] + + +def feeddown_pol_from_element( + ele: Element, verbose: bool = False, **kwargs: dict[str, any] +) -> dict[str, numpy.ndarray]: + """ + Return the feed down polynoms due to a transverse offset on an element. + + Parameters: + ele: a ring element. + verbose: prints additional info. Default: False. + kwargs: + offset:(x,y) horizontal and vertical offset in meters. + Default taken from T2 or T1. Otherwise, zero. + + Returns: + Dictionary with PolynomB and PolynomA feeddown components. + + Note: Thin lens approximation, T2=-T1. + Bending angles are ignored. + """ + # check verbose flags only once + verboseprint = print if verbose else lambda *a, **k: None + # handle cases where polynom a and b are not defined + polb = numpy.array([]) + pola = numpy.array([]) + if hasattr(ele, "PolynomB"): + polb = ele.PolynomB + else: + verboseprint(f"Element {ele.FamName} has no PolynomB") + if hasattr(ele, "PolynomA"): + pola = ele.PolynomA + else: + verboseprint(f"Element {ele.FamName} has no PolynomA") + + # check if user has set offsets + xoffset, yoffset = kwargs.get("offset", (0, 0)) + # check if element has T1 and T2. Use one. + if "offset" not in kwargs: + if hasattr(ele, "T2"): + xoffset = ele.T2[0] + yoffset = ele.T2[2] + elif hasattr(ele, "T1"): + xoffset = -ele.T1[0] + yoffset = -ele.T1[2] + else: + verboseprint(f"Element {ele.FamName} has no T1 or T2.") + verboseprint(f"Using offsets xoffset={xoffset}, yoffset={yoffset}.") + # Return the polynoms + return feeddown_polynomba(polb=polb, pola=pola, xoffset=xoffset, yoffset=yoffset) + + +def feeddown_polynomba( + polb: Sequence[float], + pola: Sequence[float], + xoffset: float = 0, + yoffset: float = 0, + verbose: bool = False, +) -> dict[str, numpy.ndarray]: + """ + Return the feeddown due to a transverse offset. + + Parameters: + polb: PolynomB. + pola: PolynomA. + xoffset: Default zero. Horizontal offset in meters. + yoffset: Default zero. Vertical offset in meters. + verbose: prints additional info. + + Returns: + Dictionary with PolynomB and PolynomA feeddown components. + + Note: + If only one polynom is passed, it is assumed to be PolynomB. + """ + # check verbose flags only once + verboseprint = print if verbose else lambda *a, **k: None + + # verify polynoms length + maxorda = len(pola) + maxordb = len(polb) + maxord = max(maxorda, maxordb) + polasum = numpy.zeros(max(0, maxord - 1)) + polbsum = numpy.zeros(max(0, maxord - 1)) + if maxorda == 0 and maxordb == 0: + verboseprint("Both polynoms are zero.") + else: + polbpad = numpy.pad( + polb, (0, maxord - maxordb), "constant", constant_values=(0, 0) + ) + polapad = numpy.pad( + pola, (0, maxord - maxorda), "constant", constant_values=(0, 0) + ) + + verboseprint(f"polb={polb},pola={pola}") + verboseprint(f"xoffset={xoffset},yoffset={yoffset}") + for ith in range(2, maxord + 1): + polbaux_b, polaaux_b = feeddown_from_nth_order( + ith, polbpad[ith - 1], xoffset, yoffset, poltype="B" + ) + polbaux_a, polaaux_a = feeddown_from_nth_order( + ith, polapad[ith - 1], xoffset, yoffset, poltype="A" + ) + polbshort = polbaux_b + polbaux_a + polashort = polaaux_b + polaaux_a + polbsum = polbsum + numpy.pad( + polbshort, (0, maxord - ith), "constant", constant_values=(0, 0) + ) + polasum = polasum + numpy.pad( + polashort, (0, maxord - ith), "constant", constant_values=(0, 0) + ) + poldict = {"PolynomB": polbsum, "PolynomA": polasum} + verboseprint(f"poldict={poldict}") + return poldict + + +def feeddown_from_nth_order( + nthorder: int, + nthpolcomp: float, + xoffset: float, + yoffset: float, + poltype: str = "B", + verbose: bool = False, +) -> tuple[numpy.ndarray, numpy.ndarray]: + """ + Return the feeddown polynoms from an nth-order magnet component transverse offset. + + Parameters: + nthorder: integer order of the magnet component, e.g. 1 for dipole, + 2 for quadrupoles, 3 for sextupoles, etc. + nthpolcomp: float. nth component of the polynom, i.e. the value of + PolynomA/B[nthorder-1]. + xoffset: float. Horizontal offset in meters. + yoffset: float. Vertical offset in meters. + poltype: Default 'B'. Could be 'B' or 'A', otherwise ignored. + verbose: print info on input and output parameters. + + Returns: + Tuple of two numpy arrays with PolynomB and PolynomA from feeddown. + """ + verboseprint = print if verbose else lambda *a, **k: None + verboseprint(f"nthorder={nthorder},nthpolcomp={nthpolcomp},poltype={poltype}") + verboseprint(f"xoffset={xoffset},yoffset={yoffset}") + fakeimag = {0: 1, 1: 0, 2: -1, 3: 0} + polbaux = numpy.zeros(nthorder - 1) + polaaux = numpy.zeros(nthorder - 1) + if nthpolcomp != 0: + for kterm in range(1, nthorder): + ichoosek = comb(nthorder - 1, kterm) + pascalsn = comb(kterm, numpy.arange(kterm + 1)) + powk = numpy.arange(kterm + 1) + powkflip = numpy.arange(kterm, -1, -1) + recoefs = numpy.array([fakeimag[numpy.mod(idx, 4)] for idx in powk]) + imcoefs = numpy.array([fakeimag[numpy.mod(idx + 3, 4)] for idx in powk]) + commonfactor = nthpolcomp * ichoosek * (-1) ** kterm + repart = ( + recoefs + * commonfactor + * (pascalsn * (xoffset**powkflip * yoffset**powk)) + ) + impart = ( + imcoefs + * commonfactor + * (pascalsn * (xoffset**powkflip * yoffset**powk)) + ) + polbaux[nthorder - kterm - 1] = polbaux[nthorder - kterm - 1] + repart.sum() + polaaux[nthorder - kterm - 1] = polaaux[nthorder - kterm - 1] + impart.sum() + verboseprint(f"polbaux={polbaux},polaaux={polaaux}") + polbout, polaout = polbaux, polaaux + # skew components swap the imaginary and real feed-down and change the sign + # of the real part, i.e. j*j = -1 + if poltype == "A": + polbout, polaout = -1 * polaaux, polbaux + return polbout, polaout