Add switch_injection_mode and set_injection

injection.py

@@ -62,17 +62,64 @@ def _get_extraction_table(bucket_list):
 
 get_extraction_table = np.vectorize(_get_extraction_table)
 
-def set_SPM_injection(bucket_list, fine_delay):
+def switch_injection_mode(injection_mode):
+    """
+    Allow to switch the injection mode.
+
+    Parameters
+    ----------
+    injection_mode : "SPM" or "LPM"
+        Injection "Short Pulse Mode" or "Long Pulse Mode".
 
+    """
+    spm = DeviceProxy("LIN/SY/LOCAL.SPM.1")
+    lpm = DeviceProxy("LIN/SY/LOCAL.LPM.1")
+    if injection_mode == "SPM":
+        if spm.spmLinacEvent != 2:
+            lpm.lpmEvent = 0
+            spm.spmLinacEvent = 0
+            time.sleep(0.1)
+            spm.spmLinacEvent = 2
+    elif injection_mode == "LPM":
+        if lpm.lpmEvent != 2:
+            spm.spmLinacEvent = 0
+            lpm.lpmEvent = 0
+            time.sleep(0.1)
+            lpm.lpmEvent = 2
+    else:
+        raise ValueError("injection_mode should be 'LPM' or 'SPM'.")
+
+def set_injection(injection_mode, bucket_list, fine_delay):
+    """
+    Set injection mode and bucket list to be filled.
+
+    Parameters
+    ----------
+    injection_mode : "SPM" or "LPM"
+        Injection "Short Pulse Mode" or "Long Pulse Mode".
+    bucket_list : array like of int or int
+        Bucket number list, buckets are numbered from 1 to 416.
+    fine_delay : float
+        Fine delay in [s].
+
+    """
+    switch_injection_mode(injection_mode)
     central = DeviceProxy("ANS/SY/CENTRAL")
     extraction = get_extraction_table(bucket_list)
     linac = get_linac_table(bucket_list, fine_delay)
     n_bucket = len(extraction)
     central.settables(np.concatenate((np.array([n_bucket]), extraction, linac)))
         
-        
 def inject_beam(booster_cav):
+    """
+    Inject beam into the storage ring using the current injection setup.
 
+    Parameters
+    ----------
+    booster_cav : int
+        Booster cavity number to be used.
+
+    """
     if booster_cav == 1:
         cav_booster = DeviceProxy('BOO/RF/LLE.1')
         rampe = DeviceProxy("BOO/RF/RAMPETENSION")