From 777c23888e0a4c66dbe735f25474acd3b6e1acac Mon Sep 17 00:00:00 2001
From: Integer Limit <103940576+IntegerLimit@users.noreply.github.com>
Date: Thu, 4 Jul 2024 17:14:26 +1000
Subject: [PATCH 1/3] Support Exporting DNA/RNA Analysis to MrBayes Block Files
 (#260)

* Support Exporting DNA/RNA Analysis to MrBayes Block Files

* Fix Formatting Issues

* Move Functions to Supplementary, Fix +R Remapping
---
 .gitignore              |   1 +
 main/phyloanalysis.cpp  |  15 ++-
 main/phylotesting.h     |   7 ++
 model/modeldna.cpp      |   4 +
 model/modeldna.h        |   6 ++
 model/modelsubst.h      |   6 ++
 tree/phylosupertree.cpp |  83 +++++++++++++----
 tree/phylosupertree.h   |  19 +++-
 tree/phylotree.cpp      | 196 +++++++++++++++++++++++++++++++++++++++-
 tree/phylotree.h        |  33 +++++++
 utils/tools.cpp         |  19 +++-
 utils/tools.h           |   6 +-
 12 files changed, 368 insertions(+), 27 deletions(-)

diff --git a/.gitignore b/.gitignore
index fe0d324c6..926678ab6 100644
--- a/.gitignore
+++ b/.gitignore
@@ -14,6 +14,7 @@ tags
 test_scripts/iqtree_test_cmds.txt
 pllrepo/
 build/
+cmake-build-debug/
 /Default-clang
 .idea/
 .idea/codeStyleSettings.xml
diff --git a/main/phyloanalysis.cpp b/main/phyloanalysis.cpp
index e712707bc..48719d28a 100644
--- a/main/phyloanalysis.cpp
+++ b/main/phyloanalysis.cpp
@@ -1156,6 +1156,12 @@ void printOutfilesInfo(Params &params, IQTree &tree) {
     if (params.optimize_linked_gtr) {
         cout << "  GTRPMIX nex file:              " << params.out_prefix << ".GTRPMIX.nex" << endl;
     }
+
+    if (params.mr_bayes_output) {
+        cout << endl << "MrBayes block written to:" << endl;
+        cout << "  Base template file:            " << params.out_prefix << ".mr_bayes_scheme.nex" << endl;
+        cout << "  Template file with parameters: " << params.out_prefix << ".mr_bayes_model.nex" << endl;
+    }
     cout << endl;
 
 }
@@ -3100,8 +3106,13 @@ void runTreeReconstruction(Params &params, IQTree* &iqtree) {
 
     }
     if (iqtree->isSuperTree()) {
-        ((PhyloSuperTree*) iqtree)->computeBranchLengths();
-        ((PhyloSuperTree*) iqtree)->printBestPartitionParams((string(params.out_prefix) + ".best_model.nex").c_str());
+        auto superTree = (PhyloSuperTree*) iqtree;
+        superTree->computeBranchLengths();
+        superTree->printBestPartitionParams((string(params.out_prefix) + ".best_model.nex").c_str());
+    }
+    if (params.mr_bayes_output) {
+        iqtree->printMrBayesBlock((string(params.out_prefix) + ".mr_bayes_scheme.nex").c_str(), false);
+        iqtree->printMrBayesBlock((string(params.out_prefix) + ".mr_bayes_model.nex").c_str(), true);
     }
 
     cout << "BEST SCORE FOUND : " << iqtree->getCurScore() << endl;
diff --git a/main/phylotesting.h b/main/phylotesting.h
index 20d7d188b..286ca377d 100644
--- a/main/phylotesting.h
+++ b/main/phylotesting.h
@@ -390,5 +390,12 @@ string convertSeqTypeToSeqTypeName(SeqType seq_type);
 
 string detectSeqTypeName(string model_name);
 
+/**
+ * get string name from a SeqType object
+ * @param seq_type input sequence type
+ * @return name
+ */
+string getSeqTypeName(SeqType seq_type);
+
 
 #endif /* PHYLOTESTING_H_ */
diff --git a/model/modeldna.cpp b/model/modeldna.cpp
index c10317c2f..f2cbc0be0 100644
--- a/model/modeldna.cpp
+++ b/model/modeldna.cpp
@@ -564,3 +564,7 @@ void ModelDNA::setVariables(double *variables) {
 //                              j++;
 //                      }
 }
+
+string ModelDNA::getModelDNACode() {
+    return param_spec;
+}
diff --git a/model/modeldna.h b/model/modeldna.h
index 6694edde1..3f9bd7a61 100644
--- a/model/modeldna.h
+++ b/model/modeldna.h
@@ -114,6 +114,12 @@ class ModelDNA : public ModelMarkov
      */
     virtual void computeTipLikelihood(PML::StateType state, double *state_lk);
 
+    /**
+     * Get the Model DNA 'code', in form 'abcdef', used with ModelDNA model
+     * Returns empty string by default (this is not a dna specific model)
+     */
+    virtual string getModelDNACode();
+
 protected:
 
 	/**
diff --git a/model/modelsubst.h b/model/modelsubst.h
index 37835a8f2..07e7444cf 100644
--- a/model/modelsubst.h
+++ b/model/modelsubst.h
@@ -392,6 +392,12 @@ class ModelSubst: public Optimization, public CheckpointFactory
     */
     virtual ModelSubst *getMutationModel() { return this; }
 
+    /**
+     * Get the Model DNA 'code', in form 'abcdef', used with ModelDNA model
+     * Returns empty string by default (this is not a dna specific model)
+     */
+     virtual string getModelDNACode() { return ""; }
+
 	/*****************************************************
 		Checkpointing facility
 	*****************************************************/
diff --git a/tree/phylosupertree.cpp b/tree/phylosupertree.cpp
index 21fabf81d..9d0213051 100644
--- a/tree/phylosupertree.cpp
+++ b/tree/phylosupertree.cpp
@@ -23,6 +23,7 @@
 #include "main/phylotesting.h"
 #include "model/partitionmodel.h"
 #include "utils/MPIHelper.h"
+#include "utils/tools.h"
 
 PhyloSuperTree::PhyloSuperTree()
  : IQTree()
@@ -1519,29 +1520,37 @@ void PhyloSuperTree::writeBranches(ostream &out) {
     }
 }
 
-void PhyloSuperTree::printBestPartitionParams(const char *filename) {
+void printCharsets(SuperAlignment* saln, size_t size, ofstream &out)
+{
+    for (int part = 0; part < size; part++) {
+        string name = saln->partitions[part]->name;
+        replace(name.begin(), name.end(), '+', '_');
+        out << "  charset " << name << " = ";
+        if (!saln->partitions[part]->aln_file.empty()) out << saln->partitions[part]->aln_file << ": ";
+        /*if (saln->partitions[part]->seq_type == SEQ_CODON)
+            out << "CODON, ";*/
+        if (!saln->partitions[part]->sequence_type.empty())
+            out << saln->partitions[part]->sequence_type << ", ";
+        string pos = saln->partitions[part]->position_spec;
+        replace(pos.begin(), pos.end(), ',' , ' ');
+        out << pos << ";" << endl;
+    }
+}
+
+void PhyloSuperTree::printBestPartitionParams(const char *filename)
+{
     try {
         ofstream out;
         out.exceptions(ios::failbit | ios::badbit);
         out.open(filename);
         out << "#nexus" << endl
         << "begin sets;" << endl;
-        int part;
-        SuperAlignment *saln = (SuperAlignment*)aln;
-        for (part = 0; part < size(); part++) {
-            string name = saln->partitions[part]->name;
-            replace(name.begin(), name.end(), '+', '_');
-            out << "  charset " << name << " = ";
-            if (!saln->partitions[part]->aln_file.empty()) out << saln->partitions[part]->aln_file << ": ";
-            /*if (saln->partitions[part]->seq_type == SEQ_CODON)
-                out << "CODON, ";*/
-            out << saln->partitions[part]->sequence_type << ", ";
-            string pos = saln->partitions[part]->position_spec;
-            replace(pos.begin(), pos.end(), ',' , ' ');
-            out << pos << ";" << endl;
-        }
+        auto *saln = (SuperAlignment*)aln;
+
+        printCharsets(saln, size(), out);
+
         out << "  charpartition mymodels =" << endl;
-        for (part = 0; part < size(); part++) {
+        for (int part = 0; part < size(); part++) {
             string name = saln->partitions[part]->name;
             replace(name.begin(), name.end(), '+', '_');
             if (part > 0) out << "," << endl;
@@ -1555,3 +1564,45 @@ void PhyloSuperTree::printBestPartitionParams(const char *filename) {
         outError(ERR_WRITE_OUTPUT, filename);
     }
 }
+
+void PhyloSuperTree::printMrBayesFreeRateReplacement(RateHeterogeneity *rate, string &charset, ofstream &out) {
+    // Call PhyloTree's function inside the correct checkpoint struct
+    checkpoint->startStruct("PartitionModelPlen");
+    checkpoint->startStruct(charset);
+    PhyloTree::printMrBayesFreeRateReplacement(rate, charset, out);
+    checkpoint->endStruct();
+    checkpoint->endStruct();
+}
+
+void PhyloSuperTree::printMrBayesBlock(const char *filename, bool inclParams)
+{
+    ofstream out = startMrBayesBlock(filename);
+    auto *saln = (SuperAlignment*)aln;
+    printCharsets(saln, size(), out);
+    out << endl;
+
+    // Create Partition
+    size_type listSize = size();
+    out << "  partition iqtree = " << listSize << ": ";
+    for (int part = 0; part < listSize; part++) {
+        string name = saln->partitions[part]->name;
+        replace(name.begin(), name.end(), '+', '_');
+        out << name;
+        if (part != listSize - 1) out << ", ";
+        else out << ";" << endl;
+    }
+
+    // Set Partition for Use
+    out << "  set partition = iqtree;" << endl;
+
+    // Set Outgroup (if available)
+    if (!rooted) out << "  outgroup " << root->name << ";" << endl << endl;
+
+    // Partition-Specific Model Information
+    for (int part = 0; part < listSize; part++) {
+        // MrBayes Partitions are 1-indexed
+        printMrBayesModelText(this, at(part), out, convertIntToString(part + 1), saln->partitions[part]->name, inclParams);
+    }
+    out << "end;" << endl;
+    out.close();
+}
diff --git a/tree/phylosupertree.h b/tree/phylosupertree.h
index 60dbdcb9a..f2304fcab 100644
--- a/tree/phylosupertree.h
+++ b/tree/phylosupertree.h
@@ -475,12 +475,25 @@ class PhyloSuperTree : public IQTree, public vector<PhyloTree* >
      */
     void printBestPartitionParams(const char *filename);
 
-    
+    /**
+        print mr bayes block with partition information and best model parameters
+        @param filename output file name
+        @param inclParams whether to include iqtree params
+      */
+    virtual void printMrBayesBlock(const char *filename, bool inclParams);
+
+    /**
+     * Prints the replacement prior settings for +R or +R+I.
+     * @param rate the heterogeneity rate
+     * @param charset the (original) charset of the current partition. An empty string if not a partitioned tree
+     * @param out the ofstream to print to
+     */
+    virtual void printMrBayesFreeRateReplacement(RateHeterogeneity* rate, string &charset, ofstream &out);
+
     /** True when mixed codon with other data type */
     bool rescale_codon_brlen;
-    
-    int totalNNIs, evalNNIs;
 
+    int totalNNIs, evalNNIs;
 };
 
 #endif
diff --git a/tree/phylotree.cpp b/tree/phylotree.cpp
index b6bfcb595..8b1f6bce8 100644
--- a/tree/phylotree.cpp
+++ b/tree/phylotree.cpp
@@ -37,7 +37,7 @@
 #include "model/modelmixture.h"
 #include "phylonodemixlen.h"
 #include "phylotreemixlen.h"
-
+#include "main/phylotesting.h"
 
 const int LH_MIN_CONST = 1;
 
@@ -6246,3 +6246,197 @@ void PhyloTree::doNNI_simple(NNIMove &move) {
     }
 }
 
+/* Private Supplementary Logic Functions */
+
+/**
+ * Find Rate Value (propinv for +I, gamma for +G/+R, invgamma for +I+G/+I+R) <br>
+ * Default is Equal (+E)
+ */
+string getMrBayesRHASName(RateHeterogeneity* rate)
+{
+    bool hasGamma = rate->getGammaShape() != 0.0;
+    bool hasInvariable = rate->getPInvar() != 0.0;
+    if (hasGamma) {
+        if (hasInvariable)
+            return "invgamma";
+        return "gamma";
+    }
+    if (hasInvariable)
+        return "propinv";
+    return "equal";
+}
+
+void printMrBayesRHASPrior(PhyloTree* originalTree, RateHeterogeneity* rate, string &charset, ofstream &out)
+{
+    bool hasI = rate->getPInvar() > 0.0;
+
+    // Freerate (+R)
+    // Get replacement gamma shape
+    if (rate->isFreeRate()) {
+        originalTree ->printMrBayesFreeRateReplacement(rate, charset, out);
+    }
+
+    // Gamma Distribution (+G/+R)
+    // Dirichlet is not available here, use fixed
+    if (rate->getGammaShape() > 0.0)
+        out << " shapepr=fixed(" << rate->getGammaShape() << ")";
+
+    // Invariable Sites (+I)
+    // Dirichlet is not available here, use fixed
+    if (hasI)
+        out << " pinvarpr=fixed(" << rate->getPInvar() << ")";
+}
+
+void printMrBayesModelTextDNA(PhyloTree* originalTree, PhyloTree* tree, ofstream &out, string &partition, string &charset, bool inclParams)
+{
+    ModelSubst* model = tree->getModel();
+    bool equalFreq = model->freq_type == FREQ_EQUAL;
+    short nst = 6;
+    string rateMod = getMrBayesRHASName(tree->getRate());
+
+    // Find NST value (1 for JC/JC69/F81, 2 for K80/K2P/HKY/HKY85, 6 for SYM/GTR)
+    // NST is set to 6 by default (above), so check for JC/JC69/F81 and K80/K2P/HKY/HKY85
+    // If it returns a valid dna code, we can extract information from there (needed for user-defined models)
+    string dnaCode = model->getModelDNACode();
+    string modelStr = model->getName();
+    if (!dnaCode.empty()) {
+        // if all equal
+        if (dnaCode.find_first_not_of(dnaCode[0]) == string::npos)
+            nst = 1;
+        // if A-G=C-T and everything else is equal to first part of dna code
+        else if (dnaCode[1] == dnaCode[4] && dnaCode[0] == dnaCode[2] && dnaCode[0] == dnaCode[3] && dnaCode[0] == dnaCode[5])
+            nst = 2;
+    } else if (!modelStr.empty()) {
+        // Check the name of the model
+        if (strcmp(modelStr.c_str(), "JC") == 0 || strcmp(modelStr.c_str(), "JC69") == 0 || strcmp(modelStr.c_str(), "F81") == 0)
+            nst = 1;
+        else if (strcmp(modelStr.c_str(), "K80") == 0 || strcmp(modelStr.c_str(), "K2P") == 0 || strcmp(modelStr.c_str(), "HKY") == 0 || strcmp(modelStr.c_str(), "HKY85") == 0)
+            nst = 2;
+    }
+
+    // NucModel = 4By4: DNA Nucleotides (4 Options, A, C, G, T)
+    out << "  lset applyto=(" << partition << ") nucmodel=4by4 nst=" << nst << " rates=" << rateMod << ";" << endl;
+
+    // Priors (apart from Fixed Freq)
+    if (!inclParams) {
+        // If not include other params, simply set fixed frequency and return
+        if (!equalFreq) return;
+
+        out << "  prset applyto=(" << partition << ") statefreqpr=fixed(equal);" << endl;
+        return;
+    }
+
+    out << "  prset applyto=(" << partition << ")";
+
+    printMrBayesRHASPrior(originalTree, tree->getRate(), charset, out);
+
+    // Frequency of Nucleotides (+F)
+    if (equalFreq)
+        out << " statefreqpr=fixed(equal)";
+    else {
+        auto* freq = new double[model->num_states];
+        model->getStateFrequency(freq);
+        out << " statefreqpr=dirichlet(";
+        for (int i = 0; i < model->num_states; ++i) {
+            if (i != 0) out << ", ";
+            out << freq[i];
+        }
+        out << ")";
+
+        delete[] freq;
+    }
+
+    // Reversible Rate Matrix
+    int numRateEntries = model->getNumRateEntries();
+    auto* rateMat = new double[numRateEntries];
+    model->getRateMatrix(rateMat);
+
+    out << " revmatpr=dirichlet(";
+    for (int i = 0; i < numRateEntries; ++i) {
+        if (i != 0) out << ", ";
+        out << rateMat[i];
+    }
+
+    // Close revmatpr + prset
+    out << ");" << endl;
+
+    delete[] rateMat;
+}
+
+/* Public Supplementary Functions (Used by PhyloSuperTree) */
+ofstream startMrBayesBlock(const char *filename) {
+    ofstream out;
+    try {
+        out.exceptions(ios::failbit | ios::badbit);
+        out.open(filename);
+
+        // Write Warning
+        out << "#nexus" << endl << endl
+            <<"[This MrBayes Block Declaration provides the retrieved values from the IQTree Run.]" << endl
+            << "[Note that MrBayes does not support a large collection of models, so defaults of 'nst=6' for DNA and 'wag' for Protein will be used if a model that does not exist in MrBayes is selected.]" << endl
+            << "[Furthermore, the Model Parameter '+R' will be replaced by '+G'.]" << endl
+            << "[This should be used as a Template Only.]" << endl << endl;
+
+        // Begin File, Print Charsets
+        out << "begin mrbayes;" << endl;
+    } catch (ios::failure &) {
+        outError(ERR_WRITE_OUTPUT, filename);
+    }
+    return out;
+}
+
+void printMrBayesModelText(PhyloTree* originalTree, PhyloTree* tree, ofstream& out, string partition, string charset, bool inclParams)
+{
+    switch(tree->aln->seq_type) {
+        case SEQ_DNA:
+            printMrBayesModelTextDNA(originalTree, tree, out, partition, charset, inclParams);
+            break;
+        case SEQ_CODON:
+            break;
+        case SEQ_PROTEIN:
+            break;
+        case SEQ_BINARY:
+            break;
+        case SEQ_MORPH:
+            break;
+        default:
+            outWarning("MrBayes Block output not supported for sequence type " + getSeqTypeName(tree->aln->seq_type));
+            return;
+    }
+    out << endl; // Add extra endline after model text for readability
+}
+
+/* Protected Overrided Function */
+void PhyloTree::printMrBayesFreeRateReplacement(RateHeterogeneity *rate, string &charset, ofstream &out) {
+    bool hasInvar = rate->getPInvar() > 0.0;
+    string key = hasInvar ? "RateGammaInvar" : "RateGamma";
+
+    double gamma_shape = 0.0;
+    double p_invar = 0.0;
+
+    checkpoint->startStruct(key);
+
+    CKP_RESTORE(gamma_shape);
+    if (hasInvar)
+        CKP_RESTORE(p_invar);
+
+    checkpoint->endStruct();
+
+    // Sanity Check
+    if (gamma_shape > 0.0)
+        out << " shapepr=fixed(" << gamma_shape << ")";
+
+    if (p_invar > 0.0)
+        out << " pinvarpr=fixed(" << p_invar << ")";
+}
+
+/* Main Function */
+void PhyloTree::printMrBayesBlock(const char *filename, bool inclParams) {
+    ofstream out = startMrBayesBlock(filename);
+    // Set Outgroup (if available)
+    if (!rooted) out << "  outgroup " << root->name << ";" << endl << endl;
+
+    printMrBayesModelText(this, this, out, "all", "", inclParams);
+    out << "end;" << endl;
+    out.close();
+}
diff --git a/tree/phylotree.h b/tree/phylotree.h
index b8a19f986..f1f2029e0 100644
--- a/tree/phylotree.h
+++ b/tree/phylotree.h
@@ -2264,6 +2264,21 @@ class PhyloTree : public MTree, public Optimization, public CheckpointFactory {
      */
     const string& getDistanceFileWritten() const;
 
+    /**
+        print MrBayes block with partition information, and best model parameters (if inclParams is true)
+        @param filename output file name
+        @param inclParams whether to include iqtree params
+      */
+    virtual void printMrBayesBlock(const char *filename, bool inclParams);
+
+    /**
+     * Prints the replacement prior settings for +R or +R+I.
+     * @param rate the heterogeneity rate
+     * @param charset the (original) charset of the current partition. An empty string if not a partitioned tree
+     * @param out the ofstream to print to
+     */
+    virtual void printMrBayesFreeRateReplacement(RateHeterogeneity* rate, string &charset, ofstream &out);
+
     
 protected:
 
@@ -2511,4 +2526,22 @@ class PhyloTree : public MTree, public Optimization, public CheckpointFactory {
     void doneProgress();
 };
 
+/**
+ * print MrBayes formatted text, of model information (and parameters), to an ofstream.
+ * @param originalTree the original tree, of which is called to retrieve values in checkpoint file (to allow for overriding)
+ * @param tree tree to print info from
+ * @param out stream to print to
+ * @param partition the partition to 'apply' to (all to apply to all partitions)
+ * @param charset the (original) charset of the current partition. Used to retrieve values in checkpoint file. Input an empty string if this is not a partitioned tree.
+ * @param inclParams whether to include the parameters of the model
+ */
+void printMrBayesModelText(PhyloTree* originalTree, PhyloTree* tree, ofstream& out, string partition, string charset, bool inclParams);
+
+/**
+ * Start a MrBayes Block, as well as printing a warning.
+ * @param filename output file name
+ * @return ofstream for other functions to print to
+ */
+ofstream startMrBayesBlock(const char *filename);
+
 #endif
diff --git a/utils/tools.cpp b/utils/tools.cpp
index 10930e466..db21f7917 100644
--- a/utils/tools.cpp
+++ b/utils/tools.cpp
@@ -1524,6 +1524,7 @@ void parseArg(int argc, char *argv[], Params &params) {
     params.include_pre_mutations = false;
     params.mutation_file = "";
     params.site_starting_index = 0;
+    params.mr_bayes_output = false;
     
     // store original params
     for (cnt = 1; cnt < argc; cnt++) {
@@ -5662,7 +5663,10 @@ void parseArg(int argc, char *argv[], Params &params) {
                 params.make_consistent = true;
                 continue;
             }
-            
+            if (strcmp(argv[cnt], "-mrbayes") == 0) {
+                params.mr_bayes_output = true;
+                continue;
+            }
             if (argv[cnt][0] == '-') {
                 string err = "Invalid \"";
                 err += argv[cnt];
@@ -5800,6 +5804,10 @@ void parseArg(int argc, char *argv[], Params &params) {
     if (params.model_name.find("LINK") != string::npos || params.model_name.find("MERGE") != string::npos)
         if (params.partition_merge == MERGE_NONE)
             params.partition_merge = MERGE_RCLUSTERF;
+
+    // Set MrBayes Block Output if -mset mrbayes
+    if (params.model_set == "mrbayes")
+        params.mr_bayes_output = true;
     
     if (params.alisim_active && !params.aln_file && !params.user_file && !params.partition_file && params.tree_gen == NONE)
         outError("A tree filepath is a mandatory input to execute AliSim when neither Inference mode nor Random mode (generating a random tree) is inactive. Use -t <TREE_FILEPATH> ; or Activate the inference mode by -s <ALIGNMENT_FILE> ; or Activate Random mode by -t RANDOM{<MODEL>,<NUM_TAXA>} where <MODEL> is yh, u, cat, bal, bd{<birth_rate>,<death_rate>} stands for Yule-Harding, Uniform, Caterpillar, Balanced, Birth-Death model respectively.");
@@ -6014,7 +6022,7 @@ void usage_iqtree(char* argv[], bool full_command) {
 //            << "  -pll                 Use phylogenetic likelihood library (PLL) (default: off)" << endl
     << "  --ninit NUM          Number of initial parsimony trees (default: 100)" << endl
     << "  --ntop NUM           Number of top initial trees (default: 20)" << endl
-    << "  --nbest NUM          Number of best trees retained during search (defaut: 5)" << endl
+    << "  --nbest NUM          Number of best trees retained during search (default: 5)" << endl
     << "  -n NUM               Fix number of iterations to stop (default: OFF)" << endl
     << "  --nstop NUM          Number of unsuccessful iterations to stop (default: 100)" << endl
     << "  --perturb NUM        Perturbation strength for randomized NNI (default: 0.5)" << endl
@@ -6071,6 +6079,8 @@ void usage_iqtree(char* argv[], bool full_command) {
     << "  -m ...+LMSS          Additionally test strand-symmetric models" << endl
     << "  --mset STRING        Restrict search to models supported by other programs" << endl
     << "                       (raxml, phyml, mrbayes, beast1 or beast2)" << endl
+    << "                       If 'mrbayes' is selected, will output a MrBayes" << endl
+    << "                       Block File if Data Type is supported." << endl
     << "  --mset STR,...       Comma-separated model list (e.g. -mset WAG,LG,JTT)" << endl
     << "  --msub STRING        Amino-acid model source" << endl
     << "                       (nuclear, mitochondrial, chloroplast or viral)" << endl
@@ -6304,8 +6314,9 @@ void usage_iqtree(char* argv[], bool full_command) {
     //			<< "  -d <outfile>         Calculate the distance matrix inferred from tree" << endl
     //			<< "  -stats <outfile>     Output some statistics about branch lengths" << endl
     //			<< "  -comp <treefile>     Compare tree with each in the input trees" << endl;
-
-
+        << "  -mrbayes             Outputs a Mr Bayes block file, to use as a template for future analysis" << endl
+        << "                       Will not output if incompatible data types are used" << endl
+        << "                       (MrBayes only supports DNA, Codon, Protein, Binary and Morphological data)" << endl
         << endl;
 
     if (full_command) {
diff --git a/utils/tools.h b/utils/tools.h
index b6f24978a..75e538027 100644
--- a/utils/tools.h
+++ b/utils/tools.h
@@ -2781,6 +2781,11 @@ class Params {
     *  site starting index (for predefined mutations in AliSim)
     */
     int site_starting_index;
+
+    /**
+    * Whether to output a MrBayes Block File
+    */
+    bool mr_bayes_output;
 };
 
 /**
@@ -3728,5 +3733,4 @@ double frob_norm (double m[], int n, double scale=1.0);
 */
 string getOutputNameWithExt(const InputType& format, const string& output_filepath);
 
-
 #endif

From 02df654ea9c47dedb6c27ac28656464dc46d01cc Mon Sep 17 00:00:00 2001
From: Integer Limit <103940576+IntegerLimit@users.noreply.github.com>
Date: Fri, 5 Jul 2024 21:39:17 +1000
Subject: [PATCH 2/3] Support Exporting Protein, Binary and Morphological Data
 to MrBayes (#264)

* Protein Model

* Morphological Models Support

* Cleanup

Move Model Specific Functions to each model class
Move other functions from phylotree and phylosupertree to phyloanalysis

* Cleanup Imports

* Binary Model Support

* Misc Cleanup

Misc Cleanup

* Output Files Readability, Default Warning & Help Message

* Fix Edge Case: Importing Values < 0.01 into MrBayes

* Fix Edge Case: Extra Characters in Charset

* Fix +G+I or +R Inputs

* Fix Issues with Binary Model

* Fix Issues with Morphology Model
---
 main/phyloanalysis.cpp    |  82 +++++++++++++++-
 model/modelbin.cpp        |  53 +++++++++++
 model/modelbin.h          |  12 +++
 model/modeldna.cpp        |  98 ++++++++++++++++++-
 model/modeldna.h          |  12 ++-
 model/modelmorphology.cpp |  51 ++++++++++
 model/modelmorphology.h   |  12 +++
 model/modelprotein.cpp    |  89 +++++++++++++++++
 model/modelprotein.h      |  12 +++
 model/modelsubst.cpp      |  41 ++++++++
 model/modelsubst.h        |  22 ++++-
 tree/phylosupertree.cpp   |  84 ++++------------
 tree/phylosupertree.h     |  19 +---
 tree/phylotree.cpp        | 196 +-------------------------------------
 tree/phylotree.h          |  33 -------
 utils/tools.cpp           |  34 +++++++
 utils/tools.h             |  17 ++++
 17 files changed, 546 insertions(+), 321 deletions(-)

diff --git a/main/phyloanalysis.cpp b/main/phyloanalysis.cpp
index 48719d28a..7dd7475d5 100644
--- a/main/phyloanalysis.cpp
+++ b/main/phyloanalysis.cpp
@@ -2602,6 +2602,82 @@ void printTrees(vector<string> trees, Params &params, string suffix) {
     treesOut.close();
 }
 
+void printMrBayesBlockFile(const char* filename, IQTree* &iqtree, bool inclParams) {
+    ofstream out;
+    try {
+        out.exceptions(ios::failbit | ios::badbit);
+        out.open(filename);
+
+        // Write Warning
+        out << "#nexus" << endl << endl
+            <<"[This MrBayes Block Declaration provides the retrieved values from the IQTree Run.]" << endl
+            << "[Note that MrBayes does not support a large collection of models, so defaults of 'nst=6' for DNA and 'wag' for Protein will be used if a model that does not exist in MrBayes is selected.]" << endl
+            << "[Furthermore, the Model Parameter '+R' will be replaced by '+G'.]" << endl
+            << "[This should be used as a Template Only.]" << endl << endl;
+
+        // Begin File, Print Charsets
+        out << "begin mrbayes;" << endl;
+    } catch (ios::failure &) {
+        outError(ERR_WRITE_OUTPUT, filename);
+    }
+
+    if (!iqtree->isSuperTree()) {
+        // Set Outgroup (if available)
+        if (!iqtree->rooted) out << "  outgroup " << iqtree->root->name << ";" << endl << endl;
+
+        out << "  [Using Model '" << iqtree->getModelName() << "']" << endl;
+        iqtree->getModel()->printMrBayesModelText(iqtree->getRate(), out, "all", "", false, inclParams);
+
+        out << endl << "end;" << endl;
+        out.close();
+        return;
+    }
+
+    auto superTree = (PhyloSuperTree*) iqtree;
+    auto saln = (SuperAlignment*) superTree->aln;
+    auto size = superTree->size();
+
+    for (int part = 0; part < size; part++) {
+        string name = saln->partitions[part]->name;
+        replace(name.begin(), name.end(), '+', '_');
+        out << "  charset " << name << " = ";
+
+        string pos = saln->partitions[part]->position_spec;
+        replace(pos.begin(), pos.end(), ',' , ' ');
+        out << pos << ";" << endl;
+    }
+
+    // Create Partition
+    out << endl << "  partition iqtree = " << size << ": ";
+    for (int part = 0; part < size; part++) {
+        if (part != 0) out << ", ";
+
+        string name = saln->partitions[part]->name;
+        replace(name.begin(), name.end(), '+', '_');
+        out << name;
+    }
+    out << ";" << endl;
+
+    // Set Partition for Use
+    out << "  set partition = iqtree;" << endl;
+
+    // Set Outgroup (if available)
+    if (!superTree->rooted) out << "  outgroup " << superTree->root->name << ";" << endl << endl;
+
+    // Partition-Specific Model Information
+    for (int part = 0; part < size; part++) {
+        PhyloTree* currentTree = superTree->at(part);
+
+        // MrBayes Partitions are 1-indexed
+        out << "  [Partition No. " << convertIntToString(part + 1) << ", Using Model '" << currentTree->getModelName() << "']" << endl;
+        currentTree->getModel()->printMrBayesModelText(currentTree->getRate(), out,
+                convertIntToString(part + 1), saln->partitions[part]->name, true, inclParams);
+        out << endl;
+    }
+    out << "end;" << endl;
+    out.close();
+}
+
 /************************************************************
  *  MAIN TREE RECONSTRUCTION
  ***********************************************************/
@@ -3111,8 +3187,10 @@ void runTreeReconstruction(Params &params, IQTree* &iqtree) {
         superTree->printBestPartitionParams((string(params.out_prefix) + ".best_model.nex").c_str());
     }
     if (params.mr_bayes_output) {
-        iqtree->printMrBayesBlock((string(params.out_prefix) + ".mr_bayes_scheme.nex").c_str(), false);
-        iqtree->printMrBayesBlock((string(params.out_prefix) + ".mr_bayes_model.nex").c_str(), true);
+        cout << endl << "Writing MrBayes Block Files..." << endl;
+        printMrBayesBlockFile((string(params.out_prefix) + ".mr_bayes_scheme.nex").c_str(), iqtree, false);
+        printMrBayesBlockFile((string(params.out_prefix) + ".mr_bayes_model.nex").c_str(), iqtree, true);
+        cout << endl;
     }
 
     cout << "BEST SCORE FOUND : " << iqtree->getCurScore() << endl;
diff --git a/model/modelbin.cpp b/model/modelbin.cpp
index e8291e00a..262cbea80 100644
--- a/model/modelbin.cpp
+++ b/model/modelbin.cpp
@@ -76,3 +76,56 @@ string ModelBIN::getNameParams(bool show_fixed_params) {
     }
     return retname.str();
 }
+
+void ModelBIN::printMrBayesModelText(RateHeterogeneity* rate, ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams) {
+    // MrBayes does not support Invariable Modifier for Binary data
+    if (rate->isFreeRate() || rate->getPInvar() > 0.0) {
+        warnLogStream("MrBayes does not support Invariable Sites with Binary Data! +I has been ignored!", out);
+    }
+
+    // Lset Parameters
+    out << "  lset applyto=(" << partition << ") rates=";
+
+    // Free Rate should be substituted by +G (+I not supported)
+    bool hasGamma = rate->getGammaShape() != 0.0 || rate->isFreeRate();
+    if (hasGamma) {
+        // Rate Categories + Gamma
+        out << "gamma ngammacat=" << rate->getNRate();
+    } else
+        out << "equal";
+
+    out << ";" << endl;
+
+    if (!inclParams) {
+        if (freq_type == FREQ_EQUAL) out << "  prset applyto=(" << partition << ") statefreqpr=fixed(equal);" << endl;
+        return;
+    }
+
+    // Prset Parameters
+    out << "  prset applyto=(" << partition << ")";
+
+    // Freerate (+R)
+    // Get replacement Gamma Shape
+    if (rate->isFreeRate()) {
+        printMrBayesFreeRateReplacement(rate, charset, out, false);
+    }
+
+    // Gamma Distribution (+G/+R)
+    // Dirichlet is not available here, use fixed
+    if (rate->getGammaShape() > 0.0)
+        out << " shapepr=fixed(" << minValueCheckMrBayes(rate->getGammaShape()) << ")";
+
+    // State Frequencies
+    if (freq_type == FREQ_EQUAL)
+        out << " statefreqpr=fixed(equal)";
+    else {
+        out << " statefreqpr=dirichlet(";
+        for (int i = 0; i < num_states; ++i) {
+            if (i != 0) out << ", ";
+            out << minValueCheckMrBayes(state_freq[i]);
+        }
+        out << ")";
+    }
+
+    out << ";" << endl;
+}
diff --git a/model/modelbin.h b/model/modelbin.h
index 9c3ab6cf2..b2c257d9e 100644
--- a/model/modelbin.h
+++ b/model/modelbin.h
@@ -55,6 +55,18 @@ class ModelBIN : public ModelMarkov
     */
     virtual void startCheckpoint();
 
+    /**
+     * Print the model information in a format that can be accepted by MrBayes, using lset and prset.<br>
+     * By default, it simply prints a warning to the log and to the stream, stating that this model is not supported by MrBayes.
+     * @param rate the rate information
+     * @param out the ofstream to print the result to
+     * @param partition the partition to apply lset and prset to
+     * @param charset the current partition's charset. Useful for getting information from the checkpoint file
+     * @param isSuperTree whether the tree is a super tree. Useful for retrieving information from the checkpoint file, which has different locations for PhyloTree and PhyloSuperTree
+     * @param inclParams whether to include IQTree optimized parameters for the model
+     */
+    virtual void printMrBayesModelText(RateHeterogeneity* rate, ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams);
+
 };
 
 #endif
diff --git a/model/modeldna.cpp b/model/modeldna.cpp
index f2cbc0be0..4a36e16be 100644
--- a/model/modeldna.cpp
+++ b/model/modeldna.cpp
@@ -565,6 +565,100 @@ void ModelDNA::setVariables(double *variables) {
 //                      }
 }
 
-string ModelDNA::getModelDNACode() {
-    return param_spec;
+void ModelDNA::printMrBayesModelText(RateHeterogeneity* rate, ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams) {
+    bool equalFreq = freq_type == FREQ_EQUAL;
+    short nst = 6;
+
+    // Find NST value (1 for JC/JC69/F81, 2 for K80/K2P/HKY/HKY85, 6 for SYM/GTR)
+    // NST is set to 6 by default (above), so check for JC/JC69/F81 and K80/K2P/HKY/HKY85
+    // If it returns a valid dna code, we can extract information from there (needed for user-defined models)
+    if (!param_spec.empty()) {
+        // if all equal
+        if (param_spec.find_first_not_of(param_spec[0]) == string::npos) {
+            nst = 1;
+            // if A-G=C-T and everything else is equal to first part of dna code
+        } else if (param_spec[1] == param_spec[4] &&
+                param_spec[0] == param_spec[2] && param_spec[0] == param_spec[3] && param_spec[0] == param_spec[5]) {
+            nst = 2;
+        }
+    } else if (!name.empty()) {
+        // Check the name of the model
+        if (strcmp(name.c_str(), "JC") == 0 || strcmp(name.c_str(), "JC69") == 0 || strcmp(name.c_str(), "F81") == 0) {
+            nst = 1;
+        } else if (strcmp(name.c_str(), "K80") == 0 || strcmp(name.c_str(), "K2P") == 0 || strcmp(name.c_str(), "HKY") == 0 || strcmp(name.c_str(), "HKY85") == 0) {
+            nst = 2;
+        }
+    }
+
+    // NucModel = 4By4: DNA Nucleotides (4 Options, A, C, G, T)
+    out << "  lset applyto=(" << partition << ") nucmodel=4by4 nst=" << nst << " rates=";
+
+    // RHAS Specification
+    // Free Rate should be substituted by +G+I
+    bool hasGamma = rate->getGammaShape() != 0.0 || rate->isFreeRate();
+    bool hasInvariable = rate->getPInvar() != 0.0 || rate->isFreeRate();
+    if (hasGamma) {
+        if (hasInvariable)
+            out << "invgamma";
+        else
+            out << "gamma";
+    } else if (hasInvariable)
+        out << "propinv";
+    else
+        out << "equal";
+
+    // Rate Categories
+    if (hasGamma)
+        out << " ngammacat=" << rate->getNRate();
+
+    out << ";" << endl;
+
+    // Priors (apart from Fixed Freq)
+    if (!inclParams) {
+        // If not include other params, simply set fixed frequency and return
+        if (!equalFreq) return;
+
+        out << "  prset applyto=(" << partition << ") statefreqpr=fixed(equal);" << endl;
+        return;
+    }
+
+    out << "  prset applyto=(" << partition << ")";
+
+    // Freerate (+R)
+    // Get replacement Gamma Shape + Invariable Sites
+    if (rate->isFreeRate()) {
+        printMrBayesFreeRateReplacement(isSuperTree, charset, out);
+    }
+
+    // Gamma Distribution (+G/+R)
+    // Dirichlet is not available here, use fixed
+    if (rate->getGammaShape() > 0.0)
+        out << " shapepr=fixed(" << minValueCheckMrBayes(rate->getGammaShape()) << ")";
+
+    // Invariable Sites (+I)
+    // Dirichlet is not available here, use fixed
+    if (rate->getPInvar() > 0.0)
+        out << " pinvarpr=fixed(" << minValueCheckMrBayes(rate->getPInvar()) << ")";
+
+    // Frequency of Nucleotides (+F)
+    if (equalFreq)
+        out << " statefreqpr=fixed(equal)";
+    else {
+        out << " statefreqpr=dirichlet(";
+        for (int i = 0; i < num_states; ++i) {
+            if (i != 0) out << ", ";
+            out << minValueCheckMrBayes(state_freq[i]);
+        }
+        out << ")";
+    }
+
+    // Reversible Rate Matrix
+    out << " revmatpr=dirichlet(";
+    for (int i = 0; i < getNumRateEntries(); ++i) {
+        if (i != 0) out << ", ";
+        out << minValueCheckMrBayes(rates[i]);
+    }
+
+    // Close revmatpr + prset
+    out << ");" << endl;
 }
diff --git a/model/modeldna.h b/model/modeldna.h
index 3f9bd7a61..0c252c178 100644
--- a/model/modeldna.h
+++ b/model/modeldna.h
@@ -115,10 +115,16 @@ class ModelDNA : public ModelMarkov
     virtual void computeTipLikelihood(PML::StateType state, double *state_lk);
 
     /**
-     * Get the Model DNA 'code', in form 'abcdef', used with ModelDNA model
-     * Returns empty string by default (this is not a dna specific model)
+     * Print the model information in a format that can be accepted by MrBayes, using lset and prset.<br>
+     * By default, it simply prints a warning to the log and to the stream, stating that this model is not supported by MrBayes.
+     * @param rate the rate information
+     * @param out the ofstream to print the result to
+     * @param partition the partition to apply lset and prset to
+     * @param charset the current partition's charset. Useful for getting information from the checkpoint file
+     * @param isSuperTree whether the tree is a super tree. Useful for retrieving information from the checkpoint file, which has different locations for PhyloTree and PhyloSuperTree
+     * @param inclParams whether to include IQTree optimized parameters for the model
      */
-    virtual string getModelDNACode();
+    virtual void printMrBayesModelText(RateHeterogeneity* rate, ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams);
 
 protected:
 
diff --git a/model/modelmorphology.cpp b/model/modelmorphology.cpp
index a4df995be..ddda51874 100644
--- a/model/modelmorphology.cpp
+++ b/model/modelmorphology.cpp
@@ -150,3 +150,54 @@ void ModelMorphology::writeInfo(ostream &out) {
         out << endl;
     }
 }
+
+void ModelMorphology::printMrBayesModelText(RateHeterogeneity* rate, ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams) {
+    warnLogStream("MrBayes only supports Morphological Data with states from {0-9}!", out);
+    warnLogStream("Morphological Data with states {A-Z} may cause errors!", out);
+    warnLogStream("Use Morphological Models in MrBayes with Caution!", out);
+
+    // MrBayes does not support Invariable Modifier for Morph data
+    if (rate->isFreeRate() || rate->getPInvar() > 0.0) {
+        warnLogStream("MrBayes does not support Invariable Sites with Morphological Data! +I has been ignored!", out);
+    }
+    // MrBayes does not support State Frequency for Morph data
+    if (freq_type != FREQ_EQUAL) {
+        warnLogStream("MrBayes does not support non-equal frequencies for Morphological Data! Frequencies are left as the default! (All equal)", out);
+    }
+
+    // Lset Parameters
+    out << "  lset applyto=(" << partition << ") rates=";
+
+    // Free Rate should be substituted by +G (+I not supported)
+    bool hasGamma = rate->getGammaShape() != 0.0 || rate->isFreeRate();
+    if (hasGamma) {
+        // Rate Categories + Gamma
+        out << "gamma ngammacat=" << rate->getNRate();
+    } else
+        out << "equal";
+
+    out << ";" << endl;
+
+    // ctype (ordered or not)
+    if (strcmp(name.c_str(), "ORDERED") == 0)
+        out << "  ctype ordered: " << charset << ";" << endl;
+
+    // Since morphological doesn't have state frequencies, if there are no parameters to prset, return
+    if (!inclParams || !rate->isFreeRate() && rate->getGammaShape() <= 0.0) return;
+
+    // Prset Parameters
+    out << "  prset applyto=(" << partition << ")";
+
+    // Freerate (+R)
+    // Get replacement Gamma Shape
+    if (rate->isFreeRate()) {
+        printMrBayesFreeRateReplacement(rate, charset, out, false);
+    }
+
+    // Gamma Distribution (+G/+R)
+    // Dirichlet is not available here, use fixed
+    if (rate->getGammaShape() > 0.0)
+        out << " shapepr=fixed(" << minValueCheckMrBayes(rate->getGammaShape()) << ")";
+
+    out << ";" << endl;
+}
diff --git a/model/modelmorphology.h b/model/modelmorphology.h
index e9d271dff..f598103bd 100644
--- a/model/modelmorphology.h
+++ b/model/modelmorphology.h
@@ -78,6 +78,18 @@ class ModelMorphology: public ModelMarkov {
     virtual void readRates(istream &in) noexcept(false);
 
     virtual ~ModelMorphology();
+
+    /**
+     * Print the model information in a format that can be accepted by MrBayes, using lset and prset.<br>
+     * By default, it simply prints a warning to the log and to the stream, stating that this model is not supported by MrBayes.
+     * @param rate the rate information
+     * @param out the ofstream to print the result to
+     * @param partition the partition to apply lset and prset to
+     * @param charset the current partition's charset. Useful for getting information from the checkpoint file
+     * @param isSuperTree whether the tree is a super tree. Useful for retrieving information from the checkpoint file, which has different locations for PhyloTree and PhyloSuperTree
+     * @param inclParams whether to include IQTree optimized parameters for the model
+     */
+    virtual void printMrBayesModelText(RateHeterogeneity* rate, ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams);
 };
 
 #endif /* MODELMORPHOLOGY_H_ */
diff --git a/model/modelprotein.cpp b/model/modelprotein.cpp
index 4e5d7196e..24f8a99ff 100644
--- a/model/modelprotein.cpp
+++ b/model/modelprotein.cpp
@@ -1162,3 +1162,92 @@ void ModelProtein::computeTipLikelihood(PML::StateType state, double *state_lk)
     state_lk[ambi_aa[cstate*2+1]] = 1.0;
 }
 
+void ModelProtein::printMrBayesModelText(RateHeterogeneity* rate, ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams) {
+    // Lset Parameters
+    out << "  lset applyto=(" << partition << ") nucmodel=protein rates=";
+
+    // RHAS Specification
+    // Free Rate should be substituted by +G+I
+    bool hasGamma = rate->getGammaShape() != 0.0 || rate->isFreeRate();
+    bool hasInvariable = rate->getPInvar() != 0.0 || rate->isFreeRate();
+    if (hasGamma) {
+        if (hasInvariable)
+            out << "invgamma";
+        else
+            out << "gamma";
+    } else if (hasInvariable)
+        out << "propinv";
+    else
+        out << "equal";
+
+    // Rate Categories
+    if (hasGamma)
+        out << " ngammacat=" << rate->getNRate();
+
+    out << ";" << endl;
+
+    out << "  prset applyto=(" << partition << ")";
+
+    // Get MrBayes Model
+    auto aaModelMap = getIqTreeToMrBayesAAModels();
+    auto iter = aaModelMap.find(name);
+    string mappedModel = "gtr";
+
+    // If model is in map, set mappedModel to the value
+    if (iter != aaModelMap.end())
+        mappedModel = iter->second;
+
+    out << " aamodelpr=fixed(" << mappedModel << ")";
+
+    // GTR Customization
+    if (strcmp(mappedModel.c_str(), "gtr") == 0) {
+        // add rate matrix and state frequencies (mandatory for setting gtr values)
+        out << " aarevmatpr=";
+
+        // if matrix is GTR20, use dirichlet, else use fixed
+        if (strcmp(name.c_str(), "GTR20") == 0)
+            out << "dirichlet(";
+        else
+            out << "fixed(";
+
+        for (int i = 0; i < getNumRateEntries(); ++i) {
+            if (i != 0) out << ", ";
+            out << minValueCheckMrBayes(rates[i]);
+        }
+        out << ")";
+
+        // Frequency type is never equal to FREQ_EQUAL, even with Poisson
+        // Frequency is also auto-set if we use a model defined by MrBayes
+        out << " statefreqpr=dirichlet(";
+        for (int i = 0; i < num_states; ++i) {
+            if (i != 0) out << ", ";
+            out << minValueCheckMrBayes(state_freq[i]);
+        }
+        out << ")";
+    }
+
+    // if not to include the parameters (for Protein, simply +I, +G, +R)
+    if (!inclParams) {
+        out << ";";
+        return;
+    }
+
+    // Freerate (+R)
+    // Get replacement Gamma Shape + Invariable Sites
+    if (rate->isFreeRate()) {
+        printMrBayesFreeRateReplacement(isSuperTree, charset, out);
+    }
+
+    // Gamma Distribution (+G/+R)
+    // Dirichlet is not available here, use fixed
+    if (rate->getGammaShape() > 0.0)
+        out << " shapepr=fixed(" << minValueCheckMrBayes(rate->getGammaShape()) << ")";
+
+    // Invariable Sites (+I)
+    // Dirichlet is not available here, use fixed
+    if (rate->getPInvar() > 0.0)
+        out << " pinvarpr=fixed(" << minValueCheckMrBayes(rate->getPInvar()) << ")";
+
+    out << ";" << endl;
+}
+
diff --git a/model/modelprotein.h b/model/modelprotein.h
index e6680eb4a..6dd1f69e6 100644
--- a/model/modelprotein.h
+++ b/model/modelprotein.h
@@ -79,6 +79,18 @@ class ModelProtein : public ModelMarkov
      */
     virtual void computeTipLikelihood(PML::StateType state, double *state_lk);
 
+    /**
+     * Print the model information in a format that can be accepted by MrBayes, using lset and prset.<br>
+     * By default, it simply prints a warning to the log and to the stream, stating that this model is not supported by MrBayes.
+     * @param rate the rate information
+     * @param out the ofstream to print the result to
+     * @param partition the partition to apply lset and prset to
+     * @param charset the current partition's charset. Useful for getting information from the checkpoint file
+     * @param isSuperTree whether the tree is a super tree. Useful for retrieving information from the checkpoint file, which has different locations for PhyloTree and PhyloSuperTree
+     * @param inclParams whether to include IQTree optimized parameters for the model
+     */
+    virtual void printMrBayesModelText(RateHeterogeneity* rate, ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams);
+
 private:
     ModelsBlock *models_block;
 
diff --git a/model/modelsubst.cpp b/model/modelsubst.cpp
index 77363edf5..5c5675bea 100644
--- a/model/modelsubst.cpp
+++ b/model/modelsubst.cpp
@@ -210,6 +210,47 @@ double *ModelSubst::newTransMatrix() {
 	return new double[num_states * num_states];
 }
 
+void ModelSubst::printMrBayesFreeRateReplacement(bool isSuperTree, string &charset, ofstream &out, bool inclInvariable)
+{
+    double gamma_shape = 0.0;
+    double p_invar = 0.0;
+
+    if (isSuperTree) {
+        checkpoint->startStruct("PartitionModelPlen");
+        checkpoint->startStruct(charset);
+    }
+    checkpoint->startStruct("RateGammaInvar");
+
+    // Safety Check
+    if (!checkpoint->hasKey("gamma_shape") || !checkpoint->hasKey("p_invar")) {
+        outWarning("Could not find replacement for Freerate Distribution in MrBayes Block!");
+        checkpoint->endStruct();
+        return;
+    }
+
+    outWarning("FreeRate Distributions (+R) are not available in MrBayes. It will be replaced by +G+I.");
+
+    // Always replace +R with +G+I
+    CKP_RESTORE(gamma_shape);
+    if (inclInvariable)
+        CKP_RESTORE(p_invar);
+
+    checkpoint->endStruct();
+    if (isSuperTree) {
+        checkpoint->endStruct();
+        checkpoint->endStruct();
+    }
+
+    out << " shapepr=fixed(" << minValueCheckMrBayes(gamma_shape) << ")";
+    if (inclInvariable)
+        out << " pinvarpr=fixed(" << minValueCheckMrBayes(p_invar) << ")";
+}
+
+void ModelSubst::printMrBayesModelText(RateHeterogeneity* rate, ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams) {
+    out << "  [MrBayes Block Output is not supported by this model!]" << endl;
+    outWarning("MrBayes Block Output is not supported by this model of name '" + name + "'!");
+}
+
 ModelSubst::~ModelSubst()
 {
     // mem space pointing to target model and thus avoid double free here
diff --git a/model/modelsubst.h b/model/modelsubst.h
index 07e7444cf..370737418 100644
--- a/model/modelsubst.h
+++ b/model/modelsubst.h
@@ -17,6 +17,7 @@
 #include "utils/optimization.h"
 #include "utils/checkpoint.h"
 #include "phylo-yaml/statespace.h"
+#include "model/rateheterogeneity.h"
 
 using namespace std;
 
@@ -393,10 +394,25 @@ class ModelSubst: public Optimization, public CheckpointFactory
     virtual ModelSubst *getMutationModel() { return this; }
 
     /**
-     * Get the Model DNA 'code', in form 'abcdef', used with ModelDNA model
-     * Returns empty string by default (this is not a dna specific model)
+     * Prints the replacement prior settings, in MrBayes, for +R or +R+I.
+     * @param isSuperTree whether the tree is a super tree
+     * @param charset the (original) charset of the current partition. An empty string if not a partitioned tree
+     * @param out the ofstream to print to
+     * @param inclInvariable whether to include invariable sites. defaults to true
      */
-     virtual string getModelDNACode() { return ""; }
+    void printMrBayesFreeRateReplacement(bool isSuperTree, string &charset, ofstream &out, bool inclInvariable = true);
+
+    /**
+     * Print the model information in a format that can be accepted by MrBayes, using lset and prset.<br>
+     * By default, it simply prints a warning to the log and to the stream, stating that this model is not supported by MrBayes.
+     * @param rate the rate information
+     * @param out the ofstream to print the result to
+     * @param partition the partition to apply lset and prset to
+     * @param charset the current partition's charset. Useful for getting information from the checkpoint file
+     * @param isSuperTree whether the tree is a super tree. Useful for retrieving information from the checkpoint file, which has different locations for PhyloTree and PhyloSuperTree
+     * @param inclParams whether to include IQTree optimized parameters for the model
+     */
+    virtual void printMrBayesModelText(RateHeterogeneity* rate, ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams);
 
 	/*****************************************************
 		Checkpointing facility
diff --git a/tree/phylosupertree.cpp b/tree/phylosupertree.cpp
index 9d0213051..ff3884c17 100644
--- a/tree/phylosupertree.cpp
+++ b/tree/phylosupertree.cpp
@@ -23,7 +23,6 @@
 #include "main/phylotesting.h"
 #include "model/partitionmodel.h"
 #include "utils/MPIHelper.h"
-#include "utils/tools.h"
 
 PhyloSuperTree::PhyloSuperTree()
  : IQTree()
@@ -1520,37 +1519,30 @@ void PhyloSuperTree::writeBranches(ostream &out) {
     }
 }
 
-void printCharsets(SuperAlignment* saln, size_t size, ofstream &out)
-{
-    for (int part = 0; part < size; part++) {
-        string name = saln->partitions[part]->name;
-        replace(name.begin(), name.end(), '+', '_');
-        out << "  charset " << name << " = ";
-        if (!saln->partitions[part]->aln_file.empty()) out << saln->partitions[part]->aln_file << ": ";
-        /*if (saln->partitions[part]->seq_type == SEQ_CODON)
-            out << "CODON, ";*/
-        if (!saln->partitions[part]->sequence_type.empty())
-            out << saln->partitions[part]->sequence_type << ", ";
-        string pos = saln->partitions[part]->position_spec;
-        replace(pos.begin(), pos.end(), ',' , ' ');
-        out << pos << ";" << endl;
-    }
-}
-
-void PhyloSuperTree::printBestPartitionParams(const char *filename)
-{
+void PhyloSuperTree::printBestPartitionParams(const char *filename) {
     try {
         ofstream out;
         out.exceptions(ios::failbit | ios::badbit);
         out.open(filename);
         out << "#nexus" << endl
         << "begin sets;" << endl;
-        auto *saln = (SuperAlignment*)aln;
-
-        printCharsets(saln, size(), out);
-
+        int part;
+        SuperAlignment *saln = (SuperAlignment*)aln;
+        for (part = 0; part < size(); part++) {
+            string name = saln->partitions[part]->name;
+            replace(name.begin(), name.end(), '+', '_');
+            out << "  charset " << name << " = ";
+            if (!saln->partitions[part]->aln_file.empty()) out << saln->partitions[part]->aln_file << ": ";
+            /*if (saln->partitions[part]->seq_type == SEQ_CODON)
+                out << "CODON, ";*/
+            if (!saln->partitions[part]->sequence_type.empty())
+                out << saln->partitions[part]->sequence_type << ", ";
+            string pos = saln->partitions[part]->position_spec;
+            replace(pos.begin(), pos.end(), ',' , ' ');
+            out << pos << ";" << endl;
+        }
         out << "  charpartition mymodels =" << endl;
-        for (int part = 0; part < size(); part++) {
+        for (part = 0; part < size(); part++) {
             string name = saln->partitions[part]->name;
             replace(name.begin(), name.end(), '+', '_');
             if (part > 0) out << "," << endl;
@@ -1564,45 +1556,3 @@ void PhyloSuperTree::printBestPartitionParams(const char *filename)
         outError(ERR_WRITE_OUTPUT, filename);
     }
 }
-
-void PhyloSuperTree::printMrBayesFreeRateReplacement(RateHeterogeneity *rate, string &charset, ofstream &out) {
-    // Call PhyloTree's function inside the correct checkpoint struct
-    checkpoint->startStruct("PartitionModelPlen");
-    checkpoint->startStruct(charset);
-    PhyloTree::printMrBayesFreeRateReplacement(rate, charset, out);
-    checkpoint->endStruct();
-    checkpoint->endStruct();
-}
-
-void PhyloSuperTree::printMrBayesBlock(const char *filename, bool inclParams)
-{
-    ofstream out = startMrBayesBlock(filename);
-    auto *saln = (SuperAlignment*)aln;
-    printCharsets(saln, size(), out);
-    out << endl;
-
-    // Create Partition
-    size_type listSize = size();
-    out << "  partition iqtree = " << listSize << ": ";
-    for (int part = 0; part < listSize; part++) {
-        string name = saln->partitions[part]->name;
-        replace(name.begin(), name.end(), '+', '_');
-        out << name;
-        if (part != listSize - 1) out << ", ";
-        else out << ";" << endl;
-    }
-
-    // Set Partition for Use
-    out << "  set partition = iqtree;" << endl;
-
-    // Set Outgroup (if available)
-    if (!rooted) out << "  outgroup " << root->name << ";" << endl << endl;
-
-    // Partition-Specific Model Information
-    for (int part = 0; part < listSize; part++) {
-        // MrBayes Partitions are 1-indexed
-        printMrBayesModelText(this, at(part), out, convertIntToString(part + 1), saln->partitions[part]->name, inclParams);
-    }
-    out << "end;" << endl;
-    out.close();
-}
diff --git a/tree/phylosupertree.h b/tree/phylosupertree.h
index f2304fcab..60dbdcb9a 100644
--- a/tree/phylosupertree.h
+++ b/tree/phylosupertree.h
@@ -475,25 +475,12 @@ class PhyloSuperTree : public IQTree, public vector<PhyloTree* >
      */
     void printBestPartitionParams(const char *filename);
 
-    /**
-        print mr bayes block with partition information and best model parameters
-        @param filename output file name
-        @param inclParams whether to include iqtree params
-      */
-    virtual void printMrBayesBlock(const char *filename, bool inclParams);
-
-    /**
-     * Prints the replacement prior settings for +R or +R+I.
-     * @param rate the heterogeneity rate
-     * @param charset the (original) charset of the current partition. An empty string if not a partitioned tree
-     * @param out the ofstream to print to
-     */
-    virtual void printMrBayesFreeRateReplacement(RateHeterogeneity* rate, string &charset, ofstream &out);
-
+    
     /** True when mixed codon with other data type */
     bool rescale_codon_brlen;
-
+    
     int totalNNIs, evalNNIs;
+
 };
 
 #endif
diff --git a/tree/phylotree.cpp b/tree/phylotree.cpp
index 8b1f6bce8..b6bfcb595 100644
--- a/tree/phylotree.cpp
+++ b/tree/phylotree.cpp
@@ -37,7 +37,7 @@
 #include "model/modelmixture.h"
 #include "phylonodemixlen.h"
 #include "phylotreemixlen.h"
-#include "main/phylotesting.h"
+
 
 const int LH_MIN_CONST = 1;
 
@@ -6246,197 +6246,3 @@ void PhyloTree::doNNI_simple(NNIMove &move) {
     }
 }
 
-/* Private Supplementary Logic Functions */
-
-/**
- * Find Rate Value (propinv for +I, gamma for +G/+R, invgamma for +I+G/+I+R) <br>
- * Default is Equal (+E)
- */
-string getMrBayesRHASName(RateHeterogeneity* rate)
-{
-    bool hasGamma = rate->getGammaShape() != 0.0;
-    bool hasInvariable = rate->getPInvar() != 0.0;
-    if (hasGamma) {
-        if (hasInvariable)
-            return "invgamma";
-        return "gamma";
-    }
-    if (hasInvariable)
-        return "propinv";
-    return "equal";
-}
-
-void printMrBayesRHASPrior(PhyloTree* originalTree, RateHeterogeneity* rate, string &charset, ofstream &out)
-{
-    bool hasI = rate->getPInvar() > 0.0;
-
-    // Freerate (+R)
-    // Get replacement gamma shape
-    if (rate->isFreeRate()) {
-        originalTree ->printMrBayesFreeRateReplacement(rate, charset, out);
-    }
-
-    // Gamma Distribution (+G/+R)
-    // Dirichlet is not available here, use fixed
-    if (rate->getGammaShape() > 0.0)
-        out << " shapepr=fixed(" << rate->getGammaShape() << ")";
-
-    // Invariable Sites (+I)
-    // Dirichlet is not available here, use fixed
-    if (hasI)
-        out << " pinvarpr=fixed(" << rate->getPInvar() << ")";
-}
-
-void printMrBayesModelTextDNA(PhyloTree* originalTree, PhyloTree* tree, ofstream &out, string &partition, string &charset, bool inclParams)
-{
-    ModelSubst* model = tree->getModel();
-    bool equalFreq = model->freq_type == FREQ_EQUAL;
-    short nst = 6;
-    string rateMod = getMrBayesRHASName(tree->getRate());
-
-    // Find NST value (1 for JC/JC69/F81, 2 for K80/K2P/HKY/HKY85, 6 for SYM/GTR)
-    // NST is set to 6 by default (above), so check for JC/JC69/F81 and K80/K2P/HKY/HKY85
-    // If it returns a valid dna code, we can extract information from there (needed for user-defined models)
-    string dnaCode = model->getModelDNACode();
-    string modelStr = model->getName();
-    if (!dnaCode.empty()) {
-        // if all equal
-        if (dnaCode.find_first_not_of(dnaCode[0]) == string::npos)
-            nst = 1;
-        // if A-G=C-T and everything else is equal to first part of dna code
-        else if (dnaCode[1] == dnaCode[4] && dnaCode[0] == dnaCode[2] && dnaCode[0] == dnaCode[3] && dnaCode[0] == dnaCode[5])
-            nst = 2;
-    } else if (!modelStr.empty()) {
-        // Check the name of the model
-        if (strcmp(modelStr.c_str(), "JC") == 0 || strcmp(modelStr.c_str(), "JC69") == 0 || strcmp(modelStr.c_str(), "F81") == 0)
-            nst = 1;
-        else if (strcmp(modelStr.c_str(), "K80") == 0 || strcmp(modelStr.c_str(), "K2P") == 0 || strcmp(modelStr.c_str(), "HKY") == 0 || strcmp(modelStr.c_str(), "HKY85") == 0)
-            nst = 2;
-    }
-
-    // NucModel = 4By4: DNA Nucleotides (4 Options, A, C, G, T)
-    out << "  lset applyto=(" << partition << ") nucmodel=4by4 nst=" << nst << " rates=" << rateMod << ";" << endl;
-
-    // Priors (apart from Fixed Freq)
-    if (!inclParams) {
-        // If not include other params, simply set fixed frequency and return
-        if (!equalFreq) return;
-
-        out << "  prset applyto=(" << partition << ") statefreqpr=fixed(equal);" << endl;
-        return;
-    }
-
-    out << "  prset applyto=(" << partition << ")";
-
-    printMrBayesRHASPrior(originalTree, tree->getRate(), charset, out);
-
-    // Frequency of Nucleotides (+F)
-    if (equalFreq)
-        out << " statefreqpr=fixed(equal)";
-    else {
-        auto* freq = new double[model->num_states];
-        model->getStateFrequency(freq);
-        out << " statefreqpr=dirichlet(";
-        for (int i = 0; i < model->num_states; ++i) {
-            if (i != 0) out << ", ";
-            out << freq[i];
-        }
-        out << ")";
-
-        delete[] freq;
-    }
-
-    // Reversible Rate Matrix
-    int numRateEntries = model->getNumRateEntries();
-    auto* rateMat = new double[numRateEntries];
-    model->getRateMatrix(rateMat);
-
-    out << " revmatpr=dirichlet(";
-    for (int i = 0; i < numRateEntries; ++i) {
-        if (i != 0) out << ", ";
-        out << rateMat[i];
-    }
-
-    // Close revmatpr + prset
-    out << ");" << endl;
-
-    delete[] rateMat;
-}
-
-/* Public Supplementary Functions (Used by PhyloSuperTree) */
-ofstream startMrBayesBlock(const char *filename) {
-    ofstream out;
-    try {
-        out.exceptions(ios::failbit | ios::badbit);
-        out.open(filename);
-
-        // Write Warning
-        out << "#nexus" << endl << endl
-            <<"[This MrBayes Block Declaration provides the retrieved values from the IQTree Run.]" << endl
-            << "[Note that MrBayes does not support a large collection of models, so defaults of 'nst=6' for DNA and 'wag' for Protein will be used if a model that does not exist in MrBayes is selected.]" << endl
-            << "[Furthermore, the Model Parameter '+R' will be replaced by '+G'.]" << endl
-            << "[This should be used as a Template Only.]" << endl << endl;
-
-        // Begin File, Print Charsets
-        out << "begin mrbayes;" << endl;
-    } catch (ios::failure &) {
-        outError(ERR_WRITE_OUTPUT, filename);
-    }
-    return out;
-}
-
-void printMrBayesModelText(PhyloTree* originalTree, PhyloTree* tree, ofstream& out, string partition, string charset, bool inclParams)
-{
-    switch(tree->aln->seq_type) {
-        case SEQ_DNA:
-            printMrBayesModelTextDNA(originalTree, tree, out, partition, charset, inclParams);
-            break;
-        case SEQ_CODON:
-            break;
-        case SEQ_PROTEIN:
-            break;
-        case SEQ_BINARY:
-            break;
-        case SEQ_MORPH:
-            break;
-        default:
-            outWarning("MrBayes Block output not supported for sequence type " + getSeqTypeName(tree->aln->seq_type));
-            return;
-    }
-    out << endl; // Add extra endline after model text for readability
-}
-
-/* Protected Overrided Function */
-void PhyloTree::printMrBayesFreeRateReplacement(RateHeterogeneity *rate, string &charset, ofstream &out) {
-    bool hasInvar = rate->getPInvar() > 0.0;
-    string key = hasInvar ? "RateGammaInvar" : "RateGamma";
-
-    double gamma_shape = 0.0;
-    double p_invar = 0.0;
-
-    checkpoint->startStruct(key);
-
-    CKP_RESTORE(gamma_shape);
-    if (hasInvar)
-        CKP_RESTORE(p_invar);
-
-    checkpoint->endStruct();
-
-    // Sanity Check
-    if (gamma_shape > 0.0)
-        out << " shapepr=fixed(" << gamma_shape << ")";
-
-    if (p_invar > 0.0)
-        out << " pinvarpr=fixed(" << p_invar << ")";
-}
-
-/* Main Function */
-void PhyloTree::printMrBayesBlock(const char *filename, bool inclParams) {
-    ofstream out = startMrBayesBlock(filename);
-    // Set Outgroup (if available)
-    if (!rooted) out << "  outgroup " << root->name << ";" << endl << endl;
-
-    printMrBayesModelText(this, this, out, "all", "", inclParams);
-    out << "end;" << endl;
-    out.close();
-}
diff --git a/tree/phylotree.h b/tree/phylotree.h
index f1f2029e0..b8a19f986 100644
--- a/tree/phylotree.h
+++ b/tree/phylotree.h
@@ -2264,21 +2264,6 @@ class PhyloTree : public MTree, public Optimization, public CheckpointFactory {
      */
     const string& getDistanceFileWritten() const;
 
-    /**
-        print MrBayes block with partition information, and best model parameters (if inclParams is true)
-        @param filename output file name
-        @param inclParams whether to include iqtree params
-      */
-    virtual void printMrBayesBlock(const char *filename, bool inclParams);
-
-    /**
-     * Prints the replacement prior settings for +R or +R+I.
-     * @param rate the heterogeneity rate
-     * @param charset the (original) charset of the current partition. An empty string if not a partitioned tree
-     * @param out the ofstream to print to
-     */
-    virtual void printMrBayesFreeRateReplacement(RateHeterogeneity* rate, string &charset, ofstream &out);
-
     
 protected:
 
@@ -2526,22 +2511,4 @@ class PhyloTree : public MTree, public Optimization, public CheckpointFactory {
     void doneProgress();
 };
 
-/**
- * print MrBayes formatted text, of model information (and parameters), to an ofstream.
- * @param originalTree the original tree, of which is called to retrieve values in checkpoint file (to allow for overriding)
- * @param tree tree to print info from
- * @param out stream to print to
- * @param partition the partition to 'apply' to (all to apply to all partitions)
- * @param charset the (original) charset of the current partition. Used to retrieve values in checkpoint file. Input an empty string if this is not a partitioned tree.
- * @param inclParams whether to include the parameters of the model
- */
-void printMrBayesModelText(PhyloTree* originalTree, PhyloTree* tree, ofstream& out, string partition, string charset, bool inclParams);
-
-/**
- * Start a MrBayes Block, as well as printing a warning.
- * @param filename output file name
- * @return ofstream for other functions to print to
- */
-ofstream startMrBayesBlock(const char *filename);
-
 #endif
diff --git a/utils/tools.cpp b/utils/tools.cpp
index db21f7917..d0081931b 100644
--- a/utils/tools.cpp
+++ b/utils/tools.cpp
@@ -6317,6 +6317,7 @@ void usage_iqtree(char* argv[], bool full_command) {
         << "  -mrbayes             Outputs a Mr Bayes block file, to use as a template for future analysis" << endl
         << "                       Will not output if incompatible data types are used" << endl
         << "                       (MrBayes only supports DNA, Codon, Protein, Binary and Morphological data)" << endl
+        << "                       There will be no substitution model data output for Lie Markov and Mixture models!" << endl
         << endl;
 
     if (full_command) {
@@ -7910,3 +7911,36 @@ string getOutputNameWithExt(const InputType& format, const string& output_filepa
             return output_filepath + ".phy";
     }
 }
+
+void warnLogStream(string warn, ofstream &out) {
+    outWarning(warn);
+    out << "[" << warn << "]" << endl;
+}
+
+double minValueCheckMrBayes(double origValue) {
+     if (origValue < 0.01) {
+         outWarning("MrBayes does not support values < 0.01! Using 0.01 instead...");
+         return 0.01;
+     }
+     return origValue;
+}
+
+// Cached Map
+unordered_map<string, string> iqTreeToMrBayesAAModels;
+
+unordered_map<string, string> getIqTreeToMrBayesAAModels() {
+    if (!iqTreeToMrBayesAAModels.empty()) return iqTreeToMrBayesAAModels;
+
+    iqTreeToMrBayesAAModels["Poisson"] = "poisson";
+    iqTreeToMrBayesAAModels["JTT"] = "jones";
+    iqTreeToMrBayesAAModels["Dayhoff"] = "dayhoff";
+    iqTreeToMrBayesAAModels["mtREV"] = "mtrev";
+    iqTreeToMrBayesAAModels["mtMAM"] = "mtmam";
+    iqTreeToMrBayesAAModels["WAG"] = "wag";
+    iqTreeToMrBayesAAModels["rtREV"] = "rtrev";
+    iqTreeToMrBayesAAModels["cpREV"] = "cprev";
+    iqTreeToMrBayesAAModels["VT"] = "vt";
+    iqTreeToMrBayesAAModels["Blosum62"] = "blosum";
+    iqTreeToMrBayesAAModels["LG"] = "lg";
+    return iqTreeToMrBayesAAModels;
+}
diff --git a/utils/tools.h b/utils/tools.h
index 75e538027..b83cd8846 100644
--- a/utils/tools.h
+++ b/utils/tools.h
@@ -3733,4 +3733,21 @@ double frob_norm (double m[], int n, double scale=1.0);
 */
 string getOutputNameWithExt(const InputType& format, const string& output_filepath);
 
+/**
+ * Prints a warning message to the log and to the ofstream, in a NEXUS format.
+ */
+void warnLogStream(string warn, ofstream &out);
+
+/**
+ * ensures a number, to be inputted into MrBayes, is larger than the minimum value for MrBayes (0.01)
+ */
+double minValueCheckMrBayes(double origValue);
+
+
+/**
+ * get a map of iqtree amino acid/protein substitution models to MrBayes amino acid/protein substitution models.<br>
+ * models which are not supported by mrbayes are not included. GTR20 is assumed as default.
+ */
+unordered_map<string, string> getIqTreeToMrBayesAAModels();
+
 #endif

From d36486ccd2a66979842f1af9b1170042da4adaa5 Mon Sep 17 00:00:00 2001
From: Integer Limit <103940576+IntegerLimit@users.noreply.github.com>
Date: Mon, 8 Jul 2024 17:47:22 +1000
Subject: [PATCH 3/3] Support Exporting Codon Analysis to MrBayes Block Files
 (#266)

* Codon Model

* Fix Compiler Error due to Merge Conflicts

* Fix Codon Model `NucModel` Parameter

* Fix Empirical Warning + Indentation of Warnings

* Improve Start-Of-File Warnings

* Fix Indentation in alignment.cpp
---
 alignment/alignment.cpp   | 66 +++++++++++++++++++++++++--------------
 alignment/alignment.h     |  6 ++++
 main/phyloanalysis.cpp    | 20 +++++++++---
 model/modelbin.cpp        |  4 ++-
 model/modelbin.h          |  2 +-
 model/modelcodon.cpp      | 66 +++++++++++++++++++++++++++++++++++++++
 model/modelcodon.h        | 11 +++++++
 model/modeldna.cpp        |  3 +-
 model/modeldna.h          |  3 +-
 model/modelmorphology.cpp |  4 ++-
 model/modelmorphology.h   |  3 +-
 model/modelprotein.cpp    |  4 ++-
 model/modelprotein.h      |  3 +-
 model/modelsubst.cpp      |  2 +-
 model/modelsubst.h        |  3 +-
 utils/tools.cpp           | 42 +++++++++++++++----------
 utils/tools.h             | 11 ++++++-
 17 files changed, 193 insertions(+), 60 deletions(-)

diff --git a/alignment/alignment.cpp b/alignment/alignment.cpp
index 16e5b7d88..ee1e9fc76 100644
--- a/alignment/alignment.cpp
+++ b/alignment/alignment.cpp
@@ -23,6 +23,7 @@
 
 #include <Eigen/LU>
 #ifdef USE_BOOST
+#include <boost/bimap.hpp>
 #include <boost/math/distributions/binomial.hpp>
 #endif
 
@@ -62,6 +63,8 @@ char genetic_code23[] = "KNKNTTTTRSRSIIMIQHQHPPPPRRRRLLLLEDEDAAAAGGGGVVVV*Y*YSSS
 char genetic_code24[] = "KNKNTTTTSSKSIIMIQHQHPPPPRRRRLLLLEDEDAAAAGGGGVVVV*Y*YSSSSWCWCLFLF"; // Pterobranchia mitochondrial
 char genetic_code25[] = "KNKNTTTTRSRSIIMIQHQHPPPPRRRRLLLLEDEDAAAAGGGGVVVV*Y*YSSSSGCWCLFLF"; // Candidate Division SR1 and Gracilibacteria
 
+boost::bimap<int, char*> genetic_code_map;
+
 Alignment::Alignment()
         : vector<Pattern>()
 {
@@ -1603,7 +1606,32 @@ void Alignment::convertStateStr(string &str, SeqType seq_type) {
         (*it) = convertState(*it, seq_type);
 }
 */
- 
+
+boost::bimap<int, char*> getGeneticCodeMap() {
+    if (!genetic_code_map.empty()) return genetic_code_map;
+
+    genetic_code_map.insert({1, genetic_code1});
+    genetic_code_map.insert({2, genetic_code2});
+    genetic_code_map.insert({3, genetic_code3});
+    genetic_code_map.insert({4, genetic_code4});
+    genetic_code_map.insert({5, genetic_code5});
+    genetic_code_map.insert({6, genetic_code6});
+    genetic_code_map.insert({9, genetic_code9});
+    genetic_code_map.insert({10, genetic_code10});
+    genetic_code_map.insert({11, genetic_code11});
+    genetic_code_map.insert({12, genetic_code12});
+    genetic_code_map.insert({13, genetic_code13});
+    genetic_code_map.insert({14, genetic_code14});
+    genetic_code_map.insert({16, genetic_code16});
+    genetic_code_map.insert({21, genetic_code21});
+    genetic_code_map.insert({22, genetic_code22});
+    genetic_code_map.insert({23, genetic_code23});
+    genetic_code_map.insert({24, genetic_code24});
+    genetic_code_map.insert({25, genetic_code25});
+
+    return genetic_code_map;
+}
+
 void Alignment::initCodon(char *gene_code_id) {
     // build index from 64 codons to non-stop codons
 	int transl_table = 1;
@@ -1613,30 +1641,13 @@ void Alignment::initCodon(char *gene_code_id) {
 		} catch (string &str) {
 			outError("Wrong genetic code ", gene_code_id);
 		}
-		switch (transl_table) {
-		case 1: genetic_code = genetic_code1; break;
-		case 2: genetic_code = genetic_code2; break;
-		case 3: genetic_code = genetic_code3; break;
-		case 4: genetic_code = genetic_code4; break;
-		case 5: genetic_code = genetic_code5; break;
-		case 6: genetic_code = genetic_code6; break;
-		case 9: genetic_code = genetic_code9; break;
-		case 10: genetic_code = genetic_code10; break;
-		case 11: genetic_code = genetic_code11; break;
-		case 12: genetic_code = genetic_code12; break;
-		case 13: genetic_code = genetic_code13; break;
-		case 14: genetic_code = genetic_code14; break;
-		case 15: genetic_code = genetic_code15; break;
-		case 16: genetic_code = genetic_code16; break;
-		case 21: genetic_code = genetic_code21; break;
-		case 22: genetic_code = genetic_code22; break;
-		case 23: genetic_code = genetic_code23; break;
-		case 24: genetic_code = genetic_code24; break;
-		case 25: genetic_code = genetic_code25; break;
-		default:
+		auto codeMap = getGeneticCodeMap();
+		auto found = codeMap.left.find(transl_table);
+		if (found == codeMap.left.end()) {
 			outError("Wrong genetic code ", gene_code_id);
-			break;
+			return;
 		}
+		genetic_code = found->second;
 	} else {
 		genetic_code = genetic_code1;
 	}
@@ -1669,6 +1680,15 @@ void Alignment::initCodon(char *gene_code_id) {
 //	cout << "num_states = " << num_states << endl;
 }
 
+int Alignment::getGeneticCodeId() {
+    if (seq_type != SEQ_CODON || genetic_code == nullptr) return 0;
+
+    auto code_map = getGeneticCodeMap();
+    auto found = code_map.right.find(genetic_code);
+    if (found == code_map.right.end()) return 0;
+    return found->second;
+}
+
 int getMorphStates(StrVector &sequences) {
 	char maxstate = 0;
 	for (StrVector::iterator it = sequences.begin(); it != sequences.end(); it++)
diff --git a/alignment/alignment.h b/alignment/alignment.h
index 934591a80..bf7d6d800 100644
--- a/alignment/alignment.h
+++ b/alignment/alignment.h
@@ -1010,6 +1010,12 @@ class Alignment : public vector<Pattern>, public CharSet, public StateSpace {
      */
     void extractMapleFile(const std::string& aln_name, const InputType& format);
 
+    /**
+     * Get the numerical id of the genetic code
+     * @return id the genetic code id, or 0 if not a codon type
+     */
+    int getGeneticCodeId();
+
 protected:
 
 
diff --git a/main/phyloanalysis.cpp b/main/phyloanalysis.cpp
index 7dd7475d5..b1addbcf8 100644
--- a/main/phyloanalysis.cpp
+++ b/main/phyloanalysis.cpp
@@ -2608,11 +2608,21 @@ void printMrBayesBlockFile(const char* filename, IQTree* &iqtree, bool inclParam
         out.exceptions(ios::failbit | ios::badbit);
         out.open(filename);
 
+        string provide = "basic models";
+        if (inclParams) provide = "optimized values";
+        else if (iqtree->isSuperTree()) provide = "basic partition structure and models";
+
         // Write Warning
         out << "#nexus" << endl << endl
-            <<"[This MrBayes Block Declaration provides the retrieved values from the IQTree Run.]" << endl
-            << "[Note that MrBayes does not support a large collection of models, so defaults of 'nst=6' for DNA and 'wag' for Protein will be used if a model that does not exist in MrBayes is selected.]" << endl
-            << "[Furthermore, the Model Parameter '+R' will be replaced by '+G'.]" << endl
+            <<"[This MrBayes Block Declaration provides the " << provide << " from the IQTree Run.]" << endl
+            << "[Note that MrBayes does not support a large collection of models, so defaults of 'nst=6' for DNA and 'wag' for Protein will be used if a model that does not exist in MrBayes is used.]" << endl;
+
+        if (inclParams)
+            out << "[However, for those cases, there will still be a rate matrix provided.]" << endl
+                << "[For DNA, this will still mean the rates may vary outside the restrictions of the model.]" << endl
+                << "[For Protein, this is essentially a perfect replacement.]" << endl;
+
+        out << "[Furthermore, the Model Parameter '+R' will be replaced by '+G'.]" << endl
             << "[This should be used as a Template Only.]" << endl << endl;
 
         // Begin File, Print Charsets
@@ -2626,7 +2636,7 @@ void printMrBayesBlockFile(const char* filename, IQTree* &iqtree, bool inclParam
         if (!iqtree->rooted) out << "  outgroup " << iqtree->root->name << ";" << endl << endl;
 
         out << "  [Using Model '" << iqtree->getModelName() << "']" << endl;
-        iqtree->getModel()->printMrBayesModelText(iqtree->getRate(), out, "all", "", false, inclParams);
+        iqtree->getModel()->printMrBayesModelText(out, "all", "", false, inclParams);
 
         out << endl << "end;" << endl;
         out.close();
@@ -2670,7 +2680,7 @@ void printMrBayesBlockFile(const char* filename, IQTree* &iqtree, bool inclParam
 
         // MrBayes Partitions are 1-indexed
         out << "  [Partition No. " << convertIntToString(part + 1) << ", Using Model '" << currentTree->getModelName() << "']" << endl;
-        currentTree->getModel()->printMrBayesModelText(currentTree->getRate(), out,
+        currentTree->getModel()->printMrBayesModelText(out,
                 convertIntToString(part + 1), saln->partitions[part]->name, true, inclParams);
         out << endl;
     }
diff --git a/model/modelbin.cpp b/model/modelbin.cpp
index 262cbea80..e0100fc26 100644
--- a/model/modelbin.cpp
+++ b/model/modelbin.cpp
@@ -77,7 +77,9 @@ string ModelBIN::getNameParams(bool show_fixed_params) {
     return retname.str();
 }
 
-void ModelBIN::printMrBayesModelText(RateHeterogeneity* rate, ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams) {
+void ModelBIN::printMrBayesModelText(ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams) {
+    RateHeterogeneity* rate = phylo_tree->getRate();
+
     // MrBayes does not support Invariable Modifier for Binary data
     if (rate->isFreeRate() || rate->getPInvar() > 0.0) {
         warnLogStream("MrBayes does not support Invariable Sites with Binary Data! +I has been ignored!", out);
diff --git a/model/modelbin.h b/model/modelbin.h
index b2c257d9e..b6a5024ec 100644
--- a/model/modelbin.h
+++ b/model/modelbin.h
@@ -65,7 +65,7 @@ class ModelBIN : public ModelMarkov
      * @param isSuperTree whether the tree is a super tree. Useful for retrieving information from the checkpoint file, which has different locations for PhyloTree and PhyloSuperTree
      * @param inclParams whether to include IQTree optimized parameters for the model
      */
-    virtual void printMrBayesModelText(RateHeterogeneity* rate, ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams);
+    virtual void printMrBayesModelText(ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams);
 
 };
 
diff --git a/model/modelcodon.cpp b/model/modelcodon.cpp
index d39c0390a..f95190743 100644
--- a/model/modelcodon.cpp
+++ b/model/modelcodon.cpp
@@ -1142,6 +1142,72 @@ double ModelCodon::optimizeParameters(double gradient_epsilon) {
 	return score;
 }
 
+void ModelCodon::printMrBayesModelText(ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams) {
+    // NST should be 1 if fix_kappa is true (no ts/tv ratio), else it should be 2
+    int nst = fix_kappa ? 1 : 2;
+    int code = phylo_tree->aln->getGeneticCodeId();
+    string mrBayesCode = getMrBayesGeneticCode(code);
+    RateHeterogeneity* rate = phylo_tree->getRate();
+
+    if (rate->isFreeRate() || rate->getGammaShape() > 0.0 || rate->getPInvar() > 0.0) {
+        warnLogStream("MrBayes Codon Models do not support any Heterogenity Rate Modifiers! (+I, +G, +R) They have been ignored!", out);
+    }
+
+    if (mrBayesCode.empty()) {
+        warnLogStream("MrBayes Does Not Support Codon Code " + convertIntToString(code) + "! Defaulting to Universal (Code 1).", out);
+        mrBayesCode = "universal";
+    }
+
+    if (num_params == 0 || name.find('_') != string::npos) { // If this model is a Empirical / Semi-Empirical Model
+        warnLogStream("MrBayes does not support Empirical or Semi-Empirical Codon Models. State Frequency will still be set, but no rate matrix will be set.", out);
+    }
+
+    out << "  lset applyto=(" << partition << ") nucmodel=codon omegavar=equal nst=" << nst << " code=" << mrBayesCode << ";" << endl;
+
+    if (!inclParams) return;
+
+    out << "  prset applyto=(" << partition << ") statefreqpr=dirichlet(";
+
+    // State Frequency Prior
+    bool isFirst = true;
+    for (int i = 0; i < num_states; i++) {
+        if (phylo_tree->aln->isStopCodon(i)) continue;
+
+        if (!isFirst) out << ", ";
+        else isFirst = false;
+
+        out << state_freq[i];
+    }
+
+    out << ")";
+
+    /*
+     * TS/TV Ratio (Kappa)
+     * Ratio Should be:
+     * `beta(kappa, 1)` when `codon_kappa_style` is `CK_ONE_KAPPA` (kappa here represents the ts/tv rate ratio)
+     * `beta(kappa, 1)` when `codon_kappa_style` is `CK_ONE_KAPPA_TS` (kappa here represents the transition rate)
+     * `beta(1, kappa)` when `codon_kappa_style` is `CK_ONE_KAPPA_TV` (kappa here represents the transversion rate)
+     * `beta(kappa, kappa2)` when `codon_kappa_style` is `CK_TWO_KAPPA` (kappa here represents the transition rate, and kappa2 represents the transversion rate)
+     */
+    if (!fix_kappa) {
+        double v1 = codon_kappa_style == CK_ONE_KAPPA_TV ? 1 : kappa;
+        double v2 = 1;
+        if (codon_kappa_style == CK_ONE_KAPPA_TV)
+            v2 = kappa;
+        else if (codon_kappa_style == CK_TWO_KAPPA)
+            v2 = kappa2;
+
+        out << " tratiopr=beta(" << v1 << ", " << v2 << ")";
+    }
+
+    // DN/DS Ratio (Omega)
+    // Ratio is set to `dirichlet(omega, 1)`
+    if (!fix_omega) {
+        out << " omegapr=dirichlet(" << omega << ", 1)";
+    }
+
+    out << ";" << endl;
+}
 
 void ModelCodon::writeInfo(ostream &out) {
     if (name.find('_') == string::npos)
diff --git a/model/modelcodon.h b/model/modelcodon.h
index 5b95c46df..a745f49c8 100644
--- a/model/modelcodon.h
+++ b/model/modelcodon.h
@@ -210,6 +210,17 @@ class ModelCodon: public ModelMarkov {
 	*/
 	virtual bool getVariables(double *variables);
 
+    /**
+     * Print the model information in a format that can be accepted by MrBayes, using lset and prset.<br>
+     * By default, it simply prints a warning to the log and to the stream, stating that this model is not supported by MrBayes.
+     * @param out the ofstream to print the result to
+     * @param partition the partition to apply lset and prset to
+     * @param charset the current partition's charset. Useful for getting information from the checkpoint file
+     * @param isSuperTree whether the tree is a super tree. Useful for retrieving information from the checkpoint file, which has different locations for PhyloTree and PhyloSuperTree
+     * @param inclParams whether to include IQTree optimized parameters for the model
+     */
+    virtual void printMrBayesModelText(ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams);
+
 };
 
 #endif /* MODELCODON_H_ */
diff --git a/model/modeldna.cpp b/model/modeldna.cpp
index 4a36e16be..5a4aae840 100644
--- a/model/modeldna.cpp
+++ b/model/modeldna.cpp
@@ -565,9 +565,10 @@ void ModelDNA::setVariables(double *variables) {
 //                      }
 }
 
-void ModelDNA::printMrBayesModelText(RateHeterogeneity* rate, ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams) {
+void ModelDNA::printMrBayesModelText(ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams) {
     bool equalFreq = freq_type == FREQ_EQUAL;
     short nst = 6;
+    RateHeterogeneity* rate = phylo_tree->getRate();
 
     // Find NST value (1 for JC/JC69/F81, 2 for K80/K2P/HKY/HKY85, 6 for SYM/GTR)
     // NST is set to 6 by default (above), so check for JC/JC69/F81 and K80/K2P/HKY/HKY85
diff --git a/model/modeldna.h b/model/modeldna.h
index 0c252c178..98fb70f1a 100644
--- a/model/modeldna.h
+++ b/model/modeldna.h
@@ -117,14 +117,13 @@ class ModelDNA : public ModelMarkov
     /**
      * Print the model information in a format that can be accepted by MrBayes, using lset and prset.<br>
      * By default, it simply prints a warning to the log and to the stream, stating that this model is not supported by MrBayes.
-     * @param rate the rate information
      * @param out the ofstream to print the result to
      * @param partition the partition to apply lset and prset to
      * @param charset the current partition's charset. Useful for getting information from the checkpoint file
      * @param isSuperTree whether the tree is a super tree. Useful for retrieving information from the checkpoint file, which has different locations for PhyloTree and PhyloSuperTree
      * @param inclParams whether to include IQTree optimized parameters for the model
      */
-    virtual void printMrBayesModelText(RateHeterogeneity* rate, ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams);
+    virtual void printMrBayesModelText(ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams);
 
 protected:
 
diff --git a/model/modelmorphology.cpp b/model/modelmorphology.cpp
index ddda51874..8312a6439 100644
--- a/model/modelmorphology.cpp
+++ b/model/modelmorphology.cpp
@@ -151,11 +151,13 @@ void ModelMorphology::writeInfo(ostream &out) {
     }
 }
 
-void ModelMorphology::printMrBayesModelText(RateHeterogeneity* rate, ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams) {
+void ModelMorphology::printMrBayesModelText(ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams) {
     warnLogStream("MrBayes only supports Morphological Data with states from {0-9}!", out);
     warnLogStream("Morphological Data with states {A-Z} may cause errors!", out);
     warnLogStream("Use Morphological Models in MrBayes with Caution!", out);
 
+    RateHeterogeneity* rate = phylo_tree->getRate();
+
     // MrBayes does not support Invariable Modifier for Morph data
     if (rate->isFreeRate() || rate->getPInvar() > 0.0) {
         warnLogStream("MrBayes does not support Invariable Sites with Morphological Data! +I has been ignored!", out);
diff --git a/model/modelmorphology.h b/model/modelmorphology.h
index f598103bd..6367b1282 100644
--- a/model/modelmorphology.h
+++ b/model/modelmorphology.h
@@ -82,14 +82,13 @@ class ModelMorphology: public ModelMarkov {
     /**
      * Print the model information in a format that can be accepted by MrBayes, using lset and prset.<br>
      * By default, it simply prints a warning to the log and to the stream, stating that this model is not supported by MrBayes.
-     * @param rate the rate information
      * @param out the ofstream to print the result to
      * @param partition the partition to apply lset and prset to
      * @param charset the current partition's charset. Useful for getting information from the checkpoint file
      * @param isSuperTree whether the tree is a super tree. Useful for retrieving information from the checkpoint file, which has different locations for PhyloTree and PhyloSuperTree
      * @param inclParams whether to include IQTree optimized parameters for the model
      */
-    virtual void printMrBayesModelText(RateHeterogeneity* rate, ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams);
+    virtual void printMrBayesModelText(ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams);
 };
 
 #endif /* MODELMORPHOLOGY_H_ */
diff --git a/model/modelprotein.cpp b/model/modelprotein.cpp
index 24f8a99ff..558e938e3 100644
--- a/model/modelprotein.cpp
+++ b/model/modelprotein.cpp
@@ -1162,10 +1162,12 @@ void ModelProtein::computeTipLikelihood(PML::StateType state, double *state_lk)
     state_lk[ambi_aa[cstate*2+1]] = 1.0;
 }
 
-void ModelProtein::printMrBayesModelText(RateHeterogeneity* rate, ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams) {
+void ModelProtein::printMrBayesModelText(ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams) {
     // Lset Parameters
     out << "  lset applyto=(" << partition << ") nucmodel=protein rates=";
 
+    RateHeterogeneity* rate = phylo_tree->getRate();
+
     // RHAS Specification
     // Free Rate should be substituted by +G+I
     bool hasGamma = rate->getGammaShape() != 0.0 || rate->isFreeRate();
diff --git a/model/modelprotein.h b/model/modelprotein.h
index 6dd1f69e6..b24cb7c1d 100644
--- a/model/modelprotein.h
+++ b/model/modelprotein.h
@@ -82,14 +82,13 @@ class ModelProtein : public ModelMarkov
     /**
      * Print the model information in a format that can be accepted by MrBayes, using lset and prset.<br>
      * By default, it simply prints a warning to the log and to the stream, stating that this model is not supported by MrBayes.
-     * @param rate the rate information
      * @param out the ofstream to print the result to
      * @param partition the partition to apply lset and prset to
      * @param charset the current partition's charset. Useful for getting information from the checkpoint file
      * @param isSuperTree whether the tree is a super tree. Useful for retrieving information from the checkpoint file, which has different locations for PhyloTree and PhyloSuperTree
      * @param inclParams whether to include IQTree optimized parameters for the model
      */
-    virtual void printMrBayesModelText(RateHeterogeneity* rate, ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams);
+    virtual void printMrBayesModelText(ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams);
 
 private:
     ModelsBlock *models_block;
diff --git a/model/modelsubst.cpp b/model/modelsubst.cpp
index 5c5675bea..fe4fc8ada 100644
--- a/model/modelsubst.cpp
+++ b/model/modelsubst.cpp
@@ -246,7 +246,7 @@ void ModelSubst::printMrBayesFreeRateReplacement(bool isSuperTree, string &chars
         out << " pinvarpr=fixed(" << minValueCheckMrBayes(p_invar) << ")";
 }
 
-void ModelSubst::printMrBayesModelText(RateHeterogeneity* rate, ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams) {
+void ModelSubst::printMrBayesModelText(ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams) {
     out << "  [MrBayes Block Output is not supported by this model!]" << endl;
     outWarning("MrBayes Block Output is not supported by this model of name '" + name + "'!");
 }
diff --git a/model/modelsubst.h b/model/modelsubst.h
index 370737418..71bc0a14f 100644
--- a/model/modelsubst.h
+++ b/model/modelsubst.h
@@ -405,14 +405,13 @@ class ModelSubst: public Optimization, public CheckpointFactory
     /**
      * Print the model information in a format that can be accepted by MrBayes, using lset and prset.<br>
      * By default, it simply prints a warning to the log and to the stream, stating that this model is not supported by MrBayes.
-     * @param rate the rate information
      * @param out the ofstream to print the result to
      * @param partition the partition to apply lset and prset to
      * @param charset the current partition's charset. Useful for getting information from the checkpoint file
      * @param isSuperTree whether the tree is a super tree. Useful for retrieving information from the checkpoint file, which has different locations for PhyloTree and PhyloSuperTree
      * @param inclParams whether to include IQTree optimized parameters for the model
      */
-    virtual void printMrBayesModelText(RateHeterogeneity* rate, ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams);
+    virtual void printMrBayesModelText(ofstream& out, string partition, string charset, bool isSuperTree, bool inclParams);
 
 	/*****************************************************
 		Checkpointing facility
diff --git a/utils/tools.cpp b/utils/tools.cpp
index d0081931b..af5865b9a 100644
--- a/utils/tools.cpp
+++ b/utils/tools.cpp
@@ -7912,9 +7912,9 @@ string getOutputNameWithExt(const InputType& format, const string& output_filepa
     }
 }
 
-void warnLogStream(string warn, ofstream &out) {
+void warnLogStream(string warn, ofstream &out, string indentation) {
     outWarning(warn);
-    out << "[" << warn << "]" << endl;
+    out << indentation << "[" << warn << "]" << endl;
 }
 
 double minValueCheckMrBayes(double origValue) {
@@ -7925,22 +7925,30 @@ double minValueCheckMrBayes(double origValue) {
      return origValue;
 }
 
-// Cached Map
-unordered_map<string, string> iqTreeToMrBayesAAModels;
+const unordered_map<string, string> iqTreeToMrBayesAAModels = {
+        {"Poisson", "poisson"},
+        {"JTT", "jones"},
+        {"Dayhoff", "dayhoff"},
+        {"mtREV", "mtrev"},
+        {"mtMAM", "mtmam"},
+        {"WAG", "wag"},
+        {"rtREV", "rtrev"},
+        {"cpREV", "cprev"},
+        {"VT", "vt"},
+        {"Blosum62", "blosum"},
+        {"LG", "lg"},
+};
+
+// Anything outside of index 10 (Code No. 11) is invalid, leave that as empty string
+const string indexedMrBayesGeneticCodes[25] = {"universal", "vertmt", "yeast", "mycoplasma", "invermt",
+                                               "ciliate", "", "", "echinoderm", "euplotid", "universal"};
 
 unordered_map<string, string> getIqTreeToMrBayesAAModels() {
-    if (!iqTreeToMrBayesAAModels.empty()) return iqTreeToMrBayesAAModels;
-
-    iqTreeToMrBayesAAModels["Poisson"] = "poisson";
-    iqTreeToMrBayesAAModels["JTT"] = "jones";
-    iqTreeToMrBayesAAModels["Dayhoff"] = "dayhoff";
-    iqTreeToMrBayesAAModels["mtREV"] = "mtrev";
-    iqTreeToMrBayesAAModels["mtMAM"] = "mtmam";
-    iqTreeToMrBayesAAModels["WAG"] = "wag";
-    iqTreeToMrBayesAAModels["rtREV"] = "rtrev";
-    iqTreeToMrBayesAAModels["cpREV"] = "cprev";
-    iqTreeToMrBayesAAModels["VT"] = "vt";
-    iqTreeToMrBayesAAModels["Blosum62"] = "blosum";
-    iqTreeToMrBayesAAModels["LG"] = "lg";
     return iqTreeToMrBayesAAModels;
 }
+
+string getMrBayesGeneticCode(int geneticCodeId) {
+    if (geneticCodeId == 0) return "";
+
+    return indexedMrBayesGeneticCodes[geneticCodeId - 1];
+}
diff --git a/utils/tools.h b/utils/tools.h
index b83cd8846..77675a54f 100644
--- a/utils/tools.h
+++ b/utils/tools.h
@@ -3735,8 +3735,11 @@ string getOutputNameWithExt(const InputType& format, const string& output_filepa
 
 /**
  * Prints a warning message to the log and to the ofstream, in a NEXUS format.
+ * @param warn the message to print
+ * @param out the ofstream to print to
+ * @param indentation optional indentation included only in ofstream. defaults to two spaces
  */
-void warnLogStream(string warn, ofstream &out);
+void warnLogStream(string warn, ofstream &out, string indentation = "  ");
 
 /**
  * ensures a number, to be inputted into MrBayes, is larger than the minimum value for MrBayes (0.01)
@@ -3750,4 +3753,10 @@ double minValueCheckMrBayes(double origValue);
  */
 unordered_map<string, string> getIqTreeToMrBayesAAModels();
 
+/**
+ * get the MrBayes equivalent of a genetic code, given the id of the code. Returns and empty string if that code
+ * is not supported in MrBayes.
+ */
+string getMrBayesGeneticCode(int geneticCodeId);
+
 #endif