fix(binomial-distribution): compute pmf in log-space to avoid overflow

Sources/StatKit/Descriptive Statistics/Distribution/Discrete/BinomialDistribution.swift

@@ -36,33 +36,39 @@ public struct BinomialDistribution: DiscreteDistribution, UnivariateDistribution
   public init(probability: Double, trials: Int) {
     precondition(
       0 <= probability && probability <= 1,
-      "The probability needs to be in (0, 1) (\(probability) provided)."
+      "The probability needs to be in [0, 1] (\(probability) provided)."
     )
     precondition(
       0 <= trials,
-      "Possibilities need to non-negative (\(trials) provided)."
+      "Number of trials need to non-negative (\(trials) provided)."
     )
 
     self.probability = probability
     self.trials = trials
   }
 
   public func pmf(x: Int, logarithmic: Bool = false) -> Double {
-    let coefficient = Double(choose(n: trials, k: x))
-    let successes: Double = .pow(probability, Double(x))
-    let failures: Double = .pow(1 - probability, Double(trials - x))
-    let result = coefficient * successes * failures
+    guard 0 <= x && x <= trials else {
+      return logarithmic ? -.infinity : 0
+    }
 
-    switch result {
-      case ...0:
-        return logarithmic ? -.infinity : 0
+    if probability == 0 {
+      let isMode = x == 0
+      return logarithmic ? (isMode ? 0 : -.infinity) : (isMode ? 1 : 0)
+    }
+    if probability == 1 {
+      let isMode = x == trials
+      return logarithmic ? (isMode ? 0 : -.infinity) : (isMode ? 1 : 0)
+    }
 
-      case 1...:
-        return logarithmic ? 0 : 1
+    let logCoefficient = Double.logGamma(Double(trials + 1))
+                       - Double.logGamma(Double(x + 1))
+                       - Double.logGamma(Double(trials - x + 1))
+    let logResult = logCoefficient
+                  + Double(x) * .log(probability)
+                  + Double(trials - x) * .log(1 - probability)
 
-      default:
-        return logarithmic ? .log(result) : result
-    }
+    return logarithmic ? logResult : .exp(logResult)
   }
 
   public func cdf(x: Int, logarithmic: Bool = false) -> Double {
@@ -97,6 +103,7 @@ public struct BinomialDistribution: DiscreteDistribution, UnivariateDistribution
   }
 
   public func sample(_ numberOfElements: Int) -> [Int] {
+    precondition(0 < numberOfElements, "The requested number of samples need to be greater than 0.")
     var rng = Xoroshiro256StarStar()
     let weights = (0 ... trials).map { pmf(x: $0) }
     let table = WalkersAliasTable(weights: weights)

Tests/StatKitTests/Descriptive Statistics Tests/Distribution Tests/Discrete/BinomialDistributionTests.swift

@@ -6,9 +6,9 @@ struct BinomialDistributionTests {
   @Test(
     "Valid distribution parameters return correct mean",
     arguments: [
-      (0, 20, 0),
-      (1, 1, 1),
-      (0.5, 200, 100),
+      (0.0,  20,   0.0),
+      (1.0,   1,   1.0),
+      (0.5, 200, 100.0),
     ]
   )
   func validInputReturnsCorrectMean(probability: Double, trials: Int, expectedMean: Double) async throws {
@@ -19,9 +19,9 @@ struct BinomialDistributionTests {
   @Test(
     "Valid distribution parameters return correct variance",
     arguments: [
-      (0, 20, 0),
-      (1, 1, 0),
-      (0.5, 200, 50),
+      (0.0,  20,  0.0),
+      (1.0,   1,  0.0),
+      (0.5, 200, 50.0),
     ]
   )
   func validInputReturnsCorrectVariance(probability: Double, trials: Int, expectedVariance: Double) async throws {
@@ -32,10 +32,10 @@ struct BinomialDistributionTests {
   @Test(
     "Valid distribution parameters return correct skewness",
     arguments: [
-      (0, 20, .infinity),
-      (1, 1, -.infinity),
-      (0.5, 200, 0),
-      (0.3, 200, 0.06172133998)
+      (0.0,  20, Double.infinity),
+      (1.0,   1, -Double.infinity),
+      (0.5, 200, 0.0),
+      (0.3, 200, 0.06172133998),
     ]
   )
   func validInputReturnsCorrectSkewness(probability: Double, trials: Int, expectedSkewness: Double) async throws {
@@ -46,39 +46,101 @@ struct BinomialDistributionTests {
   @Test(
     "Valid distribution parameters return correct excess kurtosis",
     arguments: [
-      (0, 20, .infinity),
-      (1, 1, .infinity),
-      (0.5, 20, -0.1),
-      (0.7, 20, -0.0619047619),
-      (0.05, 200, 0.07526315789)
+      (0.0,  20, Double.infinity),
+      (1.0,   1, Double.infinity),
+      (0.5,  20, -0.1),
+      (0.7,  20, -0.0619047619),
+      (0.05, 200, 0.07526315789),
     ]
   )
-  func validInputReturnsCorrectSkewness(probability: Double, trials: Int, expectedKurtosis: Double) async throws {
+  func validInputReturnsCorrectExcessKurtosis(probability: Double, trials: Int, expectedKurtosis: Double) async throws {
     let kurtosis = BinomialDistribution(probability: probability, trials: trials).excessKurtosis
     #expect(kurtosis.isApproximatelyEqual(to: expectedKurtosis, absoluteTolerance: 1e-6))
   }
 
   @Test(
     "Valid distribution parameters return correct PMF value",
     arguments: [
-      (0, 20, 10, 0),
-      (0.99, 1, 1, 0.99),
-      (0.5, 20, 5, 0.0147857666015625),
+      (0.0,   20, 10, 0.0),
+      (0.99,   1,  1, 0.99),
+      (0.5,   20,  5, 0.01478576660),
+      (0.5,  100, 50, 0.07958923739),
+      (0.3,  100, 30, 0.08678391876),
     ]
   )
-  func validInputReturnsCorrectSkewness(probability: Double, trials: Int, x: Int, expectedPMF: Double) async throws {
+  func validInputReturnsCorrectPMF(probability: Double, trials: Int, x: Int, expectedPMF: Double) async throws {
     let pmf = BinomialDistribution(probability: probability, trials: trials).pmf(x: x)
     #expect(pmf.isApproximatelyEqual(to: expectedPMF, absoluteTolerance: 1e-6))
   }
 
+  @Test(
+    "Valid distribution parameters return correct log PMF value",
+    arguments: [
+      (0.99,   1,  1, -0.01005034),
+      (0.5,   20,  5, -4.21409028),
+      (0.5,  100, 50, -2.53087640),
+      (0.3,  100, 30, -2.44433456),
+    ]
+  )
+  func validInputReturnsCorrectLogPMF(probability: Double, trials: Int, x: Int, expectedLogPMF: Double) async throws {
+    let pmf = BinomialDistribution(probability: probability, trials: trials).pmf(x: x, logarithmic: true)
+    #expect(pmf.isApproximatelyEqual(to: expectedLogPMF, absoluteTolerance: 1e-5))
+  }
+
+  @Test(
+    "PMF at boundary probability returns correct value",
+    arguments: [
+      (0.0, 10,  0,  1.0),
+      (0.0, 10,  1,  0.0),
+      (1.0, 10, 10,  1.0),
+      (1.0, 10,  9,  0.0),
+    ]
+  )
+  func pmfAtBoundaryProbabilityIsCorrect(probability: Double, trials: Int, x: Int, expectedPMF: Double) async throws {
+    let pmf = BinomialDistribution(probability: probability, trials: trials).pmf(x: x)
+    #expect(pmf.isApproximatelyEqual(to: expectedPMF, absoluteTolerance: 1e-6))
+  }
+
+  @Test(
+    "Log PMF at boundary probability returns zero",
+    arguments: [
+      (0.0, 10,  0),
+      (1.0, 10, 10),
+    ]
+  )
+  func logPMFAtBoundaryProbabilityIsZero(probability: Double, trials: Int, x: Int) async throws {
+    let pmf = BinomialDistribution(probability: probability, trials: trials).pmf(x: x, logarithmic: true)
+    #expect(pmf == 0)
+  }
+
+  @Test(
+    "PMF outside valid range returns zero",
+    arguments: [(0.5, 10), (0.3, 100)]
+  )
+  func pmfOutsideRangeIsZero(probability: Double, trials: Int) async throws {
+    let distribution = BinomialDistribution(probability: probability, trials: trials)
+    #expect(distribution.pmf(x: -1) == 0)
+    #expect(distribution.pmf(x: trials + 1) == 0)
+  }
+
+  @Test(
+    "Log PMF outside valid range returns negative infinity",
+    arguments: [(0.5, 10), (0.3, 100)]
+  )
+  func logPMFOutsideRangeIsNegativeInfinity(probability: Double, trials: Int) async throws {
+    let distribution = BinomialDistribution(probability: probability, trials: trials)
+    #expect(distribution.pmf(x: -1, logarithmic: true) == -.infinity)
+    #expect(distribution.pmf(x: trials + 1, logarithmic: true) == -.infinity)
+  }
+
   @Test(
     "Valid distribution parameters return correct CDF value",
     arguments: [
-      (0.5, 20, 1, 0.0000200271606),
-      (0.5, 20, 0, 0.0000009536743),
+      (0.5, 20,  1, 0.0000200271606),
+      (0.5, 20,  0, 0.0000009536743),
       (0.5, 20, 10, 0.5880985260010),
       (0.7, 20, 10, 0.04796190),
-      (0.7, 20, 7, 0.00127888),
+      (0.7, 20,  7, 0.00127888),
       (0.7, 20, 15, 0.76249222),
     ]
   )
@@ -88,36 +150,56 @@ struct BinomialDistributionTests {
   }
 
   @Test(
-    "Valid distribution parameter return correct log CDF value",
+    "Valid distribution parameters return correct log CDF value",
     arguments: [
-      (0.5, 20, 1, -10.8184212),
-      (0.5, 20, 0, -13.8629436),
-      (0.5, 20, 10, -0.5308608),
-      (0.7, 20, 10, -3.037348),
-      (0.7, 20, 7, -6.661771),
-      (0.7, 20, 15, -0.271163),
+      (0.5, 20,  1, -10.8184212),
+      (0.5, 20,  0, -13.8629436),
+      (0.5, 20, 10,  -0.5308608),
+      (0.7, 20, 10,  -3.037348),
+      (0.7, 20,  7,  -6.661771),
+      (0.7, 20, 15,  -0.271163),
     ]
   )
   func validInputReturnsCorrectLogCDF(probability: Double, trials: Int, x: Int, expectedLogCDF: Double) async throws {
     let cdf = BinomialDistribution(probability: probability, trials: trials).cdf(x: x, logarithmic: true)
     #expect(cdf.isApproximatelyEqual(to: expectedLogCDF, absoluteTolerance: 1e-6))
   }
 
+#if swift(>=6.2)
+  @Test("Negative trial count triggers a precondition failure")
+  func negativeTrialsTriggersPreconditionFailure() async {
+    await #expect(processExitsWith: .failure) {
+      _ = BinomialDistribution(probability: 0.5, trials: -1)
+    }
+  }
+
+  @Test("Out-of-range probability triggers a precondition failure")
+  func outOfRangeProbabilityTriggersPreconditionFailure() async {
+    await #expect(processExitsWith: .failure) {
+      _ = BinomialDistribution(probability: 1.1, trials: 10)
+    }
+  }
+
+  @Test("Non-positive sample count triggers a precondition failure")
+  func nonPositiveSampleCountTriggersPreconditionFailure() async {
+    await #expect(processExitsWith: .failure) {
+      _ = BinomialDistribution(probability: 0.5, trials: 10).sample(0)
+    }
+  }
+#endif
+
   @Test("Sampling from a distribution returns correct proportions")
   func testSampling() async throws {
-    let numberOfSamples = 1000000
+    let numberOfSamples = 1_000_000
     let numberOfTrials = 50
     let distribution = BinomialDistribution(probability: 0.7, trials: numberOfTrials)
     let samples = distribution.sample(numberOfSamples)
 
     var proportions = samples.reduce(into: [Int: Double]()) { result, number in result[Int(number), default: 0] += 1 }
-
     for key in proportions.keys { proportions[key]? /= Double(numberOfSamples) }
 
     #expect(samples.count == numberOfSamples)
-    let testRange =  0 ... numberOfTrials
-
-    for successes in testRange {
+    for successes in 0 ... numberOfTrials {
       #expect(proportions[successes, default: 0.0].isApproximatelyEqual(to: distribution.pmf(x: successes), absoluteTolerance: 1e-3))
     }
   }