Fixes and improvements for 'calcstate' analysis

src/Analysis_State.h

@@ -2,35 +2,59 @@
 #define INC_ANALYSIS_STATE_H
 #include <map>
 #include "Analysis.h"
+#include "Array1D.h"
 #include "DataSet_double.h"
 /// Analyze transitions between states
 class Analysis_State : public Analysis {
   public:
-    Analysis_State() : state_data_(0), curveOut_(0), stateOut_(0), transOut_(0), debug_(0) {}
+    Analysis_State();
     DispatchObject* Alloc() const { return (DispatchObject*)new Analysis_State(); }
     void Help() const;
 
     Analysis::RetType Setup(ArgList&, AnalysisSetup&, int);
     Analysis::RetType Analyze();
   private:
     // ----- PRIVATE CLASS DEFINITIONS ---------------------------------------------
+    // -------------------------------------------
     /// Hold the definition of a state and associated data
     class StateType {
       public:
-        StateType() : set_(0), min_(0.0), max_(0.0) {}
-        StateType(std::string const& i, DataSet_1D* d, double m, double x) :
-                  id_(i), set_(d), min_(m), max_(x) {}
+        StateType() : num_(-1) {}
+        StateType(std::string const& i, int n) : id_(i), num_(n) {}
+        /// Add criterion for state
+        void AddCriterion(DataSet_1D* d, double m, double x) {
+          Sets_.push_back( d );
+          Min_.push_back( m );
+          Max_.push_back( x );
+        }
+        /// \return smallest number of frames among all criterion DataSets
+        size_t Nframes() const;
+        /// \return State ID
         std::string const& ID() const { return id_; }
+        /// \return State ID as const char*
         const char* id()        const { return id_.c_str(); }
-        DataSet_1D const& DS()  const { return *set_; }
-        double Min()            const { return min_; }
-        double Max()            const { return max_; }
+        /// \return Unique state number
+        int Num() const { return num_; }
+        /// \return true if given frame satifies all criteria
+        bool InState(int n) const {
+          for (unsigned int idx = 0; idx != Sets_.size(); idx++) {
+            double dval = Sets_[idx]->Dval(n);
+            if (dval < Min_[idx] || dval >= Max_[idx]) // TODO periodic
+              return false;
+          }
+          return true;
+        }
+        /// Print state info to stdout
+        void PrintState() const;
       private:
+        typedef std::vector<double> Darray;
         std::string id_;
-        DataSet_1D* set_;
-        double min_;
-        double max_;
+        Array1D Sets_;   ///< DataSets used to determine if we are in State
+        Darray Min_;     ///< Above this value we are in state.
+        Darray Max_;     ///< Below this value we are in state.
+        int num_;        ///< Unique state identifier
     };
+    // -------------------------------------------
     /// Hold information about a transition from state0 to state1
     class Transition {
       public:
@@ -76,23 +100,36 @@ class Analysis_State : public Analysis {
         int Nlifetimes_; ///< Number of state0 lifetimes before going to state1
         DataSet_double* curve_; ///< Lifetime curve for state0->state1
     };
+    // -------------------------------------------
 
     typedef std::vector<StateType> StateArray;
     typedef std::pair<int,int> StatePair;
     typedef std::map<StatePair, Transition> TransMapType;
     typedef std::pair<StatePair, Transition> TransPair;
     typedef std::vector<int> Iarray;
+    typedef std::pair<std::string, int> IdxPair;
+    typedef std::map<std::string, int> IdxMapType;
+    typedef std::vector<std::string> Sarray;
 
     std::string const& StateName(int) const;
     const char* stateName(int i) const { return StateName(i).c_str(); }
 
     static std::string UNDEFINED_;
     StateArray States_;
+    IdxMapType NameMap_;
+    Sarray StateNames_;
     TransMapType TransitionMap_;
     DataSet* state_data_;
+    DataSet* state_counts_;
+    DataSet* state_fracs_;
+    DataSet* state_lifetimes_;
+    DataSet* state_avglife_;
+    DataSet* state_maxlife_;
+    DataSet* state_names_;
     DataSetList* masterDSL_;
     DataFile* curveOut_;
-    CpptrajFile* stateOut_;
+    DataFile* lifeOut_;
+    DataFile* countOut_;
     CpptrajFile* transOut_;
     int debug_;
     bool normalize_;

src/DataIO_Grace.cpp

@@ -4,6 +4,7 @@
 #include "CpptrajStdio.h"
 #include "BufferedLine.h"
 #include "DataSet_double.h"
+#include "DataSet_string.h"
 
 // TODO: Set dimension labels
 // Dont assume anything about set ordering
@@ -130,10 +131,29 @@ int DataIO_Grace::WriteDataNormal(CpptrajFile& file, DataSetList const& Sets) {
   file.Printf("@with g0\n@  xaxis label \"%s\"\n@  yaxis label \"%s\"\n"
               "@  legend 0.2, 0.995\n@  legend char size 0.60\n",
               Sets[0]->Dim(0).Label().c_str(), "");
+  // Determine if we have a single STRING DataSet. Assume these contain labels.
+  DataSet_string* labelSet = 0;
+  for (DataSetList::const_iterator set = Sets.begin(); set != Sets.end(); ++set)
+  {
+    if (labelSet == 0) {
+      if ( (*set)->Type() == DataSet::STRING )
+        labelSet = (DataSet_string*)(*set);
+    } else {
+      if ( (*set)->Type() == DataSet::STRING ) {
+        labelSet = 0;
+        break;
+      }
+    }
+  }
+  if (labelSet != 0)
+    mprintf("\tUsing string set '%s' for data point labels\n", labelSet->legend());
   // Loop over DataSets
   unsigned int setnum = 0;
   DataSet::SizeArray frame(1);
-  for (DataSetList::const_iterator set = Sets.begin(); set != Sets.end(); ++set, ++setnum) {
+  for (DataSetList::const_iterator set = Sets.begin(); set != Sets.end(); ++set, ++setnum)
+  {
+    // Skip the label set if defined.
+    if (*set == labelSet) continue;
     size_t maxFrames = (*set)->Size();
     // Set information
     file.Printf("@  s%u legend \"%s\"\n@target G0.S%u\n@type xy\n",
@@ -148,6 +168,8 @@ int DataIO_Grace::WriteDataNormal(CpptrajFile& file, DataSetList const& Sets) {
     for (frame[0] = 0; frame[0] < maxFrames; frame[0]++) {
       file.Printf(xfmt.fmt(), (*set)->Coord(0, frame[0]));
       (*set)->WriteBuffer(file, frame);
+      if (labelSet != 0)
+        file.Printf(" \"%s\"", (*labelSet)[frame[0]].c_str());
       file.Printf("\n");
     }
   }

src/DataIO_Std.cpp

@@ -24,6 +24,7 @@ DataIO_Std::DataIO_Std() :
   prec_(UNSPEC),
   indexcol_(-1),
   isInverted_(false), 
+  groupByName_(false),
   hasXcolumn_(true), 
   writeHeader_(true), 
   square2d_(false),
@@ -600,20 +601,23 @@ int DataIO_Std::Read_Mat3x3(std::string const& fname,
 
 // -----------------------------------------------------------------------------
 void DataIO_Std::WriteHelp() {
-  mprintf("\tnoheader   : Do not print header line.\n"
-          "\tinvert     : Flip X/Y axes (1D).\n"
-          "\tnoxcol     : Do not print X (index) column (1D).\n"
-          "\tsquare2d   : Write 2D data sets in matrix-like format.\n"
-          "\tnosquare2d : Write 2D data sets as '<X> <Y> <Value>'.\n"
-          "\tnosparse   : Write all 3D grid voxels (default).\n"
-          "\tsparse     : Only write 3D grid voxels with value > cutoff (default 0).\n"
+  mprintf("\tnoheader    : Do not print header line.\n"
+          "\tinvert      : Flip X/Y axes (1D).\n"
+          "\tgroupbyname : (1D) group data sets by name,\n"
+          "\tnoxcol      : Do not print X (index) column (1D).\n"
+          "\tsquare2d    : Write 2D data sets in matrix-like format.\n"
+          "\tnosquare2d  : Write 2D data sets as '<X> <Y> <Value>'.\n"
+          "\tnosparse    : Write all 3D grid voxels (default).\n"
+          "\tsparse      : Only write 3D grid voxels with value > cutoff (default 0).\n"
           "\t\tcut : Cutoff for 'sparse'; default 0.\n");
 }
 
 // DataIO_Std::processWriteArgs()
 int DataIO_Std::processWriteArgs(ArgList &argIn) {
   if (!isInverted_ && argIn.hasKey("invert"))
     isInverted_ = true;
+  if (!groupByName_ && argIn.hasKey("groupbyname"))
+    groupByName_ = true;
   if (hasXcolumn_ && argIn.hasKey("noxcol"))
     hasXcolumn_ = false;
   if (writeHeader_ && argIn.hasKey("noheader"))
@@ -677,10 +681,41 @@ int DataIO_Std::WriteData(FileName const& fname, DataSetList const& SetList)
       // Special case of 2D - may have sieved frames.
       err = WriteCmatrix(file, SetList);
     } else if (SetList[0]->Ndim() == 1) {
-      if (isInverted_)
-        err = WriteDataInverted(file, SetList);
-      else
-        err = WriteDataNormal(file, SetList);
+      if (groupByName_) {
+        DataSetList tmpdsl;
+        std::vector<bool> setIsWritten(SetList.size(), false);
+        unsigned int startIdx = 0;
+        unsigned int nWritten = 0;
+        while (nWritten < SetList.size()) {
+          std::string currentName = SetList[startIdx]->Meta().Name();
+          int firstNonMatch = -1;
+          for (unsigned int idx = startIdx; idx != SetList.size(); idx++)
+          {
+            if (!setIsWritten[idx])
+            {
+              if (currentName == SetList[idx]->Meta().Name())
+              {
+                tmpdsl.AddCopyOfSet( SetList[idx] );
+                setIsWritten[idx] = true;
+                nWritten++;
+              } else if (firstNonMatch == -1)
+                firstNonMatch = (int)idx;
+            }
+          }
+          if (firstNonMatch > -1)
+            startIdx = (unsigned int)firstNonMatch;
+          if (isInverted_)
+            err += WriteDataInverted(file, tmpdsl);
+          else
+            err += WriteDataNormal(file, tmpdsl);
+          tmpdsl.ClearAll();
+        }
+      } else {
+        if (isInverted_)
+          err = WriteDataInverted(file, SetList);
+        else
+          err = WriteDataNormal(file, SetList);
+      }
     } else if (SetList[0]->Ndim() == 2)
       err = WriteData2D(file, SetList);
     else if (SetList[0]->Ndim() == 3)

src/DataIO_Std.h

@@ -35,6 +35,7 @@ class DataIO_Std : public DataIO {
     precType prec_;    ///< 3d reads, data set precision
     int indexcol_;     ///< Read: column containing index (X) values
     bool isInverted_;  ///< For 1D writes invert X/Y.
+    bool groupByName_; ///< For 1D writes, group sets with same name together
     bool hasXcolumn_;
     bool writeHeader_;
     bool square2d_;

src/DataSetList.h

@@ -38,6 +38,8 @@ class DataSetList {
     const_iterator end()   const { return DataList_.end();   }
     /// Clear all non-Topology and non-Reference DataSets
     void Clear();
+    /// Clear entire DataSetList
+    void ClearAll();
     /// Sort sets in the DataSetList
     void Sort();
     /// True if no DataSets in list.
@@ -144,8 +146,6 @@ class DataSetList {
     void Timing() const;
 #   endif
   private:
-    /// Clear entire DataSetList
-    void ClearAll();
     /// Search for and remove specified data set if found, optionally free memory.
     DataSet* EraseSet( DataSet*, bool );
     /// Warn if DataSet not found but may be pending.