diff --git a/src/qutip_qoc/analytical_control.py b/src/qutip_qoc/analytical_control.py index daf979f..0cc15e3 100644 --- a/src/qutip_qoc/analytical_control.py +++ b/src/qutip_qoc/analytical_control.py @@ -16,7 +16,7 @@ from qutip_qoc.crab import Multi_CRAB from qutip_qoc.grape import Multi_GRAPE -__all__ = ["optimize_pulses"] +__all__ = ["global_local_optimization"] def get_init_and_bounds_from_options(lst, input): @@ -173,7 +173,7 @@ def opt_callback(self, x, f, accept): return terminate -def optimize_pulses( +def global_local_optimization( objectives, pulse_options, time_interval, @@ -182,7 +182,7 @@ def optimize_pulses( optimizer_kwargs, minimizer_kwargs, integrator_kwargs, - qtrl_optimizer=None, + qtrl_optimizers=None, ): """ Optimize a pulse sequence to implement a given target unitary by optimizing @@ -281,7 +281,7 @@ def optimize_pulses( integrator_kwargs["normalize_output"] = False integrator_kwargs.setdefault("progress_bar", False) - # extract initial and boundary values + # extract initial and boundary values for global and local optimizer x0, bounds = [], [] for key in pulse_options.keys(): get_init_and_bounds_from_options(x0, pulse_options[key].get("guess")) @@ -290,8 +290,7 @@ def optimize_pulses( get_init_and_bounds_from_options(x0, time_options.get("guess", None)) get_init_and_bounds_from_options(bounds, time_options.get("bounds", None)) - if len(x0) != 0: - optimizer_kwargs.setdefault("x0", np.concatenate(x0)) + optimizer_kwargs["x0"] = np.concatenate(x0) # algorithm specific settings if algorithm_kwargs.get("alg") == "JOPT": @@ -316,29 +315,10 @@ def optimize_pulses( **integrator_kwargs, ) elif algorithm_kwargs.get("alg") == "CRAB": - multi_objective = Multi_CRAB( - qtrl_optimizer, - objectives, - time_interval, - time_options, - pulse_options, - algorithm_kwargs, - guess_params=optimizer_kwargs["x0"], - **integrator_kwargs, - ) - minimizer_kwargs.setdefault("method", "Nelder-Mead") + multi_objective = Multi_CRAB(qtrl_optimizers) elif algorithm_kwargs.get("alg") == "GRAPE": - multi_objective = Multi_GRAPE( - qtrl_optimizer, - objectives, - time_interval, - time_options, - pulse_options, - algorithm_kwargs, - guess_params=optimizer_kwargs["x0"], - **integrator_kwargs, - ) + multi_objective = Multi_GRAPE(qtrl_optimizers) # optimizer specific settings opt_method = optimizer_kwargs.get( @@ -352,7 +332,7 @@ def optimize_pulses( optimizer_kwargs.setdefault( # or "max_iter" "niter", optimizer_kwargs.get( # use algorithm_kwargs - "max_iter", algorithm_kwargs.get("max_iter", 1000) + "max_iter", algorithm_kwargs.get("max_iter", 1) ), ) @@ -386,7 +366,7 @@ def optimize_pulses( time_interval, guess_params=x0, var_time=var_t, - qtrl_optimizer=qtrl_optimizer, + qtrl_optimizers=qtrl_optimizers, ) # Callback instance for termination and logging diff --git a/src/qutip_qoc/crab.py b/src/qutip_qoc/crab.py index 5aa5dda..229ef63 100644 --- a/src/qutip_qoc/crab.py +++ b/src/qutip_qoc/crab.py @@ -70,16 +70,11 @@ class Multi_CRAB: def __init__( self, - qtrl_optimizer, - objectives, - time_interval, - time_options, - pulse_options, - alg_kwargs, - guess_params, - **integrator_kwargs, + qtrl_optimizers, ): - for optim in list(qtrl_optimizer): + if not isinstance(qtrl_optimizers, list): + qtrl_optimizers = [qtrl_optimizers] + for optim in qtrl_optimizers: crab = copy.deepcopy(optim) crab.fid_err_func_wrapper = types.MethodType(fid_err_func_wrapper, crab) # Stack for each objective diff --git a/src/qutip_qoc/grape.py b/src/qutip_qoc/grape.py index dee6a97..077bc0f 100644 --- a/src/qutip_qoc/grape.py +++ b/src/qutip_qoc/grape.py @@ -119,18 +119,11 @@ class Multi_GRAPE: def __init__( self, - qtrl_optimizer, - objectives, - time_interval, - time_options, - pulse_options, - alg_kwargs, - guess_params, - **integrator_kwargs, + qtrl_optimizers, ): - if not isinstance(qtrl_optimizer, list): - qtrl_optimizer = [qtrl_optimizer] - for optim in qtrl_optimizer: + if not isinstance(qtrl_optimizers, list): + qtrl_optimizers = [qtrl_optimizers] + for optim in qtrl_optimizers: grape = copy.deepcopy(optim) grape.fid_err_func_wrapper = types.MethodType(fid_err_func_wrapper, grape) # Stack for each objective diff --git a/src/qutip_qoc/optimize.py b/src/qutip_qoc/optimize.py index ce5933f..322f942 100644 --- a/src/qutip_qoc/optimize.py +++ b/src/qutip_qoc/optimize.py @@ -1,12 +1,9 @@ -import time import numpy as np import qutip_qtrl.logging_utils as logging -from qutip_qtrl.pulseoptim import optimize_pulse import qutip_qtrl.pulseoptim as cpo -from qutip_qoc.analytical_control import optimize_pulses as opt_pulses -from qutip_qoc.result import Result +from qutip_qoc.analytical_control import global_local_optimization __all__ = ["optimize_pulses"] @@ -15,11 +12,11 @@ def optimize_pulses( objectives, pulse_options, time_interval, - time_options={}, - algorithm_kwargs={}, - optimizer_kwargs={}, - minimizer_kwargs={}, - integrator_kwargs={}, + time_options=None, + algorithm_kwargs=None, + optimizer_kwargs=None, + minimizer_kwargs=None, + integrator_kwargs=None, ): """ Wrapper to choose between GOAT/JOPT and GRAPE/CRAB optimization. @@ -119,17 +116,22 @@ def optimize_pulses( result : :class:`qutip_qoc.Result` Optimization result. """ - alg = algorithm_kwargs.get("alg", "GRAPE") - - if isinstance(objectives, list): - if alg == "GRAPE" and len(objectives) != 1: - raise TypeError( - "GRAPE optimization only supports one objective at a time. Please use CRAB, GOAT or JOPT for multiple objectives." - ) - else: - objectives = [objectives] + if time_options is None: + time_options = {} + if algorithm_kwargs is None: + algorithm_kwargs = {} + if optimizer_kwargs is None: + optimizer_kwargs = {} + if minimizer_kwargs is None: + minimizer_kwargs = {} + if integrator_kwargs is None: + integrator_kwargs = {} + + alg = algorithm_kwargs.get("alg", "GRAPE") # works with most input types Hd_lst, Hc_lst = [], [] + if not isinstance(objectives, list): + objectives = [objectives] for objective in objectives: # extract drift and control Hamiltonians from the objective Hd_lst.append(objective.H[0]) @@ -144,8 +146,7 @@ def optimize_pulses( for key in pulse_options.keys(): x0.append(pulse_options[key].get("guess")) bounds.append(pulse_options[key].get("bounds")) - - try: + try: # GRAPE, CRAB format lbound = [b[0][0] for b in bounds] ubound = [b[0][1] for b in bounds] except Exception: @@ -158,122 +159,54 @@ def optimize_pulses( else: log_level = logging.WARN - # low level minimizer overrides high level algorithm kwargs - max_iter = minimizer_kwargs.get("options", {}).get( - "maxiter", algorithm_kwargs.get("max_iter", 1000) - ) - - optim_method = minimizer_kwargs.get( - "method", algorithm_kwargs.get("optim_method", "BFGS") - ) - - result = Result(objectives, time_interval) - - start_time = time.time() - - if alg == "DEPRECATED": - # only allow for scalar bound - lbound = lbound[0] - ubound = ubound[0] - - min_g = minimizer_kwargs.get("gtol", 1e-10) - - algorithm_kwargs["alg_params"] = { - "init_amps": np.array(x0).T, - } - - res = optimize_pulse( - drift=Hd_lst[0], - ctrls=Hc_lst[0], - initial=init, - target=targ, - num_tslots=time_interval.n_tslots, - evo_time=time_interval.evo_time, - tau=None, # implicitly derived from tslots - amp_lbound=lbound, - amp_ubound=ubound, - fid_err_targ=algorithm_kwargs.get("fid_err_targ", 1e-10), - min_grad=min_g, - max_iter=max_iter, - max_wall_time=algorithm_kwargs.get("max_wall_time", 180), - alg=alg, - optim_method=optim_method, - method_params=minimizer_kwargs, - optim_alg=None, # deprecated - max_metric_corr=None, # deprecated - accuracy_factor=None, # deprecated - alg_params=algorithm_kwargs.get("alg_params", None), - optim_params=algorithm_kwargs.get("optim_params", None), - dyn_type=algorithm_kwargs.get("dyn_type", "GEN_MAT"), - dyn_params=algorithm_kwargs.get("dyn_params", None), - prop_type=algorithm_kwargs.get("prop_type", "DEF"), - prop_params=algorithm_kwargs.get("prop_params", None), - fid_type=algorithm_kwargs.get("fid_type", "DEF"), - fid_params=algorithm_kwargs.get("fid_params", None), - phase_option=None, # deprecated - fid_err_scale_factor=None, # deprecated - tslot_type=algorithm_kwargs.get("tslot_type", "DEF"), - tslot_params=algorithm_kwargs.get("tslot_params", None), - amp_update_mode=None, # deprecated - init_pulse_type=algorithm_kwargs.get("init_pulse_type", "DEF"), - init_pulse_params=algorithm_kwargs.get("init_pulse_params", None), - pulse_scaling=algorithm_kwargs.get("pulse_scaling", 1.0), - pulse_offset=algorithm_kwargs.get("pulse_offset", 0.0), - ramping_pulse_type=algorithm_kwargs.get("ramping_pulse_type", None), - ramping_pulse_params=algorithm_kwargs.get("ramping_pulse_params", None), - log_level=algorithm_kwargs.get("log_level", log_level), - out_file_ext=algorithm_kwargs.get("out_file_ext", None), - gen_stats=algorithm_kwargs.get("gen_stats", False), + if "options" in minimizer_kwargs: + minimizer_kwargs["options"].setdefault( + "maxiter", algorithm_kwargs.get("max_iter", 1000) ) - end_time = time.time() - - # extract runtime information - result.start_local_time = time.strftime( - "%Y-%m-%d %H:%M:%S", time.localtime(start_time) + minimizer_kwargs["options"].setdefault( + "gtol", algorithm_kwargs.get("min_grad", 0.0 if alg == "CRAB" else 1e-10) ) - result.end_local_time = time.strftime( - "%Y-%m-%d %H:%M:%S", time.localtime(end_time) - ) - total_seconds = end_time - start_time - result.iter_seconds = [res.num_iter / total_seconds] * res.num_iter - result.n_iters = res.num_iter - result.message = res.termination_reason - result.final_states = [res.evo_full_final] - result.infidelity = res.fid_err - result.guess_params = res.initial_amps.T - result._guess_controls = res.initial_amps.T - result.optimized_params = res.final_amps.T - result._optimized_controls = res.final_amps.T - - # not present in analytical results - result.stats = res.stats - return result # GRAPE result + else: + minimizer_kwargs["options"] = { + "maxiter": algorithm_kwargs.get("max_iter", 1000), + "gtol": algorithm_kwargs.get("min_grad", 0.0 if alg == "CRAB" else 1e-10), + } + # prepare qtrl optimizers qtrl_optimizers = [] if alg == "CRAB" or alg == "GRAPE": - # Check wether guess referes to amplitudes or parameters - use_as_amps = len(x0[0]) == time_interval.n_tslots - num_coeffs = algorithm_kwargs.get("num_coeffs", None) - fix_frequency = algorithm_kwargs.get("fix_frequency", False) - - # use crab algorithm: max_iter for local minimizer - # TODO: wifi: nelder-mead is called maxiter - minimizer_kwargs.setdefault( - "options", - minimizer_kwargs.get( - "options", {"maxiter": algorithm_kwargs.get("max_iter", 1000)} - ), - ) + if alg == "GRAPE": # algorithm specific kwargs + # only scalar bounds + lbound = lbound[0] + ubound = ubound[0] + use_as_amps = True + alg_params = algorithm_kwargs.get("alg_params", {}) + + elif alg == "CRAB": + minimizer_kwargs.setdefault("method", "Nelder-Mead") # no gradient + # Check wether guess referes to amplitudes (or parameters for CRAB) + use_as_amps = len(x0[0]) == time_interval.n_tslots + num_coeffs = algorithm_kwargs.get("num_coeffs", None) + fix_frequency = algorithm_kwargs.get("fix_frequency", False) + + if num_coeffs is None: + if use_as_amps: + num_coeffs = ( + 2 # default only two sets of fourier expansion coefficients + ) + else: # depending on the number of parameters given + num_coeffs = len(x0[0]) // 2 if fix_frequency else len(x0[0]) // 3 - if num_coeffs is None: - # default only two sets of fourier expansion coefficients - if use_as_amps: - num_coeffs = 2 - else: # depending on the number of parameters given - num_coeffs = len(x0[0]) // 2 if fix_frequency else len(x0[0]) // 3 + alg_params = { + "num_coeffs": num_coeffs, + "init_coeff_scaling": algorithm_kwargs.get("init_coeff_scaling"), + "crab_pulse_params": algorithm_kwargs.get("crab_pulse_params"), + "fix_frequency": fix_frequency, + } + # one optimizer for each objective for i, objective in enumerate(objectives): - alg_params = { + params = { "drift": Hd_lst[i], "ctrls": Hc_lst[i], "initial": init, @@ -284,21 +217,16 @@ def optimize_pulses( "amp_lbound": lbound, "amp_ubound": ubound, "fid_err_targ": algorithm_kwargs.get("fid_err_targ", 1e-10), - "min_grad": 0.0, - "max_iter": max_iter, + "min_grad": minimizer_kwargs["options"]["gtol"], + "max_iter": minimizer_kwargs["options"]["maxiter"], "max_wall_time": algorithm_kwargs.get("max_wall_time", 180), "alg": alg, - "optim_method": optim_method, + "optim_method": algorithm_kwargs.get("optim_method", None), "method_params": minimizer_kwargs, "optim_alg": None, # deprecated "max_metric_corr": None, # deprecated "accuracy_factor": None, # deprecated - "alg_params": { - "num_coeffs": num_coeffs, - "init_coeff_scaling": algorithm_kwargs.get("init_coeff_scaling"), - "crab_pulse_params": algorithm_kwargs.get("crab_pulse_params"), - "fix_frequency": fix_frequency, - }, + "alg_params": alg_params, "optim_params": algorithm_kwargs.get("optim_params", None), "dyn_type": algorithm_kwargs.get("dyn_type", "GEN_MAT"), "dyn_params": algorithm_kwargs.get("dyn_params", None), @@ -329,49 +257,47 @@ def optimize_pulses( "gen_stats": algorithm_kwargs.get("gen_stats", False), } - ###### code from qutip_qtrl.pulseoptim.optimize_pulse ###### - qtrl_optimizer = cpo.create_pulse_optimizer( - drift=alg_params["drift"], - ctrls=alg_params["ctrls"], - initial=alg_params["initial"], - target=alg_params["target"], - num_tslots=alg_params["num_tslots"], - evo_time=alg_params["evo_time"], - tau=alg_params["tau"], - amp_lbound=alg_params["amp_lbound"], - amp_ubound=alg_params["amp_ubound"], - fid_err_targ=alg_params["fid_err_targ"], - min_grad=alg_params["min_grad"], - max_iter=alg_params["max_iter"], - max_wall_time=alg_params["max_wall_time"], - alg=alg_params["alg"], - optim_method=alg_params["optim_method"], - method_params=alg_params["method_params"], - optim_alg=alg_params["optim_alg"], - max_metric_corr=alg_params["max_metric_corr"], - accuracy_factor=alg_params["accuracy_factor"], - alg_params=alg_params["alg_params"], - optim_params=alg_params["optim_params"], - dyn_type=alg_params["dyn_type"], - dyn_params=alg_params["dyn_params"], - prop_type=alg_params["prop_type"], - prop_params=alg_params["prop_params"], - fid_type=alg_params["fid_type"], - fid_params=alg_params["fid_params"], - phase_option=alg_params["phase_option"], - fid_err_scale_factor=alg_params["fid_err_scale_factor"], - tslot_type=alg_params["tslot_type"], - tslot_params=alg_params["tslot_params"], - amp_update_mode=alg_params["amp_update_mode"], - init_pulse_type=alg_params["init_pulse_type"], - init_pulse_params=alg_params["init_pulse_params"], - pulse_scaling=alg_params["pulse_scaling"], - pulse_offset=alg_params["pulse_offset"], - ramping_pulse_type=alg_params["ramping_pulse_type"], - ramping_pulse_params=alg_params["ramping_pulse_params"], - log_level=alg_params["log_level"], - gen_stats=alg_params["gen_stats"], + drift=params["drift"], + ctrls=params["ctrls"], + initial=params["initial"], + target=params["target"], + num_tslots=params["num_tslots"], + evo_time=params["evo_time"], + tau=params["tau"], + amp_lbound=params["amp_lbound"], + amp_ubound=params["amp_ubound"], + fid_err_targ=params["fid_err_targ"], + min_grad=params["min_grad"], + max_iter=params["max_iter"], + max_wall_time=params["max_wall_time"], + alg=params["alg"], + optim_method=params["optim_method"], + method_params=params["method_params"], + optim_alg=params["optim_alg"], + max_metric_corr=params["max_metric_corr"], + accuracy_factor=params["accuracy_factor"], + alg_params=params["alg_params"], + optim_params=params["optim_params"], + dyn_type=params["dyn_type"], + dyn_params=params["dyn_params"], + prop_type=params["prop_type"], + prop_params=params["prop_params"], + fid_type=params["fid_type"], + fid_params=params["fid_params"], + phase_option=params["phase_option"], + fid_err_scale_factor=params["fid_err_scale_factor"], + tslot_type=params["tslot_type"], + tslot_params=params["tslot_params"], + amp_update_mode=params["amp_update_mode"], + init_pulse_type=params["init_pulse_type"], + init_pulse_params=params["init_pulse_params"], + pulse_scaling=params["pulse_scaling"], + pulse_offset=params["pulse_offset"], + ramping_pulse_type=params["ramping_pulse_type"], + ramping_pulse_params=params["ramping_pulse_params"], + log_level=params["log_level"], + gen_stats=params["gen_stats"], ) dyn = qtrl_optimizer.dynamics dyn.init_timeslots() @@ -417,7 +343,7 @@ def optimize_pulses( qtrl_optimizers.append(qtrl_optimizer) - return opt_pulses( + return global_local_optimization( objectives, pulse_options, time_interval, diff --git a/src/qutip_qoc/result.py b/src/qutip_qoc/result.py index 66d5ad7..31017cb 100644 --- a/src/qutip_qoc/result.py +++ b/src/qutip_qoc/result.py @@ -83,7 +83,7 @@ def __init__( optimized_params=None, infidelity=np.inf, var_time=False, - qtrl_optimizer=None, + qtrl_optimizers=None, ): self.time_interval = time_interval self.objectives = objectives @@ -102,7 +102,7 @@ def __init__( self.final_states = final_states self.infidelity = infidelity self.var_time = var_time - self.qtrl_optimizer = qtrl_optimizer + self.qtrl_optimizers = qtrl_optimizers def __str__(self): return textwrap.dedent( @@ -185,7 +185,7 @@ def optimized_controls(self): if len(xf) == len(self.time_interval.tslots): cf = xf else: # parameterized CRAB - pgen = self.qtrl_optimizer[0].pulse_generator[j] + pgen = self.qtrl_optimizers[0].pulse_generator[j] pgen.set_optim_var_vals(np.array(self.optimized_params[j])) cf = pgen.gen_pulse() opt_ctrl.append(cf) @@ -196,8 +196,8 @@ def optimized_controls(self): @property def guess_controls(self): if self._guess_controls is None: - if self.qtrl_optimizer: - qtrl_res = self.qtrl_optimizer[0]._create_result() + if self.qtrl_optimizers: + qtrl_res = self.qtrl_optimizers[0]._create_result() gss_ctrl = qtrl_res.initial_amps.T else: gss_ctrl = [] @@ -220,7 +220,7 @@ def guess_controls(self): if len(xi) == len(self.time_interval.tslots): c0 = xi else: # parameterized CRAB - pgen = self.qtrl_optimizer[0].pulse_generator[j] + pgen = self.qtrl_optimizers[0].pulse_generator[j] pgen.set_optim_var_vals(np.array(self.guess_params[j])) c0 = pgen.gen_pulse() gss_ctrl.append(c0) @@ -261,7 +261,7 @@ def final_states(self): evo_time = self.time_interval.evo_time para_keys = [] - if not self.qtrl_optimizer: # GOAT/JOPT + if not self.qtrl_optimizers: # GOAT/JOPT # extract parameter names from control functions f(t, para_key) c_sigs = [signature(Hc[1]) for Hc in self.objectives[0].H[1:]] c_keys = [sig.parameters.keys() for sig in c_sigs]