Emittance Measurement Class

lcls_tools/common/data/model_general_calcs.py

@@ -1,3 +1,6 @@
+from meme.model import Model
+
+
 def bmag(twiss, twiss_reference):
     """Calculates BMAG from imput twiss and reference twiss"""
     beta_a, alpha_a, beta_b, alpha_b = twiss
@@ -49,3 +52,11 @@ def bdes_to_kmod(e_tot=None, effective_length=None, bdes=None,
         bp = ele["E_TOT"] / 1e9 / 299.792458 * 1e4  # kG m
         effective_length = ele["L"]
     return bdes / effective_length / bp  # kG / m / kG m = 1/m^2
+
+
+def get_optics(magnet: str, measurement_device: str, beamline: str):
+    """Get rmats and twiss for a given beamline, magnet and measurement device"""
+    model = Model(beamline)
+    rmats = model.get_rmat(from_device=magnet, to_device=measurement_device, from_device_pos='mid')
+    twiss = model.get_twiss(measurement_device)
+    return rmats, twiss

lcls_tools/common/measurements/emittance_measurement.py

@@ -0,0 +1,100 @@
+import numpy as np
+from pydantic import ConfigDict
+
+from lcls_tools.common.devices.magnet import MagnetCollection
+from lcls_tools.common.measurements.measurement import Measurement
+from lcls_tools.common.data.model_general_calcs import bdes_to_kmod, get_optics
+
+from typing import Optional
+
+
+class QuadScanEmittance(Measurement):
+    """Use a quad and profile monitor/wire scanner to perform an emittance measurement
+    ------------------------
+    Arguments:
+    beamline: beamline where the devices are located
+    energy: beam energy
+    magnet_collection: MagnetCollection object of magnets for an area of the beamline (use create_magnet())
+    magnet_name: name of magnet
+    magnet_length: length of magnet
+    scan_values: BDES values of magnet to scan over
+    device_measurement: Measurement object of profile monitor/wire scanner
+    ------------------------
+    Methods:
+    measure: does the quad scan, getting the beam sizes at each scan value,
+    gets the rmat and twiss parameters, then computes and returns the emittance and BMAG
+    measure_beamsize: take measurement from measurement device, store beam sizes
+    """
+    beamline: str
+    energy: float
+    magnet_collection: MagnetCollection
+    magnet_name: str
+    # TODO: remove magnet_length once lengths added to yaml files
+    magnet_length: float
+    scan_values: list[float]
+    device_measurement: Measurement
+
+    rmat: Optional[np.ndarray] = None
+    twiss: Optional[np.ndarray] = None
+    beam_sizes: Optional[dict] = {}
+
+    name: str = "quad_scan_emittance"
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    @property
+    def magnet_settings(self) -> list[dict]:
+        return [{self.magnet_name: value} for value in self.scan_values]
+
+    def measure(self):
+        """Returns the emittance, BMAG, x_rms and y_rms
+        Get the rmat, twiss parameters, and measured beam sizes
+        Perform the scan, measuring beam sizes at each scan value
+        Compute the emittance and BMAG using the geometric focusing strengths,
+        beam sizes squared, magnet length, rmat, and twiss betas and alphas"""
+        self.magnet_collection.scan(scan_settings=self.magnet_settings, function=self.measure_beamsize)
+        self.rmat, self.twiss = get_optics(self.magnet_name, self.device_measurement.device.name, self.beamline)
+        beamsize_squared = np.vstack((self.beam_sizes["x_rms"], self.beam_sizes["y_rms"]))**2
+        # TODO: uncomment once lengths added to yaml files
+        # magnet_length = self.magnet_collection.magnets[self.magnet_name].length
+        rmat = np.array([self.rmat[0:2, 0:2], self.rmat[2:4, 2:4]])  # x rmat and y rmat
+        twiss_betas_alphas = np.array([[self.twiss["beta_x"], self.twiss["alpha_x"]],
+                                       [self.twiss["beta_y"], self.twiss["alpha_y"]]])
+        kmod = bdes_to_kmod(self.energy, self.magnet_length, np.array(self.scan_values))
+        emittance, bmag, _, _ = compute_emit_bmag(
+            k=kmod,
+            beamsize_squared=beamsize_squared,
+            q_len=self.magnet_length,
+            rmat=rmat,
+            twiss_design=twiss_betas_alphas
+        )
+
+        results = {
+            "emittance": emittance,
+            "BMAG": bmag,
+            "x_rms": self.beam_sizes["x_rms"],
+            "y_rms": self.beam_sizes["y_rms"]
+        }
+
+        return results
+
+    def measure_beamsize(self):
+        """Take measurement from measurement device, store beam sizes in self.beam_sizes"""
+        results = self.device_measurement.measure()
+        if "x_rms" not in self.beam_sizes:
+            self.beam_sizes["x_rms"] = []
+        if "y_rms" not in self.beam_sizes:
+            self.beam_sizes["y_rms"] = []
+        self.beam_sizes["x_rms"].append(np.mean(results["Sx"]))
+        self.beam_sizes["y_rms"].append(np.mean(results["Sy"]))
+
+
+class MultiDeviceEmittance(Measurement):
+    name: str = "multi_device_emittance"
+
+    def measure(self):
+        raise NotImplementedError("Multi-device emittance not yet implemented")
+
+
+# TODO: delete and import actual compute_emit_bmag
+def compute_emit_bmag(self, k, beamsize_squared, q_len, rmat, twiss_design, thin_lens, maxiter):
+    raise NotImplementedError("compute_emit_bmag to be implemented elsewhere")

tests/unit_tests/lcls_tools/common/measurements/test_emittance_measurement.py

@@ -0,0 +1,94 @@
+from unittest import TestCase
+from unittest.mock import patch, Mock
+import numpy as np
+
+from lcls_tools.common.devices.reader import create_magnet
+from lcls_tools.common.measurements.emittance_measurement import QuadScanEmittance
+from lcls_tools.common.measurements.screen_profile import ScreenBeamProfileMeasurement
+
+
+class EmittanceMeasurementTest(TestCase):
+    def setUp(self) -> None:
+        self.options = [
+            "TRIM",
+            "PERTURB",
+            "BCON_TO_BDES",
+            "SAVE_BDES",
+            "LOAD_BDES",
+            "UNDO_BDES",
+            "DAC_ZERO",
+            "CALIB",
+            "STDZ",
+            "RESET",
+            "TURN_OFF",
+            "TURN_ON",
+            "DEGAUSS",
+        ]
+        self.ctrl_options_patch = patch("epics.PV.get_ctrlvars", new_callable=Mock)
+        self.mock_ctrl_options = self.ctrl_options_patch.start()
+        self.mock_ctrl_options.return_value = {"enum_strs": tuple(self.options)}
+        self.magnet_collection = create_magnet(area="GUNB")
+        return super().setUp()
+
+    @patch("lcls_tools.common.measurements.emittance_measurement.get_optics")
+    @patch("lcls_tools.common.measurements.emittance_measurement.compute_emit_bmag")
+    @patch(
+        "lcls_tools.common.devices.magnet.Magnet.is_bact_settled",
+        new_callable=Mock,
+    )
+    @patch("epics.PV.put", new_callable=Mock)
+    @patch("lcls_tools.common.devices.magnet.Magnet.trim", new_callable=Mock)
+    def test_quad_scan(self, mock_trim, mock_put, mock_bact_settle, mock_compute_emit_bmag, mock_get_optics):
+        """Test quad scan emittance measurement"""
+        rmat_mock = np.array([
+            [2.5, 4.2, 0.0, 0.0, 1.6, -1.0],
+            [3.6, 1.0, 0.0, 0.0, 3.8, -6.9],
+            [0.0, 0.0, -5.1, 4.1, 0.0, 0.0],
+            [0.0, 0.0, -3.6, 9.9, 0.0, 0.0],
+            [3.5, -2.5, 0.0, 0.0, 1.0, 5.9],
+            [-1.0, -1.2, 0.0, 0.0, -4.9, 1.0]
+        ])
+        twiss_dtype = np.dtype([('s', 'float32'), ('z', 'float32'), ('length', 'float32'),
+                                ('p0c', 'float32'), ('alpha_x', 'float32'), ('beta_x', 'float32'),
+                                ('eta_x', 'float32'), ('etap_x', 'float32'), ('psi_x', 'float32'),
+                                ('alpha_y', 'float32'), ('beta_y', 'float32'), ('eta_y', 'float32'),
+                                ('etap_y', 'float32'), ('psi_y', 'float32')])
+        twiss_mock = np.array(
+            [(11.9, 2027.7, 0.05, 1.3, 3.9, 5.5,
+              -6.1, 1.2, 5.9, 0.01, 5.3, 0., 0., 3.5)],
+            dtype=twiss_dtype
+        )
+        mock_get_optics.return_value = (rmat_mock, twiss_mock)
+        mock_compute_emit_bmag.return_value = (1.5e-9, 1.5, None, None)
+        mock_bact_settle.return_value = True
+        beamline = "SC_DIAG0"
+        energy = 3.0e9
+        magnet_name = "CQ01B"
+        magnet_length = 1.0
+        scan_values = [-6.0, -3.0, 0.0]
+        number_scan_values = len(scan_values)
+        profmon_measurement = Mock(spec=ScreenBeamProfileMeasurement)
+        profmon_measurement.device = Mock()
+        profmon_measurement.device.name = "YAG01B"
+        profmon_measurement.measure.return_value = {
+            "Cx": [0.0],
+            "Cy": [0.0],
+            "Sx": [50.0],
+            "Sy": [50.0],
+            "bb_penalty": [0.0],
+            "total_intensity": [1.0e6],
+        }
+        quad_scan = QuadScanEmittance(beamline=beamline, energy=energy, magnet_collection=self.magnet_collection,
+                                      magnet_name=magnet_name, magnet_length=magnet_length, scan_values=scan_values,
+                                      device_measurement=profmon_measurement)
+        result_dict = quad_scan.measure()
+        assert result_dict["emittance"] == mock_compute_emit_bmag.return_value[0]
+        assert result_dict["BMAG"] == mock_compute_emit_bmag.return_value[1]
+        assert result_dict["x_rms"] == [profmon_measurement.measure.return_value["Sx"]] * number_scan_values
+        assert result_dict["y_rms"] == [profmon_measurement.measure.return_value["Sy"]] * number_scan_values
+        mock_compute_emit_bmag.assert_called_once()
+        mock_get_optics.assert_called_once()
+        self.assertEqual(mock_put.call_count, number_scan_values)
+        self.assertEqual(mock_trim.call_count, number_scan_values)
+        self.assertEqual(mock_bact_settle.call_count, number_scan_values)
+        assert profmon_measurement.measure.call_count == number_scan_values