gradients for conditioning on survival

beast-dev · Oct 3, 2023 · 1593fd8 · 1593fd8
1 parent 359de1b
commit 1593fd8
Show file tree

Hide file tree

Showing 8 changed files with 50 additions and 30 deletions.
diff --git a/ci/TestXML/testCrssbdpGradient.xml b/ci/TestXML/testCrssbdpGradient.xml
@@ -46,7 +46,7 @@
 	<!-- ########## -->
 
 	<!--full model-->
-	<newBirthDeathSerialSampling id="bdss1" units="years" conditionOnSurvival="false">
+	<newBirthDeathSerialSampling id="bdss1" units="years" conditionOnSurvival="true">
 		<birthRate>
 			<parameter id="bdss1.birthRate" value="8.472" lower="0.0"/>
 		</birthRate>
@@ -130,7 +130,7 @@
 	<!-- ########## -->
 
 	<!-- rho = 1 case -->
-	<newBirthDeathSerialSampling id="bdss2" units="years" conditionOnSurvival="false">
+	<newBirthDeathSerialSampling id="bdss2" units="years" conditionOnSurvival="true">
 		<birthRate>
 			<parameter id="bdss2.birthRate" value="8.472" lower="0.0"/>
 		</birthRate>
@@ -211,7 +211,7 @@
 	<!-- ########## -->
 
 	<!-- rho = 0 case -->
-	<newBirthDeathSerialSampling id="bdss3" units="years" conditionOnSurvival="false">
+	<newBirthDeathSerialSampling id="bdss3" units="years" conditionOnSurvival="true">
 		<birthRate>
 			<parameter id="bdss3.birthRate" value="8.472" lower="0.0"/>
 		</birthRate>
@@ -283,7 +283,7 @@
 	<!-- ########## -->
 
 	<!-- r = 1 case -->
-	<newBirthDeathSerialSampling id="bdss4" units="years" conditionOnSurvival="false">
+	<newBirthDeathSerialSampling id="bdss4" units="years" conditionOnSurvival="true">
 		<birthRate>
 			<parameter id="bdss4.birthRate" value="8.472" lower="0.0"/>
 		</birthRate>
@@ -364,7 +364,7 @@
 	<!-- ########## -->
 
 	<!-- r = 0 case -->
-	<newBirthDeathSerialSampling id="bdss5" units="years" conditionOnSurvival="false">
+	<newBirthDeathSerialSampling id="bdss5" units="years" conditionOnSurvival="true">
 		<birthRate>
 			<parameter id="bdss5.birthRate" value="8.472" lower="0.0"/>
 		</birthRate>
@@ -435,7 +435,7 @@
 	<!-- ########## -->
 
 	<!-- mu = 0, r > 0 case -->
-	<newBirthDeathSerialSampling id="bdss6" units="years" conditionOnSurvival="false">
+	<newBirthDeathSerialSampling id="bdss6" units="years" conditionOnSurvival="true">
 		<birthRate>
 			<parameter id="bdss6.birthRate" value="8.472" lower="0.0"/>
 		</birthRate>
@@ -525,7 +525,7 @@
 	</treeModel>
 
 	<!-- psi = 0, 0 < rho < 1 case -->
-	<newBirthDeathSerialSampling id="bdss7" units="years" conditionOnSurvival="false">
+	<newBirthDeathSerialSampling id="bdss7" units="years" conditionOnSurvival="true">
 		<birthRate>
 			<parameter id="bdss7.birthRate" value="8.472" lower="0.0"/>
 		</birthRate>
@@ -590,7 +590,7 @@
 	<!-- ########## -->
 
 	<!-- psi = 0, rho = 1 case -->
-	<newBirthDeathSerialSampling id="bdss8" units="years" conditionOnSurvival="false">
+	<newBirthDeathSerialSampling id="bdss8" units="years" conditionOnSurvival="true">
 		<birthRate>
 			<parameter id="bdss8.birthRate" value="8.472" lower="0.0"/>
 		</birthRate>

diff --git a/ci/TestXML/testEssbdpGradient.xml b/ci/TestXML/testEssbdpGradient.xml
@@ -47,7 +47,7 @@
     <!-- ########## -->
 
     <!-- multiple rates model with three intervals -->
-    <newBirthDeathSerialSampling id="new.bdss1" units="years" hasFinalSample="false" conditionOnSurvival="false">
+    <newBirthDeathSerialSampling id="new.bdss1" units="years" hasFinalSample="false" conditionOnSurvival="true">
         <birthRate>
             <parameter id="bdss1.birthRate" value ="8.47 2.88 1.82" lower="0.0"/>
         </birthRate>
@@ -127,7 +127,7 @@
     <!-- ########## -->
 
     <!-- multiple rates model with three intervals and one intensive sampling event at t1 = 2 -->
-    <newBirthDeathSerialSampling id="new.bdss2" units="years" hasFinalSample="false" conditionOnSurvival="false">
+    <newBirthDeathSerialSampling id="new.bdss2" units="years" hasFinalSample="false" conditionOnSurvival="true">
         <birthRate>
             <parameter id="bdss2.birthRate" value ="8.47 2.88 1.82" lower="0.0"/>
         </birthRate>
@@ -234,7 +234,7 @@
     </treeModel>
 
 
-    <newBirthDeathSerialSampling id="bdss3" units="years" conditionOnSurvival="false">
+    <newBirthDeathSerialSampling id="bdss3" units="years" conditionOnSurvival="true">
         <birthRate>
             <parameter id="bdss3.birthRate" value ="8.47 4.66 1.82" lower="0.0"/>
         </birthRate>
@@ -307,7 +307,7 @@
     <!-- ########## -->
 
     <!-- Origin is in different interval from rest of tree -->
-    <newBirthDeathSerialSampling id="new.bdss4" units="years" hasFinalSample="false" conditionOnSurvival="false">
+    <newBirthDeathSerialSampling id="new.bdss4" units="years" hasFinalSample="false" conditionOnSurvival="true">
         <birthRate>
             <parameter id="bdss4.birthRate" value ="8.47 2.88 1.82" lower="0.0"/>
         </birthRate>

diff --git a/src/dr/evomodel/speciation/BirthDeathEpisodicSeriallySampledModel.java b/src/dr/evomodel/speciation/BirthDeathEpisodicSeriallySampledModel.java
@@ -452,9 +452,7 @@ private double ti(int i) {
     }
 
     @Override
-    public double logConditioningProbability() {
-       return 0;
-    }
+    public double logConditioningProbability(int model) {return 0;}
 
     private double partialApartialLambda(int i) {
         return (lambda(i) - mu(i) + psi(i)) / Ai[i];
@@ -733,7 +731,7 @@ public void processGradientSampling(double[] gradient, int currentModelSegment,
     }
 
     @Override
-    public void logConditioningProbability(double[] gradient) {
+    public void logConditioningProbability(int currentModelSegment, double[] gradient) {
         return;
     }
     @Override

diff --git a/src/dr/evomodel/speciation/CachedGradientDelegate.java b/src/dr/evomodel/speciation/CachedGradientDelegate.java
@@ -120,7 +120,7 @@ private double[] getGradientLogDensityImpl() {
         // origin branch is a fake branch that doesn't exist in the tree, now compute its contribution
         provider.processGradientOrigin(gradient, currentModelSegment, treeIntervals.getTotalDuration());
 
-        provider.logConditioningProbability(gradient);
+        provider.logConditioningProbability(currentModelSegment,gradient);
 
         if (MEASURE_RUN_TIME) {
             timer.stop();

diff --git a/src/dr/evomodel/speciation/EfficientSpeciationLikelihood.java b/src/dr/evomodel/speciation/EfficientSpeciationLikelihood.java
@@ -142,7 +142,7 @@ final BigFastTreeIntervals getTreeIntervals() {
         // origin branch is a fake branch that doesn't exist in the tree, now compute its contribution
         logL += speciationModel.processOrigin(currentModelSegment, treeIntervals.getTotalDuration());
 
-        logL += speciationModel.logConditioningProbability();
+        logL += speciationModel.logConditioningProbability(currentModelSegment);
 
         if (MEASURE_RUN_TIME) {
             timer.stop();

diff --git a/src/dr/evomodel/speciation/NewBirthDeathSerialSamplingModel.java b/src/dr/evomodel/speciation/NewBirthDeathSerialSamplingModel.java
@@ -288,19 +288,21 @@ private static double computeB(double lambda, double mu, double psi, double rho,
     }
 
     @Override
-    public double logConditioningProbability() {
+    public double logConditioningProbability(int model) {
         double logP = 0.0;
         if ( conditionOnSurvival ) {
             double origin = originTime.getParameterValue(0);
             double[] modelBreakPoints = getBreakPoints();
-            int idx = modelBreakPoints.length - 2;
-            double intervalStart = modelBreakPoints[idx];
-            // Origin is probably near the last index, so we do a linear search forwards in time from there
-            while ( origin < intervalStart) {
-                --idx;
-                intervalStart = modelBreakPoints[idx];
+
+            double segmentIntervalEnd = modelBreakPoints[model];
+
+            while (origin >= segmentIntervalEnd) { // TODO Maybe it's >= ?
+                ++model;
+                updateLikelihoodModelValues(model);
+                segmentIntervalEnd = modelBreakPoints[model];
             }
-            logP -= Math.log(1.0 - p(idx, origin));
+
+            logP -= Math.log(1.0 - p(model, origin));
         }
         return logP;
     }
@@ -1021,9 +1023,29 @@ public void processGradientOrigin(double[] gradient, int currentModelSegment, do
     }
 
     @Override
-    public void logConditioningProbability(double[] gradient) {
+    public void logConditioningProbability(int model, double[] gradient) {
+        double grad = 0.0;
+        double[] dPOrigin = new double[numIntervals * 4];
         if ( conditionOnSurvival ) {
-            throw new RuntimeException("Cannot yet condition ESSBDP for gradient.");
+            double origin = originTime.getParameterValue(0);
+            double[] modelBreakPoints = getBreakPoints();
+            double intervalStart = model > 0? modelBreakPoints[model-1]:0;
+            double segmentIntervalEnd = modelBreakPoints[model];
+
+            while (origin >= segmentIntervalEnd) { // TODO Maybe it's >= ?
+                intervalStart = segmentIntervalEnd;
+                ++model;
+                updateGradientModelValues(model);
+                segmentIntervalEnd = modelBreakPoints[model];
+            }
+            double eAt_Origin = Math.exp(A * (origin - intervalStart));
+            grad += 1 /(1.0 - p(eAt_Origin));
+            dPCompute(model, origin, intervalStart, eAt_Origin, dPOrigin, dG2);
+            for (int p = 0; p < 4; p++) {
+                for (int k = 0; k <= model; k++) {
+                    gradient[genericIndex(k, p, numIntervals)] += grad * dPOrigin[k * 4 + p];
+                }
+            }
         }
     }
 

diff --git a/src/dr/evomodel/speciation/SpeciationModel.java b/src/dr/evomodel/speciation/SpeciationModel.java
@@ -140,7 +140,7 @@ public void updateModelValues() {
         throw new RuntimeException("Not implemented");
     }
 
-    public double logConditioningProbability() {
+    public double logConditioningProbability(int model) {
         throw new RuntimeException("Not implemented");
     }
 }
diff --git a/src/dr/evomodel/speciation/SpeciationModelGradientProvider.java b/src/dr/evomodel/speciation/SpeciationModelGradientProvider.java
@@ -90,7 +90,7 @@ default void processGradientOrigin(double[] gradient,
                                        int currentModelSegment,
                                        double totalDuration) { throw new RuntimeException("Not yet implemented"); }
 
-    default void logConditioningProbability(double[] gradient) { throw new RuntimeException("Not yet implemented"); }
+    default void logConditioningProbability(int currentModelSegment, double[] gradient) { throw new RuntimeException("Not yet implemented"); }
 
     default void updateGradientModelValues(int currentModelSegment) { throw new RuntimeException("Not yet implemented"); }