How to use qgt in t-vmc

[whatisspin]

Hi,
I am trying to do the time evolution for the Ising Hamiltonian (from the netket paper.)

import netket as nk
import netket.nn as nknn
import flax.linen as nn
import jax.nn as jnn
import jax.numpy as jnp

L = 10
chain = nk.graph.Chain(L)
hi = nk.hilbert.Spin(s=1/2) ** L
sampler = nk.sampler.MetropolisLocal(hi)

class FullModel(nknn.Module):
    
    alpha: int
    
    @nn.compact
    def __call__(self,x):
        rho= nknn.Dense(features= self.alpha*L, dtype=float, kernel_init= jnn.initializers.normal(stddev=0.01), bias_init=jnn.initializers.normal(stddev=0.01))(x)
        rho= nknn.relu(rho)       
        
        phase= nknn.Dense(features= self.alpha*L, dtype=float, kernel_init= jnn.initializers.normal(stddev=0.01), bias_init=jnn.initializers.normal(stddev=0.01))(x)
        phase= nknn.relu(phase)

        x= jnp.sum(rho, axis=(-1))+(1.0j*jnp.sum(phase, axis=(-1)))
        return x
model= FullModel(alpha=1)

rbm = nk.models.RBM(alpha=1, dtype=complex)

vstate = nk.vqs.MCState(sampler, model, n_samples=4096)

ha0 = nk.operator.Ising(hi, chain, h=0.5)

optimizer = nk.optimizer.Sgd(0.01)

sr = nk.optimizer.SR()

print(vstate.n_parameters)
vmc = nk.VMC(ha0, optimizer, variational_state=vstate, preconditioner=sr)

vmc.run(300)

def custom_simple_tdvp(
    hamiltonian, # Hamiltonian
    vstate, # variational state
    t0: float, # initial time
    dt: float, # time step
    t_end: float, # end time
):
    t = t0
    while t < t_end:
        energy, f = vstate.expect_and_grad(hamiltonian)
        G = vstate.quantum_geometric_tensor(qgt_T=nk.optimizer.qgt.QGTJacobianDense())
        gamma_f = jax.tree_map(lambda x: -1.0j * x, f)
        dtheta, _ = G.solve(jax.scipy.sparse.linalg.cg, gamma_f)
        vstate.parameters = jax.tree_multimap(
            lambda x, y: x + dt * y, vstate.parameters, dtheta
        )
        t = t + dt
        print(t)

ha1 = nk.operator.Ising(hi, chain, h=1.0)
te_custom= custom_simple_tdvp(ha1,vstate,0.0,0.001,1.0)

While this works for the rbm, for the model it gives the following error

TypeError: dot_general requires contracting dimensions to have the same shape, got (220,) and (440,).

As I understand, the (220,) corresponds to the number of parameters. But why is the qgt twice of that?

So, while it works with complex parameters, what would be the correct way to use qgt when using real parameters?

[PhilipVinc]

Hi, I'm on holiday so I'm not going to be very responsive until the 16th (in the meantime @femtobit might help if he's not on holiday as well).

The thing you are trying to run is actually unsupported by netket (and we discuss it in the paper).
Current released netket can only evolve networks with complex parameters.
If you have real parameters, it cannot work because we are computing the (real values) gradient of the energy instead of the (complex valued) forces vector.

There's a PR open (#1261 ) to fix it, which will be merged soon after I'm back from holiday. It might contain some discussion on why it does not work right now.
If you use the version from that PR, you should replace calls to expect_and_grad with calls to expect and forces and then take the imaginary part of it.

So, the reason you see 440 instead of 220 is because internally in some QGT implementation we split in real and imaginary part. In your case, we are splitting the complex vector gamma_f in a double-sized real vector, but your parameters are real and this leads to error.

For a bunch of reasons (if you are curious, I can discuss them when I'm back) we did not want to support this operation (in short, it would generate the wrong dynamics). We should do a better job of throwing a more informative error rather than this one.

[whatisspin]

Thank you for your response. The different dimensions make sense now.
I also found the newer version of the paper with more information, thanks.

[PhilipVinc]

Hey @whatisspin , we just released a new version (3.5) that works with real/non holomorphic parameters as well.
I'm writing a blog post which will go up in a few days. In the meantime, you can use this snippet:

def custom_simple_tdvp(
    hamiltonian: AbstractOperator,  # Hamiltonian
    vstate: VariationalState,       # variational state
    t0: float,                      # initial time
    dt: float,                      # time step
    t_end: float,                   # end time
):
    t = t0
    while t < t_end:
        # compute the forces `f` generated by the Hamiltonian
        energy, f = vstate.expect_and_forces(hamiltonian)
        G = vstate.quantum_geometric_tensor()
        # multiply the gradient by -1.0j for unitary dynamics
        gamma_f = jax.tree_map(lambda x: -1.0j * x, f)
        # take the real part of -1.0*f if the parameters are real
        gamma_f = jax.tree_map(
	        lambda x, target: (x if jnp.iscomplexobj(target) else x.real),
	        gamma_f,
	        vstate.parameters,
	    )
        # Solve the linear system using any solver, such as CG
        # (or write your own regularization scheme)
        dtheta, _ = G.solve(jax.scipy.sparse.linalg.cg, gamma_f)
        # update the parameters (theta = theta + dt * dtheta)
        vs.parameters = jax.tree_map(
            lambda x, y: x + dt * y, vstate.parameters, dtheta
        )
        t = t + dt