Support for auto differentiation in `Coefficient` and `QobjEvo`.

[Ericgig]

Description Remove coercion of coefficient values to complex in Coefficient and QobjEvo allowing to use tf.Variable as coefficient to enable auto-differentiation to be used in solver. It is build on top of #1644 since it include integration method using our data object.

Coefficient creation is no longer hard coded, but use a dict of type: callable, each type is associated to a function that create a coefficient, allowing to add support for other coefficient types per project such as qutip-tensorflow. Also, function based coefficient no longer cast to complex, but the check in the coefficient function is still there. Since other type are not supported without new data layer, this feels like a proper balance.

QobjEvo cython casting of coefficient's value to complex as been removing. No tests are added since other types of coefficient are not supported by data layer. But here is an example of auto differentiation with this branch:

import qutip as qt
import numpy as np
import qutip_tensorflow as qtf
import tensorflow as tf
from qutip.core.coefficient import Coefficient, coefficient_builders
from qutip.solver.sesolve import SeSolver

# Create a new Coefficient for tf.Variable, 
# This should be added to qutip-tensorflow,
# Here I use a step interpolation.
class TfVarArgsCoefficient(Coefficient):
    def __init__(self, variable, tlist, **_):
        self.variable = variable
        self._tlist = tlist
        
    def __call__(self, t, _args=None, **_):
        if t <= self._tlist[0]:
            return self.variable[0]
        if t >= self._tlist[-1]:
            return self.variable[-1]
        idx = np.searchsorted(self._tlist, t, 'right') - 1
        return self.variable[idx]

coefficient_builders[tf.Variable] = TfVarArgsCoefficient

# Create a Coefficient as normal.
variable = tf.Variable(np.linspace(0.1,0.9,6), dtype=tf.complex128)
coeff1 = qt.coefficient(variable, tlist=np.linspace(0,0.5,6))

N = 5
a = qt.destroy(N, dtype='tftensor')
num = qt.num(N, dtype='tftensor')
H = qt.qeye(N, dtype='tftensor') + qt.QobjEvo([a+a.dag(), coeff1])

solver = SeSolver(H, options={
    'method': 'vern7',  # Only verner method support autodiff for now.
    'state_data_type': "",  # The default is 'dense', so it must be overwritten.
    'first_step': 0.05,  # Fixed step make auto-diff faster, but still slow.
    'min_step': 0.05,
    'max_step': 0.05,
    'atol': 1e-2,
})
tlist = np.linspace(0,0.5,6)

with tf.GradientTape() as tape:
    psi = qt.basis(N, N-1, dtype='tftensor')
    solver.start(psi, 0)
    # Result' expect are numpy's array, so we need to compute them ourself.
    out = solver.run(psi, tlist)
    expects = [qt.expect(num, state) for state in out.states[1:]]

%time tape.jacobian(expects[-1], variable)

Output:

20.0%. Run time:   0.03s. Est. time left: 00:00:00:00
40.0%. Run time:   0.05s. Est. time left: 00:00:00:00
60.0%. Run time:   0.08s. Est. time left: 00:00:00:00
80.0%. Run time:   0.10s. Est. time left: 00:00:00:00
100.0%. Run time:   0.12s. Est. time left: 00:00:00:00

/home/eric/miniconda3/lib/python3.9/site-packages/tensorflow/python/framework/ops.py:1128: ComplexWarning: Casting complex values to real discards the imaginary part
  return float(self._numpy())

Total run time:   0.12s
CPU times: user 15.9 s, sys: 174 ms, total: 16 s
Wall time: 16.1 s

<tf.Tensor: shape=(5,), dtype=complex128, numpy=
array([-0.17340122+2.03548130e-09j, -0.18075013+2.08120101e-09j,
       -0.19236081+1.58339267e-08j, -0.18075012+1.12229134e-08j,
       -0.17648842-7.47414642e-08j])>

So while slow, it works. Maybe adding simpler runge-kutta method and / or support for tensorflow's ode could make it more usable.

Related issues or PRs Build on top of #1644

Changelog Remove casting of coefficients to complex.

[coveralls]

Coverage increased (+1.6%) to 67.304% when pulling f4473d1b16750ba783d54d7d10c914343433819c on Ericgig:Qevo.autodiff.support into d240c4ccc4196e39982f2a4c74e723e973138fa6 on qutip:dev.major.

[Ericgig]

Will redo to also work with jax.