Port https://github.com/ericlippert/probability ( Episode ~11) to F#

_Util.fs

module FsProb.Util

open System
open System.Threading
open System.Security.Cryptography

/// <summary>
/// 暗号強度を備えた、スレッドセーフで、すべて静的な乱数生成器（RNG）。
/// まだ使いやすい API とは言えません。もっと改善できるはずです。
/// </summary>
[<RequireQualifiedAccess>]
module BetterRandom =
    let private crng = new ThreadLocal<_>(fun () -> RandomNumberGenerator.Create())
    let private bytes = new ThreadLocal<_>(fun () -> Array.zeroCreate<byte> sizeof<int>)
    let private d = 1L <<< 52 |> float

    let nextInt () =
        crng.Value.GetBytes(bytes.Value)
        BitConverter.ToInt32(bytes.Value, 0) &&& Int32.MaxValue

    let rec nextDouble () =
        let upper = (nextInt () &&& 0x001FFFFF |> int64) <<< 31
        let lower = nextInt () |> int64
        let n = upper ||| lower |> double
        let q = n / d
        if q = 1.0 then nextDouble () else q

/// <summary>
/// スレッドセーフかつすべて静的で、暗号学的にランダム化された Random のラッパー。
/// まだ完璧とは言えないものの、わずかに改善されています。
/// </summary>
[<RequireQualifiedAccess>]
module PseudoRandom =
    let private prng = new ThreadLocal<_>(fun () -> Random(BetterRandom.nextInt ()))

    let nextInt () = prng.Value.Next()
    let nextDouble () = prng.Value.NextDouble()

let histogram low high items =
    let width = 40
    let height = 20
    let sampleCount = 100000

    let toBucketIndex x =
        (float width) * (x - low) / (high - low) |> int

    let buckets =
        items
        |> Seq.take sampleCount
        |> Seq.map toBucketIndex
        |> Seq.filter (fun k -> 0 <= k && k < width)
        |> Seq.countBy id
        |> Seq.sortBy fst

    let max = buckets |> Seq.map snd |> Seq.max
    let scale = if max < height then 1.0 else (float height) / (float max)

    let hist =
        seq {
            for r in 0..height ->
                if r < height then
                    let chars =
                        buckets
                        |> Seq.map (fun (_, b) -> if float b * scale > float (height - r) then '*' else ' ')

                    String.Join("", chars)
                else
                    String('-', width)
        }

    String.Join(Environment.NewLine, hist)

let rec gcd a b = if b = 0 then a else gcd b (a % b)

let gcdMany (numbers: int seq) = Seq.reduce gcd numbers

let lcm a b = a * b / (gcd a b)

let histogramD (d: 'a seq) =
    let sampleCount = 100000
    let width = 40

    let dict =
        d
        |> Seq.take sampleCount
        |> Seq.groupBy id
        |> Seq.map (fun (k, vs) -> k, Seq.length vs)
        |> Map.ofSeq

    let labelMax = dict |> Map.keys |> Seq.map _.ToString().Length |> Seq.max
    let toLabel (t: 'a) = t.ToString().PadLeft(labelMax)
    let sup = dict |> Map.keys |> Seq.sortBy toLabel |> Seq.toList
    let max = dict |> Map.values |> Seq.max
    let scale = if max < width then 1.0 else (double width) / (double max)

    let bar (t: 'a) =
        new string ('*', int (double dict[t] * scale))

    sup
    |> Seq.map (fun s -> $"{toLabel s}|{bar s}")
    |> fun xs -> String.Join(Environment.NewLine, xs)

Distribution.fs

module FsProb.Distribution

open System
open FsProb.Util

type IDistribution<'a> =
    abstract member Sample: unit -> 'a

type IWeightedDistribution<'a> =
    inherit IDistribution<'a>
    abstract member Weight: 'a -> double

type IDiscreteDistribution<'a> =
    inherit IDistribution<'a>
    abstract member Support: unit -> 'a seq
    abstract member Weight: 'a -> int

let inline sample<'a, 'b when 'a :> IDistribution<'b>> (d: 'a) = d.Sample()

let inline weight<'a, 'b, 'c when 'a :> IDistribution<'b> and 'a: (member Weight: 'b -> 'c)> (d: 'a) x = d.Weight(x)

let inline support<'a, 'b when 'a :> IDiscreteDistribution<'b>> (d: 'a) = d.Support()

let samples<'a> (d: IDistribution<'a>) =
    d |> Seq.unfold (fun s -> Some(sample s, s))

let scu =
    { new IWeightedDistribution<double> with
        member self.Sample() = PseudoRandom.nextDouble ()

        member self.Weight(x) =
            if 0.0 <= x && x < 1.0 then 1.0 else 0.0 }

let histogram low high (d: IDistribution<double>) = d |> samples |> histogram low high

let singleton<'a when 'a: equality> (v: 'a) =
    { new IDiscreteDistribution<'a> with
        member self.Sample() = v
        member self.Support() = seq { v }
        member self.Weight(x) = if x = v then 1 else 0 }

let sdu x y =
    let dist min max =
        { new IDiscreteDistribution<int> with
            member self.Support() = seq { min..max }

            member self.Sample() =
                let scale = 1.0 + float (max - min)
                int (scu.Sample() * scale) + min

            member self.Weight(i) = if min <= i && i <= max then 1 else 0 }

    match compare x y with
    | c when c < 0 -> dist x y
    | 0 -> singleton x
    | _ -> dist y x

type NormalDistribution(mean: double, sigma: double) =
    let stdSample () =
        // Box-Muller method
        let x = sample scu
        let y = sample scu
        sqrt (-2.0 * (log x)) * cos (2.0 * Math.PI * y)

    let piroot = 1.0 / sqrt (2.0 * Math.PI)

    member this.Mean = mean
    member this.Sigma = sigma

    interface IWeightedDistribution<double> with
        member this.Sample() = mean + sigma * stdSample ()

        member this.Weight(x) =
            exp (-(x - mean) * (x - mean) / (2.0 * sigma * sigma)) * piroot / sigma

let normal mean sigma = NormalDistribution(mean, sigma)

let histogramD (d: IDiscreteDistribution<'a>) = d |> samples |> histogramD

let showWeights (d: IDiscreteDistribution<'a>) =
    let labelMax = support d |> Seq.map (fun x -> x.ToString().Length) |> Seq.max
    let toLabel (t: 'a) = t.ToString().PadLeft(labelMax)
    String.Join(Environment.NewLine, support d |> Seq.map (fun s -> $"{toLabel s}:{weight d s}"))

let empty<'a> =
    { new IDiscreteDistribution<'a> with
        member self.Sample() =
            failwith "Cannot sample from empty distribution"

        member self.Support() = Seq.empty
        member self.Weight(x) = 0 }

let bernoulli zero one =
    let dist zero one =
        { new IDiscreteDistribution<int> with
            member self.Support() =
                seq {
                    0
                    1
                }

            member self.Sample() =
                let zeroRatio = float zero / float (zero + one)
                if scu.Sample() <= zeroRatio then 0 else 1

            member self.Weight(i) =
                match i with
                | 0 -> zero
                | 1 -> one
                | _ -> 0 }

    match zero, one with
    | x, y when x < 0 && y < 0 -> failwith "Invalid Argument"
    | 0, 0 -> empty
    | 0, _ -> singleton 1
    | _, 0 -> singleton 0
    | _ ->
        let d = gcd zero one
        dist (zero / d) (one / d)

let projected (underlying: IDiscreteDistribution<'a>) (projection: 'a -> 'r) =
    let result =
        let weights =
            support underlying
            |> Seq.groupBy projection
            |> Seq.map (fun (k, vs) -> k, Seq.sumBy (weight underlying) vs)
            |> Map.ofSeq

        { new IDiscreteDistribution<'r> with
            member this.Sample() = projection (sample underlying)

            member this.Weight(x) =
                weights |> Map.tryFind x |> Option.defaultValue 0

            member this.Support() = Map.keys weights }

    match support result |> Seq.truncate 2 |> Seq.toList with
    | [] -> empty
    | [ x ] -> singleton x
    | _ -> result

let map<'a, 'r when 'r: comparison> (mapping: 'a -> 'r) (d: IDiscreteDistribution<'a>) =
    let weights =
        support d
        |> Seq.groupBy mapping
        |> Seq.map (fun (k, vs) -> k, Seq.sumBy (weight d) vs)
        |> Map.ofSeq

    let projected =
        { new IDiscreteDistribution<'r> with
            member this.Sample() = mapping (sample d)

            member this.Weight(x) =
                weights |> Map.tryFind x |> Option.defaultValue 0

            member this.Support() = Map.keys weights }

    let s = support projected

    match Seq.length s with
    | 0 -> empty
    | 1 -> singleton (Seq.head s)
    | _ -> projected

let toUniform<'a when 'a: comparison> (items: 'a seq) =
    let array = Seq.toArray items
    sdu 0 (Array.length array - 1) |> map (fun i -> array[i])

[<Struct>]
type private WithWeight = { Index: int; Weight: int }

[<Struct>]
type private Bucket =
    | Deterministic of WithWeight
    | Split of WithWeight * WithWeight

let weightedInteger (weights: int seq) =
    let buildAliasTable (weights: int seq) =
        let count = Seq.length weights
        let sum = Seq.sum weights

        let exacts, lows, highs =
            (([], [], []), Seq.mapi (fun i x -> i, x) weights)
            ||> Seq.fold (fun (es, ls, hs) (index, weight) ->
                let scaled = weight * count

                if scaled = sum then
                    ({ Index = index; Weight = scaled } :: es), ls, hs
                else if scaled < sum then
                    es, ({ Index = index; Weight = scaled } :: ls), hs
                else
                    es, ls, ({ Index = index; Weight = scaled } :: hs))

        let exactBuckets = exacts |> List.map Deterministic

        let rec loop (buckets: Bucket list) (lows: WithWeight list) (highs: WithWeight list) =
            match lows, highs with
            | [], [] -> buckets
            | low :: lowLeft, high :: highLeft ->
                let lowNeeds = sum - low.Weight

                let alias =
                    { Index = high.Index
                      Weight = lowNeeds }

                let remainingHigh =
                    { Index = high.Index
                      Weight = high.Weight - lowNeeds }

                let lowBucket = Split(low, alias)

                if remainingHigh.Weight = sum then
                    let highBucket = Deterministic remainingHigh
                    loop (highBucket :: lowBucket :: buckets) lowLeft highLeft
                else if remainingHigh.Weight < sum then
                    loop (lowBucket :: buckets) (remainingHigh :: lowLeft) highLeft
                else
                    loop (lowBucket :: buckets) lowLeft (remainingHigh :: highLeft)
            | _ -> failwith ""

        loop exactBuckets lows highs

    if Seq.exists (fun x -> x < 0) weights || not (Seq.exists (fun x -> x > 0) weights) then
        failwith ""
    else
        match weights |> Seq.truncate 3 |> Seq.toList with
        | [] -> empty
        | [ _ ] -> singleton 0
        | [ weight0; weight1 ] -> bernoulli weight0 weight1
        | _ ->
            let d = gcdMany weights
            let reducedWeights = Seq.map (fun x -> x / d) weights |> Seq.toArray

            let distributions =
                reducedWeights
                |> buildAliasTable
                |> List.map (function
                    | Deterministic x -> singleton x.Index
                    | Split(main, alias) ->
                        bernoulli main.Weight alias.Weight
                        |> map (fun x -> if x = 0 then main.Index else alias.Index))
                |> List.toArray

            { new IDiscreteDistribution<int> with
                member this.Sample() =
                    let randomIndex = sdu 0 (reducedWeights.Length - 1) |> sample
                    distributions[randomIndex] |> sample

                member this.Support() =
                    seq { 0 .. reducedWeights.Length - 1 }
                    |> Seq.filter (fun x -> reducedWeights[x] <> 0)

                member this.Weight(x) =
                    if 0 <= x && x < reducedWeights.Length then
                        reducedWeights[x]
                    else
                        0 }

let toWeighted items weights =
    let list = Seq.toArray items
    weightedInteger weights |> map (fun i -> list[i])

let filter<'a when 'a: comparison> (predicate: 'a -> bool) (d: IDiscreteDistribution<'a>) =
    let s = support d |> Seq.filter predicate |> Seq.toList
    let ws = s |> List.map (fun t -> weight d t)
    toWeighted s ws

let rec projected'<'a, 'r when 'r: comparison> (underlying: IDiscreteDistribution<'a>) (projection: 'a -> 'r) =
    let result =
        let weights =
            support underlying
            |> Seq.groupBy projection
            |> Seq.map (fun (k, vs) -> k, Seq.sumBy (weight underlying) vs)
            |> Map.ofSeq

        { new IDiscreteDistribution<'r> with
            member this.Sample() = projection (sample underlying)

            member this.Weight(x) =
                weights |> Map.tryFind x |> Option.defaultValue 0

            member this.Support() = Map.keys weights }

    match support result |> Seq.truncate 2 |> Seq.toList with
    | [] -> empty
    | [ x ] -> singleton x
    | _ -> result

and map'<'a, 'r when 'r: comparison> (projection: 'a -> 'r) (d: IDiscreteDistribution<'a>) =
    let rs =
        support d
        |> Seq.groupBy projection
        |> Seq.mapi (fun i (k, vs) -> i, (k, Seq.sumBy (weight d) vs))
        |> Map.ofSeq

    let weighted = weightedInteger (rs |> Map.values |> Seq.map snd)
    projected' weighted (fun i -> rs[i] |> fst)