diff --git a/examples/fieldexpansion/combined_function_fodo.py b/examples/fieldexpansion/combined_function_fodo.py new file mode 100644 index 000000000..3784ad4a1 --- /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, 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() diff --git a/examples/fieldexpansion/compare_dipole_fringes.py b/examples/fieldexpansion/compare_dipole_fringes.py new file mode 100644 index 000000000..87cf760a2 --- /dev/null +++ b/examples/fieldexpansion/compare_dipole_fringes.py @@ -0,0 +1,53 @@ +""" +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 + +# 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 + +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, 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: 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() + +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 diff --git a/examples/fieldexpansion/dipole_with_fringes.py b/examples/fieldexpansion/dipole_with_fringes.py new file mode 100644 index 000000000..d169976e5 --- /dev/null +++ b/examples/fieldexpansion/dipole_with_fringes.py @@ -0,0 +1,95 @@ +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 # magnetic length +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" + +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['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') + +fodo= env.new_line(name='fodo', 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],',') + + + + + 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') + + + + diff --git a/tests/test_fieldexpansion_element.py b/tests/test_fieldexpansion_element.py new file mode 100644 index 000000000..18804f8e4 --- /dev/null +++ b/tests/test_fieldexpansion_element.py @@ -0,0 +1,261 @@ +import xtrack as xt +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]]) + 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, pkin_const=True) + + 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.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) + +def test_sdep(): + 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, 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]) + line.track(p0, _force_no_end_turn_actions=True) + + 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(): + 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(q0=1, mass0=1) + tw = fodo.twiss4d() + + myfodo = xt.Line(elements=[ + xt.Drift(length=1.2), + 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.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, a=np.array([[0]]), b=np.array([[0],[-7]]), bs=np.array([0]), ny=5) + ]) + myfodo.particle_ref = xt.Particles(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) + +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) + +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*(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) + + 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) + +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) diff --git a/xtrack/beam_elements/elements.py b/xtrack/beam_elements/elements.py index 692f9d35e..a946a731f 100644 --- a/xtrack/beam_elements/elements.py +++ b/xtrack/beam_elements/elements.py @@ -4943,7 +4943,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. """ @@ -5018,3 +5018,378 @@ 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. + + Parameters + ---------- + 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 = True + 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[:], + + "pkin_const": xo.Int64, + "sstart": 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"', + ] + + _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): + 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() + 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 "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") + + 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 + + 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_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): + """ + 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, 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 + 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 = True + 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[:], + + "pkin_const": xo.Int64, + "sstart": xo.Float64, + } + + _extra_c_sources = [ + '#include "xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h"', + '#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): + 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() + 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 "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") + + 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 + + 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_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): + if 'h' in kwargs and kwargs['h'] > 1e-9: + return BentFieldExpansion(*args, **kwargs) + else: + kwargs.pop('h', None) + return StraightFieldExpansion(*args, **kwargs) 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..a0ab60cb9 --- /dev/null +++ b/xtrack/beam_elements/elements_src/create_fieldexpansion_bent.h @@ -0,0 +1,96 @@ +#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], +... 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 +*/ + + +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..3b48fb1bd --- /dev/null +++ b/xtrack/beam_elements/elements_src/create_fieldexpansion_straight.h @@ -0,0 +1,78 @@ +#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], +... 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 +*/ + +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 new file mode 100644 index 000000000..7e49ff016 --- /dev/null +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion.h @@ -0,0 +1,209 @@ +#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; + 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); +} + +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, + LocalParticle* part0) +{ + + const double nstep = CONCATDATA(DATA, _get_nstep(el)); + double ds = CONCATDATA(DATA, _get_ds(el)); + double sstart = CONCATDATA(DATA, _get_sstart)(el); + + Expansion f; + 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;} + + int pkin_const = CONCATDATA(DATA, _get_pkin_const)(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 + 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]; + 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], s, &v); + 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]); + 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); + 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..40ee6d0b7 --- /dev/null +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion_bent.h @@ -0,0 +1,115 @@ +#ifndef XTRACK_TRACK_FIELDEXPANSION_BENT_H +#define XTRACK_TRACK_FIELDEXPANSION_BENT_H + +#include "track_fieldexpansion_helpers.h" + + +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 */ + + fieldexpansion_reset_field_value(out); + + 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); + + /* 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+1)/(i+1)! */ + 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 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 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..c96638765 --- /dev/null +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion_helpers.h @@ -0,0 +1,103 @@ +#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 */ + int na, nb, deg; + int mmin, mmax, moff, nm; + int qemin, nq; + double h; + double straight; + 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 { + 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; + +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)); +} + +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; + b = b * s + a; + a = a * s + p[k]; + } + *v = a; + *d1 = b; + *d2 = c; +} + +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, + &f->V[i * f->nm + m], + &f->D1[i * f->nm + m], + &f->D2[i * f->nm + m]); + } + } +} + +GPUFUN +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); + } +} + +#endif 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..fea987490 --- /dev/null +++ b/xtrack/beam_elements/elements_src/track_fieldexpansion_straight.h @@ -0,0 +1,106 @@ +#ifndef XTRACK_TRACK_FIELDEXPANSION_STRAIGHT_H +#define XTRACK_TRACK_FIELDEXPANSION_STRAIGHT_H + +#include "track_fieldexpansion_helpers.h" + +GPUFUN +int evaluate_expansion_straight(Expansion *f, double x, double y, double s, + FieldValue *out) { + + fieldexpansion_reset_field_value(out); + + 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); + + /* 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; /* -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; /* -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 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 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,