Merge pull request #40 from pluflou/multiwire_scan

Multiwire scan feature -- base.

Author: pluflou

.github/workflows/todo_to_issue.yml

@@ -0,0 +1,9 @@
+name: "Run TODO to Issue"
+on: ["push"]
+jobs:
+    build:
+        runs-on: "ubuntu-latest"
+        steps:
+            - uses: "actions/checkout@v3"
+            - name: "TODO to Issue"
+              uses: "alstr/todo-to-issue-action@v4"

README.md

@@ -13,3 +13,6 @@ Or install by running:
 
     pip install pyemittance
 
+## Setting the correct configuration parameters for the machine/simulation
+
+For examples on how to set the correct configuration and how to use the new API, please see [`examples/tutorial.ipynb`](https://github.com/pluflou/PyEmittance/blob/main/docs/examples/tutorial.ipynb).

docs/examples/multiwire_scan.ipynb

@@ -0,0 +1,87 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "d42f27c2-676d-44af-a548-9531666931a3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pyemittance import PyEmittance\n",
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "879b56fa-2759-4f44-9227-349f88fd23aa",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "meas = PyEmittance(config_name=\"LCLS_WS02\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "dc8046f2-b673-41d6-aae5-cdc11315a44b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "meas.emit_calc_type = 'multiwire'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "d5b19318-7d3c-43bd-b2d4-2af9ac430f53",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "meas.plot = True\n",
+    "meas.save_runs = True"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "003d34ba-aa3b-4614-9b0e-952d0fa90725",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Running offline.\n",
+      "Savepaths not set. Please set them in 'configs/savepaths.json'\n",
+      "Using docs examples directory:  /Users/smiskov/Documents/SLAC/PyEmittance/docs/examples\n"
+     ]
+    }
+   ],
+   "source": [
+    "meas.measure_emittance()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/examples/quad_scan_tutorial.ipynb

@@ -773,7 +773,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.13"
+   "version": "3.7.11"
   },
   "vscode": {
    "interpreter": {

docs/examples/tutorial.ipynb

@@ -0,0 +1 @@
+timestamp,varx_cur,vary_cur,varz_cur,bact_cur,xrms,yrms,xrms_err,yrms_err

examples/summaries/beamsize_config_info.csv

@@ -0,0 +1 @@
+timestamp,ncol,nrow,roi_xmin,roi_xmax,roi_ymin,roi_ymax,resolution,bact,x_size,y_size,xrms,yrms,xrms_err,yrms_err]

examples/summaries/image_acq_quad_info.csv

@@ -1,5 +1,243 @@
 from . import _version
 __version__ = _version.get_versions()['version']
 
-from pyemittance.pyemittance import PyEmittance
+from pyemittance.observer import Observer
+from pyemittance.data_handler import adapt_range, check_symmetry, find_inflection_pnt, add_measurements_btwn_pnts
+from pyemittance.emittance_calc import EmitCalc
+from pyemittance.emittance_calc_multiwire import MultiWireCalc
+from pyemittance.load_json_configs import load_configs
 
+# For setting quad back after scan is done
+from pyemittance.beam_io import MachineIO
+
+# TODO: set up unit testing for PyEmittance, EmitCalc, MultiWireCalc, Observer
+class PyEmittance:
+
+    def __init__(self,
+                 emit_calc_type='quadscan',  # 'quadscan' or 'multiwire'
+                 config_name='LCLS_OTR2',  # make sure config corresponds to calc type
+                 config_dict=None,  # supersedes json configs
+                 meas_type='OTRS',  # only relevant if emit_calc_type=='quadscan'
+                 use_model=False,
+                 online=False
+                 ):
+
+        self.emit_calc_type = emit_calc_type
+        self.meas_type = meas_type
+
+        # if config is not provided, use LCLS-OTR2 as default
+        self.config_name = config_name
+        self.config_dict = config_dict if config_dict else load_configs(self.config_name)
+
+        # if running on machine, use_model=False
+        self.use_model = use_model
+        # only True if setting PVs
+        self.online = online
+        self.verbose = True
+
+        # injector settings (SOL, CQ, SQ) if optimizing
+        # TODO: remove injector settings options from all modules (optimizer needs to be separate)
+        self.inj_config = None
+        if self.emit_calc_type == 'quadscan':
+            # initial rough quad scan
+            self.quad_init = [-6, -4, -2, 0]
+
+        # pyemittance method options
+        self.save_runs = False
+        self.calc_bmag = False
+        self.show_plots = True
+        if self.emit_calc_type == 'quadscan':
+            self.adapt_ranges = True
+            self.num_points = 7
+            self.check_sym = True
+            self.infl_check = True
+            self.add_pnts = True
+            self.use_prev_meas = True
+            self.quad_tol = 0.05
+
+        # simulation/model options
+        # beamsize function from model
+        self.get_bs_model = None
+        self.add_noise = False
+
+        # to save total number of points queried
+        self.return_num_points = False
+
+    def measure_emittance(self):
+        if self.emit_calc_type == 'quadscan':
+           self.measure_emittance_quad_scan()
+        elif self.emit_calc_type == 'multiwire':
+            self.measure_emittance_multiwire()
+        else:
+            raise Exception("Cannot perform measurement. 'emit_calc_type' needs to be 'quadscan' or 'multiwire'.")
+
+    def measure_emittance_quad_scan(self):
+        # get initial points from the observer
+        o = Observer([], {'x': [], 'y': []}, {'x': [], 'y': []})
+        o.use_model = self.use_model
+        o.inj_config = self.inj_config
+        o.online = self.online
+        o.meas_type = self.meas_type
+        o.use_prev_meas = self.use_prev_meas
+        o.tolerance = self.quad_tol
+        o.config_name = self.config_name
+        o.config_dict = self.config_dict
+
+        # print warning
+        if self.online and self.verbose:
+            print("Running online!")
+        else:
+            print("Running offline.")
+
+        # if using sim
+        # set beamsize fn
+        o.get_beamsizes_model = self.get_bs_model
+        o.add_noise = self.add_noise
+
+        energy = o.config_dict['beamline_info']['energy']
+        l_quad = o.config_dict['beamline_info']['l']
+
+        # Save init quad value
+        io = MachineIO(self.config_name, self.config_dict, self.meas_type)
+        io.online = self.online
+        quad_init = io.getquad()
+
+        # Get initial beamsizes (rough scan)
+        bs_x_list, bs_y_list, bs_x_list_err, bs_y_list_err = o.measure_beam(self.quad_init)
+
+        quad_range_x = self.quad_init
+        quad_range_y = self.quad_init
+
+        if self.adapt_ranges:
+            quad_range_x = adapt_range(quad_range_x,
+                                       bs_x_list,
+                                       'x',
+                                       w=bs_x_list_err,
+                                       energy=energy,
+                                       l_eff=l_quad,
+                                       num_points=self.num_points
+                                       )
+            quad_range_y = adapt_range(quad_range_y,
+                                       bs_y_list,
+                                       'y',
+                                       w=bs_y_list_err,
+                                       energy=energy,
+                                       l_eff=l_quad,
+                                       num_points=self.num_points
+                                       )
+
+            new_beamsize_x = o.measure_beam(quad_range_x)
+            bs_x_list, bs_x_list_err = new_beamsize_x[0], new_beamsize_x[2]
+            new_beamsize_y = o.measure_beam(quad_range_y)
+            bs_y_list, bs_y_list_err = new_beamsize_y[1], new_beamsize_y[3]
+
+        if self.check_sym:
+            add_points_x = check_symmetry(quad_range_x, bs_x_list, bs_x_list_err, 'x',
+                                          bs_fn=o.measure_beam, add_meas=True)
+            add_points_y = check_symmetry(quad_range_y, bs_y_list, bs_y_list_err, 'y',
+                                          bs_fn=o.measure_beam, add_meas=True)
+
+            if add_points_x is not None:
+                quad_range_x = add_points_x[0]
+                bs_x_list = add_points_x[1]
+                bs_x_list_err = add_points_x[2]
+
+            if add_points_y is not None:
+                quad_range_y = add_points_y[0]
+                bs_y_list = add_points_y[1]
+                bs_y_list_err = add_points_y[2]
+
+        if self.infl_check:
+            left_x, right_x = find_inflection_pnt(quad_range_x,
+                                                  bs_x_list,
+                                                  show_plots=self.show_plots
+                                                  )
+            left_y, right_y = find_inflection_pnt(quad_range_y,
+                                                  bs_y_list,
+                                                  show_plots=self.show_plots
+                                                  )
+
+            # truncate data
+            quad_range_x = quad_range_x[left_x:right_x]
+            bs_x_list = bs_x_list[left_x:right_x]
+            bs_x_list_err = bs_x_list_err[left_x:right_x]
+
+            quad_range_y = quad_range_y[left_y:right_y]
+            bs_y_list = bs_y_list[left_y:right_y]
+            bs_y_list_err = bs_y_list_err[left_y:right_y]
+
+        if self.add_pnts:
+            quad_range_x, bs_x_list, bs_x_list_err = add_measurements_btwn_pnts(quad_range_x,
+                                                                                bs_x_list,
+                                                                                bs_x_list_err,
+                                                                                self.num_points,
+                                                                                'x',
+                                                                                bs_fn=o.measure_beam
+                                                                                )
+            quad_range_y, bs_y_list, bs_y_list_err = add_measurements_btwn_pnts(quad_range_y,
+                                                                                bs_y_list,
+                                                                                bs_y_list_err,
+                                                                                self.num_points,
+                                                                                'y',
+                                                                                bs_fn=o.measure_beam
+                                                                                )
+
+        # Put quad back
+        io.online = self.online
+        io.setquad(quad_init)
+
+        # Finally get emittance
+        ef = EmitCalc({'x': quad_range_x, 'y': quad_range_y},
+                      {'x': bs_x_list, 'y': bs_y_list},
+                      {'x': bs_x_list_err, 'y': bs_y_list_err},
+                      config_dict=o.config_dict,
+                      config_name=o.config_name
+                      )
+        ef.plot = self.show_plots
+        ef.save_runs = self.save_runs
+        ef.calc_bmag = self.calc_bmag
+
+        # get normalized transverse emittance
+        ef.get_emit()
+
+        # save total number of points queried
+        if self.return_num_points:
+            ef.out_dict["total_points_measured"] = len(o.quad_meas)
+
+        return ef.out_dict
+
+    def measure_emittance_multiwire(self):
+        # get wire measurements from the observer
+        o = Observer()
+        o.inj_config = self.inj_config
+        o.online = self.online
+        o.meas_type = 'WIRE'
+        o.emit_calc_type = self.emit_calc_type
+        o.config_name = self.config_name
+        o.config_dict = self.config_dict
+
+        # print warning
+        if self.online and self.verbose:
+            print("Running online!")
+        else:
+            print("Running offline.")
+
+        # Get beamsizes
+        # This returns lists now: xrms, yrms, xrms_err, yrms_err
+        # TODO: implement error dict to be returned to track multiwire success
+        bs_x_list, bs_y_list, bs_x_list_err, bs_y_list_err = o.multiwire_measure_beam()
+
+        # Finally get emittance
+        ef = MultiWireCalc(beam_vals={'x': bs_x_list, 'y': bs_y_list},
+                           beam_vals_err={'x': bs_x_list_err, 'y': bs_y_list_err},
+                           config_dict=o.config_dict,
+                           config_name=o.config_name
+                           )
+        ef.plot = self.show_plots
+        ef.save_runs = self.save_runs
+        ef.calc_bmag = self.calc_bmag
+
+        # get normalized transverse emittance
+        ef.get_emit()
+
+        return ef.out_dict