From 7c85dcc26d9fd7d808fe0adb1238bc2abb3705fc Mon Sep 17 00:00:00 2001
From: Joshua Brinsfield <sapphous@gmail.com>
Date: Sat, 31 Oct 2020 18:51:05 -0400
Subject: [PATCH] Add 3D quiver support

---
 matplotlibcpp.h | 86 +++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 86 insertions(+)

diff --git a/matplotlibcpp.h b/matplotlibcpp.h
index 93a72be..52b2780 100644
--- a/matplotlibcpp.h
+++ b/matplotlibcpp.h
@@ -1202,6 +1202,92 @@ bool quiver(const std::vector<NumericX>& x, const std::vector<NumericY>& y, cons
     return res;
 }
 
+template<typename NumericX, typename NumericY, typename NumericZ, typename NumericU, typename NumericW, typename NumericV>
+bool quiver(const std::vector<NumericX>& x, const std::vector<NumericY>& y, const std::vector<NumericZ>& z, const std::vector<NumericU>& u, const std::vector<NumericW>& w, const std::vector<NumericV>& v, const std::map<std::string, std::string>& keywords = {})
+{
+  //set up 3d axes stuff
+  static PyObject *mpl_toolkitsmod = nullptr, *axis3dmod = nullptr;
+  if (!mpl_toolkitsmod) {
+    detail::_interpreter::get();
+
+    PyObject* mpl_toolkits = PyString_FromString("mpl_toolkits");
+    PyObject* axis3d = PyString_FromString("mpl_toolkits.mplot3d");
+    if (!mpl_toolkits || !axis3d) { throw std::runtime_error("couldnt create string"); }
+
+    mpl_toolkitsmod = PyImport_Import(mpl_toolkits);
+    Py_DECREF(mpl_toolkits);
+    if (!mpl_toolkitsmod) { throw std::runtime_error("Error loading module mpl_toolkits!"); }
+
+    axis3dmod = PyImport_Import(axis3d);
+    Py_DECREF(axis3d);
+    if (!axis3dmod) { throw std::runtime_error("Error loading module mpl_toolkits.mplot3d!"); }
+  }
+  
+  //assert sizes match up
+  assert(x.size() == y.size() && x.size() == u.size() && u.size() == w.size() && x.size() == z.size() && x.size() == v.size() && u.size() == v.size());
+
+  //set up parameters
+  detail::_interpreter::get();
+
+  PyObject* xarray = detail::get_array(x);
+  PyObject* yarray = detail::get_array(y);
+  PyObject* zarray = detail::get_array(z);
+  PyObject* uarray = detail::get_array(u);
+  PyObject* warray = detail::get_array(w);
+  PyObject* varray = detail::get_array(v);
+
+  PyObject* plot_args = PyTuple_New(6);
+  PyTuple_SetItem(plot_args, 0, xarray);
+  PyTuple_SetItem(plot_args, 1, yarray);
+  PyTuple_SetItem(plot_args, 2, zarray);
+  PyTuple_SetItem(plot_args, 3, uarray);
+  PyTuple_SetItem(plot_args, 4, warray);
+  PyTuple_SetItem(plot_args, 5, varray);
+
+  // 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()));
+  }
+    
+  //get figure gca to enable 3d projection
+  PyObject *fig =
+      PyObject_CallObject(detail::_interpreter::get().s_python_function_figure,
+                          detail::_interpreter::get().s_python_empty_tuple);
+  if (!fig) throw std::runtime_error("Call to figure() failed.");
+
+  PyObject *gca_kwargs = PyDict_New();
+  PyDict_SetItemString(gca_kwargs, "projection", PyString_FromString("3d"));
+
+  PyObject *gca = PyObject_GetAttrString(fig, "gca");
+  if (!gca) throw std::runtime_error("No gca");
+  Py_INCREF(gca);
+  PyObject *axis = PyObject_Call(
+      gca, detail::_interpreter::get().s_python_empty_tuple, gca_kwargs);
+
+  if (!axis) throw std::runtime_error("No axis");
+  Py_INCREF(axis);
+  Py_DECREF(gca);
+  Py_DECREF(gca_kwargs);
+  
+  //plot our boys bravely, plot them strongly, plot them with a wink and clap
+  PyObject *plot3 = PyObject_GetAttrString(axis, "quiver");
+  if (!plot3) throw std::runtime_error("No 3D line plot");
+  Py_INCREF(plot3);
+  PyObject* res = PyObject_Call(
+          plot3, plot_args, kwargs);
+  if (!res) throw std::runtime_error("Failed 3D plot");
+  Py_DECREF(plot3);
+  Py_DECREF(axis);
+  Py_DECREF(kwargs);
+  Py_DECREF(plot_args);
+  if (res)
+      Py_DECREF(res);
+
+  return res;
+}
+
 template<typename NumericX, typename NumericY>
 bool stem(const std::vector<NumericX>& x, const std::vector<NumericY>& y, const std::string& s = "")
 {