Allow variables in mask expressions in nested 'for mask' loops (#941)

* Add variables to be used inside mask expressions

* Fix DoTest. Add another nested test

* Change test name

* Comment out debug

* Protect in case user uses parentheses in integer for loop

* Version 6.2.5; revision bump for allowing variables in mask expressions
in nest for mask loops.

src/ForLoop_integer.cpp

@@ -33,7 +33,7 @@ int ForLoop_integer::SetupFor(CpptrajState& State, ArgList& argIn) {
     mprinterr("Internal Error: Too many arguments for integer for loop.\n");
     return 1;
   }
-  ArgList varArg( argIn.GetStringNext(), ";" );
+  ArgList varArg( argIn.GetStringNext(), "();" );
   //varArg.PrintDebug(); // DEBUG
   if (varArg.Nargs() < 2 || varArg.Nargs() > 3) {
     mprinterr("Error: Malformed 'for' loop variable.\n"
@@ -62,6 +62,7 @@ int ForLoop_integer::SetupFor(CpptrajState& State, ArgList& argIn) {
   // Second argument: <var><OP><end>
   size_t pos0 = VarName().size();
   size_t pos1 = pos0 + 1;
+  //mprintf("DEBUG: VarName='%s' pos0 %zu pos1 %zu\n", VarName().c_str(),pos0, pos1);
   endOp_ = NO_OP;
   int iargIdx = 1;
   if (varArg.Nargs() == 3) {

src/ForLoop_mask.cpp

@@ -5,6 +5,12 @@
 #include "StringRoutines.h"
 #include "DataSetList.h"
 
+/** CONSTRUCTOR */
+ForLoop_mask::ForLoop_mask() :
+  mtype_(NTYPES),
+  currentTop_(0)
+{}
+
 void ForLoop_mask::helpText() {
   mprintf("\t{atoms|residues|molecules|molfirstres|mollastres}\n"
           "\t  <var> inmask <mask> [%s]\n", DataSetList::TopIdxArgs);
@@ -15,7 +21,7 @@ int ForLoop_mask::SetupFor(CpptrajState& State, ArgList& argIn) {
   static const char* TypeStr[NTYPES] = { "ATOMS ", "RESIDUES ", "MOLECULES ",
                                    "MOL_FIRST_RES ", "MOL_LAST_RES " };
   // {atoms|residues|molecules} <var> inmask <mask> [TOP KEYWORDS]
-  Topology* currentTop = 0;
+  currentTop_ = 0;
   mtype_ = NTYPES;
   if      (argIn.hasKey("atoms")) mtype_ = ATOMS;
   else if (argIn.hasKey("residues")) mtype_ = RESIDUES;
@@ -34,24 +40,51 @@ int ForLoop_mask::SetupFor(CpptrajState& State, ArgList& argIn) {
     mprinterr("Error: mask for loop does not contain 'inmask' keyword.\n");
     return 1;
   }
-  AtomMask currentMask;
-  if (currentMask.SetMaskString( argIn.GetStringKey("inmask") )) return 1;
+  maskExpr_ = argIn.GetStringKey("inmask");
+  if (maskExpr_.empty()) {
+    mprinterr("Error: Must set mask expression via 'inmask'.\n");
+    return 1;
+  }
   Topology* top = State.DSL().GetTopByIndex( argIn );
-  if (top != 0) currentTop = top;
-  if (currentTop == 0) return 1;
+  if (top != 0) currentTop_ = top;
+  if (currentTop_ == 0) return 1;
   // Get the variable name
   if (SetupLoopVar( State.DSL(), argIn.GetStringNext() )) return 1;
+
+  std::string description("(" + std::string(TypeStr[mtype_]) +
+                          VarName() + " inmask " + maskExpr_ + ")");
+  SetDescription( description );
+  return 0;
+}
+
+/** Set up the atom mask, obtain indices. */
+int ForLoop_mask::BeginFor(DataSetList const& DSL) {
+  if (currentTop_ == 0) {
+    mprinterr("Internal Error: ForLoop_mask::BeginFor() called before Setup().\n");
+    return LOOP_ERROR;
+  }
+  Idxs_.clear();
+  AtomMask currentMask;
+  // Do variable replacement
+  std::string maskExpr2;
+  int nReplaced = DSL.ReplaceVariables( maskExpr2, maskExpr_ );
+  if (nReplaced > 0) {
+    //mprintf("DEBUG: old mask '%s' new mask '%s'\n", maskExpr_.c_str(), maskExpr2.c_str());
+    if (currentMask.SetMaskString( maskExpr2 )) return LOOP_ERROR;
+  } else {
+    if (currentMask.SetMaskString( maskExpr_ )) return LOOP_ERROR;
+  }
   // Set up mask
-  if (currentTop->SetupIntegerMask( currentMask )) return 1;
+  if (currentTop_->SetupIntegerMask( currentMask )) return LOOP_ERROR;
   currentMask.MaskInfo();
-  if (currentMask.None()) return 1;
+  if (currentMask.None()) return 0; // No iterations
   // Set up indices
   if (mtype_ == ATOMS)
     Idxs_ = currentMask.Selected();
   else if (mtype_ == RESIDUES) {
     int curRes = -1;
     for (AtomMask::const_iterator at = currentMask.begin(); at != currentMask.end(); ++at) {
-      int res = (*currentTop)[*at].ResNum();
+      int res = (*currentTop_)[*at].ResNum();
       if (res != curRes) {
         Idxs_.push_back( res );
         curRes = res;
@@ -63,29 +96,23 @@ int ForLoop_mask::SetupFor(CpptrajState& State, ArgList& argIn) {
   {
     int curMol = -1;
     for (AtomMask::const_iterator at = currentMask.begin(); at != currentMask.end(); ++at) {
-      int mol = (*currentTop)[*at].MolNum();
+      int mol = (*currentTop_)[*at].MolNum();
       if (mol != curMol) {
         if (mtype_ == MOLECULES)
           Idxs_.push_back( mol );
         else {
           int res;
           if (mtype_ == MOLFIRSTRES)
-            res = (*currentTop)[ currentTop->Mol( mol ).MolUnit().Front() ].ResNum();
+            res = (*currentTop_)[ currentTop_->Mol( mol ).MolUnit().Front() ].ResNum();
           else // MOLLASTRES
-            res = (*currentTop)[ currentTop->Mol( mol ).MolUnit().Back()-1 ].ResNum();
+            res = (*currentTop_)[ currentTop_->Mol( mol ).MolUnit().Back()-1 ].ResNum();
           Idxs_.push_back( res );
         }
         curMol = mol;
       }
     }
   }
-  std::string description("(" + std::string(TypeStr[mtype_]) +
-                          VarName() + " inmask " + currentMask.MaskExpression() + ")");
-  SetDescription( description );
-  return 0;
-}
 
-int ForLoop_mask::BeginFor(DataSetList const& DSL) {
   idx_ = Idxs_.begin();
   return (int)Idxs_.size();
 }

src/ForLoop_mask.h

@@ -2,11 +2,12 @@
 #define INC_FORLOOP_MASK_H
 #include <vector>
 #include "ForLoop.h"
+class Topology;
 /// For loop over mask expression
 class ForLoop_mask : public ForLoop {
     enum MaskType {ATOMS=0, RESIDUES, MOLECULES, MOLFIRSTRES, MOLLASTRES, NTYPES};
   public:
-    ForLoop_mask() : mtype_(NTYPES) {}
+    ForLoop_mask();
 
     static void helpText();
 
@@ -19,5 +20,7 @@ class ForLoop_mask : public ForLoop {
     Iarray Idxs_;                ///< (MASK only) Selected atom/residue/molecule indices
     Iarray::const_iterator idx_; ///< (MASK only) Current atom/residue/molecule index
     MaskType mtype_;             ///< Hold specific mask type
+    std::string maskExpr_;       ///< The atom mask expression
+    Topology* currentTop_;       ///< Topology to use when setting up mask.
 };
 #endif

src/Version.h

@@ -12,7 +12,7 @@
  * Whenever a number that precedes <revision> is incremented, all subsequent
  * numbers should be reset to 0.
  */
-#define CPPTRAJ_INTERNAL_VERSION "V6.2.4"
+#define CPPTRAJ_INTERNAL_VERSION "V6.2.5"
 /// PYTRAJ relies on this
 #define CPPTRAJ_VERSION_STRING CPPTRAJ_INTERNAL_VERSION
 #endif

test/Test_Control/RunTest.sh

@@ -5,7 +5,7 @@
 CleanFiles for.in TRP.vec.dat TRP.rms.dat TRP.CA.dist.dat TRP.tocenter.dat \
            nh.dat rms.nofit.dat last10.dat distance.dat nested.agr \
            EndToEnd0.dat EndToEnd1.dat EndToEnd2.agr temp.*.dat \
-           DataOut.dat testset.dist.dat
+           DataOut.dat testset.dist.dat nested2.dat
 
 TESTNAME='Loop tests'
 Requires netcdf maxthreads 10
@@ -109,7 +109,7 @@ show
 list
 EOF
 RunCpptraj "$UNITNAME"
-DoTest nested.agr nested.agr.save
+DoTest nested.agr.save nested.agr
 
 UNITNAME='Test loop over data set blocks'
 CheckFor maxthreads 1
@@ -172,5 +172,24 @@ EOF
   DoTest DataOut.dat.save DataOut.dat
 fi
 
+UNITNAME='Test nested mask loops'
+cat > for.in <<EOF
+parm ../tz2.parm7
+trajin ../tz2.nc
+reference ../tz2.nc 1 1
+
+for atoms A1 inmask :2-4@N
+  for atoms A2 inmask \$A1<@3.0&!\$A1
+    distance d\$A1.\$A2 \$A1 \$A2 out nested2.dat
+  done
+done
+
+run
+show
+list
+EOF
+RunCpptraj "$UNITNAME"
+DoTest nested2.dat.save nested2.dat
+
 EndTest
 exit 0