ForestExport.cpp

#include <iostream>
#include <fstream>
#include <iomanip>
#include <vector>
#include <map>
#include <set>
#include <utility>
#include "ForestExport.h"
#include "Forest.h"

void ForestExport::exportForest(const char *aFilename, size_t aCounter) const {
  std::vector<std::pair<int, int> > node_from;
  std::vector<std::pair<int, int> > node_to;
  std::vector<double> branch_length;

  // Get all forest connections
  std::vector<ForestNode>::const_iterator ifn = mForest.mRoots.begin();
  for (; ifn != mForest.mRoots.end(); ++ifn) {
    exportForestWalker(&(*ifn), mForest.mBranchLengths, node_from, node_to,
                       branch_length);
  }

  // Remove duplicated nodes
  std::set<std::pair<int, int> > vertices;
  std::vector<std::pair<int, int> >::const_iterator ip = node_from.begin();
  for (; ip != node_from.end(); ++ip)
    vertices.insert(*ip);
  for (ip = node_to.begin(); ip != node_to.end(); ++ip)
    vertices.insert(*ip);

  // Convert to node indices
  std::map<std::pair<int, int>, int> map;
  int idx;
  std::set<std::pair<int, int> >::const_iterator iv = vertices.begin();
  for (idx = 1; iv != vertices.end(); ++iv, ++idx) {
    std::pair<int, int> p(iv->first, iv->second);
    map[p] = idx;
  }

  // Map values to branches (identified by the end node)
  std::map<std::pair<int, int>, double> map_value;
  for (size_t i = 0; i < node_to.size(); ++i) {
    map_value[node_to[i]] = branch_length[i];
  }

  // Check if the filename contains a %03d format
  char temp_filename[1024];
  if (strrchr(aFilename, '%')) {
    sprintf(temp_filename, aFilename, static_cast<unsigned int>(aCounter));
    aFilename = temp_filename;
  } else if (strrchr(aFilename, '@')) {
    char z[1024];
    strncpy(z, aFilename, 1023);
    z[1023] = '\0';
    char *p = strrchr(z, '@');
    *p = '%';
    sprintf(temp_filename, z, static_cast<unsigned int>(aCounter));
    aFilename = temp_filename;
  }

  // Open the file and write the forest
  std::ofstream net(aFilename, std::ios_base::trunc | std::ios_base::out);
  if (!net.good()) {
    std::cout << "Cannot create net file <" << aFilename << ">" << std::endl;
  } else {
    net << "graph [\n";
    net << "comment \"Created by FastCodeML\"\n";
    net << "directed 1\n";
    net << "Version 1\n";

    for (iv = vertices.begin(), idx = 1; iv != vertices.end(); ++iv, ++idx) {
      net << "node [\n";
      net << "   id " << idx << '\n';
      if (iv->second == 0) {
        net << "   label \"Root " << iv->first << "\"\n";
        net << "   type 0\n";
      } else {
        std::string s = mForest.mNodeNames[iv->second];
        if (s.empty())
          net << "   label \"" << iv->second << "\"\n";
        else
          net << "   label \"" << s << "\"\n";
        net << "   type 1\n";
      }
      net << "]\n";
    }

    for (size_t i = 0; i < node_from.size(); ++i) {
      std::pair<int, int> pf(node_from[i].first, node_from[i].second);
      std::pair<int, int> pt(node_to[i].first, node_to[i].second);
      bool same_tree = node_from[i].first == node_to[i].first;

      net << "edge [\n";
      net << "   source " << map[pf] << '\n';
      net << "   target " << map[pt] << '\n';
      net << "   label \"" << std::fixed << std::setprecision(2)
          << map_value[pt] << (same_tree ? "" : "+")
          << "\"\n"; // Trailing plus means not same tree link
      net << "]\n";
    }
    net << "]\n";
    net.close();
  }
}

void ForestExport::exportForestWalker(
    const ForestNode *aNode, const std::vector<double> &aBranchLengths,
    std::vector<std::pair<int, int> > &aNodeFrom,
    std::vector<std::pair<int, int> > &aNodeTo,
    std::vector<double> &aLength) const {
  int my_node_id = aNode->mBranchId + 1;
  int my_tree_id = aNode->mOwnTree;

  const unsigned int nc = aNode->mChildrenCount;
  for (unsigned int i = 0; i < nc; ++i) {
    ForestNode *n = aNode->mChildrenList[i];
    int your_node_id = n->mBranchId + 1;
    int your_tree_id = n->mOwnTree;

    std::pair<int, int> p_from(my_tree_id, my_node_id);
    std::pair<int, int> p_to(your_tree_id, your_node_id);

    aNodeFrom.push_back(p_from);
    aNodeTo.push_back(p_to);
    aLength.push_back(mForest.mBranchLengths[your_node_id]);

    if (your_tree_id == my_tree_id)
      exportForestWalker(n, aBranchLengths, aNodeFrom, aNodeTo, aLength);
  }
}