Add LPhyBEAST for UCLN LinguaPhylo/linguaPhylo#488

LinguaPhylo · Oct 4, 2024 · 3e3c1ef · 3e3c1ef
1 parent 8d8002e
commit 3e3c1ef
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 17 deletions.
diff --git a/lphybeast/src/main/java/lphybeast/spi/LPhyBEASTExtImpl.java b/lphybeast/src/main/java/lphybeast/spi/LPhyBEASTExtImpl.java
@@ -100,6 +100,7 @@ public List<Class<? extends GeneratorToBEAST>> getGeneratorToBEASTs() {
 //                StructuredCoalescentToMascot.class,
                 TreeLengthToBEAST.class,
                 TN93ToBEAST.class,
+                UCLNRelaxedClockToBEAST.class,
                 UniformToBEAST.class,
                 VectorizedDistributionToBEAST.class,
                 VectorizedFunctionToBEAST.class,

diff --git a/lphybeast/src/main/java/lphybeast/tobeast/generators/PhyloCTMCToBEAST.java b/lphybeast/src/main/java/lphybeast/tobeast/generators/PhyloCTMCToBEAST.java
@@ -10,7 +10,6 @@
 import beast.base.evolution.sitemodel.SiteModel;
 import beast.base.evolution.substitutionmodel.SubstitutionModel;
 import beast.base.evolution.tree.Tree;
-import beast.base.inference.distribution.Prior;
 import beast.base.inference.parameter.RealParameter;
 import beastclassic.evolution.alignment.AlignmentFromTrait;
 import beastclassic.evolution.likelihood.AncestralStateTreeLikelihood;
@@ -20,7 +19,7 @@
 import consoperators.SimpleDistance;
 import consoperators.SmallPulley;
 import lphy.base.distribution.DiscretizedGamma;
-import lphy.base.distribution.LogNormal;
+import lphy.base.distribution.UCLN;
 import lphy.base.evolution.branchrate.LocalBranchRates;
 import lphy.base.evolution.branchrate.LocalClock;
 import lphy.base.evolution.likelihood.PhyloCTMC;
@@ -174,25 +173,15 @@ public static void constructTreeAndBranchRate(PhyloCTMC phyloCTMC, GenericTreeLi
         if (branchRates != null) {
 
             Generator generator = branchRates.getGenerator();
-            if (generator instanceof IID &&
-                    ((IID<?>) generator).getBaseDistribution() instanceof LogNormal) {
+            if (generator instanceof UCLN ucln) {
 
-                // simpleRelaxedClock.lphy
-                UCRelaxedClockModel relaxedClockModel = new UCRelaxedClockModel();
+                UCRelaxedClockModel relaxedClockModel = (UCRelaxedClockModel) context.getBEASTObject(generator);
+                relaxedClockModel.setID(branchRates.getCanonicalId());
 
-                Prior logNormalPrior = (Prior) context.getBEASTObject(generator);
-
-                RealParameter beastBranchRates = context.getAsRealParameter(branchRates);
-
-                relaxedClockModel.setInputValue("rates", beastBranchRates);
-                relaxedClockModel.setInputValue("tree", tree);
-                relaxedClockModel.setInputValue("distr", logNormalPrior.distInput.get());
-                relaxedClockModel.setID(branchRates.getCanonicalId() + ".model");
-                relaxedClockModel.initAndValidate();
                 treeLikelihood.setInputValue("branchRateModel", relaxedClockModel);
 
                 if (skipBranchOperators == false) {
-                    addRelaxedClockOperators(tree, relaxedClockModel, beastBranchRates, context);
+                    addRelaxedClockOperators(tree, relaxedClockModel, context);
                 }
 
             } else if (generator instanceof LocalBranchRates) {
@@ -287,7 +276,9 @@ public static SiteModel constructSiteModel(PhyloCTMC phyloCTMC, BEASTContext con
      * @deprecated this will be replaced by ORC soon
      */
     @Deprecated
-    private static void addRelaxedClockOperators(Tree tree, UCRelaxedClockModel relaxedClockModel, RealParameter rates, BEASTContext context) {
+    private static void addRelaxedClockOperators(Tree tree, UCRelaxedClockModel relaxedClockModel, BEASTContext context) {
+
+        RealParameter rates = relaxedClockModel.rateInput.get();
 
         double tWindowSize = tree.getRoot().getHeight() / 10.0;
 

diff --git a/lphybeast/src/main/java/lphybeast/tobeast/generators/UCLNRelaxedClockToBEAST.java b/lphybeast/src/main/java/lphybeast/tobeast/generators/UCLNRelaxedClockToBEAST.java
@@ -0,0 +1,43 @@
+package lphybeast.tobeast.generators;
+
+import beast.base.core.BEASTInterface;
+import beast.base.evolution.branchratemodel.UCRelaxedClockModel;
+import beast.base.inference.distribution.LogNormalDistributionModel;
+import beast.base.inference.parameter.RealParameter;
+import lphy.base.distribution.UCLN;
+import lphy.base.evolution.tree.TimeTree;
+import lphy.core.model.Value;
+import lphybeast.BEASTContext;
+import lphybeast.GeneratorToBEAST;
+
+public class UCLNRelaxedClockToBEAST implements GeneratorToBEAST<UCLN, UCRelaxedClockModel> {
+    @Override
+    public UCRelaxedClockModel generatorToBEAST(UCLN ucln, BEASTInterface value, BEASTContext context) {
+
+        UCRelaxedClockModel ucRelaxedClockModel = new UCRelaxedClockModel();
+
+        LogNormalDistributionModel logNormDist = new LogNormalDistributionModel();
+        logNormDist.setInputValue("M", context.getAsRealParameter(ucln.getUclnMean()));
+        logNormDist.setInputValue("S", context.getAsRealParameter(ucln.getUclnSigma()));
+        logNormDist.initAndValidate();
+        ucRelaxedClockModel.setInputValue("distr", logNormDist);
+
+        Value<TimeTree> tree = ucln.getTree();
+        ucRelaxedClockModel.setInputValue("tree", context.getBEASTObject(tree));
+
+        if (value instanceof RealParameter rates) {
+            ucRelaxedClockModel.setInputValue("rates", rates);
+        } else throw new IllegalArgumentException("Value sampled from LPhy UCLN should be mapped to RealParameter ! " + value);
+
+        ucRelaxedClockModel.initAndValidate();
+        return ucRelaxedClockModel;
+    }
+
+    @Override
+    public Class<UCLN> getGeneratorClass() { return UCLN.class; }
+
+    @Override
+    public Class<UCRelaxedClockModel> getBEASTClass() {
+        return UCRelaxedClockModel.class;
+    }
+}