xtrack 0.60.0__tar.gz → 0.61.0__tar.gz

{xtrack-0.60.0/xtrack.egg-info → xtrack-0.61.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xtrack
-Version: 0.60.0
+Version: 0.61.0
 Summary: Tracking library for particle accelerators
 Home-page: https://xsuite.readthedocs.io/
 Download-URL: https://pypi.python.org/pypi/xtrack

{xtrack-0.60.0 → xtrack-0.61.0}/ducktrack/elements.py

@@ -651,14 +651,18 @@ class LinearTransferMatrix(Element):
         ("y_ref_1","","",0.0),
         ("py_ref_1","","",0.0),
         ("damping_rate_x","","",0.0),
+        ("damping_rate_px","","",0.0),
         ("damping_rate_y","","",0.0),
-        ("damping_rate_s","","",0.0),
-        ("equ_emit_x","","",0.0),
-        ("equ_emit_y","","",0.0),
-        ("equ_emit_s","","",0.0),
+        ("damping_rate_py","","",0.0),
+        ("damping_rate_zeta","","",0.0),
+        ("damping_rate_pzeta","","",0.0),
         ("gauss_noise_ampl_x","","",0.0),
+        ("gauss_noise_ampl_px","","",0.0),
         ("gauss_noise_ampl_y","","",0.0),
-        ("gauss_noise_ampl_s","","",0.0)]
+        ("gauss_noise_ampl_py","","",0.0),
+        ("gauss_noise_ampl_zeta","","",0.0),
+        ("gauss_noise_ampl_pzeta","","",0.0),
+        ("damping_matrix","","",0.0),]
 
     def track(self,p):
         sin = p._m.sin
@@ -746,37 +750,50 @@ class LinearTransferMatrix(Element):
             p.y *= geo_emit_factor
             p.py = old_py * geo_emit_factor
 
-        if self.damping_rate_x < 0.0 or self.damping_rate_y < 0.0 or self.damping_rate_s < 0.0:
-            raise ValueError("Damping rates cannot be negative")
-        if self.equ_emit_x < 0.0 or self.equ_emit_y < 0.0 or self.equ_emit_s < 0.0:
-            raise ValueError("Equilibrium emittance cannot be negative")
-        if self.damping_rate_x > 0.0:
-            factor = 1.0-self.damping_rate_x/2 
-            p.x *= factor
-            p.px *= factor 
-            if self.equ_emit_x > 0.0:
-                p.px += np.sqrt(2.0*self.equ_emit_x*self.damping_rate_x/self.beta_x_1)*np.random.randn(len(p.px))
-        if self.damping_rate_y > 0.0:
-            factor = 1.0-self.damping_rate_y/2 
-            p.y *= factor
-            p.py *= factor 
-            if self.equ_emit_y > 0.0:
-                p.py += np.sqrt(2.0*self.equ_emit_y*self.damping_rate_y/self.beta_y_1)*np.random.randn(len(p.py))
-        if self.damping_rate_s > 0.0:
-            factor = 1.0-self.damping_rate_s/2
-            p.zeta *= factor
-            p.delta *= factor
-            if self.equ_emit_s > 0.0:
-                p.delta += np.sqrt(2.0*self.equ_emit_s*self.damping_rate_s/self.beta_s)*np.random.randn(len(p.delta))
-
-        if self.gauss_noise_ampl_x < 0.0 or self.gauss_noise_ampl_y < 0.0 or self.gauss_noise_ampl_s < 0.0:
+        if not hasattr(self.damping_matrix,'__iter__'):
+            if (self.damping_rate_x < 0.0 or self.damping_rate_px < 0.0
+                    or self.damping_rate_y < 0.0 or self.damping_rate_py < 0.0
+                    or self.damping_rate_zeta < 0.0 or self.damping_rate_pzeta < 0.0):
+                raise ValueError("Damping rates cannot be negative")
+            if self.damping_rate_x > 0.0:
+                p.x *= 1.0-self.damping_rate_x
+            if self.damping_rate_px > 0.0:
+                p.px *= 1.0-self.damping_rate_px
+            if self.damping_rate_y > 0.0:
+                p.y *= 1.0-self.damping_rate_y
+            if self.damping_rate_py > 0.0:
+                p.py *= 1.0-self.damping_rate_py
+            if self.damping_rate_zeta > 0.0:
+                p.zeta *= 1.0-self.damping_rate_zeta
+            if self.damping_rate_pzeta > 0.0:
+                p.pzeta *= 1.0-self.damping_rate_pzeta
+        else:
+            assert np.shape(self.damping_matrix) == (6,6)
+            transformation = self.damping_matrix + np.identity(6)
+            p.x,p.px,p.y,p.py,p.zeta,p.pzeta = transformation[0,0]*p.x+transformation[0,1]*p.px+transformation[0,2]*p.y+transformation[0,3]*p.py+transformation[0,4]*p.zeta+transformation[0,5]*p.pzeta,\
+                                               transformation[1,0]*p.x+transformation[1,1]*p.px+transformation[1,2]*p.y+transformation[1,3]*p.py+transformation[1,4]*p.zeta+transformation[1,5]*p.pzeta,\
+                                               transformation[2,0]*p.x+transformation[2,1]*p.px+transformation[2,2]*p.y+transformation[2,3]*p.py+transformation[2,4]*p.zeta+transformation[2,5]*p.pzeta,\
+                                               transformation[3,0]*p.x+transformation[3,1]*p.px+transformation[3,2]*p.y+transformation[3,3]*p.py+transformation[3,4]*p.zeta+transformation[3,5]*p.pzeta,\
+                                               transformation[4,0]*p.x+transformation[4,1]*p.px+transformation[4,2]*p.y+transformation[4,3]*p.py+transformation[4,4]*p.zeta+transformation[4,5]*p.pzeta,\
+                                               transformation[5,0]*p.x+transformation[5,1]*p.px+transformation[5,2]*p.y+transformation[5,3]*p.py+transformation[5,4]*p.zeta+transformation[5,5]*p.pzeta
+
+
+        if (self.gauss_noise_ampl_x < 0.0 or self.gauss_noise_ampl_px < 0.0 
+                or self.gauss_noise_ampl_y < 0.0 or self.gauss_noise_ampl_py < 0.0
+                or self.gauss_noise_ampl_zeta < 0.0 or self.gauss_noise_ampl_pzeta < 0.0):
             raise ValueError("Noise amplitude cannot be negative")
+        if self.gauss_noise_ampl_x > 0.0:
+            p.x += self.gauss_noise_ampl_x*np.random.randn(len(p.x))
         if self.gauss_noise_ampl_x > 0.0:
             p.px += self.gauss_noise_ampl_x*np.random.randn(len(p.px))
         if self.gauss_noise_ampl_y > 0.0:
-            p.py += self.gauss_noise_ampl_y*np.random.randn(len(p.py))
-        if self.gauss_noise_ampl_s > 0.0:
-            p.delta += self.gauss_noise_ampl_s*np.random.randn(len(p.delta))
+            p.y += self.gauss_noise_ampl_y*np.random.randn(len(p.y))
+        if self.gauss_noise_ampl_py > 0.0:
+            p.py += self.gauss_noise_ampl_py*np.random.randn(len(p.py))
+        if self.gauss_noise_ampl_zeta > 0.0:
+            p.zeta += self.gauss_noise_ampl_zeta*np.random.randn(len(p.zeta))
+        if self.gauss_noise_ampl_pzeta > 0.0:
+            p.pzeta += self.gauss_noise_ampl_pzeta*np.random.randn(len(p.pzeta))
 
         # re-adding dispersion and closed orbit
         old_x=p.x

{xtrack-0.60.0 → xtrack-0.61.0}/tests/test_elements.py

@@ -771,8 +771,11 @@ def test_simplified_accelerator_segment_uncorrelated_damping(test_context):
     y_ref_1 = 4E-2
     py_ref_1 = 5E-4
     damping_rate_x = 5E-4
+    damping_rate_px = 2E-4
     damping_rate_y = 1E-3
-    damping_rate_s = 2E-3
+    damping_rate_py = 7E-3
+    damping_rate_zeta = 2E-3
+    damping_rate_pzeta = 1E-2
 
     dtk_particle = dtk.TestParticles(
             p0c=25.92e9,
@@ -798,26 +801,113 @@ def test_simplified_accelerator_segment_uncorrelated_damping(test_context):
         energy_increment=energy_increment,
         x_ref=(x_ref_0, x_ref_1), px_ref=(px_ref_0, px_ref_1),
         y_ref=(y_ref_0, y_ref_1), py_ref=(py_ref_0, py_ref_1),
-        damping_rate_x = damping_rate_x,
-        damping_rate_y = damping_rate_y,
-        damping_rate_s = damping_rate_s)
+        damping_rate_x = damping_rate_x,damping_rate_px = damping_rate_px,
+        damping_rate_y = damping_rate_y,damping_rate_py = damping_rate_py,
+        damping_rate_zeta = damping_rate_zeta,damping_rate_pzeta = damping_rate_pzeta)
 
     arc.track(particles)
 
     dtk_arc = dtk.elements.LinearTransferMatrix(alpha_x_0=alpha_x_0, beta_x_0=beta_x_0, disp_x_0=disp_x_0,
-    alpha_x_1=alpha_x_1, beta_x_1=beta_x_1, disp_x_1=disp_x_1,
-    alpha_y_0=alpha_y_0, beta_y_0=beta_y_0, disp_y_0=disp_y_0,
-    alpha_y_1=alpha_y_1, beta_y_1=beta_y_1, disp_y_1=disp_y_1,
-    Q_x=Q_x, Q_y=Q_y,
-    beta_s=beta_s, Q_s=Q_s,
-    chroma_x=0.0, chroma_y=0.0,
-    det_xx=0.0, det_xy=0.0, det_yy=0.0, det_yx=0.0,
-    energy_ref_increment=energy_ref_increment,energy_increment=energy_increment,
-    x_ref_0 = x_ref_0, px_ref_0 = px_ref_0, x_ref_1 = x_ref_1, px_ref_1 = px_ref_1,
-    y_ref_0 = y_ref_0, py_ref_0 = py_ref_0, y_ref_1 = y_ref_1, py_ref_1 = py_ref_1,
-    damping_rate_x = damping_rate_x,damping_rate_y = damping_rate_y,damping_rate_s = damping_rate_s)
+        alpha_x_1=alpha_x_1, beta_x_1=beta_x_1, disp_x_1=disp_x_1,
+        alpha_y_0=alpha_y_0, beta_y_0=beta_y_0, disp_y_0=disp_y_0,
+        alpha_y_1=alpha_y_1, beta_y_1=beta_y_1, disp_y_1=disp_y_1,
+        Q_x=Q_x, Q_y=Q_y,
+        beta_s=beta_s, Q_s=Q_s,
+        chroma_x=0.0, chroma_y=0.0,
+        det_xx=0.0, det_xy=0.0, det_yy=0.0, det_yx=0.0,
+        energy_ref_increment=energy_ref_increment,energy_increment=energy_increment,
+        x_ref_0 = x_ref_0, px_ref_0 = px_ref_0, x_ref_1 = x_ref_1, px_ref_1 = px_ref_1,
+        y_ref_0 = y_ref_0, py_ref_0 = py_ref_0, y_ref_1 = y_ref_1, py_ref_1 = py_ref_1,
+        damping_rate_x = damping_rate_x,damping_rate_px = damping_rate_px,
+        damping_rate_y = damping_rate_y,damping_rate_py = damping_rate_py,
+        damping_rate_zeta = damping_rate_zeta,damping_rate_pzeta = damping_rate_pzeta)
     dtk_arc.track(dtk_particle)
+    
+    assert np.isclose(test_context.nparray_from_context_array(particles.x)[0],
+                      dtk_particle.x, rtol=1e-14, atol=1e-14)
+    assert np.isclose(test_context.nparray_from_context_array(particles.px)[0],
+                      dtk_particle.px, rtol=1e-14, atol=1e-14)
+    assert np.isclose(test_context.nparray_from_context_array(particles.y)[0],
+                      dtk_particle.y, rtol=1e-14, atol=1e-14)
+    assert np.isclose(test_context.nparray_from_context_array(particles.py)[0],
+                      dtk_particle.py, rtol=1e-14, atol=1e-14)
+    assert np.isclose(test_context.nparray_from_context_array(particles.zeta)[0],
+                      dtk_particle.zeta, rtol=1e-14, atol=1e-14)
+    assert np.isclose(test_context.nparray_from_context_array(particles.delta)[0],
+                      dtk_particle.delta, rtol=1e-14, atol=1e-14)
+
+@for_all_test_contexts
+def test_simplified_accelerator_segment_correlated_damping(test_context):
+    alpha_x_0 = -0.5
+    beta_x_0 = 100.0
+    disp_x_0 = 1.8
+    alpha_x_1 = 2.1
+    beta_x_1 = 2.0
+    disp_x_1 = 3.3
+    alpha_y_0 = -0.4
+    beta_y_0 = 8.0
+    disp_y_0 = -0.2
+    alpha_y_1 = 0.7
+    beta_y_1 = 0.3
+    disp_y_1 = -1.9
+    Q_x = 0.27
+    Q_y = 0.34
+    beta_s = 856.9
+    Q_s = 0.001
+    energy_ref_increment = 1.2E9
+    energy_increment = 4.8E8
+    x_ref_0 = -5E-3
+    px_ref_0 = 6E-4
+    x_ref_1 = 2E-2
+    px_ref_1 = -5E-5
+    y_ref_0 = -9E-2
+    py_ref_0 = 1E-4
+    y_ref_1 = 4E-2
+    py_ref_1 = 5E-4
+    damping_matrix = np.reshape(np.random.randn(36),(6,6))
+
+    dtk_particle = dtk.TestParticles(
+            p0c=25.92e9,
+            x=1e-3,
+            px=1e-5,
+            y=-2e-3,
+            py=-1.5e-5,
+            zeta=2.,
+            delta=2E-4)
+
+    particles = xp.Particles.from_dict(dtk_particle.to_dict(),
+                                       _context=test_context)
+
+
+    arc = xt.LineSegmentMap(_context=test_context,
+        alfx=(alpha_x_0, alpha_x_1), betx=(beta_x_0, beta_x_1),
+        dx=(disp_x_0, disp_x_1), dpx=(0.0, 0.0),
+        alfy=(alpha_y_0, alpha_y_1), bety=(beta_y_0, beta_y_1),
+        dy=(disp_y_0, disp_y_1), dpy=(0.0, 0.0),
+        qx=Q_x, qy=Q_y,
+        bets=beta_s, qs=Q_s,
+        energy_ref_increment=energy_ref_increment,
+        energy_increment=energy_increment,
+        x_ref=(x_ref_0, x_ref_1), px_ref=(px_ref_0, px_ref_1),
+        y_ref=(y_ref_0, y_ref_1), py_ref=(py_ref_0, py_ref_1),
+        damping_matrix = damping_matrix)
+
+    arc.track(particles)
 
+    dtk_arc = dtk.elements.LinearTransferMatrix(alpha_x_0=alpha_x_0, beta_x_0=beta_x_0, disp_x_0=disp_x_0,
+        alpha_x_1=alpha_x_1, beta_x_1=beta_x_1, disp_x_1=disp_x_1,
+        alpha_y_0=alpha_y_0, beta_y_0=beta_y_0, disp_y_0=disp_y_0,
+        alpha_y_1=alpha_y_1, beta_y_1=beta_y_1, disp_y_1=disp_y_1,
+        Q_x=Q_x, Q_y=Q_y,
+        beta_s=beta_s, Q_s=Q_s,
+        chroma_x=0.0, chroma_y=0.0,
+        det_xx=0.0, det_xy=0.0, det_yy=0.0, det_yx=0.0,
+        energy_ref_increment=energy_ref_increment,energy_increment=energy_increment,
+        x_ref_0 = x_ref_0, px_ref_0 = px_ref_0, x_ref_1 = x_ref_1, px_ref_1 = px_ref_1,
+        y_ref_0 = y_ref_0, py_ref_0 = py_ref_0, y_ref_1 = y_ref_1, py_ref_1 = py_ref_1,
+        damping_matrix = damping_matrix)
+    dtk_arc.track(dtk_particle)
+    
     assert np.isclose(test_context.nparray_from_context_array(particles.x)[0],
                       dtk_particle.x, rtol=1e-14, atol=1e-14)
     assert np.isclose(test_context.nparray_from_context_array(particles.px)[0],
@@ -841,9 +931,17 @@ def test_simplified_accelerator_segment_uncorrelated_damping_equilibrium(test_co
     Q_y = 0.22
     beta_s = 856.9
     Q_s = 0.015
-    damping_rate_x = 5E-4
-    damping_rate_y = 1E-3
+    damping_rate_h = 5E-4
+    damping_rate_v = 1E-3
     damping_rate_s = 2E-3
+    
+    damping_rate_x = damping_rate_h/2
+    damping_rate_px = damping_rate_h/2
+    damping_rate_y = damping_rate_v/2
+    damping_rate_py = damping_rate_v/2
+    damping_rate_zeta = 0.0
+    damping_rate_pzeta = damping_rate_s
+    
     energy = 45.6
     equ_emit_x = 0.3E-9
     equ_emit_y = 1E-12
@@ -852,6 +950,13 @@ def test_simplified_accelerator_segment_uncorrelated_damping_equilibrium(test_co
     beta_s = equ_length/equ_delta
     equ_emit_s = equ_length*equ_delta
 
+    
+    gauss_noise_ampl_px = np.sqrt(equ_emit_x*damping_rate_h/beta_x_0)
+    gauss_noise_ampl_x = beta_x_0*gauss_noise_ampl_px
+    gauss_noise_ampl_py = np.sqrt(equ_emit_y*damping_rate_v/beta_y_0)
+    gauss_noise_ampl_y = beta_y_0*gauss_noise_ampl_py
+    gauss_noise_ampl_delta = np.sqrt(2*equ_emit_s*damping_rate_s/beta_s)
+    
     npart = int(1E3)
     particles = xp.Particles(_context=test_context,
                 x=np.random.randn(npart)*np.sqrt(equ_emit_x*beta_x_0),
@@ -863,7 +968,6 @@ def test_simplified_accelerator_segment_uncorrelated_damping_equilibrium(test_co
                 p0c=energy*1E9)
     particles._init_random_number_generator()
 
-
     arc = xt.LineSegmentMap(_context=test_context,
         alfx=(alpha_x_0, alpha_x_0), betx=(beta_x_0, beta_x_0),
         dx=(0.0, 0.0), dpx=(0.0, 0.0),
@@ -871,10 +975,12 @@ def test_simplified_accelerator_segment_uncorrelated_damping_equilibrium(test_co
         dy=(0.0, 0.0), dpy=(0.0, 0.0),
         qx=Q_x, qy=Q_y,
         bets=beta_s, qs=Q_s,
-        damping_rate_x = damping_rate_x,
-        damping_rate_y = damping_rate_y,
-        damping_rate_s = damping_rate_s,
-        equ_emit_x=equ_emit_x, equ_emit_y=equ_emit_y, equ_emit_s=equ_emit_s)
+        damping_rate_x=damping_rate_x,damping_rate_px=damping_rate_px,
+        damping_rate_y=damping_rate_y,damping_rate_py=damping_rate_py,
+        damping_rate_zeta=0.0,damping_rate_pzeta=damping_rate_pzeta,
+        gauss_noise_ampl_x = gauss_noise_ampl_x, gauss_noise_ampl_px = gauss_noise_ampl_px, 
+        gauss_noise_ampl_y = gauss_noise_ampl_y, gauss_noise_ampl_py = gauss_noise_ampl_py,
+        gauss_noise_ampl_zeta = 0.0, gauss_noise_ampl_pzeta = gauss_noise_ampl_delta)
 
     gamma_x = (1.0+alpha_x_0**2)/beta_x_0
     gamma_y = (1.0+alpha_y_0**2)/beta_y_0
@@ -903,6 +1009,38 @@ def test_simplified_accelerator_segment_uncorrelated_damping_equilibrium(test_co
     assert np.isclose(equ_emit_s,equ_emit_s_0, rtol=1e-1, atol=1e-10)
 
 
+@for_all_test_contexts
+def test_simplified_accelerator_segment_correlated_noise(test_context):
+    npart = int(1E6)
+    scale = 1E-6
+    random_matrix = np.reshape(np.random.rand(36),(6,6))
+    data = np.transpose(np.random.multivariate_normal(np.zeros(6),random_matrix,npart))
+    covariance_matrix = np.cov(data)
+
+    particles = xp.Particles(_context=test_context,
+                x=np.zeros(npart),
+                p0c=45E9)
+    particles._init_random_number_generator()
+
+    arc = xt.LineSegmentMap(_context=test_context,
+        betx=1.0, bety=1.0,bets=1.0,
+        qx=0.0, qy=0.0, qs=0.0,
+        gauss_noise_matrix=covariance_matrix*scale
+        )
+
+    arc.track(particles)
+    data = np.zeros((6,npart))
+    particles.move(_context=xo.context_default)
+    data[0,:] = particles.x
+    data[1,:] = particles.px
+    data[2,:] = particles.y
+    data[3,:] = particles.py
+    data[4,:] = particles.zeta
+    data[5,:] = particles.pzeta
+    cov = np.cov(data)/scale
+    xo.assert_allclose(cov,covariance_matrix,atol=1E-4,rtol=0.1)
+
+
 @for_all_test_contexts
 def test_nonlinearlens(test_context):
     mad = Madx(stdout=False)

{xtrack-0.60.0 → xtrack-0.61.0}/tests/test_tracker.py

@@ -719,3 +719,93 @@ def test_track_log_and_merit_function(pimms_mad, test_context):
     assert np.isclose(intensity[0], intensity_before, atol=1e-14, rtol=0)
     # The last log point is from the beginning of the last turn:
     assert np.isclose(intensity[-1], intensity_after, atol=0, rtol=1e-2)
+
+
+@for_all_test_contexts
+def test_init_io_buffer(test_context):
+    class TestElementRecord(xo.HybridClass):
+        _xofields = {
+            '_index': xt.RecordIndex,
+            'record_field': xo.Int64[:],
+            'record_at_element': xo.Int64[:],
+        }
+
+    class TestElement(xt.BeamElement):
+        _xofields={
+            'element_field': xo.Int64,
+        }
+
+        _internal_record_class = TestElementRecord
+
+        _extra_c_sources = [
+            r'''
+            /*gpufun*/
+            void TestElement_track_local_particle(TestElementData el, LocalParticle* part0){
+                // Extract the record and record_index
+                TestElementRecordData record = TestElementData_getp_internal_record(el, part0);
+                RecordIndex record_index = NULL;
+                if (record){
+                    record_index = TestElementRecordData_getp__index(record);
+                }
+
+                int64_t elem_field = TestElementData_get_element_field(el);
+
+                //start_per_particle_block (part0->part)
+                    if (record) {  // Record exists
+                        // Get a slot in the record (this is thread safe)
+                        int64_t i_slot = RecordIndex_get_slot(record_index);
+
+                        if (i_slot>=0) {  // Slot available
+                            TestElementRecordData_set_record_field(
+                                record,
+                                i_slot,
+                                elem_field
+                            );
+                            TestElementRecordData_set_record_at_element(
+                                record,
+                                i_slot,
+                                LocalParticle_get_at_element(part)
+                            );
+                        }
+                    }
+                //end_per_particle_block
+            }
+            '''
+        ]
+
+    line = xt.Line(elements=[
+        TestElement(element_field=3),
+        TestElement(element_field=4),
+    ])
+    line.build_tracker(_context=test_context)
+
+    record = line.start_internal_logging_for_elements_of_type(
+        TestElement,
+        capacity=1000,
+    )
+
+    num_turns = 100
+
+    part = xp.Particles(_context=test_context, x=[1e-3, 2e-3, 3e-3])
+    line.track(part, num_turns=num_turns)
+
+    num_recorded = record._index.num_recorded
+    num_particles = len(part.x)
+
+    record.move(_context=xo.ContextCpu())
+
+    assert num_recorded == (2 * num_particles * num_turns)
+    assert np.sum(record.record_field == 3) == num_particles * num_turns
+    assert np.sum(record.record_field == 4) == num_particles * num_turns
+    assert np.all(record.record_field[:num_recorded][
+                    record.record_at_element[:num_recorded] == 0] == 3)
+    assert np.all(record.record_field[:num_recorded][
+                    record.record_at_element[:num_recorded] == 1] == 4)
+    assert np.all(record.record_field[num_recorded:] == 0)
+
+    # Now we stop logging and manually reset to mimic the situation where the
+    # record is manually flushed.
+    line.stop_internal_logging_for_elements_of_type(TestElement)
+    line.tracker._init_io_buffer()
+
+    assert line.tracker.io_buffer is not None

{xtrack-0.60.0 → xtrack-0.61.0}/tests/test_twiss.py

@@ -538,6 +538,7 @@ def test_twiss_range(test_context, cycle_to, line_name, check, init_at_edge, col
     if collider_for_test_twiss_range is not None:
         collider = collider_for_test_twiss_range
     else:
+        raise ValueError('This should not happen')
         collider = xt.Multiline.from_json(test_data_folder /
                         'hllhc15_thick/hllhc15_collider_thick.json')
         collider.lhcb1.twiss_default['method'] = '4d'
@@ -1670,4 +1671,63 @@ def test_twiss_strength_reverse_vs_madx(test_context):
     assert_allclose(twm2['k4sl', 'mcdsxf.3r1:1'], tw.lhcb2['k4sl', 'mcdsxf.3r1'], rtol=0, atol=1e-14)
 
     assert_allclose(twm1['k5sl', 'mctsxf.3r1:1'], tw.lhcb1['k5sl', 'mctsxf.3r1'], rtol=0, atol=1e-14)
-    assert_allclose(twm2['k5sl', 'mctsxf.3r1:1'], tw.lhcb2['k5sl', 'mctsxf.3r1'], rtol=0, atol=1e-14)
+    assert_allclose(twm2['k5sl', 'mctsxf.3r1:1'], tw.lhcb2['k5sl', 'mctsxf.3r1'], rtol=0, atol=1e-14)
+
+
+@for_all_test_contexts
+@pytest.mark.parametrize('line_name', ['lhcb1'])
+@pytest.mark.parametrize('reverse', [False, True])
+@pytest.mark.parametrize('section', [
+    (xt.START, xt.END),
+    (xt.START, '_end_point'),
+    (xt.START, 'ip6'),
+    ('ip4', xt.END),
+])
+def test_twiss_range_start_end(test_context, line_name, reverse, section, collider_for_test_twiss_range):
+    collider = collider_for_test_twiss_range
+    init_at = 'ip5'
+
+    collider.vars['on_x5hs'] = 200
+    collider.vars['on_x5vs'] = 123
+    collider.vars['on_sep5h'] = 1
+    collider.vars['on_sep5v'] = 2
+
+    line = collider[line_name]
+
+    if collider.lhcb1.element_names[0] != 'ip3':
+        collider.lhcb1.cycle('ip3', inplace=True)
+    if collider.lhcb2.element_names[0] != 'ip3':
+        collider.lhcb2.cycle('ip3', inplace=True)
+
+    if isinstance(test_context, xo.ContextCpu) and (
+        test_context.omp_num_threads != line._context.omp_num_threads):
+        buffer = test_context.new_buffer()
+    elif isinstance(test_context, line._context.__class__):
+        buffer = line._buffer
+    else:
+        buffer = test_context.new_buffer()
+
+    line.build_tracker(_buffer=buffer)
+
+    tw = line.twiss()
+    tw_init = tw.get_twiss_init(init_at)
+
+    start = section[0]
+    end = section[1]
+    if not reverse:
+        tw_test = line.twiss(start=start, end=end, init=tw_init, reverse=reverse)
+    else:
+        with pytest.raises(ValueError) as excinfo:
+            line.twiss(start=start, end=end, init=tw_init, reverse=reverse)
+        assert 'reverse' in str(excinfo.value)
+        return
+
+    start_el = (line.element_names[0] if start == xt.START else start)
+    end_el = (line.element_names[-1] if (end == xt.END or end == '_end_point') else end)
+    tw_ref = line.twiss(start=start_el, end=end_el, init=tw_init, reverse=reverse)
+
+    for kk in tw_test._data.keys():
+        if kk in ('particle_on_co', '_action'):
+            continue
+        assert np.all(tw_test._data[kk] == tw_ref._data[kk]), (
+            'Issues with variable: ' + kk)

{xtrack-0.60.0 → xtrack-0.61.0}/xtrack/__init__.py

@@ -3,7 +3,7 @@
 # Copyright (c) CERN, 2021.                 #
 # ######################################### #
 
-from .general import _pkg_root, _print
+from .general import _pkg_root, _print, START, END
 
 from .particles import (Particles, PROTON_MASS_EV, ELECTRON_MASS_EV,
                         enable_pyheadtail_interface, disable_pyheadtail_interface)
@@ -15,8 +15,8 @@ from .tracker_data import TrackerData
 from .line import Line, Node, freeze_longitudinal, _temp_knobs, EnergyProgram
 from .tracker import Tracker, Log
 from .match import (Vary, Target, TargetList, VaryList, TargetInequality, Action,
-                    TargetRelPhaseAdvance, TargetSet, START, END,
-                    GreaterThan, LessThan, TargetRmatrixTerm, TargetRmatrix)
+                    TargetRelPhaseAdvance, TargetSet, GreaterThan, LessThan,
+                    TargetRmatrixTerm, TargetRmatrix)
 from .targets import (TargetLuminosity, TargetSeparationOrthogonalToCrossing,
                       TargetSeparation)
 from .twiss import TwissInit, TwissTable

xtrack-0.61.0/xtrack/_version.py

@@ -0,0 +1 @@
+__version__ = '0.61.0'