diff --git a/oqupy/contractions.py b/oqupy/contractions.py index ee116db..426eda7 100644 --- a/oqupy/contractions.py +++ b/oqupy/contractions.py @@ -26,7 +26,25 @@ from oqupy.control import Control from oqupy.dynamics import Dynamics, MeanFieldDynamics from oqupy.process_tensor import BaseProcessTensor -from oqupy.system import BaseSystem, System, TimeDependentSystem +from oqupy.system import BaseSystem, System, TimeDependentSystem, ParameterizedSystem +from oqupy.system import MeanFieldSystem +from oqupy.operators import left_super, right_super +from oqupy.util import check_convert, check_isinstance, check_true +from oqupy.util import get_progress + +from typing import List, Optional, Text, Tuple, Union, Callable + +import numpy as np +from numpy import ndarray +import tensornetwork as tn +from oqupy.system import TimeDependentSystemWithField + +from oqupy.config import NpDtype, INTEGRATE_EPSREL, SUBDIV_LIMIT +from oqupy.control import Control +from oqupy.dynamics import Dynamics, MeanFieldDynamics +from oqupy.process_tensor import BaseProcessTensor +from oqupy.system import ParameterizedSystem, BaseSystem, System, TimeDependentSystem + from oqupy.system import MeanFieldSystem from oqupy.operators import left_super, right_super from oqupy.util import check_convert, check_isinstance, check_true @@ -181,6 +199,280 @@ def controls(step: int): return Dynamics(times=list(times),states=states) + +def compute_gradient_and_dynamics( + system: ParameterizedSystem, + parameters : Optional[ndarray]=None, + initial_state: Optional[ndarray] = None, + target_derivative: Optional[Union[Callable,ndarray]] = None, + dt: Optional[float] = None, + num_steps: Optional[int] = None, + start_time: Optional[float] = 0.0, + process_tensors: Optional[Union[List[BaseProcessTensor], + BaseProcessTensor]] = None, + control: Optional[Control] = None, + record_all: Optional[bool] = True, + get_forward_and_backprop_list = False, + dynamics_only: Optional[bool] = False, + progress_type: Optional[Text] = None) -> Tuple[List,Dynamics]: + """ + Compute some objective function and calculate its gradient w.r.t. + some control parameters, accounting + (optionally) for interaction with an environment using one or more + process tensors. + + Parameters + ---------- + system: Union[System, TimeDependentSystem] + Object containing the system Hamiltonian information. + initial_state: ndarray + Initial system state. + target_derivative: + Some pure target state or derivative w.r.t. an objective functioni + dt: float + Length of a single time step. + num_steps: int + Optional number of time steps to be computed. + start_time: float + Optional start time offset. + process_tensor: Union[List[BaseProcessTensor],BaseProcessTensor] + Optional process tensor object or list of process tensor objects. + control: Control + Optional control operations. + record_all: bool + If `false` function only computes the final state. + subdiv_limit: int (default = config.SUBDIV_LIMIT) + The maximum number of subdivisions used during the adaptive + algorithm when integrating the system Liouvillian. If None + then the Liouvillian is not integrated but sampled twice to + to construct the system propagators at each timestep. + liouvillian_epsrel: float (default = config.INTEGRATE_EPSREL) + The relative error tolerance for the adaptive algorithm + when integrating the system Liouvillian. + progress_type: str (default = None) + The progress report type during the computation. Types are: + {``silent``, ``simple``, ``bar``}. If `None` then + the default progress type is used. + + Returns + ------- + dynamics: Dynamics + The system dynamics for the given system Hamiltonian + (accounting for the interaction with the environment). + """ + + # -- input parsing -- + parsed_parameters = _compute_dynamics_input_parse( + False, system, initial_state, dt, num_steps, start_time, + process_tensors, control, record_all) + system, initial_state, dt, num_steps, start_time, \ + process_tensors, control, record_all, hs_dim = parsed_parameters + + assert target_derivative is not None, \ + 'target state must be given explicitly' + + num_envs = len(process_tensors) + + # -- prepare propagators -- + propagators = system.get_propagators(dt, parameters) + + # -- prepare controls -- + def controls(step: int): + return control.get_controls( + step, + dt=dt, + start_time=start_time) + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # ~~~~ Forwardpropagation ~~~~ + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + # -- initialize computation -- + # + # Initial state including the bond legs to the environments with: + # edges 0, 1, .., num_envs-1 are the bond legs of the environments + # edge -1 is the state leg + initial_ndarray = initial_state.reshape(hs_dim**2) + initial_ndarray.shape = tuple([1]*num_envs+[hs_dim**2]) + current_node = tn.Node(initial_ndarray) + current_edges = current_node[:] + + states = [] + title = "--> Compute dynamics:" + prog_bar = get_progress(progress_type)(num_steps, title) + prog_bar.enter() + + forwardprop_derivs_list = [] + mpo_list=[] + + for step in range(num_steps+1): + + # -- apply pre measurement control -- + pre_measurement_control, post_measurement_control = controls(step) + + if pre_measurement_control is not None: + current_node, current_edges = _apply_system_superoperator( + current_node, current_edges, pre_measurement_control) + + if step == num_steps: + break + + # -- extract current state -- update field -- + if record_all: + caps = _get_caps(process_tensors, step) + state_tensor = _apply_caps(current_node, current_edges, caps) + state = state_tensor.reshape(hs_dim, hs_dim) + states.append(state) + + prog_bar.update(step) + + # -- apply post measurement control -- + if post_measurement_control is not None: + current_node, current_edges = _apply_system_superoperator( + current_node, current_edges, post_measurement_control) + + forwardprop_derivs_list.append(tn.replicate_nodes([current_node])[0]) + + # -- propagate one time step -- + first_half_prop, second_half_prop = propagators(step) + + pt_mpos = _get_pt_mpos(process_tensors, step) + mpo_list.append(pt_mpos) + + current_node, current_edges = _apply_system_superoperator( + current_node, current_edges, first_half_prop) + current_node, current_edges = _apply_pt_mpos( + current_node, current_edges, pt_mpos) + + current_node, current_edges = _apply_system_superoperator( + current_node, current_edges, second_half_prop) + + # -- extract last state -- + caps = _get_caps(process_tensors, num_steps) + state_tensor = _apply_caps(current_node, current_edges, caps) + final_state = state_tensor.reshape(hs_dim, hs_dim) + states.append(final_state) + + prog_bar.update(num_steps) + prog_bar.exit() + + # -- create dynamics object -- + if record_all: + times = start_time + np.arange(len(states))*dt + else: + times = [start_time + len(states)*dt] + + if dynamics_only: + return [], Dynamics(times=list(times),states=states) + + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # ~~~~~ Backpropagation ~~~~~~ + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + + # -- initialize computation (except backwards) -- + # + # Initial state including the bond legs to the environments with: + # edges 0, 1, .., num_envs-1 are the bond legs of the environments + # edge -1 is the state leg + + if callable(target_derivative): + target_derivative=target_derivative(states[-1]) + + target_ndarray = target_derivative + target_ndarray = target_ndarray.reshape(hs_dim**2) + target_ndarray.shape = tuple([1]*num_envs+[hs_dim**2]) + current_node = tn.Node(target_ndarray) + current_edges = current_node[:] + + combined_deriv_list = [] + + pre_measurement_control, post_measurement_control=controls(num_steps) + + if pre_measurement_control is not None: + current_node, current_edges = _apply_system_superoperator( + current_node, current_edges, pre_measurement_control.T) + + forwardprop_tensor = forwardprop_derivs_list[num_steps-1] + + if get_forward_and_backprop_list: + backprop_derivs_list.append(tn.replicate_nodes([current_node])[0]) + + pt_mpos = mpo_list[num_steps-1] + backprop_tensor = tn.replicate_nodes([current_node])[0] + + fwd_edges = forwardprop_tensor[:] + deriv_forwardprop_tensor,fwd_edges = _apply_derivative_pt_mpos(forwardprop_tensor,fwd_edges,pt_mpos) + + for i,pt_mpo in enumerate(pt_mpos): + fwd_edges[i] ^ backprop_tensor[i] + + deriv = deriv_forwardprop_tensor @ backprop_tensor + + combined_deriv_list.append(tn.replicate_nodes([deriv])[0]) + + for step in reversed(range(1,num_steps)): + + # -- now the backpropagation part -- + pre_measurement_control, post_measurement_control = controls(step) + first_half_prop, second_half_prop = propagators(step) + pt_mpos = _get_pt_mpos_backprop(mpo_list, step) + + current_node, current_edges = _apply_system_superoperator( + current_node, current_edges, second_half_prop.T) + + current_node, current_edges = _apply_pt_mpos( + current_node, current_edges, pt_mpos) + + current_node, current_edges = _apply_system_superoperator( + current_node, current_edges, first_half_prop.T) + + if post_measurement_control is not None: + current_node, current_edges = _apply_system_superoperator( + current_node, current_edges, post_measurement_control.T) + + if pre_measurement_control is not None: + current_node, current_edges = _apply_system_superoperator( + current_node, current_edges, pre_measurement_control.T) + + forwardprop_tensor = forwardprop_derivs_list[step-1] + + if get_forward_and_backprop_list: + backprop_derivs_list.append(tn.replicate_nodes([current_node])[0]) + + backprop_tensor = tn.replicate_nodes([current_node])[0] + + pt_mpos = mpo_list[step-1] + + fwd_edges = forwardprop_tensor[:] + deriv_forwardprop_tensor,fwd_edges = _apply_derivative_pt_mpos(forwardprop_tensor,fwd_edges,pt_mpos) + + for i,pt_mpo in enumerate(pt_mpos): + fwd_edges[i] ^ backprop_tensor[i] + + deriv = deriv_forwardprop_tensor @ backprop_tensor + + combined_deriv_list.append(tn.replicate_nodes([deriv])[0]) + + # deriv_list is currently in the reversed order from what you'd expect, so + # reversing the order of the list..... + + deriv_list_reversed = list(reversed(combined_deriv_list)) + + # -- create dynamics object -- + if record_all: + times = start_time + np.arange(len(states))*dt + else: + times = [start_time + len(states)*dt] + + dynamics_instance = Dynamics(times=list(times),states=states) + + if get_forward_and_backprop_list is False: + forwardprop_derivs_list = None + backprop_derivs_list = None + + return deriv_list_reversed, dynamics_instance + def compute_dynamics_with_field( mean_field_system: MeanFieldSystem, initial_field: complex, @@ -581,6 +873,49 @@ def _get_pt_mpos(process_tensors: List[BaseProcessTensor], step: int): return pt_mpos +def _get_pt_mpos_backprop(mpo_list:ndarray, step: int): + """same as above but swaps the system legs and internal bond legs + before returning MPOs. + + [forwardprop] + [1] + | [3] + | / + |---------| / + | |/ + | |\ propagate + |---------| \ ^ + | \ | + | [2] + [0] + + | + | + \ / + v + [backprop] + [0] + | [3] + | / + |---------| / + | |/ + | |\ propagate + |---------| \ ^ + | \ | + | [2] + [1] + """ + pt_mpos = mpo_list[step] + pt_mpos_rev = [] + for i,pt_mpo in enumerate(pt_mpos): + # now swap axes so propagating upwards on the PT diagram propagates + # *backwards* in time + pt_mpo = np.swapaxes(pt_mpo,0,1) # internal bond legs + pt_mpo = np.swapaxes(pt_mpo,2,3) # system propagator legs + pt_mpos_rev.append(pt_mpo) + return pt_mpos_rev + + def _apply_system_superoperator(current_node, current_edges, sup_op): """ToDo """ if sup_op is None: @@ -605,7 +940,40 @@ def _apply_caps(current_node, current_edges, caps): def _apply_pt_mpos(current_node, current_edges, pt_mpos): - """ToDo """ + """ + Apply MPO for forward propagation step + + before mpo application: + [1] + | [3] + | / + |---------| / + | |/ + | |\ + |---------| \ + | \ + | [2] + [0] + + [i] + | + | [-1] + | | + | + + after mpo application: + [0] bond edge + | [-1] system edge + | / + |---------| / + | |/ + | |\ + |---------| \ + | \ + | [] + | | + | + """ for i, pt_mpo in enumerate(pt_mpos): if pt_mpo is None: continue @@ -619,6 +987,91 @@ def _apply_pt_mpos(current_node, current_edges, pt_mpos): current_edges[-1] = new_sys_edge return current_node, current_edges +def _apply_derivative_pt_mpos(current_node,current_edges,pt_mpos): + """ + Apply MPO corresponding to the time-step of the propagators the derivative is being taken w.r.t. + + before mpo application: + [1] + | [3] + | / + |---------| / + | |/ + | i |\ + |---------| \ + | \ + | [2] + [0] + + [i] + | + | [-1] + | | + | + + after mpo application: + [i] + | [-1] post_mpo_edge + | / + |---------| / + | |/ + | |\ + |---------| \ + | \ + | [-2] pre_mpo_edge + | + | + | + | [-3] prop_edge + | | + | + """ + prev_prop_edge = current_edges[-1] + pt_mpo_node = tn.Node(pt_mpos[0]) + pre_mpo_edge = pt_mpo_node[2] + new_bond_edge=pt_mpo_node[1] + post_mpo_edge = pt_mpo_node[3] + + prev_prop_edge.name = "prop edge" + pre_mpo_edge.name = "pre mpo edge 0" + new_bond_edge.name="bond edge 0" + post_mpo_edge.name = "post mpo edge 0" + + current_edges[0] ^ pt_mpo_node[0] + current_node = current_node @ pt_mpo_node + current_edges = current_node[:] + + bond_edges=[new_bond_edge] + + for i,pt_mpo in enumerate(pt_mpos[1:]): + if pt_mpo is None: + continue + + pt_mpo_node=tn.Node(pt_mpo) + + new_bond_edge = pt_mpo_node[1] + bond_edges.append(new_bond_edge) + post_mpo_edge = pt_mpo_node[3] + + new_bond_edge.name="bond edge "+str(i+1) + post_mpo_edge.name = "post mpo edge "+str(i+1) + + current_edges[0] ^ pt_mpo_node[0] + + current_edges[-1]^pt_mpo_node[2] + + current_node = current_node @ pt_mpo_node + current_edges = current_node[:] + + current_edges[-1]=post_mpo_edge + current_edges[-2]=pre_mpo_edge + current_edges[-3]=prev_prop_edge + + for i,mpo in enumerate(pt_mpos): + current_edges[i]=bond_edges[i] + + return current_node,current_edges + # -- compute correlations ------------------------------------------------