Implement Solvent Excluded Surfaces (#972)

* Implement Solvent Excluded Surfaces

Port big optimization from EDTSurf
Now the cube calculation step is 100 - 200x faster (QElapsedTimer)

Implement Cube::fill() and other suggestions

Signed-off-by: Aritz Erkiaga <aerkiaga3@gmail.com>

avogadro/core/cube.cpp

@@ -263,19 +263,37 @@ double Cube::value(const Vector3& pos) const
          value(hC.x(), hC.y(), hC.z()) * P.x() * P.y() * P.z();
 }
 
-bool Cube::setValue(int i, int j, int k, double value_)
+bool Cube::setValue(unsigned int i, unsigned int j, unsigned int k, double value_)
 {
   unsigned int index = i * m_points.y() * m_points.z() + j * m_points.z() + k;
-  if (index < m_data.size()) {
-    m_data[index] = value_;
-    if (value_ < m_minValue)
-      m_minValue = value_;
-    else if (value_ > m_maxValue)
-      m_maxValue = value_;
-    return true;
-  } else {
+  if (index >= m_data.size())
     return false;
-  }
+  m_data[index] = value_;
+  if (value_ < m_minValue)
+    m_minValue = value_;
+  else if (value_ > m_maxValue)
+    m_maxValue = value_;
+  return true;
+}
+
+void Cube::fill(double value_)
+{
+  std::fill(m_data.begin(), m_data.end(), value_);
+  m_minValue = m_maxValue = value_;
+}
+
+bool Cube::fillStripe(
+  unsigned int i, unsigned int j, unsigned int kfirst, unsigned int klast, double value_
+) {
+  unsigned int stripeStartIndex = i * m_points.y() * m_points.z() + j * m_points.z();
+  unsigned int firstIndex = stripeStartIndex + kfirst;
+  if (firstIndex >= m_data.size())
+    return false;
+  unsigned int lastIndex = stripeStartIndex + klast;
+  if (lastIndex >= m_data.size())
+    return false;
+  std::fill(&m_data[firstIndex], &m_data[lastIndex+1], value_);
+  return true;
 }
 
 } // End Core namespace

avogadro/core/cube.h

@@ -183,14 +183,32 @@ class AVOGADROCORE_EXPORT Cube
    * @param k z component of the position.
    * @param value Value at the specified position.
    */
-  bool setValue(int i, int j, int k, double value);
+  bool setValue(unsigned int i, unsigned int j, unsigned int k, double value);
 
   /**
    * Sets the value at the specified index in the cube.
    * @param i 1-dimensional index of the point to set in the cube.
    */
   bool setValue(unsigned int i, double value);
 
+  /**
+   * Sets all indices in the cube to the specified value.
+   * @param value Value to fill the cube with.
+   */
+  void fill(double value);
+
+  /**
+   * Sets all indices in a Z stripe of the cube to the specified value.
+   * @param i x component of the position.
+   * @param j y component of the position.
+   * @param kfirst first z position to fill.
+   * @param klast last z position to fill.
+   * @param value Value to fill the stripe with.
+   */
+  bool fillStripe(
+    unsigned int i, unsigned int j, unsigned int kfirst, unsigned int klast, double value
+  );
+
   /**
    * @return The minimum  value at any point in the Cube.
    */

avogadro/qtplugins/surfaces/surfacedialog.cpp

@@ -26,8 +26,8 @@ SurfaceDialog::SurfaceDialog(QWidget* parent_, Qt::WindowFlags f)
   m_ui->surfaceCombo->addItem(tr("Van der Waals"), Surfaces::Type::VanDerWaals);
   m_ui->surfaceCombo->addItem(tr("Solvent Accessible"),
                               Surfaces::Type::SolventAccessible);
-  /*m_ui->surfaceCombo->addItem(tr("Solvent Excluded"),
-                              Surfaces::Type::SolventExcluded);*/
+  m_ui->surfaceCombo->addItem(tr("Solvent Excluded"),
+                              Surfaces::Type::SolventExcluded);
 
   connect(m_ui->surfaceCombo, SIGNAL(currentIndexChanged(int)),
           SLOT(surfaceComboChanged(int)));

avogadro/qtplugins/surfaces/surfaces.cpp

@@ -40,6 +40,7 @@ namespace {
 #include <avogadro/quantumio/nwchemlog.h>
 
 #include <QtConcurrent/QtConcurrentMap>
+#include <QtConcurrent/QtConcurrentRun>
 #include <QtCore/QBuffer>
 #include <QtCore/QCoreApplication>
 #include <QtCore/QDebug>
@@ -56,6 +57,7 @@ namespace QtPlugins {
 using Core::Array;
 using Core::Cube;
 using Core::GaussianSet;
+using Core::NeighborPerceiver;
 using QtGui::Molecule;
 
 class Surfaces::PIMPL
@@ -70,6 +72,9 @@ Surfaces::Surfaces(QObject* p) : ExtensionPlugin(p), d(new PIMPL())
   auto action = new QAction(this);
   action->setText(tr("Create Surfaces…"));
   connect(action, SIGNAL(triggered()), SLOT(surfacesActivated()));
+  connect(&m_displayMeshWatcher, SIGNAL(finished()), SLOT(displayMesh()));
+  connect(&m_performEDTStepWatcher, SIGNAL(finished()), SLOT(performEDTStep()));
+
   m_actions.push_back(action);
 
   // Register quantum file formats
@@ -163,8 +168,6 @@ void Surfaces::calculateSurface()
     case SolventAccessible:
     case SolventExcluded:
       calculateEDT();
-      // pass a molecule and return a Cube for m_cube
-      connect(&m_cubesWatcher, SIGNAL(finished()), SLOT(displayMesh()));
       break;
 
     case ElectronDensity:
@@ -181,71 +184,147 @@ void Surfaces::calculateSurface()
   }
 }
 
+float inline square(float x) { return x * x; }
+
 void Surfaces::calculateEDT()
 {
-  double probeRadius = 0.0;
-  double radiusOffset = 0.0;
-  switch (m_dialog->surfaceType()) {
-    case VanDerWaals:
-        break;
-    case SolventAccessible:
-        radiusOffset = 1.4;
-        break;
-    case SolventExcluded:
-        probeRadius = 1.4;
-        break;
-  }
-
-  // cache Van der Waals radii
-  std::vector<double> radii(m_molecule->atomCount());
-  double max_radius = probeRadius;
-  for (size_t i = 0; i < radii.size(); i++) {
-    radii[i] = Core::Elements::radiusVDW(m_molecule->atomicNumber(i));
-    if (radii[i] > max_radius)
-      max_radius = radii[i];
-  }
+  QFuture future = QtConcurrent::run([=]() {
+    double probeRadius = 0.0;
+    switch (m_dialog->surfaceType()) {
+      case VanDerWaals:
+          break;
+      case SolventAccessible:
+      case SolventExcluded:
+          probeRadius = 1.4;
+          break;
+    }
 
-  double padding = max_radius + radiusOffset;
-  m_cube->setLimits(*m_molecule, m_dialog->resolution(), padding);
-
-  auto neighborPerceiver = Core::NeighborPerceiver(
-    m_molecule->atomPositions3d(), 2.0 * max_radius
-  );
-  Vector3i cubeSize = m_cube->dimensions();
-
-  std::vector<Vector3i> *indices = new std::vector<Vector3i>;
-  indices->reserve(cubeSize[0] * cubeSize[1] * cubeSize[2]);
-  for (size_t zi = 0; zi < cubeSize[2]; zi++)
-    for (size_t yi = 0; yi < cubeSize[1]; yi++)
-      for (size_t xi = 0; xi < cubeSize[0]; xi++)
-        indices->emplace_back(xi, yi, zi);
-
-  const float res = m_dialog->resolution();
-  const Vector3 min = m_cube->min();
-  const float mdist = probeRadius;
-
-  thread_local Array<Index> *neighbors = nullptr;
-  QFuture future = QtConcurrent::map(*indices, [=](Vector3i &in) {
-    Vector3 pos(in(0), in(1), in(2));
-    pos += Vector3(0.5, 0.5, 0.5);
-    pos *= res;
-    pos += min;
-
-    double minDistance = mdist;
-    if (neighbors == nullptr)
-      neighbors = new Array<Index>;
-    neighborPerceiver.getNeighborsInclusiveInPlace(*neighbors, pos);
-    for (Index neighbor: *neighbors) {
-      Vector3 neighborPos = m_molecule->atomPosition3d(neighbor);
-      double distance = (neighborPos - pos).norm();
-      distance -= radii[neighbor] + radiusOffset;
-      minDistance = std::min(minDistance, distance);
+    // first, make a list of all atom positions and radii
+    Array<Vector3> atomPositions = m_molecule->atomPositions3d();
+    std::vector<std::pair<Vector3, double>> *atoms =
+      new std::vector<std::pair<Vector3, double>>(m_molecule->atomCount());
+    double max_radius = probeRadius;
+    for (size_t i = 0; i < atoms->size(); i++) {
+      (*atoms)[i].first = atomPositions[i];
+      (*atoms)[i].second = Core::Elements::radiusVDW(m_molecule->atomicNumber(i)) + probeRadius;
+      if ((*atoms)[i].second > max_radius)
+        max_radius = (*atoms)[i].second;
     }
 
-    m_cube->setValue(in(0), in(1), in(2), minDistance);
+    double padding = max_radius + probeRadius;
+    m_cube->setLimits(*m_molecule, m_dialog->resolution(), padding);
+    m_cube->fill(-1.0);
+
+    const float res = m_dialog->resolution();
+    const Vector3 min = m_cube->min();
+    const float mdist = probeRadius;
+
+    // then, for each atom, set cubes around it up to a certain radius
+    QFuture innerFuture = QtConcurrent::map(*atoms, [=](std::pair<Vector3, double> &in) {
+      double startPosX = in.first(0) - in.second;
+      double endPosX = in.first(0) + in.second;
+      int startIndexX = (startPosX - min(0)) / res;
+      int endIndexX = (endPosX - min(0)) / res + 1;
+      for (int indexX = startIndexX; indexX < endIndexX; indexX++) {
+        double posX = indexX * res + min(0);
+        double radiusXsq = square(in.second) - square(posX - in.first(0));
+        if (radiusXsq < 0.0)
+          continue;
+        double radiusX = sqrt(radiusXsq);
+        double startPosY = in.first(1) - radiusX;
+        double endPosY = in.first(1) + radiusX;
+        int startIndexY = (startPosY - min(1)) / res;
+        int endIndexY = (endPosY - min(1)) / res + 1;
+        for (int indexY = startIndexY; indexY < endIndexY; indexY++) {
+          double posY = indexY * res + min(1);
+          double lengthXYsq = square(radiusX) - square(posY - in.first(1));
+          if (lengthXYsq < 0.0)
+            continue;
+          double lengthXY = sqrt(lengthXYsq);
+          double startPosZ = in.first(2) - lengthXY;
+          double endPosZ = in.first(2) + lengthXY;
+          int startIndexZ = (startPosZ - min(2)) / res;
+          int endIndexZ = (endPosZ - min(2)) / res + 1;
+          m_cube->fillStripe(indexX, indexY, startIndexZ, endIndexZ - 1, 1.0f);
+        }
+      }
+    });
+
+    innerFuture.waitForFinished();
   });
+
+  // SolventExcluded requires an extra pass
+  if (m_dialog->surfaceType() == SolventExcluded) {
+    m_performEDTStepWatcher.setFuture(future);
+  } else {
+    m_displayMeshWatcher.setFuture(future);
+  }
+}
 
-  m_cubesWatcher.setFuture(future);
+void Surfaces::performEDTStep()
+{
+  QFuture future = QtConcurrent::run([=]() {
+    const double probeRadius = 1.4;
+    const double scaledProbeRadius = probeRadius / m_dialog->resolution();
+
+    // make a list of all "outside" cubes in contact with an "inside" cube
+    // these are the only ones that can be "nearest" to an "inside" cube
+    Array<Vector3> relativePositions;
+    // also make a list of all "inside" cubes
+    std::vector<Vector3i> *insideIndices = new std::vector<Vector3i>;
+    Vector3i size = m_cube->dimensions();
+    relativePositions.reserve(size(0) * size(1) * 4); // O(n^2)
+    insideIndices->reserve(size(0) * size(1) * size(2)); // O(n^3)
+    for (int z = 0; z < size(2); z++) {
+      int zp = std::max(z - 1, 0);
+      int zn = std::min(z + 1, size(2) - 1);
+      for (int y = 0; y < size(1); y++) {
+        int yp = std::max(y - 1, 0);
+        int yn = std::min(y + 1, size(1) - 1);
+        for (int x = 0; x < size(0); x++) {
+          if (m_cube->value(x, y, z) > 0.0) {
+            insideIndices->emplace_back(x, y, z);
+            continue;
+          }
+          int xp = std::max(x - 1, 0);
+          int xn = std::min(x + 1, size(0) - 1);
+          if (m_cube->value(xp, y, z) > 0.0
+            || m_cube->value(xn, y, z) > 0.0
+            || m_cube->value(x, yp, z) > 0.0
+            || m_cube->value(x, yn, z) > 0.0
+            || m_cube->value(x, y, zp) > 0.0
+            || m_cube->value(x, y, zn) > 0.0
+          ) {
+            relativePositions.push_back(Vector3(x, y, z));
+          }
+        }
+      }
+    }
+
+    // pass the list to a NeighborPerceiver so it's faster to look up
+    NeighborPerceiver perceiver(relativePositions, scaledProbeRadius);
+
+    // now, exclude all "inside" cubes too close to any "outside" cube
+    thread_local Array<Index> *neighbors = nullptr;
+    QFuture innerFuture = QtConcurrent::map(*insideIndices, [=](Vector3i &in) {
+      Vector3 pos = in.cast<double>();
+      if (neighbors == nullptr)
+        neighbors = new Array<Index>;
+      perceiver.getNeighborsInclusiveInPlace(*neighbors, pos);
+      for (Index neighbor: *neighbors) {
+        const Vector3 &npos = relativePositions[neighbor];
+        float distance = (npos - pos).norm();
+        if (distance <= scaledProbeRadius) {
+          m_cube->setValue(in(0), in(1), in(2), -1.0f);
+          break;
+        }
+      }
+    });
+
+    innerFuture.waitForFinished();
+  });
+
+  m_displayMeshWatcher.setFuture(future);
 }
 
 void Surfaces::calculateQM()

avogadro/qtplugins/surfaces/surfaces.h

@@ -71,6 +71,7 @@ private slots:
   void surfacesActivated();
   void calculateSurface();
   void calculateEDT();
+  void performEDTStep(); // EDT step for SolventExcluded
   void calculateQM();
   void calculateCube();
 
@@ -94,9 +95,13 @@ private slots:
 
   Core::Cube* m_cube = nullptr;
   std::vector<Core::Cube*> m_cubes;
-  QFutureWatcher<void> m_cubesWatcher;
+  /* One QFutureWatcher per asynchronous slot function, e.g.:*/
+  /* calculateEDT() -> [performEDTStep()] -> displayMesh() */
+  QFutureWatcher<void> m_performEDTStepWatcher;
+  QFutureWatcher<void> m_displayMeshWatcher;
   Core::Mesh* m_mesh1 = nullptr;
   Core::Mesh* m_mesh2 = nullptr;
+  /* displayMesh() -> meshFinished() */
   QtGui::MeshGenerator* m_meshGenerator1 = nullptr;
   QtGui::MeshGenerator* m_meshGenerator2 = nullptr;