Basic iterator interface for optimization drivers

[femtobit]

This PR contains a first version of the iterator interface for the VMC (and exact imaginary time) driver, addressing #107.
(The implementation is partially modeled after the description here.)

Basically, it allows to write code such as

VariationalMonteCarlo vmc(...);
for (const auto& step : vmc.Iterate()) {
    outfile << 'step=' << step.index << '\n';
    outfile << json(step.observables) << '\n';
}

In principle, it is possible to modify the state, Hamiltonian, or other objects in between iteration steps (though this is not done anywhere right now).
The Python bindings expose this function as well, allowing, e.g.,

vmc = nk.gs.vmc.Vmc(...)
for i, step in enumerate(vmc.iter()):
    print(step.observables.Energy)

In order to expose step.observables, Python bindings for the ObsManager class are provided, with an interface similar to a python dict (similar enough to make d = dict(step.observables) work).

Some open issues (or potentially points to discuss):

1. The (VMC) iterators dereference to a struct containing basic information on the state:
   

   netket/NetKet/GroundState/variational_montecarlo.hpp

   Lines 98 to 103 in a5c6d62

   	struct Step {
   	Index index;
   	Eigen::VectorXd acceptance;
   	ObsManager observables;
   	nonstd::optional<Eigen::VectorXcd> parameters;
   	};

   
   This contains a copy of the data (as opposed to referencing the corresponding internal state of the VariationalMonteCarlo class) in order to make it possible to write code such as
   steps = list(vmc.iter(100)).
   This has the downside of copying this data every step, regardless of whether it is needed. For the machine parameters, an option is provided, so that step.parameters can be left empty (== nonstd::nullopt) if storage of the machine parameters is not needed.
2. Extracting the data from ObsManager still performs an MPI reduction which has to be performed on all MPI processes, leading to code such as
   

   netket/Tutorials/PyNetKet/ground_state_iter.py

   Lines 40 to 45 in 00a7c60

   	for i, st in enumerate(vmc.iter()):
   	obs = dict(st.observables) # TODO: needs to be called on all MPI processes
   	if mpi_rank == 0:
   	print("step={}".format(i))
   	print("acceptance={}".format(list(st.acceptance)))
   	print("observables={}\n".format(obs))

   
   The code should probably be changed so that step.observables does not perform MPI calls. The question is whether it should contain sensible values only on rank 0 or on all MPI processes (the first option should be sufficient but might also lead to annoying-to-debug errors.)
3. The VariationalMonteCarlo::Run function is provided which is now implemented in terms of the iterator interface. This can be used to get essentially the behaviour of NetKet v1:
   

   netket/Tutorials/PyNetKet/ground_state.py

   Lines 42 to 50 in 00a7c60

   	vmc = nk.gs.vmc.Vmc(
   	hamiltonian=ha,
   	sampler=sa,
   	optimizer=op,
   	n_samples=1000,
   	diag_shift=0.1,
   	method='Sr')
   	vmc.run(filename_prefix="test", max_steps=300, save_params_every=10)

   We can discuss whether we want to provide this function or maybe a different kind of simplified interface.
4. This PR also contains a commit introducing a header for common type aliases (related to Add standard typedefs for matrices and vectors #86). For now, it contains Complex = std::complex<double> (because I think it is cumbersome to always spell-out that type) and Index. The Index type is added because right now there are several types used as indices (usually int or std::size_t) in different places in the netket codebase. I suggest standardizing on Index = std::ptrdiff_t (which is a 64-bit integer on my system). (Whether signed or unsigned types should be used as array indices is a somewhat controversial question in C++, but note that newer additions to the standard such as std::span uses ptrdiff_t. This is also in line with what the Google Style Guide has to say about unsigned integers.)

Let me know what you think.

[gcarleo]

Thanks a lot @femtobit, this is excellent work, and the interface is very cool!

Just let me add a few comments here, before going into the review:

The (VMC) iterators dereference to a struct containing basic information on the state:

Just to clarify: if I do something like

for st in vmc.iter(10, 2):
   obs = st.observables
   ... 

are the observables computed at every step or only every 2 steps? (The energy needs anyway to be computed at every step, but maybe we can trigger the calculation of other observables, as added with add_observable only every n steps (2 in this case) ?

On a related note, I think we could just remove parameters from the step object entirely, since those are a property of the machine. In practice you can always do print(ma.parameters) to retrieve them if you want at any point. The same is true for the acceptance, which is a property of the sampler, I would remove that too. I think we can just keep the observables in the stepper?

I understand that this would break the possibility of retrieving the parameters/acceptance doing steps = list(vmc.iter(100)), but in practice the later usage of iter doesn't seem to me very conventional, since in any case one can expect the single iterations to be time consuming.

The code should probably be changed so that step.observables does not perform MPI calls.

I think that, if possible, calling step.observables should trigger itself the calculation of the observables from the stored visible units in the Vmc (see previous comment). However I am not sure it is possible at all to remove MPI calls from there, since MPI is used to collect the averages from different processors.

[gcarleo]

@femtobit my main comments are about naming, which basically means that the PR is great :)

[gcarleo]

I think that netket.gs.vmc.Vmc is too deep of a module hierarchy, we should try to simplify it.
netket.Vmc doesn't look bad to me, but let's discuss more

[femtobit]

I think it would make sense to separate the VMC code from the ED code. So netket.vmc and netket.ed seem fine to me.

[femtobit]

Although strictly speaking ed is a bit of a misnomer for exact time evolution.

[gcarleo]

We can do netket.exact.TimeEvolution ,netket.exact.SparseDiag , netket.exact.FullDiag?

It's just that netket.vmc.Vmc seems redundant to me, since the vmc module would consist only of the Vmc class, thus I'd rather do netket.Vmc, much simpler

[femtobit]

Yes, netket.exact sounds good.

Regarding netket.vmc: I see your point, however, we might get netket.TdVmc in the future. (In principle, we could also call it netket.vmc.TimeEvolution or similar.) Also, there is netket.vmc.Iterator (which users would usually not see, but has to be somewhere).

[gcarleo]

For consistency we should name this machine and not psi (my fault, I know)

[gcarleo]

same here, we could do netket.ed.ImagTimePropagation and netket.ed.SparseDiag etc?

[gcarleo]

I think this file can be just removed

[gcarleo]

maybe filename_prefix should be output_prefix or just output as before?

[gcarleo]

do we need to explicitly convert to dict here?

[femtobit]

Currently yes, because while the ObsManager mostly works like a dict, printing it does not give you all the data.

[gcarleo]

I think that what could be useful is to do

vmc.add_observable(ma.parameters)
vmc.add_observable(sa.acceptance)

so that one can chose what optional things to measure.

One could even do

vmc.add_observable(datetime.datetime.now)

that would print the time at which the iteration has been completed, really cool!

This should be technically feasible overloading add_observable to take a callable function?

[femtobit]

On a related note, I think we could just remove parameters from the step object entirely, since those are a property of the machine. In practice you can always do print(ma.parameters) to retrieve them if you want at any point. The same is true for the acceptance, which is a property of the sampler, I would remove that too. I think we can just keep the observables in the stepper?

I understand that this would break the possibility of retrieving the parameters/acceptance doing steps = list(vmc.iter(100)), but in practice the later usage of iter doesn't seem to me very conventional, since in any case one can expect the single iterations to be time consuming.

Yes, I think this is a valid option.

I think that what could be useful is to do
vmc.add_observable(ma.parameters)
vmc.add_observable(sa.acceptance)
so that one can chose what optional things to measure.

One could even do
vmc.add_observable(datetime.datetime.now)

Looks fun but I'd prefer to keep the observables restricted to actual quantum mechanical observables, i.e., operators and do

for step in vmc.iter():
    date = datetime.datetime.now()
    # ... etc ...

What bothers me a bit is that in principle, all the information can be obtained from elsewhere, so there is no real need for the step object. (The observables can be currently only retrieved from step, but adding something like sampler.mean(operators) is possible.) But then again, in a for-loop one should iterate over something.

are the observables computed at every step or only every 2 steps? (The energy needs anyway to be computed at every step, but maybe we can trigger the calculation of other observables, as added with add_observable only every n steps (2 in this case) ?

Good point. I think they are calculated every step right now. I will change that to compute them only when necessary.

[gcarleo]

Looks fun but I'd prefer to keep the observables restricted to actual quantum mechanical observables, i.e., operators and do

Indeed that's true, that would better be some add_to_log method or similar rather than add_observable that should strictly be for quantum observables

What bothers me a bit is that in principle, all the information can be obtained from elsewhere, so there is no real need for the step object. (The observables can be currently only retrieved from step, but adding something like sampler.mean(operators) is possible.) But then again, in a for-loop one should iterate over something.

True, (to retrieve observables one could also just do vmc.observables or something, by the way) but using these add_xx methods is useful if one wants to just use vmc.run(..) and record custom things during the run/dump them somewhere on a file.

But I agree that the iterator is convenient, yes, but maybe just a range(n_iter) would have worked :) (see also my original comment in #107 )

[femtobit]

But I agree that the iterator is convenient, yes, but maybe just a range(n_iter) would have worked :) (see also my original comment in #107 )

Of course, it is essentially equivalent to something like

for i in range(1000):
    vmc.advance(step_size)
    print(vmc.observables)

Maybe we should have implemented the iterators in pure python and just provide the Advance function from C++. Then, implementing the iterator is a matter of writing something like

def iter_obs(vmc):
    while True:
        vmc.advance()
        yield vmc.observables()

for obs in iter_obs(vmc):
    print(obs)

But well, in the end the interface is similar and providing control over the simulation to the surrounding code (which was the main goal) can be achieved either way.

For now, I suggest we keep add_observables the way it is and let step contain the corresponding expectation values and if more information is needed, users can obtain it from within the loop body. (Adding convenience functions around that later is possible without too much effort.)

[gcarleo]

Thanks @femtobit , just a reminder to discuss later about adding the observables as a member of vmc, and leave the vmc iterator just as an index counting to advance the iterations

[femtobit]

Thanks @femtobit , just a reminder to discuss later about adding the observables as a member of vmc, and leave the vmc iterator just as an index counting to advance the iterations

This is done now:

netket/Tutorials/PyNetKet/ground_state_iter.py

Lines 40 to 42 in 3ddd9a2

	for step in vmc.iter():
	obs = vmc.get_observable_stats()
	if mpi_rank == 0:

The observable stats are only computed if the method is called. I chose to use this a bit more verbose function name (instead of, say, a property vmc.observables) because this call actually performs computation and the MPI reduction.

[femtobit]

Merged into feature branch, so that this code can serve as base for other PRs