spacecharge: added semantics in tools and removal of duplicate...

...default options in loader_mad.py .

pysixtrack/be_beamfields/spacecharge.py

@@ -15,7 +15,7 @@ class SpaceChargeCoasting(Element):
         ("y_co", "m", "Vertical closed orbit offset", 0.0),
     ]
     _extra = [
-        ("min_sigma_diff", "m", "Threshold to detect round beam", 1e-10),
+        ("min_sigma_diff", "m", "Threshold to detect round beam", 1e-8),
         ("enabled", "", "Switch to disable space charge effect", True),
     ]
 
@@ -79,7 +79,7 @@ class SpaceChargeQGaussianProfile(Element):
         ("y_co", "m", "Vertical closed orbit offset", 0.0),
     ]
     _extra = [
-        ("min_sigma_diff", "m", "Threshold to detect round beam", 1e-10),
+        ("min_sigma_diff", "m", "Threshold to detect round beam", 1e-8),
         ("enabled", "", "Switch to disable space charge effect", True),
         ("q_parameter", "", "q parameter of generalised Gaussian distribution (q=1 for standard Gaussian)", 1.0),
         ("b_parameter", "", "b parameter of generalised Gaussian distribution (b=1 for standard Gaussian)", 1.0),
@@ -158,7 +158,7 @@ class SpaceChargeInterpolatedProfile(Element):
         ("y_co", "m", "Vertical closed orbit offset", 0.0),
     ]
     _extra = [
-        ("min_sigma_diff", "m", "Threshold to detect round beam", 1e-10),
+        ("min_sigma_diff", "m", "Threshold to detect round beam", 1e-8),
         ("enabled", "", "Switch to disable space charge effect", True),
     ]
 

pysixtrack/be_beamfields/tools.py

@@ -11,8 +11,7 @@ def norm(v):
 
 
 def get_points_twissdata_for_elements(
-    ele_names, mad, seq_name, use_survey=True, use_twiss=True
-):
+        ele_names, mad, seq_name, use_survey=True, use_twiss=True):
 
     mad.use(sequence=seq_name)
 
@@ -77,8 +76,8 @@ def get_elements(seq, ele_type=None, slot_id=None):
 
 
 def get_points_twissdata_for_element_type(
-    mad, seq_name, ele_type=None, slot_id=None, use_survey=True, use_twiss=True
-):
+        mad, seq_name, ele_type=None, slot_id=None, 
+        use_survey=True, use_twiss=True):
 
     elements, element_names = get_elements(
         seq=mad.sequence[seq_name], ele_type=ele_type, slot_id=slot_id
@@ -116,8 +115,7 @@ def find_alpha_and_phi(dpx, dpy):
 
 
 def get_bb_names_madpoints_sigmas(
-    mad, seq_name, use_survey=True, use_twiss=True
-):
+        mad, seq_name, use_survey=True, use_twiss=True):
     (
         _,
         element_names,
@@ -136,8 +134,7 @@ def get_bb_names_madpoints_sigmas(
 
 
 def shift_strong_beam_based_on_close_ip(
-    points_weak, points_strong, IPs_survey_weak, IPs_survey_strong
-):
+        points_weak, points_strong, IPs_survey_weak, IPs_survey_strong):
 
     for i_bb, _ in enumerate(points_weak):
 
@@ -206,16 +203,15 @@ def find_bb_separations(points_weak, points_strong, names=None):
 
 
 def setup_beam_beam_in_line(
-    line,
-    bb_names,
-    bb_sigmas_strong,
-    bb_points_weak,
-    bb_points_strong,
-    beta_r_strong,
-    bunch_intensity_strong,
-    n_slices_6D,
-    bb_coupling,
-):
+        line,
+        bb_names,
+        bb_sigmas_strong,
+        bb_points_weak,
+        bb_points_strong,
+        beta_r_strong,
+        bunch_intensity_strong,
+        n_slices_6D,
+        bb_coupling):
 
     sep_x, sep_y = find_bb_separations(
         points_weak=bb_points_weak,
@@ -326,32 +322,18 @@ def get_spacecharge_names_twdata(mad, seq_name, mode):
     return mad_sc_names, twdata
 
 
-def setup_spacecharge_in_line(
-    sc_elements,
-    sc_lengths,
-    sc_twdata,
-    betagamma,
-    number_of_particles,
-    zeta_length,
-    delta_rms,
-    neps_x,
-    neps_y,
-):
+def _setup_spacecharge_in_line(
+        sc_elements,
+        sc_lengths,
+        sc_twdata,
+        betagamma,
+        number_of_particles,
+        delta_rms,
+        neps_x,
+        neps_y):
 
     for ii, ss in enumerate(sc_elements):
-
         ss.number_of_particles = number_of_particles
-        if isinstance(ss, pysixtrack.elements.SpaceChargeQGaussianProfile):
-            ss.bunchlength_rms = zeta_length
-        elif isinstance(ss, pysixtrack.elements.SpaceChargeCoasting):
-            ss.circumference = zeta_length
-        elif isinstance(ss, pysixtrack.elements.SpaceChargeInterpolatedProfile):
-            # InterpolatedProfile
-            if ii == 0:
-                print ('Warning: please add the profile data to '
-                       'SpaceChargeInterpolatedProfile nodes manually!')
-        else:
-            raise NotImplementedError('Unknown space charge node type.')
         ss.sigma_x = np.sqrt(
             sc_twdata["betx"][ii] * neps_x / betagamma
             + (sc_twdata["dispersion_x"][ii] * delta_rms) ** 2
@@ -365,10 +347,64 @@ def setup_spacecharge_in_line(
         ss.y_co = sc_twdata["y"][ii]
         ss.enabled = True
 
+def setup_spacecharge_bunched_in_line(
+        sc_elements,
+        sc_lengths,
+        sc_twdata,
+        betagamma,
+        number_of_particles,
+        delta_rms,
+        neps_x,
+        neps_y,
+        bunchlength_rms):
+
+    for ii, ss in enumerate(sc_elements):
+        ss.bunchlength_rms = bunchlength_rms
+    _setup_spacecharge_in_line(
+        sc_elements, sc_lengths, sc_twdata, betagamma, 
+        number_of_particles, delta_rms, neps_x, neps_y)
+
+def setup_spacecharge_coasting_in_line(
+        sc_elements,
+        sc_lengths,
+        sc_twdata,
+        betagamma,
+        number_of_particles,
+        delta_rms,
+        neps_x,
+        neps_y,
+        circumference):
+
+    for ii, ss in enumerate(sc_elements):
+        ss.circumference = circumference
+    _setup_spacecharge_in_line(
+        sc_elements, sc_lengths, sc_twdata, betagamma, 
+        number_of_particles, delta_rms, neps_x, neps_y)
+
+def setup_spacecharge_interpolated_in_line(
+        sc_elements,
+        sc_lengths,
+        sc_twdata,
+        betagamma,
+        number_of_particles,
+        delta_rms,
+        neps_x,
+        neps_y,
+        line_density_profile,
+        dz,
+        z0):
+
+    for ii, ss in enumerate(sc_elements):
+        ss.line_density_profile = line_density_profile
+        ss.dz = dz
+        ss.z0 = z0
+    _setup_spacecharge_in_line(
+        sc_elements, sc_lengths, sc_twdata, betagamma,
+        number_of_particles, delta_rms, neps_x, neps_y)
+
 
 def check_spacecharge_consistency(
-    sc_elements, sc_names, sc_lengths, mad_sc_names
-):
+        sc_elements, sc_names, sc_lengths, mad_sc_names):
     assert len(sc_elements) == len(mad_sc_names)
     assert len(sc_lengths) == len(mad_sc_names)
     for ii, (ss, nn) in enumerate(zip(sc_elements, sc_names)):

pysixtrack/loader_mad.py

@@ -158,40 +158,14 @@ def iter_from_madx_sequence(
                 )
         elif mad_etype == "placeholder":
             if ee.slot_id == 1:
-                newele = classes.SpaceChargeCoasting(
-                    number_of_particles=0.0,
-                    circumference=0.0,
-                    sigma_x=1.0,
-                    sigma_y=1.0,
-                    length=0.0,
-                    x_co=0.0,
-                    y_co=0.0,
-                )
+                newele = classes.SpaceChargeCoasting()
             elif ee.slot_id == 2:
-                newele = classes.SpaceChargeQGaussianProfile(
-                    number_of_particles=0.0,
-                    bunchlength_rms=0.0,
-                    sigma_x=1.0,
-                    sigma_y=1.0,
-                    length=0.0,
-                    x_co=0.0,
-                    y_co=0.0,
-                )
+                newele = classes.SpaceChargeQGaussianProfile()
             elif ee.slot_id == 3:
-                newele = classes.SpaceChargeInterpolatedProfile(
-                    number_of_particles=0.0,
-                    line_density_profile=[1.0, 1.0],
-                    z0=0.0,
-                    dz=1.0,
-                    sigma_x=1.0,
-                    sigma_y=1.0,
-                    length=0.0,
-                    x_co=0.0,
-                    y_co=0.0,
-                )
+                newele = classes.SpaceChargeInterpolatedProfile()
             else:
                 newele = myDrift(length=ee.l)
-                old_pp+=ee.l
+                old_pp += ee.l
         else:
             raise ValueError(f'MAD element "{mad_etype}" not recognized')