Substantial refactoring and cleanup, type parameters change

Author: oschulz

Project.toml

@@ -3,6 +3,7 @@ uuid = "0bbb1fad-0f24-45fe-94a4-415852c5cc3b"
 version = "0.2.1"
 
 [deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArraysOfArrays = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -11,13 +12,15 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [compat]
+Adapt = "1.0, 2.0"
 ArraysOfArrays = "0.4, 0.5"
 Distributions = "0.21.3, 0.22, 0.23, 0.24"
 StatsBase = "0.32, 0.33"
 julia = "1"
 
 [extras]
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["ForwardDiff", "Test"]

README.md

@@ -7,8 +7,8 @@
 [![Codecov](https://codecov.io/gh/oschulz/EmpiricalDistributions.jl/branch/master/graph/badge.svg)](https://codecov.io/gh/oschulz/EmpiricalDistributions.jl)
 
 A Julia package for empirical probability distributions. Currently
-implements uni- and multivariate binned distributions, backed by
-[StatsBase.jl](https://github.com/JuliaStats/StatsBase.jl) histograms.
+implements uni- and multivariate binned distributions that can be created
+from [StatsBase.jl](https://github.com/JuliaStats/StatsBase.jl) histograms.
 
 
 ## Documentation

docs/src/index.md

@@ -3,12 +3,12 @@
 A Julia package for empirical probability distributions that implement the
 [Distributions.jl](https://github.com/JuliaStats/Distributions.jl) API.
 
-This package currently provides uni- and multivariate binned distributions,
-backed by [StatsBase.jl](https://github.com/JuliaStats/StatsBase.jl)
-histograms.
+This package currently provides uni- and multivariate binned distributions
+that can be created from
+[StatsBase.jl](https://github.com/JuliaStats/StatsBase.jl) histograms.
 
-[`UvBinnedDist`](@ref) wraps a 1-dimensional histogram and presents it as
-a (binned) univariate continuous distribution:
+[`UvBinnedDist`](@ref), usually derived from a 1-dimensional histogram,
+represents a binned univariate continuous distribution:
 
 ```julia
 using Distributions, StatsBase
@@ -31,15 +31,21 @@ maximum(uvdist), minimum(uvdist)
 rand(uvdist, 5)
 ```
 
-[`MvBinnedDist`](@ref) does the same for a multi-dimensional histogram,
-and presents it as a (binned) multivariate continuous distribution:
+A binned distribution can be converted back to a histogram:
+
+```julia
+using LinearAlgebra
+Histogram(uvdist) == normalize(uvhist)
+```
+
+
+[`MvBinnedDist`](@ref), usually derived from a multidimensional histogram,
+represents a binned multivariate continuous distribution:
 
 ```julia
 X_mv = rand(MvNormal([3.5, 0.5], [2.0 0.5; 0.5 1.0]), 10^5)
 mvhist = fit(Histogram, (X_mv[1, :], X_mv[2, :]))
 
-using Distributions, EmpiricalDistributions
-
 mvdist = MvBinnedDist(mvhist)
 mvdist isa Distribution{Multivariate,Continuous}
 

src/EmpiricalDistributions.jl

@@ -8,12 +8,13 @@ using LinearAlgebra
 using Random
 using Statistics
 
+using Adapt
 using ArraysOfArrays
 using Distributions
 using StatsBase
 
 
-include("hist_funcs.jl")
+include("utils.jl")
 include("uv_binned_dist.jl")
 include("mv_binned_dist.jl")
 

src/hist_funcs.jl

@@ -2,26 +2,38 @@
 
 
 """
-    MvBinnedDist <: Distribution{Multivariate,Continuous}
+    MvBinnedDist <: StatsBase.Distribution{Multivariate,Continuous}
 
-Wraps a multi-dimensional histograms and presents it as a binned multivariate
-distribution.
+A binned multivariate distribution, usually derived from a histogram.
 
-Constructor:
+Constructors:
 
-    MvBinnedDist(h::Histogram{<:Real,N})
+```julia
+    MvBinnedDist(h::StatsBase.Histogram{<:Real,N})
+    MvBinnedDist{T<:Real}(h::StatsBase.Histogram{<:Real,N})
+```
+
+You can convert a `MvBinnedDist` back to a histogram via 
+
+```julia
+convert(StatsBase.Histogram, dist::MvBinnedDist)
+StatsBase.Histogram(dist::MvBinnedDist)
+```
 """
 struct MvBinnedDist{
     T <: Real,
     N,
-    H <: Histogram{<:Real, N},
+    U <: Real,
+    ET <: NTuple{N,AbstractVector{<:Real}},
     VT <: AbstractVector{T},
-    MT <: AbstractMatrix{T}
+    MT <: AbstractMatrix{T},
+    VU <: AbstractVector{U},
+    AU <: AbstractArray{U,N}
 } <: Distributions.Distribution{Multivariate,Continuous}
-    hist::H
-    _edges::NTuple{N, <:AbstractVector{T}}
-    _cart_inds::CartesianIndices{N, NTuple{N, Base.OneTo{Int}}}
-    _probability_edges::VT
+    _edges::ET
+    _bin_pdf::AU
+    _bin_linidx_cdf::VU
+    _closed_left::Bool
     _mean::VT
     _mode::VT
     _var::VT
@@ -31,64 +43,75 @@ end
 export MvBinnedDist
 
 
-function MvBinnedDist(h::StatsBase.Histogram{<:Real, N}, T::DataType = Float64) where {N}
+function MvBinnedDist{T}(h::Histogram{<:Real}) where {T<:Real}
     nh = normalize(h)
 
-    probabilty_widths = nh.weights * inv(sum(nh.weights))
-    probabilty_edges::Vector{T} = Vector{Float64}(undef, length(h.weights) + 1)
-    probabilty_edges[1] = 0
-    for (i, w) in enumerate(probabilty_widths)
-        v = probabilty_edges[i] + probabilty_widths[i]
-        probabilty_edges[i+1] = v > 1 ? 1 : v
-    end
+    edges = nh.edges
+    bin_pdf = nh.weights 
+
+    closed_left = nh.closed == :left
+
+    Y = nh.weights
+    X = _bin_centers.(nh.edges)
+    W = _bin_widths.(nh.edges)
+
+    bin_linidx_cdf = cumsum(broadcast(idx -> Y[idx] .* prod(map(getindex, W, idx.I)), vec(CartesianIndices(Y))))
+    @assert last(bin_linidx_cdf) ≈ 1
+    bin_linidx_cdf[end] = 1
 
-    mean_est = _mean(h)
-    mode_est = _mode(nh)
-    var_est = _var(h, mean = mean_est)
-    cov_est = _cov(h, mean = mean_est)
+    mean_est_tpl = _mean(nh)
+    mean_est = [T.(mean_est_tpl)...]
+    mode_est = [T.(_mode(nh))...]
+    var_est = [T.(_var(nh, mean_est_tpl))...]
+    cov_est = T.(_cov(nh, mean_est_tpl))
 
     return MvBinnedDist(
-        nh,
-        collect.(nh.edges),
-        CartesianIndices(nh.weights),
-        probabilty_edges,
-        mean_est,
-        mode_est,
-        var_est,
-        cov_est
+        edges, bin_pdf, bin_linidx_cdf, closed_left,
+        mean_est, mode_est, var_est, cov_est
     )
 end
 
+MvBinnedDist(h::Histogram{<:Real}) = MvBinnedDist{float(promote_type(map(eltype, h.edges)...))}(h)
 
-Base.convert(::Type{Histogram}, d::MvBinnedDist) = d.hist
+
+function Adapt.adapt_structure(to, d::MvBinnedDist)
+    MvBinnedDist(
+        map(e -> adapt(to, e), d._edges), adapt(to, d._bin_pdf), adapt(to, d._bin_linidx_cdf),
+        adapt(to, d._closed_left), adapt(to, d._mean), adapt(to, d._mode), adapt(to, d._var), adapt(to, d._cov)
+    )
+end
 
 
-Base.length(d::MvBinnedDist{T, N}) where {T, N} = N
-Base.size(d::MvBinnedDist{T, N}) where {T, N} = (N,)
-Base.eltype(d::MvBinnedDist{T, N}) where {T, N} = T
+Histogram(d::MvBinnedDist) = Histogram(map(Array, d._edges), Array(d._bin_pdf), (d._closed_left ? :left : :right), true)
+Base.convert(::Type{Histogram}, d::MvBinnedDist) = Histogram(d)
 
-Statistics.mean(d::MvBinnedDist{T, N}) where {T, N} = d._mean
-StatsBase.mode(d::MvBinnedDist{T, N}) where {T, N} = d._mode
-Statistics.var(d::MvBinnedDist{T, N}) where {T, N} = d._var
-Statistics.cov(d::MvBinnedDist{T, N}) where {T, N} = d._cov
 
+Base.length(d::MvBinnedDist{T,N}) where {T,N} = N
+Base.size(d::MvBinnedDist{T,N}) where {T,N} = (N,)
+Base.eltype(d::MvBinnedDist{T,N}) where {T,N} = T
 
-function Distributions._rand!(r::AbstractRNG, d::MvBinnedDist{T,N}, A::AbstractVector{<:Real}) where {T, N}
-    rand!(r, A)
-    next_inds::UnitRange{Int} = searchsorted(d._probability_edges::Vector{T}, A[1]::T)
-    cell_lin_index::Int = min(next_inds.start, next_inds.stop)
-    cell_car_index = d._cart_inds[cell_lin_index]
-    for idim in Base.OneTo(N)
-        i = cell_car_index[idim]
-        sub_int = d._edges[idim][i:i+1]
-        sub_int_width::T = sub_int[2] - sub_int[1]
-        A[idim] = sub_int[1] + sub_int_width * A[idim]
+Statistics.mean(d::MvBinnedDist) = d._mean
+StatsBase.mode(d::MvBinnedDist) = d._mode
+Statistics.var(d::MvBinnedDist) = d._var
+Statistics.cov(d::MvBinnedDist) = d._cov
+
+
+function Distributions._rand!(rng::AbstractRNG, d::MvBinnedDist{T,N}, A::AbstractVector{<:Real}) where {T,N}
+    @assert length(eachindex(A)) == N
+    u = rand(rng)
+    i = searchsortedfirst(d._bin_linidx_cdf, u)
+    idx_lo = CartesianIndices(d._bin_pdf)[i]
+    idx_hi = idx_lo + CartesianIndex(1, 1)
+    x_lo = map(getindex, d._edges, idx_lo.I)
+    x_hi = map(getindex, d._edges, idx_hi.I)
+    for i in 1:N
+        A[i] = _rand_uniform(rng, T, x_lo[i], x_hi[i])
     end
     return A
 end
 
-function Distributions._rand!(r::AbstractRNG, d::MvBinnedDist{T,N}, A::AbstractMatrix{<:Real}) where {T, N}
-    Distributions._rand!.((r,), (d,), nestedview(A))
+function Distributions._rand!(rng::AbstractRNG, d::MvBinnedDist{T,N}, A::AbstractMatrix{<:Real}) where {T,N}
+    Distributions._rand!.(Ref(rng), (d,), nestedview(A))
     return A
 end
 
@@ -104,17 +127,24 @@ end
 end
 
 
-function Distributions.pdf(d::MvBinnedDist{T,N}, x::AbstractVector{<:Real}) where {T,N}
+function Distributions.pdf(d::MvBinnedDist{T,N,U}, x::AbstractVector{<:Real}) where {T,N,U}
     length(eachindex(x)) == N || throw(ArgumentError("Length of variate doesn't match dimensionality of distribution"))
-    x_tpl = _unsafe_unroll_tuple(x, Val(N))
-    _pdf(d.hist, x_tpl)
+    xs = _unsafe_unroll_tuple(x, Val(N))
+
+    idxs = _find_bin(d._edges, d._closed_left, xs)
+    if checkbounds(Bool, d._bin_pdf, idxs...)
+        @inbounds r = d._bin_pdf[idxs...]
+        convert(U, r)
+    else
+        zero(U)
+    end
 end
 
 
-function Distributions.logpdf(d::MvBinnedDist{T, N}, x::AbstractArray{<:Real, 1}) where {T, N}
+function Distributions.logpdf(d::MvBinnedDist{T,N}, x::AbstractArray{<:Real, 1}) where {T,N}
     return log(pdf(d, x))
 end
 
-function Distributions._logpdf(d::MvBinnedDist{T,N}, x::AbstractArray{<:Real, 1}) where {T, N}
+function Distributions._logpdf(d::MvBinnedDist{T,N}, x::AbstractArray{<:Real, 1}) where {T,N}
     return logpdf(d, x)
 end