Exemplo n.º 1
0
// Exec_DataSetCmd::ModifyPoints()
Exec::RetType Exec_DataSetCmd::ModifyPoints(CpptrajState& State, ArgList& argIn, bool drop) {
  const char* mode;
  if (drop)
    mode = "Drop";
  else
    mode = "Kee";
  // Keywords
  std::string name = argIn.GetStringKey("name");
  int start = argIn.getKeyInt("start", 0) - 1;
  int stop = argIn.getKeyInt("stop", -1);
  int offset = argIn.getKeyInt("offset", -1);
  Range points;
  if (start < 0 && stop < 0 && offset < 0) {
    std::string rangearg = argIn.GetStringKey("range");
    if (rangearg.empty()) {
      mprinterr("Error: Must specify range or start/stop/offset.\n");
      return CpptrajState::ERR;
    }
    points.SetRange( rangearg );
    if (points.Empty()) {
      mprinterr("Error: Range '%s' is empty.\n", rangearg.c_str());
      return CpptrajState::ERR;
    }
    mprintf("\t%sping points in range %s\n", mode, rangearg.c_str());
    // User args start from 1
    points.ShiftBy(-1);
  }
  // Get data set to drop/keep points from
  // Loop over all DataSet arguments 
  std::string ds_arg = argIn.GetStringNext();
  while (!ds_arg.empty()) {
    DataSetList dsl = State.DSL().GetMultipleSets( ds_arg );
    for (DataSetList::const_iterator it = dsl.begin(); it != dsl.end(); ++it)
    {
      DataSet* DS = *it;
      if (DS->Size() < 1) {
        mprinterr("Error: Set '%s' is empty.\n", DS->legend());
        return CpptrajState::ERR;
      }
      // Restrict to 1D sets for now TODO more types
      if (DS->Group() != DataSet::SCALAR_1D) {
        mprinterr("Error: Currently only works for 1D scalar data sets.\n");
        return CpptrajState::ERR;
      }
      DataSet_1D* ds1 = (DataSet_1D*)DS;
      // Output data set
      DataSet* out = 0;
      if (name.empty()) {
        // Modifying this set. Create new temporary set.
        out = State.DSL().Allocate( ds1->Type() );
        if (out == 0) return CpptrajState::ERR;
        *out = *ds1;
        mprintf("\tOverwriting set '%s'\n", ds1->legend());
      } else {
        // Write to new set
        MetaData md = ds1->Meta();
        md.SetName( name );
        out = State.DSL().AddSet(ds1->Type(), md);
        if (out == 0) return CpptrajState::ERR;
        mprintf("\tNew set is '%s'\n", out->Meta().PrintName().c_str());
      }
      out->Allocate(DataSet::SizeArray(1, ds1->Size()));
      if (points.Empty()) {
        // Drop by start/stop/offset. Set defaults if needed
        if (start < 0)  start = 0;
        if (stop < 0)   stop = ds1->Size();
        if (offset < 0) offset = 1;
        mprintf("\t%sping points from %i to %i, step %i\n", mode, start+1, stop, offset);
        for (int idx = start; idx < stop; idx += offset)
          points.AddToRange( idx );
      } // TODO check that range values are valid?
      if (State.Debug() > 0) mprintf("DEBUG: Keeping points:");
      Range::const_iterator pt = points.begin();
      int idx = 0;
      int odx = 0;
      if (drop) {
        // Drop points
        for (; idx < (int)ds1->Size(); idx++) {
          if (pt == points.end()) break;
          if (*pt != idx) {
            if (State.Debug() > 0) mprintf(" %i", idx + 1);
            KeepPoint(ds1, out, idx, odx);
          } else
            ++pt;
        }
        // Keep all remaining points
        for (; idx < (int)ds1->Size(); idx++) {
          if (State.Debug() > 0) mprintf(" %i", idx + 1);
          KeepPoint(ds1, out, idx, odx);
        }
      } else {
        // Keep points
        for (; pt != points.end(); pt++) {
          if (*pt >= (int)ds1->Size()) break;
          if (State.Debug() > 0) mprintf(" %i", *pt + 1);
          KeepPoint(ds1, out, *pt, odx);
        }
      }
      if (State.Debug() > 0) mprintf("\n");
      if (name.empty()) {
        // Replace old set with new set
        State.DSL().RemoveSet( ds1 );
        State.DSL().AddSet( out );
      }
    } // END loop over sets
    ds_arg = argIn.GetStringNext();
  } // END loop over set args
  return CpptrajState::OK;
}