Added arrow, cla, margins, contour

JBPennington · lava · commit 374dcd032ca8 · 2020-06-17T18:51:50.000+02:00
diff --git a/matplotlibcpp.h b/matplotlibcpp.h
@@ -45,6 +45,7 @@ namespace detail {
 static std::string s_backend;
 
 struct _interpreter {
+    PyObject* s_python_function_arrow;
     PyObject *s_python_function_show;
     PyObject *s_python_function_close;
     PyObject *s_python_function_draw;
@@ -54,6 +55,7 @@ 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_semilogx;
     PyObject *s_python_function_semilogy;
     PyObject *s_python_function_loglog;
@@ -79,8 +81,10 @@ struct _interpreter {
     PyObject *s_python_function_gca;
     PyObject *s_python_function_xticks;
     PyObject *s_python_function_yticks;
+    PyObject* s_python_function_margins;
     PyObject *s_python_function_tick_params;
     PyObject *s_python_function_grid;
+    PyObject* s_python_function_cla;
     PyObject *s_python_function_clf;
     PyObject *s_python_function_errorbar;
     PyObject *s_python_function_annotate;
@@ -186,6 +190,7 @@ struct _interpreter {
         Py_DECREF(pylabname);
         if (!pylabmod) { throw std::runtime_error("Error loading module pylab!"); }
 
+        s_python_function_arrow = safe_import(pymod, "arrow");
         s_python_function_show = safe_import(pymod, "show");
         s_python_function_close = safe_import(pymod, "close");
         s_python_function_draw = safe_import(pymod, "draw");
@@ -194,6 +199,7 @@ struct _interpreter {
         s_python_function_fignum_exists = safe_import(pymod, "fignum_exists");
         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_semilogx = safe_import(pymod, "semilogx");
         s_python_function_semilogy = safe_import(pymod, "semilogy");
         s_python_function_loglog = safe_import(pymod, "loglog");
@@ -215,13 +221,15 @@ struct _interpreter {
         s_python_function_gca = safe_import(pymod, "gca");
         s_python_function_xticks = safe_import(pymod, "xticks");
         s_python_function_yticks = safe_import(pymod, "yticks");
+        s_python_function_margins = safe_import(pymod, "margins");
         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");
         s_python_function_ginput = safe_import(pymod, "ginput");
         s_python_function_save = safe_import(pylabmod, "savefig");
         s_python_function_annotate = safe_import(pymod,"annotate");
+        s_python_function_cla = safe_import(pymod, "cla");
         s_python_function_clf = safe_import(pymod, "clf");
         s_python_function_errorbar = safe_import(pymod, "errorbar");
         s_python_function_tight_layout = safe_import(pymod, "tight_layout");
@@ -709,6 +717,37 @@ bool fill_between(const std::vector<Numeric>& x, const std::vector<Numeric>& y1,
     return res;
 }
 
+template <typename Numeric>
+bool arrow(Numeric x, Numeric y, Numeric end_x, Numeric end_y, const std::string& fc = "r",
+           const std::string ec = "k", Numeric head_length = 0.25, Numeric head_width = 0.1625) {
+    PyObject* obj_x = PyFloat_FromDouble(x);
+    PyObject* obj_y = PyFloat_FromDouble(y);
+    PyObject* obj_end_x = PyFloat_FromDouble(end_x);
+    PyObject* obj_end_y = PyFloat_FromDouble(end_y);
+
+    PyObject* kwargs = PyDict_New();
+    PyDict_SetItemString(kwargs, "fc", PyString_FromString(fc.c_str()));
+    PyDict_SetItemString(kwargs, "ec", PyString_FromString(ec.c_str()));
+    PyDict_SetItemString(kwargs, "head_width", PyFloat_FromDouble(head_width));
+    PyDict_SetItemString(kwargs, "head_length", PyFloat_FromDouble(head_length));
+
+    PyObject* plot_args = PyTuple_New(4);
+    PyTuple_SetItem(plot_args, 0, obj_x);
+    PyTuple_SetItem(plot_args, 1, obj_y);
+    PyTuple_SetItem(plot_args, 2, obj_end_x);
+    PyTuple_SetItem(plot_args, 3, obj_end_y);
+
+    PyObject* res =
+            PyObject_Call(detail::_interpreter::get().s_python_function_arrow, plot_args, kwargs);
+
+    Py_DECREF(plot_args);
+    Py_DECREF(kwargs);
+    if (res)
+        Py_DECREF(res);
+
+    return res;
+}
+
 template< typename Numeric>
 bool hist(const std::vector<Numeric>& y, long bins=10,std::string color="b",
           double alpha=1.0, bool cumulative=false)
@@ -1040,6 +1079,39 @@ bool plot(const std::vector<NumericX>& x, const std::vector<NumericY>& y, const
     return res;
 }
 
+template <typename NumericX, typename NumericY, typename NumericZ>
+bool contour(const std::vector<NumericX>& x, const std::vector<NumericY>& y,
+             const std::vector<NumericZ>& z,
+             const std::map<std::string, std::string>& keywords = {}) {
+    assert(x.size() == y.size() && x.size() == z.size());
+
+    PyObject* xarray = get_array(x);
+    PyObject* yarray = get_array(y);
+    PyObject* zarray = get_array(z);
+
+    PyObject* plot_args = PyTuple_New(3);
+    PyTuple_SetItem(plot_args, 0, xarray);
+    PyTuple_SetItem(plot_args, 1, yarray);
+    PyTuple_SetItem(plot_args, 2, zarray);
+
+    // construct keyword args
+    PyObject* kwargs = PyDict_New();
+    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* res =
+            PyObject_Call(detail::_interpreter::get().s_python_function_contour, plot_args, kwargs);
+
+    Py_DECREF(kwargs);
+    Py_DECREF(plot_args);
+    if (res)
+        Py_DECREF(res);
+
+    return res;
+}
+
 template<typename NumericX, typename NumericY, typename NumericU, typename NumericW>
 bool quiver(const std::vector<NumericX>& x, const std::vector<NumericY>& y, const std::vector<NumericU>& u, const std::vector<NumericW>& w, const std::map<std::string, std::string>& keywords = {})
 {
@@ -1669,6 +1741,38 @@ inline void yticks(const std::vector<Numeric> &ticks, const std::map<std::string
     yticks(ticks, {}, keywords);
 }
 
+template <typename Numeric> inline void margins(Numeric margin)
+{
+    // construct positional args
+    PyObject* args = PyTuple_New(1);
+    PyTuple_SetItem(args, 0, PyFloat_FromDouble(margin));
+
+    PyObject* res =
+            PyObject_CallObject(detail::_interpreter::get().s_python_function_margins, args);
+    if (!res)
+        throw std::runtime_error("Call to margins() failed.");
+
+    Py_DECREF(args);
+    Py_DECREF(res);
+}
+
+template <typename Numeric> inline void margins(Numeric margin_x, Numeric margin_y)
+{
+    // construct positional args
+    PyObject* args = PyTuple_New(2);
+    PyTuple_SetItem(args, 0, PyFloat_FromDouble(margin_x));
+    PyTuple_SetItem(args, 1, PyFloat_FromDouble(margin_y));
+
+    PyObject* res =
+            PyObject_CallObject(detail::_interpreter::get().s_python_function_margins, args);
+    if (!res)
+        throw std::runtime_error("Call to margins() failed.");
+
+    Py_DECREF(args);
+    Py_DECREF(res);
+}
+
+
 inline void tick_params(const std::map<std::string, std::string>& keywords, const std::string axis = "both")
 {
   detail::_interpreter::get();
@@ -2069,6 +2173,18 @@ inline void clf() {
     Py_DECREF(res);
 }
 
+inline void cla() {
+    detail::_interpreter::get();
+
+    PyObject* res = PyObject_CallObject(detail::_interpreter::get().s_python_function_cla,
+                                        detail::_interpreter::get().s_python_empty_tuple);
+
+    if (!res)
+        throw std::runtime_error("Call to cla() failed.");
+
+    Py_DECREF(res);
+}
+
 inline void ion() {
     detail::_interpreter::get();
 
