add get_2darray for templated matrices and spy plot

Cryoris · Cryoris · commit ab37736f7d99 · 2019-09-20T16:19:17.000+02:00
diff --git a/Makefile b/Makefile
@@ -20,7 +20,7 @@ eigen_include = -I /usr/local/include/eigen3
 example_execs = minimal modern basic animation nonblock xkcd quiver bar surface subplot fill_inbetween fill update
 
 # Executable names for examples using Eigen
-eigen_execs = eigen loglog semilogx semilogy small
+eigen_execs = eigen loglog semilogx semilogy small spy
 
 # Example targets (default if just 'make' is called)
 examples: $(example_execs)
diff --git a/examples/spy.cpp b/examples/spy.cpp
@@ -0,0 +1,41 @@
+#include <vector>
+#include <Eigen/Dense>
+#include "../matplotlibcpp.h"
+namespace plt = matplotlibcpp;
+
+int main() {
+
+  const unsigned n = 100;
+  Eigen::MatrixXd A(n / 2, n);
+  std::vector<std::vector<double>> B;
+
+  for (unsigned i = 0; i < n / 2; ++i) {
+    A(i, i) = 1;
+    std::vector<double> row(n);
+    row[i] = 1;
+
+    if (i < n / 2) {
+      A(i, i + n / 2) = 1;
+      row[i + n / 2] = 1;
+    }
+    B.push_back(row);
+  }
+
+  for (unsigned i = 0; i < n / 2; ++i) {
+    for (unsigned j = 0; j < n; ++j) {
+      if (A(i, j) != B[i][j]) {
+        std::cout << i << "," << j << " differ!\n";
+      }
+    }
+  }
+
+  plt::figure();
+  plt::title("Eigen");
+  plt::spy(A);
+
+  plt::figure();
+  plt::title("vector");
+  plt::spy(B);
+  plt::show();
+  return 0;
+}
diff --git a/matplotlibcpp.h b/matplotlibcpp.h
@@ -59,6 +59,7 @@ struct _interpreter {
   PyObject *s_python_function_fill_between;
   PyObject *s_python_function_hist;
   PyObject *s_python_function_scatter;
+  PyObject *s_python_function_spy;
   PyObject *s_python_function_subplot;
   PyObject *s_python_function_legend;
   PyObject *s_python_function_xlim;
@@ -188,6 +189,7 @@ struct _interpreter {
         PyObject_GetAttrString(pymod, "fill_between");
     s_python_function_hist = PyObject_GetAttrString(pymod, "hist");
     s_python_function_scatter = PyObject_GetAttrString(pymod, "scatter");
+    s_python_function_spy = PyObject_GetAttrString(pymod, "spy");
     s_python_function_subplot = PyObject_GetAttrString(pymod, "subplot");
     s_python_function_legend = PyObject_GetAttrString(pymod, "legend");
     s_python_function_ylim = PyObject_GetAttrString(pymod, "ylim");
@@ -236,7 +238,8 @@ struct _interpreter {
         !s_python_function_errorbar || !s_python_function_tight_layout ||
         !s_python_function_stem || !s_python_function_xkcd ||
         !s_python_function_text || !s_python_function_suptitle ||
-        !s_python_function_bar || !s_python_function_subplots_adjust) {
+        !s_python_function_bar || !s_python_function_subplots_adjust ||
+        !s_python_function_spy) {
       throw std::runtime_error("Couldn't find required function!");
     }
 
@@ -253,6 +256,7 @@ struct _interpreter {
         !PyFunction_Check(s_python_function_loglog) ||
         !PyFunction_Check(s_python_function_fill) ||
         !PyFunction_Check(s_python_function_fill_between) ||
+        !PyFunction_Check(s_python_function_spy) ||
         !PyFunction_Check(s_python_function_subplot) ||
         !PyFunction_Check(s_python_function_legend) ||
         !PyFunction_Check(s_python_function_annotate) ||
@@ -396,7 +400,7 @@ PyObject *get_2darray(const std::vector<::std::vector<Numeric>> &v) {
   detail::_interpreter::get(); // interpreter needs to be initialized for the
                                // numpy commands to work
   if (v.size() < 1)
-    throw std::runtime_error("get_2d_array v too small");
+    throw std::runtime_error("get_2darray v too small");
 
   npy_intp vsize[2] = {static_cast<npy_intp>(v.size()),
                        static_cast<npy_intp>(v[0].size())};
@@ -408,14 +412,39 @@ PyObject *get_2darray(const std::vector<::std::vector<Numeric>> &v) {
 
   for (const ::std::vector<Numeric> &v_row : v) {
     if (v_row.size() != static_cast<size_t>(vsize[1]))
-      throw std::runtime_error("Missmatched array size");
+      throw std::runtime_error("mismatched array size");
     std::copy(v_row.begin(), v_row.end(), vd_begin);
     vd_begin += vsize[1];
   }
 
   return reinterpret_cast<PyObject *>(varray);
 }
 
+// suitable for Eigen matrices
+template <typename Matrix>
+PyObject *get_2darray(const Matrix &A) {
+  detail::_interpreter::get(); // interpreter needs to be initialized for the
+                               // numpy commands to work
+  if (A.size() < 1)
+    throw std::runtime_error("get_2darray A too small");
+
+  npy_intp vsize[2] = {static_cast<npy_intp>(A.rows()),
+                       static_cast<npy_intp>(A.cols())};
+
+  PyArrayObject *varray =
+      (PyArrayObject *)PyArray_SimpleNew(2, vsize, NPY_DOUBLE);
+
+  double *vd_begin = static_cast<double *>(PyArray_DATA(varray));
+
+  for (std::size_t i = 0; i < A.rows(); ++i) {
+    for (std::size_t j = 0; j < A.cols(); ++j) {
+      *(vd_begin + i * A.cols() + j) = A(i, j);
+    }
+  }
+
+  return reinterpret_cast<PyObject *>(varray);
+}
+
 #else // fallback if we don't have numpy: copy every element of the given vector
 
 template <typename Vector> PyObject *get_array(const Vector &v) {
@@ -869,6 +898,29 @@ bool scatter(const VectorX &x, const VectorY &y, const double s = 1.0) {
   return res;
 }
 
+// @brief Spy plot
+// @param A the matrix
+template <typename Matrix>
+bool spy(const Matrix &A, double precision=0) {
+  PyObject *Aarray = get_2darray(A);
+
+  PyObject *kwargs = PyDict_New();
+  PyDict_SetItemString(kwargs, "precision", PyFloat_FromDouble(precision));
+
+  PyObject *plot_args = PyTuple_New(1);
+  PyTuple_SetItem(plot_args, 0, Aarray);
+
+  PyObject *res = PyObject_Call(
+      detail::_interpreter::get().s_python_function_spy, plot_args, kwargs);
+
+  Py_DECREF(plot_args);
+  Py_DECREF(kwargs);
+  if (res)
+    Py_DECREF(res);
+
+  return res;
+}
+
 template <typename Numeric>
 bool bar(const std::vector<Numeric> &y, std::string ec = "black",
          std::string ls = "-", double lw = 1.0,
