Rainier Beta distribution

Beta.scala

import com.stripe.rainier.compute._
import com.stripe.rainier.core._

object Beta {
  def apply(a: Real, b: Real): Continuous = new Continuous {
    def realLogDensity(p: Real): Real =
      If(p <= 0,
         Real.zero.log,
         If(p >= 1,
            Real.zero.log,
            (a - 1) * p.log + (b - 1) * (1 - p).log - Combinatrics.beta(a, b)))

    def param: RandomVariable[Real] = {
      val x = new Variable
      val logistic = Real.one / (Real.one + (x * -1).exp)
      val logisticJacobian = logistic * (1 - logistic)
      val density = realLogDensity(logistic) + logisticJacobian.log
      RandomVariable(logistic, density)
    }

    def generator: Generator[Double] = Generator.require(Set(a, b)) { (r, n) =>
      val z1 = Gamma(a, 1).generator.get(r, n)
      val z2 = Gamma(b, 1).generator.get(r, n)
      z1 / (z1 + z2)
    }
  }
}