3D Animation of RNA structure

#include "visual/Whiteboard.h"

#include "base/CommandLineParser.h"
#include "base/FileParser.h"
#include "base/SVector.h"
#include "visual/Color.h"

#include "visual/Axes.h"
#include "visual/Geometry.h"
#include <math.h>
#include <iostream>


// Define a class that holds the coordinates of a 
// tertiary RNA structure
class Structure
{
public:
  Structure() {}

  // Read data from file (see sample_data/structure_data.txt)
  void Read(const string & fileName) {
    FlatFileParser parser;
  
    parser.Open(fileName);
    parser.ParseLine();
    m_data = parser.AsString(0);

    parser.ParseLine();
    m_paired.resize(m_data.size(), -1);

    while (parser.ParseLine()) {
      if (parser.AsString(0) == "</coords>")
	break;
      m_x.push_back(parser.AsFloat(0));
      m_y.push_back(0.);
      m_z.push_back(parser.AsFloat(1));
   }

    parser.ParseLine();
    while (parser.ParseLine()) {
      if (parser.AsString(0) == "</pairs>")
	break;
      int a = parser.AsInt(0) - 1;
      int b = parser.AsInt(1) - 1;
      m_paired[a] = b;
    }
 
  }
  
  // Return number of nucleotides
  int isize() const {return m_paired.isize();}

  // Return the nucleotide
  string Nuc(int i) const {
    char N[64];
    N[0] = m_data[i];
    N[1] = 0;
    return N;
  }

  // Return base coordinates (static or dynamic)
  double X(int i) const {return m_x[i];}
  double Y(int i) const {return m_y[i];}
  double Z(int i) const {return m_z[i];}
 
  // Which ones are paired up
  int Pair(int i) const {return m_paired[i];}

  // Simulate folding in the y coordinate
  void Change(int pos, double y) {
    int i, j;
    m_y[pos] += y;
    if (m_paired[pos] >= 0) {
      m_y[m_paired[pos]] += y;
    } else {
      cout << "Unpaired: " << pos << endl;
    }

    for (j=0; j<50; j++) {
      svec<double> tmp;
      tmp.resize(m_y.isize(), 0.);
      tmp[0] = (m_y[0] + m_y[1])/2.;
      tmp[m_y.isize()-1] = (m_y[m_y.isize()-1] + m_y[m_y.isize()-2])/2.;
      for (i=1; i<m_y.isize()-1; i++) {
	tmp[i] = (m_y[i-1] + m_y[i] + m_y[i+1])/3.;
	//cout << i << " " << tmp[i] << endl;
      }
      m_y = tmp;
    }
  }

private:
  string m_data;
  svec<int> m_paired;
  svec<double> m_x;
  svec<double> m_y;
  svec<double> m_z;
  
};


// Print one frame for the animation
void OneFrame(const string & o, // output file
	      double phi,         // rotation/phi
	      double theta,       // rotation/theta
	      const Structure & fold)
{
  double x_offset = 20;
  double y_offset = 20;


  int i, j;
  ns_whiteboard::whiteboard board;

 
  double x_max = 600.;
  double y_max = 800.;

  // Draw a rectangle as the background
  board.Add( new ns_whiteboard::rect( ns_whiteboard::xy_coords(0, 0), 
  				      ns_whiteboard::xy_coords(x_max + 2*x_offset, y_max + 2*y_offset),
  				      color(0.99,0.99,0.99)));
  
  // Use the 3D geometry class
  Geometry3D g;

  // Let's center 0 at 350
  g.SetOffset(350.);

  // Pass the rotation to the geometry, which will figure out the rest
  g.SetRotation(phi, theta);
  
  // First, let's plot the connecting lines between pairs (hard-coded light grey)
  for (i=0; i<fold.isize(); i++) {
    int p = fold.Pair(i);
    if (p < 0)
      continue;
    
 
    board.Add( new ns_whiteboard::line( g.Coords(fold.X(i)+5., fold.Y(i), fold.Z(i)+5.), 
					g.Coords(fold.X(p)+5., fold.Y(p), fold.Z(p)+5.),
					1.,
					color(0.6,0.6,0.6)));
    
 
  }
    

  // Let's draw the nucleotides...
  for (i=0; i<fold.isize(); i++) {
    color cc(0., 0., 0.);
    //...and color them in.
    if (fold.Nuc(i) == "A")
      cc = color(0.0, 0.8, 0.);
    if (fold.Nuc(i) == "C")
      cc = color(0., 0., 0.7);
    if (fold.Nuc(i) == "G")
      cc = color(0.8, 0., 0.);
      
    board.Add( new ns_whiteboard::text( g.Coords(fold.X(i), fold.Y(i), fold.Z(i)),
					fold.Nuc(i), cc, 12.*g.Norm(fold.X(i), fold.Y(i), fold.Z(i)), 
					"Times-Roman", 0, true));
  }


  // Display and save
  ofstream out(o.c_str());
  
  ns_whiteboard::ps_display display(out, x_max + 2 * x_offset, y_max + 2 * y_offset);
  board.DisplayOn(&display);
 
}

int main( int argc, char** argv )
{
  // COmmand line processing
  commandArg<string> aStringI("-i","input file");
  commandArg<string> aStringO("-o","outfile (post-script)");
  commandLineParser P(argc,argv);
  P.SetDescription("RNA folding animation example");

  P.registerArg(aStringI);
  P.registerArg(aStringO);

  P.parse();

  string o = P.GetStringValueFor(aStringO);
  string in = P.GetStringValueFor(aStringI);

  Structure fold;
  fold.Read(in);
  cout << "Done reading." << endl;
  

  int i;
  int k = 1000;
  double phi = 0.;
  double theta = 0.;

  // Let's make 150 frames
  for (i=0; i<150; i++) {
     char name[256];
    sprintf(name, "%s%d.ps", o.c_str(), k);
    
    phi = 0.05*i;
    theta = 0.03*i;

    // Change the structure
    fold.Change(10, 500.);
    fold.Change(36, -500.);

    //Draw one frame
    OneFrame(name, phi, theta, fold);
    k++;
  }

  //==================================================================================
  cout << " Turn output into animation via 'convert <out>*.ps animated.gif' on the command line." << endl;
 
  return 0;
}