Handles energies through scripts - still need gradients

Signed-off-by: Geoff Hutchison <[email protected]>

avogadro/calc/chargemanager.h

@@ -80,16 +80,6 @@ class AVOGADROCALC_EXPORT ChargeManager
    */
   bool removeModel(const std::string& identifier);
 
-  /**
-   * New instance of the model for the specified @p identifier. Ownership
-   * is passed to the caller.
-   * @param identifier The unique identifier of the format.
-   * @return Instance of the format, nullptr if not found. Ownership passes to
-   * the
-   * caller.
-   */
-  ChargeModel* newModelFromIdentifier(const std::string& identifier) const;
-
   /**
    * Get a list of all loaded identifiers
    */

avogadro/calc/energycalculator.cpp

@@ -5,6 +5,8 @@
 
 #include "energycalculator.h"
 
+#include <iostream>
+
 namespace Avogadro::Calc {
 
 void EnergyCalculator::gradient(const TVector& x, TVector& grad)

avogadro/calc/energymanager.cpp

@@ -54,6 +54,15 @@ bool EnergyManager::addModel(EnergyCalculator* model)
   return true;
 }
 
+EnergyCalculator* EnergyManager::model(const std::string& identifier) const
+{
+  auto it = m_identifiers.find(identifier);
+  if (it == m_identifiers.end()) {
+    return nullptr;
+  }
+  return m_models[it->second]->newInstance();
+}
+
 bool EnergyManager::removeModel(const std::string& identifier)
 {
   auto ids = m_identifiers[identifier];

avogadro/calc/energymanager.h

@@ -86,7 +86,7 @@ class AVOGADROCALC_EXPORT EnergyManager
    * @return Instance of the model, nullptr if not found. Ownership passes to
    * the caller.
    */
-  EnergyCalculator* newModelFromIdentifier(const std::string& identifier) const;
+  EnergyCalculator* model(const std::string& identifier) const;
 
   /**
    * Get a list of all loaded identifiers

avogadro/calc/lennardjones.cpp

@@ -30,6 +30,8 @@ void LennardJones::setMolecule(Core::Molecule* mol)
     return; // nothing to do
   }
 
+  m_mask = mol->frozenAtomMask();
+
   m_cell = mol->unitCell(); // could be nullptr
   int numAtoms = mol->atomCount();
 

avogadro/qtgui/pythonscript.cpp

@@ -227,6 +227,8 @@ QByteArray PythonScript::asyncWriteAndResponse(QByteArray input, const int expec
     return QByteArray(); // wait
   }
 
+  qDebug() << " writing to process: " << input << " " << expectedLines;
+
   m_process->write(input);
   QByteArray buffer;
 
@@ -243,9 +245,13 @@ QByteArray PythonScript::asyncWriteAndResponse(QByteArray input, const int expec
         buffer += m_process->readLine();
         remainingLines--;
       }
+      // clear the rest of the buffer
+      qDebug() << " remaining " << m_process->readAll();
     }
   }
 
+  qDebug() << " read from process: " << buffer;
+
   return buffer;
 }
 

avogadro/qtplugins/forcefield/forcefield.cpp

@@ -45,6 +45,8 @@ const int energyAction = 0;
 const int optimizeAction = 1;
 const int configureAction = 2;
 const int freezeAction = 3;
+const int unfreezeAction = 4;
+const int constraintAction = 5;
 
 Forcefield::Forcefield(QObject* parent_) : ExtensionPlugin(parent_)
 {
@@ -68,10 +70,17 @@ Forcefield::Forcefield(QObject* parent_) : ExtensionPlugin(parent_)
 
   action = new QAction(this);
   action->setEnabled(true);
-  action->setText(tr("Freeze Atoms"));
+  action->setText(tr("Freeze Selected Atoms"));
   action->setData(freezeAction);
   connect(action, SIGNAL(triggered()), SLOT(freezeSelected()));
   m_actions.push_back(action);
+
+  action = new QAction(this);
+  action->setEnabled(true);
+  action->setText(tr("Unfreeze Selected Atoms"));
+  action->setData(unfreezeAction);
+  connect(action, SIGNAL(triggered()), SLOT(unfreezeSelected()));
+  m_actions.push_back(action);
 }
 
 Forcefield::~Forcefield() {}
@@ -99,8 +108,6 @@ void Forcefield::setMolecule(QtGui::Molecule* mol)
     return;
 
   m_molecule = mol;
-
-  // @todo set the available method list
 }
 
 void Forcefield::optimize()
@@ -115,15 +122,18 @@ void Forcefield::optimize()
   cppoptlib::LbfgsSolver<EnergyCalculator> solver;
 
   int n = m_molecule->atomCount();
+  // we have to cast the current 3d positions into a VectorXd
   Core::Array<Vector3> pos = m_molecule->atomPositions3d();
-  // just to get the right size / shape
-  Core::Array<Vector3> forces = m_molecule->atomPositions3d();
   double* p = pos[0].data();
   Eigen::Map<Eigen::VectorXd> map(p, 3 * n);
   Eigen::VectorXd positions = map;
+
   Eigen::VectorXd gradient = Eigen::VectorXd::Zero(3 * n);
+  // just to get the right size / shape
+  // we'll use this to draw the force arrows later
+  Core::Array<Vector3> forces = m_molecule->atomPositions3d();
 
-  // Create a Criteria class to adjust stopping criteria
+  // Create a Criteria class so we can update coords every N steps
   cppoptlib::Criteria<Real> crit = cppoptlib::Criteria<Real>::defaults();
 
   // e.g., every 5 steps, update coordinates
@@ -134,20 +144,27 @@ void Forcefield::optimize()
   solver.setStopCriteria(crit);
 
   // set the method
-  //@todo check m_method for a particular calculator
-  Calc::LennardJones lj;
-  lj.setMolecule(m_molecule);
+  std::string recommended = recommendedForceField();
+  qDebug() << "Energy method: " << recommended.c_str();
+
+  if (m_method == nullptr) {
+    // we have to create the calculator
+    m_method = Calc::EnergyManager::instance().model(recommended);
+  }
+  m_method->setMolecule(m_molecule);
+  m_method->setMask(m_molecule->frozenAtomMask());
 
-  double energy = lj.value(positions);
+  Real energy = m_method->value(positions);
   for (unsigned int i = 0; i < m_maxSteps / crit.iterations; ++i) {
-    solver.minimize(lj, positions);
+    solver.minimize(*m_method, positions);
 
-    double currentEnergy = lj.value(positions);
-    lj.gradient(positions, gradient);
+    Real currentEnergy = m_method->value(positions);
+    // get the current gradient for force visualization
+    m_method->gradient(positions, gradient);
 
     // update coordinates
     const double* d = positions.data();
-    // casting would be lovely...
+    // casting back would be lovely...
     for (size_t i = 0; i < n; ++i) {
       pos[i] = Vector3(*(d), *(d + 1), *(d + 2));
       d += 3;
@@ -179,19 +196,61 @@ void Forcefield::energy()
     return;
 
   //@todo check m_method for a particular calculator
-  Calc::LennardJones lj;
-  lj.setMolecule(m_molecule);
+  std::string recommended = recommendedForceField();
+  qDebug() << "Energy method: " << recommended.c_str();
+
+  if (m_method == nullptr) {
+    // we have to create the calculator
+    m_method = Calc::EnergyManager::instance().model(recommended);
+  }
+  m_method->setMolecule(m_molecule);
 
   int n = m_molecule->atomCount();
+  // we have to cast the current 3d positions into a VectorXd
   Core::Array<Vector3> pos = m_molecule->atomPositions3d();
   double* p = pos[0].data();
   Eigen::Map<Eigen::VectorXd> map(p, 3 * n);
   Eigen::VectorXd positions = map;
 
-  QString msg(tr("Energy = %L1 kJ/mol").arg(lj.value(positions)));
+  // now get the energy
+  Real energy = m_method->value(positions);
+
+  QString msg(tr("Energy = %L1").arg(energy));
   QMessageBox::information(nullptr, tr("Avogadro"), msg);
 }
 
+std::string Forcefield::recommendedForceField() const
+{
+  // if we have a unit cell, we need to use the LJ calculator
+  // (implementing something better would be nice)
+  if (m_molecule == nullptr || m_molecule->unitCell() != nullptr)
+    return "LJ";
+
+  // otherwise, let's see what identifers are returned
+  auto list =
+    Calc::EnergyManager::instance().identifiersForMolecule(*m_molecule);
+  if (list.empty())
+    return "LJ"; // this will always work
+
+  // iterate to see what we have
+  std::string bestOption;
+  for (auto options : list) {
+    // ideally, we'd use GFN-FF but it needs tweaking
+    // everything else is a ranking
+    // GAFF is better than MMFF94 which is better than UFF
+    if (options == "UFF" && bestOption != "GAFF" || bestOption != "MMFF94")
+      bestOption = options;
+    if (options == "MMFF94" && bestOption != "GAFF")
+      bestOption = options;
+    if (options == "GAFF")
+      bestOption = options;
+  }
+  if (!bestOption.empty())
+    return bestOption;
+  else
+    return "LJ"; // this will always work
+}
+
 void Forcefield::freezeSelected()
 {
   if (!m_molecule)
@@ -201,7 +260,21 @@ void Forcefield::freezeSelected()
   // now freeze the specified atoms
   for (Index i = 0; i < numAtoms; ++i) {
     if (m_molecule->atomSelected(i)) {
-      // m_molecule->setAtomFrozen(i, true);
+      m_molecule->setFrozenAtom(i, true);
+    }
+  }
+}
+
+void Forcefield::unfreezeSelected()
+{
+  if (!m_molecule)
+    return;
+
+  int numAtoms = m_molecule->atomCount();
+  // now freeze the specified atoms
+  for (Index i = 0; i < numAtoms; ++i) {
+    if (m_molecule->atomSelected(i)) {
+      m_molecule->setFrozenAtom(i, false);
     }
   }
 }
@@ -241,6 +314,9 @@ void Forcefield::registerScripts()
          it = m_scripts.constBegin(),
          itEnd = m_scripts.constEnd();
        it != itEnd; ++it) {
+
+    qDebug() << " register " << (*it)->identifier().c_str();
+
     if (!Calc::EnergyManager::registerModel((*it)->newInstance())) {
       qDebug() << "Could not register model" << (*it)->identifier().c_str()
                << "due to name conflict.";

avogadro/qtplugins/forcefield/forcefield.h

@@ -57,6 +57,8 @@ class Forcefield : public QtGui::ExtensionPlugin
 
   void setMolecule(QtGui::Molecule* mol) override;
 
+  std::string recommendedForceField() const;
+
 public slots:
   /**
    * Scan for new scripts in the Forcefield directories.
@@ -69,8 +71,10 @@ private slots:
   void energy();
   void optimize();
   void freezeSelected();
+  void unfreezeSelected();
 
 private:
+
   QList<QAction*> m_actions;
   QtGui::Molecule* m_molecule = nullptr;
 

avogadro/qtplugins/forcefield/scriptenergy.cpp

@@ -27,7 +27,7 @@
 namespace Avogadro::QtPlugins {
 
 ScriptEnergy::ScriptEnergy(const QString& scriptFileName_)
-  : m_interpreter(new QtGui::PythonScript(scriptFileName_)), m_valid(false),
+  : m_interpreter(new QtGui::PythonScript(scriptFileName_)), m_valid(true),
     m_inputFormat(NotUsed), m_gradients(false), m_ions(false),
     m_radicals(false), m_unitCells(false)
 {
@@ -101,6 +101,8 @@ void ScriptEnergy::setMolecule(Core::Molecule* mol)
   options << "-f" << m_tempFile.fileName();
 
   // start the interpreter
+  //@ todo - check if there was a previous process running
+  // .. if so, kill it
   m_interpreter->asyncExecute(options);
 }
 
@@ -123,6 +125,7 @@ Real ScriptEnergy::value(const Eigen::VectorXd& x)
 
 void ScriptEnergy::gradient(const Eigen::VectorXd& x, Eigen::VectorXd& grad)
 {
+  //@ todo read the gradient if available from the process
   EnergyCalculator::gradient(x, grad);
 }
 
@@ -290,6 +293,8 @@ void ScriptEnergy::readMetaData()
   // get the element mask
   // (if it doesn't exist, the default is no elements anyway)
   m_valid = parseElements(metaData);
+
+  qDebug() << " finished parsing metadata ";
 }
 
 bool ScriptEnergy::parseString(const QJsonObject& ob, const QString& key,

avogadro/qtplugins/forcefield/scripts/gfnff.py

@@ -40,6 +40,9 @@ def run(filename):
     with open(filename, "r") as f:
         mol_cjson = json.load(f)
 
+    with open("/Users/ghutchis/gfnff.log", "a") as f:
+        f.write(filename + "\n")
+
     # first setup the calculator
     atoms = np.array(mol_cjson["atoms"]["elements"]["number"])
     coord_list = mol_cjson["atoms"]["coords"]["3d"]
@@ -67,6 +70,9 @@ def run(filename):
         # first print the energy of these coordinates
         print(res.get_energy())  # in Hartree
 
+        with open("/Users/ghutchis/gfnff.log", "a") as f:
+            f.write(str(res.get_energy()) + "\n")
+
         # now print the gradient
         # .. we don't want the "[]" in the output
         output = np.array2string(res.get_gradient())