Commit a7bc22fb authored by Peter Eastman's avatar Peter Eastman
Browse files

Optimized setPositions() from Python

parent 9b523076
......@@ -194,13 +194,59 @@ int Py_SequenceToVecDouble(PyObject* obj, std::vector<double>& out) {
}
%fragment("Py_SequenceToVecVec3", "header", fragment="Py_SequenceToVec3") {
int isNumpyAvailable() {
static bool initialized = false;
static bool available = false;
if (!initialized) {
initialized = true;
available = (_import_array() >= 0);
}
return available;
}
int Py_SequenceToVecVec3(PyObject* obj, std::vector<Vec3>& out) {
PyObject* stripped = Py_StripOpenMMUnits(obj); // new reference
if (isNumpyAvailable()) {
if (PyArray_Check(stripped) && PyArray_ISCARRAY_RO(stripped) && PyArray_NDIM(stripped) == 2 && PyArray_DIM(stripped, 1) == 3) {
int type = PyArray_TYPE(stripped);
int length = PyArray_DIM(stripped, 0);
void* data = PyArray_DATA((PyArrayObject*) stripped);
if (type == NPY_DOUBLE) {
out.resize(length);
memcpy(&out[0][0], data, 3*sizeof(double)*length);
Py_DECREF(stripped);
return SWIG_OK;
}
if (type == NPY_FLOAT) {
out.resize(length);
float* floatData = (float*) data;
for (int i = 0; i < length; i++)
out[i] = Vec3(floatData[3*i], floatData[3*i+1], floatData[3*i+2]);
Py_DECREF(stripped);
return SWIG_OK;
}
if (type == NPY_INT32) {
out.resize(length);
int* intData = (int*) data;
for (int i = 0; i < length; i++)
out[i] = Vec3(intData[3*i], intData[3*i+1], intData[3*i+2]);
Py_DECREF(stripped);
return SWIG_OK;
}
if (type == NPY_INT64) {
out.resize(length);
long long* longData = (long long*) data;
for (int i = 0; i < length; i++)
out[i] = Vec3(longData[3*i], longData[3*i+1], longData[3*i+2]);
Py_DECREF(stripped);
return SWIG_OK;
}
}
}
int ret = 0;
PyObject* stripped = NULL;
PyObject* item = NULL;
PyObject* item1 = NULL;
PyObject* iterator = NULL;
stripped = Py_StripOpenMMUnits(obj); // new reference
iterator = PyObject_GetIter(stripped); // new reference
if (iterator == NULL) {
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment