added (optional) numpy support; first use via 'Entity.positions' property

CMakeLists.txt

@@ -22,6 +22,8 @@ option(ENABLE_GFX "whether graphics support should be enabled"
       ON)
 option(ENABLE_IMG "whether the image processing module should be compiled"
        ON)
+option(USE_NUMPY "whether numpy support is added"
+      OFF)
 option(USE_DOUBLE_PRECISION "whether to compile in double precision" 
        OFF)
 option(ENABLE_SPNAV "whether 3DConnexion devices should be supported"
@@ -87,6 +89,11 @@ if (USE_SHADER)
 else()
   set(_SHADER OFF)
 endif()
+if (USE_NUMPY)
+  set(_NUMPY ON)
+else()
+  set(_NUMPY OFF)
+endif()
 if (COMPILE_TMTOOLS)
   set(_TM_TOOLS ON)
 else()
@@ -219,6 +226,10 @@ find_package(PNG REQUIRED)
 find_package(Eigen 2.0.0 REQUIRED)
 find_package(Python 2.4 REQUIRED)
 
+if(USE_NUMPY)
+  find_package(Numpy REQUIRED)
+endif()
+
 if (ENABLE_IMG)
   find_package(FFTW REQUIRED)
   find_package(TIFF REQUIRED)
@@ -252,6 +263,10 @@ include_directories(${Boost_INCLUDE_DIRS}
                     ${TIFF_INCLUDE_DIR}
                     ${SPNAV_INCLUDE_DIR}
                     )
+if(USE_NUMPY)
+include_directories(${PYTHON_NUMPY_INCLUDE_DIR})
+endif()
+
 if (UNIX)
   SET(CMAKE_SKIP_BUILD_RPATH  FALSE)
   SET(CMAKE_BUILD_WITH_INSTALL_RPATH FALSE)
@@ -284,6 +299,7 @@ message(STATUS
         "   OpenGL support             (-DENABLE_GFX) : ${_OPENGL}\n"
         "   Image Processing support   (-DENABLE_IMG) : ${_IMG}\n"
         "   Shader support             (-DUSE_SHADER) : ${_SHADER}\n"
+        "   Numpy support              (-DUSE_NUMPY)  : ${_NUMPY}\n"
         "   Optimize                     (-DOPTIMIZE) : ${_OPT}\n"
         "   Profiling support             (-DPROFILE) : ${_PROFILE}\n"
         "   Double Precision (-DUSE_DOUBLE_PRECISION) : ${_DOUBLE_PREC}\n"

cmake_support/FindNumpy.cmake

+if (PYTHON_NUMPY_INCLUDE_DIR)
+  # in cache already
+  set (PYTHON_NUMPY_FIND_QUIETLY TRUE)
+endif (PYTHON_NUMPY_INCLUDE_DIR)
+
+#INCLUDE(FindPython)
+
+IF(PYTHON_BINARY)
+    EXEC_PROGRAM ("${PYTHON_BINARY}"
+      ARGS "-c 'import numpy; print numpy.get_include()'"
+      OUTPUT_VARIABLE PYTHON_NUMPY_INCLUDE_DIR
+      RETURN_VALUE PYTHON_NUMPY_NOT_FOUND)
+
+    if (PYTHON_NUMPY_INCLUDE_DIR)
+      set (PYTHON_NUMPY_FOUND TRUE)
+      set (PYTHON_NUMPY_INCLUDE_DIR ${PYTHON_NUMPY_INCLUDE_DIR} CACHE STRING "Numpy include path")
+    else (PYTHON_NUMPY_INCLUDE_DIR)
+      set(PYTHON_NUMPY_FOUND FALSE)
+    endif (PYTHON_NUMPY_INCLUDE_DIR)
+ENDIF(PYTHON_BINARY)
+
+if (PYTHON_NUMPY_FOUND)
+  if (NOT PYTHON_NUMPY_FIND_QUIETLY)
+    message (STATUS "Numpy headers found")
+  endif (NOT PYTHON_NUMPY_FIND_QUIETLY)
+else (PYTHON_NUMPY_FOUND)
+  if (PYTHON_NUMPY_FIND_REQUIRED)
+    message (FATAL_ERROR "Numpy headers missing")
+  endif (PYTHON_NUMPY_FIND_REQUIRED)
+endif (PYTHON_NUMPY_FOUND)
+
+MARK_AS_ADVANCED (PYTHON_NUMPY_INCLUDE_DIR)
\ No newline at end of file

modules/config/CMakeLists.txt

@@ -5,14 +5,20 @@ dllexport.hh
 version.hh
 )
 
-
 set(animations_enabled 0)
+
 if (USE_SHADER)
   set(shader_support 1)
 else()
   set(shader_support 0)
 endif()
 
+if (USE_NUMPY)
+  set(numpy_support 1)
+else()
+  set(numpy_support 0)
+endif()
+
 if (PROFILE)
   set(profiling_enabled 1)
 else()

modules/config/config.hh.in

@@ -31,4 +31,5 @@
 #define OST_STATIC_PROPERTY_WORKAROUND @static_props@
 #define OST_SPNAV_ENABLED @spnav_enabled@
 #define OST_FFT_USE_THREADS @fftw_use_threads@
+#define OST_NUMPY_SUPPORT_ENABLED @numpy_support@
 #endif

modules/mol/base/pymod/export_entity.cc

@@ -30,6 +30,10 @@ using namespace ost::mol;
 
 #include <ost/export_helper/generic_property_def.hh>
 
+#if OST_NUMPY_SUPPORT_ENABLED
+#include <numpy/arrayobject.h>
+#endif
+
 namespace {
 EntityHandle create1() {  return CreateEntity(); }
 
@@ -65,10 +69,96 @@ ICSEditor depr_request_ics_editor(EntityHandle e, EditMode m)
 }
 
 
+PyObject* get_pos(EntityHandle& entity)
+{
+#if OST_NUMPY_SUPPORT_ENABLED
+  npy_intp dims[]={entity.GetAtomCount(),3};
+  PyObject* na = PyArray_SimpleNew(2,dims,NPY_DOUBLE);
+  npy_double* nad = reinterpret_cast<npy_double*>(PyArray_DATA(na));
+  for(AtomHandleIter it=entity.AtomsBegin();it!=entity.AtomsEnd();++it,nad+=3) {
+    geom::Vec3 pos=(*it).GetPos();
+    nad[0]=static_cast<npy_double>(pos[0]);
+    nad[1]=static_cast<npy_double>(pos[1]);
+    nad[2]=static_cast<npy_double>(pos[2]);
+  }
+  return na;
+#else
+  throw std::runtime_error("GetPositions disabled, since numpy support is not compiled in");
+  return 0;
+#endif
+}
+
+
+#if OST_NUMPY_SUPPORT_ENABLED
+template <typename T>
+void set_pos_t(PyArrayObject* na, EntityHandle& entity)
+{
+  XCSEditor ed=entity.EditXCS(BUFFERED_EDIT);
+
+  if(PyArray_ISCONTIGUOUS(na)) {
+    T* data = reinterpret_cast<T*>(PyArray_DATA(na));
+    size_t count=0;
+    for(AtomHandleIter it=entity.AtomsBegin();it!=entity.AtomsEnd();++it,++count) {
+      ed.SetAtomPos(*it,geom::Vec3(static_cast<Real>(data[count*3+0]),
+                                   static_cast<Real>(data[count*3+1]),
+                                   static_cast<Real>(data[count*3+2])));
+    }
+  } else {
+    size_t count=0;
+    for(AtomHandleIter it=entity.AtomsBegin();it!=entity.AtomsEnd();++it,++count) {
+      ed.SetAtomPos(*it,geom::Vec3(static_cast<Real>(*reinterpret_cast<T*>(PyArray_GETPTR2(na,count,0))),
+                                   static_cast<Real>(*reinterpret_cast<T*>(PyArray_GETPTR2(na,count,1))),
+                                   static_cast<Real>(*reinterpret_cast<T*>(PyArray_GETPTR2(na,count,2)))));
+    }
+  }
+}
+#endif
+
+void set_pos(EntityHandle& entity, object& pyobj)
+{
+#if OST_NUMPY_SUPPORT_ENABLED
+  size_t acount = entity.GetAtomCount();
+  
+  if(!PyArray_Check(pyobj.ptr())) {
+    throw std::runtime_error(std::string("expected a numpy array"));
+    return;
+  }
+  PyArrayObject* na=reinterpret_cast<PyArrayObject*>(pyobj.ptr());
+
+  if(PyArray_NDIM(na)!=2 || PyArray_DIM(na,0)!=acount || PyArray_DIM(na,1)!=3) {
+    throw std::runtime_error("excpected a numpy array of shape (NAtoms, 3)");
+    return;
+  }
+
+  switch(PyArray_TYPE(na)) {
+  case NPY_FLOAT:
+    set_pos_t<float>(na,entity); break;
+  case NPY_DOUBLE:
+    set_pos_t<double>(na,entity); break;
+  case NPY_INT:
+    set_pos_t<int>(na,entity); break;
+  case NPY_LONG:
+    set_pos_t<long>(na,entity); break;
+  case NPY_SHORT:
+    set_pos_t<short>(na,entity); break;
+  default:
+    throw std::runtime_error("excpected a numpy array of type float, double, int, long or short");
+    return;
+  };
+
+#else
+  throw std::runtime_error("SetPositions disabled, since numpy support is not compiled in");
+#endif
 }
 
+} // ns
+
 void export_Entity()
 {
+#if OST_NUMPY_SUPPORT_ENABLED
+  import_array();
+#endif
+
   class_<EntityBase> ent_base("EntityBase", no_init);
   ent_base
     .def(self_ns::str(self))
@@ -139,6 +229,11 @@ void export_Entity()
     .def("IsTransformationIdentity",&EntityHandle::IsTransformationIdentity)
     .def(self==self)
     .def(self!=self)
+#if OST_NUMPY_SUPPORT_ENABLED
+    .def("SetPositions",set_pos)
+    .def("GetPositions",get_pos)
+    .add_property("positions",get_pos,set_pos)
+#endif
   ;
 
   def("CreateEntity",create1);

modules/mol/base/tests/test_numpy.py

+import sys,unittest
+sys.path.append("../../../../stage/lib64/openstructure")
+from ost import *
+import numpy
+
+class TestNumpy(unittest.TestCase):
+  def setUp(self):
+    pass
+
+  def test_(self):
+    entity=mol.CreateEntity()
+    ed=entity.EditXCS()
+    ch=ed.InsertChain("X")
+    re=ed.AppendResidue(ch,"ALA")
+    a0=ed.InsertAtom(re,"A",geom.Vec3(0,0,0))
+    a1=ed.InsertAtom(re,"B",geom.Vec3(1,0,0))
+    a2=ed.InsertAtom(re,"C",geom.Vec3(2,0,0))
+    a3=ed.InsertAtom(re,"D",geom.Vec3(3,0,0))
+
+    self.assertTrue(geom.Distance(a0.pos,geom.Vec3(0,0,0))<1e-10)
+    self.assertTrue(geom.Distance(a1.pos,geom.Vec3(1,0,0))<1e-10)
+    self.assertTrue(geom.Distance(a2.pos,geom.Vec3(2,0,0))<1e-10)
+    self.assertTrue(geom.Distance(a3.pos,geom.Vec3(3,0,0))<1e-10)
+
+    entity.positions=numpy.array([[0,1,0],[0,2,0],[0,3,0],[0,4,0]])
+
+    self.assertTrue(geom.Distance(a0.pos,geom.Vec3(0,1,0))<1e-10)
+    self.assertTrue(geom.Distance(a1.pos,geom.Vec3(0,2,0))<1e-10)
+    self.assertTrue(geom.Distance(a2.pos,geom.Vec3(0,3,0))<1e-10)
+    self.assertTrue(geom.Distance(a3.pos,geom.Vec3(0,4,0))<1e-10)
+
+    na=entity.positions
+
+    self.assertTrue(geom.Distance(geom.Vec3(na[0][0],na[0][1],na[0][2]),geom.Vec3(0,1,0))<1e-10)
+    self.assertTrue(geom.Distance(geom.Vec3(na[1][0],na[1][1],na[1][2]),geom.Vec3(0,2,0))<1e-10)
+    self.assertTrue(geom.Distance(geom.Vec3(na[2][0],na[2][1],na[2][2]),geom.Vec3(0,3,0))<1e-10)
+    self.assertTrue(geom.Distance(geom.Vec3(na[3][0],na[3][1],na[3][2]),geom.Vec3(0,4,0))<1e-10)
+
+if __name__== '__main__':
+    unittest.main()
+