Changeset 8c9cce for src/Python/export_numpy.cpp

Timestamp:

May 19, 2021, 7:06:29 PM (4 years ago)

Author:

Frederik Heber <frederik.heber@…>

Branches:

Candidate_v1.7.0, stable

Children:

ecc050

Parents:

7ca7d5d

git-author:

Frederik Heber <frederik.heber@…> (04/23/19 01:11:50)

git-committer:

Frederik Heber <frederik.heber@…> (05/19/21 19:06:29)

Message:

Numpy exporters are fully operational.

• TESTS: added test on working getter and setter.

File:

• src/Python/export_numpy.cpp (modified) (3 diffs)

src/Python/export_numpy.cpp

@@ -35 +35 @@
 //#include "CodePatterns/MemDebug.hpp"
 
+#include<boost/bind.hpp>
+#include<boost/function.hpp>
 #include<boost/python/def.hpp>
-
 #include<boost/python/numpy.hpp>
 
 #include "CodePatterns/Assert.hpp"
+
+#include "World.hpp"
 
 namespace p = boost::python;
@@ -46 +49 @@
 unsigned int get_num_atoms()
 {
-  return 0;
+  return World::getInstance().countSelectedAtoms();
+}
+
+np::ndarray allocate_ndarray(const unsigned int num_atoms)
+{
+  p::tuple shape = p::make_tuple(num_atoms, 3);
+  np::dtype dtype = np::dtype::get_builtin<float>();
+  np::ndarray array = np::zeros(shape, dtype);
+  return array;
+}
+
+np::ndarray get_ndarray(boost::function<const Vector &(const atom &)> &_get_function)
+{
+  unsigned int num_atoms = get_num_atoms();
+  std::cout << num_atoms << std::endl;
+  np::ndarray positions = allocate_ndarray(num_atoms);
+
+  unsigned int ia=0;
+  for (World::AtomSelectionConstIterator iter = World::getInstance().beginAtomSelection();
+      iter != World::getInstance().endAtomSelection();
+      ++iter) {
+    const atom & current = *iter->second;
+    for (unsigned int i=0;i<NDIM;++i)
+      positions[ia][i] = _get_function(current)[i];
+    ++ia;
+    ASSERT(ia <= num_atoms, "get_positions() - more atoms selected than expected.");
+  }
+
+  return positions;
 }
 
 np::ndarray get_positions()
 {
-  unsigned int num_atoms = get_num_atoms();
-  p::tuple shape = p::make_tuple(num_atoms, 3);
-  np::dtype dtype = np::dtype::get_builtin<float>();
-  np::ndarray positions = np::zeros(shape, dtype);
-  return positions;
+  static boost::function< const Vector&(const atom&) > get_vector =
+      boost::bind(&AtomInfo::getPosition, _1);
+  return get_ndarray(get_vector);
 }
 
-void set_positions(const np::ndarray &new_positions)
+np::ndarray get_velocities()
+{
+  static boost::function< const Vector&(const atom&) > get_vector =
+      boost::bind(&AtomInfo::getAtomicVelocity, _1);
+  return get_ndarray(get_vector);
+}
+
+np::ndarray get_forces()
+{
+  static boost::function< const Vector&(const atom&) > get_vector =
+      boost::bind(&AtomInfo::getAtomicForce, _1);
+  return get_ndarray(get_vector);
+}
+
+void set_ndarray(
+    const np::ndarray &new_positions,
+    boost::function<void (atom &, const Vector &)> &_set_function)
 {
   unsigned int num_atoms = get_num_atoms();
@@ -65 +110 @@
   ASSERT( new_positions.shape()[0] == num_atoms,
     "pyMoleCuilder::set_positions() - numpy array has unexpected size.");
+
+  unsigned int ia=0;
+  Vector temp;
+  for (World::AtomSelectionIterator iter = World::getInstance().beginAtomSelection();
+      iter != World::getInstance().endAtomSelection();
+      ++iter) {
+    atom &current = *iter->second;
+    for (unsigned int i=0;i<NDIM;++i)
+      temp[i] = p::extract<double>(new_positions[ia][i]);
+    _set_function(current, temp);
+    ++ia;
+    ASSERT(ia <= num_atoms, "set_positions() - more atoms selected than expected.");
+  }
+}
+
+void set_positions(const np::ndarray &new_positions)
+{
+  static boost::function< void (atom&, const Vector&) > set_vector =
+      boost::bind(&AtomInfo::setPosition, _1, _2);
+  set_ndarray(new_positions, set_vector);
+}
+
+void set_velocities(const np::ndarray &new_positions)
+{
+  static boost::function< void (atom&, const Vector&) > set_vector =
+      boost::bind(&AtomInfo::setAtomicVelocity, _1, _2);
+  set_ndarray(new_positions, set_vector);
+}
+
+void set_forces(const np::ndarray &new_positions)
+{
+  static boost::function< void (atom&, const Vector&) > set_vector =
+      boost::bind(&AtomInfo::setAtomicForce, _1, _2);
+  set_ndarray(new_positions, set_vector);
 }
 
 void export_numpy()
 {
-  p::def("get_position", get_positions);
-  p::def("set_position", set_positions, p::args("position"));
+  p::def("get_positions", get_positions);
+  p::def("get_velocities", get_velocities);
+  p::def("get_forces", get_forces);
+  p::def("set_positions", set_positions, p::args("position"));
+  p::def("set_velocities", set_velocities, p::args("velocity"));
+  p::def("set_forces", set_forces, p::args("force"));
 }