From fea2e5a21989d95a31e7700c27e787a219e9de2e Mon Sep 17 00:00:00 2001 From: Ben Chen Date: Wed, 1 Oct 2025 00:07:13 +0200 Subject: [PATCH 1/2] Fixed a bug in Runnable.optimize() by updating the import command from ax. Also improved Runnable.is_scan(), made Runnable.extract_beam_function() work for multi-shot runs that are not a scan and edited some doc strings. --- abel/classes/runnable.py | 55 ++++++++++++++++++++++++++++------------ 1 file changed, 39 insertions(+), 16 deletions(-) diff --git a/abel/classes/runnable.py b/abel/classes/runnable.py index c82497e4..c7dd1cd7 100644 --- a/abel/classes/runnable.py +++ b/abel/classes/runnable.py @@ -46,7 +46,7 @@ def perform_shot(self, shot): ## SCAN FUNCTIONALITY def is_scan(self): - return self.scan_fcn is not None + return hasattr(self, 'scan_fcn') and self.scan_fcn is not None # scan function def scan(self, run_name=None, fcn=None, vals=[None], label=None, scale=1, num_shots_per_step=1, step_filter=None, shot_filter=None, savedepth=2, verbose=None, overwrite=False, parallel=False, max_cores=16): @@ -322,7 +322,7 @@ def evaluation_function(params): return val_mean # perform optimization - from ax import optimize as bayes_opt + from ax.service.managed_loop import optimize as bayes_opt best_parameters, best_values, experiment, model = bayes_opt( parameters = parameters, evaluation_function = evaluation_function, @@ -336,7 +336,8 @@ def evaluation_function(params): def extract_beam_function(self, beam_fcn, index=-1, clean=False): """ Extract mean and standard deviation value of beam parameters across a - scan. + scan. For each scan step, the mean and standard deviation are calculated + across all shots in the scan step. Parameters ---------- @@ -356,9 +357,13 @@ def extract_beam_function(self, beam_fcn, index=-1, clean=False): Returns ---------- - vals : 2D ndarray - Each column in ``vals`` corresponds to a scan step, each element in the - column corresponds to a shot. + val_mean : float or 1D ndarray + The mean values of ``beam_fcn``. Each column in ``val_mean`` + corresponds to a scan step. + + val_std : float or 1D ndarray + The standard deviations of ``beam_fcn``. Each column in ``val_std`` + corresponds to a scan step. """ import inspect @@ -374,35 +379,53 @@ def extract_beam_function(self, beam_fcn, index=-1, clean=False): def extract_function(self, fcn, clean=False): import inspect - - # Prepare the arrays - val_mean = np.empty(self.num_steps) - val_std = np.empty(self.num_steps) + input_list = inspect.signature(fcn).parameters # Extract values if self.is_scan(): + + # Prepare the arrays + val_mean = np.empty(self.num_steps) + val_std = np.empty(self.num_steps) + for step in range(self.num_steps): # Get values for this step val_output = np.empty(self.num_shots_per_step) for shot_in_step in range(self.num_shots_per_step): - input_list = inspect.signature(fcn).parameters if 'clean' in input_list: # Check if the input list contains clean. val_output[shot_in_step] = fcn(self[step, shot_in_step], clean=clean) else: val_output[shot_in_step] = fcn(self[step, shot_in_step]) - - # Get step mean and error - val_mean[step] = np.mean(val_output) - val_std[step] = np.std(val_output) + # Get step mean and error + val_mean[step] = np.mean(val_output) + val_std[step] = np.std(val_output) + + elif not self.is_scan() and self.num_shots > 0: + val_output = np.empty(self.num_shots) + for shot in range(self.num_shots): + + if 'clean' in input_list: # Check if the input list contains clean. + val_output[shot] = fcn(self[shot], clean=clean) + else: + val_output[shot] = fcn(self[shot]) + + # Get shot mean and error + val_mean = np.mean(val_output) + val_std = np.std(val_output) + + else: + raise ValueError('Shots not found.') + return val_mean, val_std def scan_steps_beam_params(self, beam_fcn, clean=False): """ - Extract beam parameter data across all scan steps and shots. + Extract output beam parameter data for all shots in each scan step and + store these separately. Parameters ---------- From 1577e5373a240061aeb6a9aaf34a46693dd71696 Mon Sep 17 00:00:00 2001 From: Ben Chen <6263622+ben-c-2013@users.noreply.github.com> Date: Fri, 6 Mar 2026 14:13:43 +0100 Subject: [PATCH 2/2] Saving attempt at updating Runnable.optimize(). Does not work yet. --- abel/classes/runnable.py | 95 +++++++++++++++++++++++++++++++++++----- 1 file changed, 84 insertions(+), 11 deletions(-) diff --git a/abel/classes/runnable.py b/abel/classes/runnable.py index c7dd1cd7..4f659310 100644 --- a/abel/classes/runnable.py +++ b/abel/classes/runnable.py @@ -15,14 +15,25 @@ class Runnable(ABC): # shot tracking function (to be repeated) def perform_shot(self, shot): - + self.shot = shot self.step = self.steps[shot] + + # apply scan function if it exists if self.scan_fcn is not None: - vals_all = np.repeat(self.vals,self.num_shots_per_step) + vals_all = np.repeat(self.vals, self.num_shots_per_step) + + #print('inside perform_shot') #TODO: delete self.scan_fcn(self, vals_all[shot]) + + print('self.scan_fcn(self, vals_all[shot]): ', self.scan_fcn(self, vals_all[shot])) + + print(self.vals) + print(self.plasma_lens.offset_x) # TODO:delete + + #print('inside perform_shot') #TODO: delete # check if object exists if not self.overwrite and os.path.exists(self.object_path(shot)): @@ -105,6 +116,8 @@ def scan(self, run_name=None, fcn=None, vals=[None], label=None, scale=1, num_sh # recalculate number of cores used num_cores = min(max_cores, len(shots_to_perform)) + + #print('inside scan') #TODO: delete # perform parallel tracking with joblib_progress('Tracking shots ('+str(num_cores)+' in parallel)', len(shots_to_perform)): @@ -256,8 +269,12 @@ def final_beam(self): ## Support for Bayesian optimisation through Ax.optimize def set_parameters(self, params): - for key in params.keys(): - self.set_attr(key, params[key]) + print(type(params)) # TODO: delete + print(params[0]) # TODO: delete + + for param in params: + for key in param.keys(): + self.set_attr(key, param[key]) # Setting attribute def set_attr(self, attr, val): @@ -299,6 +316,9 @@ def evaluation_function(params): # set the parameters opt_fcn = lambda obj, _: obj.set_parameters(params) + + + print('hello inside evaluation_function') # TODO: delete # run the simulations self.scan(run_name=self.run_name, fcn=opt_fcn, vals=np.arange(num_steps), step_filter=self.step, num_shots_per_step=num_shots_per_step, savedepth=savedepth, verbose=verbose, overwrite=overwrite, parallel=parallel, max_cores=max_cores, label=label) @@ -321,14 +341,67 @@ def evaluation_function(params): return val_mean + + # Make the list of dict parameters compatible with ax.api.client.configure_experiment() + from ax.api.client import Client + from ax.api.configs import RangeParameterConfig + + def config_parameters(parameters): + exp_params = [None for _ in range(len(parameters))] + for i in range(len(parameters)): + + if 'step_size' not in parameters[i]: + step_size = None + if 'scaling' not in parameters[i]: + scaling = None + + exp_params[i] = [ + RangeParameterConfig( + name=parameters[i]['name'], + bounds=parameters[i]['bounds'], + parameter_type=parameters[i]['type'], + step_size=step_size, + scaling=scaling + ) + ] + return exp_params + + + # perform optimization - from ax.service.managed_loop import optimize as bayes_opt - best_parameters, best_values, experiment, model = bayes_opt( - parameters = parameters, - evaluation_function = evaluation_function, - minimize = True, - total_trials = num_steps - ) + + # from ax.service.managed_loop import optimize as bayes_opt + # best_parameters, best_values, experiment, model = bayes_opt( + # parameters = parameters, + # evaluation_function = evaluation_function, + # minimize = True, + # total_trials = num_steps + # ) + + + client = Client() + + out = evaluation_function(parameters) + print(config_parameters(parameters)) #TODO: delete + + + client.configure_experiment(parameters=config_parameters(parameters)) + + client.configure_optimization(objective="-merit_fcn") + for _ in range(num_steps): + for trial_index, parameters in client.get_next_trials(max_trials=1).items(): + + result = evaluation_function(parameters) + + # Set raw_data as a dictionary with metric names as keys and results as values + raw_data = {merit_fcn: result} + + client.complete_trial( + trial_index=trial_index, + raw_data=raw_data, + ) + + best_parameters, best_values, index, name = client.get_best_parameterization() return best_parameters, best_values