add option -savek

candidateset.cpp

@@ -150,7 +150,7 @@ bool CandidateSet::update(string tree, double score) {
 }
 
 vector<double> CandidateSet::getBestScores(int numBestScore) {
-	if (numBestScore == 0)
+	if (numBestScore == 0 || numBestScore > size())
 		numBestScore = size();
 	vector<double> res;
 	for (reverse_iterator rit = rbegin(); rit != rend() && numBestScore > 0; rit++, numBestScore--) {
@@ -221,6 +221,20 @@ string CandidateSet::getRandCandVecTree(){
 	return "";
 }
 
+string CandidateSet::getCandidateSetAsSyncTrees(int limit) {
+	string syncTrees;
+	int sz = min((int) size(), limit);
+
+	auto it = rbegin();
+	for(int i = 0; i < sz; ++it, ++i) {
+		syncTrees += to_string((int) -(it->first));
+		syncTrees += ' ';
+		syncTrees += (it->second).tree;
+		syncTrees += '#';
+	}
+	return syncTrees;
+}
+
 string CandidateSet::getSyncTrees(int forWorker) {
 	string syncTrees;
 	int sz = min((int) size(), popSize);

candidateset.h

@@ -171,6 +171,8 @@ class CandidateSet : public multimap<double, CandidateTree> {
 
     string getSyncTrees(int forWorker=5);
 
+    string getCandidateSetAsSyncTrees(int limit);
+
     void updateSingleSyncTree(string singleTree);
 
     void updateSyncTrees(string syncString);

iqtree.cpp

@@ -1573,7 +1573,7 @@ string IQTree::optimizeModelParameters(bool printInfo) {
 }
 
 void IQTree::printBestScores(int numBestScore) {
-	vector<double> bestScores = candidateTrees.getBestScores(candidateTrees.popSize);
+	vector<double> bestScores = candidateTrees.getBestScores(numBestScore);
 	for (vector<double>::iterator it = bestScores.begin(); it != bestScores.end(); it++)
 		mpiout << (params->maximum_parsimony ? -(*it) : (*it)) << " ";
 	mpiout << endl;
@@ -1787,6 +1787,26 @@ void IQTree::afterTreeSearch() {
     mpiout << "$$$$$END cleaning" << endl;
     mpiout << "CPU time used for tree search (MPI): " << getCPUTime() - checkpointTime << endl;
     MPI_Barrier(MPI_COMM_WORLD);
+
+    if (params->savek) {
+        // sync the candidate sets over processes
+        if (MPIHelper::getInstance().isMaster()) {
+            for(int worker = 1; worker < MPIHelper::getInstance().getNumProcesses(); ++worker) {
+                string treeStrings;
+                MPIHelper::getInstance().recvString(treeStrings, worker, MPIHelper::TREE_STRINGS);
+                candidateTrees.updateSyncTrees(treeStrings);
+            }
+        } else {
+            string sendTrees = candidateTrees.getCandidateSetAsSyncTrees(params->savek);
+            MPIHelper::getInstance().asyncSendString(
+                sendTrees, 
+                PROC_MASTER, 
+                MPIHelper::TREE_STRINGS, 
+                &reqs[0][1]
+            );
+        }
+        MPI_Barrier(MPI_COMM_WORLD);
+    }
 }
 
 double IQTree::doTreeSearch() {

phyloanalysis.cpp

@@ -1889,8 +1889,13 @@ void runTreeReconstruction(Params &params, string &original_model, IQTree &iqtre
 	if (iqtree.isSuperTree())
 			((PhyloSuperTree*) &iqtree)->mapTrees();
 	if (params.snni && params.min_iterations && isAllowedToPrint) {
-		mpiout << (params.maximum_parsimony ? "Scores" : "Log-likelihoods") << " of best " << params.popSize << " trees: " << endl;
-		iqtree.printBestScores(iqtree.candidateTrees.popSize);
+		if (params.savek == 0) {
+			mpiout << (params.maximum_parsimony ? "Scores" : "Log-likelihoods") << " of best " << params.popSize << " trees: " << endl;
+			iqtree.printBestScores(iqtree.candidateTrees.popSize);
+		} else {
+			mpiout << (params.maximum_parsimony ? "Scores" : "Log-likelihoods") << " of best " << params.savek << " trees: " << endl;
+			iqtree.printBestScores(params.savek);
+		}
 	}
 
 	/******** Performs final model parameters optimization ******************/

tools.cpp

@@ -805,6 +805,7 @@ void parseArg(int argc, char *argv[], Params &params) {
     params.do_first_rell = false;
     params.remove_dup_seq = false;
     params.test_mode = false;
+	params.savek = 0;
 
 #ifdef _OPENMP
     params.num_threads = 0;
@@ -2263,6 +2264,15 @@ void parseArg(int argc, char *argv[], Params &params) {
 				assert(params.popSize < params.numParsTrees);
 				continue;
 			}
+			if (strcmp(argv[cnt], "-savek") == 0) {
+				cnt++;
+				if (cnt >= argc)
+					throw "Use -savek <number_of_saving_trees>";
+				params.savek = convert_int(argv[cnt]);
+				params.numParsTrees = max(params.savek, params.numParsTrees);
+				assert(params.savek > 0);
+				continue;
+			}
 			if (strcmp(argv[cnt], "-beststart") == 0) {
 				params.bestStart = true;
 				cnt++;

tools.h

@@ -1628,6 +1628,11 @@ struct Params {
 	 */
 	bool test_mode;
 
+        /**
+         * number of search trees saved
+         */
+        int savek;
+
 	/*
 	 * Diep:
 	 * Use with -test_mode