Modified time evolution of $U|\psi\rangle$ with TDVP

[jackaraz]

Hi, I’m trying to figure out how to run a modified version of TDVP. I need to evolve a MCState ($|\psi\rangle$) with respect to a Hamiltonian $H$, but instead of evolving only $|\psi\rangle$, I want to evolve $U|\psi\rangle$ under $H$, i.e., compute $e^{-iHT}U|\psi\rangle$, where $U$ is a unitary operator and $[U, H] \neq 0$.

I see that odefun_tdvp calls state.expect_and_forces, which I believe computes the expectation value $\langle \psi | H | \psi \rangle$ and the corresponding forces. Would it make sense to modify the expect_and_forces function to incorporate the unitary operator $U$ together with the Hamiltonian? I’m concerned that doing so might effectively compute $e^{-iHUT}$, which is not what I want.

I also thought about adding U to the hamiltonian during state preparation but since I'll need to change where U is acting on at every step, finding the state would become quite time consuming.

Any guidance on how to approach this correctly would be greatly appreciated.

Thanks
Best regards

[PhilipVinc]

The most straightforward thing you can do is to simply define the new variational state $$\ket{U\psi}$$ and use it.

A way to do it is the following:

...
nn_model = nk.model.RBM()
vs = nk.vqs.MCState(sampler, nn_model)
U = nk.operator.spin.sigmax(hi, 1).to_jax_operator()

Unn_fun, Unn_variables = make_logpsi_U_afun(vs._apply_fun, vs.variables, U)
Uvs = nk.vqs.MCState(sampler, apply_fun=Unn_fun, variables=Unn_variables)

the relevant function can be found in here
https://gist.github.com/PhilipVinc/5fcc9fbba0e5d6ed6d56f0a4325e2808

Then everything will work out of th box.
Be careful that when you create a new MCState, the sampler chains will be initialised completely at random, so if you do this at the beginning it won't be a problem, but if you do it with a state that has already been initialised you might have to carry over your old chains by doing something lke

Uvs.sampler_state = Uvs.sampler_state.replace(σ= vs.sampler_state.σ)

(or possibly something more sophisticated).

Notice that this approach will increase the cost of your calculations by a factor of U.max_conn_size, but will do exactly what you want.

There are other schemes to do sort of what you want, but at a lower cost, for example by importance sampling the forces, but nothing that has been carefully ironed out. In a different setting, we discuss them in https://arxiv.org/pdf/2410.10720 (look for importance sampling)

[jackaraz]

Thanks @PhilipVinc this was very useful!