stat for Simon

beast-dev · Dec 28, 2024 · fb0d1de · fb0d1de
1 parent 41a001c
commit fb0d1de
Showing 1 changed file with 39 additions and 1 deletion.
diff --git a/src/dr/evomodel/continuous/ContinuousDiffusionStatistic.java b/src/dr/evomodel/continuous/ContinuousDiffusionStatistic.java
@@ -74,6 +74,7 @@ public class ContinuousDiffusionStatistic extends Statistic.Abstract {
     public static final String SPEARMAN = "spearman";
     public static final String CORRELATION_COEFFICIENT = "correlationCoefficient";
     public static final String DISTANCE_TIME_CORRELATION = "distanceTimeCorrelation";
+    public static final String SQUAREDDISTANCE_TIME4_CORRELATION = "squaredDistanceTimeFourCorrelation";
     public static final String R_SQUARED = "Rsquared";
     public static final String STATISTIC = "statistic";
     public static final String TRAIT = "trait";
@@ -458,6 +459,18 @@ public double getStatisticValue(int dim) {
                 Regression r = new Regression(convertDoubles(times),convertDoubles(distances));
                 return r.getCorrelationCoefficient();
             }
+        }  else if (summaryStat == summaryStatistic.SQUAREDDISTANCE_TIME4_CORRELATION)  {
+            List<Double> squareddistances = squareElements(distances);
+            List<Double> timesFour = elementsTimesFour(times);
+            if (summaryMode == Mode.SPEARMAN) {
+                return getSpearmanRho(convertDoubles(timesFour),convertDoubles(squareddistances));
+            } else if (summaryMode == Mode.R_SQUARED) {
+                Regression r = new Regression(convertDoubles(timesFour), convertDoubles(squareddistances));
+                return r.getRSquared();
+            } else {
+                Regression r = new Regression(convertDoubles(timesFour),convertDoubles(squareddistances));
+                return r.getCorrelationCoefficient();
+            }
          }  else {
             return treeLength;
         }
@@ -490,6 +503,22 @@ private double[] toArray(List<Double> list) {
         return returnArray;
     }
 
+    private static List<Double> squareElements(List<Double> inputList) {
+        List<Double> squaredList = new ArrayList<>();
+        for (Double number : inputList) {
+            squaredList.add(number * number);
+        }
+        return squaredList;
+    }
+
+    private static List<Double> elementsTimesFour (List<Double> inputList) {
+        List<Double> returnList = new ArrayList<>();
+        for (Double number : inputList) {
+            returnList.add(number * 4);
+        }
+        return returnList;
+    }
+
     private double[] imputeValue(double[] nodeValue, double[] parentValue, double time, double nodeHeight, double parentHeight, double[] precisionArray, double rate, boolean trueNoise) {
 
         final double scaledTimeChild = (time - nodeHeight) * rate;
@@ -932,7 +961,8 @@ enum summaryStatistic {
         WAVEFRONT_DISTANCE,
         WAVEFRONT_DISTANCE_PHYLO,
         WAVEFRONT_RATE,
-        DISTANCE_TIME_CORRELATION
+        DISTANCE_TIME_CORRELATION,
+        SQUAREDDISTANCE_TIME4_CORRELATION
     }
 
     enum BranchSet {
@@ -1023,6 +1053,12 @@ public Object parseXMLObject(XMLObject xo) throws XMLParseException {
                     System.err.println(name+": mode = "+mode+" ignored for "+DISTANCE_TIME_CORRELATION+", reverting to correlation coefficient mode");
                     statMode = Mode.CORRELATION_COEFFICIENT;
                 }
+            } else if (statistic.equals(SQUAREDDISTANCE_TIME4_CORRELATION)) {
+                summaryStat = summaryStatistic.SQUAREDDISTANCE_TIME4_CORRELATION;
+                if (mode.equals(AVERAGE) || mode.equals(WEIGHTED_AVERAGE) || mode.equals(COEFFICIENT_OF_VARIATION) || mode.equals(MEDIAN)){
+                    System.err.println(name+": mode = "+mode+" ignored for "+SQUAREDDISTANCE_TIME4_CORRELATION+", reverting to correlation coefficient mode");
+                    statMode = Mode.CORRELATION_COEFFICIENT;
+                }
             } else if (statistic.equals(WAVEFRONT_DISTANCE)) {
                 summaryStat = summaryStatistic.WAVEFRONT_DISTANCE;
                 if (!mode.equals(WEIGHTED_AVERAGE)) {
@@ -1065,6 +1101,8 @@ public Object parseXMLObject(XMLObject xo) throws XMLParseException {
                 summaryStat = summaryStatistic.DIFFUSION_COEFFICIENT;
             } else if (statistic.equals(DISTANCE_TIME_CORRELATION)) {
                 summaryStat = summaryStatistic.DISTANCE_TIME_CORRELATION;
+            } else if (statistic.equals(SQUAREDDISTANCE_TIME4_CORRELATION)) {
+                summaryStat = summaryStatistic.SQUAREDDISTANCE_TIME4_CORRELATION;
             } else {
                 System.err.println(name+": unknown statistic: "+statistic+". Reverting to diffusion rate.");
                 summaryStat = summaryStatistic.DIFFUSION_RATE;