Ver Fonte

Add 3D quiver support

Joshua Brinsfield há 4 anos atrás
pai
commit
80bc9cd84d
1 ficheiros alterados com 86 adições e 0 exclusões
  1. 86 0
      matplotlibcpp.h

+ 86 - 0
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 = "")
 {