implement colorbar; requires returned PyObject of mappable object

tkphd · lava · commit f1af639c4271 · 2020-03-28T00:51:41.000+01:00
catch call without mappable object

stray {

add keywords (shrink is particularly useful)
diff --git a/Makefile b/Makefile
@@ -14,7 +14,8 @@ WITHOUT_NUMPY   := $(findstring $(CXXFLAGS), WITHOUT_NUMPY)
 
 # Examples requiring numpy support to compile
 EXAMPLES_NUMPY  := surface
-EXAMPLES        := minimal basic modern animation nonblock xkcd quiver bar fill_inbetween fill update subplot2grid \
+EXAMPLES        := minimal basic modern animation nonblock xkcd quiver bar \
+	           fill_inbetween fill update subplot2grid colorbar \
                    $(if WITHOUT_NUMPY,,$(EXAMPLES_NUMPY))
 
 # Prefix every example with 'examples/build/'
diff --git a/examples/colorbar.cpp b/examples/colorbar.cpp
@@ -0,0 +1,32 @@
+#define _USE_MATH_DEFINES
+#include <cmath>
+#include <iostream>
+#include "../matplotlibcpp.h"
+
+using namespace std;
+namespace plt = matplotlibcpp;
+
+int main()
+{
+    // Prepare data
+    int ncols = 500, nrows = 300;
+    std::vector<float> z(ncols * nrows);
+    for (int j=0; j<nrows; ++j) {
+        for (int i=0; i<ncols; ++i) {
+            z.at(ncols * j + i) = std::sin(std::hypot(i - ncols/2, j - nrows/2));
+        }
+    }
+
+    const float* zptr = &(z[0]);
+    const int colors = 1;
+
+    plt::title("My matrix");
+    PyObject* mat;
+    plt::imshow(zptr, nrows, ncols, colors, {}, &mat);
+    plt::colorbar(mat);
+
+    // Show plots
+    plt::show();
+    plt::close();
+    Py_DECREF(mat);
+}
diff --git a/matplotlibcpp.h b/matplotlibcpp.h
@@ -87,6 +87,7 @@ struct _interpreter {
     PyObject *s_python_function_text;
     PyObject *s_python_function_suptitle;
     PyObject *s_python_function_bar;
+    PyObject *s_python_function_colorbar;
     PyObject *s_python_function_subplots_adjust;
 
 
@@ -205,7 +206,7 @@ struct _interpreter {
         s_python_function_ylabel = safe_import(pymod, "ylabel");
         s_python_function_xticks = safe_import(pymod, "xticks");
         s_python_function_yticks = safe_import(pymod, "yticks");
-    	s_python_function_tick_params = safe_import(pymod, "tick_params");
+        s_python_function_tick_params = safe_import(pymod, "tick_params");
         s_python_function_grid = safe_import(pymod, "grid");
         s_python_function_xlim = safe_import(pymod, "xlim");
         s_python_function_ion = safe_import(pymod, "ion");
@@ -220,11 +221,11 @@ struct _interpreter {
         s_python_function_text = safe_import(pymod, "text");
         s_python_function_suptitle = safe_import(pymod, "suptitle");
         s_python_function_bar = safe_import(pymod,"bar");
+        s_python_function_colorbar = PyObject_GetAttrString(pymod, "colorbar");
         s_python_function_subplots_adjust = safe_import(pymod,"subplots_adjust");
 #ifndef WITHOUT_NUMPY
         s_python_function_imshow = safe_import(pymod, "imshow");
 #endif
-
         s_python_empty_tuple = PyTuple_New(0);
     }
 
@@ -589,7 +590,7 @@ bool hist(const std::vector<Numeric>& y, long bins=10,std::string color="b",
 
 #ifndef WITHOUT_NUMPY
     namespace internal {
-        inline void imshow(void *ptr, const NPY_TYPES type, const int rows, const int columns, const int colors, const std::map<std::string, std::string> &keywords)
+        inline void imshow(void *ptr, const NPY_TYPES type, const int rows, const int columns, const int colors, const std::map<std::string, std::string> &keywords, PyObject** out)
         {
             assert(type == NPY_UINT8 || type == NPY_FLOAT);
             assert(colors == 1 || colors == 3 || colors == 4);
@@ -613,18 +614,21 @@ bool hist(const std::vector<Numeric>& y, long bins=10,std::string color="b",
             Py_DECREF(kwargs);
             if (!res)
                 throw std::runtime_error("Call to imshow() failed");
-            Py_DECREF(res);
+            if (out)
+                *out = res;
+            else
+                Py_DECREF(res);
         }
     }
 
-    inline void imshow(const unsigned char *ptr, const int rows, const int columns, const int colors, const std::map<std::string, std::string> &keywords = {})
+    inline void imshow(const unsigned char *ptr, const int rows, const int columns, const int colors, const std::map<std::string, std::string> &keywords = {}, PyObject** out = nullptr)
     {
-        internal::imshow((void *) ptr, NPY_UINT8, rows, columns, colors, keywords);
+        internal::imshow((void *) ptr, NPY_UINT8, rows, columns, colors, keywords, out);
     }
 
-    inline void imshow(const float *ptr, const int rows, const int columns, const int colors, const std::map<std::string, std::string> &keywords = {})
+    inline void imshow(const float *ptr, const int rows, const int columns, const int colors, const std::map<std::string, std::string> &keywords = {}, PyObject** out = nullptr)
     {
-        internal::imshow((void *) ptr, NPY_FLOAT, rows, columns, colors, keywords);
+        internal::imshow((void *) ptr, NPY_FLOAT, rows, columns, colors, keywords, out);
     }
 
 #ifdef WITH_OPENCV
@@ -1136,6 +1140,27 @@ void text(Numeric x, Numeric y, const std::string& s = "")
     Py_DECREF(res);
 }
 
+void colorbar(PyObject* mappable = NULL, const std::map<std::string, float>& keywords = {})
+{
+    if (mappable == NULL)
+        throw std::runtime_error("Must call colorbar with PyObject* returned from an image, contour, surface, etc.");
+    PyObject* args = PyTuple_New(1);
+    PyTuple_SetItem(args, 0, mappable);
+
+    PyObject* kwargs = PyDict_New();
+    for(std::map<std::string, float>::const_iterator it = keywords.begin(); it != keywords.end(); ++it)
+    {
+        PyDict_SetItemString(kwargs, it->first.c_str(), PyFloat_FromDouble(it->second));
+    }
+
+    PyObject* res = PyObject_Call(detail::_interpreter::get().s_python_function_colorbar, args, kwargs);
+    if(!res) throw std::runtime_error("Call to colorbar() failed.");
+
+    Py_DECREF(args);
+    Py_DECREF(kwargs);
+    Py_DECREF(res);
+}
+
 
 inline long figure(long number = -1)
 {
