Added more stateless-compute methods and tests

sirepo_bluesky/sirepo_bluesky.py

@@ -146,15 +146,7 @@ def compute_crystal_init(self, crystal_element):
         return res
 
     def compute_crystal_orientation(self, crystal_element):
-        res_init = self._post_json(
-            "stateless-compute",
-            {
-                "method": "crystal_init",
-                "optical_element": crystal_element,
-                "simulationId": self.sim_id,
-                "simulationType": self.sim_type,
-            },
-        )
+        res_init = self.compute_crystal_init(crystal_element)
         if res_init.pop("state") != "completed":
             raise SirepoBlueskyClientException("crystal_init returned error state")
         res = self._post_json(
@@ -211,6 +203,40 @@ def get_datafile(self, file_index=-1):
         self._assert_success(response, url)
         return response.content
 
+    def process_beam_parameters(self):
+        res = self._post_json(
+            "stateless-compute",
+            {
+                "ebeam": self.data["models"]["electronBeam"],
+                "ebeam_position": self.data["models"]["electronBeamPosition"],
+                "method": "process_beam_parameters",
+                "simulationId": self.sim_id,
+                "simulationType": self.sim_type,
+                "source_type": self.data["models"]["simulation"]["sourceType"],
+                "undulator_length": self.data["models"]["undulator"]["length"],
+                "undulator_period": self.data["models"]["undulator"]["period"] / 1000.0,
+                "undulator_type": self.data["models"]["tabulatedUndulator"]["undulatorType"],
+            },
+        )
+        return res
+
+    def process_undulator_definition(self):
+        res = self._post_json(
+            "stateless-compute",
+            {
+                "amplitude": self.data["models"]["undulator"]["verticalAmplitude"],
+                "method": "process_undulator_definition",
+                "methodSignature": "process_undulator_definitionverticalDeflectingParameter",
+                "simulationId": self.sim_id,
+                "simulationType": self.sim_type,
+                # Could not locate undulator_definition in source code, using "B" as a placeholder
+                "undulator_definition": "B",
+                "undulator_parameter": self.data["models"]["undulator"]["horizontalDeflectingParameter"],
+                "undulator_period": self.data["models"]["undulator"]["period"] / 1000.0,
+            },
+        )
+        return res
+
     def simulation_list(self):
         """Returns a list of simulations for the authenticated user."""
         return self._post_json("simulation-list", dict(simulationType=self.sim_type))

sirepo_bluesky/sirepo_ophyd.py

@@ -311,6 +311,46 @@ def set(self, value):
         return NullStatus()
 
 
+class SirepoSignalCRL(SirepoSignal):
+    def set(self, value):
+        super().set(value)
+        ret = self.parent.connection.compute_crl_characteristics(self._sirepo_dict)
+        # State is added to the ret dict from compute_crl_characteristics and we
+        # want to make sure the crl element is updated properly when parameters are changed.
+        ret.pop("state")
+        # Update crl element
+        for cpt in ["absoluteFocusPosition", "focalDistance"]:
+            getattr(self.parent, cpt).put(ret[cpt])
+        return NullStatus()
+
+
+class SirepoSignalCrystal(SirepoSignal):
+    def set(self, value):
+        super().set(value)
+        ret = self.parent.connection.compute_crystal_orientation(self._sirepo_dict)
+        # State is added to the ret dict from compute_crl_characteristics and we
+        # want to make sure the crl element is updated properly when parameters are changed.
+        ret.pop("state")
+        # Update crl element
+        for cpt in [
+            "dSpacing",
+            "grazingAngle",
+            "nvx" "nvy",
+            "nvz",
+            "outframevx" "outframevy" "outoptvx" "outoptvy",
+            "outoptvz",
+            "psi0i",
+            "psi0r",
+            "psiHBi",
+            "psiHBr",
+            "psiHi",
+            "psiHr",
+            "tvx" "tvy",
+        ]:
+            getattr(self.parent, cpt).put(ret[cpt])
+        return NullStatus()
+
+
 def create_classes(sirepo_data, connection, create_objects=True, extra_model_fields=[]):
     classes = {}
     objects = {}

sirepo_bluesky/tests/test_stateless_compute.py

@@ -1,16 +1,20 @@
 import json
+import os
 import pprint
 
 import bluesky.plan_stubs as bps
 import dictdiffer
 import numpy as np
+import vcr
 
-from sirepo_bluesky.sirepo_bluesky import SirepoBluesky
+import sirepo_bluesky.tests
 from sirepo_bluesky.sirepo_ophyd import create_classes
 from sirepo_bluesky.utils.json_yaml_converter import dict_to_file
 
+cassette_location = os.path.join(os.path.dirname(sirepo_bluesky.tests.__file__), "vcr_cassettes")
 
-def test_stateless_compute_basic(srw_chx_simulation, RE):
+
+def test_stateless_compute_crl_characteristics_basic(srw_chx_simulation, RE):
     browser_request = {
         "absoluteFocusPosition": -8.7725,
         "attenuationLength": 0.007313,
@@ -84,65 +88,109 @@ def test_stateless_compute_basic(srw_chx_simulation, RE):
     assert np.allclose(newresponse1, newresponse2)
 
 
-def test_stateless_compute_advanced(srw_chx_simulation, tmp_path):
+@vcr.use_cassette(f"{cassette_location}/test_crl_characteristics.yml")
+def test_stateless_compute_crl_characteristics_advanced(srw_chx_simulation, tmp_path):
     classes, objects = create_classes(srw_chx_simulation.data, connection=srw_chx_simulation)
 
-    fname = tmp_path / "test.json"
-
-    _generate_test_file(fname)
+    _generate_test_crl_file(tmp_path / "test_crl_characteristics.json", objects["crl1"], srw_chx_simulation)
 
-    with open(fname, "r") as fp:
+    with open(tmp_path / "test_crl_characteristics.json", "r") as fp:
         combined_dict = json.load(fp)
     assert combined_dict
 
     for key in combined_dict["request"]:
         actual_response = srw_chx_simulation.compute_crl_characteristics(combined_dict["request"][key])
         expected_response = combined_dict["response"][key]
-        assert actual_response.pop("state") == "completed"
+        actual_response.pop("state")
         assert np.allclose(_remove_dict_strings(actual_response), _remove_dict_strings(expected_response))
         print(f"{key} passed")
 
 
+@vcr.use_cassette(f"{cassette_location}/test_crystal_characteristics.yml")
+def test_stateless_compute_crystal(srw_tes_simulation, tmp_path):
+    classes, objects = create_classes(srw_tes_simulation.data, connection=srw_tes_simulation)
+
+    _generate_test_crystal_file(
+        tmp_path / "test_compute_crystal.json", objects["mono_crystal1"], srw_tes_simulation
+    )
+
+    with open(tmp_path / "test_compute_crystal.json", "r") as fp:
+        combined_dict = json.load(fp)
+    assert combined_dict
+
+    for key in combined_dict["request"]:
+        actual_init = srw_tes_simulation.compute_crystal_init(combined_dict["request"][key])
+        expected_init = combined_dict["init"][key]
+        actual_init.pop("state")
+        assert np.allclose(_remove_dict_strings(actual_init), _remove_dict_strings(expected_init))
+
+        actual_orientation = srw_tes_simulation.compute_crystal_orientation(combined_dict["request"][key])
+        expected_orientation = combined_dict["orientation"][key]
+        actual_orientation.pop("state")
+        assert np.allclose(_remove_dict_strings(actual_orientation), _remove_dict_strings(expected_orientation))
+        print(f"{key} passed")
+
+
 def _remove_dict_strings(dict):
     newarr = np.array([], dtype=np.float64)
     for key in dict:
-        try:
-            newarr = np.append(newarr, np.float64(dict[key]))
-        except ValueError:
-            pass
+        if dict[key] is not None:
+            try:
+                newarr = np.append(newarr, np.float64(dict[key]))
+            except ValueError:
+                pass
     return newarr
 
 
-def _generate_test_file(fname):
-    srw_chx_simulation = SirepoBluesky("http://localhost:8000")
-    data, _ = srw_chx_simulation.auth("srw", "HXV1JQ5c")
-    classes, objects = create_classes(srw_chx_simulation.data, connection=srw_chx_simulation)
-    crl1 = objects["crl1"]
-
+def _generate_test_crl_file(fname, crl, simulation):
     combined_dict = {"request": {}, "response": {}}
 
     for radius in range(1, 201):
-        crl1.id._sirepo_dict["tipRadius"] = radius
-        expected_response = srw_chx_simulation.compute_crl_characteristics(crl1.id._sirepo_dict)
-        combined_dict["request"][f"radius{radius}"] = crl1.id._sirepo_dict.copy()
+        crl.id._sirepo_dict["tipRadius"] = radius
+        expected_response = simulation.compute_crl_characteristics(crl.id._sirepo_dict)
+        combined_dict["request"][f"radius{radius}"] = crl.id._sirepo_dict.copy()
         combined_dict["response"][f"radius{radius}"] = expected_response
         print(f"Radius {radius} added")
-    crl1.id._sirepo_dict["tipRadius"] = 1500
+    crl.id._sirepo_dict["tipRadius"] = 1500
 
     for lenses in range(1, 201):
-        crl1.id._sirepo_dict["numberOfLenses"] = lenses
-        expected_response = srw_chx_simulation.compute_crl_characteristics(crl1.id._sirepo_dict)
-        combined_dict["request"][f"lenses{lenses}"] = crl1.id._sirepo_dict.copy()
+        crl.id._sirepo_dict["numberOfLenses"] = lenses
+        expected_response = simulation.compute_crl_characteristics(crl.id._sirepo_dict)
+        combined_dict["request"][f"lenses{lenses}"] = crl.id._sirepo_dict.copy()
         combined_dict["response"][f"lenses{lenses}"] = expected_response
         print(f"Lenses {lenses} added")
-    crl1.id._sirepo_dict["numberOfLenses"] = 1
+    crl.id._sirepo_dict["numberOfLenses"] = 1
 
     for thickness in range(1, 201):
-        crl1.id._sirepo_dict["wallThickness"] = thickness
-        expected_response = srw_chx_simulation.compute_crl_characteristics(crl1.id._sirepo_dict)
-        combined_dict["request"][f"thickness{thickness}"] = crl1.id._sirepo_dict.copy()
+        crl.id._sirepo_dict["wallThickness"] = thickness
+        expected_response = simulation.compute_crl_characteristics(crl.id._sirepo_dict)
+        combined_dict["request"][f"thickness{thickness}"] = crl.id._sirepo_dict.copy()
         combined_dict["response"][f"thickness{thickness}"] = expected_response
         print(f"Thickness {thickness} added")
-    crl1.id._sirepo_dict["wallThickness"] = 80
+    crl.id._sirepo_dict["wallThickness"] = 80
+
+    dict_to_file(combined_dict, fname)
+
+
+def _generate_test_crystal_file(fname, crystal, simulation):
+    combined_dict = {"request": {}, "init": {}, "orientation": {}}
+
+    for energy in range(1, 101):
+        crystal.id._sirepo_dict["energy"] = energy
+        combined_dict["request"][f"Si energy{energy}"] = crystal.id._sirepo_dict.copy()
+        combined_dict["init"][f"Si energy{energy}"] = simulation.compute_crystal_init(crystal.id._sirepo_dict)
+        combined_dict["orientation"][f"Si energy{energy}"] = simulation.compute_crystal_orientation(
+            crystal.id._sirepo_dict
+        )
+        print(f"Si Energy {energy} added")
+    crystal.id._sirepo_dict["material"] = "Germanium (X0h)"
+    for energy in range(1, 21):
+        crystal.id._sirepo_dict["energy"] = energy
+        combined_dict["request"][f"Ge energy{energy}"] = crystal.id._sirepo_dict.copy()
+        combined_dict["init"][f"Ge energy{energy}"] = simulation.compute_crystal_init(crystal.id._sirepo_dict)
+        combined_dict["orientation"][f"Ge energy{energy}"] = simulation.compute_crystal_orientation(
+            crystal.id._sirepo_dict
+        )
+        print(f"Ge Energy {energy} added")
 
     dict_to_file(combined_dict, fname)