Probabilistic functional programming F#

pfp.fs

/// A probability is a number between 0 and 1 inclusive.
type Prob = float

/// A distribution is a sequence of outcome-probability pairs such that the
/// probabilities sum to 1.
type Dist<'a> = D of seq<'a * Prob>

/// A spread takes a sequence of elements and assigns probabilities.
type Spread<'a> = 'a list -> Dist<'a>


/// An event is a subset of a sequence of items (outcomes).
type Event<'a> = 'a -> bool

/// Operations on events.
module Event =
  
  /// Creates an event consisting of the specified outcomes.
  let oneOf (c:'a list) : Event<'a> =
    fun a -> List.exists ((=) a) c

  let just (a:'a) : Event<'a> =
    oneOf [a]


/// Operations on distributions.
module Dist =

  [<AutoOpen>]
  module private Util =

    /// Returns the cartesian product of two sequences.
    let cartesian (a:seq<'a>) (b:seq<'b>) : seq<'a * 'b> =
      Seq.collect (fun a -> b |> Seq.map (fun b -> a,b)) a

    let rec iterate (f:'a -> 'a) (z:'a) : seq<'a> =
      seq {
        yield z
        yield! iterate f (f z) }


  let private unD (D(xs)) = xs

  let private onD (f:seq<'a * Prob> -> seq<'a * Prob>) : Dist<'a> -> Dist<'a> =
    unD >> f >> D
          
  let isZero (d:Dist<'a>) : bool =
    Seq.isEmpty (unD d)

  let extract (d:Dist<'a>) : 'a seq =
    unD d |> Seq.map fst

  /// Creates a spread which creates distributions by pairing probabilities with outcomes, pointwise.
  let enum (ps:Prob list) (os:'a list) : Dist<'a> =
    List.zip os ps |> List.toSeq |> D

  /// Creates a spread which creates distributions with probabilities based on the relative values
  /// of the elements of the argument list.
  let relative (xs:int list) : Spread<'a> =
    let s = float <| List.sum xs in
    enum (xs |> List.map (fun n -> (float n) / s))
    
  /// Computes the probability of an event given a distribution.
  let prob (e:Event<'a>) (d:Dist<'a>) : Prob =
    d |> unD |> Seq.filter (fst >> e) |> Seq.map snd |> Seq.sum

  let map (f:'a -> 'b) (da:Dist<'a>) : Dist<'b> =
    unD da |> Seq.map (fun (a,p) -> f a,p) |> D

  /// Joins two independent distributions by taking all pairs of outcomes and 
  /// multiplying their probabilities.
  let zipWith2 (f:'a -> 'b -> 'c) (da:Dist<'a>) (db:Dist<'b>) : Dist<'c> =
    cartesian (unD da) (unD db) 
    |> Seq.map (fun ((a,p),(b,q)) -> f a b, p * q)
    |> D
  
  /// Returns a dependent distribution.
//  let bind (f:'a -> Dist<'b>) (da:Dist<'a>) : Dist<'b> =
//    (unD da)
//    |> Seq.collect (fun (a,p) -> unD (f a) |> Seq.map (fun (b,q) -> b, p * q)) 
//    |> D

  let bind (f:'a -> Dist<'b>) (da:Dist<'a>) : Dist<'b> =
    D <| seq { 
      for (a,p) in (unD da) do 
        for (b,q) in (unD (f a)) -> (b, p * q) }

  /// Unfolds a distribution of distributions into a distribution.
  /// Given a distribution where the outcomes are themselves distributions,
  /// we can derive the distribution of outcomes of these distributions.
  let join (d:Dist<Dist<'a>>) : Dist<'a> =
    bind id d

  let zipWith (f:'a -> 'b -> 'c) (da:Dist<'a>) (db:Dist<'b>) : Dist<'c> =
    bind (fun b -> da |> map (fun a -> f a b)) db

  let prod (da:Dist<'a>) (db:Dist<'b>) : Dist<'a * 'b> =
    zipWith (fun a b -> a,b) da db
  
  /// Creates a distribution with one outcome with probability 1.0.  
  let certainly (a:'a) : Dist<'a> =
    D [(a,1.0)]

  /// Returns a distribution describing an impossible outcome.
  /// - bind impossible f = impossible
  /// - seqn x impossible = impossible
  let impossible<'a> : Dist<'a> =
    D []
    
  let compose (f:'a -> Dist<'b>) (g:'b -> Dist<'c>) : 'a -> Dist<'c> =
    f >> bind g

  let seqn (da:Dist<'a>) (db:Dist<'b>) : Dist<'b> =
    da |> bind (fun _ -> db)
        
  let sequ (fs:seq<'a -> Dist<'a>>) : 'a -> Dist<'a> =
    Seq.fold compose certainly fs

  /// choose x impossible = x
  /// choose impossible x = x
  let mplus (da:Dist<'a>) (db:Dist<'a>) : Dist<'a> =
//    match isZero da, isZero db with
//    | true, true -> impossible
//    | true, _ -> db
//    | _, true -> da
//    | _ -> failwith ""
    Seq.append (unD da) (unD db) 
    |> Seq.groupBy fst
    |> Seq.map (fun (a,ps) -> a, ps |> Seq.map snd |> Seq.sum)
    |> D

  /// Creates a uniform distribution consisting of two outcomes where the first
  /// has the specified probability and the other the complement.
  let choose (p:Prob) (a:'a) (b:'a) : Dist<'a> =
    enum [ p ; 1.0 - p ] [ a ; b ]

  /// Returns the probability of truth, given a distribution of booleans.
  let truth (d:Dist<bool>) : Prob =
    let (b,p) = unD d |> Seq.item 0 in
    if b then p else 1.0 - p

  /// Computes the probability of truth with respect to the boolean distribution,
  /// and then creates a distribution where the outcomes are the given distributions
  /// such that the first has the probability of truh. Then unfolds/joins the resulting
  /// distribution.
  let cond (d:Dist<bool>) (da:Dist<'a>) (db:Dist<'a>) : Dist<'a> =
    let p = truth d in
    join (choose p da db)
    
  let inline accumBy (key:'a -> _) (ls:('a * 'b) list) : ('a * 'b) list =
    let rec go key ls =
      match ls with    
      | (x,p)::((y,q)::xs as ys) ->
        if (key x) = (key y) then go key ((x,p+q)::xs)
        else (x,p)::go key ys
      | _ -> ls
    go key ls

  let private onSeq (f:'a list -> 'a list) : 'a seq -> 'a seq =
    Seq.toList >> f >> List.toSeq

  /// Normalizes a distribution by accumulating outcumes by the specified key.
  let normBy (k:'a -> _) : Dist<'a> -> Dist<'a> =
    onD (onSeq (accumBy k >> List.sort))
    
  let norm d = normBy id d

  /// Scales a list of outcome probability pairs such that the 
  /// probabilities sum to 1.0.
  let scale (xs:seq<'a * Prob>) : Dist<'a> =
    let q = xs |> Seq.sumBy snd in
    D (xs |> Seq.map (fun (a,p) -> a, p/q))

  /// Creates a spread which creates distributions with the following shape.
  let shape (f:float -> float) : Spread<'a> =
    function
    | [] -> impossible
    | xs ->
      let incr = 1.0 / float (List.length xs - 1)
      let ps = Seq.map f (iterate ((+) incr) 0.0)
      scale (Seq.zip xs ps)
      
  /// Removes outcomes which return false and scales.
  let filter (f:'a -> bool) (da:Dist<'a>) : Dist<'a> =
    unD da |> Seq.filter (fst >> f) |> scale

  /// Creates a spread where the distribution follows a linear shape.
  let linear (c:float) : Spread<'a> =
    shape ((*) c)

  /// A uniform spread creates distributions with a constant shape.
  let uniformSpread<'a> : Spread<'a> =
    shape (fun _ -> 1.0)

  /// Creates a curve with the following mean and standard deviation.
  /// 
  let normalCurve (mean:float) (stdev:float) (x:float) : float =
    let u = (x - mean) / stdev in
    1.0 / sqrt (2.0 * System.Math.PI) * exp (float (-1/2) * pown u 2)

  /// Returns a normal distribution given a set of outcomes.
  let normal<'a> : Spread<'a> =
    shape (normalCurve 0.5 0.5)      

  /// Creates a uniform distribution.
  let uniform : Spread<'a> =
    fun (xs:'a list) ->
      let p = 1.0 / float (List.length xs)
      D (xs |> Seq.map (fun a -> a,p))

  let bernoulli (hd:'a) (tl:'a) : Dist<'a> =
    uniform [hd;tl]

  /// The distribution representing a 6-sided die.
  let die : Dist<int> = 
    uniformSpread [1..6]

  /// Returns the distribution of rolling N dice.
  /// Proceeds by joining the distribution of a 6-sided die with 
  /// the distribution of rolling N-1 dice where the distribution of
  /// rolling nice dice is an empty outcome with probability 1.0.
  let rec dice (n:int) : Dist<int list> =
    match n with
    | 0 -> certainly []
    | n -> zipWith (fun a b -> a::b) die (dice (n - 1))

  /// Returns a distribution where the outcomes are selections of individual
  /// elements from the specified list along with the list without the selected element.
  let selectOne (c:'a list) : Dist<'a * 'a list> =
    uniform [ for v in c -> (v, List.except [v] c) ]

  /// Returns a distribution where outcomes are selections of N elements from the list.
  /// Proceeds by selecting a single element and selecting N-1 elements from the remaining
  /// elements.
  let rec selectMany (n:int) (c:'a list) : Dist<'a list * 'a list> =
    match n with
    | 0 -> certainly ([],c)
    | n -> 
      selectOne c 
      |> bind (fun (x,c1) -> selectMany (n-1) c1 |> map (fun (xs,c2) -> x::xs,c2))
  
  let select (n:int) (c:'a list) : Dist<'a list> =
    selectMany n c |> map (fst >> List.rev)

  

  // statistical analyses
  //

  /// Returns the expected value for the distribution given an expected value
  /// for the outcomes.
  let expectedDist (expected:'a -> float) (d:Dist<'a>) : float =
    (unD d) |> Seq.map (fun (x,p) -> expected x * p) |> Seq.sum

  let variance (expected:'a -> float) (da:Dist<'a>) : float =
    let ps = unD da
    let sq x = x * x
    let ex = expectedDist expected da
    ps |> Seq.map (fun (x,p) -> p * sq (expected x - ex)) |> Seq.sum

  let stdev ex d = variance ex d |> sqrt
  

    

  
(*
let p1 = Dist.prob (Seq.filter ((=) 6) >> Seq.length >> ((<=) 2)) (Dist.dice 4)
printfn "p=%A" p1


type Marble = R | G | B

let p2 = Dist.prob ((=) [R;G;B]) (Dist.select 3 [R;R;G;G;B])
printfn "p=%A" p2
*)