feat: Support Newick tree format in output

Author: danieledler

src/core/InfomapBase.cpp

@@ -2101,6 +2101,26 @@ void InfomapBase::writeResult()
   }
 
 
+  if (printNewick) {
+    std::string filename = outDirectory + outName + ".tre";
+
+    if (!haveMemory()) {
+      Log() << "Write Newick tree to " << filename << "... ";
+      writeNewickTree(filename);
+      Log() << "done!\n";
+    } else {
+      // Write both physical and state level
+      Log() << "Write physical Newick tree to " << filename << "... ";
+      writeNewickTree(filename, false);
+      Log() << "done!\n";
+      std::string filenameStates = outDirectory + outName + "_states.tre";
+      Log() << "Write state Newick tree to " << filenameStates << "... ";
+      writeNewickTree(filenameStates, true);
+      Log() << "done!\n";
+    }
+  }
+
+
   if (printClu) {
     std::string filename = outDirectory + outName + ".clu";
     if (!haveMemory()) {
@@ -2153,6 +2173,16 @@ std::string InfomapBase::writeFlowTree(std::string filename, bool states)
   return outputFilename;
 }
 
+std::string InfomapBase::writeNewickTree(std::string filename, bool states)
+{
+  std::string outputFilename = filename.empty() ? outDirectory + outName + (haveMemory() && states ? "_states.tre" : ".tre") : filename;
+
+  SafeOutFile outFile(outputFilename);
+  writeNewickTree(outFile, states);
+
+  return outputFilename;
+}
+
 std::string InfomapBase::writeClu(std::string filename, bool states, int moduleIndexLevel)
 {
   // unsigned int maxModuleLevel = maxTreeDepth();
@@ -2400,6 +2430,60 @@ void InfomapBase::writeTreeLinks(std::ostream& outStream, bool states)
 }
 
 
+void InfomapBase::writeNewickTree(std::ostream& outStream, bool states)
+{
+  auto oldPrecision = outStream.precision();
+  // outStream << std::setprecision(9);
+  // outStream << getOutputFileHeader() << "\n";
+  outStream << std::setprecision(6);
+
+  auto isRoot = true;
+  unsigned int lastDepth = 0;
+  std::vector<double> flowStack;
+
+  auto writeNewickNode = [&](std::ostream& o, const InfoNode& node, unsigned int depth) {
+    if (depth > lastDepth || isRoot) {
+      outStream << "(";
+      flowStack.push_back(node.data.flow);
+      if (node.isLeaf())
+        outStream << (states ? node.stateId : node.physicalId) << ":" << node.data.flow;
+    } else if (depth == lastDepth) {
+      outStream << ",";
+      flowStack[flowStack.size()-1] = node.data.flow;
+      if (node.isLeaf()) {
+        outStream << (states ? node.stateId : node.physicalId) << ":" << node.data.flow;
+      }
+    } else {
+      // depth < lastDepth
+      while (flowStack.size() > depth + 1) {
+        flowStack.pop_back();
+        outStream << "):" << flowStack.back();
+      }
+      flowStack[flowStack.size()-1] = node.data.flow;
+      outStream << ",";
+    }
+    lastDepth = depth;
+    isRoot = false;
+  };
+  
+  // TODO: Make a general iterator where merging physical nodes depend on a parameter rather than type to be able to DRY here
+  if (haveMemory() && !states) {
+    for (auto it(iterTreePhysical()); !it.isEnd(); ++it) {
+      writeNewickNode(outStream, *it, it.depth());
+    }
+  } else {
+    for (auto it(iterTree()); !it.isEnd(); ++it) {
+      writeNewickNode(outStream, *it, it.depth());
+    }
+  }
+  while (flowStack.size() > 1) {
+    flowStack.pop_back();
+    outStream << "):" << flowStack.back();
+  }
+  outStream << ");\n";
+  outStream << std::setprecision(oldPrecision);
+}
+
 unsigned int InfomapBase::printPerLevelCodelength(std::ostream& out)
 {
   std::vector<PerLevelStat> perLevelStats;

src/core/InfomapBase.h

@@ -424,6 +424,15 @@ class InfomapBase : public InfomapConfig<InfomapBase> {
    */
   std::string writeFlowTree(std::string filename = "", bool states = false);
 
+	/**
+	 * Write Newick tree to a .tre file.
+	 * @param filename the filename for the output file. If empty, use default 
+	 * based on output directory and input file name
+	 * @param states if memory network, print the state-level network without merging physical nodes within modules
+	 * @return the filename written to
+	 */
+	std::string writeNewickTree(std::string filename = "", bool states = false);
+
   /**
    * Write tree to a .clu file.
    * @param filename the filename for the output file. If empty, use default
@@ -470,6 +479,12 @@ class InfomapBase : public InfomapConfig<InfomapBase> {
    */
   void writeTreeLinks(std::ostream& outStream, bool states = false);
 
+	/**
+	 * Write Newick tree to output stream
+	 * @param states, write state-level tree, else aggregate physical nodes within modules
+	 */
+	void writeNewickTree(std::ostream& outStream, bool states = false);
+
   InfoNode m_root;
   std::vector<InfoNode*> m_leafNodes;
   std::vector<InfoNode*> m_moduleNodes;

src/io/Config.cpp

@@ -79,7 +79,7 @@ Config::Config(std::string flags, bool isCLI) : isCLI(isCLI)
   // 		"Print the network with calculated flow values.", "Output", true);
 
   // -o network,states,clu,ftree
-  api.addOptionArgument(outputFormats, 'o', "output", "Comma-separated output formats without spaces, e.g. -o clu,tree,ftree. Options: clu, tree, ftree, network, states.", ArgType::list, "Output", true);
+  api.addOptionArgument(outputFormats, 'o', "output", "Comma-separated output formats without spaces, e.g. -o clu,tree,ftree. Options: clu, tree, ftree, newick, network, states.", ArgType::list, "Output", true);
 
   api.addOptionArgument(hideBipartiteNodes, "hide-bipartite-nodes", "Project bipartite solution to unipartite.", "Output", true);
 
@@ -239,6 +239,8 @@ void Config::adaptDefaults()
       printTree = true;
     } else if (o == "ftree") {
       printFlowTree = true;
+    } else if (o == "newick") {
+      printNewick = true;
     } else if (o == "network") {
       printPajekNetwork = true;
     } else if (o == "states") {

src/io/Config.h

@@ -142,6 +142,7 @@ struct Config {
   bool originallyUndirected = false;
   bool printTree = false;
   bool printFlowTree = false;
+  bool printNewick = false;
   bool printMap = false;
   bool printClu = false;
   int cluLevel = 1; // Write modules at specified depth from root. 1, 2, ... or -1 for bottom level
@@ -249,6 +250,7 @@ struct Config {
     originallyUndirected = other.originallyUndirected;
     // printTree = other.printTree;
     // printFlowTree = other.printFlowTree;
+    // printNewick = other.printNewick;
     // printMap = other.printMap;
     // printClu = other.printClu;
     // printNodeRanks = other.printNodeRanks;
@@ -345,7 +347,7 @@ struct Config {
 
   bool haveModularResultOutput() const
   {
-    return printTree || printFlowTree || printMap || printClu || printBinaryTree || printBinaryFlowTree;
+    return printTree || printFlowTree || printNewick || printMap || printClu || printBinaryTree || printBinaryFlowTree;
   }
 };