on general variational dynamics in variational_dynamics.py

[gaox-qd]

Hi, I found the example in variational_dynamics.py is very interesting.
I'm trying to extend it to more general form as shown in eq13-14 of arXiv:1812.08767v4, while the result is not good(as follows).

fig, ax = plt.subplots()
ax.plot(t,np.array(pop_ext)/2.0+0.5, c='r', label='exact')
ax.plot(t,np.array(pop_var)/2.0+0.5, c='b', label='varational')
ax.legend()
ry=np.arange(0.0,1.1,0.2)

Here are the only changes in function of lhs_matrix and rhs_vector in example/variational_dynamics.py:
variationa_dynamics_test.zip

@tc.backend.jit
def lhs_matrix(theta,psi0):
    
    def _vdot(i, j):
        return tc.backend.tensordot(tc.backend.conj(i), j, 1)

    def ij(i, j):
        """
        Inner product
        """
        # return _vdot(i, j) - _vdot(i, psi) * _vdot(psi, j)
        return _vdot(i, j) + _vdot(i, psi0) * _vdot(j, psi0)

    vij = tc.backend.vmap(ij, vectorized_argnums=0)
    vvij = tc.backend.vmap(vij, vectorized_argnums=1)
    jacobian = ppsioverptheta(theta, psi0=psi0)
    jacobian = tc.backend.transpose(jacobian)
    fim = vvij(jacobian, jacobian) 
    fim = tc.backend.real(fim)
    return fim


@tc.backend.jit
def rhs_vector(theta, psi0):
    def energy(theta, psi0):
        w = variational_wfn(theta, psi0)
        wl = tc.backend.stop_gradient(w)
        wl = tc.backend.conj(wl)
        wr = w
        wl = tc.backend.reshape(wl, [1, -1])
        wr = tc.backend.reshape(wr, [-1, 1])
        e = wl @ h @ wr 
        wl0 = tc.backend.conj(w)
        wl0 = tc.backend.reshape(wl0, [1, -1])
        e0 = wl0 @ h @ wr # <psi0| H | psi0>
        w0 = wl @ wr #<\partial psi0/\partial theta | psi0>
        return tc.backend.real(e-w0*e0)[0, 0]

    eg = tc.backend.grad(energy, argnums=0)
    rhs = eg(theta, psi0)
    rhs = tc.backend.imag(rhs)
    return rhs

[refraction-ray]

The rhs build in the implementation seems not correct to me.
For the original implementation, we build $\langle \bot \psi\vert H \vert \psi\rangle$, this is why when we use tc.backend.grad to differentiate it we get the result we want $\langle \psi \vert H \vert \partial \psi\rangle$.
For the second new term in V now, there are two approaches to add it. The first one is to add the corresponding term as the very last line rhs += blahblah. The second one, as you do now, adding terms within energy function which will be differentiated. Then for every wavefunction vector except one, we need to wrap them with tc.backend.stop_gradient, so that they are not differentiated against grad. Namely, wl0, wr should first be wrapped with stop_gradient to make the gradient correct for V term.

[gaox-qd]

Thanks for your reply.
I tried it again. It looks the result becomes better but not good either. BTW, the original code in the example works perfectly.

Here is the updated code:

@tc.backend.jit
def lhs_matrix(theta,psi0):
    
    def _vdot(i, j):
        return tc.backend.tensordot(tc.backend.conj(i), j, 1)

    def ij(i, j):
        """
        Inner product
        """
        # return _vdot(i, j) - _vdot(i, psi) * _vdot(psi, j)
        return _vdot(i, j) + _vdot(i, psi0) * _vdot(j, psi0)

    vij = tc.backend.vmap(ij, vectorized_argnums=0)
    vvij = tc.backend.vmap(vij, vectorized_argnums=1)
    jacobian = ppsioverptheta(theta, psi0=psi0)
    jacobian = tc.backend.transpose(jacobian)
    fim = vvij(jacobian, jacobian) 
    fim = tc.backend.real(fim)
    return fim


@tc.backend.jit
def rhs_vector(theta, psi0):
    def energy(theta, psi0):
        w = variational_wfn(theta, psi0)
        # wl = w
        wl = tc.backend.conj(w)
        wr = tc.backend.stop_gradient(w)
        wl = tc.backend.reshape(wl, [1, -1])
        wr = tc.backend.reshape(wr, [-1, 1])
        e = wl @ h @ wr 

        wl0 = tc.backend.stop_gradient(w)
        wl0 = tc.backend.conj(wl0)
        wl0 = tc.backend.reshape(wl0, [1, -1])
        e0 = wl0 @ h @ wr # <psi0| H | psi0>
        w0 = wl @ wr #<\partial psi0| psi0>
        return tc.backend.real(e-w0*e0)[0, 0]
    
    eg = tc.backend.grad(energy, argnums=0)
    rhs = eg(theta, psi0)
    rhs = tc.backend.imag(rhs)
    return rhs

[refraction-ray]

I suggest you split the two terms calculation in rhs vector, (say two energy function and two gradients, and please align the $\partial$ on bra or ket, make them consistent with the formula). Because according to the formula，only the first part $C$ is taken imaginary part. For current implementation, it seems to me both part are taken imag as the very last line indicates.

The new implementation for rhs would be something like:

def energy1(theta, psi0):
        w = variational_wfn(theta, psi0)
        # wl = w
        wl = tc.backend.conj(w)
        wr = tc.backend.stop_gradient(w)
        wl = tc.backend.reshape(wl, [1, -1])
        wr = tc.backend.reshape(wr, [-1, 1])
        e = wl @ h @ wr 
        return tc.backend.real(e)[0, 0]

# note I dont check the consistency of partial position 
def energy2(theta, psi0)
        w = variational_wfn(theta, psi0)
        wr = tc.backend.reshape(w, [-1, 1])
        wl0 = tc.backend.stop_gradient(w)
        wl0 = tc.backend.conj(wl0)
        wl0 = tc.backend.reshape(wl0, [1, -1])
        e0 = wl0 @ h @ w # <psi0| H | psi0>
        w0 = wl @ wr #<\partial psi0| psi0>
        return w0*e0[0, 0]
 
eg1 = tc.backend.grad(energy1, argnums=0)
eg2 = tc.backend.grad(energy2, argnums=0)

rhs1 = eg1(theta, psi0)
rhs1 = tc.backend.imag(rhs)
rhs2 = eg2(theta, psi0)
rhs = rhs1 + rhs2

If the above template works, I can try to think about better and faster way to implement the generalized rhs, since there seems to be some redundancy calculation above.

And possible other resources for the mismatch:

1. maybe we need a smaller time step
2. I am not sure about the global phase problem and how does it matter or when should we use the generalized formula. Since in usual case, the simple equation works fine.

[gaox-qd]

Hi Shixin, thanks for your suggestion.
I tried the following code. Unfortunately, the results are the same as my last post.

@tc.backend.jit
def lhs_matrix(theta,psi0):
    
    def _vdot(i, j):
        return tc.backend.tensordot(tc.backend.conj(i), j, 1)

    def ij(i, j):
        """
        Inner product
        """
        # return _vdot(i, j) - _vdot(i, psi0) * _vdot(psi0, j)
        return _vdot(i, j) + _vdot(i, psi0) * _vdot(j, psi0)
        #the above two expressions shoud give the same results,
        #since <\partial psi0| psi0> is imaginary,
        #while they're very different from each other!


    vij = tc.backend.vmap(ij, vectorized_argnums=0)
    vvij = tc.backend.vmap(vij, vectorized_argnums=1)
    jacobian = ppsioverptheta(theta, psi0=psi0)
    jacobian = tc.backend.transpose(jacobian)
    fim = vvij(jacobian, jacobian) 
    fim = tc.backend.real(fim)
    return fim


@tc.backend.jit
def rhs_vector(theta, psi0):
    def energy1(theta, psi0):
        w = variational_wfn(theta, psi0)
        # wl = w
        wl = tc.backend.conj(w)
        wr = tc.backend.stop_gradient(w)
        wl = tc.backend.reshape(wl, [1, -1])
        wr = tc.backend.reshape(wr, [-1, 1])
        e = wl @ h @ wr #<\partial psi0|H| psi0>
        return tc.backend.real(e)[0, 0]

    
    def energy2(theta, psi0):
        w = variational_wfn(theta, psi0)
        wr0 = tc.backend.stop_gradient(w)
        wr0 = tc.backend.reshape(wr0, [-1, 1])
        wl0 = tc.backend.stop_gradient(w)
        wl0 = tc.backend.conj(wl0)
        wl0 = tc.backend.reshape(wl0, [1, -1])
        e0 = wl0 @ h @ wr0 # <psi0| H | psi0>

        wl = tc.backend.conj(w)
        wl = tc.backend.reshape(wl, [1, -1])

        w0 = wl @ wr0 #<\partial psi0| psi0>
        return w0*e0[0, 0]
    
    eg1 = tc.backend.grad(energy1, argnums=0)
    eg2 = tc.backend.grad(energy2, argnums=0)

    rhs1 = eg1(theta, psi0)
    rhs1 = tc.backend.imag(rhs1) 
    rhs2 = eg2(theta, psi0) #should be a imaginary number
    rhs = rhs1 + tc.backend.real(1.0j*rhs2) 
    # rhs = rhs1 + tf.cast(1.0j*rhs2, np.float64) 
    return rhs

1. maybe we need a smaller time step
   The smaller time step, such as changed from 0.005 to 0.001, doesn't improve the results.

2. I am not sure about the global phase problem and how does it matter or when should we use the generalized formula. Since in usual case, the simple equation works fine.
   I agree with you. Here I only would like to check whether they both give the same results.

Another point on lhs_matrix function, I'm not sure:
return _vdot(i, j) - _vdot(i, psi0) * _vdot(psi0, j) should give the same as return _vdot(i, j) + _vdot(i, psi0) * _vdot(j, psi0) since <\partial psi0| psi0> is imaginary, while their results are very different from each other.

[refraction-ray]

the main issue of the current implementation is in ij, we need psi instead of psi0

[refraction-ray]

Hi, based on our discussion, I have now a correct implementation (hopefully), please check the script here: https://github.com/tencent-quantum-lab/tensorcircuit/blob/master/examples/variational_dynamics_generalized.py

[refraction-ray]

In terms of speed, jacfwd should be must faster than jacrev since the output dimension (wavefunction) is larger than the input dimension (circuit parameters). In terms of the results, I checked that they give the same Jacobian as long as the input and the output are both complex dtypes, which is the case in our script.

[gaox-qd]

I currently can not use jax backend.
Could you please check the results for jacrev when using tensorflow backend?

[refraction-ray]

Thanks for your suggestion. You are right, results from jacrev and jacfwd are different upto complex conjugate on tensorflow backend. This can be due to the different differentiation rules for complex-valued function between jax and tf. I'll further investigate this mismatch

[refraction-ray]

TF handles complex calculus different from jax: see this issue: tensorflow/tensorflow#3348 (where I also have discussion). And see this new example script for jacobian difference: https://github.com/tencent-quantum-lab/tensorcircuit/blob/master/examples/jacobian_cal.py

[gaox-qd]

Great, thanks.