GenericCategoricals.jl

using Distributions
using Random: AbstractRNG
using StatsBase: proportionmap
using Test

# Generic categoricals are also known as empirical measures.
struct GenericCategorical{T}
  mapping :: Dict{T, Float64}
  values  :: Vector{T}
  probs   :: Vector{Float64}
  function GenericCategorical(weights::Dict{T, Float64}) where T
    mapping = Dict(k => v/sum(values(weights)) for (k,v) in weights)
    new{T}(mapping, collect(keys(mapping)), collect(values(mapping)))
  end
end
GenericCategorical(xs::AbstractVector) = GenericCategorical(proportionmap(xs))

# We implement the simplest interface for distributions, consisting of
# random sampling and the evaluation of probabilities.
Distributions.logpdf(gc::GenericCategorical, x) = log(gc.mapping[x])
function Base.rand(rng::AbstractRNG, gc::GenericCategorical)
  i = rand(rng, Categorical(gc.probs))
  gc.values[i]
end

@testset "Generic categorical" begin
  gc = GenericCategorical([1, 1, 1, 2, 2, 3])
  @test rand(gc) in [1, 2, 3]
  @test rand(gc) isa Int
  @test logpdf(gc, 1) > logpdf(gc, 2) > logpdf(gc, 3)
end