xtrack 0.89.4__tar.gz → 0.90.1__tar.gz

{xtrack-0.89.4/xtrack.egg-info → xtrack-0.90.1}/PKG-INFO

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

{xtrack-0.89.4 → xtrack-0.90.1}/tests/test_collimation.py

@@ -411,7 +411,6 @@ def test_losslocationrefinement_skip_refinement_for_collimators():
     #################
 
     line.build_tracker()
-    line.config.XTRACK_GLOBAL_XY_LIMIT = 1e3
 
     ##########################
     # Build particles object #
@@ -442,3 +441,140 @@ def test_losslocationrefinement_skip_refinement_for_collimators():
     loss_loc_refinement.refine_loss_location(particles)
 
     assert np.all(part_before.s == particles.s)
+
+def test_memory_dealloc_repeated_refinements():
+
+    def _occupied_size(buffer):
+        return buffer.capacity - buffer.get_free()
+
+    n_part=10000
+    shift_x = 0.3e-2
+    shift_y = 0.5e-2
+    sandwitch_aper = True
+
+    ctx = xo.context_default
+    buf = ctx.new_buffer()
+
+    logger = logging.getLogger('xtrack')
+    logger.setLevel(logging.DEBUG)
+
+    # Define aper_0
+    aper_0 = xt.LimitEllipse(_buffer=buf, a=2e-2, b=2e-2)
+    shift_aper_0 = (shift_x, shift_y)
+    rot_deg_aper_0 = 10.
+
+    # Define aper_1
+    aper_1 = xt.LimitEllipse(_buffer=buf, a=1e-2, b=1e-2)
+    shift_aper_1 = (shift_x, shift_y)
+    rot_deg_aper_1 = 10.
+
+    # aper_0_sandwitch
+    if sandwitch_aper:
+        line_aper_0 = xt.Line(
+            elements=[xt.XYShift(_buffer=buf, dx=shift_aper_0[0], dy=shift_aper_0[1]),
+                    xt.SRotation(_buffer=buf, angle=rot_deg_aper_0),
+                    aper_0,
+                    xt.Multipole(_buffer=buf, knl=[0.00]),
+                    xt.SRotation(_buffer=buf, angle=-rot_deg_aper_0),
+                    xt.XYShift(_buffer=buf, dx=-shift_aper_0[0], dy=-shift_aper_0[1])])
+        line_aper_0.build_tracker(_buffer=buf)
+        # aper_1_sandwitch
+        line_aper_1 = xt.Line(
+            elements=[xt.XYShift(_buffer=buf, dx=shift_aper_1[0], dy=shift_aper_1[1]),
+                    xt.SRotation(_buffer=buf, angle=rot_deg_aper_1),
+                    aper_1,
+                    xt.Multipole(_buffer=buf, knl=[0.00]),
+                    xt.SRotation(_buffer=buf, angle=-rot_deg_aper_1),
+                    xt.XYShift(_buffer=buf, dx=-shift_aper_1[0], dy=-shift_aper_1[1])])
+    else:
+        aper_0.shift_x = shift_aper_0[0]
+        aper_0.shift_y = shift_aper_0[1]
+        aper_0.rot_s_rad = np.deg2rad(rot_deg_aper_0)
+        line_aper_0 = xt.Line(
+            elements=[aper_0, xt.Multipole(_buffer=buf, knl=[0.0])])
+        line_aper_0.build_tracker(_buffer=buf)
+        aper_1.shift_x = shift_aper_1[0]
+        aper_1.shift_y = shift_aper_1[1]
+        aper_1.rot_s_rad = np.deg2rad(rot_deg_aper_1)
+        line_aper_1 = xt.Line(
+            elements=[aper_1, xt.Multipole(_buffer=buf, knl=[0.00])])
+
+    line_aper_1.build_tracker(_buffer=buf)
+
+    # Build example line
+    line=xt.Line(
+        elements = ((xt.Drift(_buffer=buf, length=0.5),)
+                    + line_aper_0.elements
+                    + (xt.Drift(_buffer=buf, length=1),
+                        xt.Multipole(_buffer=buf, knl=[0.]),
+                        xt.Quadrupole(_buffer=buf, length=1),
+                        xt.Cavity(_buffer=buf, voltage=3e6, frequency=400e6),
+                        xt.ParticlesMonitor(_buffer=buf,
+                            start_at_turn=0, stop_at_turn=10, num_particles=3),
+                        xt.Drift(_buffer=buf, length=1.),
+                        xt.Marker())
+                    + line_aper_1.elements))
+    line.build_tracker(_buffer=buf)
+
+    # Test on full line
+    r = np.linspace(0, 0.018, n_part)
+    theta = np.linspace(0, 8*np.pi, n_part)
+    particles0 = xp.Particles(_context=ctx,
+            p0c=6500e9,
+            x=r*np.cos(theta)+shift_x,
+            y=r*np.sin(theta)+shift_y)
+
+
+
+    print('occupied size before refinement', _occupied_size(buf))
+
+    loss_loc_refinement = xt.LossLocationRefinement(line,
+                            n_theta = 360,
+                            r_max = 0.5, # m
+                            dr = 50e-6,
+                            ds = 0.1,
+                            save_refine_lines=False,
+                            allowed_backtrack_types=[
+                                xt.Multipole,
+                                xt.Cavity
+                                ])
+    print('occupied size after init refinement', _occupied_size(buf))
+
+    n_repetitions = 20
+
+    import time
+    for i_iter in range(n_repetitions):
+
+        particles = particles0.copy()
+
+        line.track(particles)
+
+        t0 = time.time()
+
+        loss_loc_refinement.refine_loss_location(particles)
+        print('occupied size after refinement', _occupied_size(buf))
+
+        if i_iter==0:
+            occupied_size_first_iter = _occupied_size(buf)
+        else:
+            assert _occupied_size(buf) == occupied_size_first_iter, \
+                'Buffer size increased after refinement! Memory leak?'
+
+        t1 = time.time()
+        print(f'Took\t{(t1-t0)*1e3:.2f} ms')
+
+    # Automatic checks
+    mask_lost = particles.state == 0
+    r_calc = np.sqrt((particles.x-shift_x)**2 + (particles.y-shift_y)**2)
+    assert np.all(r_calc[~mask_lost]<1e-2)
+    assert np.all(r_calc[mask_lost]>1e-2)
+    i_aper_1 = line.elements.index(aper_1)
+    assert np.all(particles.at_element[mask_lost]==i_aper_1)
+    assert np.all(particles.at_element[~mask_lost]==0)
+    s0 = line.get_s_elements()[line.elements.index(aper_0)]
+    s1 = line.get_s_elements()[line.elements.index(aper_1)]
+    r0 = np.sqrt(aper_0.a_squ)
+    r1 = np.sqrt(aper_1.a_squ)
+    s_expected = s0 + (r_calc-r0)/(r1 - r0)*(s1 - s0)
+    # TODO This threshold is a bit large
+    xo.assert_allclose(particles.s[mask_lost], s_expected[mask_lost], atol=0.11)

{xtrack-0.89.4 → xtrack-0.90.1}/tests/test_elements.py

@@ -214,6 +214,40 @@ def test_drift(test_context):
                       dtk_particle.zeta,
                       rtol=1e-14, atol=1e-14)
 
+@for_all_test_contexts
+def test_drift_exact_and_expanded(test_context):
+
+    line = xt.Line(elements=[xt.Drift(length=1.), xt.Drift(length=2.), xt.Drift(length=3.)])
+    ltot = line.get_length()
+    line.build_tracker(_context=test_context)
+
+    assert line['e2'].model == 'adaptive'
+
+    p0 = xp.Particles(p0c=1e9, px=0.3, _context=test_context)
+    x_prime_expanded = p0.px / (1 + p0.delta)
+    x_prime_exact = p0.px / np.sqrt((1 + p0.delta)**2 - p0.px**2)
+
+    p = p0.copy(_context=test_context)
+    line.track(p)
+    xo.assert_allclose(p.x, x_prime_expanded*ltot, rtol=1e-14, atol=1e-14)
+
+    line.configure_drift_model(model='exact')
+    assert line['e2'].model == 'exact'
+    p = p0.copy(_context=test_context)
+    line.track(p)
+    xo.assert_allclose(p.x, x_prime_exact*ltot, rtol=1e-14, atol=1e-14)
+
+    line.configure_drift_model(model='expanded')
+    assert line['e2'].model == 'expanded'
+    p = p0.copy(_context=test_context)
+    line.track(p)
+    xo.assert_allclose(p.x, x_prime_expanded*ltot, rtol=1e-14, atol=1e-14)
+
+    line.config.XTRACK_USE_EXACT_DRIFTS = True
+    p = p0.copy(_context=test_context)
+    line.track(p)
+    xo.assert_allclose(p.x, x_prime_exact*ltot, rtol=1e-14, atol=1e-14)
+
 
 @for_all_test_contexts
 def test_drift_exact(test_context):
@@ -232,7 +266,7 @@ def test_drift_exact(test_context):
     drift = xt.Drift(_context=test_context, length=10.)
     line = xt.Line(elements=[drift])
     line.build_tracker(compile=False, _context=test_context)
-    line.config.XTRACK_USE_EXACT_DRIFTS = True
+    line.configure_drift_model(model='exact')
     line.track(particles)
 
     dtk_drift = dtk.elements.DriftExact(length=10.)
@@ -1162,7 +1196,7 @@ def test_nonlinearlens(test_context):
     """)
 
     line = xt.Line.from_madx_sequence(mad.sequence.ss)
-    line.config.XTRACK_USE_EXACT_DRIFTS = True # to be consistent with madx
+    line.configure_drift_model('exact') # to be consistent with madx
     line.build_tracker(_context=test_context)
 
     num_p_test = 10

{xtrack-0.89.4 → xtrack-0.90.1}/tests/test_elements_thick.py

@@ -59,7 +59,7 @@ def test_combined_function_dipole_against_ptc(test_context, k0, k1, k2, length,
 
     ml = MadLoader(mad.sequence.ss, allow_thick=True)
     line_thick = ml.make_line()
-    line_thick.config.XTRACK_USE_EXACT_DRIFTS = True # to be consistent with mad
+    line_thick.configure_drift_model('exact') # to be consistent with madx
     line_thick.build_tracker(_context=test_context)
     line_thick.configure_bend_model(core=model, edge='dipole-only')
 
@@ -1283,7 +1283,7 @@ def test_solenoid_against_madx(test_context, ks, ksi, length):
     ml = MadLoader(mad.sequence.ss, allow_thick=True)
     line_thick = ml.make_line()
     line_thick.build_tracker(_context=test_context)
-    line_thick.config.XTRACK_USE_EXACT_DRIFTS = True  # to be consistent with madx
+    line_thick.configure_drift_model('exact')  # to be consistent with madx
 
     for ii in range(len(p0.x)):
         mad.input(f"""
@@ -1320,7 +1320,7 @@ def test_solenoid_against_madx(test_context, ks, ksi, length):
 def test_solenoid_thick_drift_like(test_context):
     solenoid = xt.UniformSolenoid(ks=1.001e-9, length=1, _context=test_context)
     l_drift = xt.Line(elements=[xt.Drift(length=1)])
-    l_drift.config.XTRACK_USE_EXACT_DRIFTS = True
+    l_drift.configure_drift_model('exact')
     l_drift.build_tracker(_context=test_context)
 
     p0 = xp.Particles(
@@ -1743,7 +1743,6 @@ def test_solenoid_multipole_rotations():
 def test_drift_like_solenoid_with_kicks_radiation(radiation_mode, config):
     test_context = xo.ContextCpu()
 
-    config['XTRACK_USE_EXACT_DRIFTS'] = True
     knl = [0.1, 0.4, 0.5]
     ksl = [0.2, 0.3, 0.6]
 
@@ -1752,6 +1751,7 @@ def test_drift_like_solenoid_with_kicks_radiation(radiation_mode, config):
         xt.Multipole(knl=knl, ksl=ksl),
         xt.Drift(length=0.5),
     ])
+    line_test.configure_drift_model('exact')
 
     line_ref = xt.Line(elements=[
         xt.UniformSolenoid(ks=0, length=1, knl=knl, ksl=ksl, num_multipole_kicks=1)
@@ -1803,8 +1803,6 @@ def test_drift_like_solenoid_with_kicks_radiation(radiation_mode, config):
 def test_solenoid_with_kicks_radiation(radiation_mode, config):
     test_context = xo.ContextCpu()
 
-    config['XTRACK_USE_EXACT_DRIFTS'] = True
-
     ks = 0.4
     l = 1.1
     knl = [0.1, 0.4, 0.5]
@@ -1818,6 +1816,9 @@ def test_solenoid_with_kicks_radiation(radiation_mode, config):
     line_1 = xt.Line(elements=[sol_1])
     line_3 = xt.Line(elements=[sol_3])
 
+    for ll in (line_ref, line_1, line_3):
+        ll.configure_drift_model('exact')
+
     coords = np.linspace(-0.05, 0.05, 10)
     coords_6d = np.array(list(itertools.product(*(coords,) * 6))).T
 

{xtrack-0.89.4 → xtrack-0.90.1}/tests/test_environment.py

@@ -620,19 +620,19 @@ def test_assemble_ring():
     assert 'cell3_copy' in env.lines
     assert cell3_select.particle_ref is not None
     assert env.lines['cell3_copy'] is cell3_select
-    assert cell3_select._element_dict is env.element_dict
+    assert cell3_select.element_dict is env.element_dict
     assert cell3_select.element_names[0] == 'start.cell.3.arc.2'
     assert cell3_select.element_names[-1] == 'end.cell.3.arc.2'
     assert (np.array(cell3_select.element_names) == np.array(
         tw.rows['start.cell.3.arc.2':'end.cell.3.arc.2'].name)).all()
 
-    # Check that they share the _element_dict
-    assert cell._element_dict is env.element_dict
-    assert halfcell._element_dict is env.element_dict
-    assert halfcell_ss._element_dict is env.element_dict
-    assert cell_ss._element_dict is env.element_dict
-    assert insertion._element_dict is env.element_dict
-    assert ring2._element_dict is env.element_dict
+    # Check that they share the element_dict
+    assert cell.element_dict is env.element_dict
+    assert halfcell.element_dict is env.element_dict
+    assert halfcell_ss.element_dict is env.element_dict
+    assert cell_ss.element_dict is env.element_dict
+    assert insertion.element_dict is env.element_dict
+    assert ring2.element_dict is env.element_dict
 
     cell3_select.twiss4d()
 
@@ -1039,19 +1039,19 @@ def test_assemble_ring_builders():
                                 name='cell3_copy')
     assert 'cell3_copy' in env.lines
     assert env.lines['cell3_copy'] is cell3_select
-    assert cell3_select._element_dict is env.element_dict
+    assert cell3_select.element_dict is env.element_dict
     assert cell3_select.element_names[0] == 'start.cell.3.arc.2'
     assert cell3_select.element_names[-1] == 'end.cell.3.arc.2'
     assert (np.array(cell3_select.element_names) == np.array(
         tw.rows['start.cell.3.arc.2':'end.cell.3.arc.2'].name)).all()
 
-    # Check that they share the _element_dict
-    assert cell._element_dict is env.element_dict
-    assert halfcell._element_dict is env.element_dict
-    assert halfcell_ss._element_dict is env.element_dict
-    assert cell_ss._element_dict is env.element_dict
-    assert insertion._element_dict is env.element_dict
-    assert ring2._element_dict is env.element_dict
+    # Check that they share the element_dict
+    assert cell.element_dict is env.element_dict
+    assert halfcell.element_dict is env.element_dict
+    assert halfcell_ss.element_dict is env.element_dict
+    assert cell_ss.element_dict is env.element_dict
+    assert insertion.element_dict is env.element_dict
+    assert ring2.element_dict is env.element_dict
 
     cell3_select.twiss4d()
 
@@ -1393,13 +1393,13 @@ def test_assemble_ring_repeated_elements():
     xo.assert_allclose(tw_ring2.betx[0], tw_cell.betx[0], atol=0, rtol=5e-4)
     xo.assert_allclose(tw_ring2.bety[0], tw_cell.bety[0], atol=0, rtol=5e-4)
 
-    # Check that they share the _element_dict
-    assert cell._element_dict is env.element_dict
-    assert halfcell._element_dict is env.element_dict
-    assert halfcell_ss._element_dict is env.element_dict
-    assert cell_ss._element_dict is env.element_dict
-    assert insertion._element_dict is env.element_dict
-    assert ring2._element_dict is env.element_dict
+    # Check that they share the element_dict
+    assert cell.element_dict is env.element_dict
+    assert halfcell.element_dict is env.element_dict
+    assert halfcell_ss.element_dict is env.element_dict
+    assert cell_ss.element_dict is env.element_dict
+    assert insertion.element_dict is env.element_dict
+    assert ring2.element_dict is env.element_dict
 
     xo.assert_allclose(tw_ring2['betx', 'ip::0'], tw_half_insertion['betx', 'ip'], atol=0, rtol=5e-4)
     xo.assert_allclose(tw_ring2['bety', 'ip::0'], tw_half_insertion['bety', 'ip'], atol=0, rtol=5e-4)
@@ -3193,3 +3193,236 @@ def test_enviroment_from_two_lines():
     assert np.allclose(tw2.s, tw2i.s, atol=0, rtol=1e-15)
     assert np.allclose(tw2.betx, tw2i.betx, atol=0, rtol=1e-15)
     assert np.allclose(tw2.bety, tw2i.bety, atol=0, rtol=1e-15)
+
+def test_particle_ref_from_particles_container():
+
+    env = xt.Environment()
+    env['a'] = 4.
+
+    env.new_particle('my_particle', p0c=['1e12 * a'])
+    assert 'my_particle' in env.particles
+    xo.assert_allclose(env['my_particle'].p0c, 4e12, rtol=0, atol=1e-9)
+    env['a'] = 5.
+    xo.assert_allclose(env['my_particle'].p0c, 5e12, rtol=0, atol=1e-9)
+
+    env.particle_ref = 'my_particle'
+
+    xo.assert_allclose(env.particle_ref.p0c, 5e12, rtol=0, atol=1e-9)
+    assert env.particle_ref.__class__.__name__ == 'EnvParticleRef'
+    assert env._particle_ref == 'my_particle'
+    assert env.ref['my_particle']._value is env.get('my_particle')
+    env.particle_ref.p0c = '2e12 * a'
+    xo.assert_allclose(env.particle_ref.p0c, 10e12, rtol=0, atol=1e-9)
+    env['my_particle'].p0c = '1e12 * a'
+    xo.assert_allclose(env.particle_ref.p0c, 5e12, rtol=0, atol=1e-9)
+
+    env2 = xt.Environment.from_dict(env.to_dict())
+    assert 'my_particle' in env2.particles
+    assert isinstance(env2.get('my_particle'), xt.Particles)
+    assert env2.get('my_particle') is not env.get('my_particle')
+    assert env2._particle_ref == "my_particle"
+    assert env2.ref['my_particle']._value is env2.get('my_particle')
+    xo.assert_allclose(env2['my_particle'].p0c, 5e12, rtol=0, atol=1e-9)
+    env2['a'] = 6.
+    xo.assert_allclose(env2['my_particle'].p0c, 6e12, rtol=0, atol=1e-9)
+    env2['a'] = 5.
+
+    assert env2.particle_ref.__class__.__name__ == 'EnvParticleRef'
+    env2.particle_ref.p0c = '2e12 * a'
+    xo.assert_allclose(env2.particle_ref.p0c, 10e12, rtol=0, atol=1e-9)
+    env2['my_particle'].p0c = '1e12 * a'
+    xo.assert_allclose(env2.particle_ref.p0c, 5e12, rtol=0, atol=1e-9)
+
+    env2 = env.copy()
+    assert 'my_particle' in env2.particles
+    assert env2._particle_ref == "my_particle"
+    assert env.ref['my_particle']._value is env.get('my_particle')
+    assert isinstance(env2.get('my_particle'), xt.Particles)
+    assert env2.get('my_particle') is not env.get('my_particle')
+    xo.assert_allclose(env2['my_particle'].p0c, 5e12, rtol=0, atol=1e-9)
+    env2['a'] = 6.
+    xo.assert_allclose(env2['my_particle'].p0c, 6e12, rtol=0, atol=1e-9)
+    env2['a'] = 5.
+
+    assert env2.particle_ref.__class__.__name__ == 'EnvParticleRef'
+    env2.particle_ref.p0c = '2e12 * a'
+    xo.assert_allclose(env2.particle_ref.p0c, 10e12, rtol=0, atol=1e-9)
+    env2['my_particle'].p0c = '1e12 * a'
+    xo.assert_allclose(env2.particle_ref.p0c, 5e12, rtol=0, atol=1e-9)
+
+    ll = env.new_line(name='my_line', components=[])
+    assert ll._particle_ref == 'my_particle'
+
+    xo.assert_allclose(ll.particle_ref.p0c, 5e12, rtol=0, atol=1e-9)
+    assert ll.particle_ref.__class__.__name__ == 'LineParticleRef'
+    ll.particle_ref.p0c = '2e12 * a'
+    xo.assert_allclose(env.particle_ref.p0c, 10e12, rtol=0, atol=1e-9)
+    env['my_particle'].p0c = '1e12 * a'
+    xo.assert_allclose(ll.particle_ref.p0c, 5e12, rtol=0, atol=1e-9)
+
+    ll2 = xt.Line.from_dict(ll.to_dict())
+    assert 'my_particle' in ll2.env.particles
+    assert ll2.env.particle_ref is None
+    assert ll2.particle_ref.__class__.__name__ == 'LineParticleRef'
+    assert ll2._particle_ref == 'my_particle'
+    xo.assert_allclose(ll2.env['my_particle'].p0c, 5e12, rtol=0, atol=1e-9)
+    ll2['a'] = 7.
+    xo.assert_allclose(ll2.env['my_particle'].p0c, 7e12, rtol=0, atol=1e-9)
+    ll2['a'] = 5.
+
+    ll2.particle_ref.p0c = '2e12 * a'
+    xo.assert_allclose(ll2.particle_ref.p0c, 10e12, rtol=0, atol=1e-9)
+    ll2.env['my_particle'].p0c = '1e12 * a'
+    xo.assert_allclose(ll2.particle_ref.p0c, 5e12, rtol=0, atol=1e-9)
+
+
+    ll2 = ll.copy()
+    assert 'my_particle' in ll2.env.particles
+    assert ll2.env.particle_ref is None
+    assert ll2.particle_ref.__class__.__name__ == 'LineParticleRef'
+    assert ll2._particle_ref == 'my_particle'
+    xo.assert_allclose(ll2.env['my_particle'].p0c, 5e12, rtol=0, atol=1e-9)
+    ll2['a'] = 7.
+    xo.assert_allclose(ll2.env['my_particle'].p0c, 7e12, rtol=0, atol=1e-9)
+    ll2['a'] = 5.
+
+    ll2.particle_ref.p0c = '2e12 * a'
+    xo.assert_allclose(ll2.particle_ref.p0c, 10e12, rtol=0, atol=1e-9)
+    ll2.env['my_particle'].p0c = '1e12 * a'
+    xo.assert_allclose(ll2.particle_ref.p0c, 5e12, rtol=0, atol=1e-9)
+
+def test_particle_ref_as_object():
+
+    env = xt.Environment()
+    env['a'] = 4.
+
+    env.new_particle('my_particle', p0c=['1e12 * a'])
+    assert 'my_particle' in env.particles
+    xo.assert_allclose(env['my_particle'].p0c, 4e12, rtol=0, atol=1e-9)
+    env['a'] = 5.
+    xo.assert_allclose(env['my_particle'].p0c, 5e12, rtol=0, atol=1e-9)
+
+    part = env['my_particle'].copy()
+    env.particle_ref = part
+
+    xo.assert_allclose(env.particle_ref.p0c, 5e12, rtol=0, atol=1e-9)
+    assert env.particle_ref.__class__.__name__ == 'EnvParticleRef'
+    assert env._particle_ref is part
+    env['my_particle'].p0c = '2e12 * a'
+    xo.assert_allclose(env.eval('2e12 * a'), 10e12, rtol=0, atol=1e-9)
+    xo.assert_allclose(env.particle_ref.p0c, 5e12, rtol=0, atol=1e-9)
+
+    env2 = xt.Environment.from_dict(env.to_dict())
+    assert 'my_particle' in env2.particles
+    assert isinstance(env2.get('my_particle'), xt.Particles)
+    assert env2.get('my_particle') is not env.get('my_particle')
+    assert isinstance(env2._particle_ref, xt.Particles)
+    assert env2.particle_ref.__class__.__name__ == 'EnvParticleRef'
+
+    env2 = env.copy()
+    assert 'my_particle' in env2.particles
+    assert isinstance(env2._particle_ref, xt.Particles)
+    assert isinstance(env2.get('my_particle'), xt.Particles)
+    assert env2.get('my_particle') is not env.get('my_particle')
+    assert env2.particle_ref.__class__.__name__ == 'EnvParticleRef'
+
+    ll = env.new_line(name='my_line', components=[])
+    assert isinstance(ll._particle_ref, xt.Particles)
+    assert ll._particle_ref is env._particle_ref
+    xo.assert_allclose(ll.particle_ref.p0c, 5e12, rtol=0, atol=1e-9)
+    assert ll.particle_ref.__class__.__name__ == 'LineParticleRef'
+
+    ll2 = xt.Line.from_dict(ll.to_dict())
+    assert ll2.env.particle_ref is None
+    assert ll2.particle_ref.__class__.__name__ == 'LineParticleRef'
+    assert isinstance(ll2._particle_ref, xt.Particles)
+    assert ll2._particle_ref is not ll._particle_ref
+
+
+    ll2 = ll.copy()
+    assert ll2.env.particle_ref is None
+    assert ll2.particle_ref.__class__.__name__ == 'LineParticleRef'
+    assert isinstance(ll2._particle_ref, xt.Particles)
+    assert ll2._particle_ref is not ll._particle_ref
+
+def test_line_set_particle_ref():
+
+    line = xt.Line()
+    line.set_particle_ref('electron', beta0=0.9)
+
+    xo.assert_allclose(line.particle_ref.q0 , -1, rtol=0, atol=1e-14)
+    xo.assert_allclose(line.particle_ref.mass0 , xt.ELECTRON_MASS_EV)
+    xo.assert_allclose(line.particle_ref.beta0 , 0.9, rtol=0, atol=1e-14)
+
+    env = xt.Environment()
+    env['my_beta0'] = 0.1
+    env.new_particle('my_part', 'proton', beta0='my_beta0')
+
+    line = env.new_line()
+    line.set_particle_ref('my_part')
+    line['my_beta0'] = 0.6
+
+    xo.assert_allclose(line.particle_ref.q0 , 1, rtol=0, atol=1e-14)
+    xo.assert_allclose(line.particle_ref.mass0 , xt.PROTON_MASS_EV)
+    xo.assert_allclose(line.particle_ref.beta0 , 0.6, rtol=0, atol=1e-14)
+
+    p_ref = xt.Particles('Pb208', p0c=7e12/82)
+    line = xt.Line()
+    line.set_particle_ref(p_ref)
+
+    xo.assert_allclose(line.particle_ref.q0 , 82, rtol=0, atol=1e-14)
+    xo.assert_allclose(line.particle_ref.mass0, xt.particles.masses.Pb208_MASS_EV)
+    xo.assert_allclose(line.particle_ref.p0c , 7e12/82, rtol=0, atol=1e-14)
+
+
+def test_env_set_particle_ref():
+
+    env = xt.Environment()
+    env.set_particle_ref('electron', beta0=0.9)
+
+    xo.assert_allclose(env.particle_ref.q0 , -1, rtol=0, atol=1e-14)
+    xo.assert_allclose(env.particle_ref.mass0 , xt.ELECTRON_MASS_EV)
+    xo.assert_allclose(env.particle_ref.beta0 , 0.9, rtol=0, atol=1e-14)
+
+    env.new_line(name='line1', components=[env.new('d1', 'Drift', length=1)])
+    env.line1.set_particle_ref('electron', beta0=0.9)
+
+    xo.assert_allclose(env.line1.particle_ref.q0 , -1, rtol=0, atol=1e-14)
+    xo.assert_allclose(env.line1.particle_ref.mass0 , xt.ELECTRON_MASS_EV)
+    xo.assert_allclose(env.line1.particle_ref.beta0 , 0.9, rtol=0, atol=1e-14)
+
+    env.set_particle_ref('proton', beta0=0.8)
+    xo.assert_allclose(env.particle_ref.q0 , 1, rtol=0, atol=1e-14)
+    xo.assert_allclose(env.particle_ref.mass0 , xt.PROTON_MASS_EV)
+    xo.assert_allclose(env.particle_ref.beta0 , 0.8, rtol=0, atol=1e-14)
+    xo.assert_allclose(env.line1.particle_ref.q0 , 1, rtol=0, atol=1e-14)
+    xo.assert_allclose(env.line1.particle_ref.mass0 , xt.PROTON_MASS_EV)
+    xo.assert_allclose(env.line1.particle_ref.beta0 , 0.8, rtol=0, atol=1e-14)
+
+    env['my_beta0'] = 0.1
+    env.new_particle('my_part', 'proton', beta0='my_beta0')
+    env.set_particle_ref('my_part')
+
+    env['my_beta0'] = 0.6
+    xo.assert_allclose(env.line1.particle_ref.q0 , 1, rtol=0, atol=1e-14)
+    xo.assert_allclose(env.line1.particle_ref.mass0 , xt.PROTON_MASS_EV)
+    xo.assert_allclose(env.line1.particle_ref.beta0 , 0.6, rtol=0, atol=1e-14)
+
+    p_ref = xt.Particles('Pb208', p0c=7e12/82)
+    env.set_particle_ref(p_ref)
+
+    xo.assert_allclose(env.particle_ref.q0 , 82, rtol=0, atol=1e-14)
+    xo.assert_allclose(env.particle_ref.mass0, xt.particles.masses.Pb208_MASS_EV)
+    xo.assert_allclose(env.particle_ref.p0c , 7e12/82, rtol=0, atol=1e-14)
+    xo.assert_allclose(env.line1.particle_ref.q0 , 82, rtol=0, atol=1e-14)
+    xo.assert_allclose(env.line1.particle_ref.mass0, xt.particles.masses.Pb208_MASS_EV)
+    xo.assert_allclose(env.line1.particle_ref.p0c , 7e12/82, rtol=0, atol=1e-14)
+
+    env.set_particle_ref('positron', beta0=0.7, lines=None)
+    xo.assert_allclose(env.particle_ref.q0 , 1, rtol=0, atol=1e-14)
+    xo.assert_allclose(env.particle_ref.mass0 , xt.ELECTRON_MASS_EV)
+    xo.assert_allclose(env.particle_ref.beta0 , 0.7, rtol=0, atol=1e-14)
+    # line unchanged
+    xo.assert_allclose(env.line1.particle_ref.q0 , 82, rtol=0, atol=1e-14)
+    xo.assert_allclose(env.line1.particle_ref.mass0, xt.particles.masses.Pb208_MASS_EV)
+    xo.assert_allclose(env.line1.particle_ref.p0c , 7e12/82, rtol=0, atol=1e-14)

{xtrack-0.89.4 → xtrack-0.90.1}/tests/test_line.py

@@ -65,7 +65,7 @@ def test_simplification_methods():
     line.remove_zero_length_drifts(inplace=True)
 
     # Test merging of multipoles
-    line._var_management = None
+    line.env._var_management = None
     line.insert_element(element=xt.Multipole(knl=[1, 0, 3], ksl=[0, 20, 0]), name='m1', at_s=3.3)
     line.insert_element(element=xt.Multipole(knl=[4, 2], ksl=[10, 40]), name='m2', at_s=3.3)
     line.insert_element(element=xt.Multipole(knl=[0, 3, 8], ksl=[2, 0, 17]), name='m3', at_s=3.3)