Initial vector feature added

Author: mrcdr

cpp/pylanczos.cpp

@@ -20,7 +20,12 @@ void declare_template_lanczos(pybind11::module& m, const std::string& suffix) {
       .def_readwrite("find_maximum", &PyLanczosCpp<T>::find_maximum)
       .def_readwrite("eigenvalue_offset", &PyLanczosCpp<T>::eigenvalue_offset)
       .def(pybind11::init<std::function<void(pybind11::array_t<T>, pybind11::array_t<T>)>, size_t, bool, size_t>())
-      .def("run", &PyLanczosCpp<T>::run);
+      .def("run",
+           static_cast<std::tuple<pybind11::array_t<T>, pybind11::array_t<T>, std::vector<size_t>> (PyLanczosCpp<T>::*)(
+               void)>(&PyLanczosCpp<T>::run))
+      .def("run",
+           static_cast<std::tuple<pybind11::array_t<T>, pybind11::array_t<T>, std::vector<size_t>> (PyLanczosCpp<T>::*)(
+               const std::vector<T>&)>(&PyLanczosCpp<T>::run));
 }
 
 template <typename T>

include/pylanczos.hpp

@@ -30,6 +30,21 @@ class PyLanczosCpp : public lambda_lanczos::LambdaLanczos<T> {
 
   // This tuple will be interpreted as multiple-value return (python tuple) in Python.
   std::tuple<pybind11::array_t<T>, pybind11::array_t<T>, std::vector<size_t>> run() {
+    this->lambda_lanczos::LambdaLanczos<T>::init_vector = lambda_lanczos::VectorRandomInitializer<T>::init;
+
+    return this->run_impl();
+  };
+
+  std::tuple<pybind11::array_t<T>, pybind11::array_t<T>, std::vector<size_t>> run(const std::vector<T>& initial_vec) {
+    this->lambda_lanczos::LambdaLanczos<T>::init_vector = [&initial_vec](std::vector<T>& out_vec) {
+      std::copy(initial_vec.begin(), initial_vec.end(), out_vec.begin());
+    };
+
+    return this->run_impl();
+  };
+
+ private:
+  std::tuple<pybind11::array_t<T>, pybind11::array_t<T>, std::vector<size_t>> run_impl() {
     using RT = lambda_lanczos::util::real_t<T>;
     std::vector<RT> eigenvalues;
     std::vector<std::vector<T>> eigenvectors;

pylanczos/pylanczos.py

@@ -116,8 +116,13 @@ def create_custom(mv_mul, n, dtype, find_maximum, num_eigs):
 
         return pylanczos_obj
 
-    def run(self):
+    def run(self, initial_vector=None):
         """Executes the Lanczos algorithm.
+        Parameters
+        ----------
+        initial_vector : numpy.ndarray, optional
+            Initial vector to start the Lanczos iteration.
+            If not specified, a random vector will be used.
 
         Returns
         -------
@@ -130,7 +135,10 @@ def run(self):
         klass = PyLanczos._dtype_to_suffix[self._dtype]
         engine = klass(self._mv_mul, self._n, self._find_maximum, self._num_eigs)
 
-        eigenvalues, eigenvectors, iteration_counts = engine.run()
+        if initial_vector is not None:
+            eigenvalues, eigenvectors, iteration_counts = engine.run(initial_vector)
+        else:
+            eigenvalues, eigenvectors, iteration_counts = engine.run()
 
         self._iteration_counts = iteration_counts
 

test/test_pylanczos.py

@@ -128,6 +128,28 @@ def mv_mul(v_in, v_out):
 
         np.testing.assert_allclose(eigvec, correct_eigvec)
 
+    def test_sparse_matrix_with_initial_vector(self):
+        n = 10
+
+        matrix = np.zeros((n, n), dtype='float64')
+        for i in range(n-1):
+            matrix[i, i+1] = -1
+            matrix[i+1, i] = -1
+
+        correct_eigvec = np.sin((1+np.array(range(n)))*np.pi/(n+1))
+        correct_eigvec /= np.linalg.norm(correct_eigvec)
+
+        engine = PyLanczos(matrix, False, 1)
+        eigenvalues, eigenvectors = engine.run(correct_eigvec)
+        eigval = eigenvalues[0]
+        eigvec = eigenvectors[:, 0]
+
+        np.testing.assert_almost_equal(eigval, -2.0*np.cos(np.pi/(n+1)))
+        sign = np.sign(eigvec[0])
+        correct_eigvec *= sign
+
+        np.testing.assert_allclose(eigvec, correct_eigvec)
+        np.testing.assert_equal(engine.iteration_counts[0], 1)
 
     def test_complex_matrix(self):
         matrix = np.array([[  0, 1j,  1],