Source code for finesse.analysis.actions.locks

"""Lock Actions."""

import logging

import numpy as np

import finesse.config

from ...env import is_interactive, warn
from ...parameter import Parameter, deref
from ...solutions import BaseSolution
from . import elements_to_name
from .base import Action, convert_str_to_parameter
from .random import Change
from .sensing import SensingMatrixDC, SensingMatrixSolution
from ...simulations.sparse.simulation import SparseMatrixSimulation
from tqdm.auto import tqdm
from finesse.utilities import OrderedSet

LOGGER = logging.getLogger(__name__)


[docs]class RunLocksSolution(BaseSolution): """Solution from applying the :class:`.RunLocks` action. Attributes ---------- iters : int Number of steps lock has required max_iterations : int Maximum number of iterations this lock can do error_signals : array_like error signals during locking steps, shape [num_locks, max_iterations] control_signals : array_like Control signals during locking steps, shape [num_locks, max_iterations] lock_names : tuple[str] Names of locks being controlled, shape [num_locks] feedback_names : tuple[str] Names of feedback for each lock, shape [num_locks] feedback_names : tuple[str] Names of error signals for each lock, shape [num_locks] final : arrary_like Final control signals, shape [num_locks] sensing_matrix : SensingMatrixSolution, optional The sensing matrix used when running the locks with Newton's method. """
[docs] def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.iters = 0 self.error_signals = None self.control_signals = None self.lock_names = () self.feedback_names = () self.error_signal_names = () self.final = None self.max_iterations = 0 self.num_locks = 0 self.sensing_matrix = None
[docs] def plot_error_signals(self, ax=None): """Plots how the error signals vary during this lock attempt. Parameters ---------- ax : Matplotlib.Axes, optional Axes to plot on, if no current axis is set then a new one is generated """ import matplotlib.pyplot as plt if ax is not None: plt.sca(ax) plt.semilogy(abs(self.error_signals.T[: self.iters, :])) plt.legend( tuple(f"{a}:{b}" for a, b in zip(self.lock_names, self.error_signal_names)) ) plt.xlabel("steps") plt.ylabel("Error signal [arb]")
plot = plot_error_signals # Default plot option
[docs] def plot_control_signals(self, ax=None): """Plots how the controls signals vary during this lock attempt. If `0` gaps will be shown when no change has been made to that degree of freedom for that step (As it was within the locks accuracy setting). Parameters ---------- ax : Matplotlib.Axes, optional Axes to plot on, if no current axis is set then a new one is generated """ import matplotlib.pyplot as plt if ax is not None: plt.sca(ax) plt.semilogy(abs(self.control_signals.T[: self.iters, :])) plt.legend( tuple(f"{a}:{b}" for a, b in zip(self.lock_names, self.feedback_names)) ) plt.xlabel("steps") plt.ylabel("Control signal [arb]")
# IMPORTANT: renaming this class impacts the katscript spec and should be avoided!
[docs]class RunLocks(Action): """An action that iteratively moves the system to lock. Currently, lock error signals must be readouts, not detectors, for use in this action. Parameters ---------- *locks : list, optional A list of locks to use in each RunLocks step. If not provided, all locks in model are used. method : str, either "newton" or "proportional" Which method to use in the locking iterations. scale_factor : float Factor by which to multiply all DOF changes. Should be set below 1 if it is desired to minimize overshooting. sensing_matrix : SensingMatrixSolution or None Sensing matrix of gains used in locking, of the type that would be returned by state.apply(SensingMatrixDC(lock_dof_names, readout_names) If None, the sensing matrix is recalculated. Recommended to be None except when locking multiple times in a row, e.g. with DragLocks. max_iterations : int The maximum number of locking steps in each execution of RunLocks. display_progress : boolean When true, displays the status of the error signals during locking iterations. optimize_phase : boolean Deprecated: Use an action like OptimiseRFReadoutPhaseDC instead. d_dof_phase : float Step size to use when optimizing the demodulation phase for each error signal/DOF pair. set_gains : boolean Only applies if method is "proportional". If true, sets the gains for each error signal/DOF pair. If false, uses pre-set gains. d_dof_gain : float Step size to use when calculating the gain for every pair of error signals and DOFs. exception_on_fail : boolean When true, raise exception if maximum iterations are surpassed. no_warning : boolean When true, don't even raise a warning if maximum iterations are reached. Recommended to be false unless repeatedly testing locking. pre_step : :class:`.Action` Action to apply on each step of the lock show_progress_bar : boolean Will enable the progress bar when true. name : str Name of the action. """
[docs] def __init__( self, *locks, method="proportional", scale_factor=1, sensing_matrix=None, max_iterations=10000, display_progress=False, optimize_phase=None, d_dof_phase=1e-9, set_gains=True, d_dof_gain=1e-9, exception_on_fail=True, no_warning=False, pre_step=None, show_progress_bar=None, name="run locks", ): super().__init__(name) self.locks = tuple((l if isinstance(l, str) else l.name) for l in locks) self.max_iterations = max_iterations self.method = method self.scale_factor = scale_factor self.sensing_matrix = sensing_matrix self.display_progress = display_progress self.optimize_phase = optimize_phase if optimize_phase is not None: finesse.utilities.misc.deprecation_warning( "DragLocks: `optimize_phase` is deprecated, consider using the :class:`.OptimiseRFReadoutPhaseDC` action before this one instead.", "3.0.0", ) self.d_dof_phase = d_dof_phase self.set_gains = set_gains self.d_dof_gain = d_dof_gain self.exception_on_fail = exception_on_fail self.no_warning = no_warning self.pre_step = pre_step # use local flag if provided, otherwise default to global setting if show_progress_bar is not None: self.show_progress_bar = show_progress_bar else: self.show_progress_bar = finesse.config.show_progress_bars self.pbar = None
[docs] def init_pbar(self, locks): # define template for lock status display lock_format = " ".join( [f"{lock.name} {{postfix[0][{lock.name}]}}" for lock in locks] ) # create the progress bar self.pbar = tqdm( range(self.max_iterations), initial=self.max_iterations, desc="", bar_format=f"{lock_format} |{{bar}}| {{n_fmt}}/{{total_fmt}}", colour="red", dynamic_ncols=True, postfix=[dict((lock.name, "✗") for lock in locks), len(locks)], disable=not self.show_progress_bar, )
[docs] def update_pbar(self): # decrement the counter self.pbar.update(-1)
[docs] def update_pbar_lock(self, lock_name, is_locked): # update the lock's status if self.show_progress_bar: self.pbar.postfix[0][lock_name] = "✔" if is_locked else "✘"
[docs] def complete_pbar(self): # runs when the progress bar is considered complete self.pbar.colour = "green" self.pbar.refresh()
def _do(self, state): # we need a carrier signal simulation to run the locks if state.sim is None: raise Exception("Simulation has not been built") if not isinstance(state.sim, SparseMatrixSimulation): raise NotImplementedError() # gather locks from the model if len(self.locks) > 0: # use specified locks if they are enabled locks = tuple(state.model.elements[name] for name in self.locks) else: # otherwise use all enabled locks locks = tuple(lck for lck in state.model.locks if lck.enabled) # collect all lock related workspaces dws = tuple( next( filter( lambda x: x.oinfo.name == lock.error_signal.name, OrderedSet( # workspaces can be in both lists (*state.sim.readout_workspaces, *state.sim.detector_workspaces) ), ), None, ) for lock in locks ) # Store initial parameters in case of failure so we can reset the model initial_feedback = tuple(float(lock.feedback) for lock in locks) # initialize the solution sol = RunLocksSolution(self.name) sol.max_iterations = self.max_iterations sol.num_locks = len(locks) sol.iters = -1 sol.error_signals = np.zeros((len(locks), self.max_iterations + 1)) sol.control_signals = np.zeros((len(locks), self.max_iterations + 1)) sol.lock_names = tuple(lock.name for lock in locks) sol.feedback_names = tuple(deref(lock.feedback).full_name for lock in locks) sol.error_signal_names = tuple(deref(lock.error_signal).name for lock in locks) # set up the progress bar using all enabled locks self.init_pbar(locks) if self.display_progress: GREEN = "\033[92m" if is_interactive() else "" RED = "\033[91m" if is_interactive() else "" # BOLD = "\033[1m" not used END = "\033[0m" if is_interactive() else "" print("Error Signal Residuals at Each Iteration (W):") print(format("", "23s"), end="") for lock in locks: print(format(lock.name, "^15s"), end="") # ---------------------------------------------------------------------- # Proportional method # ---------------------------------------------------------------------- if self.method in "proportional": # use the sensing matrix to set the gains? # TODO: allow this method to set gains from the sensing matrix # if self.set_gains: # for idx, _ in enumerate(lock_dof_names): # if "_Q" in err_sig_names[idx]: # locks[idx].gain = -1 / sensing_matrix.out[idx, idx].imag # else: # locks[idx].gain = -1 / sensing_matrix.out[idx, idx].real # compute as needed or until max iterations have been reached recompute = True while recompute and sol.iters < self.max_iterations: sol.iters += 1 if self.display_progress: print( format("\nIteration Number ", "<20s") + format(sol.iters, "<3d"), end="", ) # run the pre-step action if self.pre_step: state.apply(self.pre_step) # calculate the readout values state.sim.run_carrier() # compute as needed or until max iterations have been reached recompute = False for i in range(len(locks)): # read the error err = dws[i].get_output() - locks[i].offset sol.error_signals[i, sol.iters] = err # recompute if the error is too large acc = locks[i].accuracy if abs(err) >= acc: # adjust the feedback feedback = float(locks[i].gain) * err * self.scale_factor deref(locks[i].feedback).value += feedback # store it sol.control_signals[i, sol.iters] = feedback # and go again recompute = True is_locked = abs(err) < acc if self.display_progress: str_color = GREEN if is_locked else RED print(str_color + format(err, "^ 15.2e") + END, end="") # update the lock status self.update_pbar_lock(locks[i].name, is_locked) # update the bar status self.update_pbar() # ---------------------------------------------------------------------- # Newton method # ---------------------------------------------------------------------- elif self.method == "newton": # this method requires the use of readouts # TODO: make sure this can only be done with readouts (not pds) err_sigs = [lock.error_signal for lock in locks] err_sig_names = [sig.name for sig in err_sigs] readout_names = [sig.readout.name for sig in err_sigs] # fails if pd lock_dof_names = [deref(lock.feedback).owner.name for lock in locks] if self.display_progress: print("\n" + format("", "23s"), end="") for idx in range(len(locks)): print(format(err_sig_names[idx] + "1", "^15s"), end="") # a sensing matrix is required if self.sensing_matrix is not None: if isinstance(self.sensing_matrix, SensingMatrixSolution): sensing_matrix = self.sensing_matrix else: raise Exception( "Locks failed: invalid type of sensing matrix specified" ) else: sensing_matrix = state.apply( SensingMatrixDC(lock_dof_names, readout_names) ) # store the sensing maxtrix sol.sensing_matrix = sensing_matrix # Matrix of gains only for readout phases that are actually used in # locks. Also transposes the sensing matrix, so that rows rather # than columns correspond to error signals. N = len(locks) gain_matrix = np.zeros((N, N)) for dof_idx, _ in enumerate(lock_dof_names): for rd_idx, _ in enumerate(readout_names): # get the sensing matrix value val = sensing_matrix.out[dof_idx, rd_idx] # take imag or real depending on the type of signal if "_Q" in err_sig_names[rd_idx]: gain_matrix[rd_idx, dof_idx] = val.imag else: gain_matrix[rd_idx, dof_idx] = val.real # Evaluate the inverse of the gain matrix/Jacobian. Assuming # that we stay in the linear region for all DOFs/readouts, we evaluate # the inverted Jacobian only once but use it in all iterations. jacobian_inv = np.linalg.inv(gain_matrix) * self.scale_factor # compute as needed or until max iterations have been reached recompute = True while recompute and sol.iters < self.max_iterations: # set up the run sol.iters += 1 recompute = False if self.display_progress: print() print( format("Iteration Number ", "<20s") + format(sol.iters, "<3d"), end="", ) # run the pre-step action if self.pre_step: state.apply(self.pre_step) # recalculate the readout values state.sim.run_carrier() # gather the accuracy from the locks and error from the readouts acc_vect = np.array([lock.accuracy for lock in locks]) err_vect = np.array( [dws[i].get_output() - locks[i].offset for i in range(N)] ) # calculate the new feedbacks using the inverted jacobian feedback_vect = -1 * np.matmul(jacobian_inv, err_vect) # for each lock results = [None] * N for i in range(N): # store the error sol.error_signals[i, sol.iters] = err_vect[i] # if any error is too high, we need to recompute if any(np.greater(abs(err_vect), acc_vect)): # store the feedback increment sol.control_signals[i, sol.iters] = feedback_vect[i] # adjust the feedback deref(locks[i].feedback).value += feedback_vect[i] # let's do it again recompute = True results[i] = f"{locks[i].name} {err_vect[i]:.2g}" is_locked = abs(err_vect[i]) < acc_vect[i] if self.display_progress: str_color = GREEN if is_locked else RED print(str_color + format(err_vect[i], "^ 15.2e") + END, end="") # update the lock status self.update_pbar_lock(locks[i].name, is_locked) # update the bar status self.update_pbar() # method not found! else: raise Exception("Locks failed: invalid method provided") # if the locks still need to be recomputed then we've failed... if recompute: # reset the locks for lock, value in zip(locks, initial_feedback): deref(lock.feedback).value = value # throw an exception? if self.exception_on_fail: raise finesse.exceptions.LostLock( "Locks failed: max iterations reached" ) # display a warning? if not self.no_warning: warn( "Locks failed to converge, try increasing maximum iterations, check gains/error signals. See RunLocks exception_on_fail=False to inspect solution for locking problems." ) else: # locks have successfully locked self.complete_pbar() # store the final feedback values in the solution sol.final = np.array( tuple(deref(lock.feedback).value for lock in locks), dtype=float ) return sol def _requests(self, model, memo, first=True): # gather locks from the model if len(self.locks) > 0: # use specified locks if they are not enabled locks = tuple(model.elements[name] for name in self.locks) else: # otherwise use all enabled locks locks = tuple(lck for lck in model.locks if lck.enabled) for lock in locks: # the lock feedback values will be changing memo["changing_parameters"].append(deref(lock.feedback).full_name) if self.pre_step: self.pre_step._requests(model, memo)
# IMPORTANT: renaming this class impacts the katscript spec and should be avoided!
[docs]class DragLocks(Action): """An action that incrementally changes model parameter values, reaching lock at each step, until lock is reached at the desired final parameter values. Parameters ---------- *locks : list, optional A list of locks to use in each RunLocks step. Acts like *locks parameter in RunLocks: if not provided, all locks in model are used. parameters : list A list of strings. Each element should correspond to a parameter in the model. stop_points : list The final parameter values that locks move towards incrementally. relative : boolean If true, stop_points are relative to the initial parameter values. max_recursions : int The number of times that the step size is allowed to decreased by a factor of ten when locks fail. method : str, either "newton" or "proportional" The method to use in each locking step. scale_factor : float Factor by which to multiply all DOF changes. Should be set below 1 if it is desired to minimize overshooting. num_steps : int Number of steps to calculate, starting at the initial point and ending at the stop point. never_optimize_phase : boolean Deprecated: When true, never optimize readout phases. When false, phases will be optimized anytime the previous step required more than 10 iterations. exception_on_fail : boolean When true, raise exception if max_recursions is surpassed. max_iterations : int The maximum number of locking steps in each execution of RunLocks. If surpassed, step size is decreased. display_progress : boolean When true, displays the status of the lock dragging. name : str Name of the action. """ def __init__( self, *locks, parameters, stop_points, relative=False, method="proportional", scale_factor=1, num_steps=11, never_optimize_phase=None, exception_on_fail=True, max_recursions=5, max_iterations=1000, display_progress=False, show_progress_bar=False, name="drag locks", ): super().__init__(name) self.locks = tuple((l if isinstance(l, str) else l.name) for l in locks) self.parameters = parameters self.stop_points = np.array(stop_points) if len(self.parameters) != len(self.stop_points): raise ValueError("Unequal number of parameters and stopping points") self.relative = relative self.max_recursions = max_recursions self.method = method self.scale_factor = scale_factor self.num_steps = num_steps self.exception_on_fail = exception_on_fail self.max_iterations = max_iterations self.show_progress_bar = show_progress_bar self.display_progress = display_progress if never_optimize_phase is not None: finesse.utilities.misc.deprecation_warning( "DragLocks: `never_optimize_phase` is deprecated, consider using the :class:`.OptimiseRFReadoutPhaseDC` action before this one instead.", "3.0.0", ) def _do(self, state): def TryLocking(state, steps, recursion_num=0): sensing_matrix = None for step_ind, step_vals in enumerate(steps): # Change each parameter to its value at this step. for p_ind, param_val in enumerate(step_vals): state.apply(Change({self.parameters[p_ind]: param_val})) # Run locks at this step. try: step_vals_str = str([format(val, "4.3e") for val in step_vals]) if self.display_progress: print( "\t" * recursion_num + f"Step {step_ind:2d} of {len(steps) - 1}: ", end="", ) sol = state.apply( RunLocks( method=self.method, scale_factor=self.scale_factor, sensing_matrix=sensing_matrix, exception_on_fail=True, max_iterations=self.max_iterations, display_progress=False, show_progress_bar=self.show_progress_bar, ) ) # Print status of locking steps. if self.display_progress: print( "Reached lock with", self.parameters, "= " + step_vals_str + " in " + str(sol.iters) + " iterations.", ) # If the previous step converged very quickly, don't bother # optimizing phases or recalculating the sensing matrix at # the next step. if sol.iters <= 10: sensing_matrix = sol.sensing_matrix else: sensing_matrix = None if self.display_progress: print( "\t" * recursion_num + "Step required more than 10 iterations." + " Recalculating sensing matrix in next step." ) except Exception: recursion_num += 1 if self.display_progress: print( "Failed to lock with", self.parameters, "= " + step_vals_str + ". Decreasing step size.", ) if recursion_num >= self.max_recursions: raise Exception("Maximum recursion level exceeded.") if step_ind == 0: new_step_vals = np.linspace( steps[step_ind], steps[step_ind + 1], self.num_steps ) else: new_step_vals = np.linspace( steps[step_ind - 1], steps[step_ind], self.num_steps ) TryLocking(state, new_step_vals, recursion_num=recursion_num) recursion_num -= 1 return sol # Find the model parameters corresponding to the strings provided p = [convert_str_to_parameter(state.model, param) for param in self.parameters] p_vals = np.array([param.value for param in p]) # The parameter values that will be stepped through and locked to. if not self.relative: step_vals_list = np.linspace(p_vals, self.stop_points, self.num_steps) else: step_vals_list = np.linspace( p_vals, p_vals + self.stop_points, self.num_steps ) sol = TryLocking(state, step_vals_list) return sol def _requests(self, model, memo, first=True): for param in self.parameters: p = convert_str_to_parameter(model, param) if isinstance(p, Parameter): memo["changing_parameters"].append(param) if len(self.locks) == 0: # If none given lock everything for lock in model.locks: memo["changing_parameters"].append(deref(lock.feedback).full_name) rd_name = lock.error_signal.name if "_DC" not in rd_name: memo["changing_parameters"].append( lock.error_signal.readout.name + ".phase" ) else: for name in self.locks: if name not in model.elements: raise Exception(f"Model {model} does not have a lock called {name}") memo["changing_parameters"].append( deref(model.elements[name].feedback).full_name ) rd_name = model.elements[name].error_signal.name if "_DC" not in rd_name: memo["changing_parameters"].append( model.elements[name].error_signal.readout.name + ".phase" )
# IMPORTANT: renaming this class impacts the katscript spec and should be avoided!
[docs]class SetLockGains(Action): """An action that computes the optimal lock gains using the sensing matrix found with :class:`.SensingMatrixDC`. This action computes the error signal gradient for each lock with respect to its drive and sets the gain as `-gain_scale/sensing`. Parameters ---------- *locks : list, optional A list of locks for which to set the gain. If none provided, all enabled locks in model are used. Disabled locks that are explicitly listed will have their gains set. d_dof_gain : float, optional Step size to use when calculating the gain for each error signal/DOF pair. gain_scale : float, optional Extra gain scaling factor applied to the gain calculation: `-gain_scale/sensing` In multiple lock models where the locks are cross coupled using a `gain_scale` < 1 can improve the stability of the locking algorithm to stop excessively large steps. optimize_phase : bool, optional, Deprecated feature: Use :class:`.OptimiseRFReadoutPhaseDC` instead name : str Name of the action. """ def __init__( self, *locks, d_dof_gain=1e-10, gain_scale=1, name="set gains", optimize_phase=None, ): super().__init__(name) self.locks = elements_to_name(locks) self.d_dof_gain = d_dof_gain self.gain_scale = gain_scale self.sensing_action = None self.optimize_phase = optimize_phase if optimize_phase is not None: finesse.utilities.misc.deprecation_warning( "SetLockGains: `optimize_phase` is deprecated, consider using the :class:`.OptimiseRFReadoutPhaseDC` action instead.", "3.0.0", ) def _do(self, state): if state.sim is None: raise Exception("Simulation has not been built") if not isinstance(state.sim, SparseMatrixSimulation): raise NotImplementedError() if len(self.locks) == 0: locks = tuple(lck for lck in state.model.locks) else: locks = tuple( state.model.elements[lock.name] for lock in self.locks if lock.enabled ) err_sigs = [lck.error_signal for lck in locks] err_sig_names = [sig.name for sig in err_sigs] gain_matrix = state.apply(self.sensing_action) for idx, lock in enumerate(locks): err_sig = err_sig_names[idx] val = gain_matrix.out[idx, idx] if "_Q" in err_sig: gain = val.imag else: gain = val.real if gain == 0: raise ZeroDivisionError(f"{lock.name} found a gain of zero") lock.gain = -self.gain_scale / gain def _requests(self, model, memo, first=True): if len(self.locks) == 0: self.locks = tuple(lck for lck in model.locks if lck.enabled) else: self.locks = tuple(model.get(name) for name in self.locks) self.sensing_action = SensingMatrixDC( [deref(lock.feedback).owner.name for lock in self.locks], [lock.error_signal.readout.name for lock in self.locks], d_dof=self.d_dof_gain, ) self.sensing_action._requests(model, memo) for lock in self.locks: memo["changing_parameters"].append(lock.gain.full_name)