From 62749aed6d8164edab1c5cb038d7aee74bf4fac8 Mon Sep 17 00:00:00 2001
From: liuzzo <simone.liuzzo@esrf.fr>
Date: Thu, 28 Sep 2023 14:40:54 +0200
Subject: [PATCH 1/3] swithc_cavity_and_radiation makes use of
 enable_6d/disable_6d (untested)

---
 pySC/core/lattice_setting.py | 51 +++++++++++++++++++++++++++++-------
 1 file changed, 41 insertions(+), 10 deletions(-)

diff --git a/pySC/core/lattice_setting.py b/pySC/core/lattice_setting.py
index bc1df09..a0e4165 100644
--- a/pySC/core/lattice_setting.py
+++ b/pySC/core/lattice_setting.py
@@ -49,20 +49,51 @@ def switch_cavity_and_radiation(ring: Lattice, *args: str) -> Lattice:  # TODO s
     if invalid_args := [arg for arg in args if arg not in valid_args]:
         raise ValueError(f"Unknown arguments found: {invalid_args}"
                          f"Available options are: {valid_args}")
-    non_rad_pass_methods = ['BndMPoleSymplectic4Pass', 'BndMPoleSymplectic4E2Pass', 'StrMPoleSymplectic4Pass']
-    rad_pass_methods = [method.replace("Pass", "RadPass") for method in non_rad_pass_methods]
+
+    # present radiation state
+    radstate = ring.radiation
+
+    # present cavity state
+    ords = np.arange(len(ring))
+    cavs = [i for i in np.ravel(np.array([ords], dtype=int)) if hasattr(ring[i], 'Frequency')]
+
+    cavstate = False
+    pm=[]
+    for ind in cavs:
+        pm.append(ring[ind].PassMethod)
+    cavpm=np.unique(pm)
+    if len(cavpm) > 1:
+        print(f'cavity pass methods are inconsistent. Assume cavities are off.')
+        cavstate=False
+    elif len(cavpm) == 1:
+        if cavpm == 'RFCavityPass':
+            cavstate = True
+        else:
+            cavstate = False
 
     if 'radiationoff' in args:
-        for ind in range(len(ring)):
-            if ring[ind].PassMethod in rad_pass_methods:
-                ring[ind].PassMethod = ring[ind].PassMethod.replace("Rad", "")
+        if cavstate:  # rad off, cav on
+            ring.disable_6d(cavity_pass='RFCavityPass')
+        else:  # rad off, cav off
+            ring.disable_6d()
+
     elif 'radiationon' in args:
-        for ind in range(len(ring)):
-            if ring[ind].PassMethod in non_rad_pass_methods:
-                ring[ind].PassMethod = ring[ind].PassMethod.replace("Pass", "RadPass")
+        if cavstate:  # rad on, cav on
+            ring.enable_6d()
+        else:  # rad on, cav off
+            ring.enable_6d(cavity_pass='IdentityPass')
+
     if 'cavityoff' in args:
-        return switch_rf(ring, np.arange(len(ring)), False)
+        if radstate:  # rad on, cav off
+            ring.enable_6d(cavity_pass='IdentityPass')
+        else:  # rad off, cav off
+            ring.disable_6d()
+
     elif 'cavityon' in args:
-        return switch_rf(ring, np.arange(len(ring)), True)
+        if radstate:  # rad on, cav on
+            ring.enable_6d()
+        else:  # rad off, cav on
+            ring.disable_6d(cavity_pass='RFCavityPass')
+
     return ring
 

From f68182879e8b81156e90d58da728f68e5dbd0212 Mon Sep 17 00:00:00 2001
From: liuzzo <simone.liuzzo@esrf.fr>
Date: Thu, 28 Sep 2023 15:04:43 +0200
Subject: [PATCH 2/3] swithc_cavity_and_radiation makes use of
 enable_6d/disable_6d, updated

---
 pySC/core/lattice_setting.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/pySC/core/lattice_setting.py b/pySC/core/lattice_setting.py
index a0e4165..0148b26 100644
--- a/pySC/core/lattice_setting.py
+++ b/pySC/core/lattice_setting.py
@@ -83,6 +83,9 @@ def switch_cavity_and_radiation(ring: Lattice, *args: str) -> Lattice:  # TODO s
         else:  # rad on, cav off
             ring.enable_6d(cavity_pass='IdentityPass')
 
+    # update radiation state
+    radstate=ring.radiation
+
     if 'cavityoff' in args:
         if radstate:  # rad on, cav off
             ring.enable_6d(cavity_pass='IdentityPass')

From fe91295dff1100ff1b54340ec1ab57764a1fa556 Mon Sep 17 00:00:00 2001
From: liuzzo <simone.liuzzo@esrf.fr>
Date: Thu, 28 Sep 2023 15:17:20 +0200
Subject: [PATCH 3/3] swithc_cavity_and_radiation makes use of
 enable_6d/disable_6d, updated

---
 pySC/core/lattice_setting.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/pySC/core/lattice_setting.py b/pySC/core/lattice_setting.py
index 0148b26..ba070b5 100644
--- a/pySC/core/lattice_setting.py
+++ b/pySC/core/lattice_setting.py
@@ -50,9 +50,6 @@ def switch_cavity_and_radiation(ring: Lattice, *args: str) -> Lattice:  # TODO s
         raise ValueError(f"Unknown arguments found: {invalid_args}"
                          f"Available options are: {valid_args}")
 
-    # present radiation state
-    radstate = ring.radiation
-
     # present cavity state
     ords = np.arange(len(ring))
     cavs = [i for i in np.ravel(np.array([ords], dtype=int)) if hasattr(ring[i], 'Frequency')]
@@ -83,7 +80,7 @@ def switch_cavity_and_radiation(ring: Lattice, *args: str) -> Lattice:  # TODO s
         else:  # rad on, cav off
             ring.enable_6d(cavity_pass='IdentityPass')
 
-    # update radiation state
+    # get radiation state
     radstate=ring.radiation
 
     if 'cavityoff' in args: