DiscretizeGamma needs to renormalise category rates #514

LinguaPhylo · Aug 7, 2024 · 5008d3d · 5008d3d
1 parent 334cb56
commit 5008d3d
Showing 1 changed file with 12 additions and 6 deletions.
diff --git a/lphy-base/src/main/java/lphy/base/distribution/DiscretizedGamma.java b/lphy-base/src/main/java/lphy/base/distribution/DiscretizedGamma.java
@@ -51,12 +51,18 @@ protected void constructDistribution(RandomGenerator random) {
             description = "The discretized gamma probability distribution with mean = 1.")
     public RandomVariable<Double> sample() {
 
-        for (int i = 0; i < rates.length; i++) {
-            double q = (2.0 * i + 1.0) / (2.0 * rates.length);
-            rates[i] = gammaDistribution.inverseCumulativeProbability(q);
-        }
-
-        return new RandomVariable<>(null, rates[random.nextInt(rates.length)], this);
+        double meanRate = 0;
+		for (int i = 0; i < rates.length; i++) {
+			double q = (2.0 * i + 1.0) / (2.0 * rates.length);
+			rates[i] = gammaDistribution.inverseCumulativeProbability(q);
+			meanRate += rates[i];
+		}
+		// renormalise cat rates
+		meanRate /= rates.length;
+		for (int i = 0; i < rates.length; i++) {
+			rates[i] /= meanRate;
+		}
+		return new RandomVariable<>(null, rates[random.nextInt(rates.length)], this);
     }
 
     public double logDensity(Double[] x) {