Quantum Optimal Control using TensorFlow ============================================== - ``David Schuster`` lab (2017): Speedup for quantum optimal control from automatic differentiation based on graphics processing units - ``David Schuster`` lab (2019): Gradient-based optimal control of open quantum systems using quantum trajectories and automatic differentiation - **Benefits of TensorFlow** - Arbitrary cost function, which would be hard to implement in GRAPE - Best code: Using TenserFlow developed by Google - Easy GPU integration - Follows the qutip code standard :: optimizedResult = pulseoptim.optimize_pulse_unitary( H,controlsList,initialPsi0,targetPsi0,Ntime,totalTime, init_pulse_type='LIN', ) My code ------------- .. code:: python import tensorflow as tf def qtp2tf(x): return tf.constant(x.full()) def expm(A): return tf.eye(A.shape[0],dtype=A.dtype)+A class QuantumOptimalControlModel(object): def __init__( self,H,controlsList,Ntime,totalTime,pulseType='linear', ): # constants self.Hdrift = qtp2tf(H) self.dtype = self.Hdrift.dtype self.controlsList = [ qtp2tf(c) for c in controlsList ] # variables pulse = np.linspace(-1,1,Ntime) self.initialAmpsList = [ tf.constant(pulse,dtype=self.dtype) for c in controlsList ] self.ampsList = [ tf.Variable(pulse,dtype=self.dtype) for c in controlsList ] # others self.dim = H.shape[0] self.dt = totalTime/Ntime self.Ntime = Ntime def __call__(self, initialPsi): U = tf.eye(self.dim, dtype=self.dtype) for nTime in range(self.Ntime): Ht = self.Hdrift + sum([ self.ampsList[nControl][nTime] * control for nControl,control in enumerate(self.controlsList) ]) Ut = tf.linalg.expm(-1j * Ht * self.dt) #Ut = expm(-1j * Ht * self.dt) U = tf.linalg.matmul(Ut, U) finalPsi = tf.linalg.matmul(U, initialPsi) return finalPsi def fidelityTensorFlow(a,b): return tf.abs(tf.tensordot(tf.squeeze(a), tf.squeeze(b), 1)) def loss(targetPsi, finalPsi): return 1 - fidelityTensorFlow(targetPsi, finalPsi) def train(model,initial,target,learningRate=0.1): optimizer = tf.keras.optimizers.SGD(learning_rate=learningRate) with tf.GradientTape() as t: currentLoss = loss(target,model(initial)) gradients = t.gradient(currentLoss, model.ampsList) optimizer.apply_gradients(zip(gradients, model.ampsList)) return currentLoss ########################################### import numpy as np %matplotlib inline from qutip import * import qutip.control.pulseoptim as pulseoptim from yarn.qutipHelpers import \ jaynesCummingsHamiltonian, \ plotExpectation, plotWigners, plotOptimalControl, \ cat dimCav = 20 H, controlsDict, observables, operators = jaynesCummingsHamiltonian( cavityDimension=dimCav, omegaCavity=1, omegaQubit=1, g=0.1 ) controlsList = list(controlsDict.values()) controlsName = list(controlsDict.keys()) zPlus = (cat(dimCav,2)+cat(dimCav,2j)).unit() zMinus = (cat(dimCav,2)-cat(dimCav,2j)).unit() initialPsi0 = tensor(coherent(dimCav,2), fock(2,0)) targetPsi0 = tensor(zPlus, fock(2,0)) Ntime = 500; totalTime = 50 time = np.linspace(0, totalTime, Ntime) ########################################### initialPsi0 = qtp2tf(initialPsi0) targetPsi0 = qtp2tf(targetPsi0) model = QuantumOptimalControlModel(H,controlsList,Ntime,totalTime) currentLoss = 1 maxIteration=500 for i in range(maxIteration): currentLoss = train(model,initialPsi0,targetPsi0) percent10 = int(maxIteration/10) if i % percent10 == 0: print("{}% fidelity error {}".format( int(i*100/maxIteration),currentLoss )) .. parsed-literal:: 0% fidelity error 0.6635763552719809 :: --------------------------------------------------------------------------- InvalidArgumentError Traceback (most recent call last) in 85 maxIteration=500 86 for i in range(maxIteration): ---> 87 currentLoss = train(model,initialPsi0,targetPsi0) 88 percent10 = int(maxIteration/10) 89 if i % percent10 == 0: in train(model, initial, target, learningRate) 42 optimizer = tf.keras.optimizers.SGD(learning_rate=learningRate) 43 with tf.GradientTape() as t: ---> 44 currentLoss = loss(target,model(initial)) 45 gradients = t.gradient(currentLoss, model.ampsList) 46 optimizer.apply_gradients(zip(gradients, model.ampsList)) in __call__(self, initialPsi) 28 for nControl,control in enumerate(self.controlsList) 29 ]) ---> 30 Ut = tf.linalg.expm(-1j * Ht * self.dt) 31 #Ut = expm(-1j * Ht * self.dt) 32 U = tf.linalg.matmul(Ut, U) c:\python38\lib\site-packages\tensorflow\python\ops\linalg\linalg_impl.py in matrix_exponential(input, name) 322 numer = u + v 323 denom = -u + v --> 324 result = linalg_ops.matrix_solve(denom, numer) 325 max_squarings = math_ops.reduce_max(squarings) 326 c:\python38\lib\site-packages\tensorflow\python\ops\gen_linalg_ops.py in matrix_solve(matrix, rhs, adjoint, name) 1563 raise 1564 except _core._NotOkStatusException as e: -> 1565 _ops.raise_from_not_ok_status(e, name) 1566 # Add nodes to the TensorFlow graph. 1567 if adjoint is None: c:\python38\lib\site-packages\tensorflow\python\framework\ops.py in raise_from_not_ok_status(e, name) 6651 message = e.message + (" name: " + name if name is not None else "") 6652 # pylint: disable=protected-access -> 6653 six.raise_from(core._status_to_exception(e.code, message), None) 6654 # pylint: enable=protected-access 6655 ~\AppData\Roaming\Python\Python38\site-packages\six.py in raise_from(value, from_value) InvalidArgumentError: Input matrix is not invertible. [Op:MatrixSolve]