Time-evolution: new design

Created by: PhilipVinc

Now that everything is pythonically ~~ugly~~ nice, writing a variational time integrator that plugs into scipy is quite easy. I would like to take the occasion to unify the interface among time evolution and standard drivers, so this is an issue to discuss this.

Requirements:

The interface should be consistent with the vmc optimization drivers as much as possible
The interface should support open and hamiltonian evolution, otherwise people will break my beloved lindblad without me noticing
Should reuse other package's time-integrators. I am lazy and don't like to for the millionth time a RK integrator
Should be somewhat easy to plug another package integrator in, if needed.

Wishes:

I like the idea of it being easy to first define a problem (machine, initial state, time interval) and then solve it repeatedly with different solvers/timesteps/options. For anyone, this is essentially the DiffEq/SciML interface, where you first define a problem and then you solve it. I don't know if this is feasible but that would be nice.

f(u,p,t) = 1.01*u
u0 = 1/2
tspan = (0.0,1.0)
prob = ODEProblem(f,u0,tspan)
sol = solve(prob, rk45(), reltol=1e-8, abstol=1e-8, dt=0.001)

Conceptually, this means:

Define as usual the machine, the sampler
You MUST define the SR object at some point, because it is always required
Define the (potentially time dependent) hamiltonian/lindblad you wish to integrate.
Define the integrator, and all it's related options (in order to be generic, those arguments should be forwarded...)
Define the interval in time we wish to integrate upon.

possible implementation: in order to be somewhat portable, we need to be able to output a lambda function that takes a flattened vector of parameters and computes the gradient, which is equivalent to defining the function that defines an ode:

def create_odefun(op, sampler, *args, **kwargs):
    machine = sampler._machine
    if not isinstance(machine, AbstractDensityMatrix):
        driver = Vmc(op, sampler, *args, optimizer=None, **kwargs)
    else:
        driver = SteadyState(op, sampler, *args, optimizer=None, **kwargs)

    def _fun(t, w):
        machine.parameters = unflatten(w)
        dp = driver._time_evo_step()
        return flatten(dp)

   return _fun, driver

It would be a worthwhile discussion whever this time_evo_step method should live in vmc.py and steadystate.py or it should be a separate class, or just a standalone method that does dispatch.

Having this function I can already plug into scipy, diffeqpy or any other integrator module easily. Everything else is fancy logging and parameter handling. My main preoccupation for this is that we want to compute observables only ever once in a while (otherwise it really gets slow fast). And unfortunately scion does not allow callback hooks like that.

Since we like to be fancy, we would like to have an higher level TimeEvolution driver. An initial proposal might be something along the lines of the following.

H = ...
ma = nk.machine.rbmspin(hilbert=hi, alpha=1)
ma.parameters = my_physical_state
sa = nk.sampler.MetropolisLocal(machine=ma, n_chains=8)
sr = nk.optimizer.SR(diag_shift=0.001, use_iterative=True, sparse_tol=1e-6, sparse_maxiter=1000)

tevo = TimeEvolution(H, sa, sr=sr, n_samples=2000) # incase of a density matrix also, sampler_obs=sa_obs, n_samples_obs=500)
tevo.solver("rk45", t_0=(1.0, 2.0), dt=0.01, adaptive=False)

tevo.run(1.0, out="test", obs=obs, step_size=0.02) # goes from 1.0 to 2.0, saving logs every 0.02
 48%|██▎      | 0.4800/1.0 [05:28<18:19, 723.68s/time, Energy=-25.3 ± 0.01 [var=4.6e+01, R_hat=0.9993]] 
pars_end = copy(ma.parameters)


# now if I want to test out a different tilmestep I should do...
ma.parameters = my_physical_state
tevo.solver("rk45", t_0=1.0, dt=0.005, adaptive=False)
tevo.run(1.0, out="test", obs=obs, step_size=0.02) # goes from 1.0 to 2.0, saving logs every 0.02
...

# that's better! now I continue my integration to t=3.0
tevo.run(1.0, out="test", obs=obs, step_size=0.02) # goes from 2.0 to 3.0, saving logs every 0.02
...

I already implemented the above, and it's nice that it uses the same run interface of Abstractdriver with only minimal changes. Ill admit this is not the best interface ever. But it allows to have a nice driver interface while also being able to call iterate manually for only a few iterations if one so whishes.

One thing I don't like is that if you randomly interrupt this, you will have random parameters for a random Time stored in the machine you have referenced outside... but this is a core design of netket, so I'm not sure we can do much about it.

Alternatively we might more go for something like

H = ...
ma = nk.machine.rbmspin(hilbert=hi, alpha=1)
ma.parameters = my_physical_state
sa = nk.sampler.MetropolisLocal(machine=ma, n_chains=8)
sr = nk.optimizer.SR(diag_shift=0.001, use_iterative=True, sparse_tol=1e-6, sparse_maxiter=1000)

tevo_prob = TimeEvolutionProblem(H, sa, time_interval, sr=sr, n_samples=2000)
final_state, logged_obs = solve(tevo_prob, "rk45", dt=0.01, adaptive=False, obs=obs)

# want to try out other dt
final_state_dt2, logged_obs_dt2 = solve(tevo_prob, "rk45", dt=0.005, adaptive=False, obs=obs)

# that's better! now let's continue the time evolution...
tevo_prob_2 = TimeEvolutionProblem(H, sa, time_interval_2, sr=sr, n_samples=2000, state= final_state_dt2)
# state= will simply assign to the machine..
final_state_3, logged_obs_3 = solve(tevo_prob_2, "rk45", dt=0.005, adaptive=False, obs=obs)

I like this because it splits conceptually the time integrator (the solver) from the definition of the problem itself. However the first one is simpler to use in general if you don't play around much with solving the same thing with different solvers, but only integrate once.

In my experience when working with time evolution one does many tests at first, but then settles for one thing.... So i dunno. What do you think?

Miscellaneous stuff:

SciPy does not support fixed time-step integration!
- kinda. But I've been playing around recently with it, and if you set abstol, reltol= ∞ and max_dt, effectively you are fixing the dt to that value. It works well.
I don't like your jokes
- Meh. me neither.