add missing logging in XML for relaxed clock #161 LinguaPhylo/linguaP…

…hylo#488

lphybeast/src/main/java/lphybeast/BEASTContext.java

@@ -105,6 +105,8 @@ public class BEASTContext {
     //*** operators ***//
     // a list of extra loggables in 3 default loggers: parameter logger, screen logger, tree logger.
     private List<Loggable> extraLoggables = new ArrayList<>();
+    // a list of extra loggables in 3 default loggers: parameter logger, screen logger, tree logger.
+    private List<Loggable> skipLoggables = new ArrayList<>();
     // helper to create extra loggers from extensions
     private List<LoggerHelper> extraLoggers = new ArrayList<>();
 

lphybeast/src/main/java/lphybeast/tobeast/LoggerUtils.java

@@ -0,0 +1,19 @@
+package lphybeast.tobeast;
+
+import beast.base.evolution.RateStatistic;
+import beast.base.evolution.branchratemodel.BranchRateModel;
+import beast.base.evolution.tree.TreeInterface;
+
+public class LoggerUtils {
+
+    public static RateStatistic createRateStatistic(String id, BranchRateModel branchRateModel, TreeInterface tree) {
+        // <log id="" ... branchratemodel="@" tree="@"/>
+        RateStatistic rateStatistic = new RateStatistic();
+        rateStatistic.setInputValue("branchratemodel", branchRateModel);
+        rateStatistic.setInputValue("tree", tree);
+        rateStatistic.initAndValidate();
+        rateStatistic.setID(id);
+        return rateStatistic;
+    }
+
+}

lphybeast/src/main/java/lphybeast/tobeast/generators/UCLNRelaxedClockToBEAST.java

@@ -1,7 +1,9 @@
 package lphybeast.tobeast.generators;
 
 import beast.base.core.BEASTInterface;
+import beast.base.evolution.RateStatistic;
 import beast.base.evolution.branchratemodel.UCRelaxedClockModel;
+import beast.base.evolution.tree.TreeInterface;
 import beast.base.inference.distribution.LogNormalDistributionModel;
 import beast.base.inference.distribution.Prior;
 import beast.base.inference.parameter.RealParameter;
@@ -12,6 +14,8 @@
 import lphy.core.model.Value;
 import lphybeast.BEASTContext;
 import lphybeast.GeneratorToBEAST;
+import lphybeast.tobeast.LoggerUtils;
+import lphybeast.tobeast.loggers.MetaDataTreeLogger;
 
 /**
  * For ORC package.
@@ -31,7 +35,8 @@ public UCRelaxedClockModel generatorToBEAST(UCLNMean1 ucln, BEASTInterface value
         ucRelaxedClockModel.setInputValue("distr", logNormDist);
 
         Value<TimeTree> tree = ucln.getTree();
-        ucRelaxedClockModel.setInputValue("tree", context.getBEASTObject(tree));
+        TreeInterface beastTree = (TreeInterface) context.getBEASTObject(tree);
+        ucRelaxedClockModel.setInputValue("tree", beastTree);
 
         // set clock.rate to UclnMean
         for (GraphicalModelNode treeOut : tree.getOutputs()) {
@@ -58,8 +63,13 @@ public UCRelaxedClockModel generatorToBEAST(UCLNMean1 ucln, BEASTInterface value
         // in case to duplicate with RandomValue id also called "branchRates"
         ucRelaxedClockModel.setID(branchRates.getUniqueId() + ".model");
 
-        //TODO rm rates from log, replaced by
-        // <log id="ORCRatesStat" spec="beast.base.evolution.RateStatistic" branchratemodel="@OptimisedRelaxedClock" tree="@Tree"/>
+        //TODO rm rates from log, replaced by <log id="" ... branchratemodel="@" tree="@"/>
+        RateStatistic rateStatistic = LoggerUtils.createRateStatistic("RatesStat." + branchRates.getUniqueId(), ucRelaxedClockModel, beastTree);
+        context.addExtraLoggable(rateStatistic);
+
+        // Extra Logger <log ... branchratemodel="@..." tree="@..."/>
+        MetaDataTreeLogger metaDataTreeLogger = new MetaDataTreeLogger(ucRelaxedClockModel, beastTree, context);
+        context.addExtraLogger(metaDataTreeLogger);
 
         return ucRelaxedClockModel;
     }

lphybeast/src/main/java/lphybeast/tobeast/loggers/MetaDataTreeLogger.java

@@ -0,0 +1,71 @@
+package lphybeast.tobeast.loggers;
+
+import beast.base.core.BEASTInterface;
+import beast.base.evolution.TreeWithMetaDataLogger;
+import beast.base.evolution.branchratemodel.BranchRateModel;
+import beast.base.evolution.tree.TreeInterface;
+import beast.base.inference.Logger;
+import com.google.common.collect.Multimap;
+import lphy.core.model.GraphicalModelNode;
+import lphybeast.BEASTContext;
+
+import java.util.Objects;
+
+import static java.lang.Math.toIntExact;
+
+/**
+ * Log branch rates with tree.
+ * @see TreeWithMetaDataLogger
+ * @author Walter Xie
+ */
+public class MetaDataTreeLogger implements TreeLoggerHelper {
+    // add extra tree logger for AncestralStateTreeLikelihood
+    final protected BranchRateModel branchRateModel;
+    final protected BEASTContext context;
+    final TreeInterface tree;
+    String fileName;
+
+    public MetaDataTreeLogger(BranchRateModel branchRateModel, TreeInterface tree, BEASTContext context) {
+        this.branchRateModel = Objects.requireNonNull(branchRateModel);
+        this.tree = Objects.requireNonNull(tree);
+        this.context = context;
+    }
+
+    @Override
+    public Logger createLogger(long logEvery, Multimap<BEASTInterface, GraphicalModelNode<?>> elements) {
+        Logger logger = new Logger();
+        // Must convert to int
+        logger.setInputValue("logEvery", toIntExact(logEvery));
+        logger.setInputValue("fileName", getFileName());
+        logger.setInputValue("mode", "tree");
+
+        TreeWithMetaDataLogger treeWithMetaDataLogger = new TreeWithMetaDataLogger();
+        treeWithMetaDataLogger.setInputValue("tree", tree);
+        treeWithMetaDataLogger.setInputValue("branchratemodel", branchRateModel);
+
+        logger.setInputValue("log", treeWithMetaDataLogger);
+        logger.initAndValidate();
+
+        logger.setID("TreeWithMetaDataLogger." + tree.getID());
+        elements.put(logger, null);
+        return logger;
+    }
+
+    @Override
+    public TreeInterface getTree() {
+        return tree;
+    }
+
+    @Override
+    public String getFileName() {
+        return fileName;
+    }
+
+    @Override
+    public void setFileName(String fileStem, boolean isMultiple) {
+        if (isMultiple) // multi-partitions and unlink trees
+            fileName = fileStem + "_with_rates." + getTree().getID() + ".trees";
+        else
+            fileName = fileStem + "_with_rates.trees";
+    }
+}