Basic iterator interface for optimization drivers

.travis.yml

@@ -155,5 +155,5 @@ install:
 script:
   - python -c 'from mpi4py import MPI'
   - python -c 'from mpi4py import MPI; import netket; g = netket.graph.Hypercube(4, 1); h = netket.hilbert.Spin(g, 0.5); m = netket.machine.RbmSpin(h, 10)'
-  - python -m pytest --verbose
+  - python -m pytest --verbose Test/
 

CMakeLists.txt

@@ -134,13 +134,17 @@ endif()
 # option(BENCHMARK_DOWNLOAD_DEPENDENCIES "" ON)
 # add_subdirectory(External/benchmark)
 
-
 if(NOT CMAKE_BUILD_TYPE)
   message(STATUS "[NetKet] CMAKE_BUILD_TYPE not specified, setting it to "
                  "Release. Use `-DCMAKE_BUILD_TYPE=...` to overwrite.")
   set(CMAKE_BUILD_TYPE Release)
 endif()
 
+# Set -O3 in debug mode unless explicitly disabled
+if((${CMAKE_BUILD_TYPE} STREQUAL "Debug") AND NOT NETKET_DISABLE_OPTIMIZATION)
+    set(CMAKE_CXX_FLAGS  "${CMAKE_CXX_FLAGS} -O3")
+endif()
+
 if(MSVC)
   # Force to always compile with W4
   if(CMAKE_CXX_FLAGS MATCHES "/W[0-4]")

NetKet/GroundState/exact_diagonalization.hpp

@@ -22,6 +22,7 @@
 #include <ietl/lanczos.h>
 #include <ietl/randomgenerator.h>
 
+#include "common_types.hpp"
 #include "Operator/MatrixWrapper/matrix_wrapper.hpp"
 #include "Operator/operator.hpp"
 
@@ -31,7 +32,7 @@ namespace eddetail {
 
 using eigenvalues_t = std::vector<double>;
 using eigenvectors_t =
-    std::vector<Eigen::Matrix<std::complex<double>, Eigen::Dynamic, 1>>;
+    std::vector<Eigen::Matrix<Complex, Eigen::Dynamic, 1>>;
 
 struct result_t {
   eigenvalues_t eigenvalues;
@@ -42,8 +43,7 @@ struct result_t {
 template <class matrix_t, class iter_t, class random_t>
 result_t lanczos_run(const matrix_t& matrix, const random_t& random_gen,
                      iter_t& iter, int which_eigenvector = -1) {
-  using complex = std::complex<double>;
-  using vectorspace_t = ietl::vectorspace<complex>;
+  using vectorspace_t = ietl::vectorspace<Complex>;
   using lanczos_t = ietl::lanczos<matrix_t, vectorspace_t>;
 
   size_t dimension = matrix.Dimension();
@@ -105,7 +105,7 @@ eddetail::result_t lanczos_ed(const AbstractOperator& hamiltonian,
 eddetail::result_t full_ed(const AbstractOperator& hamiltonian, int first_n = 1,
                            bool get_groundstate = false) {
   using eigen_solver_t =
-      Eigen::SelfAdjointEigenSolver<Eigen::SparseMatrix<std::complex<double>>>;
+      Eigen::SelfAdjointEigenSolver<Eigen::SparseMatrix<Complex>>;
 
   SparseMatrixWrapper<> matrix(hamiltonian);
 

NetKet/GroundState/imaginary_time.hpp

@@ -9,29 +9,26 @@
 #include "Output/json_output_writer.hpp"
 #include "Stats/stats.hpp"
 
-// TODO remove Observable and replace with AbstractOperator+name
-// Provide a method AddObservable, as in VariationalMonteCarlo
 namespace netket {
 
-class ImaginaryTimeDriver {
+class ImagTimePropagation {
  public:
-  using State = Eigen::VectorXcd;
-  using Stepper = ode::AbstractTimeStepper<State>;
+  class Iterator;
+
+  using StateVector = Eigen::VectorXcd;
+  using Stepper = ode::AbstractTimeStepper<StateVector>;
   using Matrix = AbstractMatrixWrapper<>;
 
   using ObsEntry = std::pair<std::string, std::unique_ptr<Matrix>>;
-  using ObservableVector = std::vector<ObsEntry>;
 
-  ImaginaryTimeDriver(Matrix& matrix, Stepper& stepper,
-                      JsonOutputWriter& output, double tmin, double tmax,
-                      double dt)
+  ImagTimePropagation(Matrix& matrix, Stepper& stepper, double t0,
+                      StateVector initial_state)
       : matrix_(matrix),
         stepper_(stepper),
-        output_(output),
-        range_(ode::TimeRange{tmin, tmax, dt}) {
-    ode_system_ = [this](const State& x, State& dxdt, double /*t*/) {
-      dxdt.noalias() = -matrix_.Apply(x);
-    };
+        t_(t0),
+        state_(std::move(initial_state)) {
+    ode_system_ = [this](const StateVector& x, StateVector& dxdt,
+                         double /*t*/) { dxdt.noalias() = -matrix_.Apply(x); };
   }
 
   void AddObservable(const AbstractOperator& observable,
@@ -41,32 +38,28 @@ class ImaginaryTimeDriver {
     observables_.emplace_back(name, std::move(wrapper));
   }
 
-  void Run(State& initial_state) {
-    assert(initial_state.size() == Dimension());
-
-    int step = 0;
-    double t = range_.tmin;
-    while (t < range_.tmax) {
-      double next_dt =
-          (t + range_.dt <= range_.tmax) ? range_.dt : range_.tmax - t;
-
-      stepper_.Propagate(ode_system_, initial_state, t, next_dt);
-
-      // renormalize the state to prevent unbounded growth of the norm
-      initial_state.normalize();
-
-      ComputeObservables(initial_state);
-      auto obs_data = json(obsmanager_);
-      output_.WriteLog(step, obs_data, t);
+  void Advance(double dt) {
+    // Propagate the state
+    stepper_.Propagate(ode_system_, state_, t_, dt);
+    // renormalize the state to prevent unbounded growth of the norm
+    state_.normalize();
+    ComputeObservables(state_);
+    t_ += dt;
+  }
 
-      output_.WriteState(step, initial_state);
+  /*void Run(StateVector& initial_state, ) {
+    assert(initial_state.size() == Dimension());
+    state_ = initial_state;
+    for (const auto& step : Iterate(range.t0, )) {
+    }
+  }*/
 
-      step++;
-      t = range_.tmin + step * range_.dt;
-    };
+  Iterator Iterate(double dt,
+                   nonstd::optional<Index> max_steps = nonstd::nullopt) {
+    return Iterator(*this, dt, std::move(max_steps));
   }
 
-  void ComputeObservables(const State& state) {
+  void ComputeObservables(const StateVector& state) {
     const auto mean_variance = matrix_.MeanVariance(state);
     obsmanager_.Reset("Energy");
     obsmanager_.Push("Energy", mean_variance[0].real());
@@ -83,19 +76,64 @@ class ImaginaryTimeDriver {
     }
   }
 
-  int Dimension() const { return matrix_.Dimension(); }
+  const ObsManager& GetObsManager() const { return obsmanager_; }
+
+  double GetTime() const { return t_; }
+  void SetTime(double t) { t_ = t; }
+
+  class Iterator {
+   public:
+    // typedefs required for iterators
+    using iterator_category = std::input_iterator_tag;
+    using difference_type = Index;
+    using value_type = Index;
+    using pointer_type = Index*;
+    using reference_type = Index&;
+
+   private:
+    ImagTimePropagation& driver_;
+    nonstd::optional<Index> max_iter_;
+    double dt_;
+
+    Index cur_iter_;
+
+   public:
+    Iterator(ImagTimePropagation& driver, double dt,
+             nonstd::optional<Index> max_steps)
+        : driver_(driver),
+          max_iter_(std::move(max_steps)),
+          dt_(dt),
+          cur_iter_(0) {}
+
+    Index operator*() const { return cur_iter_; };
+    Iterator& operator++() {
+      driver_.Advance(dt_);
+      cur_iter_ += 1;
+      return *this;
+    }
+
+    // TODO(C++17): Replace with comparison to special Sentinel type, since
+    // C++17 allows end() to return a different type from begin().
+    bool operator!=(const Iterator&) {
+      return !max_iter_.has_value() || cur_iter_ < max_iter_.value();
+    }
+    // pybind11::make_iterator requires operator==
+    bool operator==(const Iterator& other) { return !(*this != other); }
+
+    Iterator begin() const { return *this; }
+    Iterator end() const { return *this; }
+  };
 
  private:
   Matrix& matrix_;
   Stepper& stepper_;
-  ode::OdeSystemFunction<State> ode_system_;
-
-  ObservableVector observables_;
-  ObsManager obsmanager_;
+  ode::OdeSystemFunction<StateVector> ode_system_;
 
-  JsonOutputWriter& output_;
+  double t_;
+  StateVector state_;
 
-  ode::TimeRange range_;
+  std::vector<ObsEntry> observables_;
+  ObsManager obsmanager_;
 };
 
 }  // namespace netket

NetKet/GroundState/pyground_state.hpp

@@ -30,45 +30,75 @@ namespace py = pybind11;
 namespace netket {
 
 void AddGroundStateModule(py::module &m) {
-  auto subm = m.def_submodule("gs");
+  auto m_exact = m.def_submodule("exact");
+  auto m_vmc = m.def_submodule("variational");
 
-  py::class_<VariationalMonteCarlo>(subm, "Vmc")
+  py::class_<VariationalMonteCarlo>(m_vmc, "Vmc")
       .def(py::init<const AbstractOperator &, SamplerType &,
-                    AbstractOptimizer &, int, int, std::string, int, int,
-                    std::string, double, bool, bool, bool, int>(),
+                    AbstractOptimizer &, int, int, int, const std::string &,
+                    double, bool, bool, bool>(),
            py::keep_alive<1, 2>(), py::keep_alive<1, 3>(),
            py::keep_alive<1, 4>(), py::arg("hamiltonian"), py::arg("sampler"),
-           py::arg("optimizer"), py::arg("n_samples"), py::arg("niter_opt"),
-           py::arg("output_file"), py::arg("discarded_samples") = -1,
+           py::arg("optimizer"), py::arg("n_samples"),
+           py::arg("discarded_samples") = -1,
            py::arg("discarded_samples_on_init") = 0, py::arg("method") = "Sr",
            py::arg("diag_shift") = 0.01, py::arg("rescale_shift") = false,
-           py::arg("use_iterative") = false, py::arg("use_cholesky") = true,
-           py::arg("save_every") = 50)
+           py::arg("use_iterative") = false, py::arg("use_cholesky") = true)
+      .def_property_readonly("machine", &VariationalMonteCarlo::GetMachine)
       .def("add_observable", &VariationalMonteCarlo::AddObservable,
            py::keep_alive<1, 2>())
-      .def("run", &VariationalMonteCarlo::Run);
+      .def("run", &VariationalMonteCarlo::Run, py::arg("output_prefix"),
+           py::arg("max_steps") = nonstd::nullopt, py::arg("step_size") = 1,
+           py::arg("save_params_every") = 50)
+      .def("iter", &VariationalMonteCarlo::Iterate,
+           py::arg("max_steps") = nonstd::nullopt, py::arg("step_size") = 1)
+      .def("get_observable_stats", [](VariationalMonteCarlo &self) {
+        py::dict data;
+        self.ComputeObservables();
+        self.GetObsManager().InsertAllStats(data);
+        return data;
+      });
 
-  py::class_<ImaginaryTimeDriver>(subm, "ImaginaryTimeDriver")
-      .def(py::init<ImaginaryTimeDriver::Matrix &,
-                    ImaginaryTimeDriver::Stepper &, JsonOutputWriter &, double,
-                    double, double>(),
-           py::arg("hamiltonian"), py::arg("stepper"), py::arg("output_writer"),
-           py::arg("tmin"), py::arg("tmax"), py::arg("dt"))
-      .def("add_observable", &ImaginaryTimeDriver::AddObservable,
+  py::class_<VariationalMonteCarlo::Iterator>(m_vmc, "VmcIterator")
+      .def("__iter__", [](VariationalMonteCarlo::Iterator &self) {
+        return py::make_iterator(self.begin(), self.end());
+      });
+
+  py::class_<ImagTimePropagation>(m_exact, "ImagTimePropagation")
+      .def(py::init<ImagTimePropagation::Matrix &,
+                    ImagTimePropagation::Stepper &, double,
+                    ImagTimePropagation::StateVector>(),
+           py::arg("hamiltonian"), py::arg("stepper"), py::arg("t0"),
+           py::arg("initial_state"))
+      .def("add_observable", &ImagTimePropagation::AddObservable,
            py::keep_alive<1, 2>(), py::arg("observable"), py::arg("name"),
            py::arg("matrix_type") = "Sparse")
-      .def("run", &ImaginaryTimeDriver::Run, py::arg("initial_state"));
+      .def("iter", &ImagTimePropagation::Iterate, py::arg("dt"),
+           py::arg("max_steps") = nonstd::nullopt)
+      .def_property("t", &ImagTimePropagation::GetTime,
+                    &ImagTimePropagation::SetTime)
+      .def("get_observable_stats", [](const ImagTimePropagation &self) {
+        py::dict data;
+        self.GetObsManager().InsertAllStats(data);
+        return data;
+      });
+
+  py::class_<ImagTimePropagation::Iterator>(m_exact, "ImagTimeIterator")
+      .def("__iter__", [](ImagTimePropagation::Iterator &self) {
+        return py::make_iterator(self.begin(), self.end());
+      });
 
-  py::class_<eddetail::result_t>(subm, "EdResult")
+  py::class_<eddetail::result_t>(m_exact, "EdResult")
       .def_readwrite("eigenvalues", &eddetail::result_t::eigenvalues)
       .def_readwrite("eigenvectors", &eddetail::result_t::eigenvectors)
       .def_readwrite("which_eigenvector",
                      &eddetail::result_t::which_eigenvector);
 
-  subm.def("LanczosEd", &lanczos_ed, py::arg("operator"),
-           py::arg("matrix_free") = false, py::arg("first_n") = 1,
-           py::arg("max_iter") = 1000, py::arg("seed") = 42,
-           py::arg("precision") = 1.0e-14, py::arg("get_groundstate") = false);
+  m_exact.def("LanczosEd", &lanczos_ed, py::arg("operator"),
+              py::arg("matrix_free") = false, py::arg("first_n") = 1,
+              py::arg("max_iter") = 1000, py::arg("seed") = 42,
+              py::arg("precision") = 1.0e-14,
+              py::arg("get_groundstate") = false);
 }
 
 }  // namespace netket