path: root/lib/stats.dx

'# Stats
Probability distributions and other functions useful for statistical computing.

'## Log-space floating point numbers
When working with probability densities, mass functions, distributions,
likelihoods, etc., we often work on a logarithmic scale to prevent floating
point overflow/underflow. Working on the log scale makes `product` operations
safe, since they correspond to addition on the log scale. However, when it
comes to `sum` operations on the raw probability scale, care must be taken to
avoid underflow by using a safe
[log-sum-exp](https://en.wikipedia.org/wiki/LogSumExp) function.
`LogSpace` stores the log value internally, but is "safe" for both `sum` and
`product`, since addition is implemented using the log-sum-exp trick.

'Several of the functions in this library return values as a `LogSpace Float`.
The internally stored logarithmic value can be extracted with `ln`, and the
actual value being represented (the raw probability) can be computed with
`ls_to_f`. Although it is safe to use `sum` on a value of type
`n=>LogSpace f`, it may be slightly more accurate and/or performant to instead
 use `ls_sum`, which applies the standard max-sweep log-sum-exp trick directly,
rather than relying on monoid reduction using `add`.

struct LogSpace(a:Type) =
  log : a

def Exp(x:a) -> LogSpace a given (a:Type) = LogSpace x

instance Mul(LogSpace f) given (f|Add)
  def (*)(x, y) = Exp $ x.log + y.log
  one = Exp zero

instance Fractional(LogSpace f) given (f|Sub)
  def divide(x, y) = Exp $ x.log - y.log

instance Arbitrary(LogSpace Float)
  def arb(k) = Exp (randn k)

def is_infinite(x:f) -> Bool given (f|Fractional|Sub|Mul|Ord) =
  # Note: According to this function, nan is not infinite.
  # Todo: Make polymorphic versions of these and put them in the prelude.
  infinity     : f = divide one zero
  neg_infinity : f = zero - infinity
  x == infinity || x == neg_infinity

def log_add_exp(la:f, lb:f) -> f
    given (f|Floating|Add|Sub|Mul|Fractional|Ord) =
  infinity     : f = divide(one, zero)
  neg_infinity : f = zero - infinity
  if la == infinity && lb == infinity
    then infinity
    else if la == neg_infinity && lb == neg_infinity
      then neg_infinity
      else if (la > lb)
        then la + log1p (exp (lb - la))
        else lb + log1p (exp (la - lb))

instance Add(LogSpace f) given (f|Floating|Add|Sub|Mul|Fractional|Ord)
  def (+)(x, y) = Exp $ log_add_exp x.log y.log
  zero = Exp (zero - (divide one zero))

def f_to_ls(a:f) -> LogSpace f given (f|Floating)  = Exp $ log a

def ls_to_f(x:LogSpace f) -> f given (f|Floating) = exp x.log

def ln(x:LogSpace f) -> f given (f|Floating) = x.log

def log_sum_exp(xs:n=>f) -> f
    given(n|Ix, f|Fractional|Sub|Floating|Mul|Ord) =
  m_raw = maximum(xs)
  m = case is_infinite(m_raw) of
    False -> m_raw
    True  -> zero
  m + (log $ sum $ each xs \x. exp(x) - m)

def ls_sum(x:n=>LogSpace f) -> LogSpace f
    given (n|Ix, f|Fractional|Sub|Floating|Mul|Ord) =
  lx = each(x, ln)
  Exp $ log_sum_exp lx

'## Probability distributions
Simulation and evaluation of [probability distributions](https://en.wikipedia.org/wiki/Probability_distribution). Probability distributions which can be sampled should implement `Random`, and those which can be evaluated should implement the `Dist` interface. Distributions over an ordered space (such as typical *univariate* distributions) should also implement `OrderedDist`.

# TODO: use an associated type for the `a` parameter
interface Random(d:Type, a:Type)
  draw : (d, Key) -> a                # function for random draws

# TODO: use an associated type for the `a` parameter
interface Dist(d:Type, a:Type, f:Type)
  density : (d, a) -> LogSpace f      # either the density function or mass function

interface OrderedDist(d:Type, a:Type, f:Type, given () (Dist d a f))
  cumulative : (d, a) -> LogSpace f   # the cumulative distribution function (CDF)
  survivor   : (d, a) -> LogSpace f   # the survivor function (complement of CDF)
  quantile   : (d, f) -> a            # the quantile function (inverse CDF)

'## Univariate probability distributions
Some commonly encountered univariate distributions.
### Bernoulli distribution
The [Bernoulli distribution](https://en.wikipedia.org/wiki/Bernoulli_distribution) is parameterised by its "success" probability, `prob`.

struct Bernoulli =
  prob: Float

instance Random(Bernoulli ,Bool)
  def draw(bern, k) =
    rand k < bern.prob

instance Dist(Bernoulli, Bool, Float)
  def density(bern, b) =
    if b
     then Exp $ log bern.prob
     else Exp $ log1p (-bern.prob)

instance OrderedDist(Bernoulli, Bool, Float)
  def cumulative(bern, b) =
    if b
      then Exp 0.0
      else Exp $ log1p (-bern.prob)
  def survivor(bern, b) =
    if b
      then Exp (-infinity)
      else Exp $ log bern.prob
  def quantile(bern, q) =
    q > (1 - bern.prob)


'### Binomial distribution
The [binomial distribution](https://en.wikipedia.org/wiki/Binomial_distribution) is parameterised by the number of trials, `trials`, and the success probability, `prob`. Mean is `trials*prob`.

struct Binomial =
  trials : Nat
  prob   : Float

instance Random(Binomial, Nat)
  def draw(d, k) =
    sum $ each (rand_vec d.trials (\k_. draw(a=Bool, Bernoulli(d.prob), k_)) k) b_to_n

instance Dist(Binomial, Nat, Float)
  def density(d, x) =
    if (d.prob == 0.0)
      then
        if (x == 0)
          then Exp 0.0
          else Exp (-infinity)
      else
        if (d.prob == 1.0)
          then
            if (x == d.trials)
              then Exp 0.0
              else Exp (-infinity)
          else
            tf = n_to_f d.trials
            xf = n_to_f x
            if (xf > tf)
              then Exp (-infinity)
              else Exp ( (xf * log d.prob) + ((tf - xf) * log1p (-d.prob)) +
                lgamma (tf + 1) - lgamma (xf + 1) - lgamma (tf + 1 - xf) )

instance OrderedDist(Binomial, Nat, Float)
  def cumulative(d, x) =
    xp1:Nat = x + 1
    ls_sum $ for i:(Fin xp1). density(f=Float, d, ordinal i)
  def survivor(d, x) =
    tmx = d.trials -| x
    ls_sum $ for i:(Fin tmx). density(f=Float, d, x + 1 + ordinal i)
  def quantile(d, q) =
    tp1:Nat = d.trials + 1
    lpdf = for i:(Fin tp1). ln $ density(f=Float, d, ordinal i)
    cdf = cdf_for_categorical lpdf
    mi = search_sorted cdf q
    ordinal $ from_just $ left_fence mi


'### Exponential distribution
The [exponential distribution](https://en.wikipedia.org/wiki/Exponential_distribution) is parameterised by its *rate*, `rate > 0`, so that the mean of the distribution is `1/rate`.

struct Exponential =
  rate : Float

instance Random(Exponential, Float)
  def draw(d, k) = log1p (-rand k) / -d.rate

instance Dist(Exponential, Float, Float)
  def density(d, x) =
    if (x < 0.0)
      then Exp (-infinity)
      else Exp $ log d.rate - (d.rate * x)

instance OrderedDist(Exponential, Float, Float)
  def cumulative(d, x) =
    if (x < 0.0)
      then Exp (-infinity)
      else Exp $ log1p (-exp (-d.rate * x))
  def survivor(d, x) =
    if (x < 0.0)
      then Exp 0.0
      else Exp $ -d.rate * x
  def quantile(d, q) = log1p (-q) / -d.rate


'### Geometric distribution
This [geometric distribution](https://en.wikipedia.org/wiki/Geometric_distribution) is defined as the number of trials until a success (**not** the number of failures). Parameterised by success probability, `prob`. Mean is `1/prob`.

struct Geometric =
  prob : Float

instance Random(Geometric, Nat)
  def draw(d, k) = f_to_n $ ceil $ log1p (-rand k) / log1p (-d.prob)

instance Dist(Geometric, Nat, Float)
  def density(d, n) =
    if (d.prob == 1)
      then
        if (n == 1)
          then Exp 0
          else Exp (-infinity)
      else
        nf = n_to_f n
        Exp $ ((nf-1)*log1p (-d.prob)) + log d.prob

instance OrderedDist(Geometric, Nat, Float)
  def cumulative(d, n) =
    nf = n_to_f n
    Exp $ log1p (-pow (1-d.prob) nf)
  def survivor(d, n) =
    Exp $ n_to_f n * log1p (-d.prob)
  def quantile(d, q) =
    f_to_n $ ceil $ log1p (-q) / log1p (-d.prob)


'### Normal distribution
The Gaussian, or [normal distribution](https://en.wikipedia.org/wiki/Normal_distribution) is parameterised by its *mean*, `loc`, and *standard deviation*, `scale`. ie. **not** variance or precision.

struct Normal =
  loc : Float
  scale : Float

instance Random(Normal, Float)
  def draw(d, k) = d.loc + (d.scale * randn k)

instance Dist(Normal, Float, Float)
  def density(d, x) =
    Exp $ -0.5 * (log (2 * pi * (sq d.scale)) + (sq ((x - d.loc) / d.scale)))

instance OrderedDist(Normal, Float, Float)
  def cumulative(d, x) =
    Exp $ log (0.5 * erfc ((d.loc - x) / (d.scale * sqrt(2.0))))
  def survivor(d, x) =
    Exp $ log (0.5 * erfc ((x - d.loc) / (d.scale * sqrt(2.0))))
  def quantile(_, _)  = todo  # Add `erfinv`.


'### Poisson distribution
The [Poisson distribution](https://en.wikipedia.org/wiki/Poisson_distribution) is parameterised by its rate, `rate > 0`. Mean is `rate`.

struct Poisson =
  rate : Float

instance Random(Poisson, Nat)
  def draw(d, k) =
    exp_neg_rate = exp (-d.rate)
    [current_k, next_k] = split_key(n=2, k)
    yield_state 0 \ans.
      with_state (rand current_k) \p. with_state next_k \k'.
        while \.
          if get p > exp_neg_rate
            then
              [ck, nk] = split_key(n=2, get k')
              p := (get p) * rand ck
              ans := (get ans) + 1
              k' := nk
              True
            else
              False

instance Dist(Poisson, Nat, Float)
  def density(d, n) =
    if ((d.rate == 0.0)&&(n == 0))
      then
        Exp 0.0
      else
        nf = n_to_f n
        Exp $ (nf * log d.rate) - d.rate - lgamma (nf + 1)

instance OrderedDist(Poisson, Nat, Float)
  def cumulative(d, x) =
    xp1:Nat = x + 1
    ls_sum $ for i:(Fin xp1). density(f=Float, d, ordinal i)
  def survivor(d, x) =
    xp1:Nat = x + 1
    cdf = ls_sum $ for i:(Fin xp1). density(f=Float, d, ordinal i)
    Exp $ log1p (-ls_to_f cdf)
  def quantile(d, q) =
    yield_state (0::Nat) \x.
      with_state 0.0 \cdf.
        while \.
          cdf := (get cdf) + ls_to_f (density(f=Float, d ,get x))
          if (get cdf) > q
            then
              False
            else
              x := (get x) + 1
              True


'### Uniform distribution
The [uniform distribution](https://en.wikipedia.org/wiki/Continuous_uniform_distribution) is continuous on an interval determined by a lower bound, `low`, and upper bound, `high > low`. Mean is `(low+high)/2`.

struct Uniform =
  low  : Float
  high : Float

instance Random(Uniform, Float)
  def draw(d, k) = d.low + ((d.high - d.low) * rand k)

instance Dist(Uniform, Float, Float)
  def density(d, x) =
    if (x >= d.low)&&(x <= d.high)
      then Exp $ -log (d.high - d.low)
      else Exp (-infinity)

instance OrderedDist(Uniform, Float, Float)
  def cumulative(d, x) =
    if (x < d.low)
      then Exp (-infinity)
      else if (x > d.high)
        then Exp 0.0
        else Exp $ log (x - d.low) - log (d.high - d.low)
  def survivor(d, x) =
    if (x < d.low)
      then Exp 0.0
      else if (x > d.high)
        then Exp (-infinity)
        else Exp $ log (d.high - x) - log (d.high - d.low)
  def quantile(d, q) = d.low + ((d.high - d.low) * q)


'## Data summaries
Some data summary functions. Note that `mean` is provided by the prelude.

'### Summaries for vector quantities

def mean_and_variance(xs:n=>a) -> (a, a) given (n|Ix, a|VSpace|Mul) =
  m = mean xs
  ss = sum $ each xs \x. sq(x - m)
  v = ss / (n_to_f (size n) - 1)
  (m, v)

def variance(xs:n=>a) -> a given (n|Ix, a|VSpace|Mul) =
  snd $ mean_and_variance xs

def std_dev(xs:n=>a) -> a given (n|Ix, a|VSpace|Mul|Floating) =
  sqrt $ variance xs

def covariance(xs:n=>a, ys:n=>a) -> a given (n|Ix, a|VSpace|Mul) =
  xm = mean xs
  ym = mean ys
  ss = sum for i:n. (xs[i] - xm) * (ys[i] - ym)
  ss / (n_to_f (size n) - 1)

def correlation(xs:n=>a, ys:n=>a) -> a
    given (n|Ix, a|VSpace|Mul|Floating|Fractional) =
  divide (covariance xs ys) (std_dev xs * std_dev ys)

'### Summaries for matrix quantities

def mean_and_variance_matrix(xs:n=>d=>a) -> (d=>a, d=>d=>a)
    given (n|Ix, d|Ix, a|Mul|VSpace) =
  xsMean:d=>a = (for i:d. sum for j:n. xs[j,i]) / n_to_f (size n)
  xsCov:d=>d=>a = (for i:d i':d. sum for j:n.
                         (xs[j,i'] - xsMean[i']) *
                         (xs[j,i]  - xsMean[i] ) ) / (n_to_f (size n) - 1)
  (xsMean, xsCov)

def variance_matrix(xs:n=>d=>a) -> d=>d=>a
    given (n|Ix, d|Ix, a|VSpace|Mul)=
  snd $ mean_and_variance_matrix xs

def correlation_matrix(xs:n=>d=>a) -> d=>d=>a
    given (n|Ix, d|Ix, a|Mul|VSpace|Floating|Fractional) =
  cvm = variance_matrix xs
  for i. for j. divide cvm[i,j] (sqrt cvm[i,i] * sqrt cvm[j,j])