Add some simple minimization code. (#813)

* DRR - Start adding potential terms

* Add setup routine

* Add the force calc

* Update comments

* Add EnergyArray class

* Use EnergyArray

* Start setting up potential function

* Add setup and force calc functions.

* Add a steepest descent min class

* Add total energy. Put mask setup and DoF calc in the potential function.

* Default unknown traj format

* Update dependecies

* Add emin command

* Enable emin command. Ensure space allocated for forces in frame

* Allow initial step size to be varied.

* Add dx0 and out keywords

* Fix output, change default rms tol to match sander

* Add min test

* Enable basic minimization test.

* Add code comments.

* Revision bump for adding 'emin' command (still hidden)

* Print a warning that this command is still under development in case someone stumbles onto it

* No need to write an output traj

src/Command.cpp

@@ -40,6 +40,7 @@
 #include "Exec_CombineCoords.h"
 #include "Exec_CrdAction.h"
 #include "Exec_CrdOut.h"
+#include "Exec_Emin.h"
 #include "Exec_LoadCrd.h"
 #include "Exec_LoadTraj.h"
 #include "Exec_PermuteDihedrals.h"
@@ -246,6 +247,7 @@ void Command::Init() {
   Command::AddCmd( new Exec_CombineCoords(),    Cmd::EXE, 1, "combinecrd" ); 
   Command::AddCmd( new Exec_CrdAction(),        Cmd::EXE, 1, "crdaction" );
   Command::AddCmd( new Exec_CrdOut(),           Cmd::EXE, 1, "crdout" );
+  Command::AddCmd( new Exec_Emin(),             Cmd::EXE, 1, "emin");
   Command::AddCmd( new Exec_LoadCrd(),          Cmd::EXE, 1, "loadcrd" );
   Command::AddCmd( new Exec_LoadTraj(),         Cmd::EXE, 1, "loadtraj" );
   Command::AddCmd( new Exec_PermuteDihedrals(), Cmd::EXE, 1, "permutedihedrals" );

src/EnergyArray.cpp

@@ -0,0 +1,24 @@
+#include "EnergyArray.h"
+#include "CpptrajStdio.h"
+
+/** CONSTRUCTOR */
+EnergyArray::EnergyArray() :
+  ene_((int)N_E_TERMS, 0.0)
+{}
+
+const char* EnergyArray::TypeStr_[] = { "Bond", 0 };
+
+/** Specify that an energy term will be calculated.
+  * \return Pointer to position in energy array of specified term.
+  */
+double* EnergyArray::AddType(Type typeIn) {
+  for (Tarray::const_iterator it = activeTerms_.begin(); it != activeTerms_.end(); ++it)
+  {
+    if (*it == typeIn) {
+      mprinterr("Error: Energy term %s already present.\n", TypeStr_[typeIn]);
+      return 0;
+    }
+  }
+  activeTerms_.push_back( typeIn );
+  return ( (&ene_[0]) + ((int)typeIn) );
+}

src/EnergyArray.h

@@ -0,0 +1,35 @@
+#ifndef INC_ENERGYARRAY_H
+#define INC_ENERGYARRAY_H
+#include <vector>
+class EnergyArray {
+  public:
+    EnergyArray();
+    /// Energy term types
+    enum Type { E_BOND = 0, N_E_TERMS };
+    /// \return Pointer to specified part of the energy array.
+    double* AddType(Type);
+    /// Clear all terms.
+    void clear() { activeTerms_.clear(); }
+    /// Zero all active terms.
+    void zero() {
+      for (Tarray::const_iterator it = activeTerms_.begin(); it != activeTerms_.end(); ++it)
+        ene_[*it] = 0.0;
+    }
+    /// \return Total of all active terms
+    double Total() const {
+      double etotal = 0.0;
+      for (Tarray::const_iterator it = activeTerms_.begin(); it != activeTerms_.end(); ++it)
+        etotal += ene_[*it];
+      return etotal;
+    }
+  private:
+    static const char* TypeStr_[];
+
+    typedef std::vector<double> Darray;
+    typedef std::vector<Type> Tarray;
+
+    Darray ene_;         ///< Energy array
+    Tarray activeTerms_; ///< Active energy terms
+
+};
+#endif

src/Exec_Emin.cpp

@@ -0,0 +1,91 @@
+#include "Exec_Emin.h"
+#include "CpptrajStdio.h"
+#include "PotentialFunction.h"
+#include "Minimize_SteepestDescent.h"
+
+// Exec_Emin::Help()
+void Exec_Emin::Help() const
+{
+  mprintf("\tcrdset <name> [trajoutname <name>] [rmstol <tol>] [nsteps <#>]\n"
+          "\t[<mask>] [frame <#>] [dx0 <step0>] [out <file>]\n");
+}
+
+// Exec_Emin::Execute()
+Exec::RetType Exec_Emin::Execute(CpptrajState& State, ArgList& argIn)
+{
+  mprintf("Warning: THIS COMMAND IS STILL UNDER DEVELOPMENT.\n");
+  PotentialFunction potential;
+  potential.AddTerm( PotentialTerm::BOND );
+  Minimize_SteepestDescent SD;
+
+  std::string setname = argIn.GetStringKey("crdset");
+  if (setname.empty()) {
+    mprinterr("Error: Specify COORDS set to minimize with 'crdset'\n");
+    return CpptrajState::ERR;
+  }
+  DataSet_Coords* crdset = (DataSet_Coords*)State.DSL().FindSetOfGroup(setname, DataSet::COORDINATES);
+  if (crdset == 0) {
+    mprinterr("Error: No COORDS set found with name '%s'\n", setname.c_str());
+    return CpptrajState::ERR;
+  }
+
+  // Get the frame # to be minimized.
+  int framenum = argIn.getKeyInt("frame", 0);
+  mprintf("\tMinimizing COORDS set '%s' frame %i\n", crdset->legend(), framenum+1);
+
+  if (framenum < 0 || framenum >= (int)crdset->Size()) {
+    mprinterr("Error: Frame %i is out of range.\n", framenum+1);
+    return CpptrajState::ERR;
+  }
+
+  // Get the frame. Instead of using AllocateFrame, allocate manually because
+  // we need to ensure space for forces.
+  //Frame frameIn = crdset->AllocateFrame();
+  Frame frameIn;
+  CoordinateInfo cinfo = crdset->CoordsInfo();
+  cinfo.SetForce( true );
+  frameIn.SetupFrameV(crdset->Top().Atoms(), cinfo);
+  crdset->GetFrame(framenum, frameIn);
+
+  std::string trajoutname = argIn.GetStringKey("trajoutname");
+  if (!trajoutname.empty())
+    mprintf("\tOutput trajectory: %s\n", trajoutname.c_str());
+
+  CpptrajFile* outfile = State.DFL().AddCpptrajFile(argIn.GetStringKey("out"),
+                                                    "Min. Out", DataFileList::TEXT, true);
+  if (outfile == 0) {
+    mprinterr("Internal Error: Could not allocate output file for minimization.\n");
+    return CpptrajState::ERR;
+  }
+  mprintf("\tOutput to %s\n", outfile->Filename().full());
+
+  double min_tol = argIn.getKeyDouble("rmstol", 1E-4);
+  mprintf("\tMin RMS tolerance: %g\n", min_tol);
+
+  double dx0 = argIn.getKeyDouble("dx0", 0.01);
+  mprintf("\tInitial step size: %g\n", dx0);
+
+  int nMinSteps = argIn.getKeyInt("nsteps", 1);
+  mprintf("\t%i minimization steps.\n", nMinSteps);
+
+  std::string maskexpr = argIn.GetMaskNext();
+  if (!maskexpr.empty())
+    mprintf("\tMask expression: %s\n", maskexpr.c_str());
+
+  if (potential.SetupPotential( crdset->Top(), maskexpr )) {
+    mprinterr("Error: Could not set up potential.\n");
+    return CpptrajState::ERR;
+  }
+
+  if (SD.SetupMin(trajoutname, min_tol, dx0, nMinSteps)) {
+    mprinterr("Error: Could not set up minimizer.\n");
+    return CpptrajState::ERR;
+  }
+
+  if (SD.RunMin(potential, frameIn, *outfile)) {
+    mprinterr("Error: Minimization failed.\n");
+    return CpptrajState::ERR;
+  }
+
+  return CpptrajState::OK;
+}

src/Exec_Emin.h

@@ -0,0 +1,12 @@
+#ifndef INC_EXEC_EMIN_H
+#define INC_EXEC_EMIN_H
+#include "Exec.h"
+/// <Enter description of Exec_Emin here>
+class Exec_Emin : public Exec {
+  public:
+    Exec_Emin() : Exec(HIDDEN) {}
+    void Help() const;
+    DispatchObject* Alloc() const { return (DispatchObject*)new Exec_Emin(); }
+    RetType Execute(CpptrajState&, ArgList&);
+};
+#endif

src/Minimize_SteepestDescent.cpp

@@ -0,0 +1,112 @@
+#include "Minimize_SteepestDescent.h"
+#include <algorithm> // std::fill
+#include <cmath> // sqrt
+#include "PotentialFunction.h"
+#include "Frame.h"
+#include "Trajout_Single.h"
+#include "CpptrajStdio.h"
+#include "ArgList.h"
+#include "DataSetList.h"
+#include "CpptrajFile.h"
+
+/** CONSTRUCTOR */
+Minimize_SteepestDescent::Minimize_SteepestDescent() :
+  min_tol_(1.0E-5),
+  dx0_(0.01),
+  nMinSteps_(1)
+{}
+
+/** Set up minimization. */
+int Minimize_SteepestDescent::SetupMin(std::string const& nameIn, double tolIn, double dx0In,
+                                       int stepsIn)
+{
+  trajoutName_ = nameIn;
+  min_tol_ = tolIn;
+  dx0_ = dx0In;
+  nMinSteps_ = stepsIn;
+  return 0;
+}
+
+/** Run minimization. */
+int Minimize_SteepestDescent::RunMin(PotentialFunction& potential, Frame& frameIn,
+                                     CpptrajFile& outfile)
+const
+{
+  // Output trajectory
+  int iteration = 0;
+  Trajout_Single trajOut; // TODO change type
+  if (!trajoutName_.empty()) {
+    if (trajOut.InitTrajWrite(trajoutName_, ArgList(), DataSetList(), TrajectoryFile::UNKNOWN_TRAJ))
+      return 1;
+    if (trajOut.SetupTrajWrite(potential.CurrentTop(), CoordinateInfo(), 0))
+      return 1;
+    if (trajOut.WriteSingle(iteration, frameIn)) return 1;
+  }
+
+  // Zero forces
+  if (!frameIn.HasForce()) {
+    mprinterr("Internal Error: Frame not set up with forces.\n");
+    return 1;
+  }
+  double* fxyz = frameIn.fAddress();
+  std::fill(fxyz, fxyz + frameIn.size(), 0.0);
+
+  // Degrees of freedom
+  double deg_of_freedom = potential.DegreesOfFreedom(); 
+  double fnq = sqrt(deg_of_freedom);
+  // Main loop for steepest descent
+  const double dxstm = 1.0E-5;
+  const double crits = 1.0E-6;
+  double rms = 1.0;
+  double dxst = dx0_;
+  double last_e = 0.0;
+  outfile.Printf("%-8s %12s %12s\n", "#Iter.", "ENE", "RMS");
+  while (rms > min_tol_ && iteration < nMinSteps_) {
+    if (potential.CalculateForce( frameIn )) {
+      mprinterr("Error: Could not calculate force.\n");
+      return 1;
+    }
+    double e_total = potential.Energy().Total();
+
+    // Calculate the magnitude of the force vector.
+    double sum = 0.0;
+    fxyz = frameIn.fAddress();
+    for (int idx = 0; idx < frameIn.Natom(); idx++, fxyz += 3)
+      sum += (fxyz[0]*fxyz[0] + fxyz[1]*fxyz[1] + fxyz[2]*fxyz[2]);
+    rms = sqrt( sum ) / fnq;
+    // Adjust search step size
+    if (dxst < crits) dxst = dxstm;
+    dxst = dxst / 2.0;
+    if (e_total < last_e) dxst = dxst * 2.4;
+    double dxsth = dxst / sqrt( sum );
+    last_e = e_total;
+    // Update positions and reset force array.
+    double* Xptr = frameIn.xAddress();
+    fxyz = frameIn.fAddress();
+    for (int idx = 0; idx != frameIn.Natom(); idx++, Xptr += 3, fxyz += 3)
+    {
+      //mprintf("xyz0= %g %g %g  Fxyz= %g %g %g\n", Xptr[0], Xptr[1], Xptr[2], fxyz[0], fxyz[1], fxyz[2]);
+      Xptr[0] += fxyz[0] * dxsth;
+      Xptr[1] += fxyz[1] * dxsth;
+      Xptr[2] += fxyz[2] * dxsth;
+      fxyz[0] = 0.0;
+      fxyz[1] = 0.0;
+      fxyz[2] = 0.0;
+      //mprintf("xyz1= %g %g %g\n", Xptr[0], Xptr[1], Xptr[2]);
+      //*XV += (*FV * dxsth);
+      //*FV = 0.0;
+    }
+    // Write out current E.
+    //outfile.Printf("Iteration:\t%8i %12.4E %12.4E\n", iteration, e_total, rms);
+    outfile.Printf("%-8i %12.4E %12.4E\n", iteration+1, e_total, rms);
+    //mprintf("Iteration:\t%8i %12.4E %12.4E EB=%12.4E EV=%12.4E EC=%12.4E\n",
+    //        iteration, e_total, rms, E_bond, E_vdw, E_elec);
+    iteration++;
+    // Write out current coords
+    if (trajOut.IsInitialized()) {
+      if (trajOut.WriteSingle(iteration, frameIn)) return 1;
+    }
+  } // END minimization loop
+
+  return 0;
+}

src/Minimize_SteepestDescent.h

@@ -0,0 +1,20 @@
+#ifndef INC_MINIMIZE_STEEPESTDESCENT_H
+#define INC_MINIMIZE_STEEPESTDESCENT_H
+#include <string>
+class PotentialFunction;
+class Frame;
+class CpptrajFile;
+class Minimize_SteepestDescent {
+  public:
+    Minimize_SteepestDescent();
+    /** Set up minimization with optional output trajectory, RMS tolerance, step size, # steps. */
+    int SetupMin(std::string const&, double, double,int);
+    /** Run minimization with given potential function and coordinates. */
+    int RunMin(PotentialFunction&, Frame&, CpptrajFile&) const;
+  private:
+    std::string trajoutName_;
+    double min_tol_; ///< Min RMS tolerance
+    double dx0_;     ///< Initial step size
+    int nMinSteps_;  ///< Number of minimization steps.
+};
+#endif

src/PotentialFunction.cpp

@@ -0,0 +1,63 @@
+#include "PotentialFunction.h"
+#include "CpptrajStdio.h"
+#include "Topology.h"
+// ----- All potential terms -----------
+#include "PotentialTerm_Bond.h"
+
+/** Add a term to the potential function. */
+int PotentialFunction::AddTerm(PotentialTerm::Type typeIn) {
+  PotentialTerm* term = 0;
+  switch (typeIn) {
+    case PotentialTerm::BOND : term = (PotentialTerm*)new PotentialTerm_Bond(); break;
+    default :
+      mprinterr("Internal Error: No allocator type for potential term.\n");
+      return 1;
+  }
+  if (term == 0) {
+    mprinterr("Internal Error: Could not allocate potential term.\n");
+    return 1;
+  }
+  terms_.push_back( term );
+  return 0;
+}
+
+/** Set up each term of the potential function. */
+int PotentialFunction::SetupPotential(Topology const& topIn, std::string const& maskExprIn) {
+  // First set up the mask
+  mask_.ResetMask();
+  if (mask_.SetMaskString( maskExprIn )) {
+    mprinterr("Error: Could not set up mask expression.\n");
+    return 1;
+  }
+  if (topIn.SetupCharMask( mask_ )) {
+    mprinterr("Error: Could not set up mask.\n");
+    return 1;
+  }
+  mask_.MaskInfo();
+
+  // Determine degrees of freedom
+  // TODO depending on what terms are present and how they are set up the DoF calc may change
+  deg_of_freedom_ = 3 * mask_.Nselected();
+  mprintf("\t%i degrees of freedom.\n", deg_of_freedom_);
+
+  earray_.clear();
+  for (Parray::const_iterator it = terms_.begin(); it != terms_.end(); ++it)
+  {
+    if ( (*it)->SetupTerm( topIn, mask_, earray_ ) ) {
+      mprinterr("Error: Could not set up energy term.\n");
+      return 1;
+    }
+  }
+  current_ = (Topology*)&(topIn);
+  return 0;
+}
+
+/** Calculate force for each term. */
+int PotentialFunction::CalculateForce(Frame& frameIn) {
+  earray_.zero();
+  for (Parray::const_iterator it = terms_.begin(); it != terms_.end(); ++it)
+  {
+    (*it)->CalcForce( frameIn, mask_ );
+  }
+  return 0;
+}