Bugfixes and improvements to emittance measurement class

lcls_tools/common/data/emittance.py

@@ -177,12 +177,10 @@ def loss(params):
         bmag = bmag_func(beta, alpha, beta_design, alpha_design)  # result batchshape x nsteps
     else:
         bmag = None
+        twiss_at_screen = None
 
-    results = {}
-    results["emittance"] = emit
-    results["BMAG"] = bmag
-    results["beam_matrix"] = beam_matrix
-    results["twiss_at_screen"] = twiss_at_screen
+    results = {"emittance": emit, "BMAG": bmag, "beam_matrix": beam_matrix,
+               "twiss_at_screen": twiss_at_screen}
     return results
 
 

lcls_tools/common/measurements/emittance_measurement.py

@@ -1,8 +1,9 @@
+import time
 from typing import Optional
 
 import numpy as np
 from numpy import ndarray
-from pydantic import ConfigDict, PositiveInt, field_validator
+from pydantic import ConfigDict, PositiveInt, field_validator, PositiveFloat
 
 from lcls_tools.common.data.emittance import compute_emit_bmag
 from lcls_tools.common.data.model_general_calcs import bdes_to_kmod, get_optics
@@ -12,18 +13,48 @@
 
 class QuadScanEmittance(Measurement):
     """Use a quad and profile monitor/wire scanner to perform an emittance measurement
-    ------------------------
+
     Arguments:
-    energy: beam energy
-    scan_values: BDES values of magnet to scan over
-    magnet: Magnet object used to conduct scan
-    beamsize_measurement: BeamsizeMeasurement object from profile monitor/wire scanner
-    n_measurement_shots: number of beamsize measurements to make per individual quad
-    strength
     ------------------------
+    energy: float
+        Beam energy in GeV
+
+    scan_values: List[float]
+        BDES values of magnet to scan over
+
+    magnet: Magnet
+        Magnet object used to conduct scan
+
+    beamsize_measurement: BeamsizeMeasurement
+        Beamsize measurement object from profile monitor/wire scanner
+
+    n_measurement_shots: int
+        number of beamsize measurements to make per individual quad strength
+
+    rmat: ndarray, optional
+        Transport matricies for the horizontal and vertical phase space from
+        the end of the scanning magnet to the screen, array shape should be 2 x 2 x 2 (
+        first element is the horizontal transport matrix, second is the vertical),
+        if not provided meme is used to calculate the transport matricies
+
+    design_twiss: dict[str, float], optional
+        Dictionary containing design twiss values with the following keys (`beta_x`,
+        `beta_y`, `alpha_x`, `alpha_y`) where the beta/alpha values are in units of [m]/[]
+        respectively
+
+    beam_sizes, dict[str, list[float]], optional
+        Dictionary contraining X-rms and Y-rms beam sizes (keys:`x_rms`,`y_rms`)
+        measured during the quadrupole scan in units of [m].
+
+    wait_time, float, optional
+        Wait time in seconds between changing quadrupole settings and making beamsize
+        measurements.
+
     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
     """
     energy: float
@@ -32,25 +63,38 @@ class QuadScanEmittance(Measurement):
     beamsize_measurement: Measurement
     n_measurement_shots: PositiveInt = 1
 
-    rmat: Optional[ndarray] = None  # 4 x 4 beam transport matrix
+    rmat: Optional[ndarray] = None
     design_twiss: Optional[dict] = None  # design twiss values
     beam_sizes: Optional[dict] = {}
 
+    wait_time: PositiveFloat = 5.0
+
     name: str = "quad_scan_emittance"
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
     @field_validator("rmat")
     def validate_rmat(cls, v, info):
-        assert v.shape == (4, 4)
+        assert v.shape == (2, 2, 2)
         return v
 
     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 (l_eff), rmat, and twiss betas and alphas"""
-
+        """
+        Conduct quadrupole scan to measure the beam phase space.
+
+        Returns:
+        -------
+        result : dict
+            Dictonary containing the following keys
+            - `emittance`: geometric emittance in x/y in units of [mm.mrad]
+            - `BMAG`: Twiss mismatch parameter for each quadrupole strength (unitless)
+            - `twiss_parameters`: Twiss parameters (beta, alpha, gamma) calculated at
+                the screen for each quadrupole strength in each plane
+            - `x_rms`: Measured beam sizes in horizontal direction in [m]
+            - `y_rms`: Measured beam sizes in vertical direction in [m]
+            - `info`: Measurement information for each beamsize measurement
+            """
+
+        self._info = []
         # scan magnet strength and measure beamsize
         self.magnet.scan(
             scan_settings=self.scan_values,
@@ -68,49 +112,69 @@ def measure(self):
             self.rmat = optics["rmat"]
             self.design_twiss = optics["design_twiss"]
 
+        # calculate beam size squared in units of mm
         beamsize_squared = np.vstack((
-            self.beam_sizes["x_rms"], self.beam_sizes["y_rms"]
+            np.array(self.beam_sizes["x_rms"]) * 1e3,
+            np.array(self.beam_sizes["y_rms"]) * 1e3
         )) ** 2
 
-        magnet_length = self.magnet.l_eff
+        magnet_length = self.magnet.metadata.l_eff
         if magnet_length is None:
             raise ValueError("magnet length needs to be specified for magnet "
                              f"{self.magnet.name} to be used in emittance measurement")
 
         # organize data into arrays for use in `compute_emit_bmag`
-        rmat = np.stack([self.rmat[0:2, 0:2], self.rmat[2:4, 2:4]])
+        # rmat = np.stack([self.rmat[0:2, 0:2], self.rmat[2:4, 2:4]])
         twiss_betas_alphas = np.array(
             [[self.design_twiss["beta_x"], self.design_twiss["alpha_x"]],
              [self.design_twiss["beta_y"], self.design_twiss["alpha_y"]]]
         )
+
+        # compute quadrupole focusing strengths
+        # note: need to create negative k values for vertical dimension
         kmod = bdes_to_kmod(self.energy, magnet_length, np.array(self.scan_values))
+        kmod = np.stack((kmod, -kmod))
 
         # compute emittance and bmag
         results = compute_emit_bmag(
             k=kmod,
             beamsize_squared=beamsize_squared,
             q_len=magnet_length,
-            rmat=rmat,
+            rmat=self.rmat,
             twiss_design=twiss_betas_alphas,
         )
 
         results.update({
             "x_rms": self.beam_sizes["x_rms"],
             "y_rms": self.beam_sizes["y_rms"]
         })
+        results.update({"info": self._info})
 
         return results
 
     def measure_beamsize(self):
         """Take measurement from measurement device,
         store beam sizes in self.beam_sizes"""
+        time.sleep(self.wait_time)
+
         results = self.beamsize_measurement.measure(self.n_measurement_shots)
         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"]))
+
+        sigmas = []
+        for ele in results["fit_results"]:
+            sigmas += [ele.rms_size]
+        sigmas = np.array(sigmas)
+
+        # note beamsizes here are in m
+        self.beam_sizes["x_rms"].append(
+            np.mean(sigmas[:, 0]) * self.beamsize_measurement.device.resolution * 1e-6)
+        self.beam_sizes["y_rms"].append(
+            np.mean(sigmas[:, 1]) * self.beamsize_measurement.device.resolution * 1e-6)
+
+        self._info += [results]
 
 
 class MultiDeviceEmittance(Measurement):

lcls_tools/common/measurements/screen_profile.py

@@ -25,11 +25,8 @@ class ScreenBeamProfileMeasurement(Measurement):
     model_config = ConfigDict(arbitrary_types_allowed=True)
     name: str = "beam_profile"
     device: Screen
-    beam_fit: ImageFit = ImageProjectionFit()  # actually want an additional
-    # layer before GaussianFit so that beam_fit_tool is a generic fit tool
-    # not constrained to be of gaussian fit type
+    beam_fit: ImageFit = ImageProjectionFit()
     fit_profile: bool = True
-    # return_images: bool = True
 
     def measure(self, n_shots: int = 1) -> dict:
         """
@@ -46,23 +43,13 @@ def measure(self, n_shots: int = 1) -> dict:
             images.append(self.device.image)
             # TODO: need to add a wait statement in here for images to update
 
-        results = {"raw_images": images}
+        results = {"raw_images": images, "fit_results": None}
 
         if self.fit_profile:
+            fit_results = []
             for image in images:
-                results.update(self.beam_fit.fit_image(image))
-            '''
-            results = {}
-            for image_measurement in images:
-                for key, val in image_measurement.items():
-                    if key in results:
-                        results[key].append(val)
-                    else:
-                        results[key] = [val]
-            '''
-            results = {
-                key: [d.get(key) for d in images]
-                for key in {k for meas in images for k in meas}
-            }
+                fit_results += [self.beam_fit.fit_image(image)]
+
+            results["fit_results"] = fit_results
 
         return results

requirements.txt

@@ -12,3 +12,4 @@ sphinx
 sphinx_rtd_theme
 myst_parser
 scikit-learn
+meme@ git+https://github.com/slaclab/meme.git