A few discrete probability distributions for Rainier

Discrete.scala

/**
  * Bernoulli distribution with expectation `p`
  *
  * @param p The probability of success
  */
final case class Bernoulli(p: Real) extends Discrete {
  val generator: Generator[Int] =
    Generator.require(Set(p)) { (r, n) =>
      val u = r.standardUniform
      val l = n.toDouble(p)
      if (u <= l) 1 else 0
    }

  def logDensity(v: Real) =
    If(v.eqEps(.001)(1.0), p.log, (1- p).log)
}

/**
  * Discrete Uniform with expectation `(n+m)/2`
  *
  * @param m The lower bound Integer
  * @param n The upper bound Integer
  */

final case class DiscreteUniform(m: Real, n: Real) extends Discrete {
  val generator: Generator[Int] =
    Generator.require(Set(m, n)) { (r, j) =>
      val u = r.standardUniform
      val o = j.toInt(m)
      val p = j.toInt(n)
      o + Math.floor((p - o + 1) * u).toInt
    }

  def logDensity(v: Real) =
    -1 * (n - m + 1).log
}

/**
  * Geometric distribution with expectation `1/p`
  *
  * @param p The probability of success
  */

final case class Geometric(p: Real) extends Discrete {
  val generator: Generator[Int] =
    Generator.require(Set(p)) {(r, n) =>
      val u = r.standardUniform
      val q = n.toDouble(p)
      Math.floor(Math.log(u) / Math.log(1-q)).toInt
    }

  def logDensity(v: Real) =
    p.log + v * (1-p).log
}

/**
  * Binomial Distribution with expectation `np`
  *
  * @param k Total number of trials
  * @param p The probability of success
  */

final case class Binomial(k: Real, p: Real) extends Discrete {
  val generator: Generator[Int] =
    Generator.require(Set(k, p)) { (r, n) =>
      (1 to n.toInt(k)).map({ x =>
        Bernoulli(p).generator.get(r, n)
      }).sum
    }

  def logDensity(v: Real) =
    Combinatorics.factorial(k) - Combinatorics.factorial(v) + v * p.log + (k - v) * (1 - p).log
}

/**
  * Negative Binomial distribution with expectation `n(1-p)/p`
  *
  * @param n Total number of failures
  * @param p Probability of success
  */

final case class NegativeBinomial(n: Real, p: Real) extends Discrete {
  val generator: Generator[Int] =
    Generator.require(Set(n, p)) { (r, m) =>
      (1 to m.toInt(n)).map({ x =>
        Geometric(p).generator.get(r, m)
      }).sum
    }

  def logDensity(v: Real) =
    Combinatorics.factorial(n + v - 1) - Combinatorics.factorial(v) + n * p.log + v * (1 - p).log
}