Enable matplotlib contourf function

Cbrower · Cbrower · commit 18eab4c78a1e · 2022-04-22T19:53:55.000-07:00
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -96,6 +96,9 @@ if(Python3_NumPy_FOUND)
   add_executable(contour examples/contour.cpp)
   target_link_libraries(contour PRIVATE matplotlib_cpp)
   set_target_properties(contour PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/bin")
+  add_executable(contourf examples/contourf.cpp)
+  target_link_libraries(contourf PRIVATE matplotlib_cpp)
+  set_target_properties(contourf PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/bin")
 
   add_executable(spy examples/spy.cpp)
   target_link_libraries(spy PRIVATE matplotlib_cpp)
diff --git a/examples/contourf.cpp b/examples/contourf.cpp
@@ -0,0 +1,24 @@
+#include "../matplotlibcpp.h"
+
+#include <cmath>
+
+namespace plt = matplotlibcpp;
+
+int main()
+{
+    std::vector<std::vector<double>> x, y, z;
+    for (double i = -5; i <= 5;  i += 0.25) {
+        std::vector<double> x_row, y_row, z_row;
+        for (double j = -5; j <= 5; j += 0.25) {
+            x_row.push_back(i);
+            y_row.push_back(j);
+            z_row.push_back(::std::sin(::std::hypot(i, j)));
+        }
+        x.push_back(x_row);
+        y.push_back(y_row);
+        z.push_back(z_row);
+    }
+
+    plt::contourf(x, y, z);
+    plt::show();
+}
diff --git a/matplotlibcpp.h b/matplotlibcpp.h
@@ -57,6 +57,7 @@ struct _interpreter {
     PyObject *s_python_function_plot;
     PyObject *s_python_function_quiver;
     PyObject* s_python_function_contour;
+    PyObject* s_python_function_contourf;
     PyObject *s_python_function_semilogx;
     PyObject *s_python_function_semilogy;
     PyObject *s_python_function_loglog;
@@ -234,6 +235,7 @@ struct _interpreter {
         s_python_function_plot = safe_import(pymod, "plot");
         s_python_function_quiver = safe_import(pymod, "quiver");
         s_python_function_contour = safe_import(pymod, "contour");
+        s_python_function_contourf = safe_import(pymod, "contourf");
         s_python_function_semilogx = safe_import(pymod, "semilogx");
         s_python_function_semilogy = safe_import(pymod, "semilogy");
         s_python_function_loglog = safe_import(pymod, "loglog");
@@ -625,6 +627,48 @@ void contour(const std::vector<::std::vector<Numeric>> &x,
   if (res) Py_DECREF(res);
 }
 
+template <typename Numeric>
+void contourf(const std::vector<::std::vector<Numeric>> &x,
+             const std::vector<::std::vector<Numeric>> &y,
+             const std::vector<::std::vector<Numeric>> &z,
+             const std::map<std::string, std::string> &keywords = {})
+{
+  detail::_interpreter::get();
+
+  // using numpy arrays
+  PyObject *xarray = detail::get_2darray(x);
+  PyObject *yarray = detail::get_2darray(y);
+  PyObject *zarray = detail::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_contourf, args, kwargs);
+  if (!res)
+    throw std::runtime_error("failed contourf");
+
+  Py_DECREF(args);
+  Py_DECREF(kwargs);
+  if (res) Py_DECREF(res);
+}
+
 template <typename Numeric>
 void spy(const std::vector<::std::vector<Numeric>> &x,
          const double markersize = -1,  // -1 for default matplotlib size
