SLAC-ML · wyrnat · Dec 14, 2023
diff --git a/environments/xfel_sase1_launch/README.md b/environments/xfel_sase1_launch/README.md
@@ -0,0 +1,5 @@
+# XFEL_SASE1 Environment for Badger
+
+## Prerequisites
+
+## Usage
diff --git a/environments/xfel_sase1_launch/__init__.py b/environments/xfel_sase1_launch/__init__.py
@@ -0,0 +1,80 @@
+import time
+import numpy as np
+from badger import environment
+from badger.interface import Interface
+from badger.factory import get_intf
+
+
+class Environment(environment.Environment):
+
+    name = 'XFEL Sase1 Launch'
+
+    def __init__(self, interface: Interface, params):
+        super().__init__(interface, params)
+
+        # set interface according to params
+        try:
+            new_interface, _ = get_intf(params['interface'])
+            self.interface = new_interface()
+        except KeyError:
+            print(f'Cannot load interface defined in params')
+
+        self.limits_undulators = {}
+        aircoil_list = ["XFEL.MAGNETS/MAGNET.ML/QF.2162.T2/CURRENT.SP",
+                        "XFEL.MAGNETS/MAGNET.ML/QF.2177.T2/CURRENT.SP",
+                        "XFEL.MAGNETS/MAGNET.ML/QF.2192.T2/CURRENT.SP",
+                        "XFEL.MAGNETS/MAGNET.ML/QF.2207.T2/CURRENT.SP",
+                        "XFEL.MAGNETS/MAGNET.ML/QF.2218.T2/CURRENT.SP"
+                        ]
+        for addr in aircoil_list:
+            self.limits_undulators[addr]  = [20, 60]
+
+
+    def _get_vrange(self, var):
+        return self.limits_undulators[var]
+
+    @staticmethod
+    def list_vars():
+        return ["XFEL.MAGNETS/MAGNET.ML/QF.2162.T2/CURRENT.SP",
+                "XFEL.MAGNETS/MAGNET.ML/QF.2177.T2/CURRENT.SP",
+                "XFEL.MAGNETS/MAGNET.ML/QF.2192.T2/CURRENT.SP",
+                "XFEL.MAGNETS/MAGNET.ML/QF.2207.T2/CURRENT.SP",
+                "XFEL.MAGNETS/MAGNET.ML/QF.2218.T2/CURRENT.SP"
+                ]
+
+
+    # TODO: add losses
+    @staticmethod
+    def list_obses():
+        return ['sases_average']
+
+    @staticmethod
+    def get_default_params():
+        return {
+            'waiting_time': 15,
+            'interface': 'doocs'
+        }
+
+    def _get_var(self, var):
+        return self.interface.get_value(var)
+
+    def _set_var(self, var, x):
+        self.interface.set_value(var, x)
+
+    def _get_obs(self, obs):
+        try:
+            dt = self.params['waiting_time']
+        except KeyError:
+            dt = 0
+        time.sleep(dt)
+
+
+
+        if obs == 'sases_average':
+            values = []
+            for i in range(30):
+                sa1 = self.interface.get_value(
+                    "XFEL.FEL/XGM/XGM.2643.T9/INTENSITY.SA1.RAW.TRAIN")
+                values.append(sa1)
+                time.sleep(0.1)
+            return np.mean(values)
diff --git a/environments/xfel_sase1_launch/configs.yaml b/environments/xfel_sase1_launch/configs.yaml
@@ -0,0 +1,8 @@
+---
+name: xfel_sase1
+version: "0.1"
+dependencies:
+  - numpy
+  - badger-opt
+interface:
+  - doocs
diff --git a/environments/xfel_sase2_launch/README.md b/environments/xfel_sase2_launch/README.md
@@ -0,0 +1,5 @@
+# XFEL_SASE1 Environment for Badger
+
+## Prerequisites
+
+## Usage
diff --git a/environments/xfel_sase2_launch/__init__.py b/environments/xfel_sase2_launch/__init__.py
@@ -0,0 +1,88 @@
+import time
+import numpy as np
+from badger import environment
+from badger.interface import Interface
+from badger.factory import get_intf
+
+
+class Environment(environment.Environment):
+
+    name = 'XFEL Sase2 Launch'
+
+    def __init__(self, interface: Interface, params):
+        super().__init__(interface, params)
+
+        # set interface according to params
+        try:
+            new_interface, _ = get_intf(params['interface'])
+            self.interface = new_interface()
+        except KeyError:
+            print(f'Cannot load interface defined in params')
+
+        self.limits_undulators = {}
+        aircoil_list = [
+                        "XFEL.MAGNETS/MAGNET.ML/QF.2083.T1/CURRENT.SP",
+                        "XFEL.MAGNETS/MAGNET.ML/QF.2098.T1/CURRENT.SP",
+                        "XFEL.MAGNETS/MAGNET.ML/QF.2110.T1/CURRENT.SP",
+                        "XFEL.MAGNETS/MAGNET.ML/QF.2124.T1/CURRENT.SP",
+                        "XFEL.MAGNETS/MAGNET.ML/QF.2139.T1/CURRENT.SP",
+                        "XFEL.MAGNETS/MAGNET.ML/QF.2153.T1/CURRENT.SP",
+                        "XFEL.MAGNETS/MAGNET.ML/QF.2168.T1/CURRENT.SP",
+                        "XFEL.MAGNETS/MAGNET.ML/QF.2180.T1/CURRENT.SP"
+                        ]
+        for addr in aircoil_list:
+            self.limits_undulators[addr]  = [20, 60]
+
+
+    def _get_vrange(self, var):
+        return self.limits_undulators[var]
+
+    @staticmethod
+    def list_vars():
+        return [
+                "XFEL.MAGNETS/MAGNET.ML/QF.2083.T1/CURRENT.SP",
+                "XFEL.MAGNETS/MAGNET.ML/QF.2098.T1/CURRENT.SP",
+                "XFEL.MAGNETS/MAGNET.ML/QF.2110.T1/CURRENT.SP",
+                "XFEL.MAGNETS/MAGNET.ML/QF.2124.T1/CURRENT.SP",
+                "XFEL.MAGNETS/MAGNET.ML/QF.2139.T1/CURRENT.SP",
+                "XFEL.MAGNETS/MAGNET.ML/QF.2153.T1/CURRENT.SP",
+                "XFEL.MAGNETS/MAGNET.ML/QF.2168.T1/CURRENT.SP",
+                "XFEL.MAGNETS/MAGNET.ML/QF.2180.T1/CURRENT.SP"
+                ]
+
+
+    # TODO: add losses
+    @staticmethod
+    def list_obses():
+        return ['sases_average']
+
+    @staticmethod
+    def get_default_params():
+        return {
+            'waiting_time': 15,
+            'interface': 'doocs'
+        }
+
+    def _get_var(self, var):
+        return self.interface.get_value(var)
+
+    def _set_var(self, var, x):
+        self.interface.set_value(var, x)
+
+    def _get_obs(self, obs):
+        try:
+            dt = self.params['waiting_time']
+        except KeyError:
+            dt = 0
+        time.sleep(dt)
+
+
+
+        if obs == 'sases_average':
+            values = []
+            for i in range(30):
+                sa = self.interface.get_value(
+                    "XFEL.FEL/XGM/XGM.2595.T6/INTENSITY.RAW.TRAIN")
+                values.append(sa)
+                time.sleep(0.1)
+            return np.mean(values)
diff --git a/environments/xfel_sase2_launch/configs.yaml b/environments/xfel_sase2_launch/configs.yaml
@@ -0,0 +1,14 @@
+---
+---
+name: xfel_sase2_launch
+version: "0.1"
+dependencies:
+  - numpy
+  - badger-opt
+interface:
+  - doocs
+  - doocs_mock
+  - doocs_real2sim
+params:
+  - waiting_time: 1
+  - interface: doocs
diff --git a/environments/xfel_sase3_launch/README.md b/environments/xfel_sase3_launch/README.md
@@ -0,0 +1,5 @@
+# XFEL_SASE1 Environment for Badger
+
+## Prerequisites
+
+## Usage
diff --git a/environments/xfel_sase3_launch/__init__.py b/environments/xfel_sase3_launch/__init__.py
@@ -0,0 +1,81 @@
+import time
+import numpy as np
+from badger import environment
+from badger.interface import Interface
+from badger.factory import get_intf
+
+
+class Environment(environment.Environment):
+
+    name = 'XFEL Sase2 Launch'
+
+    def __init__(self, interface: Interface, params):
+        super().__init__(interface, params)
+
+        # set interface according to params
+        try:
+            new_interface, _ = get_intf(params['interface'])
+            self.interface = new_interface()
+        except KeyError:
+            print(f'Cannot load interface defined in params')
+
+        self.limits_undulators = {}
+        aircoil_list = [
+                        "XFEL.MAGNETS/MAGNET.ML/QF.2749.T4/CURRENT.SP",
+                        "XFEL.MAGNETS/MAGNET.ML/QF.2761.T4/CURRENT.SP",
+                        "XFEL.MAGNETS/MAGNET.ML/QF.2773.T4/CURRENT.SP",
+                        "XFEL.MAGNETS/MAGNET.ML/QF.2785.T4/CURRENT.SP"
+                        ]
+
+        for addr in aircoil_list:
+            self.limits_undulators[addr]  = [20, 60]
+
+
+    def _get_vrange(self, var):
+        return self.limits_undulators[var]
+
+    @staticmethod
+    def list_vars():
+        return [
+                "XFEL.MAGNETS/MAGNET.ML/QF.2749.T4/CURRENT.SP",
+                "XFEL.MAGNETS/MAGNET.ML/QF.2761.T4/CURRENT.SP",
+                "XFEL.MAGNETS/MAGNET.ML/QF.2773.T4/CURRENT.SP",
+                "XFEL.MAGNETS/MAGNET.ML/QF.2785.T4/CURRENT.SP"
+                ]
+
+
+    # TODO: add losses
+    @staticmethod
+    def list_obses():
+        return ['sases_average']
+
+    @staticmethod
+    def get_default_params():
+        return {
+            'waiting_time': 15,
+            'interface': 'doocs'
+        }
+
+    def _get_var(self, var):
+        return self.interface.get_value(var)
+
+    def _set_var(self, var, x):
+        self.interface.set_value(var, x)
+
+    def _get_obs(self, obs):
+        try:
+            dt = self.params['waiting_time']
+        except KeyError:
+            dt = 0
+        time.sleep(dt)
+
+
+
+        if obs == 'sases_average':
+            values = []
+            for i in range(30):
+                sa = self.interface.get_value(
+                    "XFEL.FEL/XGM/XGM.3130.T10/INTENSITY.SA3.RAW.TRAIN")
+                values.append(sa)
+                time.sleep(0.1)
+            return np.mean(values)
diff --git a/environments/xfel_sase3_launch/configs.yaml b/environments/xfel_sase3_launch/configs.yaml
@@ -0,0 +1,14 @@
+---
+---
+name: xfel_sase3_launch
+version: "0.1"
+dependencies:
+  - numpy
+  - badger-opt
+interface:
+  - doocs
+  - doocs_mock
+  - doocs_real2sim
+params:
+  - waiting_time: 1
+  - interface: doocs