add contour plot, add support for keywords in scatter

Cryoris · Cryoris · commit b568e872c89a · 2019-11-06T23:05:43.000+01:00
diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt
@@ -16,6 +16,7 @@ endfunction(add_example)
 add_example(axlines)
 add_example(bar)
 add_example(basic)
+add_example(contour)
 add_example(eigen)
 add_example(errorbar)
 add_example(fill_inbetween)
@@ -26,6 +27,7 @@ add_example(minimal)
 add_example(modern)
 add_example(nonblock)
 add_example(quiver)
+add_example(scatter)
 add_example(semilogx)
 add_example(semilogy)
 add_example(small)
diff --git a/examples/contour.cpp b/examples/contour.cpp
@@ -0,0 +1,35 @@
+#include <Eigen/Dense>
+#include "../matplotlibcpp.h"
+namespace plt = matplotlibcpp;
+
+template <typename Vector>
+std::pair<Eigen::MatrixXd, Eigen::MatrixXd> meshgrid(const Vector& x, const Vector& y) {
+  const unsigned m = x.size(), n = y.size();
+  // todo fix type
+  Eigen::MatrixXd X(m, n), Y(m, n);
+  for (unsigned i = 0; i < m; ++i) {
+    for (unsigned j = 0; j < n; ++ j) {
+      X(i, j) = *(x.data() + i);
+      Y(i, j) = *(y.data() + j);
+    }
+  }
+  std::pair<Eigen::MatrixXd, Eigen::MatrixXd> res = std::make_pair(X, Y);
+  return res;
+}
+
+void plot_contour() {
+  const unsigned n = 100;
+  Eigen::VectorXd x = Eigen::VectorXd::LinSpaced(n, 0, 1),
+                  y = Eigen::VectorXd::LinSpaced(n, 0, 1);
+  std::pair<Eigen::MatrixXd, Eigen::MatrixXd> XY = meshgrid(x, y);
+  Eigen::MatrixXd Z = XY.first + 2 * XY.second;
+
+  plt::contour(XY.first, XY.second, Z);
+  plt::show();
+
+}
+
+int main() {
+  plot_contour();
+  return 0;
+}
diff --git a/examples/scatter.cpp b/examples/scatter.cpp
@@ -0,0 +1,25 @@
+#define _USE_MATH_DEFINES
+#include <cmath>
+#include <vector>
+#include "../matplotlibcpp.h"
+namespace plt = matplotlibcpp;
+
+void plot() {
+  const unsigned n = 100;
+  std::vector<double> x(n), y(n);
+
+  #include <iostream>
+  for (unsigned i = 0; i < n; ++i) {
+    x[i] = sin(2 * M_PI * i / n);
+    y[i] = cos(2 * M_PI * i / n);
+  }
+
+  plt::scatter(x, y, {{"color", "red"}, {"label", "a circle!"}});
+  plt::legend();
+  plt::show();
+}
+
+int main() {
+  plot();
+  return 0;
+}
diff --git a/matplotlibcpp.h b/matplotlibcpp.h
@@ -53,6 +53,8 @@ struct _interpreter {
   PyObject *s_python_function_fignum_exists;
   PyObject *s_python_function_plot;
   PyObject *s_python_function_quiver;
+  PyObject *s_python_function_contour;
+  PyObject *s_python_function_colormap;
   PyObject *s_python_function_axhline;
   PyObject *s_python_function_axvline;
   PyObject *s_python_function_semilogx;
@@ -184,6 +186,7 @@ struct _interpreter {
         PyObject_GetAttrString(pymod, "fignum_exists");
     s_python_function_plot = PyObject_GetAttrString(pymod, "plot");
     s_python_function_quiver = PyObject_GetAttrString(pymod, "quiver");
+    s_python_function_contour = PyObject_GetAttrString(pymod, "contour");
     s_python_function_axhline = PyObject_GetAttrString(pymod, "axhline");
     s_python_function_axvline = PyObject_GetAttrString(pymod, "axvline");
     s_python_function_semilogx = PyObject_GetAttrString(pymod, "semilogx");
@@ -228,6 +231,7 @@ struct _interpreter {
         !s_python_function_draw || !s_python_function_pause ||
         !s_python_function_figure || !s_python_function_fignum_exists ||
         !s_python_function_plot || !s_python_function_quiver ||
+        !s_python_function_contour ||
         !s_python_function_semilogx || !s_python_function_semilogy ||
         !s_python_function_loglog || !s_python_function_fill ||
         !s_python_function_fill_between || !s_python_function_subplot ||
@@ -256,6 +260,7 @@ struct _interpreter {
         !PyFunction_Check(s_python_function_fignum_exists) ||
         !PyFunction_Check(s_python_function_plot) ||
         !PyFunction_Check(s_python_function_quiver) ||
+        !PyFunction_Check(s_python_function_contour) ||
         !PyFunction_Check(s_python_function_semilogx) ||
         !PyFunction_Check(s_python_function_semilogy) ||
         !PyFunction_Check(s_python_function_loglog) ||
@@ -760,6 +765,53 @@ void plot_surface(const Matrix &x, const Matrix& y, const Matrix& z,
     Py_DECREF(res);
 }
 
+
+// @brief plot_surface for datapoints (x_ij, y_ij, z_ij) with i,j = 0..n
+// @param x The x values of the datapoints in a matrix
+// @param y The y values of the datapoints in a matrix
+// @param z The function value of the datapoints in a matrix
+// @param keywords Additional keywords
+template <typename Matrix>
+void contour(const Matrix &x, const Matrix& y, const Matrix& z,
+             const std::map<std::string, std::string> &keywords = {}) {
+  detail::_interpreter::get();
+
+  // using numpy arrays
+  PyObject *xarray = get_2darray(x);
+  PyObject *yarray = get_2darray(y);
+  PyObject *zarray = get_2darray(z);
+
+  // construct positional args
+  PyObject *args = PyTuple_New(3);
+  PyTuple_SetItem(args, 0, xarray);
+  PyTuple_SetItem(args, 1, yarray);
+  PyTuple_SetItem(args, 2, zarray);
+
+  // Build up the kw args.
+  PyObject *kwargs = PyDict_New();
+
+  PyObject *python_colormap_coolwarm = PyObject_GetAttrString(
+      detail::_interpreter::get().s_python_colormap, "coolwarm");
+
+  PyDict_SetItemString(kwargs, "cmap", python_colormap_coolwarm);
+
+  for (std::map<std::string, std::string>::const_iterator it = keywords.begin();
+       it != keywords.end(); ++it) {
+    PyDict_SetItemString(kwargs, it->first.c_str(),
+                         PyString_FromString(it->second.c_str()));
+  }
+
+  PyObject *res = PyObject_Call(detail::_interpreter::get().s_python_function_contour,
+  args, kwargs);
+  if (!res)
+    throw std::runtime_error("failed surface");
+
+  Py_DECREF(args);
+  Py_DECREF(kwargs);
+  if (res)
+    Py_DECREF(res);
+}
+
 template <typename Numeric>
 bool stem(const std::vector<Numeric> &x, const std::vector<Numeric> &y,
           const std::map<std::string, std::string> &keywords) {
@@ -893,16 +945,23 @@ bool hist(const VectorY &y, long bins = 10, std::string color = "b",
 // @brief Scatter plot
 // @param x x-coordinates of the 2d points
 // @param y y-coordinates of the 2d points
-// @param s the marker size  in points**2
+// @param s the marker size in points**2
+// @param keywords Additional keywords
 template <typename VectorX, typename VectorY>
-bool scatter(const VectorX &x, const VectorY &y, const double s = 1.0) {
+bool scatter(const VectorX &x, const VectorY &y, const double s = 1.0,
+             const std::map<std::string, std::string> keywords = {}) {
   assert(x.size() == y.size());
 
   PyObject *xarray = get_array(x);
   PyObject *yarray = get_array(y);
 
   PyObject *kwargs = PyDict_New();
-  PyDict_SetItemString(kwargs, "s", PyLong_FromLong(s));
+  PyDict_SetItemString(kwargs, "s", PyFloat_FromDouble(s));
+  for (std::map<std::string, std::string>::const_iterator it = keywords.begin();
+       it != keywords.end(); ++it) {
+    PyDict_SetItemString(kwargs, it->first.c_str(),
+                         PyUnicode_FromString(it->second.c_str()));
+  }
 
   PyObject *plot_args = PyTuple_New(2);
   PyTuple_SetItem(plot_args, 0, xarray);
@@ -919,6 +978,13 @@ bool scatter(const VectorX &x, const VectorY &y, const double s = 1.0) {
   return res;
 }
 
+template <typename VectorX, typename VectorY>
+bool scatter(const VectorX &x, const VectorY &y,
+             const std::map<std::string, std::string>& keywords) {
+  return scatter(x, y, 1.0, keywords);
+}
+
+
 // @brief Spy plot
 // @param A the matrix
 // @param precision Plot all elements above `|precision|`
