Merge pull request #15 from soblin/feature/colorbar

Feature/colorbar

Author: soblin

gallery/subplots_axes_and_figures/CMakeLists.txt

@@ -1,13 +1,15 @@
 add_demo(align_labels_demo align_labels_demo.cpp)
 add_demo(gridspec_multicolumn gridspec_multicolumn.cpp)
 add_demo(multiple_figs_demo multiple_figs_demo.cpp)
+add_demo(colorbar_placement colorbar_placement.cpp)
 add_demo(subplots subplots.cpp)
 
 add_custom_target(subplots_axes_and_figures
-  DEPENDS align_labels_demo gridspec_multicolumn multiple_figs_demo
+  DEPENDS align_labels_demo gridspec_multicolumn multiple_figs_demo colorbar_placement
   COMMAND align_labels_demo
   COMMAND gridspec_multicolumn
   COMMAND multiple_figs_demo
+  COMMAND colorbar_placement
   COMMENT "subplots_axes_and_figures"
   WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
 )

gallery/subplots_axes_and_figures/colorbar_placement.cpp

@@ -0,0 +1,48 @@
+// example from
+// https://matplotlib.org/stable/gallery/subplots_axes_and_figures/colorbar_placement.html
+
+#include <pybind11/embed.h>
+#include <pybind11/stl.h>
+
+#include <matplotlibcpp17/pyplot.h>
+
+#include <xtensor/xrandom.hpp>
+
+#include <vector>
+
+namespace py = pybind11;
+using namespace py::literals;
+using namespace std;
+using namespace matplotlibcpp17;
+
+int main1() {
+  auto plt = matplotlibcpp17::pyplot::import();
+  auto [fig, axs] = plt.subplots(2, 2);
+  const vector<string> cmaps = {"RdBu_r", "viridis"};
+  for (auto col : {0, 1}) {
+    for (auto row : {0, 1}) {
+      auto x_ = xt::random::randn<double>({20, 20}) * (col + 1.0);
+      vector<vector<double>> x(20);
+      for (int i = 0; i < 20; ++i) {
+        x[i].resize(20);
+        for (int j = 0; j < 20; ++j)
+          x[i][j] = x_(i, j);
+      }
+      auto &ax = axs[col + row * 2];
+      auto pcm = ax.pcolormesh(Args(x), Kwargs("cmap"_a = cmaps[col]));
+      fig.colorbar(Args(pcm.unwrap()),
+                   Kwargs("ax"_a = ax.unwrap(), "shrink"_a = 0.6));
+    }
+  }
+#if USE_GUI
+  plt.show();
+#else
+  plt.savefig(Args("colorbar_placement1.png"));
+#endif
+  return 0;
+}
+
+int main() {
+  py::scoped_interpreter guard{};
+  main1();
+}

include/matplotlibcpp17/axes.h

@@ -123,6 +123,11 @@ struct DECL_STRUCT_ATTR Axes : public BaseWrapper {
   legend::Legend legend(const pybind11::tuple &args = pybind11::tuple(),
                         const pybind11::dict &kwargs = pybind11::dict());
 
+  // pcolormesh
+  collections::QuadMesh
+  pcolormesh(const pybind11::tuple &args = pybind11::tuple(),
+             const pybind11::dict &kwargs = pybind11::dict());
+
   // plot
   pybind11::object plot(const pybind11::tuple &args = pybind11::tuple(),
                         const pybind11::dict &kwargs = pybind11::dict());
@@ -230,6 +235,7 @@ struct DECL_STRUCT_ATTR Axes : public BaseWrapper {
     LOAD_FUNC_ATTR(hist2d, self);
     LOAD_FUNC_ATTR(invert_yaxis, self);
     LOAD_FUNC_ATTR(legend, self);
+    LOAD_FUNC_ATTR(pcolormesh, self);
     LOAD_FUNC_ATTR(plot, self);
     // NOTE: only when called with projection='3d', `plot_surface`, `plot_wireframe`, `set_zlabel` prop exists.
     try {
@@ -278,6 +284,7 @@ struct DECL_STRUCT_ATTR Axes : public BaseWrapper {
   pybind11::object hist2d_attr;
   pybind11::object invert_yaxis_attr;
   pybind11::object legend_attr;
+  pybind11::object pcolormesh_attr;
   pybind11::object plot_attr;
   pybind11::object plot_surface_attr;
   pybind11::object plot_wireframe_attr;
@@ -478,6 +485,13 @@ legend::Legend Axes::legend(const pybind11::tuple &args,
   return legend::Legend(obj);
 }
 
+// pcolormesh
+collections::QuadMesh Axes::pcolormesh(const pybind11::tuple &args,
+                                       const pybind11::dict &kwargs) {
+  pybind11::object ret = pcolormesh_attr(*args, **kwargs);
+  return collections::QuadMesh(ret);
+}
+
 // plot
 pybind11::object Axes::plot(const pybind11::tuple &args,
                             const pybind11::dict &kwargs) {

include/matplotlibcpp17/collections.h

@@ -73,6 +73,15 @@ pybind11::object PatchCollection::set_array(const pybind11::tuple &args,
   pybind11::object ret = set_array_attr(*args, **kwargs);
   return ret;
 }
+
+/**
+ * @brief A wrapper class for matplotlib.collections.QuadMesh
+ **/
+struct DECL_STRUCT_ATTR QuadMesh : public BaseWrapper {
+public:
+  QuadMesh(pybind11::object quadmesh) { self = quadmesh; }
+};
+
 } // namespace matplotlibcpp17::collections
 
 #endif /* MATPLOTLIBCPP17_COLLECTIONS_H */

include/matplotlibcpp17/pyplot.h

@@ -302,6 +302,7 @@ std::tuple<figure::Figure, std::vector<axes::Axes>>
 PyPlot::subplots(int r, int c, const pybind11::dict &kwargs) {
   // subplots() returns [][] (if r > 1 && c > 1) else []
   // return []axes in row-major
+  // NOTE: equal to Axes.flat
   pybind11::tuple args = pybind11::make_tuple(r, c);
   pybind11::list ret = subplots_attr(*args, **kwargs);
   std::vector<axes::Axes> axes;