From ef332184bf1806993e6617f0f0e89f9d9eb24deb Mon Sep 17 00:00:00 2001 From: Silke Date: Mon, 8 Jun 2026 13:40:29 +0200 Subject: [PATCH 01/38] Finally compiling and finding tracker --- test_fieldexpansion_element.py | 14 ++ xtrack/beam_elements/elements.py | 63 ++++++++- .../elements_src/create_fieldexpansion.h | 124 ++++++++++++++++++ .../elements_src/track_fieldexpansion.h | 95 ++++++++++++++ 4 files changed, 295 insertions(+), 1 deletion(-) create mode 100644 test_fieldexpansion_element.py create mode 100644 xtrack/beam_elements/elements_src/create_fieldexpansion.h create mode 100644 xtrack/beam_elements/elements_src/track_fieldexpansion.h diff --git a/test_fieldexpansion_element.py b/test_fieldexpansion_element.py new file mode 100644 index 000000000..540d38bed --- /dev/null +++ b/test_fieldexpansion_element.py @@ -0,0 +1,14 @@ +import xtrack as xt +import numpy as np + +h = 0.2 +a = np.array([[1,1]]) +b = np.array([[1,2]]) +bs = np.array([1,5]) +ny = 5 +length=1 + +line = xt.Line(elements=[xt.FieldExpansion(length=length, h=h, aa=a, bb=b, bs=bs, ny=ny)]) +p = xt.Particles(x=0.01) +line.track(p) +print(p.x) \ No newline at end of file diff --git a/xtrack/beam_elements/elements.py b/xtrack/beam_elements/elements.py index 5f55dcc20..71facae99 100644 --- a/xtrack/beam_elements/elements.py +++ b/xtrack/beam_elements/elements.py @@ -4908,7 +4908,7 @@ class ElectronCooler(BeamElement): longitudinal component of the magnetic field. This is a measure of the magnetic field quality. With the ideal magnetic field quality being 0. - space_charge : float, optional + space_charge : float, optional length, h, bs, a, b Whether space charge of electron beam is enabled. 0 is off and 1 is on. """ @@ -4983,3 +4983,64 @@ def get_backtrack_element(self, _context=None, _buffer=None, _offset=None): class ThinSliceNotNeededError(Exception): pass + + +class FieldExpansion(BeamElement): + """ + Specifies the field expansion in general derivatives on axis in straight or curved frame + + Parameters + ---------- + h : float + Curvature of the element, in 1/m. For straight elements, h=0. + a : array, shape na, deg+1, floats + describing the polynomial coefficients for the skew multipoles. First index is multipole order, second index is polynomial coefficient. + b : array, shape nb, deg+1, floats + describing the polynomial coefficients for the normal multipoles. First index is multipole order, second index is polynomial coefficient. + bs : array, shape deg+1, floats + describing the polynomial coefficients for the longitudinal field component. Index is polynomial coefficient, + ny : int + number of powers in y to include, + + """ + isthick = True + behaves_like_drift = True + has_backtrack = False + allow_loss_refinement = False + allow_rot_and_shift = False + + _xofields = { + "length" : xo.Float64, + "h": xo.Float64, + + "ny": xo.Int64, + "deg": xo.Int64, + + "nstep": xo.Int64, + "ds": xo.Float64, + } + + _extra_c_sources = [ + '#include "xtrack/beam_elements/elements_src/track_fieldexpansion.h"', + ] + + def __init__(self, length, h, aa, bb, bs, ny, nstep=10, **kwargs): + super().__init__(**kwargs) + + self.length = length + self.h = h + self.nstep = nstep + self.ds = length/nstep + + self.aa = np.asarray(aa, dtype=np.float64) + self.bb = np.asarray(bb, dtype=np.float64) + self.bs = np.asarray(bs, dtype=np.float64) + + self.na = aa.shape[0] + self.nb = bb.shape[0] + + self.deg = aa.shape[1] - 1 + + if bb.shape[1] != self.deg + 1 or bs.shape[0] != self.deg + 1: + raise ValueError("Invalid input shapes") + \ No newline at end of file diff --git a/xtrack/beam_elements/elements_src/create_fieldexpansion.h b/xtrack/beam_elements/elements_src/create_fieldexpansion.h new file mode 100644 index 000000000..52cc36b46 --- /dev/null +++ b/xtrack/beam_elements/elements_src/create_fieldexpansion.h @@ -0,0 +1,124 @@ +#ifndef create_fieldexpansion_H +#define create_fieldexpansion_H + +typedef struct { + int ny; /* requested output order in y */ + int ncoef; /* stored phi_i coefficients: 0..ny+1 */ + int na, nb, deg; + int mmin, mmax, moff, nm; + int qemin, nq; + double h; + double *c; /* c[i,m,k], polynomial coeff of s^k in q^m term */ + double *V; /* scratch: c[i,m](s) */ + double *D1; /* scratch: d_s c[i,m] */ + double *D2; /* scratch: d2_s c[i,m] */ + double *Q; /* scratch: q^e, e=qemin.. */ +} Expansion; + +void build_expansion( + double h, + int64_t ny, + int64_t na, + int64_t nb, + int64_t deg, + const double *bs, + const double *a, + const double *b, +) +{ + Expansion *f = (Expansion *)xcalloc(1, sizeof(*f)); + f->h = h; + f->ny = ny; + f->ncoef = ny + 2; /* store phi_0..phi_{ny+1} so By is also order ny */ + f->na = na; + f->nb = nb; + f->deg = deg; + f->mmax = (na > nb - 1) ? na : (nb - 1); + f->mmin = -2 * ((f->ncoef - 1) / 2); + f->moff = -f->mmin; + f->nm = f->mmax - f->mmin + 1; + f->qemin = f->mmin - 1; + f->nq = (f->mmax + 2) - f->qemin + 1; + f->c = (double *)xcalloc((size_t)f->ncoef * (size_t)f->nm * (size_t)(deg + 1), sizeof(double)); + f->V = (double *)xcalloc((size_t)f->ncoef * (size_t)f->nm, sizeof(double)); + f->D1 = (double *)xcalloc((size_t)f->ncoef * (size_t)f->nm, sizeof(double)); + f->D2 = (double *)xcalloc((size_t)f->ncoef * (size_t)f->nm, sizeof(double)); + f->Q = (double *)xcalloc((size_t)f->nq, sizeof(double)); + + int nmax = (na > nb) ? na : nb; + double invfact[nmax + 1]; + double invhpow[nmax + 1]; + invfact[0] = 1.0; + invhpow[0] = 1.0; + for (int n = 1; n <= nmax; ++n) { + invfact[n] = invfact[n - 1] / (double)n; + invhpow[n] = invhpow[n - 1] / h; + } + + /* phi_0(s) = sum_m c[0,m](s) q^m + c[0,m] = - sum_(n>=m) (-1)^(n-m) / (h^n m! (n-m)!) a_n(s) */ + for (int m = 0; m <= na; ++m) { + double *dst = cptr(f, 0, m + f->moff); + /* a_0(s)=int_0^s b_s(u)du contributes -a_0 to phi_0. + CAREFUL: this will neglect the highest order in the polynomial, + only up to given degree in a0 is kept */ + if (m==0) { + for (int k = 0; k < deg; ++k) + dst[k + 1] = bs[k] / (double)(k + 1); + } + for (int n = (m > 1 ? m : 1); n <= na; ++n) { + double sgn = ((n - m) & 1) ? -1.0 : 1.0; + double fac = -sgn * invhpow[n] * invfact[m] * invfact[n - m]; + const double *an = a + (size_t)(n - 1) * (size_t)(deg + 1); + for (int k = 0; k <= deg; ++k) dst[k] += fac * an[k]; + } + } + + /* phi_1(q,s) = sum_m c[1,m](s) q^m + c[1,m] = - sum_(n>=m+1) (-1)^(n-1-m) / (h^(n-1) m! (n-1-m)!) b_n(s) */ + if (f->ncoef > 1) { + for (int m = 0; m <= nb - 1; ++m) { + double *dst = cptr(f, 1, m + f->moff); + for (int n = m + 1; n <= nb; ++n) { + double sgn = ((n - 1 - m) & 1) ? -1.0 : 1.0; + double fac = -sgn * invhpow[n - 1] * invfact[m] * invfact[n - 1 - m]; + const double *bn = b + (size_t)(n - 1) * (size_t)(deg + 1); + for (int k = 0; k <= deg; ++k) dst[k] += fac * bn[k]; + } + } + } + + /* Recursion: c[i+2,m] = -(d_s^2 + h^2 (m+2)^2) c[i,m+2] + implemented for polynomial expansion of c[i,m] in powers of s + C[i+2,m,k] = -(C[i,m+2,k+2]*(k+2)*(k+1) + C[i,m+2,k]*h^2*(m+2)^2) */ + for (int i = 0; i + 2 < f->ncoef; ++i) { + for (int m = f->mmin; m <= f->mmax - 2; ++m) { + const double *src = ccptr(f, i, (m + 2) + f->moff); + double *dst = cptr(f, i + 2, m + f->moff); + double lam = h * h * (double)(m + 2) * (double)(m + 2); + for (int k = 0; k <= deg; ++k) { + double v = lam * src[k]; + if (k + 2 <= deg) v += (double)(k + 2) * (double)(k + 1) * src[k + 2]; + dst[k] = -v; + } + } + } + + return f; + +} + + +void free_expansion(Expansion *f) { + if (!f) return; + free(f->c); + free(f->V); + free(f->D1); + free(f->D2); + free(f->Q); + free(f); +} + + +#endif + diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion.h b/xtrack/beam_elements/elements_src/track_fieldexpansion.h new file mode 100644 index 000000000..c7b993218 --- /dev/null +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion.h @@ -0,0 +1,95 @@ +#ifndef XTRACK_TRACK_FIELDEXPANSION_H +#define XTRACK_TRACK_FIELDEXPANSION_H + +typedef struct { + double phi; + double Bx, By, Bs; + double Ax, Ay, As; + double dAx_dx, dAx_dy, dAx_ds; + double dAs_dx, dAs_dy, dAs_ds; +} FieldValue; + +typedef struct { + double H; + double delta; + double one_plus_delta; + double radicand; + double root; + double grad[6]; /* dH/d{x,px,y,py,tau,ptau} */ + double rhs[6]; /* canonical flow dz/ds */ + double dH_ds; /* explicit derivative at fixed canonical variables */ + FieldValue pot; +} HamiltonianFlow; + + +void evaluate_expansion(Expansion *f, double x, double y, double s, FieldValue *out) { +} + +void hamiltonian_flow(Expansion *f, const double beta0, + double s, const double z[6], HamiltonianFlow *flow) { + memset(flow, 0, sizeof(*flow)); +} + +void FieldExpansion_track_local_particle( + FieldExpansionData el, + LocalParticle* part) +{ + + // HOW CAN I INITIALIZE AND KEEP THIS? + Expansion f; + + HamiltonianFlow flow; + + const double nstep = FieldExpansionData_get_nstep(el); + const double ds = FieldExpansionData_get_ds(el); + printf("nstep = %e\n", nstep); + const double beta0 = LocalParticle_get_beta0(part); + + const double x = LocalParticle_get_x(part); + const double px = LocalParticle_get_px(part); + const double y = LocalParticle_get_y(part); + const double py = LocalParticle_get_py(part); + const double tau = LocalParticle_get_zeta(part) / beta0; + const double ptau = LocalParticle_get_ptau(part); + + // TO BE DONE + // Momentum has to be continuous, vector potential discontinuous, update canonical momentum + // FieldValue v; + // evaluate_expansion(f, z[0], z[2], s, &v); + // z[1] += v.Ax; + // z[3] += v.Ay; + + double s = 0; + double z[6] = {x, px, y, py, tau, ptau}; + double ztmp[6]; + for (int step = 0; step < nstep; ++step) { + double k1[6], k2[6], k3[6], k4[6]; + + hamiltonian_flow(&f, beta0, s, z, &flow); + memcpy(k1, flow.rhs, sizeof(k1)); + for (int i = 0; i < 6; ++i) ztmp[i] = z[i] + 0.5 * ds * k1[i]; + + hamiltonian_flow(&f, beta0, s + 0.5 * ds, ztmp, &flow); + memcpy(k2, flow.rhs, sizeof(k2)); + for (int i = 0; i < 6; ++i) ztmp[i] = z[i] + 0.5 * ds * k2[i]; + + hamiltonian_flow(&f, beta0, s + 0.5 * ds, ztmp, &flow); + memcpy(k3, flow.rhs, sizeof(k3)); + for (int i = 0; i < 6; ++i) ztmp[i] = z[i] + ds * k3[i]; + + hamiltonian_flow(&f, beta0, s + ds, ztmp, &flow); + memcpy(k4, flow.rhs, sizeof(k4)); + for (int i = 0; i < 6; ++i) z[i] += ds * (k1[i] + 2.0*k2[i] + 2.0*k3[i] + k4[i]) / 6.0; + + s += ds; + } + + LocalParticle_set_x(part, z[0]); + LocalParticle_set_px(part, z[1]); + LocalParticle_set_y(part, z[2]); + LocalParticle_set_py(part, z[3]); + LocalParticle_set_zeta(part, z[4]*beta0); + LocalParticle_set_ptau(part, z[5]); +} + +#endif \ No newline at end of file From ec100f5d523d438a86687dbefec165db30a1bf5f Mon Sep 17 00:00:00 2001 From: Silke Date: Mon, 8 Jun 2026 13:45:13 +0200 Subject: [PATCH 02/38] Temporarely moved Expansion struct to tracker --- .../elements_src/create_fieldexpansion.h | 14 -------------- .../elements_src/track_fieldexpansion.h | 14 ++++++++++++++ 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/xtrack/beam_elements/elements_src/create_fieldexpansion.h b/xtrack/beam_elements/elements_src/create_fieldexpansion.h index 52cc36b46..e7909fade 100644 --- a/xtrack/beam_elements/elements_src/create_fieldexpansion.h +++ b/xtrack/beam_elements/elements_src/create_fieldexpansion.h @@ -1,20 +1,6 @@ #ifndef create_fieldexpansion_H #define create_fieldexpansion_H -typedef struct { - int ny; /* requested output order in y */ - int ncoef; /* stored phi_i coefficients: 0..ny+1 */ - int na, nb, deg; - int mmin, mmax, moff, nm; - int qemin, nq; - double h; - double *c; /* c[i,m,k], polynomial coeff of s^k in q^m term */ - double *V; /* scratch: c[i,m](s) */ - double *D1; /* scratch: d_s c[i,m] */ - double *D2; /* scratch: d2_s c[i,m] */ - double *Q; /* scratch: q^e, e=qemin.. */ -} Expansion; - void build_expansion( double h, int64_t ny, diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion.h b/xtrack/beam_elements/elements_src/track_fieldexpansion.h index c7b993218..766d40f01 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion.h +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion.h @@ -1,6 +1,20 @@ #ifndef XTRACK_TRACK_FIELDEXPANSION_H #define XTRACK_TRACK_FIELDEXPANSION_H +typedef struct { + int ny; /* requested output order in y */ + int ncoef; /* stored phi_i coefficients: 0..ny+1 */ + int na, nb, deg; + int mmin, mmax, moff, nm; + int qemin, nq; + double h; + double *c; /* c[i,m,k], polynomial coeff of s^k in q^m term */ + double *V; /* scratch: c[i,m](s) */ + double *D1; /* scratch: d_s c[i,m] */ + double *D2; /* scratch: d2_s c[i,m] */ + double *Q; /* scratch: q^e, e=qemin.. */ +} Expansion; + typedef struct { double phi; double Bx, By, Bs; From 917754aba0b2d4b61478ca3a7a8f1968536c5a44 Mon Sep 17 00:00:00 2001 From: Silke Date: Mon, 8 Jun 2026 13:58:02 +0200 Subject: [PATCH 03/38] Remove print and helper function --- .../elements_src/create_fieldexpansion.h | 12 ++++++------ .../elements_src/track_fieldexpansion.h | 1 - 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/xtrack/beam_elements/elements_src/create_fieldexpansion.h b/xtrack/beam_elements/elements_src/create_fieldexpansion.h index e7909fade..2d54c706f 100644 --- a/xtrack/beam_elements/elements_src/create_fieldexpansion.h +++ b/xtrack/beam_elements/elements_src/create_fieldexpansion.h @@ -12,7 +12,7 @@ void build_expansion( const double *b, ) { - Expansion *f = (Expansion *)xcalloc(1, sizeof(*f)); + Expansion *f = (Expansion *)calloc(1, sizeof(*f)); f->h = h; f->ny = ny; f->ncoef = ny + 2; /* store phi_0..phi_{ny+1} so By is also order ny */ @@ -25,11 +25,11 @@ void build_expansion( f->nm = f->mmax - f->mmin + 1; f->qemin = f->mmin - 1; f->nq = (f->mmax + 2) - f->qemin + 1; - f->c = (double *)xcalloc((size_t)f->ncoef * (size_t)f->nm * (size_t)(deg + 1), sizeof(double)); - f->V = (double *)xcalloc((size_t)f->ncoef * (size_t)f->nm, sizeof(double)); - f->D1 = (double *)xcalloc((size_t)f->ncoef * (size_t)f->nm, sizeof(double)); - f->D2 = (double *)xcalloc((size_t)f->ncoef * (size_t)f->nm, sizeof(double)); - f->Q = (double *)xcalloc((size_t)f->nq, sizeof(double)); + f->c = (double *)calloc((size_t)f->ncoef * (size_t)f->nm * (size_t)(deg + 1), sizeof(double)); + f->V = (double *)calloc((size_t)f->ncoef * (size_t)f->nm, sizeof(double)); + f->D1 = (double *)calloc((size_t)f->ncoef * (size_t)f->nm, sizeof(double)); + f->D2 = (double *)calloc((size_t)f->ncoef * (size_t)f->nm, sizeof(double)); + f->Q = (double *)calloc((size_t)f->nq, sizeof(double)); int nmax = (na > nb) ? na : nb; double invfact[nmax + 1]; diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion.h b/xtrack/beam_elements/elements_src/track_fieldexpansion.h index 766d40f01..489985f8d 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion.h +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion.h @@ -56,7 +56,6 @@ void FieldExpansion_track_local_particle( const double nstep = FieldExpansionData_get_nstep(el); const double ds = FieldExpansionData_get_ds(el); - printf("nstep = %e\n", nstep); const double beta0 = LocalParticle_get_beta0(part); const double x = LocalParticle_get_x(part); From a97cb42cc16263ee23dd8b43e43bdd4eac7c129e Mon Sep 17 00:00:00 2001 From: Silke Date: Mon, 8 Jun 2026 18:05:47 +0200 Subject: [PATCH 04/38] Progress on evaluating field expansion from python, keeping c values to be solved --- test_fieldexpansion_element.py | 2 +- xtrack/beam_elements/elements.py | 73 ++++++++++++---- .../elements_src/create_fieldexpansion.h | 83 +++++++++---------- 3 files changed, 95 insertions(+), 63 deletions(-) diff --git a/test_fieldexpansion_element.py b/test_fieldexpansion_element.py index 540d38bed..4feb2062e 100644 --- a/test_fieldexpansion_element.py +++ b/test_fieldexpansion_element.py @@ -8,7 +8,7 @@ ny = 5 length=1 -line = xt.Line(elements=[xt.FieldExpansion(length=length, h=h, aa=a, bb=b, bs=bs, ny=ny)]) +line = xt.Line(elements=[xt.FieldExpansion(length=length, h=h, a=a, b=b, bs=bs, ny=ny)]) p = xt.Particles(x=0.01) line.track(p) print(p.x) \ No newline at end of file diff --git a/xtrack/beam_elements/elements.py b/xtrack/beam_elements/elements.py index 71facae99..f9eac53d1 100644 --- a/xtrack/beam_elements/elements.py +++ b/xtrack/beam_elements/elements.py @@ -5003,6 +5003,7 @@ class FieldExpansion(BeamElement): number of powers in y to include, """ + isthick = True behaves_like_drift = True has_backtrack = False @@ -5012,35 +5013,73 @@ class FieldExpansion(BeamElement): _xofields = { "length" : xo.Float64, "h": xo.Float64, - "ny": xo.Int64, "deg": xo.Int64, - "nstep": xo.Int64, "ds": xo.Float64, + + "a": xo.Float64[:], + "b": xo.Float64[:], + "bs": xo.Float64[:], + + "na": xo.Int64, + "nb": xo.Int64, + "deg": xo.Int64, + + "_ncoef": xo.Int64, + "_mmax": xo.Int64, + "_mmin": xo.Int64, + "_moff": xo.Int64, + "_nm": xo.Int64, + + "_qemin": xo.Int64, + "_nq": xo.Int64, + + "_c": xo.Float64[:], + } _extra_c_sources = [ '#include "xtrack/beam_elements/elements_src/track_fieldexpansion.h"', + '#include "xtrack/beam_elements/elements_src/create_fieldexpansion.h"', ] - def __init__(self, length, h, aa, bb, bs, ny, nstep=10, **kwargs): - super().__init__(**kwargs) - - self.length = length - self.h = h - self.nstep = nstep - self.ds = length/nstep + _kernels = {'build_expansion': xo.Kernel( + c_name='build_expansion', + args=[xo.Arg(xo.ThisClass, name='el')] + ) + } + + def __init__(self, length, h, a, b, bs, ny, nstep=10, **kwargs): + kwargs['length'] = length + kwargs['h'] = h + kwargs['nstep'] = nstep + kwargs['ds'] = length/nstep - self.aa = np.asarray(aa, dtype=np.float64) - self.bb = np.asarray(bb, dtype=np.float64) - self.bs = np.asarray(bs, dtype=np.float64) + kwargs['a'] = np.asarray(a, dtype=np.float64).flatten() + kwargs['b'] = np.asarray(b, dtype=np.float64).flatten() + kwargs['bs'] = np.asarray(bs, dtype=np.float64).flatten() - self.na = aa.shape[0] - self.nb = bb.shape[0] + kwargs['na'] = a.shape[0] + kwargs['nb'] = b.shape[0] + kwargs['ny'] = ny - self.deg = aa.shape[1] - 1 + kwargs['deg'] = a.shape[1] - 1 - if bb.shape[1] != self.deg + 1 or bs.shape[0] != self.deg + 1: + if b.shape[1] != kwargs['deg'] + 1 or bs.shape[0] != kwargs['deg'] + 1: raise ValueError("Invalid input shapes") - \ No newline at end of file + + kwargs['_ncoef'] = kwargs['ny'] + 2 # store phi_0..phi_{ny+1} so By is also order ny + kwargs['_mmax'] = kwargs['na'] if kwargs['na'] > (kwargs['nb'] - 1) else (kwargs['nb'] - 1) + kwargs['_mmin'] = -2 * ((kwargs['_ncoef'] - 1) // 2) + kwargs['_moff'] = -kwargs['_mmin'] + kwargs['_nm'] = kwargs['_mmax'] - kwargs['_mmin'] + 1 + + kwargs['_qemin'] = kwargs['_mmin'] - 1 + kwargs['_nq'] = (kwargs['_mmax'] + 2) - kwargs['_qemin'] + 1 + + kwargs.setdefault("_c", np.zeros(kwargs['_ncoef'] * kwargs['_nm'] * (kwargs['deg'] + 1))) + + super().__init__(**kwargs) + + self.build_expansion(el=self) \ No newline at end of file diff --git a/xtrack/beam_elements/elements_src/create_fieldexpansion.h b/xtrack/beam_elements/elements_src/create_fieldexpansion.h index 2d54c706f..1b154736c 100644 --- a/xtrack/beam_elements/elements_src/create_fieldexpansion.h +++ b/xtrack/beam_elements/elements_src/create_fieldexpansion.h @@ -1,35 +1,36 @@ #ifndef create_fieldexpansion_H #define create_fieldexpansion_H + + void build_expansion( - double h, - int64_t ny, - int64_t na, - int64_t nb, - int64_t deg, - const double *bs, - const double *a, - const double *b, + FieldExpansionData el ) { - Expansion *f = (Expansion *)calloc(1, sizeof(*f)); - f->h = h; - f->ny = ny; - f->ncoef = ny + 2; /* store phi_0..phi_{ny+1} so By is also order ny */ - f->na = na; - f->nb = nb; - f->deg = deg; - f->mmax = (na > nb - 1) ? na : (nb - 1); - f->mmin = -2 * ((f->ncoef - 1) / 2); - f->moff = -f->mmin; - f->nm = f->mmax - f->mmin + 1; - f->qemin = f->mmin - 1; - f->nq = (f->mmax + 2) - f->qemin + 1; - f->c = (double *)calloc((size_t)f->ncoef * (size_t)f->nm * (size_t)(deg + 1), sizeof(double)); - f->V = (double *)calloc((size_t)f->ncoef * (size_t)f->nm, sizeof(double)); - f->D1 = (double *)calloc((size_t)f->ncoef * (size_t)f->nm, sizeof(double)); - f->D2 = (double *)calloc((size_t)f->ncoef * (size_t)f->nm, sizeof(double)); - f->Q = (double *)calloc((size_t)f->nq, sizeof(double)); + const double h = FieldExpansionData_get_h(el); + const int ncoef = FieldExpansionData_get__ncoef(el); + const int na = FieldExpansionData_get_na(el); + const int nb = FieldExpansionData_get_nb(el); + const int deg = FieldExpansionData_get_deg(el); + double a[na * (deg + 1)]; + for (int i = 0; i < na*(deg+1); ++i){ + a[i] = FieldExpansionData_get_a(el,i); + } + double b[nb * (deg + 1)]; + for (int i = 0; i < nb*(deg+1); ++i){ + b[i] = FieldExpansionData_get_b(el,i); + } + double bs[deg + 1]; + for (int i = 0; i < deg + 1; ++i){ + bs[i] = FieldExpansionData_get_bs(el,i); + } + + const int mmax = FieldExpansionData_get__mmax(el); + const int mmin = FieldExpansionData_get__mmin(el); + const int moff = FieldExpansionData_get__moff(el); + const int nm = FieldExpansionData_get__nm(el); + + double c[ncoef * nm * (deg + 1)]; int nmax = (na > nb) ? na : nb; double invfact[nmax + 1]; @@ -44,7 +45,7 @@ void build_expansion( /* phi_0(s) = sum_m c[0,m](s) q^m c[0,m] = - sum_(n>=m) (-1)^(n-m) / (h^n m! (n-m)!) a_n(s) */ for (int m = 0; m <= na; ++m) { - double *dst = cptr(f, 0, m + f->moff); + double *dst = cptr(f, 0, m + moff); /* a_0(s)=int_0^s b_s(u)du contributes -a_0 to phi_0. CAREFUL: this will neglect the highest order in the polynomial, only up to given degree in a0 is kept */ @@ -62,9 +63,9 @@ void build_expansion( /* phi_1(q,s) = sum_m c[1,m](s) q^m c[1,m] = - sum_(n>=m+1) (-1)^(n-1-m) / (h^(n-1) m! (n-1-m)!) b_n(s) */ - if (f->ncoef > 1) { + if (ncoef > 1) { for (int m = 0; m <= nb - 1; ++m) { - double *dst = cptr(f, 1, m + f->moff); + double *dst = cptr(f, 1, m + moff); for (int n = m + 1; n <= nb; ++n) { double sgn = ((n - 1 - m) & 1) ? -1.0 : 1.0; double fac = -sgn * invhpow[n - 1] * invfact[m] * invfact[n - 1 - m]; @@ -77,10 +78,11 @@ void build_expansion( /* Recursion: c[i+2,m] = -(d_s^2 + h^2 (m+2)^2) c[i,m+2] implemented for polynomial expansion of c[i,m] in powers of s C[i+2,m,k] = -(C[i,m+2,k+2]*(k+2)*(k+1) + C[i,m+2,k]*h^2*(m+2)^2) */ - for (int i = 0; i + 2 < f->ncoef; ++i) { - for (int m = f->mmin; m <= f->mmax - 2; ++m) { - const double *src = ccptr(f, i, (m + 2) + f->moff); - double *dst = cptr(f, i + 2, m + f->moff); + for (int i = 0; i + 2 < ncoef; ++i) { + for (int m = mmin; m <= mmax - 2; ++m) { + + const double *src = ccptr(f, i, (m + 2) + moff); /* memory location of C[i,m+2,0], taking into account that the minimal value is not zero by moff */ + double *dst = cptr(f, i + 2, m + moff); /* memory location of C[i+2,m,0] */ double lam = h * h * (double)(m + 2) * (double)(m + 2); for (int k = 0; k <= deg; ++k) { double v = lam * src[k]; @@ -90,19 +92,10 @@ void build_expansion( } } - return f; - -} - + for (int i = 0; i < ncoef * nm * (deg+1); ++i){ + FieldExpansionData_set__c(el, c[i], i); + } -void free_expansion(Expansion *f) { - if (!f) return; - free(f->c); - free(f->V); - free(f->D1); - free(f->D2); - free(f->Q); - free(f); } From 2289cfed23402c23bbb433ec60f37b75471337c2 Mon Sep 17 00:00:00 2001 From: Silke Date: Tue, 9 Jun 2026 12:44:02 +0200 Subject: [PATCH 05/38] Fieldexpansion creation is called and compiles, values to be checked --- test_fieldexpansion_element.py | 10 +++--- xtrack/beam_elements/elements.py | 2 +- .../elements_src/create_fieldexpansion.h | 35 +++++++++++-------- 3 files changed, 27 insertions(+), 20 deletions(-) diff --git a/test_fieldexpansion_element.py b/test_fieldexpansion_element.py index 4feb2062e..d68d78c9b 100644 --- a/test_fieldexpansion_element.py +++ b/test_fieldexpansion_element.py @@ -2,13 +2,15 @@ import numpy as np h = 0.2 -a = np.array([[1,1]]) -b = np.array([[1,2]]) -bs = np.array([1,5]) +a = np.array([[1,0.1], [0.2, 0], [0.3, 0]]) +b = np.array([[1,0.1], [0.5, 0]]) +bs = np.array([0.1,0]) ny = 5 length=1 +fexp = xt.FieldExpansion(length=length, h=h, a=a, b=b, bs=bs, ny=ny) +fexp -line = xt.Line(elements=[xt.FieldExpansion(length=length, h=h, a=a, b=b, bs=bs, ny=ny)]) +line = xt.Line(elements=[fexp]) p = xt.Particles(x=0.01) line.track(p) print(p.x) \ No newline at end of file diff --git a/xtrack/beam_elements/elements.py b/xtrack/beam_elements/elements.py index f9eac53d1..a0c6a2cfc 100644 --- a/xtrack/beam_elements/elements.py +++ b/xtrack/beam_elements/elements.py @@ -5058,7 +5058,7 @@ def __init__(self, length, h, a, b, bs, ny, nstep=10, **kwargs): kwargs['a'] = np.asarray(a, dtype=np.float64).flatten() kwargs['b'] = np.asarray(b, dtype=np.float64).flatten() - kwargs['bs'] = np.asarray(bs, dtype=np.float64).flatten() + kwargs['bs'] = np.asarray(bs, dtype=np.float64) kwargs['na'] = a.shape[0] kwargs['nb'] = b.shape[0] diff --git a/xtrack/beam_elements/elements_src/create_fieldexpansion.h b/xtrack/beam_elements/elements_src/create_fieldexpansion.h index 1b154736c..f167b2773 100644 --- a/xtrack/beam_elements/elements_src/create_fieldexpansion.h +++ b/xtrack/beam_elements/elements_src/create_fieldexpansion.h @@ -1,6 +1,17 @@ #ifndef create_fieldexpansion_H #define create_fieldexpansion_H +/* Index of c[i,m,k] in the c array, ordered as +c[0,mmin,0] ... c[0,mmin,deg], c[0,mmin+1,0] ... c[0,mmin+1,deg], ..., c[0,mmin+nm-1,0] ... c[0,mmin+nm-1,deg], +c[1,mmin,0] ... c[1,mmin,deg], c[1,mmin+1,0] ... c[1,mmin+1,deg], ..., c[1,mmin+nm-1,0] ... c[1,mmin+nm-1,deg], +... c[ncoef-1, mmin+nm-1, deg] +moff=-mmin is the offset to be added to m to get the correct index, +since m does not necessarily start at 0 +*/ + +const int cidx(int i, int m, int k, int nm, int moff, int deg) { + return (i * nm + (m+moff)) * (deg + 1) + k; +} void build_expansion( @@ -31,6 +42,7 @@ void build_expansion( const int nm = FieldExpansionData_get__nm(el); double c[ncoef * nm * (deg + 1)]; + memset(c, 0, sizeof(c)); int nmax = (na > nb) ? na : nb; double invfact[nmax + 1]; @@ -45,32 +57,30 @@ void build_expansion( /* phi_0(s) = sum_m c[0,m](s) q^m c[0,m] = - sum_(n>=m) (-1)^(n-m) / (h^n m! (n-m)!) a_n(s) */ for (int m = 0; m <= na; ++m) { - double *dst = cptr(f, 0, m + moff); /* a_0(s)=int_0^s b_s(u)du contributes -a_0 to phi_0. CAREFUL: this will neglect the highest order in the polynomial, only up to given degree in a0 is kept */ if (m==0) { for (int k = 0; k < deg; ++k) - dst[k + 1] = bs[k] / (double)(k + 1); + c[cidx(0,0,k+1,nm,moff,deg)] = bs[k] / (double)(k + 1); } for (int n = (m > 1 ? m : 1); n <= na; ++n) { double sgn = ((n - m) & 1) ? -1.0 : 1.0; double fac = -sgn * invhpow[n] * invfact[m] * invfact[n - m]; const double *an = a + (size_t)(n - 1) * (size_t)(deg + 1); - for (int k = 0; k <= deg; ++k) dst[k] += fac * an[k]; + for (int k = 0; k <= deg; ++k) c[cidx(0,m,k,nm,moff,deg)] += fac * an[k]; } } - + /* phi_1(q,s) = sum_m c[1,m](s) q^m c[1,m] = - sum_(n>=m+1) (-1)^(n-1-m) / (h^(n-1) m! (n-1-m)!) b_n(s) */ if (ncoef > 1) { for (int m = 0; m <= nb - 1; ++m) { - double *dst = cptr(f, 1, m + moff); for (int n = m + 1; n <= nb; ++n) { double sgn = ((n - 1 - m) & 1) ? -1.0 : 1.0; double fac = -sgn * invhpow[n - 1] * invfact[m] * invfact[n - 1 - m]; const double *bn = b + (size_t)(n - 1) * (size_t)(deg + 1); - for (int k = 0; k <= deg; ++k) dst[k] += fac * bn[k]; + for (int k = 0; k <= deg; ++k) c[cidx(1,m,k,nm,moff,deg)] += fac * bn[k]; } } } @@ -80,24 +90,19 @@ void build_expansion( C[i+2,m,k] = -(C[i,m+2,k+2]*(k+2)*(k+1) + C[i,m+2,k]*h^2*(m+2)^2) */ for (int i = 0; i + 2 < ncoef; ++i) { for (int m = mmin; m <= mmax - 2; ++m) { - - const double *src = ccptr(f, i, (m + 2) + moff); /* memory location of C[i,m+2,0], taking into account that the minimal value is not zero by moff */ - double *dst = cptr(f, i + 2, m + moff); /* memory location of C[i+2,m,0] */ double lam = h * h * (double)(m + 2) * (double)(m + 2); for (int k = 0; k <= deg; ++k) { - double v = lam * src[k]; - if (k + 2 <= deg) v += (double)(k + 2) * (double)(k + 1) * src[k + 2]; - dst[k] = -v; + double v = lam * c[cidx(i,m+2,k,nm,moff,deg)]; + if (k + 2 <= deg) v += (double)(k + 2) * (double)(k + 1) * c[cidx(i,m+2,k+2,nm,moff,deg)]; + c[cidx(i+2,m,k,nm,moff,deg)] = -v; } } } for (int i = 0; i < ncoef * nm * (deg+1); ++i){ - FieldExpansionData_set__c(el, c[i], i); + FieldExpansionData_set__c(el, i, c[i]); } - } - #endif From b31fbea94521d956640b40ad9ca54179110b6c34 Mon Sep 17 00:00:00 2001 From: Silke Date: Tue, 9 Jun 2026 18:03:38 +0200 Subject: [PATCH 06/38] Seems to work! --- test_fieldexpansion_element.py | 11 +- xtrack/beam_elements/elements.py | 9 + .../elements_src/create_fieldexpansion.h | 11 +- .../elements_src/track_fieldexpansion.h | 220 ++++++++++++++++-- 4 files changed, 221 insertions(+), 30 deletions(-) diff --git a/test_fieldexpansion_element.py b/test_fieldexpansion_element.py index d68d78c9b..de346bcee 100644 --- a/test_fieldexpansion_element.py +++ b/test_fieldexpansion_element.py @@ -2,15 +2,14 @@ import numpy as np h = 0.2 -a = np.array([[1,0.1], [0.2, 0], [0.3, 0]]) -b = np.array([[1,0.1], [0.5, 0]]) -bs = np.array([0.1,0]) +a = np.array([[1,0.1,-0.2],[1,0.1,-0.2]]) +b = np.array([[1,0.1,-0.2],[1,0.1,-0.2]]) +bs = np.array([0.1,0.2,0.1]) ny = 5 -length=1 +length=0.01 fexp = xt.FieldExpansion(length=length, h=h, a=a, b=b, bs=bs, ny=ny) -fexp line = xt.Line(elements=[fexp]) p = xt.Particles(x=0.01) -line.track(p) +line.track(p, _force_no_end_turn_actions=True) print(p.x) \ No newline at end of file diff --git a/xtrack/beam_elements/elements.py b/xtrack/beam_elements/elements.py index a0c6a2cfc..825d1a33f 100644 --- a/xtrack/beam_elements/elements.py +++ b/xtrack/beam_elements/elements.py @@ -5036,6 +5036,10 @@ class FieldExpansion(BeamElement): "_nq": xo.Int64, "_c": xo.Float64[:], + "_V": xo.Float64[:], + "_D1": xo.Float64[:], + "_D2": xo.Float64[:], + "_Q": xo.Float64[:], } @@ -5080,6 +5084,11 @@ def __init__(self, length, h, a, b, bs, ny, nstep=10, **kwargs): kwargs.setdefault("_c", np.zeros(kwargs['_ncoef'] * kwargs['_nm'] * (kwargs['deg'] + 1))) + kwargs.setdefault("_V", np.zeros(kwargs['_ncoef'] * kwargs['_nm'])) + kwargs.setdefault("_D1", np.zeros(kwargs['_ncoef'] * kwargs['_nm'])) + kwargs.setdefault("_D2", np.zeros(kwargs['_ncoef'] * kwargs['_nm'])) + kwargs.setdefault("_Q", np.zeros(kwargs['_nq'])) + super().__init__(**kwargs) self.build_expansion(el=self) \ No newline at end of file diff --git a/xtrack/beam_elements/elements_src/create_fieldexpansion.h b/xtrack/beam_elements/elements_src/create_fieldexpansion.h index f167b2773..193b674c0 100644 --- a/xtrack/beam_elements/elements_src/create_fieldexpansion.h +++ b/xtrack/beam_elements/elements_src/create_fieldexpansion.h @@ -41,8 +41,9 @@ void build_expansion( const int moff = FieldExpansionData_get__moff(el); const int nm = FieldExpansionData_get__nm(el); - double c[ncoef * nm * (deg + 1)]; - memset(c, 0, sizeof(c)); + //double c[ncoef * nm * (deg + 1)]; + double *c = (double *) FieldExpansionData_getp__c(el); + //memset(c, 0, sizeof(c)); int nmax = (na > nb) ? na : nb; double invfact[nmax + 1]; @@ -99,9 +100,9 @@ void build_expansion( } } - for (int i = 0; i < ncoef * nm * (deg+1); ++i){ - FieldExpansionData_set__c(el, i, c[i]); - } + // for (int i = 0; i < ncoef * nm * (deg+1); ++i){ + // FieldExpansionData_set__c(el, i, c[i]); + // } } #endif diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion.h b/xtrack/beam_elements/elements_src/track_fieldexpansion.h index 489985f8d..ba5475b63 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion.h +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion.h @@ -35,13 +35,173 @@ typedef struct { FieldValue pot; } HamiltonianFlow; +static inline const double *ccptr(const Expansion *f, int i, int m) { + return f->c + (((size_t)i * (size_t)f->nm + (size_t)m) * (size_t)(f->deg + 1)); +} + +static inline void poly_eval_d2(const double *p, int deg, double s, double *v, double *d1, double *d2) { + double a = p[deg], b = 0.0, c=0.0; + for (int k = deg - 1; k >= 0; --k) { + c = c * s + 2.0 * b; + b = b * s + a; + a = a * s + p[k]; + } + *v = a; + *d1 = b; + *d2 = c; +} + +static void fs_prepare_s(Expansion *f, double s) { + for (int i = 0; i < f->ncoef; ++i) { + for (int m = 0; m < f->nm; ++m) { + poly_eval_d2(ccptr(f, i, m), f->deg, s, + &f->V[i * f->nm + m], + &f->D1[i * f->nm + m], + &f->D2[i * f->nm + m]); + } + } +} -void evaluate_expansion(Expansion *f, double x, double y, double s, FieldValue *out) { +void delta_from_ptau(const double beta0, double ptau, + double *delta, double *delta1, double *ddelta1) { + { + const double r = 1.0 + 2.0 * ptau / beta0 + ptau * ptau; + *delta1 = sqrt(r); + *delta = *delta1 - 1.0; + *ddelta1 = (1.0 / beta0 + ptau) / (*delta1); + } +} + + + +int evaluate_expansion(Expansion *f, double x, double y, double s, FieldValue *out) { + const double q = 1.0 + f->h * x; + if (q == 0.0) return -1; /* singular chart */ + + memset(out, 0, sizeof(*out)); + + double *V = f->V; + double *D1 = f->D1; + double *D2 = f->D2; + double *Q = f->Q; + + fs_prepare_s(f, s); + + /* q powers from e = mmin-1 .. mmax+2 */ + Q[0] = pow(q, (double)f->qemin); + for (int t = 1; t < f->nq; ++t) Q[t] = Q[t - 1] * q; + #define QPOW(E) Q[(E) - f->qemin] + + /* As(x,0,s) + = 1/(1+hx) int_0^x dx' *(1+hx) By(x',0,s) + = 1/qh int_1^q dq' q' phi_1(q',s) + = 1/qh sum_m c[1,m] q^(m+2)/(m+2) - 1/q sum_m c[1,m] 1/(m+2) */ + if (f->ncoef > 1) { + for (int m = 0; m <= f->mmax; ++m) { + int j = m + f->moff; + const double c1m = V[1 * f->nm + j]; /* c[1,m] */ + const double dc1m = f->D1[1 * f->nm + j]; /* c[1,m]'*/ + const double g = QPOW(m + 1) - QPOW(-1); + const double den = f->h * (double)(m + 2); + if (c1m != 0.0) { + out->As += c1m * g / den; + out->dAs_dx += c1m * (((double)(m + 1)) * QPOW(m) + QPOW(-2)) / (double)(m + 2); + } + if (dc1m != 0.0) out->dAs_ds += dc1m * g / den; + } + } + + double yi = 1.0; /* y^i / i! */ + for (int i = 0; i <= f->ny; ++i) { + double sphi = 0.0, gx = 0.0, gs = 0.0, gy = 0.0; + double dgx_dx = 0.0, dgx_ds = 0.0; + double dgs_dx = 0.0, dgs_ds = 0.0; + + for (int m = f->mmin; m <= f->mmax; ++m) { + const int j = m + f->moff; + const double cim = V[i * f->nm + j]; /* c[i,m] */ + const double ci1m = V[(i + 1) * f->nm + j]; /* c[i+1,m] */ + const double dcim = D1[i * f->nm + j]; /* c[i,m]' */ + const double ddcim = D2[i * f->nm + j]; /* c[i,m]'' */ + const double qm = QPOW(m); /* q^m */ + const double qm1 = QPOW(m - 1); /* q^(m-1) */ + const double qm2 = QPOW(m - 2); /* q^(m-2) */ + + sphi += cim * qm; /* c[i,m] q^m */ + gx += f->h * (double)m * cim * qm1; /* h m c[i,m] q^(m-1) */ + gy += ci1m * qm; /* c[i+1,m] q^m */ + gs += dcim * qm1; /* c[i,m]' q^(m-1) */ + + dgx_dx += f->h * f->h * (double)m * (double)(m-1) * cim * qm2; + dgx_ds += f->h * (double)m * dcim * qm1; + dgs_dx += f->h * (double)(m-1) * dcim * qm2; + dgs_ds += ddcim * qm1; + } + + out->phi += sphi * yi; /* c[i,m] q^m y^i/i!*/ + out->Bx -= gx * yi; /* -h m c[i,m] q^(m-1) y^i/i! */ + out->By -= gy * yi; /* -c[i+1,m] q^m y^i/i! */ + out->Bs -= gs * yi; /* -c[i,m]' q^(m-1) y^i/i! */ + + /* A_x, A_s through order ny in y: need i = 0..ny-1 */ + if (i < f->ny) { + double yi1 = yi * y / (double)(i + 1); /* y^(i+1)/(i+1)! */ + out->Ax += gs * yi1; /* -c[i,m]' q^(m-1) y^(i+1)/(i+1)! */ + out->As -= gx * yi1; /* -h m c[i,m] q^(m-1) y^i/i! *//* -h m c[i,m] q^(m-1) y^i/i! */ + out->dAx_dx += dgs_dx * yi1; + out->dAx_ds += dgs_ds * yi1; + out->dAs_dx += -dgx_dx * yi1; + out->dAs_ds += -dgx_ds * yi1; + } + + yi *= y / (double)(i + 1); + } + + out->dAx_dy = -out->Bs; + out->dAs_dy = out->Bx; + + return 0; + #undef QPOW } void hamiltonian_flow(Expansion *f, const double beta0, double s, const double z[6], HamiltonianFlow *flow) { - memset(flow, 0, sizeof(*flow)); + double delta1, delta, ddelta1; + double q, pix, piy, rad, root; + + memset(flow, 0, sizeof(*flow)); + + evaluate_expansion(f, z[0], z[2], s, &flow->pot); + delta_from_ptau(beta0, z[5], &delta, &delta1, &ddelta1); + + q = 1.0 + f->h * z[0]; + pix = z[1] - flow->pot.Ax; + piy = z[3] - flow->pot.Ay; /* A_y is zero in this gauge. */ + rad = delta1 * delta1 - pix * pix - piy * piy; + root = sqrt(rad); + + flow->delta = delta; + flow->one_plus_delta = delta1; + flow->radicand = rad; + flow->root = root; + flow->H = z[5] / beta0 - q * (root + flow->pot.As); + + flow->rhs[0] = q * pix / root; // dx/ds = dH/dpx + flow->rhs[2] = q * piy / root; // dy/ds = dH/dpy + flow->rhs[4] = 1.0 / beta0 - q * delta1 * ddelta1 / root; // dtau/ds = dH/dptau + + flow->rhs[1] = f->h * (root + flow->pot.As) + + q * (pix * flow->pot.dAx_dx / root + flow->pot.dAs_dx); // dpx/ds = -dH/dx + flow->rhs[3] = q * (pix * flow->pot.dAx_dy / root + flow->pot.dAs_dy); // dpy/ds = -dH/dy + flow->rhs[5] = 0.0; // tptau/ds = -dH/dtau, H has no tau-dependence for these static fields. + + flow->grad[0] = -flow->rhs[1]; // dH/dx + flow->grad[1] = flow->rhs[0]; // dH/dpx + flow->grad[2] = -flow->rhs[3]; // dH/dy + flow->grad[3] = flow->rhs[2]; // dH/dpy + flow->grad[4] = -flow->rhs[5]; // dH/dtau + flow->grad[5] = flow->rhs[4]; // dH/dptau + flow->dH_ds = -q * (pix * flow->pot.dAx_ds / root + flow->pot.dAs_ds); } void FieldExpansion_track_local_particle( @@ -49,31 +209,52 @@ void FieldExpansion_track_local_particle( LocalParticle* part) { - // HOW CAN I INITIALIZE AND KEEP THIS? + const double nstep = FieldExpansionData_get_nstep(el); + const double ds = FieldExpansionData_get_ds(el); + const double beta0 = LocalParticle_get_beta0(part); + + const double x = LocalParticle_get_x(part); + const double px = LocalParticle_get_px(part); + const double y = LocalParticle_get_y(part); + const double py = LocalParticle_get_py(part); + const double tau = LocalParticle_get_zeta(part) / beta0; + const double ptau = LocalParticle_get_ptau(part); + double z[6] = {x, px, y, py, tau, ptau}; + Expansion f; - HamiltonianFlow flow; + f.ny = FieldExpansionData_get_ny(el); + f.ncoef = FieldExpansionData_get__ncoef(el); + + f.na = FieldExpansionData_get_na(el); + f.nb = FieldExpansionData_get_nb(el); + f.deg = FieldExpansionData_get_deg(el); - const double nstep = FieldExpansionData_get_nstep(el); - const double ds = FieldExpansionData_get_ds(el); - const double beta0 = LocalParticle_get_beta0(part); + f.mmin = FieldExpansionData_get__mmin(el); + f.mmax = FieldExpansionData_get__mmax(el); + f.moff = FieldExpansionData_get__moff(el); + f.nm = FieldExpansionData_get__nm(el); - const double x = LocalParticle_get_x(part); - const double px = LocalParticle_get_px(part); - const double y = LocalParticle_get_y(part); - const double py = LocalParticle_get_py(part); - const double tau = LocalParticle_get_zeta(part) / beta0; - const double ptau = LocalParticle_get_ptau(part); + f.qemin = FieldExpansionData_get__qemin(el); + f.nq = FieldExpansionData_get__nq(el); + + f.h = FieldExpansionData_get_h(el); + f.c = (double *)FieldExpansionData_getp__c(el); + + f.V = (double *)FieldExpansionData_getp__V(el); + f.D1 = (double *)FieldExpansionData_getp__D1(el); + f.D2 = (double *)FieldExpansionData_getp__D2(el); + f.Q = (double *)FieldExpansionData_getp__Q(el); + + HamiltonianFlow flow; - // TO BE DONE // Momentum has to be continuous, vector potential discontinuous, update canonical momentum - // FieldValue v; - // evaluate_expansion(f, z[0], z[2], s, &v); - // z[1] += v.Ax; - // z[3] += v.Ay; + FieldValue v; + evaluate_expansion(&f, z[0], z[2], 0, &v); + z[1] += v.Ax; + z[3] += v.Ay; double s = 0; - double z[6] = {x, px, y, py, tau, ptau}; double ztmp[6]; for (int step = 0; step < nstep; ++step) { double k1[6], k2[6], k3[6], k4[6]; @@ -103,6 +284,7 @@ void FieldExpansion_track_local_particle( LocalParticle_set_py(part, z[3]); LocalParticle_set_zeta(part, z[4]*beta0); LocalParticle_set_ptau(part, z[5]); + LocalParticle_add_to_s(part, ds*nstep); } #endif \ No newline at end of file From c09bdbbbb7a9ad5d2360551fb3880a443c1a8cf2 Mon Sep 17 00:00:00 2001 From: Silke Date: Tue, 9 Jun 2026 18:14:51 +0200 Subject: [PATCH 07/38] Taking into account vector potential discontinuities, added per particle block --- .../elements_src/track_fieldexpansion.h | 110 +++++++++--------- 1 file changed, 58 insertions(+), 52 deletions(-) diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion.h b/xtrack/beam_elements/elements_src/track_fieldexpansion.h index ba5475b63..b7c229702 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion.h +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion.h @@ -206,41 +206,26 @@ void hamiltonian_flow(Expansion *f, const double beta0, void FieldExpansion_track_local_particle( FieldExpansionData el, - LocalParticle* part) + LocalParticle* part0) { const double nstep = FieldExpansionData_get_nstep(el); const double ds = FieldExpansionData_get_ds(el); - const double beta0 = LocalParticle_get_beta0(part); - - const double x = LocalParticle_get_x(part); - const double px = LocalParticle_get_px(part); - const double y = LocalParticle_get_y(part); - const double py = LocalParticle_get_py(part); - const double tau = LocalParticle_get_zeta(part) / beta0; - const double ptau = LocalParticle_get_ptau(part); - double z[6] = {x, px, y, py, tau, ptau}; Expansion f; - f.ny = FieldExpansionData_get_ny(el); f.ncoef = FieldExpansionData_get__ncoef(el); - f.na = FieldExpansionData_get_na(el); f.nb = FieldExpansionData_get_nb(el); f.deg = FieldExpansionData_get_deg(el); - f.mmin = FieldExpansionData_get__mmin(el); f.mmax = FieldExpansionData_get__mmax(el); f.moff = FieldExpansionData_get__moff(el); f.nm = FieldExpansionData_get__nm(el); - f.qemin = FieldExpansionData_get__qemin(el); f.nq = FieldExpansionData_get__nq(el); - f.h = FieldExpansionData_get_h(el); f.c = (double *)FieldExpansionData_getp__c(el); - f.V = (double *)FieldExpansionData_getp__V(el); f.D1 = (double *)FieldExpansionData_getp__D1(el); f.D2 = (double *)FieldExpansionData_getp__D2(el); @@ -248,43 +233,64 @@ void FieldExpansion_track_local_particle( HamiltonianFlow flow; - // Momentum has to be continuous, vector potential discontinuous, update canonical momentum - FieldValue v; - evaluate_expansion(&f, z[0], z[2], 0, &v); - z[1] += v.Ax; - z[3] += v.Ay; - - double s = 0; - double ztmp[6]; - for (int step = 0; step < nstep; ++step) { - double k1[6], k2[6], k3[6], k4[6]; - - hamiltonian_flow(&f, beta0, s, z, &flow); - memcpy(k1, flow.rhs, sizeof(k1)); - for (int i = 0; i < 6; ++i) ztmp[i] = z[i] + 0.5 * ds * k1[i]; - - hamiltonian_flow(&f, beta0, s + 0.5 * ds, ztmp, &flow); - memcpy(k2, flow.rhs, sizeof(k2)); - for (int i = 0; i < 6; ++i) ztmp[i] = z[i] + 0.5 * ds * k2[i]; - - hamiltonian_flow(&f, beta0, s + 0.5 * ds, ztmp, &flow); - memcpy(k3, flow.rhs, sizeof(k3)); - for (int i = 0; i < 6; ++i) ztmp[i] = z[i] + ds * k3[i]; - - hamiltonian_flow(&f, beta0, s + ds, ztmp, &flow); - memcpy(k4, flow.rhs, sizeof(k4)); - for (int i = 0; i < 6; ++i) z[i] += ds * (k1[i] + 2.0*k2[i] + 2.0*k3[i] + k4[i]) / 6.0; - - s += ds; - } + START_PER_PARTICLE_BLOCK(part0, part); + const double beta0 = LocalParticle_get_beta0(part); + + const double x = LocalParticle_get_x(part); + const double px = LocalParticle_get_px(part); + const double y = LocalParticle_get_y(part); + const double py = LocalParticle_get_py(part); + const double tau = LocalParticle_get_zeta(part) / beta0; + const double ptau = LocalParticle_get_ptau(part); + const double ax = LocalParticle_get_ax(part); + const double ay = LocalParticle_get_ay(part); + double z[6] = {x, px, y, py, tau, ptau}; + + // Momentum has to be continuous, vector potential discontinuous, update canonical momentum + FieldValue v; + evaluate_expansion(&f, z[0], z[2], 0, &v); + z[1] += v.Ax - ax; + z[3] += v.Ay - ay; + + double s = 0; + double ztmp[6]; + for (int step = 0; step < nstep; ++step) { + double k1[6], k2[6], k3[6], k4[6]; + + hamiltonian_flow(&f, beta0, s, z, &flow); + memcpy(k1, flow.rhs, sizeof(k1)); + for (int i = 0; i < 6; ++i) ztmp[i] = z[i] + 0.5 * ds * k1[i]; + + hamiltonian_flow(&f, beta0, s + 0.5 * ds, ztmp, &flow); + memcpy(k2, flow.rhs, sizeof(k2)); + for (int i = 0; i < 6; ++i) ztmp[i] = z[i] + 0.5 * ds * k2[i]; + + hamiltonian_flow(&f, beta0, s + 0.5 * ds, ztmp, &flow); + memcpy(k3, flow.rhs, sizeof(k3)); + for (int i = 0; i < 6; ++i) ztmp[i] = z[i] + ds * k3[i]; + + hamiltonian_flow(&f, beta0, s + ds, ztmp, &flow); + memcpy(k4, flow.rhs, sizeof(k4)); + for (int i = 0; i < 6; ++i) z[i] += ds * (k1[i] + 2.0*k2[i] + 2.0*k3[i] + k4[i]) / 6.0; + + s += ds; + } - LocalParticle_set_x(part, z[0]); - LocalParticle_set_px(part, z[1]); - LocalParticle_set_y(part, z[2]); - LocalParticle_set_py(part, z[3]); - LocalParticle_set_zeta(part, z[4]*beta0); - LocalParticle_set_ptau(part, z[5]); - LocalParticle_add_to_s(part, ds*nstep); + // Back to zero vector potential for next element + evaluate_expansion(&f, z[0], z[2], 0, &v); + z[1] -= v.Ax; + z[3] -= v.Ay; + + LocalParticle_set_x(part, z[0]); + LocalParticle_set_px(part, z[1]); + LocalParticle_set_y(part, z[2]); + LocalParticle_set_py(part, z[3]); + LocalParticle_set_zeta(part, z[4]*beta0); + LocalParticle_set_ptau(part, z[5]); + LocalParticle_set_ax(part, 0); + LocalParticle_set_ay(part, 0); + LocalParticle_add_to_s(part, ds*nstep); + END_PER_PARTICLE_BLOCK } #endif \ No newline at end of file From 0a08b911821c7b1dbba0dafbb458a12e7bdba6ba Mon Sep 17 00:00:00 2001 From: Silke Date: Wed, 10 Jun 2026 13:10:34 +0200 Subject: [PATCH 08/38] Implemented cartesian case, to be benchmarked in more detail --- test_fieldexpansion_element.py | 10 +- xtrack/beam_elements/elements.py | 19 ++- .../elements_src/create_fieldexpansion.h | 91 +++++++++-- .../elements_src/track_fieldexpansion.h | 151 ++++++++++++++---- 4 files changed, 215 insertions(+), 56 deletions(-) diff --git a/test_fieldexpansion_element.py b/test_fieldexpansion_element.py index de346bcee..e8bf93c59 100644 --- a/test_fieldexpansion_element.py +++ b/test_fieldexpansion_element.py @@ -1,13 +1,13 @@ import xtrack as xt import numpy as np -h = 0.2 -a = np.array([[1,0.1,-0.2],[1,0.1,-0.2]]) -b = np.array([[1,0.1,-0.2],[1,0.1,-0.2]]) -bs = np.array([0.1,0.2,0.1]) +h = 0 +a = np.array([[0]]) +b = np.array([[1]]) +bs = np.array([0]) ny = 5 length=0.01 -fexp = xt.FieldExpansion(length=length, h=h, a=a, b=b, bs=bs, ny=ny) +fexp = xt.FieldExpansion(length=length, h=h, a=a, b=b, bs=bs, ny=ny, nstep=10) line = xt.Line(elements=[fexp]) p = xt.Particles(x=0.01) diff --git a/xtrack/beam_elements/elements.py b/xtrack/beam_elements/elements.py index 825d1a33f..5cdb712c7 100644 --- a/xtrack/beam_elements/elements.py +++ b/xtrack/beam_elements/elements.py @@ -5013,6 +5013,7 @@ class FieldExpansion(BeamElement): _xofields = { "length" : xo.Float64, "h": xo.Float64, + "cartesian": xo.Int64, "ny": xo.Int64, "deg": xo.Int64, "nstep": xo.Int64, @@ -5051,12 +5052,13 @@ class FieldExpansion(BeamElement): _kernels = {'build_expansion': xo.Kernel( c_name='build_expansion', args=[xo.Arg(xo.ThisClass, name='el')] - ) + ), } def __init__(self, length, h, a, b, bs, ny, nstep=10, **kwargs): kwargs['length'] = length kwargs['h'] = h + kwargs['cartesian'] = h == 0 kwargs['nstep'] = nstep kwargs['ds'] = length/nstep @@ -5074,13 +5076,18 @@ def __init__(self, length, h, a, b, bs, ny, nstep=10, **kwargs): raise ValueError("Invalid input shapes") kwargs['_ncoef'] = kwargs['ny'] + 2 # store phi_0..phi_{ny+1} so By is also order ny + kwargs['_mmax'] = kwargs['na'] if kwargs['na'] > (kwargs['nb'] - 1) else (kwargs['nb'] - 1) - kwargs['_mmin'] = -2 * ((kwargs['_ncoef'] - 1) // 2) - kwargs['_moff'] = -kwargs['_mmin'] - kwargs['_nm'] = kwargs['_mmax'] - kwargs['_mmin'] + 1 - - kwargs['_qemin'] = kwargs['_mmin'] - 1 + if kwargs['cartesian']: + kwargs['_mmin'] = 0 + kwargs['_moff'] = 0 + kwargs['_qemin'] = 0 # CHECK + else: + kwargs['_mmin'] = -2 * ((kwargs['_ncoef'] - 1) // 2) + kwargs['_moff'] = -kwargs['_mmin'] + kwargs['_qemin'] = kwargs['_mmin'] - 1 #CHECK kwargs['_nq'] = (kwargs['_mmax'] + 2) - kwargs['_qemin'] + 1 + kwargs['_nm'] = kwargs['_mmax'] - kwargs['_mmin'] + 1 kwargs.setdefault("_c", np.zeros(kwargs['_ncoef'] * kwargs['_nm'] * (kwargs['deg'] + 1))) diff --git a/xtrack/beam_elements/elements_src/create_fieldexpansion.h b/xtrack/beam_elements/elements_src/create_fieldexpansion.h index 193b674c0..410703c1d 100644 --- a/xtrack/beam_elements/elements_src/create_fieldexpansion.h +++ b/xtrack/beam_elements/elements_src/create_fieldexpansion.h @@ -13,11 +13,71 @@ const int cidx(int i, int m, int k, int nm, int moff, int deg) { return (i * nm + (m+moff)) * (deg + 1) + k; } +void build_cartesian_expansion(FieldExpansionData el){ + const double h = FieldExpansionData_get_h(el); + const int ncoef = FieldExpansionData_get__ncoef(el); + const int na = FieldExpansionData_get_na(el); + const int nb = FieldExpansionData_get_nb(el); + const int deg = FieldExpansionData_get_deg(el); + double a[na * (deg + 1)]; + for (int i = 0; i < na*(deg+1); ++i){ + a[i] = FieldExpansionData_get_a(el,i); + } + double b[nb * (deg + 1)]; + for (int i = 0; i < nb*(deg+1); ++i){ + b[i] = FieldExpansionData_get_b(el,i); + } + double bs[deg + 1]; + for (int i = 0; i < deg + 1; ++i){ + bs[i] = FieldExpansionData_get_bs(el,i); + } + + const int mmax = FieldExpansionData_get__mmax(el); + const int moff = FieldExpansionData_get__moff(el); + const int nm = FieldExpansionData_get__nm(el); + + double *c = (double *) FieldExpansionData_getp__c(el); + + int nmax = (na > nb) ? na : nb; + double invfact[nmax + 1]; + double invhpow[nmax + 1]; + invfact[0] = 1.0; + invhpow[0] = 1.0; + for (int n = 1; n <= nmax; ++n) { + invfact[n] = invfact[n - 1] / (double)n; + invhpow[n] = invhpow[n - 1] / h; + } + + for (int k = 0; k < deg; ++k) c[cidx(0,0,k+1,nm,moff,deg)] = bs[k] / (double)(k + 1); -void build_expansion( - FieldExpansionData el -) -{ + for (int n = 1; n <= na; ++n) { + const double fac = -invfact[n]; + const double *an = a + (size_t)(n - 1) * (size_t)(deg + 1); + for (int k = 0; k <= deg; ++k) c[cidx(0,n,k,nm,moff,deg)] += fac * an[k]; + } + if (ncoef > 1) { + for (int n = 1; n <= nb; ++n) { + const double fac = -invfact[n - 1]; + const double *bn = b + (size_t)(n - 1) * (size_t)(deg + 1); + for (int k = 0; k <= deg; ++k) c[cidx(1,n-1,k,nm,moff,deg)] += fac * bn[k]; + } + } + + for (int i = 0; i + 2 < ncoef; ++i) { + for (int m = 0; m <= mmax; ++m) { + for (int k = 0; k <= deg; ++k) { + double v = 0.0; + if (m + 2 < mmax) + v += (double)(m + 2) * (double)(m + 1) * c[cidx(i,m+2,k,nm,moff,deg)]; + if (k + 2 <= deg) + v += (double)(k + 2) * (double)(k + 1) * c[cidx(i,m,k+2,nm,moff,deg)]; + c[cidx(i+2,m,k,nm,moff,deg)] = -v; + } + } + } +} + +void build_bent_expansion(FieldExpansionData el){ const double h = FieldExpansionData_get_h(el); const int ncoef = FieldExpansionData_get__ncoef(el); const int na = FieldExpansionData_get_na(el); @@ -41,9 +101,7 @@ void build_expansion( const int moff = FieldExpansionData_get__moff(el); const int nm = FieldExpansionData_get__nm(el); - //double c[ncoef * nm * (deg + 1)]; double *c = (double *) FieldExpansionData_getp__c(el); - //memset(c, 0, sizeof(c)); int nmax = (na > nb) ? na : nb; double invfact[nmax + 1]; @@ -55,16 +113,13 @@ void build_expansion( invhpow[n] = invhpow[n - 1] / h; } + /* a_0(s)=int_0^s b_s(u)du contributes -a_0 to phi_0. + CAREFUL: this will neglect the highest order in the polynomial, + only up to given degree in a0 is kept */ + for (int k = 0; k < deg; ++k) c[cidx(0,0,k+1,nm,moff,deg)] = bs[k] / (double)(k + 1); /* phi_0(s) = sum_m c[0,m](s) q^m c[0,m] = - sum_(n>=m) (-1)^(n-m) / (h^n m! (n-m)!) a_n(s) */ for (int m = 0; m <= na; ++m) { - /* a_0(s)=int_0^s b_s(u)du contributes -a_0 to phi_0. - CAREFUL: this will neglect the highest order in the polynomial, - only up to given degree in a0 is kept */ - if (m==0) { - for (int k = 0; k < deg; ++k) - c[cidx(0,0,k+1,nm,moff,deg)] = bs[k] / (double)(k + 1); - } for (int n = (m > 1 ? m : 1); n <= na; ++n) { double sgn = ((n - m) & 1) ? -1.0 : 1.0; double fac = -sgn * invhpow[n] * invfact[m] * invfact[n - m]; @@ -99,10 +154,14 @@ void build_expansion( } } } +} - // for (int i = 0; i < ncoef * nm * (deg+1); ++i){ - // FieldExpansionData_set__c(el, i, c[i]); - // } +void build_expansion(FieldExpansionData el){ + if (FieldExpansionData_get_cartesian(el)) { + build_cartesian_expansion(el); + } else { + build_bent_expansion(el); + } } #endif diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion.h b/xtrack/beam_elements/elements_src/track_fieldexpansion.h index b7c229702..68f8f5a0d 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion.h +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion.h @@ -8,6 +8,7 @@ typedef struct { int mmin, mmax, moff, nm; int qemin, nq; double h; + double cartesian; double *c; /* c[i,m,k], polynomial coeff of s^k in q^m term */ double *V; /* scratch: c[i,m](s) */ double *D1; /* scratch: d_s c[i,m] */ @@ -72,11 +73,93 @@ void delta_from_ptau(const double beta0, double ptau, } } +int evaluate_cartesian_expansion(Expansion *f, double x, double y, double s, FieldValue *out) { + memset(out, 0, sizeof(*out)); -int evaluate_expansion(Expansion *f, double x, double y, double s, FieldValue *out) { + double *V = f->V; + double *D1 = f->D1; + double *D2 = f->D2; + double *X = f->Q; + + fs_prepare_s(f, s); + + /* x powers */ + X[0] = pow(x, (double)f->qemin); + for (int t = 1; t < f->nq; ++t) X[t] = X[t - 1] * x; + #define XPOW(E) X[(E) - f->qemin] + + /* As(x,0,s) */ + if (f->ncoef > 1) { + for (int m = 0; m <= f->mmax; ++m) { + int j = m + f->moff; + const double c1m = V[1 * f->nm + j]; /* c[1,m] */ + const double dc1m = f->D1[1 * f->nm + j]; /* c[1,m]'*/ + const double xp = XPOW(m + 1) / (double)(m + 1); + out->As += c1m * xp; + out->dAs_ds += dc1m * xp; + out->dAs_dx += c1m * XPOW(m); + } + } + + double yi = 1.0; /* y^i / i! */ + for (int i = 0; i <= f->ny; ++i) { + double sphi = 0.0, gx = 0.0, gs = 0.0, gy = 0.0; + double dgx_dx = 0.0, dgx_ds = 0.0; + double dgs_dx = 0.0, dgs_ds = 0.0; + + for (int m = f->mmin; m <= f->mmax; ++m) { + const int j = m + f->moff; + const double cim = V[i * f->nm + j]; /* c[i,m] */ + const double ci1m = V[(i + 1) * f->nm + j]; /* c[i+1,m] */ + const double dcim = D1[i * f->nm + j]; /* c[i,m]' */ + const double ddcim = D2[i * f->nm + j]; /* c[i,m]'' */ + + const double xm = XPOW(m); /* x^m */ + const double xm1 = XPOW(m - 1); /* x^(m-1) */ + const double xm2 = XPOW(m - 2); /* x^(m-2) */ + + sphi += cim * xm; + gx += (double)m * cim * xm1; + gy += ci1m * xm; + gs += dcim * xm; + + dgx_dx += (double)m * (double)(m - 1) * cim * xm2; + dgx_ds += (double)m * dcim * xm1; + dgs_dx += (double)m * dcim * xm1; + dgs_ds += ddcim * xm; + } + + out->phi += sphi * yi; /* c[i,m] x^m y^i/i!*/ + out->Bx -= gx * yi; /* -h m c[i,m] x^(m-1) y^i/i! */ + out->By -= gy * yi; /* -c[i+1,m] x^m y^i/i! */ + out->Bs -= gs * yi; /* -c[i,m]' x^(m-1) y^i/i! */ + + /* A_x, A_s through order ny in y: need i = 0..ny-1 */ + if (i < f->ny) { + double yi1 = yi * y / (double)(i + 1); /* y^(i+1)/(i+1)! */ + out->Ax += gs * yi1; /* -c[i,m]' x^(m-1) y^(i+1)/(i+1)! */ + out->As -= gx * yi1; /* -h m c[i,m] x^(m-1) y^i/i! */ + out->dAx_dx += dgs_dx * yi1; + out->dAx_ds += dgs_ds * yi1; + out->dAs_dx += -dgx_dx * yi1; + out->dAs_ds += -dgx_ds * yi1; + } + + yi *= y / (double)(i + 1); + } + + out->dAx_dy = -out->Bs; + out->dAs_dy = out->Bx; + + return 0; + #undef QPOW +} + + +int evaluate_bent_expansion(Expansion *f, double x, double y, double s, FieldValue *out) { const double q = 1.0 + f->h * x; - if (q == 0.0) return -1; /* singular chart */ + if (q == 0.0) return -1; /* singular chart */ memset(out, 0, sizeof(*out)); @@ -104,10 +187,10 @@ int evaluate_expansion(Expansion *f, double x, double y, double s, FieldValue *o const double g = QPOW(m + 1) - QPOW(-1); const double den = f->h * (double)(m + 2); if (c1m != 0.0) { - out->As += c1m * g / den; + out->As += c1m * g / den; out->dAs_dx += c1m * (((double)(m + 1)) * QPOW(m) + QPOW(-2)) / (double)(m + 2); } - if (dc1m != 0.0) out->dAs_ds += dc1m * g / den; + if (dc1m != 0.0) out->dAs_ds += dc1m * g / den; } } @@ -119,17 +202,18 @@ int evaluate_expansion(Expansion *f, double x, double y, double s, FieldValue *o for (int m = f->mmin; m <= f->mmax; ++m) { const int j = m + f->moff; - const double cim = V[i * f->nm + j]; /* c[i,m] */ - const double ci1m = V[(i + 1) * f->nm + j]; /* c[i+1,m] */ - const double dcim = D1[i * f->nm + j]; /* c[i,m]' */ - const double ddcim = D2[i * f->nm + j]; /* c[i,m]'' */ - const double qm = QPOW(m); /* q^m */ + const double cim = V[i * f->nm + j]; /* c[i,m] */ + const double ci1m = V[(i + 1) * f->nm + j]; /* c[i+1,m] */ + const double dcim = D1[i * f->nm + j]; /* c[i,m]' */ + const double ddcim = D2[i * f->nm + j]; /* c[i,m]'' */ + + const double qm = QPOW(m); /* q^m */ const double qm1 = QPOW(m - 1); /* q^(m-1) */ const double qm2 = QPOW(m - 2); /* q^(m-2) */ sphi += cim * qm; /* c[i,m] q^m */ gx += f->h * (double)m * cim * qm1; /* h m c[i,m] q^(m-1) */ - gy += ci1m * qm; /* c[i+1,m] q^m */ + gy += ci1m * qm; /* c[i+1,m] q^m */ gs += dcim * qm1; /* c[i,m]' q^(m-1) */ dgx_dx += f->h * f->h * (double)m * (double)(m-1) * cim * qm2; @@ -146,8 +230,8 @@ int evaluate_expansion(Expansion *f, double x, double y, double s, FieldValue *o /* A_x, A_s through order ny in y: need i = 0..ny-1 */ if (i < f->ny) { double yi1 = yi * y / (double)(i + 1); /* y^(i+1)/(i+1)! */ - out->Ax += gs * yi1; /* -c[i,m]' q^(m-1) y^(i+1)/(i+1)! */ - out->As -= gx * yi1; /* -h m c[i,m] q^(m-1) y^i/i! *//* -h m c[i,m] q^(m-1) y^i/i! */ + out->Ax += gs * yi1; /* -c[i,m]' q^(m-1) y^(i+1)/(i+1)! */ + out->As -= gx * yi1; /* -h m c[i,m] q^(m-1) y^i/i! */ out->dAx_dx += dgs_dx * yi1; out->dAx_ds += dgs_ds * yi1; out->dAs_dx += -dgx_dx * yi1; @@ -164,6 +248,14 @@ int evaluate_expansion(Expansion *f, double x, double y, double s, FieldValue *o #undef QPOW } +int evaluate_expansion(Expansion *f, double x, double y, double s, FieldValue *out) { + if (f->cartesian) { + return evaluate_cartesian_expansion(f, x, y, s, out); + } else { + return evaluate_bent_expansion(f, x, y, s, out); + } +} + void hamiltonian_flow(Expansion *f, const double beta0, double s, const double z[6], HamiltonianFlow *flow) { double delta1, delta, ddelta1; @@ -213,23 +305,24 @@ void FieldExpansion_track_local_particle( const double ds = FieldExpansionData_get_ds(el); Expansion f; - f.ny = FieldExpansionData_get_ny(el); - f.ncoef = FieldExpansionData_get__ncoef(el); - f.na = FieldExpansionData_get_na(el); - f.nb = FieldExpansionData_get_nb(el); - f.deg = FieldExpansionData_get_deg(el); - f.mmin = FieldExpansionData_get__mmin(el); - f.mmax = FieldExpansionData_get__mmax(el); - f.moff = FieldExpansionData_get__moff(el); - f.nm = FieldExpansionData_get__nm(el); - f.qemin = FieldExpansionData_get__qemin(el); - f.nq = FieldExpansionData_get__nq(el); - f.h = FieldExpansionData_get_h(el); - f.c = (double *)FieldExpansionData_getp__c(el); - f.V = (double *)FieldExpansionData_getp__V(el); - f.D1 = (double *)FieldExpansionData_getp__D1(el); - f.D2 = (double *)FieldExpansionData_getp__D2(el); - f.Q = (double *)FieldExpansionData_getp__Q(el); + f.ny = FieldExpansionData_get_ny(el); + f.ncoef = FieldExpansionData_get__ncoef(el); + f.na = FieldExpansionData_get_na(el); + f.nb = FieldExpansionData_get_nb(el); + f.deg = FieldExpansionData_get_deg(el); + f.mmin = FieldExpansionData_get__mmin(el); + f.mmax = FieldExpansionData_get__mmax(el); + f.moff = FieldExpansionData_get__moff(el); + f.nm = FieldExpansionData_get__nm(el); + f.qemin = FieldExpansionData_get__qemin(el); + f.nq = FieldExpansionData_get__nq(el); + f.h = FieldExpansionData_get_h(el); + f.cartesian = FieldExpansionData_get_cartesian(el); + f.c = (double *)FieldExpansionData_getp__c(el); + f.V = (double *)FieldExpansionData_getp__V(el); + f.D1 = (double *)FieldExpansionData_getp__D1(el); + f.D2 = (double *)FieldExpansionData_getp__D2(el); + f.Q = (double *)FieldExpansionData_getp__Q(el); HamiltonianFlow flow; From 7ea7b01c1eb20336232a5c8f4f8775df0ae0e0e8 Mon Sep 17 00:00:00 2001 From: Silke Date: Wed, 10 Jun 2026 17:36:24 +0200 Subject: [PATCH 09/38] Added test for curved tracker with many derivatives --- test_fieldexpansion_element.py | 15 ----- tests/test_fieldexpansion_element.py | 55 +++++++++++++++++++ xtrack/beam_elements/elements.py | 4 +- .../elements_src/track_fieldexpansion.h | 6 +- 4 files changed, 60 insertions(+), 20 deletions(-) delete mode 100644 test_fieldexpansion_element.py create mode 100644 tests/test_fieldexpansion_element.py diff --git a/test_fieldexpansion_element.py b/test_fieldexpansion_element.py deleted file mode 100644 index e8bf93c59..000000000 --- a/test_fieldexpansion_element.py +++ /dev/null @@ -1,15 +0,0 @@ -import xtrack as xt -import numpy as np - -h = 0 -a = np.array([[0]]) -b = np.array([[1]]) -bs = np.array([0]) -ny = 5 -length=0.01 -fexp = xt.FieldExpansion(length=length, h=h, a=a, b=b, bs=bs, ny=ny, nstep=10) - -line = xt.Line(elements=[fexp]) -p = xt.Particles(x=0.01) -line.track(p, _force_no_end_turn_actions=True) -print(p.x) \ No newline at end of file diff --git a/tests/test_fieldexpansion_element.py b/tests/test_fieldexpansion_element.py new file mode 100644 index 000000000..46aca80cf --- /dev/null +++ b/tests/test_fieldexpansion_element.py @@ -0,0 +1,55 @@ +import xtrack as xt +import numpy as np + + +def test_sdep(): + h = 0.1 + a = np.array([[1, 0.1],[0.2, 0], [0.3, 0.1]]) + b = np.array([[0.1, 0.1],[0.5, 0]]) + bs = np.array([0.1, 0]) + ny = 5 + length=0.2 + fexp = xt.FieldExpansion(length=length, h=h, a=a, b=b, bs=bs, ny=ny, nstep=100) + + p0 = xt.Particles(x=0.01, y=0.007, tau=0.002, beta0=0.7) + line = xt.Line(elements=[fexp]) + line.track(p0, _force_no_end_turn_actions=True) + + assert np.isclose(p0.x[0], 0.00968437) + assert np.isclose(p0.px[0], -0.00162025) + assert np.isclose(p0.y[0], 0.02753048) + assert np.isclose(p0.py[0], 0.20394359) + assert np.isclose(p0.zeta[0], -0.00020263) + assert np.isclose(p0.ptau[0], 0) + +def test_twiss(): + fodo = xt.Line(elements=[ + xt.Drift(length=1.2), + xt.Quadrupole(k1=7, length=0.1), + xt.Drift(length=0.5), + xt.Bend(length=0.2, k0=0.1, angle=0.1*0.2), + xt.Drift(length=0.5), + xt.Quadrupole(k1=-7, length=0.1)] + ) + fodo.particle_ref = xt.Particles(particle_id=11, q0=1, mass0=1) + tw = fodo.twiss4d() + + myfodo = xt.Line(elements=[ + xt.Drift(length=1.2), + xt.FieldExpansion(length=0.1, h=0, a=np.array([[0]]), b=np.array([[0],[7]]), bs=np.array([0]), ny=5), + xt.Drift(length=0.5), + xt.FieldExpansion(length=0.2, h=0.1, a=np.array([[0]]), b=np.array([[0.1]]), bs=np.array([0]), ny=5), + xt.Drift(length=0.5), + xt.FieldExpansion(length=0.1, h=0, a=np.array([[0]]), b=np.array([[0],[-7]]), bs=np.array([0]), ny=5) + ]) + myfodo.particle_ref = xt.Particles(particle_id=11, q0=1, mass0=1) + mytw = myfodo.twiss4d() + + assert np.allclose(tw.betx, mytw.betx) + assert np.allclose(tw.bety, mytw.bety) + assert np.allclose(tw.alfx, mytw.alfx) + assert np.allclose(tw.alfy, mytw.alfy) + assert np.allclose(tw.dx, mytw.dx) + assert np.allclose(tw.dy, mytw.dy) + + \ No newline at end of file diff --git a/xtrack/beam_elements/elements.py b/xtrack/beam_elements/elements.py index 5cdb712c7..1c6486838 100644 --- a/xtrack/beam_elements/elements.py +++ b/xtrack/beam_elements/elements.py @@ -5081,11 +5081,11 @@ def __init__(self, length, h, a, b, bs, ny, nstep=10, **kwargs): if kwargs['cartesian']: kwargs['_mmin'] = 0 kwargs['_moff'] = 0 - kwargs['_qemin'] = 0 # CHECK + kwargs['_qemin'] = 0 else: kwargs['_mmin'] = -2 * ((kwargs['_ncoef'] - 1) // 2) kwargs['_moff'] = -kwargs['_mmin'] - kwargs['_qemin'] = kwargs['_mmin'] - 1 #CHECK + kwargs['_qemin'] = kwargs['_mmin'] - 1 kwargs['_nq'] = (kwargs['_mmax'] + 2) - kwargs['_qemin'] + 1 kwargs['_nm'] = kwargs['_mmax'] - kwargs['_mmin'] + 1 diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion.h b/xtrack/beam_elements/elements_src/track_fieldexpansion.h index 68f8f5a0d..0fcecf34a 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion.h +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion.h @@ -370,9 +370,9 @@ void FieldExpansion_track_local_particle( } // Back to zero vector potential for next element - evaluate_expansion(&f, z[0], z[2], 0, &v); - z[1] -= v.Ax; - z[3] -= v.Ay; + // evaluate_expansion(&f, z[0], z[2], 0, &v); + // z[1] -= v.Ax; + // z[3] -= v.Ay; LocalParticle_set_x(part, z[0]); LocalParticle_set_px(part, z[1]); From 3f90a17b97140ab9835fc5d7089ed9568e1e06ff Mon Sep 17 00:00:00 2001 From: Silke Date: Wed, 10 Jun 2026 17:42:31 +0200 Subject: [PATCH 10/38] Added straight test, currently failing. Remember to put vector potential correction back and update tests --- tests/test_fieldexpansion_element.py | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/tests/test_fieldexpansion_element.py b/tests/test_fieldexpansion_element.py index 46aca80cf..e4c08d829 100644 --- a/tests/test_fieldexpansion_element.py +++ b/tests/test_fieldexpansion_element.py @@ -1,8 +1,7 @@ import xtrack as xt import numpy as np - -def test_sdep(): +def test_h_sdep(): h = 0.1 a = np.array([[1, 0.1],[0.2, 0], [0.3, 0.1]]) b = np.array([[0.1, 0.1],[0.5, 0]]) @@ -21,6 +20,26 @@ def test_sdep(): assert np.isclose(p0.py[0], 0.20394359) assert np.isclose(p0.zeta[0], -0.00020263) assert np.isclose(p0.ptau[0], 0) + +def test_sdep(): + h = 0 + a = np.array([[1, 0.1],[0.2, 0], [0.3, 0.1]]) + b = np.array([[0.1, 0.1],[0.5, 0]]) + bs = np.array([0.1, 0]) + ny = 5 + length=0.2 + fexp = xt.FieldExpansion(length=length, h=h, a=a, b=b, bs=bs, ny=ny, nstep=100) + + p0 = xt.Particles(x=0.01, y=0.007, tau=0.002, beta0=0.7) + line = xt.Line(elements=[fexp]) + line.track(p0, _force_no_end_turn_actions=True) + + assert np.isclose(p0.x[0], 0.00765466) + assert np.isclose(p0.px[0], -0.02167055) + assert np.isclose(p0.y[0], 0.02747711) + assert np.isclose(p0.py[0], 0.20351503) + assert np.isclose(p0.zeta[0], -1.64816266e-05) + assert np.isclose(p0.ptau[0], 0) def test_twiss(): fodo = xt.Line(elements=[ From 158f72c49d321c3ae71bbd4503b637786b356a63 Mon Sep 17 00:00:00 2001 From: Silke Date: Thu, 11 Jun 2026 16:06:03 +0200 Subject: [PATCH 11/38] Max index corrected --- xtrack/beam_elements/elements_src/create_fieldexpansion.h | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/xtrack/beam_elements/elements_src/create_fieldexpansion.h b/xtrack/beam_elements/elements_src/create_fieldexpansion.h index 410703c1d..190b1acae 100644 --- a/xtrack/beam_elements/elements_src/create_fieldexpansion.h +++ b/xtrack/beam_elements/elements_src/create_fieldexpansion.h @@ -14,6 +14,7 @@ const int cidx(int i, int m, int k, int nm, int moff, int deg) { } void build_cartesian_expansion(FieldExpansionData el){ + printf("cartesian"); const double h = FieldExpansionData_get_h(el); const int ncoef = FieldExpansionData_get__ncoef(el); const int na = FieldExpansionData_get_na(el); @@ -67,9 +68,9 @@ void build_cartesian_expansion(FieldExpansionData el){ for (int m = 0; m <= mmax; ++m) { for (int k = 0; k <= deg; ++k) { double v = 0.0; - if (m + 2 < mmax) + if (m + 2 <= mmax) v += (double)(m + 2) * (double)(m + 1) * c[cidx(i,m+2,k,nm,moff,deg)]; - if (k + 2 <= deg) + if (k + 2 <= deg) v += (double)(k + 2) * (double)(k + 1) * c[cidx(i,m,k+2,nm,moff,deg)]; c[cidx(i+2,m,k,nm,moff,deg)] = -v; } From 41518186cdf0bd9b60e8bec6647a32ee5ae13a9e Mon Sep 17 00:00:00 2001 From: Silke Date: Thu, 11 Jun 2026 16:06:57 +0200 Subject: [PATCH 12/38] put fieldvalue creation outside of particle loop for speed --- xtrack/beam_elements/elements_src/track_fieldexpansion.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion.h b/xtrack/beam_elements/elements_src/track_fieldexpansion.h index 0fcecf34a..1160595a1 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion.h +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion.h @@ -325,6 +325,7 @@ void FieldExpansion_track_local_particle( f.Q = (double *)FieldExpansionData_getp__Q(el); HamiltonianFlow flow; + FieldValue v; START_PER_PARTICLE_BLOCK(part0, part); const double beta0 = LocalParticle_get_beta0(part); @@ -340,7 +341,6 @@ void FieldExpansion_track_local_particle( double z[6] = {x, px, y, py, tau, ptau}; // Momentum has to be continuous, vector potential discontinuous, update canonical momentum - FieldValue v; evaluate_expansion(&f, z[0], z[2], 0, &v); z[1] += v.Ax - ax; z[3] += v.Ay - ay; From b5bb6b6b074040b79213a83660598fd2053024f1 Mon Sep 17 00:00:00 2001 From: Silke Date: Thu, 11 Jun 2026 16:18:48 +0200 Subject: [PATCH 13/38] Test working! --- tests/test_fieldexpansion_element.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/tests/test_fieldexpansion_element.py b/tests/test_fieldexpansion_element.py index e4c08d829..5f2f1244d 100644 --- a/tests/test_fieldexpansion_element.py +++ b/tests/test_fieldexpansion_element.py @@ -3,9 +3,9 @@ def test_h_sdep(): h = 0.1 - a = np.array([[1, 0.1],[0.2, 0], [0.3, 0.1]]) - b = np.array([[0.1, 0.1],[0.5, 0]]) - bs = np.array([0.1, 0]) + a = np.array([[1.0, 0.1], [0.2, 0.0], [0.3, 0.1]]) + b = np.array([[0.1, 0.1], [0.5, 0.0]]) + bs = np.array([0.1, 0.0]) ny = 5 length=0.2 fexp = xt.FieldExpansion(length=length, h=h, a=a, b=b, bs=bs, ny=ny, nstep=100) @@ -20,17 +20,18 @@ def test_h_sdep(): assert np.isclose(p0.py[0], 0.20394359) assert np.isclose(p0.zeta[0], -0.00020263) assert np.isclose(p0.ptau[0], 0) + assert np.isclose(p0.s[0], length) def test_sdep(): h = 0 - a = np.array([[1, 0.1],[0.2, 0], [0.3, 0.1]]) - b = np.array([[0.1, 0.1],[0.5, 0]]) - bs = np.array([0.1, 0]) + a = np.array([[1.0, 0.1], [0.2, 0.0], [0.3, 0.1]]) + b = np.array([[0.1, 0.1], [0.5, 0.0]]) + bs = np.array([0.1, 0.0]) ny = 5 length=0.2 fexp = xt.FieldExpansion(length=length, h=h, a=a, b=b, bs=bs, ny=ny, nstep=100) - p0 = xt.Particles(x=0.01, y=0.007, tau=0.002, beta0=0.7) + p0 = xt.Particles(x=0.01, y=0.007, tau=0.002, beta0=1) line = xt.Line(elements=[fexp]) line.track(p0, _force_no_end_turn_actions=True) @@ -38,7 +39,7 @@ def test_sdep(): assert np.isclose(p0.px[0], -0.02167055) assert np.isclose(p0.y[0], 0.02747711) assert np.isclose(p0.py[0], 0.20351503) - assert np.isclose(p0.zeta[0], -1.64816266e-05) + assert np.isclose(p0.zeta[0], 0.00058352) assert np.isclose(p0.ptau[0], 0) def test_twiss(): From f40317658d5dac59a41a2873abdd3e9562801693 Mon Sep 17 00:00:00 2001 From: Silke Date: Thu, 11 Jun 2026 16:23:07 +0200 Subject: [PATCH 14/38] Updated such that vector potential zero in particles after element --- tests/test_fieldexpansion_element.py | 4 ++-- xtrack/beam_elements/elements_src/track_fieldexpansion.h | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_fieldexpansion_element.py b/tests/test_fieldexpansion_element.py index 5f2f1244d..3906d43f7 100644 --- a/tests/test_fieldexpansion_element.py +++ b/tests/test_fieldexpansion_element.py @@ -15,7 +15,7 @@ def test_h_sdep(): line.track(p0, _force_no_end_turn_actions=True) assert np.isclose(p0.x[0], 0.00968437) - assert np.isclose(p0.px[0], -0.00162025) + assert np.isclose(p0.px[0], -0.00430618) assert np.isclose(p0.y[0], 0.02753048) assert np.isclose(p0.py[0], 0.20394359) assert np.isclose(p0.zeta[0], -0.00020263) @@ -36,7 +36,7 @@ def test_sdep(): line.track(p0, _force_no_end_turn_actions=True) assert np.isclose(p0.x[0], 0.00765466) - assert np.isclose(p0.px[0], -0.02167055) + assert np.isclose(p0.px[0], -0.02435948) assert np.isclose(p0.y[0], 0.02747711) assert np.isclose(p0.py[0], 0.20351503) assert np.isclose(p0.zeta[0], 0.00058352) diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion.h b/xtrack/beam_elements/elements_src/track_fieldexpansion.h index 1160595a1..29516b94b 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion.h +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion.h @@ -370,9 +370,9 @@ void FieldExpansion_track_local_particle( } // Back to zero vector potential for next element - // evaluate_expansion(&f, z[0], z[2], 0, &v); - // z[1] -= v.Ax; - // z[3] -= v.Ay; + evaluate_expansion(&f, z[0], z[2], 0, &v); + z[1] -= v.Ax; + z[3] -= v.Ay; LocalParticle_set_x(part, z[0]); LocalParticle_set_px(part, z[1]); From 5441775c82531c4e2c72b561fe470e439dc09b3c Mon Sep 17 00:00:00 2001 From: Silke Date: Thu, 11 Jun 2026 17:54:29 +0200 Subject: [PATCH 15/38] Splitted in two elements --- tests/test_fieldexpansion_element.py | 11 +- xtrack/beam_elements/elements.py | 129 +++++- .../elements_src/create_fieldexpansion.h | 169 -------- .../elements_src/create_fieldexpansion_bent.h | 97 +++++ .../create_fieldexpansion_straight.h | 80 ++++ .../elements_src/track_fieldexpansion.h | 389 ------------------ .../elements_src/track_fieldexpansion.inc | 131 ++++++ .../elements_src/track_fieldexpansion_bent.h | 111 +++++ .../track_fieldexpansion_helpers.h | 71 ++++ .../track_fieldexpansion_straight.h | 102 +++++ 10 files changed, 719 insertions(+), 571 deletions(-) delete mode 100644 xtrack/beam_elements/elements_src/create_fieldexpansion.h create mode 100644 xtrack/beam_elements/elements_src/create_fieldexpansion_bent.h create mode 100644 xtrack/beam_elements/elements_src/create_fieldexpansion_straight.h delete mode 100644 xtrack/beam_elements/elements_src/track_fieldexpansion.h create mode 100644 xtrack/beam_elements/elements_src/track_fieldexpansion.inc create mode 100644 xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h create mode 100644 xtrack/beam_elements/elements_src/track_fieldexpansion_helpers.h create mode 100644 xtrack/beam_elements/elements_src/track_fieldexpansion_straight.h diff --git a/tests/test_fieldexpansion_element.py b/tests/test_fieldexpansion_element.py index 3906d43f7..2b04cb681 100644 --- a/tests/test_fieldexpansion_element.py +++ b/tests/test_fieldexpansion_element.py @@ -8,7 +8,7 @@ def test_h_sdep(): bs = np.array([0.1, 0.0]) ny = 5 length=0.2 - fexp = xt.FieldExpansion(length=length, h=h, a=a, b=b, bs=bs, ny=ny, nstep=100) + fexp = xt.BentFieldExpansion(length=length, h=h, a=a, b=b, bs=bs, ny=ny, nstep=100) p0 = xt.Particles(x=0.01, y=0.007, tau=0.002, beta0=0.7) line = xt.Line(elements=[fexp]) @@ -23,13 +23,12 @@ def test_h_sdep(): assert np.isclose(p0.s[0], length) def test_sdep(): - h = 0 a = np.array([[1.0, 0.1], [0.2, 0.0], [0.3, 0.1]]) b = np.array([[0.1, 0.1], [0.5, 0.0]]) bs = np.array([0.1, 0.0]) ny = 5 length=0.2 - fexp = xt.FieldExpansion(length=length, h=h, a=a, b=b, bs=bs, ny=ny, nstep=100) + fexp = xt.StraightFieldExpansion(length=length, a=a, b=b, bs=bs, ny=ny, nstep=100) p0 = xt.Particles(x=0.01, y=0.007, tau=0.002, beta0=1) line = xt.Line(elements=[fexp]) @@ -56,11 +55,11 @@ def test_twiss(): myfodo = xt.Line(elements=[ xt.Drift(length=1.2), - xt.FieldExpansion(length=0.1, h=0, a=np.array([[0]]), b=np.array([[0],[7]]), bs=np.array([0]), ny=5), + xt.StraightFieldExpansion(length=0.1, a=np.array([[0]]), b=np.array([[0],[7]]), bs=np.array([0]), ny=5), xt.Drift(length=0.5), - xt.FieldExpansion(length=0.2, h=0.1, a=np.array([[0]]), b=np.array([[0.1]]), bs=np.array([0]), ny=5), + xt.BentFieldExpansion(length=0.2, h=0.1, a=np.array([[0]]), b=np.array([[0.1]]), bs=np.array([0]), ny=5), xt.Drift(length=0.5), - xt.FieldExpansion(length=0.1, h=0, a=np.array([[0]]), b=np.array([[0],[-7]]), bs=np.array([0]), ny=5) + xt.StraightFieldExpansion(length=0.1, a=np.array([[0]]), b=np.array([[0],[-7]]), bs=np.array([0]), ny=5) ]) myfodo.particle_ref = xt.Particles(particle_id=11, q0=1, mass0=1) mytw = myfodo.twiss4d() diff --git a/xtrack/beam_elements/elements.py b/xtrack/beam_elements/elements.py index 1c6486838..1260fd834 100644 --- a/xtrack/beam_elements/elements.py +++ b/xtrack/beam_elements/elements.py @@ -4984,8 +4984,122 @@ def get_backtrack_element(self, _context=None, _buffer=None, _offset=None): class ThinSliceNotNeededError(Exception): pass +class StraightFieldExpansion(BeamElement): + """ + Specifies the field expansion in general derivatives on axis in straight or curved frame + + Parameters + ---------- + h : float + Curvature of the element, in 1/m. For straight elements, h=0. + a : array, shape na, deg+1, floats + describing the polynomial coefficients for the skew multipoles. First index is multipole order, second index is polynomial coefficient. + b : array, shape nb, deg+1, floats + describing the polynomial coefficients for the normal multipoles. First index is multipole order, second index is polynomial coefficient. + bs : array, shape deg+1, floats + describing the polynomial coefficients for the longitudinal field component. Index is polynomial coefficient, + ny : int + number of powers in y to include, + + """ + + isthick = True + behaves_like_drift = True + has_backtrack = False + allow_loss_refinement = False + allow_rot_and_shift = False + + _xofields = { + "length" : xo.Float64, + "h": xo.Float64, + "straight": xo.Int64, + "ny": xo.Int64, + "deg": xo.Int64, + "nstep": xo.Int64, + "ds": xo.Float64, + + "a": xo.Float64[:], + "b": xo.Float64[:], + "bs": xo.Float64[:], + + "na": xo.Int64, + "nb": xo.Int64, + "deg": xo.Int64, + + "_ncoef": xo.Int64, + "_mmax": xo.Int64, + "_mmin": xo.Int64, + "_moff": xo.Int64, + "_nm": xo.Int64, + + "_qemin": xo.Int64, + "_nq": xo.Int64, + + "_c": xo.Float64[:], + "_V": xo.Float64[:], + "_D1": xo.Float64[:], + "_D2": xo.Float64[:], + "_Q": xo.Float64[:], + + } + + _extra_c_sources = [ + '#include "xtrack/beam_elements/elements_src/track_fieldexpansion_straight.h"', + '#include "xtrack/beam_elements/elements_src/create_fieldexpansion_straight.h"', + ] + + _kernels = {'build_expansion': xo.Kernel( + c_name='build_expansion_straight', + args=[xo.Arg(xo.ThisClass, name='el')] + ), + } + + def __init__(self, length, a, b, bs, ny, nstep=10, **kwargs): + kwargs['length'] = length + kwargs['h'] = 0 + kwargs['straight'] = 1 + kwargs['nstep'] = nstep + kwargs['ds'] = length/nstep + + kwargs['a'] = np.asarray(a, dtype=np.float64).flatten() + kwargs['b'] = np.asarray(b, dtype=np.float64).flatten() + kwargs['bs'] = np.asarray(bs, dtype=np.float64) + + kwargs['na'] = a.shape[0] + kwargs['nb'] = b.shape[0] + kwargs['ny'] = ny + + kwargs['deg'] = a.shape[1] - 1 + + if b.shape[1] != kwargs['deg'] + 1 or bs.shape[0] != kwargs['deg'] + 1: + raise ValueError("Invalid input shapes") + + kwargs['_ncoef'] = kwargs['ny'] + 2 # store phi_0..phi_{ny+1} so By is also order ny + + kwargs['_mmax'] = kwargs['na'] if kwargs['na'] > (kwargs['nb'] - 1) else (kwargs['nb'] - 1) + if kwargs['straight']: + kwargs['_mmin'] = 0 + kwargs['_moff'] = 0 + kwargs['_qemin'] = 0 + else: + kwargs['_mmin'] = -2 * ((kwargs['_ncoef'] - 1) // 2) + kwargs['_moff'] = -kwargs['_mmin'] + kwargs['_qemin'] = kwargs['_mmin'] - 1 + kwargs['_nq'] = (kwargs['_mmax'] + 2) - kwargs['_qemin'] + 1 + kwargs['_nm'] = kwargs['_mmax'] - kwargs['_mmin'] + 1 -class FieldExpansion(BeamElement): + kwargs.setdefault("_c", np.zeros(kwargs['_ncoef'] * kwargs['_nm'] * (kwargs['deg'] + 1))) + + kwargs.setdefault("_V", np.zeros(kwargs['_ncoef'] * kwargs['_nm'])) + kwargs.setdefault("_D1", np.zeros(kwargs['_ncoef'] * kwargs['_nm'])) + kwargs.setdefault("_D2", np.zeros(kwargs['_ncoef'] * kwargs['_nm'])) + kwargs.setdefault("_Q", np.zeros(kwargs['_nq'])) + + super().__init__(**kwargs) + + self.build_expansion(el=self) + +class BentFieldExpansion(BeamElement): """ Specifies the field expansion in general derivatives on axis in straight or curved frame @@ -5013,7 +5127,7 @@ class FieldExpansion(BeamElement): _xofields = { "length" : xo.Float64, "h": xo.Float64, - "cartesian": xo.Int64, + "straight": xo.Int64, "ny": xo.Int64, "deg": xo.Int64, "nstep": xo.Int64, @@ -5045,20 +5159,21 @@ class FieldExpansion(BeamElement): } _extra_c_sources = [ - '#include "xtrack/beam_elements/elements_src/track_fieldexpansion.h"', - '#include "xtrack/beam_elements/elements_src/create_fieldexpansion.h"', + '#include "xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h"', + '#include "xtrack/beam_elements/elements_src/create_fieldexpansion_bent.h"', ] _kernels = {'build_expansion': xo.Kernel( - c_name='build_expansion', + c_name='build_expansion_bent', args=[xo.Arg(xo.ThisClass, name='el')] ), } def __init__(self, length, h, a, b, bs, ny, nstep=10, **kwargs): + assert h > 1e-4, "Use straight element with h=0!" kwargs['length'] = length kwargs['h'] = h - kwargs['cartesian'] = h == 0 + kwargs['straight'] = 0 kwargs['nstep'] = nstep kwargs['ds'] = length/nstep @@ -5078,7 +5193,7 @@ def __init__(self, length, h, a, b, bs, ny, nstep=10, **kwargs): kwargs['_ncoef'] = kwargs['ny'] + 2 # store phi_0..phi_{ny+1} so By is also order ny kwargs['_mmax'] = kwargs['na'] if kwargs['na'] > (kwargs['nb'] - 1) else (kwargs['nb'] - 1) - if kwargs['cartesian']: + if kwargs['straight']: kwargs['_mmin'] = 0 kwargs['_moff'] = 0 kwargs['_qemin'] = 0 diff --git a/xtrack/beam_elements/elements_src/create_fieldexpansion.h b/xtrack/beam_elements/elements_src/create_fieldexpansion.h deleted file mode 100644 index 190b1acae..000000000 --- a/xtrack/beam_elements/elements_src/create_fieldexpansion.h +++ /dev/null @@ -1,169 +0,0 @@ -#ifndef create_fieldexpansion_H -#define create_fieldexpansion_H - -/* Index of c[i,m,k] in the c array, ordered as -c[0,mmin,0] ... c[0,mmin,deg], c[0,mmin+1,0] ... c[0,mmin+1,deg], ..., c[0,mmin+nm-1,0] ... c[0,mmin+nm-1,deg], -c[1,mmin,0] ... c[1,mmin,deg], c[1,mmin+1,0] ... c[1,mmin+1,deg], ..., c[1,mmin+nm-1,0] ... c[1,mmin+nm-1,deg], -... c[ncoef-1, mmin+nm-1, deg] -moff=-mmin is the offset to be added to m to get the correct index, -since m does not necessarily start at 0 -*/ - -const int cidx(int i, int m, int k, int nm, int moff, int deg) { - return (i * nm + (m+moff)) * (deg + 1) + k; -} - -void build_cartesian_expansion(FieldExpansionData el){ - printf("cartesian"); - const double h = FieldExpansionData_get_h(el); - const int ncoef = FieldExpansionData_get__ncoef(el); - const int na = FieldExpansionData_get_na(el); - const int nb = FieldExpansionData_get_nb(el); - const int deg = FieldExpansionData_get_deg(el); - double a[na * (deg + 1)]; - for (int i = 0; i < na*(deg+1); ++i){ - a[i] = FieldExpansionData_get_a(el,i); - } - double b[nb * (deg + 1)]; - for (int i = 0; i < nb*(deg+1); ++i){ - b[i] = FieldExpansionData_get_b(el,i); - } - double bs[deg + 1]; - for (int i = 0; i < deg + 1; ++i){ - bs[i] = FieldExpansionData_get_bs(el,i); - } - - const int mmax = FieldExpansionData_get__mmax(el); - const int moff = FieldExpansionData_get__moff(el); - const int nm = FieldExpansionData_get__nm(el); - - double *c = (double *) FieldExpansionData_getp__c(el); - - int nmax = (na > nb) ? na : nb; - double invfact[nmax + 1]; - double invhpow[nmax + 1]; - invfact[0] = 1.0; - invhpow[0] = 1.0; - for (int n = 1; n <= nmax; ++n) { - invfact[n] = invfact[n - 1] / (double)n; - invhpow[n] = invhpow[n - 1] / h; - } - - for (int k = 0; k < deg; ++k) c[cidx(0,0,k+1,nm,moff,deg)] = bs[k] / (double)(k + 1); - - for (int n = 1; n <= na; ++n) { - const double fac = -invfact[n]; - const double *an = a + (size_t)(n - 1) * (size_t)(deg + 1); - for (int k = 0; k <= deg; ++k) c[cidx(0,n,k,nm,moff,deg)] += fac * an[k]; - } - if (ncoef > 1) { - for (int n = 1; n <= nb; ++n) { - const double fac = -invfact[n - 1]; - const double *bn = b + (size_t)(n - 1) * (size_t)(deg + 1); - for (int k = 0; k <= deg; ++k) c[cidx(1,n-1,k,nm,moff,deg)] += fac * bn[k]; - } - } - - for (int i = 0; i + 2 < ncoef; ++i) { - for (int m = 0; m <= mmax; ++m) { - for (int k = 0; k <= deg; ++k) { - double v = 0.0; - if (m + 2 <= mmax) - v += (double)(m + 2) * (double)(m + 1) * c[cidx(i,m+2,k,nm,moff,deg)]; - if (k + 2 <= deg) - v += (double)(k + 2) * (double)(k + 1) * c[cidx(i,m,k+2,nm,moff,deg)]; - c[cidx(i+2,m,k,nm,moff,deg)] = -v; - } - } - } -} - -void build_bent_expansion(FieldExpansionData el){ - const double h = FieldExpansionData_get_h(el); - const int ncoef = FieldExpansionData_get__ncoef(el); - const int na = FieldExpansionData_get_na(el); - const int nb = FieldExpansionData_get_nb(el); - const int deg = FieldExpansionData_get_deg(el); - double a[na * (deg + 1)]; - for (int i = 0; i < na*(deg+1); ++i){ - a[i] = FieldExpansionData_get_a(el,i); - } - double b[nb * (deg + 1)]; - for (int i = 0; i < nb*(deg+1); ++i){ - b[i] = FieldExpansionData_get_b(el,i); - } - double bs[deg + 1]; - for (int i = 0; i < deg + 1; ++i){ - bs[i] = FieldExpansionData_get_bs(el,i); - } - - const int mmax = FieldExpansionData_get__mmax(el); - const int mmin = FieldExpansionData_get__mmin(el); - const int moff = FieldExpansionData_get__moff(el); - const int nm = FieldExpansionData_get__nm(el); - - double *c = (double *) FieldExpansionData_getp__c(el); - - int nmax = (na > nb) ? na : nb; - double invfact[nmax + 1]; - double invhpow[nmax + 1]; - invfact[0] = 1.0; - invhpow[0] = 1.0; - for (int n = 1; n <= nmax; ++n) { - invfact[n] = invfact[n - 1] / (double)n; - invhpow[n] = invhpow[n - 1] / h; - } - - /* a_0(s)=int_0^s b_s(u)du contributes -a_0 to phi_0. - CAREFUL: this will neglect the highest order in the polynomial, - only up to given degree in a0 is kept */ - for (int k = 0; k < deg; ++k) c[cidx(0,0,k+1,nm,moff,deg)] = bs[k] / (double)(k + 1); - /* phi_0(s) = sum_m c[0,m](s) q^m - c[0,m] = - sum_(n>=m) (-1)^(n-m) / (h^n m! (n-m)!) a_n(s) */ - for (int m = 0; m <= na; ++m) { - for (int n = (m > 1 ? m : 1); n <= na; ++n) { - double sgn = ((n - m) & 1) ? -1.0 : 1.0; - double fac = -sgn * invhpow[n] * invfact[m] * invfact[n - m]; - const double *an = a + (size_t)(n - 1) * (size_t)(deg + 1); - for (int k = 0; k <= deg; ++k) c[cidx(0,m,k,nm,moff,deg)] += fac * an[k]; - } - } - - /* phi_1(q,s) = sum_m c[1,m](s) q^m - c[1,m] = - sum_(n>=m+1) (-1)^(n-1-m) / (h^(n-1) m! (n-1-m)!) b_n(s) */ - if (ncoef > 1) { - for (int m = 0; m <= nb - 1; ++m) { - for (int n = m + 1; n <= nb; ++n) { - double sgn = ((n - 1 - m) & 1) ? -1.0 : 1.0; - double fac = -sgn * invhpow[n - 1] * invfact[m] * invfact[n - 1 - m]; - const double *bn = b + (size_t)(n - 1) * (size_t)(deg + 1); - for (int k = 0; k <= deg; ++k) c[cidx(1,m,k,nm,moff,deg)] += fac * bn[k]; - } - } - } - - /* Recursion: c[i+2,m] = -(d_s^2 + h^2 (m+2)^2) c[i,m+2] - implemented for polynomial expansion of c[i,m] in powers of s - C[i+2,m,k] = -(C[i,m+2,k+2]*(k+2)*(k+1) + C[i,m+2,k]*h^2*(m+2)^2) */ - for (int i = 0; i + 2 < ncoef; ++i) { - for (int m = mmin; m <= mmax - 2; ++m) { - double lam = h * h * (double)(m + 2) * (double)(m + 2); - for (int k = 0; k <= deg; ++k) { - double v = lam * c[cidx(i,m+2,k,nm,moff,deg)]; - if (k + 2 <= deg) v += (double)(k + 2) * (double)(k + 1) * c[cidx(i,m+2,k+2,nm,moff,deg)]; - c[cidx(i+2,m,k,nm,moff,deg)] = -v; - } - } - } -} - -void build_expansion(FieldExpansionData el){ - if (FieldExpansionData_get_cartesian(el)) { - build_cartesian_expansion(el); - } else { - build_bent_expansion(el); - } -} - -#endif - diff --git a/xtrack/beam_elements/elements_src/create_fieldexpansion_bent.h b/xtrack/beam_elements/elements_src/create_fieldexpansion_bent.h new file mode 100644 index 000000000..8e8a969ee --- /dev/null +++ b/xtrack/beam_elements/elements_src/create_fieldexpansion_bent.h @@ -0,0 +1,97 @@ +#ifndef create_fieldexpansion_bent_H +#define create_fieldexpansion_bent_H + +/* Index of c[i,m,k] in the c array, ordered as +c[0,mmin,0] ... c[0,mmin,deg], c[0,mmin+1,0] ... c[0,mmin+1,deg], ..., c[0,mmin+nm-1,0] ... c[0,mmin+nm-1,deg], +c[1,mmin,0] ... c[1,mmin,deg], c[1,mmin+1,0] ... c[1,mmin+1,deg], ..., c[1,mmin+nm-1,0] ... c[1,mmin+nm-1,deg], +... c[ncoef-1, mmin+nm-1, deg] +moff=-mmin is the offset to be added to m to get the correct index, +since m does not necessarily start at 0 +*/ + +const int cidx(int i, int m, int k, int nm, int moff, int deg) { + return (i * nm + (m+moff)) * (deg + 1) + k; +} + +void build_expansion_bent(BentFieldExpansionData el){ + const double h = BentFieldExpansionData_get_h(el); + const int ncoef = BentFieldExpansionData_get__ncoef(el); + const int na = BentFieldExpansionData_get_na(el); + const int nb = BentFieldExpansionData_get_nb(el); + const int deg = BentFieldExpansionData_get_deg(el); + double a[na * (deg + 1)]; + for (int i = 0; i < na*(deg+1); ++i){ + a[i] = BentFieldExpansionData_get_a(el,i); + } + double b[nb * (deg + 1)]; + for (int i = 0; i < nb*(deg+1); ++i){ + b[i] = BentFieldExpansionData_get_b(el,i); + } + double bs[deg + 1]; + for (int i = 0; i < deg + 1; ++i){ + bs[i] = BentFieldExpansionData_get_bs(el,i); + } + + const int mmax = BentFieldExpansionData_get__mmax(el); + const int mmin = BentFieldExpansionData_get__mmin(el); + const int moff = BentFieldExpansionData_get__moff(el); + const int nm = BentFieldExpansionData_get__nm(el); + + double *c = (double *) BentFieldExpansionData_getp__c(el); + + int nmax = (na > nb) ? na : nb; + double invfact[nmax + 1]; + double invhpow[nmax + 1]; + invfact[0] = 1.0; + invhpow[0] = 1.0; + for (int n = 1; n <= nmax; ++n) { + invfact[n] = invfact[n - 1] / (double)n; + invhpow[n] = invhpow[n - 1] / h; + } + + /* a_0(s)=int_0^s b_s(u)du contributes -a_0 to phi_0. + CAREFUL: this will neglect the highest order in the polynomial, + only up to given degree in a0 is kept */ + for (int k = 0; k < deg; ++k) c[cidx(0,0,k+1,nm,moff,deg)] = bs[k] / (double)(k + 1); + /* phi_0(s) = sum_m c[0,m](s) q^m + c[0,m] = - sum_(n>=m) (-1)^(n-m) / (h^n m! (n-m)!) a_n(s) */ + for (int m = 0; m <= na; ++m) { + for (int n = (m > 1 ? m : 1); n <= na; ++n) { + double sgn = ((n - m) & 1) ? -1.0 : 1.0; + double fac = -sgn * invhpow[n] * invfact[m] * invfact[n - m]; + const double *an = a + (size_t)(n - 1) * (size_t)(deg + 1); + for (int k = 0; k <= deg; ++k) c[cidx(0,m,k,nm,moff,deg)] += fac * an[k]; + } + } + + /* phi_1(q,s) = sum_m c[1,m](s) q^m + c[1,m] = - sum_(n>=m+1) (-1)^(n-1-m) / (h^(n-1) m! (n-1-m)!) b_n(s) */ + if (ncoef > 1) { + for (int m = 0; m <= nb - 1; ++m) { + for (int n = m + 1; n <= nb; ++n) { + double sgn = ((n - 1 - m) & 1) ? -1.0 : 1.0; + double fac = -sgn * invhpow[n - 1] * invfact[m] * invfact[n - 1 - m]; + const double *bn = b + (size_t)(n - 1) * (size_t)(deg + 1); + for (int k = 0; k <= deg; ++k) c[cidx(1,m,k,nm,moff,deg)] += fac * bn[k]; + } + } + } + + /* Recursion: c[i+2,m] = -(d_s^2 + h^2 (m+2)^2) c[i,m+2] + implemented for polynomial expansion of c[i,m] in powers of s + C[i+2,m,k] = -(C[i,m+2,k+2]*(k+2)*(k+1) + C[i,m+2,k]*h^2*(m+2)^2) */ + for (int i = 0; i + 2 < ncoef; ++i) { + for (int m = mmin; m <= mmax - 2; ++m) { + double lam = h * h * (double)(m + 2) * (double)(m + 2); + for (int k = 0; k <= deg; ++k) { + double v = lam * c[cidx(i,m+2,k,nm,moff,deg)]; + if (k + 2 <= deg) v += (double)(k + 2) * (double)(k + 1) * c[cidx(i,m+2,k+2,nm,moff,deg)]; + c[cidx(i+2,m,k,nm,moff,deg)] = -v; + } + } + } +} + +#endif + + diff --git a/xtrack/beam_elements/elements_src/create_fieldexpansion_straight.h b/xtrack/beam_elements/elements_src/create_fieldexpansion_straight.h new file mode 100644 index 000000000..e164b4cfa --- /dev/null +++ b/xtrack/beam_elements/elements_src/create_fieldexpansion_straight.h @@ -0,0 +1,80 @@ +#ifndef create_fieldexpansion_straight_H +#define create_fieldexpansion_straight_H + +/* Index of c[i,m,k] in the c array, ordered as +c[0,mmin,0] ... c[0,mmin,deg], c[0,mmin+1,0] ... c[0,mmin+1,deg], ..., c[0,mmin+nm-1,0] ... c[0,mmin+nm-1,deg], +c[1,mmin,0] ... c[1,mmin,deg], c[1,mmin+1,0] ... c[1,mmin+1,deg], ..., c[1,mmin+nm-1,0] ... c[1,mmin+nm-1,deg], +... c[ncoef-1, mmin+nm-1, deg] +moff=-mmin is the offset to be added to m to get the correct index, +since m does not necessarily start at 0 +*/ + +const int cidx(int i, int m, int k, int nm, int moff, int deg) { + return (i * nm + (m+moff)) * (deg + 1) + k; +} + +void build_expansion_straight(StraightFieldExpansionData el){ + const double h = StraightFieldExpansionData_get_h(el); + const int ncoef = StraightFieldExpansionData_get__ncoef(el); + const int na = StraightFieldExpansionData_get_na(el); + const int nb = StraightFieldExpansionData_get_nb(el); + const int deg = StraightFieldExpansionData_get_deg(el); + double a[na * (deg + 1)]; + for (int i = 0; i < na*(deg+1); ++i){ + a[i] = StraightFieldExpansionData_get_a(el,i); + } + double b[nb * (deg + 1)]; + for (int i = 0; i < nb*(deg+1); ++i){ + b[i] = StraightFieldExpansionData_get_b(el,i); + } + double bs[deg + 1]; + for (int i = 0; i < deg + 1; ++i){ + bs[i] = StraightFieldExpansionData_get_bs(el,i); + } + + const int mmax = StraightFieldExpansionData_get__mmax(el); + const int moff = StraightFieldExpansionData_get__moff(el); + const int nm = StraightFieldExpansionData_get__nm(el); + + double *c = (double *) StraightFieldExpansionData_getp__c(el); + + int nmax = (na > nb) ? na : nb; + double invfact[nmax + 1]; + double invhpow[nmax + 1]; + invfact[0] = 1.0; + invhpow[0] = 1.0; + for (int n = 1; n <= nmax; ++n) { + invfact[n] = invfact[n - 1] / (double)n; + invhpow[n] = invhpow[n - 1] / h; + } + + for (int k = 0; k < deg; ++k) c[cidx(0,0,k+1,nm,moff,deg)] = bs[k] / (double)(k + 1); + + for (int n = 1; n <= na; ++n) { + const double fac = -invfact[n]; + const double *an = a + (size_t)(n - 1) * (size_t)(deg + 1); + for (int k = 0; k <= deg; ++k) c[cidx(0,n,k,nm,moff,deg)] += fac * an[k]; + } + if (ncoef > 1) { + for (int n = 1; n <= nb; ++n) { + const double fac = -invfact[n - 1]; + const double *bn = b + (size_t)(n - 1) * (size_t)(deg + 1); + for (int k = 0; k <= deg; ++k) c[cidx(1,n-1,k,nm,moff,deg)] += fac * bn[k]; + } + } + + for (int i = 0; i + 2 < ncoef; ++i) { + for (int m = 0; m <= mmax; ++m) { + for (int k = 0; k <= deg; ++k) { + double v = 0.0; + if (m + 2 <= mmax) + v += (double)(m + 2) * (double)(m + 1) * c[cidx(i,m+2,k,nm,moff,deg)]; + if (k + 2 <= deg) + v += (double)(k + 2) * (double)(k + 1) * c[cidx(i,m,k+2,nm,moff,deg)]; + c[cidx(i+2,m,k,nm,moff,deg)] = -v; + } + } + } +} + +#endif \ No newline at end of file diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion.h b/xtrack/beam_elements/elements_src/track_fieldexpansion.h deleted file mode 100644 index 29516b94b..000000000 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion.h +++ /dev/null @@ -1,389 +0,0 @@ -#ifndef XTRACK_TRACK_FIELDEXPANSION_H -#define XTRACK_TRACK_FIELDEXPANSION_H - -typedef struct { - int ny; /* requested output order in y */ - int ncoef; /* stored phi_i coefficients: 0..ny+1 */ - int na, nb, deg; - int mmin, mmax, moff, nm; - int qemin, nq; - double h; - double cartesian; - double *c; /* c[i,m,k], polynomial coeff of s^k in q^m term */ - double *V; /* scratch: c[i,m](s) */ - double *D1; /* scratch: d_s c[i,m] */ - double *D2; /* scratch: d2_s c[i,m] */ - double *Q; /* scratch: q^e, e=qemin.. */ -} Expansion; - -typedef struct { - double phi; - double Bx, By, Bs; - double Ax, Ay, As; - double dAx_dx, dAx_dy, dAx_ds; - double dAs_dx, dAs_dy, dAs_ds; -} FieldValue; - -typedef struct { - double H; - double delta; - double one_plus_delta; - double radicand; - double root; - double grad[6]; /* dH/d{x,px,y,py,tau,ptau} */ - double rhs[6]; /* canonical flow dz/ds */ - double dH_ds; /* explicit derivative at fixed canonical variables */ - FieldValue pot; -} HamiltonianFlow; - -static inline const double *ccptr(const Expansion *f, int i, int m) { - return f->c + (((size_t)i * (size_t)f->nm + (size_t)m) * (size_t)(f->deg + 1)); -} - -static inline void poly_eval_d2(const double *p, int deg, double s, double *v, double *d1, double *d2) { - double a = p[deg], b = 0.0, c=0.0; - for (int k = deg - 1; k >= 0; --k) { - c = c * s + 2.0 * b; - b = b * s + a; - a = a * s + p[k]; - } - *v = a; - *d1 = b; - *d2 = c; -} - -static void fs_prepare_s(Expansion *f, double s) { - for (int i = 0; i < f->ncoef; ++i) { - for (int m = 0; m < f->nm; ++m) { - poly_eval_d2(ccptr(f, i, m), f->deg, s, - &f->V[i * f->nm + m], - &f->D1[i * f->nm + m], - &f->D2[i * f->nm + m]); - } - } -} - -void delta_from_ptau(const double beta0, double ptau, - double *delta, double *delta1, double *ddelta1) { - { - const double r = 1.0 + 2.0 * ptau / beta0 + ptau * ptau; - *delta1 = sqrt(r); - *delta = *delta1 - 1.0; - *ddelta1 = (1.0 / beta0 + ptau) / (*delta1); - } -} - -int evaluate_cartesian_expansion(Expansion *f, double x, double y, double s, FieldValue *out) { - - memset(out, 0, sizeof(*out)); - - double *V = f->V; - double *D1 = f->D1; - double *D2 = f->D2; - double *X = f->Q; - - fs_prepare_s(f, s); - - /* x powers */ - X[0] = pow(x, (double)f->qemin); - for (int t = 1; t < f->nq; ++t) X[t] = X[t - 1] * x; - #define XPOW(E) X[(E) - f->qemin] - - /* As(x,0,s) */ - if (f->ncoef > 1) { - for (int m = 0; m <= f->mmax; ++m) { - int j = m + f->moff; - const double c1m = V[1 * f->nm + j]; /* c[1,m] */ - const double dc1m = f->D1[1 * f->nm + j]; /* c[1,m]'*/ - const double xp = XPOW(m + 1) / (double)(m + 1); - out->As += c1m * xp; - out->dAs_ds += dc1m * xp; - out->dAs_dx += c1m * XPOW(m); - } - } - - double yi = 1.0; /* y^i / i! */ - for (int i = 0; i <= f->ny; ++i) { - double sphi = 0.0, gx = 0.0, gs = 0.0, gy = 0.0; - double dgx_dx = 0.0, dgx_ds = 0.0; - double dgs_dx = 0.0, dgs_ds = 0.0; - - for (int m = f->mmin; m <= f->mmax; ++m) { - const int j = m + f->moff; - const double cim = V[i * f->nm + j]; /* c[i,m] */ - const double ci1m = V[(i + 1) * f->nm + j]; /* c[i+1,m] */ - const double dcim = D1[i * f->nm + j]; /* c[i,m]' */ - const double ddcim = D2[i * f->nm + j]; /* c[i,m]'' */ - - const double xm = XPOW(m); /* x^m */ - const double xm1 = XPOW(m - 1); /* x^(m-1) */ - const double xm2 = XPOW(m - 2); /* x^(m-2) */ - - sphi += cim * xm; - gx += (double)m * cim * xm1; - gy += ci1m * xm; - gs += dcim * xm; - - dgx_dx += (double)m * (double)(m - 1) * cim * xm2; - dgx_ds += (double)m * dcim * xm1; - dgs_dx += (double)m * dcim * xm1; - dgs_ds += ddcim * xm; - } - - out->phi += sphi * yi; /* c[i,m] x^m y^i/i!*/ - out->Bx -= gx * yi; /* -h m c[i,m] x^(m-1) y^i/i! */ - out->By -= gy * yi; /* -c[i+1,m] x^m y^i/i! */ - out->Bs -= gs * yi; /* -c[i,m]' x^(m-1) y^i/i! */ - - /* A_x, A_s through order ny in y: need i = 0..ny-1 */ - if (i < f->ny) { - double yi1 = yi * y / (double)(i + 1); /* y^(i+1)/(i+1)! */ - out->Ax += gs * yi1; /* -c[i,m]' x^(m-1) y^(i+1)/(i+1)! */ - out->As -= gx * yi1; /* -h m c[i,m] x^(m-1) y^i/i! */ - out->dAx_dx += dgs_dx * yi1; - out->dAx_ds += dgs_ds * yi1; - out->dAs_dx += -dgx_dx * yi1; - out->dAs_ds += -dgx_ds * yi1; - } - - yi *= y / (double)(i + 1); - } - - out->dAx_dy = -out->Bs; - out->dAs_dy = out->Bx; - - return 0; - #undef QPOW -} - - -int evaluate_bent_expansion(Expansion *f, double x, double y, double s, FieldValue *out) { - const double q = 1.0 + f->h * x; - if (q == 0.0) return -1; /* singular chart */ - - memset(out, 0, sizeof(*out)); - - double *V = f->V; - double *D1 = f->D1; - double *D2 = f->D2; - double *Q = f->Q; - - fs_prepare_s(f, s); - - /* q powers from e = mmin-1 .. mmax+2 */ - Q[0] = pow(q, (double)f->qemin); - for (int t = 1; t < f->nq; ++t) Q[t] = Q[t - 1] * q; - #define QPOW(E) Q[(E) - f->qemin] - - /* As(x,0,s) - = 1/(1+hx) int_0^x dx' *(1+hx) By(x',0,s) - = 1/qh int_1^q dq' q' phi_1(q',s) - = 1/qh sum_m c[1,m] q^(m+2)/(m+2) - 1/q sum_m c[1,m] 1/(m+2) */ - if (f->ncoef > 1) { - for (int m = 0; m <= f->mmax; ++m) { - int j = m + f->moff; - const double c1m = V[1 * f->nm + j]; /* c[1,m] */ - const double dc1m = f->D1[1 * f->nm + j]; /* c[1,m]'*/ - const double g = QPOW(m + 1) - QPOW(-1); - const double den = f->h * (double)(m + 2); - if (c1m != 0.0) { - out->As += c1m * g / den; - out->dAs_dx += c1m * (((double)(m + 1)) * QPOW(m) + QPOW(-2)) / (double)(m + 2); - } - if (dc1m != 0.0) out->dAs_ds += dc1m * g / den; - } - } - - double yi = 1.0; /* y^i / i! */ - for (int i = 0; i <= f->ny; ++i) { - double sphi = 0.0, gx = 0.0, gs = 0.0, gy = 0.0; - double dgx_dx = 0.0, dgx_ds = 0.0; - double dgs_dx = 0.0, dgs_ds = 0.0; - - for (int m = f->mmin; m <= f->mmax; ++m) { - const int j = m + f->moff; - const double cim = V[i * f->nm + j]; /* c[i,m] */ - const double ci1m = V[(i + 1) * f->nm + j]; /* c[i+1,m] */ - const double dcim = D1[i * f->nm + j]; /* c[i,m]' */ - const double ddcim = D2[i * f->nm + j]; /* c[i,m]'' */ - - const double qm = QPOW(m); /* q^m */ - const double qm1 = QPOW(m - 1); /* q^(m-1) */ - const double qm2 = QPOW(m - 2); /* q^(m-2) */ - - sphi += cim * qm; /* c[i,m] q^m */ - gx += f->h * (double)m * cim * qm1; /* h m c[i,m] q^(m-1) */ - gy += ci1m * qm; /* c[i+1,m] q^m */ - gs += dcim * qm1; /* c[i,m]' q^(m-1) */ - - dgx_dx += f->h * f->h * (double)m * (double)(m-1) * cim * qm2; - dgx_ds += f->h * (double)m * dcim * qm1; - dgs_dx += f->h * (double)(m-1) * dcim * qm2; - dgs_ds += ddcim * qm1; - } - - out->phi += sphi * yi; /* c[i,m] q^m y^i/i!*/ - out->Bx -= gx * yi; /* -h m c[i,m] q^(m-1) y^i/i! */ - out->By -= gy * yi; /* -c[i+1,m] q^m y^i/i! */ - out->Bs -= gs * yi; /* -c[i,m]' q^(m-1) y^i/i! */ - - /* A_x, A_s through order ny in y: need i = 0..ny-1 */ - if (i < f->ny) { - double yi1 = yi * y / (double)(i + 1); /* y^(i+1)/(i+1)! */ - out->Ax += gs * yi1; /* -c[i,m]' q^(m-1) y^(i+1)/(i+1)! */ - out->As -= gx * yi1; /* -h m c[i,m] q^(m-1) y^i/i! */ - out->dAx_dx += dgs_dx * yi1; - out->dAx_ds += dgs_ds * yi1; - out->dAs_dx += -dgx_dx * yi1; - out->dAs_ds += -dgx_ds * yi1; - } - - yi *= y / (double)(i + 1); - } - - out->dAx_dy = -out->Bs; - out->dAs_dy = out->Bx; - - return 0; - #undef QPOW -} - -int evaluate_expansion(Expansion *f, double x, double y, double s, FieldValue *out) { - if (f->cartesian) { - return evaluate_cartesian_expansion(f, x, y, s, out); - } else { - return evaluate_bent_expansion(f, x, y, s, out); - } -} - -void hamiltonian_flow(Expansion *f, const double beta0, - double s, const double z[6], HamiltonianFlow *flow) { - double delta1, delta, ddelta1; - double q, pix, piy, rad, root; - - memset(flow, 0, sizeof(*flow)); - - evaluate_expansion(f, z[0], z[2], s, &flow->pot); - delta_from_ptau(beta0, z[5], &delta, &delta1, &ddelta1); - - q = 1.0 + f->h * z[0]; - pix = z[1] - flow->pot.Ax; - piy = z[3] - flow->pot.Ay; /* A_y is zero in this gauge. */ - rad = delta1 * delta1 - pix * pix - piy * piy; - root = sqrt(rad); - - flow->delta = delta; - flow->one_plus_delta = delta1; - flow->radicand = rad; - flow->root = root; - flow->H = z[5] / beta0 - q * (root + flow->pot.As); - - flow->rhs[0] = q * pix / root; // dx/ds = dH/dpx - flow->rhs[2] = q * piy / root; // dy/ds = dH/dpy - flow->rhs[4] = 1.0 / beta0 - q * delta1 * ddelta1 / root; // dtau/ds = dH/dptau - - flow->rhs[1] = f->h * (root + flow->pot.As) - + q * (pix * flow->pot.dAx_dx / root + flow->pot.dAs_dx); // dpx/ds = -dH/dx - flow->rhs[3] = q * (pix * flow->pot.dAx_dy / root + flow->pot.dAs_dy); // dpy/ds = -dH/dy - flow->rhs[5] = 0.0; // tptau/ds = -dH/dtau, H has no tau-dependence for these static fields. - - flow->grad[0] = -flow->rhs[1]; // dH/dx - flow->grad[1] = flow->rhs[0]; // dH/dpx - flow->grad[2] = -flow->rhs[3]; // dH/dy - flow->grad[3] = flow->rhs[2]; // dH/dpy - flow->grad[4] = -flow->rhs[5]; // dH/dtau - flow->grad[5] = flow->rhs[4]; // dH/dptau - flow->dH_ds = -q * (pix * flow->pot.dAx_ds / root + flow->pot.dAs_ds); -} - -void FieldExpansion_track_local_particle( - FieldExpansionData el, - LocalParticle* part0) -{ - - const double nstep = FieldExpansionData_get_nstep(el); - const double ds = FieldExpansionData_get_ds(el); - - Expansion f; - f.ny = FieldExpansionData_get_ny(el); - f.ncoef = FieldExpansionData_get__ncoef(el); - f.na = FieldExpansionData_get_na(el); - f.nb = FieldExpansionData_get_nb(el); - f.deg = FieldExpansionData_get_deg(el); - f.mmin = FieldExpansionData_get__mmin(el); - f.mmax = FieldExpansionData_get__mmax(el); - f.moff = FieldExpansionData_get__moff(el); - f.nm = FieldExpansionData_get__nm(el); - f.qemin = FieldExpansionData_get__qemin(el); - f.nq = FieldExpansionData_get__nq(el); - f.h = FieldExpansionData_get_h(el); - f.cartesian = FieldExpansionData_get_cartesian(el); - f.c = (double *)FieldExpansionData_getp__c(el); - f.V = (double *)FieldExpansionData_getp__V(el); - f.D1 = (double *)FieldExpansionData_getp__D1(el); - f.D2 = (double *)FieldExpansionData_getp__D2(el); - f.Q = (double *)FieldExpansionData_getp__Q(el); - - HamiltonianFlow flow; - FieldValue v; - - START_PER_PARTICLE_BLOCK(part0, part); - const double beta0 = LocalParticle_get_beta0(part); - - const double x = LocalParticle_get_x(part); - const double px = LocalParticle_get_px(part); - const double y = LocalParticle_get_y(part); - const double py = LocalParticle_get_py(part); - const double tau = LocalParticle_get_zeta(part) / beta0; - const double ptau = LocalParticle_get_ptau(part); - const double ax = LocalParticle_get_ax(part); - const double ay = LocalParticle_get_ay(part); - double z[6] = {x, px, y, py, tau, ptau}; - - // Momentum has to be continuous, vector potential discontinuous, update canonical momentum - evaluate_expansion(&f, z[0], z[2], 0, &v); - z[1] += v.Ax - ax; - z[3] += v.Ay - ay; - - double s = 0; - double ztmp[6]; - for (int step = 0; step < nstep; ++step) { - double k1[6], k2[6], k3[6], k4[6]; - - hamiltonian_flow(&f, beta0, s, z, &flow); - memcpy(k1, flow.rhs, sizeof(k1)); - for (int i = 0; i < 6; ++i) ztmp[i] = z[i] + 0.5 * ds * k1[i]; - - hamiltonian_flow(&f, beta0, s + 0.5 * ds, ztmp, &flow); - memcpy(k2, flow.rhs, sizeof(k2)); - for (int i = 0; i < 6; ++i) ztmp[i] = z[i] + 0.5 * ds * k2[i]; - - hamiltonian_flow(&f, beta0, s + 0.5 * ds, ztmp, &flow); - memcpy(k3, flow.rhs, sizeof(k3)); - for (int i = 0; i < 6; ++i) ztmp[i] = z[i] + ds * k3[i]; - - hamiltonian_flow(&f, beta0, s + ds, ztmp, &flow); - memcpy(k4, flow.rhs, sizeof(k4)); - for (int i = 0; i < 6; ++i) z[i] += ds * (k1[i] + 2.0*k2[i] + 2.0*k3[i] + k4[i]) / 6.0; - - s += ds; - } - - // Back to zero vector potential for next element - evaluate_expansion(&f, z[0], z[2], 0, &v); - z[1] -= v.Ax; - z[3] -= v.Ay; - - LocalParticle_set_x(part, z[0]); - LocalParticle_set_px(part, z[1]); - LocalParticle_set_y(part, z[2]); - LocalParticle_set_py(part, z[3]); - LocalParticle_set_zeta(part, z[4]*beta0); - LocalParticle_set_ptau(part, z[5]); - LocalParticle_set_ax(part, 0); - LocalParticle_set_ay(part, 0); - LocalParticle_add_to_s(part, ds*nstep); - END_PER_PARTICLE_BLOCK -} - -#endif \ No newline at end of file diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion.inc b/xtrack/beam_elements/elements_src/track_fieldexpansion.inc new file mode 100644 index 000000000..353d6d975 --- /dev/null +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion.inc @@ -0,0 +1,131 @@ +void HAMILTONIAN_FLOW(Expansion *f, const double beta0, + double s, const double z[6], HamiltonianFlow *flow) { + double delta1, delta, ddelta1; + double q, pix, piy, rad, root; + + memset(flow, 0, sizeof(*flow)); + + EVALUATE_EXPANSION(f, z[0], z[2], s, &flow->pot); + delta_from_ptau(beta0, z[5], &delta, &delta1, &ddelta1); + q = 1.0 + f->h * z[0]; + pix = z[1] - flow->pot.Ax; + piy = z[3] - flow->pot.Ay; /* A_y is zero in this gauge. */ + rad = delta1 * delta1 - pix * pix - piy * piy; + root = sqrt(rad); + + flow->delta = delta; + flow->one_plus_delta = delta1; + flow->radicand = rad; + flow->root = root; + flow->H = z[5] / beta0 - q * (root + flow->pot.As); + + flow->rhs[0] = q * pix / root; // dx/ds = dH/dpx + flow->rhs[2] = q * piy / root; // dy/ds = dH/dpy + flow->rhs[4] = 1.0 / beta0 - q * delta1 * ddelta1 / root; // dtau/ds = dH/dptau + + flow->rhs[1] = f->h * (root + flow->pot.As) + + q * (pix * flow->pot.dAx_dx / root + flow->pot.dAs_dx); // dpx/ds = -dH/dx + flow->rhs[3] = q * (pix * flow->pot.dAx_dy / root + flow->pot.dAs_dy); // dpy/ds = -dH/dy + flow->rhs[5] = 0.0; // tptau/ds = -dH/dtau, H has no tau-dependence for these static fields. + + flow->grad[0] = -flow->rhs[1]; // dH/dx + flow->grad[1] = flow->rhs[0]; // dH/dpx + flow->grad[2] = -flow->rhs[3]; // dH/dy + flow->grad[3] = flow->rhs[2]; // dH/dpy + flow->grad[4] = -flow->rhs[5]; // dH/dtau + flow->grad[5] = flow->rhs[4]; // dH/dptau + flow->dH_ds = -q * (pix * flow->pot.dAx_ds / root + flow->pot.dAs_ds); +} + +#define CONCATDATA2(a,b) a##b +#define CONCATDATA(a,b) CONCATDATA2(a,b) + +void TRACK_EXPANSION( + FIELDEXPANSIONDATA el, + LocalParticle* part0) +{ + + const double nstep = CONCATDATA(DATA, _get_nstep(el)); + const double ds = CONCATDATA(DATA, _get_ds(el)); + + Expansion f; + f.ny = CONCATDATA(DATA, _get_ny)(el); + f.ncoef = CONCATDATA(DATA, _get__ncoef)(el); + f.na = CONCATDATA(DATA, _get_na)(el); + f.nb = CONCATDATA(DATA, _get_nb)(el); + f.deg = CONCATDATA(DATA, _get_deg)(el); + f.mmin = CONCATDATA(DATA, _get__mmin)(el); + f.mmax = CONCATDATA(DATA, _get__mmax)(el); + f.moff = CONCATDATA(DATA, _get__moff)(el); + f.nm = CONCATDATA(DATA, _get__nm)(el); + f.qemin = CONCATDATA(DATA, _get__qemin)(el); + f.nq = CONCATDATA(DATA, _get__nq)(el); + f.h = CONCATDATA(DATA, _get_h)(el); + f.straight = CONCATDATA(DATA, _get_straight)(el); + f.c = (double *)CONCATDATA(DATA, _getp__c)(el); + f.V = (double *)CONCATDATA(DATA, _getp__V)(el); + f.D1 = (double *)CONCATDATA(DATA, _getp__D1)(el); + f.D2 = (double *)CONCATDATA(DATA, _getp__D2)(el); + f.Q = (double *)CONCATDATA(DATA, _getp__Q)(el); + + HamiltonianFlow flow; + FieldValue v; + + START_PER_PARTICLE_BLOCK(part0, part); + const double beta0 = LocalParticle_get_beta0(part); + + const double x = LocalParticle_get_x(part); + const double px = LocalParticle_get_px(part); + const double y = LocalParticle_get_y(part); + const double py = LocalParticle_get_py(part); + const double tau = LocalParticle_get_zeta(part) / beta0; + const double ptau = LocalParticle_get_ptau(part); + const double ax = LocalParticle_get_ax(part); + const double ay = LocalParticle_get_ay(part); + double z[6] = {x, px, y, py, tau, ptau}; + + // Momentum has to be continuous, vector potential discontinuous, update canonical momentum + EVALUATE_EXPANSION(&f, z[0], z[2], 0, &v); + z[1] += v.Ax - ax; + z[3] += v.Ay - ay; + + double s = 0; + double ztmp[6]; + for (int step = 0; step < nstep; ++step) { + double k1[6], k2[6], k3[6], k4[6]; + + HAMILTONIAN_FLOW(&f, beta0, s, z, &flow); + memcpy(k1, flow.rhs, sizeof(k1)); + for (int i = 0; i < 6; ++i) ztmp[i] = z[i] + 0.5 * ds * k1[i]; + + HAMILTONIAN_FLOW(&f, beta0, s + 0.5 * ds, ztmp, &flow); + memcpy(k2, flow.rhs, sizeof(k2)); + for (int i = 0; i < 6; ++i) ztmp[i] = z[i] + 0.5 * ds * k2[i]; + + HAMILTONIAN_FLOW(&f, beta0, s + 0.5 * ds, ztmp, &flow); + memcpy(k3, flow.rhs, sizeof(k3)); + for (int i = 0; i < 6; ++i) ztmp[i] = z[i] + ds * k3[i]; + + HAMILTONIAN_FLOW(&f, beta0, s + ds, ztmp, &flow); + memcpy(k4, flow.rhs, sizeof(k4)); + for (int i = 0; i < 6; ++i) z[i] += ds * (k1[i] + 2.0*k2[i] + 2.0*k3[i] + k4[i]) / 6.0; + + s += ds; + } + + // Back to zero vector potential for next element + EVALUATE_EXPANSION(&f, z[0], z[2], 0, &v); + z[1] -= v.Ax; + z[3] -= v.Ay; + + LocalParticle_set_x(part, z[0]); + LocalParticle_set_px(part, z[1]); + LocalParticle_set_y(part, z[2]); + LocalParticle_set_py(part, z[3]); + LocalParticle_set_zeta(part, z[4]*beta0); + LocalParticle_set_ptau(part, z[5]); + LocalParticle_set_ax(part, 0); + LocalParticle_set_ay(part, 0); + LocalParticle_add_to_s(part, ds*nstep); + END_PER_PARTICLE_BLOCK +} diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h b/xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h new file mode 100644 index 000000000..e5d0c7b3f --- /dev/null +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h @@ -0,0 +1,111 @@ +#ifndef XTRACK_TRACK_FIELDEXPANSION_BENT_H +#define XTRACK_TRACK_FIELDEXPANSION_BENT_H + +#include "track_fieldexpansion_helpers.h" + + +int evaluate_expansion_bent(Expansion *f, double x, double y, double s, FieldValue *out) { + const double q = 1.0 + f->h * x; + if (q == 0.0) return -1; /* singular chart */ + + memset(out, 0, sizeof(*out)); + + double *V = f->V; + double *D1 = f->D1; + double *D2 = f->D2; + double *Q = f->Q; + + fs_prepare_s(f, s); + + /* q powers from e = mmin-1 .. mmax+2 */ + Q[0] = pow(q, (double)f->qemin); + for (int t = 1; t < f->nq; ++t) Q[t] = Q[t - 1] * q; + #define QPOW(E) Q[(E) - f->qemin] + + /* As(x,0,s) + = 1/(1+hx) int_0^x dx' *(1+hx) By(x',0,s) + = 1/qh int_1^q dq' q' phi_1(q',s) + = 1/qh sum_m c[1,m] q^(m+2)/(m+2) - 1/q sum_m c[1,m] 1/(m+2) */ + if (f->ncoef > 1) { + for (int m = 0; m <= f->mmax; ++m) { + int j = m + f->moff; + const double c1m = V[1 * f->nm + j]; /* c[1,m] */ + const double dc1m = f->D1[1 * f->nm + j]; /* c[1,m]'*/ + const double g = QPOW(m + 1) - QPOW(-1); + const double den = f->h * (double)(m + 2); + if (c1m != 0.0) { + out->As += c1m * g / den; + out->dAs_dx += c1m * (((double)(m + 1)) * QPOW(m) + QPOW(-2)) / (double)(m + 2); + } + if (dc1m != 0.0) out->dAs_ds += dc1m * g / den; + } + } + + double yi = 1.0; /* y^i / i! */ + for (int i = 0; i <= f->ny; ++i) { + double sphi = 0.0, gx = 0.0, gs = 0.0, gy = 0.0; + double dgx_dx = 0.0, dgx_ds = 0.0; + double dgs_dx = 0.0, dgs_ds = 0.0; + + for (int m = f->mmin; m <= f->mmax; ++m) { + const int j = m + f->moff; + const double cim = V[i * f->nm + j]; /* c[i,m] */ + const double ci1m = V[(i + 1) * f->nm + j]; /* c[i+1,m] */ + const double dcim = D1[i * f->nm + j]; /* c[i,m]' */ + const double ddcim = D2[i * f->nm + j]; /* c[i,m]'' */ + + const double qm = QPOW(m); /* q^m */ + const double qm1 = QPOW(m - 1); /* q^(m-1) */ + const double qm2 = QPOW(m - 2); /* q^(m-2) */ + + sphi += cim * qm; /* c[i,m] q^m */ + gx += f->h * (double)m * cim * qm1; /* h m c[i,m] q^(m-1) */ + gy += ci1m * qm; /* c[i+1,m] q^m */ + gs += dcim * qm1; /* c[i,m]' q^(m-1) */ + + dgx_dx += f->h * f->h * (double)m * (double)(m-1) * cim * qm2; + dgx_ds += f->h * (double)m * dcim * qm1; + dgs_dx += f->h * (double)(m-1) * dcim * qm2; + dgs_ds += ddcim * qm1; + } + + out->phi += sphi * yi; /* c[i,m] q^m y^i/i!*/ + out->Bx -= gx * yi; /* -h m c[i,m] q^(m-1) y^i/i! */ + out->By -= gy * yi; /* -c[i+1,m] q^m y^i/i! */ + out->Bs -= gs * yi; /* -c[i,m]' q^(m-1) y^i/i! */ + + /* A_x, A_s through order ny in y: need i = 0..ny-1 */ + if (i < f->ny) { + double yi1 = yi * y / (double)(i + 1); /* y^(i+1)/(i+1)! */ + out->Ax += gs * yi1; /* -c[i,m]' q^(m-1) y^(i+1)/(i+1)! */ + out->As -= gx * yi1; /* -h m c[i,m] q^(m-1) y^i/i! */ + out->dAx_dx += dgs_dx * yi1; + out->dAx_ds += dgs_ds * yi1; + out->dAs_dx += -dgx_dx * yi1; + out->dAs_ds += -dgx_ds * yi1; + } + + yi *= y / (double)(i + 1); + } + + out->dAx_dy = -out->Bs; + out->dAs_dy = out->Bx; + + return 0; + #undef QPOW +} + + +#define TRACK_EXPANSION BentFieldExpansion_track_local_particle +#define HAMILTONIAN_FLOW hamiltonian_flow_bent +#define EVALUATE_EXPANSION evaluate_expansion_bent +#define FIELDEXPANSIONDATA BentFieldExpansionData +#define DATA BentFieldExpansionData +#include "track_fieldexpansion.inc" +#undef TRACK_EXPANSION +#undef HAMILTONIAN_FLOW +#undef EVALUATE_EXPANSION +#undef FIELDEXPANSIONDATA +#undef DATA + +#endif \ No newline at end of file diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion_helpers.h b/xtrack/beam_elements/elements_src/track_fieldexpansion_helpers.h new file mode 100644 index 000000000..802f2dc0a --- /dev/null +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion_helpers.h @@ -0,0 +1,71 @@ +typedef struct { + int ny; /* requested output order in y */ + int ncoef; /* stored phi_i coefficients: 0..ny+1 */ + int na, nb, deg; + int mmin, mmax, moff, nm; + int qemin, nq; + double h; + double straight; + double *c; /* c[i,m,k], polynomial coeff of s^k in q^m term */ + double *V; /* scratch: c[i,m](s) */ + double *D1; /* scratch: d_s c[i,m] */ + double *D2; /* scratch: d2_s c[i,m] */ + double *Q; /* scratch: q^e, e=qemin.. */ +} Expansion; + +typedef struct { + double phi; + double Bx, By, Bs; + double Ax, Ay, As; + double dAx_dx, dAx_dy, dAx_ds; + double dAs_dx, dAs_dy, dAs_ds; +} FieldValue; + +typedef struct { + double H; + double delta; + double one_plus_delta; + double radicand; + double root; + double grad[6]; /* dH/d{x,px,y,py,tau,ptau} */ + double rhs[6]; /* canonical flow dz/ds */ + double dH_ds; /* explicit derivative at fixed canonical variables */ + FieldValue pot; +} HamiltonianFlow; + +static inline const double *ccptr(const Expansion *f, int i, int m) { + return f->c + (((size_t)i * (size_t)f->nm + (size_t)m) * (size_t)(f->deg + 1)); +} + +static inline void poly_eval_d2(const double *p, int deg, double s, double *v, double *d1, double *d2) { + double a = p[deg], b = 0.0, c=0.0; + for (int k = deg - 1; k >= 0; --k) { + c = c * s + 2.0 * b; + b = b * s + a; + a = a * s + p[k]; + } + *v = a; + *d1 = b; + *d2 = c; +} + +static void fs_prepare_s(Expansion *f, double s) { + for (int i = 0; i < f->ncoef; ++i) { + for (int m = 0; m < f->nm; ++m) { + poly_eval_d2(ccptr(f, i, m), f->deg, s, + &f->V[i * f->nm + m], + &f->D1[i * f->nm + m], + &f->D2[i * f->nm + m]); + } + } +} + +void delta_from_ptau(const double beta0, double ptau, + double *delta, double *delta1, double *ddelta1) { + { + const double r = 1.0 + 2.0 * ptau / beta0 + ptau * ptau; + *delta1 = sqrt(r); + *delta = *delta1 - 1.0; + *ddelta1 = (1.0 / beta0 + ptau) / (*delta1); + } +} diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion_straight.h b/xtrack/beam_elements/elements_src/track_fieldexpansion_straight.h new file mode 100644 index 000000000..609f76c41 --- /dev/null +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion_straight.h @@ -0,0 +1,102 @@ +#ifndef XTRACK_TRACK_FIELDEXPANSION_STRAIGHT_H +#define XTRACK_TRACK_FIELDEXPANSION_STRAIGHT_H + +#include "track_fieldexpansion_helpers.h" + +int evaluate_expansion_straight(Expansion *f, double x, double y, double s, FieldValue *out) { + + memset(out, 0, sizeof(*out)); + + double *V = f->V; + double *D1 = f->D1; + double *D2 = f->D2; + double *X = f->Q; + + fs_prepare_s(f, s); + + /* x powers */ + X[0] = pow(x, (double)f->qemin); + for (int t = 1; t < f->nq; ++t) X[t] = X[t - 1] * x; + #define XPOW(E) X[(E) - f->qemin] + + /* As(x,0,s) */ + if (f->ncoef > 1) { + for (int m = 0; m <= f->mmax; ++m) { + int j = m + f->moff; + const double c1m = V[1 * f->nm + j]; /* c[1,m] */ + const double dc1m = f->D1[1 * f->nm + j]; /* c[1,m]'*/ + const double xp = XPOW(m + 1) / (double)(m + 1); + out->As += c1m * xp; + out->dAs_ds += dc1m * xp; + out->dAs_dx += c1m * XPOW(m); + } + } + + double yi = 1.0; /* y^i / i! */ + for (int i = 0; i <= f->ny; ++i) { + double sphi = 0.0, gx = 0.0, gs = 0.0, gy = 0.0; + double dgx_dx = 0.0, dgx_ds = 0.0; + double dgs_dx = 0.0, dgs_ds = 0.0; + + for (int m = f->mmin; m <= f->mmax; ++m) { + const int j = m + f->moff; + const double cim = V[i * f->nm + j]; /* c[i,m] */ + const double ci1m = V[(i + 1) * f->nm + j]; /* c[i+1,m] */ + const double dcim = D1[i * f->nm + j]; /* c[i,m]' */ + const double ddcim = D2[i * f->nm + j]; /* c[i,m]'' */ + + const double xm = XPOW(m); /* x^m */ + const double xm1 = XPOW(m - 1); /* x^(m-1) */ + const double xm2 = XPOW(m - 2); /* x^(m-2) */ + + sphi += cim * xm; + gx += (double)m * cim * xm1; + gy += ci1m * xm; + gs += dcim * xm; + + dgx_dx += (double)m * (double)(m - 1) * cim * xm2; + dgx_ds += (double)m * dcim * xm1; + dgs_dx += (double)m * dcim * xm1; + dgs_ds += ddcim * xm; + } + + out->phi += sphi * yi; /* c[i,m] x^m y^i/i!*/ + out->Bx -= gx * yi; /* -h m c[i,m] x^(m-1) y^i/i! */ + out->By -= gy * yi; /* -c[i+1,m] x^m y^i/i! */ + out->Bs -= gs * yi; /* -c[i,m]' x^(m-1) y^i/i! */ + + /* A_x, A_s through order ny in y: need i = 0..ny-1 */ + if (i < f->ny) { + double yi1 = yi * y / (double)(i + 1); /* y^(i+1)/(i+1)! */ + out->Ax += gs * yi1; /* -c[i,m]' x^(m-1) y^(i+1)/(i+1)! */ + out->As -= gx * yi1; /* -h m c[i,m] x^(m-1) y^i/i! */ + out->dAx_dx += dgs_dx * yi1; + out->dAx_ds += dgs_ds * yi1; + out->dAs_dx += -dgx_dx * yi1; + out->dAs_ds += -dgx_ds * yi1; + } + + yi *= y / (double)(i + 1); + } + + out->dAx_dy = -out->Bs; + out->dAs_dy = out->Bx; + + return 0; + #undef QPOW +} + + +#define TRACK_EXPANSION StraightFieldExpansion_track_local_particle +#define HAMILTONIAN_FLOW hamiltonian_flow_straight +#define EVALUATE_EXPANSION evaluate_expansion_straight +#define FIELDEXPANSIONDATA StraightFieldExpansionData +#define DATA StraightFieldExpansionData +#include "track_fieldexpansion.inc" +#undef TRACK_EXPANSION +#undef HAMILTONIAN_FLOW +#undef EVALUATE_EXPANSION +#undef FIELDEXPANSIONDATA +#undef DATA + +#endif \ No newline at end of file From 8ee654fdf5ab22041320549217d865a644c56f8e Mon Sep 17 00:00:00 2001 From: Silke Date: Thu, 11 Jun 2026 18:18:33 +0200 Subject: [PATCH 16/38] Twiss not working --- .../elements_src/create_fieldexpansion_bent.h | 5 ++--- .../elements_src/create_fieldexpansion_straight.h | 6 ++---- .../elements_src/track_fieldexpansion_helpers.h | 9 +++++++++ 3 files changed, 13 insertions(+), 7 deletions(-) diff --git a/xtrack/beam_elements/elements_src/create_fieldexpansion_bent.h b/xtrack/beam_elements/elements_src/create_fieldexpansion_bent.h index 8e8a969ee..e2e3705ed 100644 --- a/xtrack/beam_elements/elements_src/create_fieldexpansion_bent.h +++ b/xtrack/beam_elements/elements_src/create_fieldexpansion_bent.h @@ -1,6 +1,8 @@ #ifndef create_fieldexpansion_bent_H #define create_fieldexpansion_bent_H +#include "track_fieldexpansion_helpers.h" + /* Index of c[i,m,k] in the c array, ordered as c[0,mmin,0] ... c[0,mmin,deg], c[0,mmin+1,0] ... c[0,mmin+1,deg], ..., c[0,mmin+nm-1,0] ... c[0,mmin+nm-1,deg], c[1,mmin,0] ... c[1,mmin,deg], c[1,mmin+1,0] ... c[1,mmin+1,deg], ..., c[1,mmin+nm-1,0] ... c[1,mmin+nm-1,deg], @@ -9,9 +11,6 @@ moff=-mmin is the offset to be added to m to get the correct index, since m does not necessarily start at 0 */ -const int cidx(int i, int m, int k, int nm, int moff, int deg) { - return (i * nm + (m+moff)) * (deg + 1) + k; -} void build_expansion_bent(BentFieldExpansionData el){ const double h = BentFieldExpansionData_get_h(el); diff --git a/xtrack/beam_elements/elements_src/create_fieldexpansion_straight.h b/xtrack/beam_elements/elements_src/create_fieldexpansion_straight.h index e164b4cfa..c4032a9e9 100644 --- a/xtrack/beam_elements/elements_src/create_fieldexpansion_straight.h +++ b/xtrack/beam_elements/elements_src/create_fieldexpansion_straight.h @@ -1,6 +1,8 @@ #ifndef create_fieldexpansion_straight_H #define create_fieldexpansion_straight_H +#include "track_fieldexpansion_helpers.h" + /* Index of c[i,m,k] in the c array, ordered as c[0,mmin,0] ... c[0,mmin,deg], c[0,mmin+1,0] ... c[0,mmin+1,deg], ..., c[0,mmin+nm-1,0] ... c[0,mmin+nm-1,deg], c[1,mmin,0] ... c[1,mmin,deg], c[1,mmin+1,0] ... c[1,mmin+1,deg], ..., c[1,mmin+nm-1,0] ... c[1,mmin+nm-1,deg], @@ -9,10 +11,6 @@ moff=-mmin is the offset to be added to m to get the correct index, since m does not necessarily start at 0 */ -const int cidx(int i, int m, int k, int nm, int moff, int deg) { - return (i * nm + (m+moff)) * (deg + 1) + k; -} - void build_expansion_straight(StraightFieldExpansionData el){ const double h = StraightFieldExpansionData_get_h(el); const int ncoef = StraightFieldExpansionData_get__ncoef(el); diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion_helpers.h b/xtrack/beam_elements/elements_src/track_fieldexpansion_helpers.h index 802f2dc0a..f1e49c8c9 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion_helpers.h +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion_helpers.h @@ -1,3 +1,10 @@ +#ifndef TRACK_FIELDEXPANSION_HELPERS_H +#define TRACK_FIELDEXPANSION_HELPERS_H + +const int cidx(int i, int m, int k, int nm, int moff, int deg) { + return (i * nm + (m+moff)) * (deg + 1) + k; +} + typedef struct { int ny; /* requested output order in y */ int ncoef; /* stored phi_i coefficients: 0..ny+1 */ @@ -69,3 +76,5 @@ void delta_from_ptau(const double beta0, double ptau, *ddelta1 = (1.0 / beta0 + ptau) / (*delta1); } } + +#endif \ No newline at end of file From ab5e9f1f579a47864878583cf02d88c7a10647bb Mon Sep 17 00:00:00 2001 From: Silke Van der Schueren Date: Thu, 11 Jun 2026 21:57:16 +0200 Subject: [PATCH 17/38] fix macro expansion, rename file --- .../{track_fieldexpansion.inc => track_fieldexpansion.h} | 4 ++++ xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h | 2 +- .../elements_src/track_fieldexpansion_straight.h | 2 +- 3 files changed, 6 insertions(+), 2 deletions(-) rename xtrack/beam_elements/elements_src/{track_fieldexpansion.inc => track_fieldexpansion.h} (99%) diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion.inc b/xtrack/beam_elements/elements_src/track_fieldexpansion.h similarity index 99% rename from xtrack/beam_elements/elements_src/track_fieldexpansion.inc rename to xtrack/beam_elements/elements_src/track_fieldexpansion.h index 353d6d975..d8ac9da28 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion.inc +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion.h @@ -37,8 +37,12 @@ void HAMILTONIAN_FLOW(Expansion *f, const double beta0, flow->dH_ds = -q * (pix * flow->pot.dAx_ds / root + flow->pot.dAs_ds); } +#ifndef CONCATDATA2 #define CONCATDATA2(a,b) a##b +#endif +#ifndef CONCATDATA #define CONCATDATA(a,b) CONCATDATA2(a,b) +#endif void TRACK_EXPANSION( FIELDEXPANSIONDATA el, diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h b/xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h index e5d0c7b3f..16d565eae 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h @@ -101,7 +101,7 @@ int evaluate_expansion_bent(Expansion *f, double x, double y, double s, FieldVal #define EVALUATE_EXPANSION evaluate_expansion_bent #define FIELDEXPANSIONDATA BentFieldExpansionData #define DATA BentFieldExpansionData -#include "track_fieldexpansion.inc" +#include "track_fieldexpansion.h" #undef TRACK_EXPANSION #undef HAMILTONIAN_FLOW #undef EVALUATE_EXPANSION diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion_straight.h b/xtrack/beam_elements/elements_src/track_fieldexpansion_straight.h index 609f76c41..7e9a9b25c 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion_straight.h +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion_straight.h @@ -92,7 +92,7 @@ int evaluate_expansion_straight(Expansion *f, double x, double y, double s, Fiel #define EVALUATE_EXPANSION evaluate_expansion_straight #define FIELDEXPANSIONDATA StraightFieldExpansionData #define DATA StraightFieldExpansionData -#include "track_fieldexpansion.inc" +#include "track_fieldexpansion.h" #undef TRACK_EXPANSION #undef HAMILTONIAN_FLOW #undef EVALUATE_EXPANSION From 48ef7c8412853d3a1e5b1e67ff9fed2b1fe34694 Mon Sep 17 00:00:00 2001 From: Silke Van der Schueren Date: Thu, 11 Jun 2026 22:12:10 +0200 Subject: [PATCH 18/38] Bug fix in canonical momentum shift --- xtrack/beam_elements/elements_src/track_fieldexpansion.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion.h b/xtrack/beam_elements/elements_src/track_fieldexpansion.h index d8ac9da28..4397b1157 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion.h +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion.h @@ -118,7 +118,7 @@ void TRACK_EXPANSION( } // Back to zero vector potential for next element - EVALUATE_EXPANSION(&f, z[0], z[2], 0, &v); + EVALUATE_EXPANSION(&f, z[0], z[2], s, &v); z[1] -= v.Ax; z[3] -= v.Ay; From 2bcee1ab866950cd86894541948e6b00f863cd05 Mon Sep 17 00:00:00 2001 From: Silke Date: Fri, 12 Jun 2026 16:24:53 +0200 Subject: [PATCH 19/38] Element factory added --- tests/test_fieldexpansion_element.py | 10 +++++----- xtrack/beam_elements/elements.py | 11 ++++++++++- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/tests/test_fieldexpansion_element.py b/tests/test_fieldexpansion_element.py index 2b04cb681..112770289 100644 --- a/tests/test_fieldexpansion_element.py +++ b/tests/test_fieldexpansion_element.py @@ -8,7 +8,7 @@ def test_h_sdep(): bs = np.array([0.1, 0.0]) ny = 5 length=0.2 - fexp = xt.BentFieldExpansion(length=length, h=h, a=a, b=b, bs=bs, ny=ny, nstep=100) + fexp = xt.FieldExpansion(length=length, h=h, a=a, b=b, bs=bs, ny=ny, nstep=100) p0 = xt.Particles(x=0.01, y=0.007, tau=0.002, beta0=0.7) line = xt.Line(elements=[fexp]) @@ -28,7 +28,7 @@ def test_sdep(): bs = np.array([0.1, 0.0]) ny = 5 length=0.2 - fexp = xt.StraightFieldExpansion(length=length, a=a, b=b, bs=bs, ny=ny, nstep=100) + fexp = xt.FieldExpansion(length=length, a=a, b=b, bs=bs, ny=ny, nstep=100) p0 = xt.Particles(x=0.01, y=0.007, tau=0.002, beta0=1) line = xt.Line(elements=[fexp]) @@ -55,11 +55,11 @@ def test_twiss(): myfodo = xt.Line(elements=[ xt.Drift(length=1.2), - xt.StraightFieldExpansion(length=0.1, a=np.array([[0]]), b=np.array([[0],[7]]), bs=np.array([0]), ny=5), + xt.FieldExpansion(length=0.1, a=np.array([[0]]), b=np.array([[0],[7]]), bs=np.array([0]), ny=5), xt.Drift(length=0.5), - xt.BentFieldExpansion(length=0.2, h=0.1, a=np.array([[0]]), b=np.array([[0.1]]), bs=np.array([0]), ny=5), + xt.FieldExpansion(length=0.2, h=0.1, a=np.array([[0]]), b=np.array([[0.1]]), bs=np.array([0]), ny=5), xt.Drift(length=0.5), - xt.StraightFieldExpansion(length=0.1, a=np.array([[0]]), b=np.array([[0],[-7]]), bs=np.array([0]), ny=5) + xt.FieldExpansion(length=0.1, a=np.array([[0]]), b=np.array([[0],[-7]]), bs=np.array([0]), ny=5) ]) myfodo.particle_ref = xt.Particles(particle_id=11, q0=1, mass0=1) mytw = myfodo.twiss4d() diff --git a/xtrack/beam_elements/elements.py b/xtrack/beam_elements/elements.py index 1260fd834..5e142bea0 100644 --- a/xtrack/beam_elements/elements.py +++ b/xtrack/beam_elements/elements.py @@ -5213,4 +5213,13 @@ def __init__(self, length, h, a, b, bs, ny, nstep=10, **kwargs): super().__init__(**kwargs) - self.build_expansion(el=self) \ No newline at end of file + self.build_expansion(el=self) + + +class FieldExpansion(BeamElement): + def __new__(cls, *args, **kwargs): + if 'h' in kwargs and kwargs['h'] > 1e-9: + return BentFieldExpansion(*args, **kwargs) + else: + kwargs.pop('h', None) + return StraightFieldExpansion(*args, **kwargs) From 4717a6e0fc6816b824368d3c88eb19315fdea76d Mon Sep 17 00:00:00 2001 From: Silke Date: Tue, 14 Jul 2026 14:36:58 +0200 Subject: [PATCH 20/38] Added comparison of thin fringe (backwards drift, fringe, drift) to Etienne fringe --- examples/fringes/compare_dipole_fringes.py | 47 ++++++++++++++++++++++ 1 file changed, 47 insertions(+) create mode 100644 examples/fringes/compare_dipole_fringes.py diff --git a/examples/fringes/compare_dipole_fringes.py b/examples/fringes/compare_dipole_fringes.py new file mode 100644 index 000000000..e6bae084f --- /dev/null +++ b/examples/fringes/compare_dipole_fringes.py @@ -0,0 +1,47 @@ +import xtrack as xt +import numpy as np + +# Fringe from 0 to 0.1 over length 0.05, with derivatives zero at the endpoints +bfringe=np.array([[0,0,120,-1600]]) +length = 0.05 +bmax = 0.1 +gap = 0.04 +ss = np.linspace(0,length,100) +bvals = bfringe[0,0]+bfringe[0,1]*ss+bfringe[0,2]*ss**2+bfringe[0,3]*ss**3 +fint = np.trapezoid((bmax - bvals)*bvals / bmax**2 / gap, ss) + +exactFringe = xt.FieldExpansion(length=length, h=0, a=np.array([[0,0,0,0]]), b=bfringe, bs=np.array([0,0,0,0]), ny=10) +invDrift = xt.FieldExpansion(length=-length/2, a=np.array([[0]]), b=np.array([[0]]), bs=np.array([0]), ny=5) +invBend = xt.FieldExpansion(length=-length/2, h=0, a=np.array([[0]]), b=np.array([[0]]), bs=np.array([0]), ny=5) +thinFringe = xt.Line(elements=[invDrift, exactFringe, invBend]) + +# Xsuite fringe with same parameters +PTCfringe = xt.Bend(length=0, k0=0.1, edge_entry_model="full", edge_entry_fint=fint, edge_entry_hgap=gap/2, edge_exit_active=0) + +# Check effect of vertical offset including "SAD"-term +yy = np.linspace(-0.1, 0.1, 100) +p0 = xt.Particles(y=yy, x=0, px=0, py=0, zeta=0, delta=0, beta0=1) +p1 = p0.copy() + +thinFringe.track(p0, _force_no_end_turn_actions=True) +PTCfringe.track(p1) + +plt.plot(yy, p0.py, label='Thin fringe') +plt.plot(yy, p1.py, label='PTC fringe') +plt.legend() +plt.xlabel('y') +plt.ylabel('py') + +# Effect of delta at given y +dd = np.linspace(-0.1, 0.1, 100) +p0 = xt.Particles(y=0.002, x=0, px=0, py=0, zeta=0, delta=dd, beta0=1) +p1 = p0.copy() + +thinFringe.track(p0, _force_no_end_turn_actions=True) +PTCfringe.track(p1) + +plt.plot(dd, p0.py, label='Thin fringe') +plt.plot(dd, p1.py, label='PTC fringe') +plt.legend() +plt.xlabel('delta') +plt.ylabel('py') \ No newline at end of file From f03708562cc768987992f543da7b7310d0b7be08 Mon Sep 17 00:00:00 2001 From: Silke Date: Tue, 14 Jul 2026 16:16:44 +0200 Subject: [PATCH 21/38] Added import matplotlib --- examples/fringes/compare_dipole_fringes.py | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/fringes/compare_dipole_fringes.py b/examples/fringes/compare_dipole_fringes.py index e6bae084f..d45ec10a9 100644 --- a/examples/fringes/compare_dipole_fringes.py +++ b/examples/fringes/compare_dipole_fringes.py @@ -1,5 +1,6 @@ import xtrack as xt import numpy as np +import matplotlib.pyplot as plt # Fringe from 0 to 0.1 over length 0.05, with derivatives zero at the endpoints bfringe=np.array([[0,0,120,-1600]]) From 0f6993f2b3f57cb5f36d1fd83fb8113b05de48e2 Mon Sep 17 00:00:00 2001 From: Silke Date: Tue, 14 Jul 2026 17:47:11 +0200 Subject: [PATCH 22/38] changed names of kernels for creating expansion, solves memory issue but now gives tracking error --- tests/test_fieldexpansion_element.py | 13 +++++++++++++ xtrack/beam_elements/elements.py | 16 +++++++--------- 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/tests/test_fieldexpansion_element.py b/tests/test_fieldexpansion_element.py index 112770289..b17477679 100644 --- a/tests/test_fieldexpansion_element.py +++ b/tests/test_fieldexpansion_element.py @@ -40,6 +40,19 @@ def test_sdep(): assert np.isclose(p0.py[0], 0.20351503) assert np.isclose(p0.zeta[0], 0.00058352) assert np.isclose(p0.ptau[0], 0) + +def test_repeatability(): + fexp1 = xt.FieldExpansion(length=0.2, h=0.1, a=np.array([[1]]), b=np.array([[0]]), bs=np.array([0]), ny=5, nstep=100) + fexp2 = xt.FieldExpansion(length=0.2, h=0, a=np.array([[0, 0.1]]), b=np.array([[1,0]]), bs=np.array([0,0]), ny=5, nstep=100) + fexp3 = xt.FieldExpansion(length=0.2, h=0.1, a=np.array([[1]]), b=np.array([[0]]), bs=np.array([0]), ny=5, nstep=100) + p1 = xt.Particles(x=0.1) + p2 = xt.Particles(x=0.1) + p3 = xt.Particles(x=0.1) + fexp1.track(p1) + fexp2.track(p2) + fexp3.track(p3) + + assert np.isclose(p1.x[0], p3.x[0]) def test_twiss(): fodo = xt.Line(elements=[ diff --git a/xtrack/beam_elements/elements.py b/xtrack/beam_elements/elements.py index 5e142bea0..bf3434d07 100644 --- a/xtrack/beam_elements/elements.py +++ b/xtrack/beam_elements/elements.py @@ -4986,12 +4986,10 @@ class ThinSliceNotNeededError(Exception): class StraightFieldExpansion(BeamElement): """ - Specifies the field expansion in general derivatives on axis in straight or curved frame + Specifies the field expansion in general derivatives on axis in straight frame. Parameters ---------- - h : float - Curvature of the element, in 1/m. For straight elements, h=0. a : array, shape na, deg+1, floats describing the polynomial coefficients for the skew multipoles. First index is multipole order, second index is polynomial coefficient. b : array, shape nb, deg+1, floats @@ -5048,7 +5046,7 @@ class StraightFieldExpansion(BeamElement): '#include "xtrack/beam_elements/elements_src/create_fieldexpansion_straight.h"', ] - _kernels = {'build_expansion': xo.Kernel( + _kernels = {'build_expansion_straight': xo.Kernel( c_name='build_expansion_straight', args=[xo.Arg(xo.ThisClass, name='el')] ), @@ -5097,16 +5095,16 @@ def __init__(self, length, a, b, bs, ny, nstep=10, **kwargs): super().__init__(**kwargs) - self.build_expansion(el=self) + self.build_expansion_straight(el=self) class BentFieldExpansion(BeamElement): """ - Specifies the field expansion in general derivatives on axis in straight or curved frame + Specifies the field expansion in general derivatives on axis in curved frame. Parameters ---------- h : float - Curvature of the element, in 1/m. For straight elements, h=0. + Curvature of the element, in 1/m. For straight elements, h=0, use StraightFieldExpansion. a : array, shape na, deg+1, floats describing the polynomial coefficients for the skew multipoles. First index is multipole order, second index is polynomial coefficient. b : array, shape nb, deg+1, floats @@ -5163,7 +5161,7 @@ class BentFieldExpansion(BeamElement): '#include "xtrack/beam_elements/elements_src/create_fieldexpansion_bent.h"', ] - _kernels = {'build_expansion': xo.Kernel( + _kernels = {'build_expansion_bent': xo.Kernel( c_name='build_expansion_bent', args=[xo.Arg(xo.ThisClass, name='el')] ), @@ -5213,7 +5211,7 @@ def __init__(self, length, h, a, b, bs, ny, nstep=10, **kwargs): super().__init__(**kwargs) - self.build_expansion(el=self) + self.build_expansion_bent(el=self) class FieldExpansion(BeamElement): From d44a3fcaadf3448ccf9271d9f56c85e224bbaa25 Mon Sep 17 00:00:00 2001 From: Silke Date: Wed, 15 Jul 2026 10:38:06 +0200 Subject: [PATCH 23/38] Changed beta of test --- tests/test_fieldexpansion_element.py | 23 ++++++----------------- 1 file changed, 6 insertions(+), 17 deletions(-) diff --git a/tests/test_fieldexpansion_element.py b/tests/test_fieldexpansion_element.py index b17477679..59740da60 100644 --- a/tests/test_fieldexpansion_element.py +++ b/tests/test_fieldexpansion_element.py @@ -1,5 +1,7 @@ import xtrack as xt import numpy as np +import xobjects +xobjects.context_cpu.allow_no_prebuilt_kernel = True def test_h_sdep(): h = 0.1 @@ -30,7 +32,7 @@ def test_sdep(): length=0.2 fexp = xt.FieldExpansion(length=length, a=a, b=b, bs=bs, ny=ny, nstep=100) - p0 = xt.Particles(x=0.01, y=0.007, tau=0.002, beta0=1) + p0 = xt.Particles(x=0.01, y=0.007, tau=0.002, beta0=0.7) line = xt.Line(elements=[fexp]) line.track(p0, _force_no_end_turn_actions=True) @@ -38,21 +40,8 @@ def test_sdep(): assert np.isclose(p0.px[0], -0.02435948) assert np.isclose(p0.y[0], 0.02747711) assert np.isclose(p0.py[0], 0.20351503) - assert np.isclose(p0.zeta[0], 0.00058352) + assert np.isclose(p0.zeta[0], -1.64816267e-05) assert np.isclose(p0.ptau[0], 0) - -def test_repeatability(): - fexp1 = xt.FieldExpansion(length=0.2, h=0.1, a=np.array([[1]]), b=np.array([[0]]), bs=np.array([0]), ny=5, nstep=100) - fexp2 = xt.FieldExpansion(length=0.2, h=0, a=np.array([[0, 0.1]]), b=np.array([[1,0]]), bs=np.array([0,0]), ny=5, nstep=100) - fexp3 = xt.FieldExpansion(length=0.2, h=0.1, a=np.array([[1]]), b=np.array([[0]]), bs=np.array([0]), ny=5, nstep=100) - p1 = xt.Particles(x=0.1) - p2 = xt.Particles(x=0.1) - p3 = xt.Particles(x=0.1) - fexp1.track(p1) - fexp2.track(p2) - fexp3.track(p3) - - assert np.isclose(p1.x[0], p3.x[0]) def test_twiss(): fodo = xt.Line(elements=[ @@ -63,7 +52,7 @@ def test_twiss(): xt.Drift(length=0.5), xt.Quadrupole(k1=-7, length=0.1)] ) - fodo.particle_ref = xt.Particles(particle_id=11, q0=1, mass0=1) + fodo.particle_ref = xt.Particles(q0=1, mass0=1) tw = fodo.twiss4d() myfodo = xt.Line(elements=[ @@ -74,7 +63,7 @@ def test_twiss(): xt.Drift(length=0.5), xt.FieldExpansion(length=0.1, a=np.array([[0]]), b=np.array([[0],[-7]]), bs=np.array([0]), ny=5) ]) - myfodo.particle_ref = xt.Particles(particle_id=11, q0=1, mass0=1) + myfodo.particle_ref = xt.Particles(q0=1, mass0=1) mytw = myfodo.twiss4d() assert np.allclose(tw.betx, mytw.betx) From 6d70de0ba51a7506eac389ebc258122577d80f1c Mon Sep 17 00:00:00 2001 From: Silke Date: Wed, 15 Jul 2026 11:45:38 +0200 Subject: [PATCH 24/38] Renamed folder, comments in example --- .../compare_dipole_fringes.py | 20 +++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) rename examples/{fringes => fieldexpansion}/compare_dipole_fringes.py (71%) diff --git a/examples/fringes/compare_dipole_fringes.py b/examples/fieldexpansion/compare_dipole_fringes.py similarity index 71% rename from examples/fringes/compare_dipole_fringes.py rename to examples/fieldexpansion/compare_dipole_fringes.py index d45ec10a9..e505e2a6a 100644 --- a/examples/fringes/compare_dipole_fringes.py +++ b/examples/fieldexpansion/compare_dipole_fringes.py @@ -1,25 +1,29 @@ +""" +Compare an exact thin fringe created as inverse drift - exact fringe - inverse bend to the Forest fringe +""" + import xtrack as xt import numpy as np import matplotlib.pyplot as plt -# Fringe from 0 to 0.1 over length 0.05, with derivatives zero at the endpoints +# Exact fringe from 0 to 0.1 over length 0.05, with derivatives zero at the endpoints bfringe=np.array([[0,0,120,-1600]]) length = 0.05 bmax = 0.1 -gap = 0.04 -ss = np.linspace(0,length,100) -bvals = bfringe[0,0]+bfringe[0,1]*ss+bfringe[0,2]*ss**2+bfringe[0,3]*ss**3 -fint = np.trapezoid((bmax - bvals)*bvals / bmax**2 / gap, ss) -exactFringe = xt.FieldExpansion(length=length, h=0, a=np.array([[0,0,0,0]]), b=bfringe, bs=np.array([0,0,0,0]), ny=10) +exactFringe = xt.FieldExpansion(length=length, a=np.array([[0,0,0,0]]), b=bfringe, bs=np.array([0,0,0,0]), ny=10) invDrift = xt.FieldExpansion(length=-length/2, a=np.array([[0]]), b=np.array([[0]]), bs=np.array([0]), ny=5) -invBend = xt.FieldExpansion(length=-length/2, h=0, a=np.array([[0]]), b=np.array([[0]]), bs=np.array([0]), ny=5) +invBend = xt.FieldExpansion(length=-length/2, a=np.array([[0]]), b=np.array([[bmax]]), bs=np.array([0]), ny=5) thinFringe = xt.Line(elements=[invDrift, exactFringe, invBend]) # Xsuite fringe with same parameters +gap = 0.04 +ss = np.linspace(0,length,100) +bvals = bfringe[0,0]+bfringe[0,1]*ss+bfringe[0,2]*ss**2+bfringe[0,3]*ss**3 +fint = np.trapezoid((bmax - bvals)*bvals / bmax**2 / gap, ss) PTCfringe = xt.Bend(length=0, k0=0.1, edge_entry_model="full", edge_entry_fint=fint, edge_entry_hgap=gap/2, edge_exit_active=0) -# Check effect of vertical offset including "SAD"-term +# Check effect of vertical offset including "SAD"-term: one can observe a deviation at large y values yy = np.linspace(-0.1, 0.1, 100) p0 = xt.Particles(y=yy, x=0, px=0, py=0, zeta=0, delta=0, beta0=1) p1 = p0.copy() From 2f96363e2e55d2a53f44e1c13554e9d4b2c12f1e Mon Sep 17 00:00:00 2001 From: Silke Date: Wed, 15 Jul 2026 12:10:29 +0200 Subject: [PATCH 25/38] Example how to use for quadrupole --- .../fieldexpansion/compare_dipole_fringes.py | 1 + examples/fieldexpansion/quadrupole_fringe.py | 47 +++++++++++++++++++ 2 files changed, 48 insertions(+) create mode 100644 examples/fieldexpansion/quadrupole_fringe.py diff --git a/examples/fieldexpansion/compare_dipole_fringes.py b/examples/fieldexpansion/compare_dipole_fringes.py index e505e2a6a..87cf760a2 100644 --- a/examples/fieldexpansion/compare_dipole_fringes.py +++ b/examples/fieldexpansion/compare_dipole_fringes.py @@ -7,6 +7,7 @@ import matplotlib.pyplot as plt # Exact fringe from 0 to 0.1 over length 0.05, with derivatives zero at the endpoints +# Fringe shape is specified as polynomial in s, lowest coefficient first bfringe=np.array([[0,0,120,-1600]]) length = 0.05 bmax = 0.1 diff --git a/examples/fieldexpansion/quadrupole_fringe.py b/examples/fieldexpansion/quadrupole_fringe.py new file mode 100644 index 000000000..797d1cd0f --- /dev/null +++ b/examples/fieldexpansion/quadrupole_fringe.py @@ -0,0 +1,47 @@ +import xtrack as xt +import numpy as np +import matplotlib.pyplot as plt + +""" +Example for a quadrupole thin fringe and a full magnet with fringe fields +""" + +# Fringe from 0 to 0.1 over length 0.05, with derivatives zero at the endpoints +# Fringe shape is specified as polynomial in s, lowest coefficient first +# First row is dipole coefficients, second row is quadrupole coefficients +bentrance=np.array([[0,0,0,0], [0,0,120,-1600]]) +fringelength = 0.05 +bmax = 0.1 + +entranceFringe = xt.FieldExpansion(length=fringelength, a=np.array([[0,0,0,0]]), b=bentrance, bs=np.array([0,0,0,0]), ny=10) +invDrift = xt.FieldExpansion(length=-fringelength/2, a=np.array([[0]]), b=np.array([[0]]), bs=np.array([0]), ny=5) +invQuad = xt.FieldExpansion(length=-fringelength/2, a=np.array([[0]]), b=np.array([[0], [bmax]]), bs=np.array([0]), ny=5) +thinFringe = xt.Line(elements=[invDrift, entranceFringe, invQuad]) + + +# Full quadrupole including fringes +bexit = np.array([[0,0,0,0], [0.1,0,-120,1600]]) +exitFringe = xt.FieldExpansion(length=fringelength, a=np.array([[0,0,0,0]]), b=bexit, bs=np.array([0,0,0,0]), ny=10) + +magnlength = 1 +bodylength = magnlength - fringelength # Magnetic length up to center of fringe fields (symmetry) +body = xt.Quadrupole(k1=bmax, length=bodylength) # Can also be FieldExpansion element +fringeQuad = xt.Line(elements=[entranceFringe, body, exitFringe]) + +# Check if does what is expected +Quad = xt.Quadrupole(k1=bmax, length=magnlength) +driftQuad = xt.Line(elements=[xt.Drift(length=fringelength/2), Quad, xt.Drift(length=fringelength/2)]) + +p0 = xt.Particles(x=np.linspace(-0.01, 0.01, 10)) +p1 = p0.copy() + +fringeQuad.track(p0) +driftQuad.track(p1) + +# Both have same magnetic length, but one does not include fringe fields. The focussing is very similar. +plt.scatter(p0.x, p0.px, label='Fringe quad') +plt.scatter(p1.x, p1.px, label='Drift quad', marker='x') + + + + From 52754d2d86570b7a301d804953790793c2c373a9 Mon Sep 17 00:00:00 2001 From: Silke Date: Wed, 15 Jul 2026 13:54:53 +0200 Subject: [PATCH 26/38] Element now also supports backtracking --- tests/test_fieldexpansion_element.py | 26 ++++++++++++++++--- xtrack/beam_elements/elements.py | 4 +-- .../elements_src/track_fieldexpansion.h | 8 ++++-- 3 files changed, 31 insertions(+), 7 deletions(-) diff --git a/tests/test_fieldexpansion_element.py b/tests/test_fieldexpansion_element.py index 59740da60..338507271 100644 --- a/tests/test_fieldexpansion_element.py +++ b/tests/test_fieldexpansion_element.py @@ -1,7 +1,7 @@ import xtrack as xt import numpy as np -import xobjects -xobjects.context_cpu.allow_no_prebuilt_kernel = True +import xobjects as xo +xo.context_cpu.allow_no_prebuilt_kernel = True def test_h_sdep(): h = 0.1 @@ -73,4 +73,24 @@ def test_twiss(): assert np.allclose(tw.dx, mytw.dx) assert np.allclose(tw.dy, mytw.dy) - \ No newline at end of file +def test_backtrack(): + h = 0.1 + a = np.array([[1.0, 0.1], [0.2, 0.0], [0.3, 0.1]]) + b = np.array([[0.1, 0.1], [0.5, 0.0]]) + bs = np.array([0.1, 0.0]) + ny = 5 + length=0.2 + fexp = xt.FieldExpansion(length=length, h=h, a=a, b=b, bs=bs, ny=ny, nstep=100) + + p0 = xt.Particles(x=0.01, y=0.007, tau=0.002, beta0=0.7) + line = xt.Line(elements=[fexp]) + + p_test = p0.copy() + line.track(p_test) + line.track(p_test, backtrack=True) + + assert np.all(p_test.state == 1) + for coordinate in ['x', 'px', 'y', 'py', 'zeta', 'delta', 's']: + xo.assert_allclose( + getattr(p_test, coordinate), getattr(p0, coordinate), + rtol=0, atol=1e-12) \ No newline at end of file diff --git a/xtrack/beam_elements/elements.py b/xtrack/beam_elements/elements.py index 538de8e86..12449d032 100644 --- a/xtrack/beam_elements/elements.py +++ b/xtrack/beam_elements/elements.py @@ -5038,7 +5038,7 @@ class StraightFieldExpansion(BeamElement): isthick = True behaves_like_drift = True - has_backtrack = False + has_backtrack = True allow_loss_refinement = False allow_rot_and_shift = False @@ -5153,7 +5153,7 @@ class BentFieldExpansion(BeamElement): isthick = True behaves_like_drift = True - has_backtrack = False + has_backtrack = True allow_loss_refinement = False allow_rot_and_shift = False diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion.h b/xtrack/beam_elements/elements_src/track_fieldexpansion.h index 4397b1157..0dc844db5 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion.h +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion.h @@ -50,7 +50,8 @@ void TRACK_EXPANSION( { const double nstep = CONCATDATA(DATA, _get_nstep(el)); - const double ds = CONCATDATA(DATA, _get_ds(el)); + double ds = CONCATDATA(DATA, _get_ds(el)); + double sstart = 0; Expansion f; f.ny = CONCATDATA(DATA, _get_ny)(el); @@ -72,6 +73,9 @@ void TRACK_EXPANSION( f.D2 = (double *)CONCATDATA(DATA, _getp__D2)(el); f.Q = (double *)CONCATDATA(DATA, _getp__Q)(el); + int64_t const backtrack = LocalParticle_check_track_flag(part0, XS_FLAG_BACKTRACK); + if (backtrack) {sstart = ds * nstep; ds = -ds;} + HamiltonianFlow flow; FieldValue v; @@ -93,7 +97,7 @@ void TRACK_EXPANSION( z[1] += v.Ax - ax; z[3] += v.Ay - ay; - double s = 0; + double s = sstart; double ztmp[6]; for (int step = 0; step < nstep; ++step) { double k1[6], k2[6], k3[6], k4[6]; From 48f2a44260a8041a6a6738c92452bc5bb18c6432 Mon Sep 17 00:00:00 2001 From: Silke Date: Fri, 17 Jul 2026 11:07:47 +0200 Subject: [PATCH 27/38] Added FieldExpansion to precompiled kernels --- tests/test_fieldexpansion_element.py | 1 - .../prebuilt_kernel_definitions/element_inits.py | 15 +++++++++++++++ .../prebuilt_kernel_definitions/element_types.py | 2 ++ 3 files changed, 17 insertions(+), 1 deletion(-) diff --git a/tests/test_fieldexpansion_element.py b/tests/test_fieldexpansion_element.py index 338507271..2920e37c4 100644 --- a/tests/test_fieldexpansion_element.py +++ b/tests/test_fieldexpansion_element.py @@ -1,7 +1,6 @@ import xtrack as xt import numpy as np import xobjects as xo -xo.context_cpu.allow_no_prebuilt_kernel = True def test_h_sdep(): h = 0.1 diff --git a/xtrack/prebuilt_kernel_definitions/element_inits.py b/xtrack/prebuilt_kernel_definitions/element_inits.py index 131fe08c3..6cfe56672 100644 --- a/xtrack/prebuilt_kernel_definitions/element_inits.py +++ b/xtrack/prebuilt_kernel_definitions/element_inits.py @@ -44,4 +44,19 @@ 'data': (1,1,1,1), 'obs_names': ['dummy'] }, + 'BentFieldExpansion': { + 'length': 0, + 'a': np.array([[0]]), + 'b': np.array([[0]]), + 'bs': np.array([0]), + 'ny': 0, + 'h': 1 + }, + 'StraightFieldExpansion': { + 'length': 0, + 'a': np.array([[0]]), + 'b': np.array([[0]]), + 'bs': np.array([0]), + 'ny': 0 + }, } diff --git a/xtrack/prebuilt_kernel_definitions/element_types.py b/xtrack/prebuilt_kernel_definitions/element_types.py index a6f07ce30..aae68ff7b 100644 --- a/xtrack/prebuilt_kernel_definitions/element_types.py +++ b/xtrack/prebuilt_kernel_definitions/element_types.py @@ -44,6 +44,8 @@ DriftExact, Misalignment, SplineBoris, + StraightFieldExpansion, + BentFieldExpansion, # Drift Slices DriftSlice, DriftExactSlice, From 6856d5cfbb660e1434ecf55da8f4aeef97acde7c Mon Sep 17 00:00:00 2001 From: Silke Date: Fri, 17 Jul 2026 13:29:24 +0200 Subject: [PATCH 28/38] Test against boris integrator --- tests/test_fieldexpansion_element.py | 31 +++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/tests/test_fieldexpansion_element.py b/tests/test_fieldexpansion_element.py index 2920e37c4..e80e709be 100644 --- a/tests/test_fieldexpansion_element.py +++ b/tests/test_fieldexpansion_element.py @@ -92,4 +92,33 @@ def test_backtrack(): for coordinate in ['x', 'px', 'y', 'py', 'zeta', 'delta', 's']: xo.assert_allclose( getattr(p_test, coordinate), getattr(p0, coordinate), - rtol=0, atol=1e-12) \ No newline at end of file + rtol=0, atol=1e-12) + +def test_against_boris(): + p0 = xt.Particles(x=0.01, y=0.005, tau=0.001, px=0.003, py=0.004, ptau=0.002, beta0=0.7) + p1 = p0.copy() + length = 1 + + a = np.array([[0.04, 0.2, 0.08], [0, 0, 0.1]]) + b = np.array([[0.05, 0.04, 0.07], [0.01, 0, 0]]) + bs = np.array([0.1, 0.02, 0]) + + def fieldvalue(x,y,z): + # Determined with bpmeth + return ((0.1*z**2*x + 0.08*z**2 + 0.2*z - y**2*(4.8*x + 3.84)/48 + 0.01*y + 0.04) * p0.rigidity0[0], + (0.07*z**2 + 0.04*z + 0.01*x + 0.2/3*y**3 - 0.07*y**2 - y*(2.4*z**2 + 2.4*x**2 + 3.84*x - 0.48)/24 + 0.05) * p0.rigidity0[0], + (0.1*z*x**2 - 0.1*z*y**2 - 0.02*z + x*(0.16*z + 0.2) + y*(0.84*z + 0.24)/6 - 0.1) * p0.rigidity0[0]) + + boris = xt.BorisSpatialIntegrator(fieldmap_callable=fieldvalue, s_start=0, s_end=length, n_steps=500) + fexp = xt.FieldExpansion(length=length, a=a, b=b, bs=bs, ny=5, nstep=50) + + boris.track(p0) + fexp.track(p1) + + assert np.isclose(p0.x, p1.x) + assert np.isclose(p0.px, p1.px) + assert np.isclose(p0.y, p1.y) + assert np.isclose(p0.py, p1.py) + assert np.isclose(p0.zeta, p1.zeta) + assert np.isclose(p0.ptau, p1.ptau) + assert np.isclose(p0.s, p1.s) \ No newline at end of file From 3b28954a98052b5636856745a4e4b617925ae340 Mon Sep 17 00:00:00 2001 From: Silke Date: Mon, 20 Jul 2026 17:46:39 +0200 Subject: [PATCH 29/38] Example fodo with combined function rk4 model --- .../fieldexpansion/combined_function_fodo.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 examples/fieldexpansion/combined_function_fodo.py diff --git a/examples/fieldexpansion/combined_function_fodo.py b/examples/fieldexpansion/combined_function_fodo.py new file mode 100644 index 000000000..0f7c1bc6c --- /dev/null +++ b/examples/fieldexpansion/combined_function_fodo.py @@ -0,0 +1,21 @@ +import numpy as np +import xtrack as xt + +k0 = np.pi/200 +k1 = 0.05 +length = 1 + +k1quad = 1.6 +lengthquad = 0.2 + +lengthdrift = 0.5 + +bcoeffs = np.array([[k0], [k1]]) +combinedfunction = xt.FieldExpansion(length=length, a=np.array([[0]]), b=bcoeffs, bs=np.array([0]), ny=5) +quad1 = xt.Quadrupole(k1=k1quad, length=lengthquad) +quad2 = xt.Quadrupole(k1=-k1quad, length=lengthquad) +drift = xt.Drift(length=lengthdrift) + +fodo = xt.Line(elements=[quad1, drift, combinedfunction, drift, quad2, drift, combinedfunction, drift]) +fodo.set_particle_ref() +tw = fodo.twiss4d() \ No newline at end of file From 51b42d65a50890d2feb7ad39a99d863597956508 Mon Sep 17 00:00:00 2001 From: Silke Date: Mon, 20 Jul 2026 17:59:01 +0200 Subject: [PATCH 30/38] Added nstep in example --- examples/fieldexpansion/combined_function_fodo.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/fieldexpansion/combined_function_fodo.py b/examples/fieldexpansion/combined_function_fodo.py index 0f7c1bc6c..3784ad4a1 100644 --- a/examples/fieldexpansion/combined_function_fodo.py +++ b/examples/fieldexpansion/combined_function_fodo.py @@ -11,11 +11,11 @@ lengthdrift = 0.5 bcoeffs = np.array([[k0], [k1]]) -combinedfunction = xt.FieldExpansion(length=length, a=np.array([[0]]), b=bcoeffs, bs=np.array([0]), ny=5) +combinedfunction = xt.FieldExpansion(length=length, a=np.array([[0]]), b=bcoeffs, bs=np.array([0]), ny=5, nstep=10) quad1 = xt.Quadrupole(k1=k1quad, length=lengthquad) quad2 = xt.Quadrupole(k1=-k1quad, length=lengthquad) drift = xt.Drift(length=lengthdrift) fodo = xt.Line(elements=[quad1, drift, combinedfunction, drift, quad2, drift, combinedfunction, drift]) fodo.set_particle_ref() -tw = fodo.twiss4d() \ No newline at end of file +tw = fodo.twiss4d() From 6019c81b440918e0ffb4f09e6f366d5a3ce99454 Mon Sep 17 00:00:00 2001 From: Silke Date: Fri, 24 Jul 2026 14:06:57 +0200 Subject: [PATCH 31/38] Fix typo in comment --- xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h b/xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h index 16d565eae..d955fb3f6 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h @@ -78,7 +78,7 @@ int evaluate_expansion_bent(Expansion *f, double x, double y, double s, FieldVal if (i < f->ny) { double yi1 = yi * y / (double)(i + 1); /* y^(i+1)/(i+1)! */ out->Ax += gs * yi1; /* -c[i,m]' q^(m-1) y^(i+1)/(i+1)! */ - out->As -= gx * yi1; /* -h m c[i,m] q^(m-1) y^i/i! */ + out->As -= gx * yi1; /* -h m c[i,m] q^(m-1) y^(i+1)/(i+1)! */ out->dAx_dx += dgs_dx * yi1; out->dAx_ds += dgs_ds * yi1; out->dAs_dx += -dgx_dx * yi1; From 9e8ad63de82287e03027a936e46f15b48b030656 Mon Sep 17 00:00:00 2001 From: Silke Date: Fri, 24 Jul 2026 14:43:13 +0200 Subject: [PATCH 32/38] Typo in comment --- .../elements_src/track_fieldexpansion_straight.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion_straight.h b/xtrack/beam_elements/elements_src/track_fieldexpansion_straight.h index 7e9a9b25c..71960d44e 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion_straight.h +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion_straight.h @@ -61,7 +61,7 @@ int evaluate_expansion_straight(Expansion *f, double x, double y, double s, Fiel } out->phi += sphi * yi; /* c[i,m] x^m y^i/i!*/ - out->Bx -= gx * yi; /* -h m c[i,m] x^(m-1) y^i/i! */ + out->Bx -= gx * yi; /* -m c[i,m] x^(m-1) y^i/i! */ out->By -= gy * yi; /* -c[i+1,m] x^m y^i/i! */ out->Bs -= gs * yi; /* -c[i,m]' x^(m-1) y^i/i! */ @@ -69,7 +69,7 @@ int evaluate_expansion_straight(Expansion *f, double x, double y, double s, Fiel if (i < f->ny) { double yi1 = yi * y / (double)(i + 1); /* y^(i+1)/(i+1)! */ out->Ax += gs * yi1; /* -c[i,m]' x^(m-1) y^(i+1)/(i+1)! */ - out->As -= gx * yi1; /* -h m c[i,m] x^(m-1) y^i/i! */ + out->As -= gx * yi1; /* -m c[i,m] x^(m-1) y^i/i! */ out->dAx_dx += dgs_dx * yi1; out->dAx_ds += dgs_ds * yi1; out->dAs_dx += -dgx_dx * yi1; From e077e83d6cb0066c112f87efb9ea693469e49a2c Mon Sep 17 00:00:00 2001 From: Silke Date: Fri, 31 Jul 2026 15:08:38 +0200 Subject: [PATCH 33/38] Added option to choose between kinetic and canonical momentum constant --- tests/test_fieldexpansion_element.py | 6 ++--- xtrack/beam_elements/elements.py | 21 +++++++++++++-- .../elements_src/track_fieldexpansion.h | 26 +++++++++++++------ 3 files changed, 40 insertions(+), 13 deletions(-) diff --git a/tests/test_fieldexpansion_element.py b/tests/test_fieldexpansion_element.py index e80e709be..b81eaf3e2 100644 --- a/tests/test_fieldexpansion_element.py +++ b/tests/test_fieldexpansion_element.py @@ -9,7 +9,7 @@ def test_h_sdep(): bs = np.array([0.1, 0.0]) ny = 5 length=0.2 - fexp = xt.FieldExpansion(length=length, h=h, a=a, b=b, bs=bs, ny=ny, nstep=100) + fexp = xt.FieldExpansion(length=length, h=h, a=a, b=b, bs=bs, ny=ny, nstep=100, pkin_const=True) p0 = xt.Particles(x=0.01, y=0.007, tau=0.002, beta0=0.7) line = xt.Line(elements=[fexp]) @@ -29,7 +29,7 @@ def test_sdep(): bs = np.array([0.1, 0.0]) ny = 5 length=0.2 - fexp = xt.FieldExpansion(length=length, a=a, b=b, bs=bs, ny=ny, nstep=100) + fexp = xt.FieldExpansion(length=length, a=a, b=b, bs=bs, ny=ny, nstep=100, pkin_const=True) p0 = xt.Particles(x=0.01, y=0.007, tau=0.002, beta0=0.7) line = xt.Line(elements=[fexp]) @@ -110,7 +110,7 @@ def fieldvalue(x,y,z): (0.1*z*x**2 - 0.1*z*y**2 - 0.02*z + x*(0.16*z + 0.2) + y*(0.84*z + 0.24)/6 - 0.1) * p0.rigidity0[0]) boris = xt.BorisSpatialIntegrator(fieldmap_callable=fieldvalue, s_start=0, s_end=length, n_steps=500) - fexp = xt.FieldExpansion(length=length, a=a, b=b, bs=bs, ny=5, nstep=50) + fexp = xt.FieldExpansion(length=length, a=a, b=b, bs=bs, ny=5, nstep=50, pkin_const=True) boris.track(p0) fexp.track(p1) diff --git a/xtrack/beam_elements/elements.py b/xtrack/beam_elements/elements.py index 12449d032..a97995c64 100644 --- a/xtrack/beam_elements/elements.py +++ b/xtrack/beam_elements/elements.py @@ -5074,6 +5074,8 @@ class StraightFieldExpansion(BeamElement): "_D2": xo.Float64[:], "_Q": xo.Float64[:], + "pkin_const": xo.Int64, + "sstart": xo.Float64, } _extra_c_sources = [ @@ -5087,12 +5089,13 @@ class StraightFieldExpansion(BeamElement): ), } - def __init__(self, length, a, b, bs, ny, nstep=10, **kwargs): + def __init__(self, length, a, b, bs, ny, nstep=10, sstart=0, **kwargs): kwargs['length'] = length kwargs['h'] = 0 kwargs['straight'] = 1 kwargs['nstep'] = nstep kwargs['ds'] = length/nstep + kwargs['sstart'] = sstart kwargs['a'] = np.asarray(a, dtype=np.float64).flatten() kwargs['b'] = np.asarray(b, dtype=np.float64).flatten() @@ -5104,6 +5107,11 @@ def __init__(self, length, a, b, bs, ny, nstep=10, **kwargs): kwargs['deg'] = a.shape[1] - 1 + if "pkin_const" in kwargs: + kwargs["pkin_const"] = int(kwargs["pkin_const"]) + else: + kwargs["pkin_const"] = 0 # Default is symplectic option + if b.shape[1] != kwargs['deg'] + 1 or bs.shape[0] != kwargs['deg'] + 1: raise ValueError("Invalid input shapes") @@ -5189,6 +5197,8 @@ class BentFieldExpansion(BeamElement): "_D2": xo.Float64[:], "_Q": xo.Float64[:], + "pkin_const": xo.Int64, + "sstart": xo.Float64, } _extra_c_sources = [ @@ -5202,13 +5212,14 @@ class BentFieldExpansion(BeamElement): ), } - def __init__(self, length, h, a, b, bs, ny, nstep=10, **kwargs): + def __init__(self, length, h, a, b, bs, ny, nstep=10, sstart=0, **kwargs): assert h > 1e-4, "Use straight element with h=0!" kwargs['length'] = length kwargs['h'] = h kwargs['straight'] = 0 kwargs['nstep'] = nstep kwargs['ds'] = length/nstep + kwargs['sstart'] = sstart kwargs['a'] = np.asarray(a, dtype=np.float64).flatten() kwargs['b'] = np.asarray(b, dtype=np.float64).flatten() @@ -5219,6 +5230,12 @@ def __init__(self, length, h, a, b, bs, ny, nstep=10, **kwargs): kwargs['ny'] = ny kwargs['deg'] = a.shape[1] - 1 + + if "pkin_const" in kwargs: + kwargs["pkin_const"] = int(kwargs["pkin_const"]) + else: + kwargs["pkin_const"] = 0 # Default symplectic option + if b.shape[1] != kwargs['deg'] + 1 or bs.shape[0] != kwargs['deg'] + 1: raise ValueError("Invalid input shapes") diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion.h b/xtrack/beam_elements/elements_src/track_fieldexpansion.h index 0dc844db5..a1b887096 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion.h +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion.h @@ -51,7 +51,7 @@ void TRACK_EXPANSION( const double nstep = CONCATDATA(DATA, _get_nstep(el)); double ds = CONCATDATA(DATA, _get_ds(el)); - double sstart = 0; + double sstart = CONCATDATA(DATA, _get_sstart)(el); Expansion f; f.ny = CONCATDATA(DATA, _get_ny)(el); @@ -76,6 +76,8 @@ void TRACK_EXPANSION( int64_t const backtrack = LocalParticle_check_track_flag(part0, XS_FLAG_BACKTRACK); if (backtrack) {sstart = ds * nstep; ds = -ds;} + int pkin_const = CONCATDATA(DATA, _get_pkin_const)(el); + HamiltonianFlow flow; FieldValue v; @@ -93,9 +95,11 @@ void TRACK_EXPANSION( double z[6] = {x, px, y, py, tau, ptau}; // Momentum has to be continuous, vector potential discontinuous, update canonical momentum - EVALUATE_EXPANSION(&f, z[0], z[2], 0, &v); - z[1] += v.Ax - ax; - z[3] += v.Ay - ay; + if (pkin_const) { + EVALUATE_EXPANSION(&f, z[0], z[2], sstart, &v); + z[1] += v.Ax - ax; + z[3] += v.Ay - ay; + } double s = sstart; double ztmp[6]; @@ -123,8 +127,16 @@ void TRACK_EXPANSION( // Back to zero vector potential for next element EVALUATE_EXPANSION(&f, z[0], z[2], s, &v); - z[1] -= v.Ax; - z[3] -= v.Ay; + if (pkin_const) { + z[1] -= v.Ax; + z[3] -= v.Ay; + LocalParticle_set_ax(part, 0); + LocalParticle_set_ay(part, 0); + } + else { + LocalParticle_set_ax(part, v.Ax); + LocalParticle_set_ay(part, v.Ay); + } LocalParticle_set_x(part, z[0]); LocalParticle_set_px(part, z[1]); @@ -132,8 +144,6 @@ void TRACK_EXPANSION( LocalParticle_set_py(part, z[3]); LocalParticle_set_zeta(part, z[4]*beta0); LocalParticle_set_ptau(part, z[5]); - LocalParticle_set_ax(part, 0); - LocalParticle_set_ay(part, 0); LocalParticle_add_to_s(part, ds*nstep); END_PER_PARTICLE_BLOCK } From b8e9f40fc0f715c0ebe7ccf023e9fad901721c81 Mon Sep 17 00:00:00 2001 From: Silke Date: Thu, 27 Aug 2026 18:19:25 +0200 Subject: [PATCH 34/38] Fixed sign mistake that was also present in bpmeth! Added additional test straight vs curved for a1, b1, bs that should spot this type of bugs --- tests/test_fieldexpansion_element.py | 65 +++++++++++++++---- .../elements_src/create_fieldexpansion_bent.h | 2 +- .../create_fieldexpansion_straight.h | 2 +- 3 files changed, 53 insertions(+), 16 deletions(-) diff --git a/tests/test_fieldexpansion_element.py b/tests/test_fieldexpansion_element.py index b81eaf3e2..45bdc6712 100644 --- a/tests/test_fieldexpansion_element.py +++ b/tests/test_fieldexpansion_element.py @@ -15,11 +15,11 @@ def test_h_sdep(): line = xt.Line(elements=[fexp]) line.track(p0, _force_no_end_turn_actions=True) - assert np.isclose(p0.x[0], 0.00968437) - assert np.isclose(p0.px[0], -0.00430618) - assert np.isclose(p0.y[0], 0.02753048) - assert np.isclose(p0.py[0], 0.20394359) - assert np.isclose(p0.zeta[0], -0.00020263) + assert np.isclose(p0.x[0], 0.00995953) + assert np.isclose(p0.px[0], -0.00021054) + assert np.isclose(p0.y[0], 0.02753286) + assert np.isclose(p0.py[0], 0.2039826) + assert np.isclose(p0.zeta[0], -0.00020395) assert np.isclose(p0.ptau[0], 0) assert np.isclose(p0.s[0], length) @@ -35,11 +35,11 @@ def test_sdep(): line = xt.Line(elements=[fexp]) line.track(p0, _force_no_end_turn_actions=True) - assert np.isclose(p0.x[0], 0.00765466) - assert np.isclose(p0.px[0], -0.02435948) - assert np.isclose(p0.y[0], 0.02747711) - assert np.isclose(p0.py[0], 0.20351503) - assert np.isclose(p0.zeta[0], -1.64816267e-05) + assert np.isclose(p0.x[0], 0.00792934) + assert np.isclose(p0.px[0], -0.02026721) + assert np.isclose(p0.y[0], 0.02750655) + assert np.isclose(p0.py[0], 0.20395936) + assert np.isclose(p0.zeta[0], -1.64370788e-05) assert np.isclose(p0.ptau[0], 0) def test_twiss(): @@ -105,9 +105,9 @@ def test_against_boris(): def fieldvalue(x,y,z): # Determined with bpmeth - return ((0.1*z**2*x + 0.08*z**2 + 0.2*z - y**2*(4.8*x + 3.84)/48 + 0.01*y + 0.04) * p0.rigidity0[0], - (0.07*z**2 + 0.04*z + 0.01*x + 0.2/3*y**3 - 0.07*y**2 - y*(2.4*z**2 + 2.4*x**2 + 3.84*x - 0.48)/24 + 0.05) * p0.rigidity0[0], - (0.1*z*x**2 - 0.1*z*y**2 - 0.02*z + x*(0.16*z + 0.2) + y*(0.84*z + 0.24)/6 - 0.1) * p0.rigidity0[0]) + return ((0.1*z**2*x + 0.08*z**2 + 0.2*z - y**2*(0.4*x + 0.32)/4 + 0.01*y + 0.04) * p0.rigidity0[0], + (0.07*z**2 + 0.04*z + 0.01*x + 0.0666666666666667*y**3 - 0.07*y**2 - y*(0.2*z**2 + 0.2*x**2 + 0.32*x + 0.04)/2 + 0.05) * p0.rigidity0[0], + (0.1*z*x**2 - 0.1*z*y**2 + 0.02*z + x*(0.16*z + 0.2) + y*(0.14*z + 0.04) + 0.1) * p0.rigidity0[0]) boris = xt.BorisSpatialIntegrator(fieldmap_callable=fieldvalue, s_start=0, s_end=length, n_steps=500) fexp = xt.FieldExpansion(length=length, a=a, b=b, bs=bs, ny=5, nstep=50, pkin_const=True) @@ -121,4 +121,41 @@ def fieldvalue(x,y,z): assert np.isclose(p0.py, p1.py) assert np.isclose(p0.zeta, p1.zeta) assert np.isclose(p0.ptau, p1.ptau) - assert np.isclose(p0.s, p1.s) \ No newline at end of file + assert np.isclose(p0.s, p1.s) + +def test_straighttocurved(): + def curved_to_straight(p, h): + return { + "x" : (1/h + p.x) * np.cos((p.s)*h) - 1/h, + "s" : (1/h + p.x) * np.sin((p.s)*h), + "y" : p.y, + } + + def straight_to_curved(p, h): + return { + "x" : np.sqrt((1/h + p.x)**2 + (p.s)**2) - 1/h, + "s" : 1/h * np.arctan((p.s)/(1/h + p.x)), + "y" : p.y, + } + + b1st = 1 + a1st = 0.1 + bsst = 0.5 + h = 0.4 + length = 0.5 + + bcu = np.array([[b1st, 0, 0, 0, 0, 0, 0, 0, 0, 0]]) + acu = np.array([[a1st, 0, - a1st/2*h**2, 0, a1st/24*h**4, 0, - a1st/720*h**6, 0, a1st/40320*h**8, 0]]) + bscu = np.array([0, -a1st*h, 0, a1st/6*h**3, 0, -a1st/120*h**5, 0, a1st/5040*h**7, 0, -a1st/362880*h**9]) + + p0 = xt.Particles(x=0.01, y=0.005, tau=0.001, px=0.003, py=0.004, ptau=0.002, beta0=0.7) + p1 = p0.copy() + + line_straight = xt.Line(elements=[xt.FieldExpansion(length=length, h=0, a=np.array([[a1st]]), b=np.array([[b1st]]), bs=np.array([bsst]), ny=5, nstep=100)]) + line_straight.track(p0, _force_no_end_turn_actions=True) + line_curved = xt.Line(elements=[xt.FieldExpansion(length=straight_to_curved(p0, h)["s"], h=h, a=acu, b=bcu, bs=bscu, ny=5, nstep=100)]) + line_curved.track(p1, _force_no_end_turn_actions=True) + + assert np.isclose(curved_to_straight(p1, h)["x"], p0.x) + assert np.isclose(curved_to_straight(p1, h)["s"], p0.s) + assert np.isclose(curved_to_straight(p1, h)["y"], p0.y) \ No newline at end of file diff --git a/xtrack/beam_elements/elements_src/create_fieldexpansion_bent.h b/xtrack/beam_elements/elements_src/create_fieldexpansion_bent.h index e2e3705ed..a0ab60cb9 100644 --- a/xtrack/beam_elements/elements_src/create_fieldexpansion_bent.h +++ b/xtrack/beam_elements/elements_src/create_fieldexpansion_bent.h @@ -51,7 +51,7 @@ void build_expansion_bent(BentFieldExpansionData el){ /* a_0(s)=int_0^s b_s(u)du contributes -a_0 to phi_0. CAREFUL: this will neglect the highest order in the polynomial, only up to given degree in a0 is kept */ - for (int k = 0; k < deg; ++k) c[cidx(0,0,k+1,nm,moff,deg)] = bs[k] / (double)(k + 1); + for (int k = 0; k < deg; ++k) c[cidx(0,0,k+1,nm,moff,deg)] = -bs[k] / (double)(k + 1); /* phi_0(s) = sum_m c[0,m](s) q^m c[0,m] = - sum_(n>=m) (-1)^(n-m) / (h^n m! (n-m)!) a_n(s) */ for (int m = 0; m <= na; ++m) { diff --git a/xtrack/beam_elements/elements_src/create_fieldexpansion_straight.h b/xtrack/beam_elements/elements_src/create_fieldexpansion_straight.h index c4032a9e9..3b48fb1bd 100644 --- a/xtrack/beam_elements/elements_src/create_fieldexpansion_straight.h +++ b/xtrack/beam_elements/elements_src/create_fieldexpansion_straight.h @@ -46,7 +46,7 @@ void build_expansion_straight(StraightFieldExpansionData el){ invhpow[n] = invhpow[n - 1] / h; } - for (int k = 0; k < deg; ++k) c[cidx(0,0,k+1,nm,moff,deg)] = bs[k] / (double)(k + 1); + for (int k = 0; k < deg; ++k) c[cidx(0,0,k+1,nm,moff,deg)] = -bs[k] / (double)(k + 1); for (int n = 1; n <= na; ++n) { const double fac = -invfact[n]; From 6d567b8884731ccf05297bddcfbf8dc1706db5d9 Mon Sep 17 00:00:00 2001 From: Silke Date: Mon, 31 Aug 2026 11:33:43 +0200 Subject: [PATCH 35/38] Added example for dipole with fringes --- .../fieldexpansion/dipole_with_fringes.py | 61 +++++++++++++++++++ 1 file changed, 61 insertions(+) create mode 100644 examples/fieldexpansion/dipole_with_fringes.py diff --git a/examples/fieldexpansion/dipole_with_fringes.py b/examples/fieldexpansion/dipole_with_fringes.py new file mode 100644 index 000000000..b3cfa3a96 --- /dev/null +++ b/examples/fieldexpansion/dipole_with_fringes.py @@ -0,0 +1,61 @@ +import numpy as np +import xtrack as xt +import matplotlib.pyplot as plt + +def get_bshape4(L, bb0, bp0, bbL, bpL, bint): + """ + Calculate coefficients of a fourth order polynomial with the given values and derivatives + at the edges, and the given integral + :param L: length of the segment + :param bb0: value at 0 + :param bp0: derivative at 0 + :param bbL: value at L + :param bpL: derivative at L + :param bint: integral between 0 and L + :return: set of coefficients, lowest order first + """ + + t1 = np.array([1, 0, -18, 32, -15 ]) # p(0)=1 + t2 = np.array([0, 1, -9/2, 6, -5/2]) # p'(0)=1 + t3 = np.array([0, 0, -12, 28, -15 ]) # p(L)=1 + t4 = np.array([0, 0, 3/2, -4, 5/2]) # p'(L)=1 + t5 = np.array([0, 0, 30, -60, 30 ]) # int_0^1 dx p(x)=1 + + return (bb0*t1 + L*bp0*t2 + bbL*t3 + L*bpL*t4 + bint/L*t5) * L**np.arange(0, -5, -1) + +def calc_value(coeffs, s): + order = len(coeffs) - 1 + return np.sum(coeffs[:, None] * s[None, :]**np.arange(order + 1)[:, None], axis=0) + +gap = 0.05 +length = 0.5 +bmax = 0.5 +h = bmax # 1 / bending radius + +fringe_length = 3*gap +body_length = length - fringe_length + +b1_in = get_bshape4(fringe_length, 0, 0, bmax, 0, fringe_length * bmax / 2) +b1_body = np.array([bmax]) +b1_out = get_bshape4(fringe_length, bmax, 0, 0, 0, fringe_length * bmax / 2) + +fig, ax = plt.subplots() +s1 = np.linspace(0, fringe_length, 100) +ax.plot(s1, calc_value(b1_in, s1)) +s2 = np.linspace(0, body_length, 100) +ax.plot(s2+fringe_length, calc_value(b1_body, s2)) +s3 = np.linspace(0, fringe_length, 100) +ax.plot(s3+fringe_length+body_length, calc_value(b1_out, s3)) + +assert body_length >= 0, "Different shape needed to describe such short magnets" + +dipole = xt.Line(elements = [ + xt.FieldExpansion(length=fringe_length/2, b=np.array([b1_in]), a=0*np.array([b1_in]), bs=0*b1_in, ny=5, nstep=10), + xt.FieldExpansion(length=fringe_length/2, b=np.array([b1_in]), a=0*np.array([b1_in]), bs=0*b1_in, ny=5, nstep=10, h=h), + xt.FieldExpansion(length=body_length, b=np.array([b1_body]), a=0*np.array([b1_body]), bs=0*b1_body, ny=5, nstep=10, h=h), + xt.FieldExpansion(length=fringe_length/2, b=np.array([b1_out]), a=0*np.array([b1_out]), bs=0*b1_out, ny=5, nstep=10, h=h), + xt.FieldExpansion(length=fringe_length/2, b=np.array([b1_out]), a=0*np.array([b1_out]), bs=0*b1_out, ny=5, nstep=10) +]) + +p0 = xt.Particles() +dipole.track(p0) \ No newline at end of file From 287759e429ba1b81fa4cd98b930990ac886d561f Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Mon, 31 Aug 2026 13:46:58 +0200 Subject: [PATCH 36/38] update example --- .../fieldexpansion/dipole_with_fringes.py | 46 +++++++++++++++++-- 1 file changed, 43 insertions(+), 3 deletions(-) diff --git a/examples/fieldexpansion/dipole_with_fringes.py b/examples/fieldexpansion/dipole_with_fringes.py index b3cfa3a96..135bd4c96 100644 --- a/examples/fieldexpansion/dipole_with_fringes.py +++ b/examples/fieldexpansion/dipole_with_fringes.py @@ -14,7 +14,7 @@ def get_bshape4(L, bb0, bp0, bbL, bpL, bint): :param bint: integral between 0 and L :return: set of coefficients, lowest order first """ - + t1 = np.array([1, 0, -18, 32, -15 ]) # p(0)=1 t2 = np.array([0, 1, -9/2, 6, -5/2]) # p'(0)=1 t3 = np.array([0, 0, -12, 28, -15 ]) # p(L)=1 @@ -28,7 +28,7 @@ def calc_value(coeffs, s): return np.sum(coeffs[:, None] * s[None, :]**np.arange(order + 1)[:, None], axis=0) gap = 0.05 -length = 0.5 +length = 0.5 # magnetic length bmax = 0.5 h = bmax # 1 / bending radius @@ -58,4 +58,44 @@ def calc_value(coeffs, s): ]) p0 = xt.Particles() -dipole.track(p0) \ No newline at end of file +dipole.track(p0) + +### Fodo +env=xt.Environment() +env['k1quad'] = 1.6 +env['lengthquad'] = 0.2 +env.new("quad1", xt.Quadrupole, k1='k1quad', length='lengthquad') +env.new("quad2", xt.Quadrupole, k1='-k1quad', length='lengthquad') +nstep=5 +env.elements['e0']=xt.FieldExpansion(length=fringe_length/2, b=np.array([b1_in]), a=0*np.array([b1_in]), bs=0*b1_in, ny=5, nstep=nstep) +env.elements['e1']=xt.FieldExpansion(length=fringe_length/2, b=np.array([b1_in]), a=0*np.array([b1_in]), bs=0*b1_in, ny=5, nstep=nstep, h=h) +env.elements['e2']=xt.FieldExpansion(length=body_length, b=np.array([b1_body]), a=0*np.array([b1_body]), bs=0*b1_body, ny=5, nstep=nstep, h=h) +env.elements['e3']=xt.FieldExpansion(length=fringe_length/2, b=np.array([b1_out]), a=0*np.array([b1_out]), bs=0*b1_out, ny=5, nstep=nstep, h=h) +env.elements['e4']=xt.FieldExpansion(length=fringe_length/2, b=np.array([b1_out]), a=0*np.array([b1_out]), bs=0*b1_out, ny=5, nstep=nstep) +env.new_line(name="dipole",components=['e0', 'e1', 'e2', 'e3', 'e4']) +fodo= env.new_line(name='myline', length=5, + components=[ + env.place('quad1', at=0.1), + env.place('dipole', at=1.0), + env.place('quad2', at=2.2), + env.place('dipole', at=3.5), + ]) +fodo.set_particle_ref() +tw = fodo.twiss4d() +tw.plot() +x0,px0=tw.x[0],tw.px[0] +part=fodo.build_particles(x=x0+np.array([0.0,0.01,0.0]), + px=px0+np.array([0,0,0]), + y=np.array([0.0,0.0,0.01])) +fodo.track(part,num_turns=100000,turn_by_turn_monitor=True) +out=fodo.record_last_track +fig,(a1,a2)=plt.subplots(2,1) + +for i in range(len(out.x)): + a1.plot(fodo.record_last_track.x[i],fodo.record_last_track.px[i],',') + a2.plot(fodo.record_last_track.y[i],fodo.record_last_track.py[i],',') + + + + + From 666bd049d8ea79ffc9a2d052874d8a771ac2d6ea Mon Sep 17 00:00:00 2001 From: Silke Date: Mon, 31 Aug 2026 17:04:52 +0200 Subject: [PATCH 37/38] Corrections in the example --- .../fieldexpansion/dipole_with_fringes.py | 26 +++++++------------ 1 file changed, 10 insertions(+), 16 deletions(-) diff --git a/examples/fieldexpansion/dipole_with_fringes.py b/examples/fieldexpansion/dipole_with_fringes.py index 135bd4c96..d169976e5 100644 --- a/examples/fieldexpansion/dipole_with_fringes.py +++ b/examples/fieldexpansion/dipole_with_fringes.py @@ -49,31 +49,25 @@ def calc_value(coeffs, s): assert body_length >= 0, "Different shape needed to describe such short magnets" -dipole = xt.Line(elements = [ - xt.FieldExpansion(length=fringe_length/2, b=np.array([b1_in]), a=0*np.array([b1_in]), bs=0*b1_in, ny=5, nstep=10), - xt.FieldExpansion(length=fringe_length/2, b=np.array([b1_in]), a=0*np.array([b1_in]), bs=0*b1_in, ny=5, nstep=10, h=h), - xt.FieldExpansion(length=body_length, b=np.array([b1_body]), a=0*np.array([b1_body]), bs=0*b1_body, ny=5, nstep=10, h=h), - xt.FieldExpansion(length=fringe_length/2, b=np.array([b1_out]), a=0*np.array([b1_out]), bs=0*b1_out, ny=5, nstep=10, h=h), - xt.FieldExpansion(length=fringe_length/2, b=np.array([b1_out]), a=0*np.array([b1_out]), bs=0*b1_out, ny=5, nstep=10) -]) +env=xt.Environment() +nstep=5 +env.elements['e0']=xt.FieldExpansion(pkin_const=0, length=fringe_length/2, b=np.array([b1_in]), a=0*np.array([b1_in]), bs=0*b1_in, ny=5, nstep=nstep) +env.elements['e1']=xt.FieldExpansion(pkin_const=0, length=fringe_length/2, b=np.array([b1_in]), a=0*np.array([b1_in]), bs=0*b1_in, ny=5, nstep=nstep, h=h, sstart=fringe_length/2) +env.elements['e2']=xt.FieldExpansion(pkin_const=0, length=body_length, b=np.array([b1_body]), a=0*np.array([b1_body]), bs=0*b1_body, ny=5, nstep=nstep, h=h) +env.elements['e3']=xt.FieldExpansion(pkin_const=0, length=fringe_length/2, b=np.array([b1_out]), a=0*np.array([b1_out]), bs=0*b1_out, ny=5, nstep=nstep, h=h) +env.elements['e4']=xt.FieldExpansion(pkin_const=0, length=fringe_length/2, b=np.array([b1_out]), a=0*np.array([b1_out]), bs=0*b1_out, ny=5, nstep=nstep, sstart=fringe_length/2) +dipole = env.new_line(name="dipole",components=['e0', 'e1', 'e2', 'e3', 'e4']) p0 = xt.Particles() dipole.track(p0) ### Fodo -env=xt.Environment() env['k1quad'] = 1.6 env['lengthquad'] = 0.2 env.new("quad1", xt.Quadrupole, k1='k1quad', length='lengthquad') env.new("quad2", xt.Quadrupole, k1='-k1quad', length='lengthquad') -nstep=5 -env.elements['e0']=xt.FieldExpansion(length=fringe_length/2, b=np.array([b1_in]), a=0*np.array([b1_in]), bs=0*b1_in, ny=5, nstep=nstep) -env.elements['e1']=xt.FieldExpansion(length=fringe_length/2, b=np.array([b1_in]), a=0*np.array([b1_in]), bs=0*b1_in, ny=5, nstep=nstep, h=h) -env.elements['e2']=xt.FieldExpansion(length=body_length, b=np.array([b1_body]), a=0*np.array([b1_body]), bs=0*b1_body, ny=5, nstep=nstep, h=h) -env.elements['e3']=xt.FieldExpansion(length=fringe_length/2, b=np.array([b1_out]), a=0*np.array([b1_out]), bs=0*b1_out, ny=5, nstep=nstep, h=h) -env.elements['e4']=xt.FieldExpansion(length=fringe_length/2, b=np.array([b1_out]), a=0*np.array([b1_out]), bs=0*b1_out, ny=5, nstep=nstep) -env.new_line(name="dipole",components=['e0', 'e1', 'e2', 'e3', 'e4']) -fodo= env.new_line(name='myline', length=5, + +fodo= env.new_line(name='fodo', length=5, components=[ env.place('quad1', at=0.1), env.place('dipole', at=1.0), From ace0beb8375df6e5009654e74efdfc2f1bd1106c Mon Sep 17 00:00:00 2001 From: Riccardo De Maria Date: Tue, 1 Sep 2026 11:35:14 +0200 Subject: [PATCH 38/38] Expose FieldExpansion field evaluation API --- tests/test_fieldexpansion_element.py | 102 +++++++++++++- xtrack/beam_elements/elements.py | 124 +++++++++++++++++- .../elements_src/track_fieldexpansion.h | 108 +++++++++++---- .../elements_src/track_fieldexpansion_bent.h | 18 ++- .../track_fieldexpansion_helpers.h | 41 ++++-- .../track_fieldexpansion_straight.h | 18 ++- 6 files changed, 361 insertions(+), 50 deletions(-) diff --git a/tests/test_fieldexpansion_element.py b/tests/test_fieldexpansion_element.py index 45bdc6712..18804f8e4 100644 --- a/tests/test_fieldexpansion_element.py +++ b/tests/test_fieldexpansion_element.py @@ -2,6 +2,106 @@ import numpy as np import xobjects as xo + +def test_get_field_straight(): + a = np.array([[0.04, 0.2, 0.08], [0, 0, 0.1]]) + b = np.array([[0.05, 0.04, 0.07], [0.01, 0, 0]]) + bs = np.array([0.1, 0.02, 0]) + element = xt.StraightFieldExpansion( + length=1, a=a, b=b, bs=bs, ny=5) + + x = np.array([[0.01], [0.02]]) + y = np.array([0.005, -0.004, 0.003]) + s = np.array([0.001, 0.4, 0.9]) + x_broadcast, y_broadcast, s_broadcast = np.broadcast_arrays(x, y, s) + + bx_expected = ( + 0.1 * s_broadcast**2 * x_broadcast + + 0.08 * s_broadcast**2 + + 0.2 * s_broadcast + - y_broadcast**2 * (0.4 * x_broadcast + 0.32) / 4 + + 0.01 * y_broadcast + + 0.04 + ) + by_expected = ( + 0.07 * s_broadcast**2 + + 0.04 * s_broadcast + + 0.01 * x_broadcast + + y_broadcast**3 / 15 + - 0.07 * y_broadcast**2 + - y_broadcast * ( + 0.2 * s_broadcast**2 + + 0.2 * x_broadcast**2 + + 0.32 * x_broadcast + + 0.04 + ) / 2 + + 0.05 + ) + bs_expected = ( + 0.1 * s_broadcast * x_broadcast**2 + - 0.1 * s_broadcast * y_broadcast**2 + + 0.02 * s_broadcast + + x_broadcast * (0.16 * s_broadcast + 0.2) + + y_broadcast * (0.14 * s_broadcast + 0.04) + + 0.1 + ) + + field = element.get_field(x=x, y=y, s=s) + + assert isinstance(field, np.ndarray) + assert field.shape == x_broadcast.shape + assert field.dtype.names == ( + 'phi', + 'Bx', 'By', 'Bs', + 'Ax', 'Ay', 'As', + 'dAx_dx', 'dAx_dy', 'dAx_ds', + 'dAs_dx', 'dAs_dy', 'dAs_ds', + ) + xo.assert_allclose(field['Bx'], bx_expected, rtol=0, atol=1e-14) + xo.assert_allclose(field['By'], by_expected, rtol=0, atol=1e-14) + xo.assert_allclose(field['Bs'], bs_expected, rtol=0, atol=1e-14) + xo.assert_allclose(field['Ay'], 0, rtol=0, atol=1e-14) + xo.assert_allclose(field['dAx_dy'], -field['Bs'], rtol=0, atol=1e-14) + xo.assert_allclose(field['dAs_dy'], field['Bx'], rtol=0, atol=1e-14) + xo.assert_allclose( + field['dAx_ds'] - field['dAs_dx'], + field['By'], + rtol=0, + atol=1e-14, + ) + + scalar_field = element.get_field(x=x[0, 0], y=y[0], s=s[0]) + assert isinstance(scalar_field, np.ndarray) + assert scalar_field.shape == () + assert scalar_field.dtype == field.dtype + xo.assert_allclose( + scalar_field['Bx'], bx_expected[0, 0], rtol=0, atol=1e-14) + xo.assert_allclose( + scalar_field['By'], by_expected[0, 0], rtol=0, atol=1e-14) + xo.assert_allclose( + scalar_field['Bs'], bs_expected[0, 0], rtol=0, atol=1e-14) + + +def test_get_field_bent(): + element = xt.BentFieldExpansion( + length=1, + h=0.1, + a=np.array([[0.0]]), + b=np.array([[0.5]]), + bs=np.array([0.0]), + ny=5, + ) + x = np.array([-0.2, 0.0, 0.3]) + y = np.array([0.01, -0.03, 0.02]) + s = np.array([0.0, 0.5, 1.0]) + + field = element.get_field(x=x, y=y, s=s) + + xo.assert_allclose(field['Bx'], np.zeros(3), rtol=0, atol=1e-14) + xo.assert_allclose(field['By'], np.full(3, 0.5), rtol=0, atol=1e-14) + xo.assert_allclose(field['Bs'], np.zeros(3), rtol=0, atol=1e-14) + + def test_h_sdep(): h = 0.1 a = np.array([[1.0, 0.1], [0.2, 0.0], [0.3, 0.1]]) @@ -158,4 +258,4 @@ def straight_to_curved(p, h): assert np.isclose(curved_to_straight(p1, h)["x"], p0.x) assert np.isclose(curved_to_straight(p1, h)["s"], p0.s) - assert np.isclose(curved_to_straight(p1, h)["y"], p0.y) \ No newline at end of file + assert np.isclose(curved_to_straight(p1, h)["y"], p0.y) diff --git a/xtrack/beam_elements/elements.py b/xtrack/beam_elements/elements.py index a97995c64..a946a731f 100644 --- a/xtrack/beam_elements/elements.py +++ b/xtrack/beam_elements/elements.py @@ -5019,6 +5019,72 @@ def get_backtrack_element(self, _context=None, _buffer=None, _offset=None): class ThinSliceNotNeededError(Exception): pass + +_FIELD_VALUE_NAMES = ( + 'phi', + 'Bx', 'By', 'Bs', + 'Ax', 'Ay', 'As', + 'dAx_dx', 'dAx_dy', 'dAx_ds', + 'dAs_dx', 'dAs_dy', 'dAs_ds', +) +_FIELD_VALUE_DTYPE = np.dtype([ + (name, np.float64) for name in _FIELD_VALUE_NAMES +]) + + +def _field_expansion_get_field(element, x, y, s): + """Run a field-expansion evaluation kernel on broadcast input arrays.""" + context = element._context + + def _to_numpy(value): + if isinstance(value, context.nplike_array_type): + value = context.nparray_from_context_array(value) + return np.asarray(value, dtype=np.float64) + + x_arr, y_arr, s_arr = np.broadcast_arrays( + _to_numpy(x), _to_numpy(y), _to_numpy(s)) + output_shape = x_arr.shape + n_points = x_arr.size + + if not element.straight and np.any(1.0 + element.h * x_arr == 0.0): + raise ValueError("x contains a point on the singular curved coordinate axis") + + field_values = context.zeros( + n_points * len(_FIELD_VALUE_NAMES), dtype=np.float64) + + if n_points: + x_context = context.nparray_to_context_array( + np.ascontiguousarray(x_arr).reshape(-1)) + y_context = context.nparray_to_context_array( + np.ascontiguousarray(y_arr).reshape(-1)) + s_context = context.nparray_to_context_array( + np.ascontiguousarray(s_arr).reshape(-1)) + + n_values = int(element._ncoef) * int(element._nm) + work_v = context.zeros(n_points * n_values, dtype=np.float64) + work_d1 = context.zeros(n_points * n_values, dtype=np.float64) + work_d2 = context.zeros(n_points * n_values, dtype=np.float64) + work_q = context.zeros(n_points * int(element._nq), dtype=np.float64) + + element.compile_kernels(only_if_needed=True) + kernel = context.kernels[element._field_evaluation_kernel_name] + kernel( + el=element, + x=x_context, + y=y_context, + s=s_context, + n_points=n_points, + field_values=field_values, + work_v=work_v, + work_d1=work_d1, + work_d2=work_d2, + work_q=work_q, + ) + + field_values = context.nparray_from_context_array(field_values) + return np.asarray(field_values).view(_FIELD_VALUE_DTYPE).reshape(output_shape) + + class StraightFieldExpansion(BeamElement): """ Specifies the field expansion in general derivatives on axis in straight frame. @@ -5083,10 +5149,28 @@ class StraightFieldExpansion(BeamElement): '#include "xtrack/beam_elements/elements_src/create_fieldexpansion_straight.h"', ] + _field_evaluation_kernel_name = 'StraightFieldExpansion_get_field' + _kernels = {'build_expansion_straight': xo.Kernel( c_name='build_expansion_straight', args=[xo.Arg(xo.ThisClass, name='el')] - ), + ), + 'StraightFieldExpansion_get_field': xo.Kernel( + c_name='StraightFieldExpansion_get_field', + args=[ + xo.Arg(xo.ThisClass, name='el'), + xo.Arg(xo.Float64, pointer=True, const=True, name='x'), + xo.Arg(xo.Float64, pointer=True, const=True, name='y'), + xo.Arg(xo.Float64, pointer=True, const=True, name='s'), + xo.Arg(xo.Int64, name='n_points'), + xo.Arg(xo.Float64, pointer=True, name='field_values'), + xo.Arg(xo.Float64, pointer=True, name='work_v'), + xo.Arg(xo.Float64, pointer=True, name='work_d1'), + xo.Arg(xo.Float64, pointer=True, name='work_d2'), + xo.Arg(xo.Float64, pointer=True, name='work_q'), + ], + n_threads='n_points', + ), } def __init__(self, length, a, b, bs, ny, nstep=10, sstart=0, **kwargs): @@ -5139,6 +5223,15 @@ def __init__(self, length, a, b, bs, ny, nstep=10, sstart=0, **kwargs): super().__init__(**kwargs) self.build_expansion_straight(el=self) + + def get_field(self, x, y, s): + """Evaluate ``FieldValue`` at broadcastable ``x, y, s``. + + Returns a structured NumPy array with the broadcast input shape and + fields ``phi``, ``Bx``, ``By``, ``Bs``, ``Ax``, ``Ay``, ``As``, and + all derivatives stored by the C ``FieldValue`` structure. + """ + return _field_expansion_get_field(self, x=x, y=y, s=s) class BentFieldExpansion(BeamElement): """ @@ -5206,10 +5299,28 @@ class BentFieldExpansion(BeamElement): '#include "xtrack/beam_elements/elements_src/create_fieldexpansion_bent.h"', ] + _field_evaluation_kernel_name = 'BentFieldExpansion_get_field' + _kernels = {'build_expansion_bent': xo.Kernel( c_name='build_expansion_bent', args=[xo.Arg(xo.ThisClass, name='el')] - ), + ), + 'BentFieldExpansion_get_field': xo.Kernel( + c_name='BentFieldExpansion_get_field', + args=[ + xo.Arg(xo.ThisClass, name='el'), + xo.Arg(xo.Float64, pointer=True, const=True, name='x'), + xo.Arg(xo.Float64, pointer=True, const=True, name='y'), + xo.Arg(xo.Float64, pointer=True, const=True, name='s'), + xo.Arg(xo.Int64, name='n_points'), + xo.Arg(xo.Float64, pointer=True, name='field_values'), + xo.Arg(xo.Float64, pointer=True, name='work_v'), + xo.Arg(xo.Float64, pointer=True, name='work_d1'), + xo.Arg(xo.Float64, pointer=True, name='work_d2'), + xo.Arg(xo.Float64, pointer=True, name='work_q'), + ], + n_threads='n_points', + ), } def __init__(self, length, h, a, b, bs, ny, nstep=10, sstart=0, **kwargs): @@ -5265,6 +5376,15 @@ def __init__(self, length, h, a, b, bs, ny, nstep=10, sstart=0, **kwargs): self.build_expansion_bent(el=self) + def get_field(self, x, y, s): + """Evaluate ``FieldValue`` at broadcastable ``x, y, s``. + + Returns a structured NumPy array with the broadcast input shape and + fields ``phi``, ``Bx``, ``By``, ``Bs``, ``Ax``, ``Ay``, ``As``, and + all derivatives stored by the C ``FieldValue`` structure. + """ + return _field_expansion_get_field(self, x=x, y=y, s=s) + class FieldExpansion(BeamElement): def __new__(cls, *args, **kwargs): diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion.h b/xtrack/beam_elements/elements_src/track_fieldexpansion.h index a1b887096..7e49ff016 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion.h +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion.h @@ -1,3 +1,37 @@ +#ifndef CONCATDATA2 +#define CONCATDATA2(a,b) a##b +#endif +#ifndef CONCATDATA +#define CONCATDATA(a,b) CONCATDATA2(a,b) +#endif + +#ifndef FIELDEXPANSION_FIELD_VALUE_SIZE +#define FIELDEXPANSION_FIELD_VALUE_SIZE 13 +#endif + +GPUFUN +void CONCATDATA(DATA, _init_expansion)(FIELDEXPANSIONDATA el, Expansion *f) { + f->ny = CONCATDATA(DATA, _get_ny)(el); + f->ncoef = CONCATDATA(DATA, _get__ncoef)(el); + f->na = CONCATDATA(DATA, _get_na)(el); + f->nb = CONCATDATA(DATA, _get_nb)(el); + f->deg = CONCATDATA(DATA, _get_deg)(el); + f->mmin = CONCATDATA(DATA, _get__mmin)(el); + f->mmax = CONCATDATA(DATA, _get__mmax)(el); + f->moff = CONCATDATA(DATA, _get__moff)(el); + f->nm = CONCATDATA(DATA, _get__nm)(el); + f->qemin = CONCATDATA(DATA, _get__qemin)(el); + f->nq = CONCATDATA(DATA, _get__nq)(el); + f->h = CONCATDATA(DATA, _get_h)(el); + f->straight = CONCATDATA(DATA, _get_straight)(el); + f->c = (GPUGLMEM const double *)CONCATDATA(DATA, _getp__c)(el); + f->V = (GPUGLMEM double *)CONCATDATA(DATA, _getp__V)(el); + f->D1 = (GPUGLMEM double *)CONCATDATA(DATA, _getp__D1)(el); + f->D2 = (GPUGLMEM double *)CONCATDATA(DATA, _getp__D2)(el); + f->Q = (GPUGLMEM double *)CONCATDATA(DATA, _getp__Q)(el); +} + +GPUFUN void HAMILTONIAN_FLOW(Expansion *f, const double beta0, double s, const double z[6], HamiltonianFlow *flow) { double delta1, delta, ddelta1; @@ -37,12 +71,55 @@ void HAMILTONIAN_FLOW(Expansion *f, const double beta0, flow->dH_ds = -q * (pix * flow->pot.dAx_ds / root + flow->pot.dAs_ds); } -#ifndef CONCATDATA2 -#define CONCATDATA2(a,b) a##b -#endif -#ifndef CONCATDATA -#define CONCATDATA(a,b) CONCATDATA2(a,b) -#endif +GPUKERN +void GET_FIELD( + FIELDEXPANSIONDATA el, + GPUGLMEM const double *x, + GPUGLMEM const double *y, + GPUGLMEM const double *s, + const int64_t n_points, + GPUGLMEM double *field_values, + GPUGLMEM double *work_v, + GPUGLMEM double *work_d1, + GPUGLMEM double *work_d2, + GPUGLMEM double *work_q) +{ + VECTORIZE_OVER(ii, n_points); + Expansion f; + CONCATDATA(DATA, _init_expansion)(el, &f); + + const int64_t n_values = (int64_t)f.ncoef * (int64_t)f.nm; + f.V = work_v + ii * n_values; + f.D1 = work_d1 + ii * n_values; + f.D2 = work_d2 + ii * n_values; + f.Q = work_q + ii * (int64_t)f.nq; + + FieldValue field; + const int status = EVALUATE_EXPANSION( + &f, x[ii], y[ii], s[ii], &field); + const int64_t offset = ii * FIELDEXPANSION_FIELD_VALUE_SIZE; + if (status == 0) { + field_values[offset + 0] = field.phi; + field_values[offset + 1] = field.Bx; + field_values[offset + 2] = field.By; + field_values[offset + 3] = field.Bs; + field_values[offset + 4] = field.Ax; + field_values[offset + 5] = field.Ay; + field_values[offset + 6] = field.As; + field_values[offset + 7] = field.dAx_dx; + field_values[offset + 8] = field.dAx_dy; + field_values[offset + 9] = field.dAx_ds; + field_values[offset + 10] = field.dAs_dx; + field_values[offset + 11] = field.dAs_dy; + field_values[offset + 12] = field.dAs_ds; + } + else { + for (int jj = 0; jj < FIELDEXPANSION_FIELD_VALUE_SIZE; ++jj) { + field_values[offset + jj] = NAN; + } + } + END_VECTORIZE; +} void TRACK_EXPANSION( FIELDEXPANSIONDATA el, @@ -54,24 +131,7 @@ void TRACK_EXPANSION( double sstart = CONCATDATA(DATA, _get_sstart)(el); Expansion f; - f.ny = CONCATDATA(DATA, _get_ny)(el); - f.ncoef = CONCATDATA(DATA, _get__ncoef)(el); - f.na = CONCATDATA(DATA, _get_na)(el); - f.nb = CONCATDATA(DATA, _get_nb)(el); - f.deg = CONCATDATA(DATA, _get_deg)(el); - f.mmin = CONCATDATA(DATA, _get__mmin)(el); - f.mmax = CONCATDATA(DATA, _get__mmax)(el); - f.moff = CONCATDATA(DATA, _get__moff)(el); - f.nm = CONCATDATA(DATA, _get__nm)(el); - f.qemin = CONCATDATA(DATA, _get__qemin)(el); - f.nq = CONCATDATA(DATA, _get__nq)(el); - f.h = CONCATDATA(DATA, _get_h)(el); - f.straight = CONCATDATA(DATA, _get_straight)(el); - f.c = (double *)CONCATDATA(DATA, _getp__c)(el); - f.V = (double *)CONCATDATA(DATA, _getp__V)(el); - f.D1 = (double *)CONCATDATA(DATA, _getp__D1)(el); - f.D2 = (double *)CONCATDATA(DATA, _getp__D2)(el); - f.Q = (double *)CONCATDATA(DATA, _getp__Q)(el); + CONCATDATA(DATA, _init_expansion)(el, &f); int64_t const backtrack = LocalParticle_check_track_flag(part0, XS_FLAG_BACKTRACK); if (backtrack) {sstart = ds * nstep; ds = -ds;} diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h b/xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h index d955fb3f6..40ee6d0b7 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h @@ -4,16 +4,18 @@ #include "track_fieldexpansion_helpers.h" -int evaluate_expansion_bent(Expansion *f, double x, double y, double s, FieldValue *out) { +GPUFUN +int evaluate_expansion_bent(Expansion *f, double x, double y, double s, + FieldValue *out) { const double q = 1.0 + f->h * x; if (q == 0.0) return -1; /* singular chart */ - memset(out, 0, sizeof(*out)); + fieldexpansion_reset_field_value(out); - double *V = f->V; - double *D1 = f->D1; - double *D2 = f->D2; - double *Q = f->Q; + GPUGLMEM double *V = f->V; + GPUGLMEM double *D1 = f->D1; + GPUGLMEM double *D2 = f->D2; + GPUGLMEM double *Q = f->Q; fs_prepare_s(f, s); @@ -99,13 +101,15 @@ int evaluate_expansion_bent(Expansion *f, double x, double y, double s, FieldVal #define TRACK_EXPANSION BentFieldExpansion_track_local_particle #define HAMILTONIAN_FLOW hamiltonian_flow_bent #define EVALUATE_EXPANSION evaluate_expansion_bent +#define GET_FIELD BentFieldExpansion_get_field #define FIELDEXPANSIONDATA BentFieldExpansionData #define DATA BentFieldExpansionData #include "track_fieldexpansion.h" #undef TRACK_EXPANSION #undef HAMILTONIAN_FLOW #undef EVALUATE_EXPANSION +#undef GET_FIELD #undef FIELDEXPANSIONDATA #undef DATA -#endif \ No newline at end of file +#endif diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion_helpers.h b/xtrack/beam_elements/elements_src/track_fieldexpansion_helpers.h index f1e49c8c9..c96638765 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion_helpers.h +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion_helpers.h @@ -13,11 +13,11 @@ typedef struct { int qemin, nq; double h; double straight; - double *c; /* c[i,m,k], polynomial coeff of s^k in q^m term */ - double *V; /* scratch: c[i,m](s) */ - double *D1; /* scratch: d_s c[i,m] */ - double *D2; /* scratch: d2_s c[i,m] */ - double *Q; /* scratch: q^e, e=qemin.. */ + GPUGLMEM const double *c; /* c[i,m,k], polynomial coeff of s^k in q^m term */ + GPUGLMEM double *V; /* scratch: c[i,m](s) */ + GPUGLMEM double *D1; /* scratch: d_s c[i,m] */ + GPUGLMEM double *D2; /* scratch: d2_s c[i,m] */ + GPUGLMEM double *Q; /* scratch: q^e, e=qemin.. */ } Expansion; typedef struct { @@ -40,11 +40,32 @@ typedef struct { FieldValue pot; } HamiltonianFlow; -static inline const double *ccptr(const Expansion *f, int i, int m) { +GPUFUN +void fieldexpansion_reset_field_value(FieldValue *out) { + out->phi = 0.0; + out->Bx = 0.0; + out->By = 0.0; + out->Bs = 0.0; + out->Ax = 0.0; + out->Ay = 0.0; + out->As = 0.0; + out->dAx_dx = 0.0; + out->dAx_dy = 0.0; + out->dAx_ds = 0.0; + out->dAs_dx = 0.0; + out->dAs_dy = 0.0; + out->dAs_ds = 0.0; +} + +GPUFUN +GPUGLMEM const double *ccptr(const Expansion *f, int i, int m) { return f->c + (((size_t)i * (size_t)f->nm + (size_t)m) * (size_t)(f->deg + 1)); } -static inline void poly_eval_d2(const double *p, int deg, double s, double *v, double *d1, double *d2) { +GPUFUN +void poly_eval_d2(GPUGLMEM const double *p, int deg, double s, + GPUGLMEM double *v, GPUGLMEM double *d1, + GPUGLMEM double *d2) { double a = p[deg], b = 0.0, c=0.0; for (int k = deg - 1; k >= 0; --k) { c = c * s + 2.0 * b; @@ -56,7 +77,8 @@ static inline void poly_eval_d2(const double *p, int deg, double s, double *v, d *d2 = c; } -static void fs_prepare_s(Expansion *f, double s) { +GPUFUN +void fs_prepare_s(Expansion *f, double s) { for (int i = 0; i < f->ncoef; ++i) { for (int m = 0; m < f->nm; ++m) { poly_eval_d2(ccptr(f, i, m), f->deg, s, @@ -67,6 +89,7 @@ static void fs_prepare_s(Expansion *f, double s) { } } +GPUFUN void delta_from_ptau(const double beta0, double ptau, double *delta, double *delta1, double *ddelta1) { { @@ -77,4 +100,4 @@ void delta_from_ptau(const double beta0, double ptau, } } -#endif \ No newline at end of file +#endif diff --git a/xtrack/beam_elements/elements_src/track_fieldexpansion_straight.h b/xtrack/beam_elements/elements_src/track_fieldexpansion_straight.h index 71960d44e..fea987490 100644 --- a/xtrack/beam_elements/elements_src/track_fieldexpansion_straight.h +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion_straight.h @@ -3,14 +3,16 @@ #include "track_fieldexpansion_helpers.h" -int evaluate_expansion_straight(Expansion *f, double x, double y, double s, FieldValue *out) { +GPUFUN +int evaluate_expansion_straight(Expansion *f, double x, double y, double s, + FieldValue *out) { - memset(out, 0, sizeof(*out)); + fieldexpansion_reset_field_value(out); - double *V = f->V; - double *D1 = f->D1; - double *D2 = f->D2; - double *X = f->Q; + GPUGLMEM double *V = f->V; + GPUGLMEM double *D1 = f->D1; + GPUGLMEM double *D2 = f->D2; + GPUGLMEM double *X = f->Q; fs_prepare_s(f, s); @@ -90,13 +92,15 @@ int evaluate_expansion_straight(Expansion *f, double x, double y, double s, Fiel #define TRACK_EXPANSION StraightFieldExpansion_track_local_particle #define HAMILTONIAN_FLOW hamiltonian_flow_straight #define EVALUATE_EXPANSION evaluate_expansion_straight +#define GET_FIELD StraightFieldExpansion_get_field #define FIELDEXPANSIONDATA StraightFieldExpansionData #define DATA StraightFieldExpansionData #include "track_fieldexpansion.h" #undef TRACK_EXPANSION #undef HAMILTONIAN_FLOW #undef EVALUATE_EXPANSION +#undef GET_FIELD #undef FIELDEXPANSIONDATA #undef DATA -#endif \ No newline at end of file +#endif