Merge pull request #39 from ChristopherMayes/rework_output

Rework output

Author: pluflou

docs/examples/tutorial.ipynb

@@ -1,5 +1,5 @@
-# conda env create -f environment-dev.yml
-name: pyemittance-dev
+# conda env create -f environment.yml
+name: pyemittance
 channels:
   - conda-forge
 dependencies:

environment.yml

@@ -1,182 +1,5 @@
 from . import _version
 __version__ = _version.get_versions()['version']
 
-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.load_json_configs import load_configs
-
-
-class PyEmittance:
-
-    def __init__(self,
-                 config_name='LCLS_OTR2',
-                 config_dict=None,  # supersedes json configs
-                 meas_type='OTRS',
-                 use_model=False,
-                 online=False
-                 ):
-
-        # 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)
-
-        self.meas_type = meas_type
-        # 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
-        self.inj_config = None
-        # initial rough quad scan
-        self.quad_init = [-6, -4, -2, 0]
-
-        # pyemittance method options
-        self.adapt_ranges = True
-        self.num_points = 7
-        self.check_sym = True
-        self.infl_check = True
-        self.add_pnts = True
-        self.show_plots = True
-        self.use_prev_meas = True
-        self.quad_tol = 0.05
-        self.save_runs = False
-        self.calc_bmag = False
-
-        # 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):
-        # 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
-
-        # 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
-
-        o.config_name = self.config_name
-        o.config_dict = self.config_dict
-
-        energy = o.config_dict['beamline_info']['energy']
-        l_quad = o.config_dict['beamline_info']['l']
-
-        # 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
-                                                                                )
-
-        # 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
+from pyemittance.pyemittance import PyEmittance
 

pyemittance/__init__.py

@@ -108,9 +108,9 @@ def eval_emit_machine(inj_config=None,
     ef.get_emit()
 
     # save total number of points queried
-    ef.out_dict["total_points_measured"] = len(o.quad_meas)
+    ef.output["total_points_measured"] = len(o.quad_meas)
 
-    return ef.out_dict
+    return ef.output
 
 
 def eval_emit_surrogate(get_bs_model,
@@ -214,6 +214,6 @@ def eval_emit_surrogate(get_bs_model,
     ef.get_emit()
 
     # save total number of points queried
-    ef.out_dict["total_points_measured"] = len(o.quad_meas)
+    ef.output["total_points_measured"] = len(o.quad_meas)
 
-    return ef.out_dict
+    return ef.output

pyemittance/emit_eval_example.py

@@ -57,22 +57,7 @@ def __init__(self,
         self.init_saving()
 
         # Main output of emittance calc
-        self.out_dict = {'nemitx': None,
-                         'nemity': None,
-                         'nemitx_err': None,
-                         'nemity_err': None,
-                         'nemit': None,
-                         'nemit_err': None,
-                         'bmagx': None,
-                         'bmagy': None,
-                         'bmagx_err': None,
-                         'bmagy_err': None,
-                         'bmag_emit': None,
-                         'bmag_emit_err': None,
-                         'opt_q_x': None,
-                         'opt_q_y': None,
-                         'total_points_measured': None
-                         }
+        self.output = {}
 
     def load_config(self):
         # if path exists, load from path
@@ -121,48 +106,55 @@ def get_emit(self):
 
             weights = self.weighting_func(bs, bs_err) # 1/sigma
 
-            # Storing quadvals and beamsizes in self.out_dict for plotting purposes
-            self.out_dict[f'quadvals{dim}'] = list(q)
-            self.out_dict[f'beamsizes{dim}'] = list(bs)
-            self.out_dict[f'beamsizeserr{dim}'] =  list(bs_err)
+            # Storing quadvals and beamsizes in self.output for plotting purposes
+            self.output[f'quadvals{dim}'] = list(q)
+            self.output[f'beamsizes{dim}'] = list(bs)
+            self.output[f'beamsizeserr{dim}'] =  list(bs_err)
 
             res = estimate_sigma_mat_thick_quad(bs, kL, bs_err, weights, dim=dim, Lquad=self.quad_len,
-                                                energy=self.energy, rmat=self.rmat, calc_bmag=self.calc_bmag,
+                                                energy=self.energy, rmat=self.rmat,
                                                 plot=self.plot, verbose=self.verbose)
-            if np.isnan(res[0]):
-                self.out_dict['nemitx'], self.out_dict['nemity'] = np.nan, np.nan
-                self.out_dict['nemitx_err'], self.out_dict['nemity_err'] = np.nan, np.nan
-                self.out_dict['bmagx'], self.out_dict['bmagy'] = np.nan, np.nan
-                self.out_dict['bmagx_err'], self.out_dict['bmagy_err'] = np.nan, np.nan
-                return self.out_dict
-            else:
-                emit, emit_err, beta_rel_err, alpha_rel_err = res[0:4]
-                if self.calc_bmag:
-                    sig_11, sig_12, sig_22 = res[4:]
-
-            norm_emit_res = normalize_emit(emit, emit_err, self.energy)
-            self.out_dict[f'nemit{dim}'] = normalize_emit(emit, emit_err, self.energy)[0]
-            self.out_dict[f'nemit{dim}_err'] = normalize_emit(emit, emit_err, self.energy)[1]
+            # Add all results
+            self.output.update(res)            
+            
+            # Skip further calcs if there was an error
+            if res[f'error_{dim}']:        
+                continue
 
             if self.calc_bmag:
-                self.sig_mat_screen[dim] = [sig_11, sig_12, sig_22]
+                if dim =='x':
+                    sig_11 = res['screen_sigma_11']
+                    sig_12 = res['screen_sigma_12']
+                    sig_22 = res['screen_sigma_22']
+                    
+                else:
+                    sig_11 = res['screen_sigma_33']
+                    sig_12 = res['screen_sigma_34']
+                    sig_22 = res['screen_sigma_44']
+                self.sig_mat_screen[dim] = [sig_11, sig_12, sig_22] 
+                    
+                beta_rel_err  = res[f'beta_{dim}_rel_err']   
+                alpha_rel_err = res[f'alpha_{dim}_rel_err']   
+                
+                
                 self.beta_err = beta_rel_err
                 self.alpha_err = alpha_rel_err
 
                 bmag_calc_res = self.get_twiss_bmag(dim=dim)
                 # Get bmag and bmag_err
-                self.out_dict[f'bmag{dim}'] = bmag_calc_res[0]
-                self.out_dict[f'bmag{dim}_err'] = bmag_calc_res[1]
+                self.output[f'screen_bmag{dim}'] = bmag_calc_res[0]
+                self.output[f'screen_bmag{dim}_err'] = bmag_calc_res[1]
                 # Get best value for scanning quad
-                self.out_dict[f'opt_q_{dim}'] = q[bmag_calc_res[2]]
+                self.output[f'optimal_quadval_{dim}'] = q[bmag_calc_res[2]]
 
-        # get geometric mean (set to None when x and y are not calculated)
-        self.get_gmean_emit()
+        # get geometric mean if possible
+        if (not self.output['error_x']) and  (not self.output['error_y']) :
+            self.get_gmean_emit()
 
         if self.save_runs:
             self.save_run()
 
-        return self.out_dict
+        return self.output
 
     def get_twiss_bmag(self, dim='x'):
 
@@ -185,31 +177,31 @@ def get_twiss_bmag(self, dim='x'):
     def get_gmean_emit(self):
 
         try:
-            nemit = np.sqrt( self.out_dict['nemitx'] * self.out_dict['nemity'] )
-            nemit_err = nemit * ( (self.out_dict['nemitx_err']/self.out_dict['nemitx'])**2 +
-                                  (self.out_dict['nemity_err']/self.out_dict['nemity'])**2 )**0.5
+            nemit = np.sqrt( self.output['norm_emit_x'] * self.output['norm_emit_y'] )
+            nemit_err = nemit * ( (self.output['norm_emit_x_err']/self.output['norm_emit_x'])**2 +
+                                  (self.output['norm_emit_y_err']/self.output['norm_emit_y'])**2 )**0.5
 
-            self.out_dict['nemit'] = nemit
-            self.out_dict['nemit_err'] = nemit_err
+            self.output['sqrt_norm_emit_4d'] = nemit
+            self.output['sqrt_norm_emit_4d_err'] = nemit_err
 
-            if self.out_dict['bmagx'] is not None and self.out_dict['bmagy'] is not None:
-                nbmag = np.sqrt( self.out_dict['bmagx'] * self.out_dict['bmagy'] )
+            if 'bmag_x' in self.output and 'bmag_y' in self.output:
+                nbmag = np.sqrt( self.output['bmag_x'] * self.output['bmag_y'] )
                 bmag_emit_err = nemit*nbmag * (
-                    (self.out_dict['nemitx_err']/self.out_dict['nemitx'])**2 +
-                    (self.out_dict['nemity_err']/self.out_dict['nemity'])**2 +
-                    (self.out_dict['bmagx_err']/self.out_dict['bmagx'])**2 +
-                    (self.out_dict['bmagy_err']/self.out_dict['bmagy'])**2)**0.5
-                self.out_dict['bmag_emit'] = nemit * nbmag
-                self.out_dict['bmag_emit_err'] = bmag_emit_err
+                    (self.output['norm_emit_x_err']/self.output['norm_emit_x'])**2 +
+                    (self.output['norm_emit_y_err']/self.output['norm_emit_y'])**2 +
+                    (self.output['bmag_x_err']/self.output['bmag_x'])**2 +
+                    (self.output['bmag_y_err']/self.output['bmag_y'])**2)**0.5
+                self.output['bmag_emit'] = nemit * nbmag
+                self.output['bmag_emit_err'] = bmag_emit_err
 
         except TypeError:
-            self.out_dict['nemit'] = np.nan
-            self.out_dict['nemit_err'] = np.nan
-            self.out_dict['bmag_emit'] = np.nan
-            self.out_dict['bmag_emit_err'] = np.nan
+            self.output['sqrt_norm_emit_4d'] = np.nan
+            self.output['sqrt_norm_emit_4d_err'] = np.nan
+            self.output['bmag_emit'] = np.nan
+            self.output['bmag_emit_err'] = np.nan
 
     def save_run(self):
-        save_emit_run(self.out_dict, path=self.config_dict['savepaths']['fits'])
+        save_emit_run(self.output, path=self.config_dict['savepaths']['fits'])
 
     def init_saving(self):
         """Initialize dirs and files for saving"""