Source code for krotov.optimize

import copy
import inspect
import logging
import time
from functools import partial

import numpy as np
from qutip import Qobj
from qutip.parallel import serial_map

from .conversions import (
from .info_hooks import chain
from .mu import derivative_wrt_pulse
from .propagators import Propagator, expm
from .result import Result
from .second_order import _overlap
from .shapes import one_shape, zero_shape

__all__ = ['optimize_pulses']

[docs]def optimize_pulses( objectives, pulse_options, tlist, *, propagator, chi_constructor, mu=None, sigma=None, iter_start=0, iter_stop=5000, check_convergence=None, info_hook=None, modify_params_after_iter=None, storage='array', parallel_map=None, store_all_pulses=False, continue_from=None, skip_initial_forward_propagation=False, norm=None, overlap=None ): r"""Use Krotov's method to optimize towards the given `objectives`. Optimize all time-dependent controls found in the Hamiltonians or Liouvillians of the given `objectives`. Args: objectives (list[Objective]): List of objectives pulse_options (dict): Mapping of time-dependent controls found in the Hamiltonians of the objectives to a dictionary of options for that control. There must be options given for *every* control. As numpy arrays are unhashable and thus cannot be used as dict keys, the options for a control that is an array must be set using the key ``id(control)`` (see the example below). The options of any particular control *must* contain the following keys: * ``'lambda_a'``: the Krotov step size (float value). This governs the overall magnitude of the pulse update. Large values result in small updates. Small values may lead to sharp spikes and numerical instability. * ``'update_shape'`` : Function S(t) in the range [0, 1] that scales the pulse update for the pulse value at t. This can be used to ensure boundary conditions (S(0) = S(T) = 0), and enforce smooth switch-on and switch-off. This can be a callable that takes a single argument `t`; or the values 1 or 0 for a constant update-shape. The value 0 disables the optimization of that particular control. In addition, the following keys *may* occur: * ``'args'``: If the control is a callable with arguments ``(t, args)`` (as required by QuTiP), a dict of argument values to pass as `args`. If ``'args'`` is not specified via the `pulse_options`, controls will be discretized using the default ``args=None``. For example, for `objectives` that contain a Hamiltonian of the form ``[H0, [H1, u], [H2, g]]``, where ``H0``, ``H1``, and ``H2`` are :class:`~qutip.Qobj` instances, ``u`` is a numpy array .. doctest:: >>> u = numpy.zeros(1000) and ``g`` is a control function .. doctest:: >>> def g(t, args): ... E0 = args.get('E0', 0.0) ... return E0 then a possible value for `pulse_options` would look like this: .. doctest:: >>> from krotov.shapes import flattop >>> from functools import partial >>> pulse_options = { ... id(u): {'lambda_a': 1.0, 'update_shape': 1}, ... g: dict( ... lambda_a=1.0, ... update_shape=partial( ... flattop, t_start=0, t_stop=10, t_rise=1.5 ... ), ... args=dict(E0=1.0) ... ) ... } The use of :class:`dict` and the ``{...}`` syntax are completely equivalent, but :class:`dict` is better for nested indentation. tlist (numpy.ndarray): Array of time grid values, cf. :func:`~qutip.mesolve.mesolve` propagator (callable or list[callable]): Function that propagates the state backward or forwards in time by a single time step, between two points in `tlist`. Alternatively, a list of functions, one for each objective. If the propagator is stateful, it should be an instance of :class:`krotov.propagators.Propagator`. See :mod:`krotov.propagators` for details. chi_constructor (callable): Function that calculates the boundary condition for the backward propagation. This is where the final-time functional (indirectly) enters the optimization. See :mod:`krotov.functionals` for details. mu (None or callable): Function that calculates the derivative $\frac{\partial H}{\partial\epsilon}$ for an equation of motion $\dot{\phi}(t) = -i H[\phi(t)]$ of an abstract operator $H$ and an abstract state $\phi$. If None, defaults to :func:``, which covers the standard Schrödinger and master equations. See :mod:`` for a full explanation of the role of `mu` in the optimization, and the required function signature. sigma (None or krotov.second_order.Sigma): Function (instance of a :class:`.Sigma` subclass) that calculates the second-order contribution. If None, the first-order Krotov method is used. iter_start (int): The formal iteration number at which to start the optimization iter_stop (int): The iteration number after which to end the optimization, whether or not convergence has been reached check_convergence (None or callable): Function that determines whether the optimization has converged. If None, the optimization will only end when `iter_stop` is reached. See :mod:`krotov.convergence` for details. info_hook (None or callable): Function that is called after each iteration of the optimization, for the purpose of analysis. Any value returned by `info_hook` (e.g. an evaluated functional :math:`J_T`) will be stored, for each iteration, in the `info_vals` attribute of the returned :class:`.Result`. The `info_hook` must have the same signature as :func:`krotov.info_hooks.print_debug_information`. It should not modify its arguments in any way, except for `shared_data`. modify_params_after_iter (None or callable): Function that is called after each iteration, which may modify its arguments for certain advanced use cases, such as dynamically adjusting `lambda_vals`, or applying spectral filters to the `optimized_pulses`. It has the same interface as `info_hook` but should not return anything. The `modify_params_after_iter` function is called immediately before `info_hook`, and can transfer arbitrary data to any subsequent `info_hook` via the `shared_data` argument. storage (callable): Storage constructor for the storage of propagated states. Must accept an integer parameter `N` and return an empty array-like container of length `N`. The default value 'array' is equivalent to ``functools.partial(numpy.empty, dtype=object)``. parallel_map (callable or tuple or None): Parallel function evaluator. If given as a callable, the argument must have the same specification as :func:`qutip.parallel.serial_map`. A value of None is the same as passing :func:`qutip.parallel.serial_map`. If given as a tuple, that tuple must contain three callables, each of which has the same specification as :func:`qutip.parallel.serial_map`. These three callables are used to parallelize (1) the initial forward-propagation, (2) the backward-propagation under the guess pulses, and (3) the forward-propagation by a single time step under the optimized pulses. See :mod:`krotov.parallelization` for details. store_all_pulses (bool): Whether or not to store the optimized pulses from *all* iterations in :class:`.Result`. continue_from (None or Result): If given, continue an optimization from a previous :class:`.Result`. The result must have identical `objectives`. skip_initial_forward_propagation (bool): If given as `True` together with `continue_from`, skip the initial forward propagation ("zeroth iteration"), and take the forward-propagated states from :attr:`.Result.states` instead. norm (callable or None): A single-argument function to calculate the norm of states. If None, delegate to the :meth:`~qutip.Qobj.norm` method of the states. overlap (callable or None): A two-argument function to calculate the complex overlap of two states. If None, delegate to :meth:`qutip.Qobj.overlap` for Hilbert space states and to the Hilbert-Schmidt norm $\tr[\rho_1^\dagger \rho2]$ for density matrices or operators. Returns: Result: The result of the optimization. Raises: ValueError: If any controls are not real-valued, or if any update shape is not a real-valued function in the range [0, 1]; if using `continue_from` with a :class:`.Result` with differing `objectives`; if there are any required keys missing in `pulse_options`. """ logger = logging.getLogger('krotov') # Initialization"Initializing optimization with Krotov's method") if mu is None: mu = derivative_wrt_pulse second_order = sigma is not None if norm is None: norm = lambda state: state.norm() if overlap is None: overlap = _overlap if modify_params_after_iter is not None: # From a technical perspective, the `modify_params_after_iter` is # really just another info_hook, the only difference is the # convention that info_hooks shouldn't modify the parameters. if info_hook is None: info_hook = modify_params_after_iter else: info_hook = chain(modify_params_after_iter, info_hook) if isinstance(propagator, list): propagators = propagator assert len(propagators) == len(objectives) else: propagators = [copy.deepcopy(propagator) for _ in objectives] # copy.deepcopy will only do something on Propagator objects. For # functions (even with closures), it just returns the same function. _check_propagators_interface(propagators, logger) adjoint_objectives = [obj.adjoint() for obj in objectives] if storage == 'array': storage = partial(np.empty, dtype=object) if parallel_map is None: parallel_map = serial_map if not isinstance(parallel_map, (tuple, list)): parallel_map = (parallel_map, parallel_map, parallel_map) ( guess_controls, # "controls": sampled on the time grid guess_pulses, # "pulses": sampled on the time grid intervals pulses_mapping, # keep track of where to plug in pulse values lambda_vals, # Krotov step width λₐ, for each control shape_arrays, # update shape S(t), per control, sampled on intervals ) = _initialize_krotov_controls(objectives, pulse_options, tlist) if continue_from is not None: guess_controls, guess_pulses = _restore_from_previous_result( continue_from, objectives, tlist, store_all_pulses ) g_a_integrals = np.zeros(len(guess_pulses)) # ∫gₐ(t)dt is a very useful measure of whether λₐ is too small (large # ∫gₐ(t)dt, relative to the pulse amplitude), and whether we're approaching # convergence ("speeding up" for increasing values, "slowing down" for # decreasing values) if continue_from is None: result = Result() result.start_local_time = time.localtime() else: result = copy.deepcopy(continue_from) # Initial forward-propagation tic = time.time() if skip_initial_forward_propagation: forward_states = _skip_initial_forward_propagation( objectives, continue_from, sigma, logger ) else: forward_states = parallel_map[0]( _forward_propagation, list(range(len(objectives))), ( objectives, guess_pulses, pulses_mapping, tlist, propagators, storage, ), ) toc = time.time() fw_states_T = [states[-1] for states in forward_states] tau_vals = np.array( [ overlap(, state_T) for (state_T, obj) in zip(fw_states_T, objectives) ] ) if second_order: forward_states0 = forward_states # ∀t: Δϕ=0, for iteration 0 else: # the forward-propagated states only need to be stored for the second # order update forward_states0 = forward_states = None info = None optimized_pulses = copy.deepcopy(guess_pulses) info_hook_static_args = dict( # these do no change between iterations (although # `modify_params_after_iter` may modify any of these, to # the extent that they're mutable) objectives=objectives, adjoint_objectives=adjoint_objectives, lambda_vals=lambda_vals, shape_arrays=shape_arrays, tlist=tlist, propagator=propagator, chi_constructor=chi_constructor, mu=mu, sigma=sigma, iter_start=iter_start, iter_stop=iter_stop, ) if info_hook is not None: info = info_hook( backward_states=None, forward_states=forward_states, forward_states0=forward_states0, guess_pulses=guess_pulses, optimized_pulses=optimized_pulses, g_a_integrals=g_a_integrals, fw_states_T=fw_states_T, tau_vals=tau_vals, start_time=tic, stop_time=toc, iteration=0, info_vals=[], shared_data={}, **info_hook_static_args ) # Initialize Result object result.tlist = tlist result.objectives = objectives result.guess_controls = guess_controls result.optimized_controls = optimized_pulses result.controls_mapping = pulses_mapping if continue_from is None: # we only store information about the "0" iteration if we're starting a # new optimization if info is not None: result.info_vals.append(info) result.iters.append(0) result.iter_seconds.append(int(toc - tic)) if not np.all(tau_vals == None): # noqa result.tau_vals.append(tau_vals) if store_all_pulses: result.all_pulses.append(guess_pulses) else: iter_start = continue_from.iters[-1] "Continuing from previous result, with iteration %d", iter_start + 1, ) result.states = fw_states_T # Main optimization loop for krotov_iteration in range(iter_start + 1, iter_stop + 1):"Started Krotov iteration %d", krotov_iteration) tic = time.time() # Boundary condition for the backward propagation # -- this is where the functional enters the optimization. # `fw_states_T` are the states forward-propagated under the guess pulse # of the current iteration, which is the optimized pulse of the # previous iteration. This is how we get the `fw_states_T` here: they # are left over from the forward-propagation in the previous iteration. chi_states = chi_constructor( fw_states_T=fw_states_T, objectives=objectives, tau_vals=tau_vals ) chi_norms = [norm(chi) for chi in chi_states] # normalizing χ improves numerical stability; the norm then has to be # taken into account when calculating Δϵ chi_states = [chi / nrm for (chi, nrm) in zip(chi_states, chi_norms)] # Backward propagation backward_states = parallel_map[1]( _backward_propagation, list(range(len(objectives))), ( chi_states, adjoint_objectives, guess_pulses, pulses_mapping, tlist, propagators, storage, ), ) # Forward propagation and pulse update"Started forward propagation/pulse update") if second_order: forward_states = [ storage(len(tlist)) for _ in range(len(objectives)) ] g_a_integrals[:] = 0.0 if second_order: # In the update for the pulses in the first time interval, we use # the states at t=0. Hence, Δϕ(t=0) = 0 delta_phis = [ Qobj(np.zeros(shape=chi_states[k].shape)) for k in range(len(objectives)) ] if second_order: for i_obj in range(len(objectives)): forward_states[i_obj][0] = objectives[i_obj].initial_state delta_eps = [ np.zeros(len(tlist) - 1, dtype=np.complex128) for _ in guess_pulses ] optimized_pulses = copy.deepcopy(guess_pulses) fw_states = [obj.initial_state for obj in objectives] for time_index in range(len(tlist) - 1): # iterate over time intervals dt = tlist[time_index + 1] - tlist[time_index] if second_order: σ = sigma(tlist[time_index] + 0.5 * dt) # pulse update for (i_pulse, guess_pulse) in enumerate(guess_pulses): for (i_obj, objective) in enumerate(objectives): χ = backward_states[i_obj][time_index] μ = mu( objectives, i_obj, guess_pulses, pulses_mapping, i_pulse, time_index, ) Ψ = fw_states[i_obj] update = overlap(χ, μ(Ψ)) # ⟨χ|μ|Ψ⟩ ∈ ℂ update *= chi_norms[i_obj] if second_order: update += 0.5 * σ * overlap(delta_phis[i_obj], μ(Ψ)) delta_eps[i_pulse][time_index] += update λₐ = lambda_vals[i_pulse] S_t = shape_arrays[i_pulse][time_index] Δϵ = (S_t / λₐ) * delta_eps[i_pulse][time_index].imag # ∈ ℝ g_a_integrals[i_pulse] += abs(Δϵ) ** 2 * dt # dt may vary! optimized_pulses[i_pulse][time_index] += Δϵ # forward propagation fw_states = parallel_map[2]( _forward_propagation_step, list(range(len(objectives))), ( fw_states, objectives, optimized_pulses, pulses_mapping, tlist, time_index, propagators, ), ) if second_order: # Δϕ(t + dt), to be used for the update in the next interval delta_phis = [ fw_states[k] - forward_states0[k][time_index + 1] for k in range(len(objectives)) ] # storage for i_obj in range(len(objectives)): forward_states[i_obj][time_index + 1] = fw_states[i_obj]"Finished forward propagation/pulse update") fw_states_T = fw_states tau_vals = np.array( [ overlap(, fw_state_T) for (fw_state_T, obj) in zip(fw_states_T, objectives) ] ) toc = time.time() # Display information about iteration if info_hook is not None: info = info_hook( backward_states=backward_states, forward_states=forward_states, forward_states0=forward_states0, fw_states_T=fw_states_T, guess_pulses=guess_pulses, optimized_pulses=optimized_pulses, g_a_integrals=g_a_integrals, tau_vals=tau_vals, start_time=tic, stop_time=toc, info_vals=result.info_vals, shared_data={}, iteration=krotov_iteration, **info_hook_static_args ) # Update optimization `result` with info from finished iteration result.iters.append(krotov_iteration) result.iter_seconds.append(int(toc - tic)) if info is not None: result.info_vals.append(info) if not np.all(tau_vals == None): # noqa result.tau_vals.append(tau_vals) result.optimized_controls = optimized_pulses # pulses (time intervals) will be converted to controls (time grid # points) farther below in "Finalize" if store_all_pulses: # we need to make a copy, so that the conversion in "Finalize" # doesn't affect `all_pulses` as well. result.all_pulses.append(copy.deepcopy(optimized_pulses)) result.states = fw_states_T"Finished Krotov iteration %d", krotov_iteration) # Convergence check msg = None if check_convergence is not None: msg = check_convergence(result) if krotov_iteration >= info_hook_static_args['iter_stop']: # modify_params_after_iter may change iter_stop! iter_stop = info_hook_static_args['iter_stop'] result.message = "Reached %d iterations" % iter_stop break if bool(msg) is True: # this is not an anti-pattern! result.message = "Reached convergence" if isinstance(msg, str): result.message += ": " + msg break else: # prepare for next iteration guess_pulses = optimized_pulses if second_order: sigma.refresh( forward_states=forward_states, forward_states0=forward_states0, chi_states=chi_states, chi_norms=chi_norms, optimized_pulses=optimized_pulses, guess_pulses=guess_pulses, objectives=objectives, result=result, ) forward_states0 = forward_states else: # optimization finished without `check_convergence` break result.message = "Reached %d iterations" % max(iter_start, iter_stop) # Finalize result.end_local_time = time.localtime() for i, pulse in enumerate(optimized_pulses): result.optimized_controls[i] = pulse_onto_tlist(pulse) return result
def _shape_val_to_callable(val): if val == 1: return one_shape elif val == 0: return zero_shape else: if callable(val): return val else: raise ValueError("update_shape must be a callable") def _enforce_shape_array_range(shape_array): """Enforce values ∈ [0, 1] in shape array, with some room for rounding errors that will be clipped away. """ assert not np.iscomplexobj(shape_array) # `discretize` should catch this # the rounding errors can be introduced by control_onto_interval, and # result in values slightly below 0 or above 1. We allow a generous margin # of ±0.01; if something nonsensical is passed as a shape, we can be pretty # sure that it will deviate by a significantly larger error. if np.min(shape_array) < -0.01 or np.max(shape_array) > 1.01: raise ValueError( "Update shapes ('update_shape' in pulse options-dict) must have " "values in the range [0, 1], not [%s, %s]" % (np.min(shape_array), np.max(shape_array)) ) return np.clip(shape_array, a_min=0.0, a_max=1.0) def _check_propagators_interface(propagators, logger): """Warn if any of the propagators do not have the expected interface. In order to pass muster, each propagator must either have the same interface as :func:`krotov.propagators.expm`, or :meth:`krotov.propagators.Propagator.__call__` """ for propagator in propagators: spec = inspect.getfullargspec(propagator) spec_func = inspect.getfullargspec(expm) spec_obj = inspect.getfullargspec(Propagator.__call__) if spec != spec_func and spec != spec_obj: logger.warning( "The propagator %s does not have the expected interface.", propagator, ) def _initialize_krotov_controls(objectives, pulse_options, tlist): """Extract discretized guess controls and pulses from `objectives`, and return them with the associated mapping and option data""" guess_controls = extract_controls(objectives) pulses_mapping = extract_controls_mapping(objectives, guess_controls) options_list = pulse_options_dict_to_list(pulse_options, guess_controls) try: guess_controls = [ discretize( control, tlist, args=(pulse_options[control].get('args', None),), ) for control in guess_controls ] except (TypeError, np.ComplexWarning) as exc_info: raise ValueError( "Cannot discretize controls: %s. Note that " "all controls must be real-valued. Complex controls must be " "split into an independent real and imaginary part in the " "objectives before passing them to the optimization" % exc_info ) guess_pulses = [ # defined on the tlist intervals control_onto_interval(control) for control in guess_controls ] try: lambda_vals = np.array( [float(options['lambda_a']) for options in options_list] ) except KeyError: raise ValueError( "Each value in pulse_options must be a dict that contains " "the key 'lambda_a'." ) shape_arrays = [] for options in options_list: try: S = discretize( _shape_val_to_callable(options['update_shape']), tlist, args=() ) except KeyError: raise ValueError( "Each value in pulse_options must be a dict that contains " "the key 'update_shape'." ) except (TypeError, np.ComplexWarning) as exc_info: raise ValueError( "Update shapes ('update_shape' in pulse options-dict) must be " "real-valued: %s" % exc_info ) shape_arrays.append( _enforce_shape_array_range(control_onto_interval(S)) ) return ( guess_controls, guess_pulses, pulses_mapping, lambda_vals, shape_arrays, ) def _restore_from_previous_result(result, objectives, tlist, store_all_pulses): """Load `guess_controls` and `guess_pulses` from the given Result object. Raises: ValueError: if `result` is incompatible with the given `objectives` and `tlist`. """ if not isinstance(result, Result): raise ValueError("Continuation is only possible from a Result object") if len(objectives) != len(result.objectives): raise ValueError( "When continuing from a previous Result, the number of " "objectives must be the same" ) for (a, b) in zip(objectives, result.objectives): if a != b: raise ValueError( "When continuing from a previous Result, the objectives must " "remain unchanged" ) if store_all_pulses: if len(result.all_pulses) == 0: raise ValueError( "The store_all_pulses parameter cannot be changed when " "continuing from a previous Result. Pass it as False." ) else: if len(result.all_pulses) > 0: raise ValueError( "The store_all_pulses parameter cannot be changed when " "continuing from a previous Result. Pass it as True." ) try: if np.max(np.abs(np.array(tlist) - np.array(result.tlist))) > 1e-5: # we're ok with pretty significant rounding errors: if someone is # doing something stupid (like changing physical units), the error # should be large raise ValueError("tlist does not match") except ValueError: # if the size of the tlist has changed, the "if" statement will already # raise a ValueError raise ValueError( "When continuing from a previous Result, the controls must be " "defined on the same time grid" ) nt = len(tlist) guess_controls = [] for control in result.optimized_controls: if len(control) == nt - 1: # the Result was dumped before the optimization was complete # (see krotov.konvergence.dump_result), so that the # `optimized_controls` are actually pulses (defined on the # intervals) guess_controls.append(pulse_onto_tlist(control)) elif len(control) == nt: # the Result was dumped after the optimization was complete, so # that the `optmized_controls` are in fact controls (defined on the # points of tlist) guess_controls.append(control) else: # this should never happen raise ValueError( "Invalid Result: optimized_controls and tlist are incongruent" ) guess_pulses = [ # defined on the tlist intervals control_onto_interval(control) for control in guess_controls ] return guess_controls, guess_pulses def _skip_initial_forward_propagation(objectives, result, sigma, logger): """Extract forward_states from existing result""" if sigma is not None: raise ValueError( "skip_initial_forward_propagation is incompatible with " "second order Krotov (sigma is not None)" ) if result is not None: # forwards_states is normally a list of storage objects that store # the *entire* forward propagation for each objective. The stored # states are only needed for second order Krotov. When skipping the # forward propagation, we'll use a fake "storage object" that # contains only the final state. That'll be sufficient for setting # fw_states_T and tau_vals below forward_states = [[state] for state in result.states] else: logger.warning( "You should not use `skip_initial_forward_propagation` unless " "you are also passing `continue_from`" ) # If the chi_constructor does not depend on the fw_states_T (like # chis_re), and you're using first order Krotov, you don't actually # have to do the initial forward propagation. It's not really worth # it to skip that propagation, though, as we usually have to do # hundreds of iterations. So, this is an undocumented feature. forward_states = [[None] for _ in objectives] return forward_states def _forward_propagation( i_objective, objectives, pulses, pulses_mapping, tlist, propagators, storage, store_all=True, ): """Forward propagation of the initial state of a single objective over the entire `tlist`""" logger = logging.getLogger('krotov') "Started initial forward propagation of objective %d", i_objective ) obj = objectives[i_objective] state = obj.initial_state if store_all: storage_array = storage(len(tlist)) storage_array[0] = state mapping = pulses_mapping[i_objective] for time_index in range(len(tlist) - 1): # index over intervals H = plug_in_pulse_values(obj.H, pulses, mapping[0], time_index) c_ops = [ plug_in_pulse_values(c_op, pulses, mapping[ic + 1], time_index) for (ic, c_op) in enumerate(obj.c_ops) ] dt = tlist[time_index + 1] - tlist[time_index] state = propagators[i_objective]( H, state, dt, c_ops, initialize=(time_index == 0) ) if store_all: storage_array[time_index + 1] = state "Finished initial forward propagation of objective %d", i_objective ) if store_all: return storage_array else: return state def _backward_propagation( i_state, chi_states, adjoint_objectives, pulses, pulses_mapping, tlist, propagators, storage, ): """Backward propagation of chi_states[i_state] over the entire `tlist`""" logger = logging.getLogger('krotov')"Started backward propagation of state %d", i_state) state = chi_states[i_state] obj = adjoint_objectives[i_state] storage_array = storage(len(tlist)) storage_array[-1] = state mapping = pulses_mapping[i_state] for time_index in range(len(tlist) - 2, -1, -1): # index bw over intervals H = plug_in_pulse_values( obj.H, pulses, mapping[0], time_index, conjugate=True ) c_ops = [ plug_in_pulse_values(c_op, pulses, mapping[ic + 1], time_index) for (ic, c_op) in enumerate(obj.c_ops) ] dt = tlist[time_index + 1] - tlist[time_index] state = propagators[i_state]( H, state, dt, c_ops, backwards=True, initialize=(time_index == len(tlist) - 2), ) storage_array[time_index] = state"Finished backward propagation of state %d", i_state) return storage_array def _forward_propagation_step( i_state, states, objectives, pulses, pulses_mapping, tlist, time_index, propagators, ): """Forward-propagate states[i_state] by a single time step""" state = states[i_state] obj = objectives[i_state] mapping = pulses_mapping[i_state] H = plug_in_pulse_values(obj.H, pulses, mapping[0], time_index) c_ops = [ plug_in_pulse_values(c_op, pulses, mapping[ic + 1], time_index) for (ic, c_op) in enumerate(obj.c_ops) ] dt = tlist[time_index + 1] - tlist[time_index] return propagators[i_state]( H, state, dt, c_ops, initialize=(time_index == 0) )