Implement KeyedDistribution and KeyedSampleable #1

bencottier · 2021-03-05T19:08:38Z

This proof of concept for KeyedDistributions extends the AxisKeys.jl ecosystem with KeyedSampleable and KeyedDistribution types. The types thinly wrap Sampleables with a vector of keys, corresponding to the variates of the Sampleable. This is analogous to KeyedArray wrapping an array.

Not implemented:

NamedDims support Support named dims #2
Callable lookup syntax Support lookup syntax (round brackets) #3
Methods to marginalize a distribution Add marginalize method for certain distributions #4
MixtureModel support Support MixtureModels #5

rofinn · 2021-03-05T19:18:07Z

Callable lookup syntax (not sure what this would do)

I feel like that would likely just return the same thing as marginalise?

glennmoy · 2021-03-08T12:29:00Z

Should Sampleable have all the statistical functions e.g. mean, in addition to Distribution? This doesn't mention it https://juliastats.org/Distributions.jl/latest/extends/#Create-a-Distribution

To subtype Sampleable I believe you only have to implement rand (univariate) or length and _rand! (multivariate).

src/KeyedDistributions.jl

test/runtests.jl

glennmoy · 2021-03-08T12:51:37Z

One of the equality tests is failing, so there may be a problem with my struct implementation

Has it anything to do with using ==() instead of isequal(), or constructing the KeyedDistribution inside the function?

glennmoy · 2021-03-08T12:57:49Z

Callable lookup syntax (not sure what this would do)

I feel like that would likely just return the same thing as marginalise?

Yeah I think this means extracting the relevant elements by calling the keyed distribution like a function (as AxisKeys can do)

julia> A = KeyedArray([0.1, 0.2, 0.3, 0.4, 0.5], :obj=>[:a, :b, :c, :d, :e])

julia> d = MvNormal(A)

julia> kd = KeyedDistribution(d)

julia> kd(obj=[:a, :c, :d])  # returns another KD marginalised on [a, c, d]

glennmoy · 2021-03-08T12:58:45Z

I think we can open issues for the remaining tasks and in this PR nail down the Distributions API at least.

What's your sense of the time/effort this would take?

Co-authored-by: Glenn Moynihan <glenn.moynihan@invenialabs.co.uk>

bencottier · 2021-03-08T18:53:33Z

Should Sampleable have all the statistical functions e.g. mean, in addition to Distribution? This doesn't mention it https://juliastats.org/Distributions.jl/latest/extends/#Create-a-Distribution

To subtype Sampleable I believe you only have to implement rand (univariate) or length and _rand! (multivariate).

I thought so, it's just that IndexedDistributions implemented the statistical functions for both, so I wonder if it's expected somewhere in our code.

Also I believe we have to overload rand regardless to return a KeyedArray. It was a regular array without.

bencottier · 2021-03-08T18:58:52Z

Callable lookup syntax (not sure what this would do)

I feel like that would likely just return the same thing as marginalise?

Yeah I think this means extracting the relevant elements by calling the keyed distribution like a function (as AxisKeys can do)
julia> A = KeyedArray([0.1, 0.2, 0.3, 0.4, 0.5], :obj=>[:a, :b, :c, :d, :e])

julia> d = MvNormal(A)

julia> kd = KeyedDistribution(d)

julia> kd(obj=[:a, :c, :d])  # returns another KD marginalised on [a, c, d]

That makes sense, just note it only corresponds to marginalising for some distributions (I think in the exponential family e.g. normal, t distribution). With this syntax, we could have a marginalize method for clarity that simply uses it under the hood.

bencottier · 2021-03-08T19:01:15Z

I think we can open issues for the remaining tasks and in this PR nail down the Distributions API at least.

What's your sense of the time/effort this would take?

I estimate at most 1 day spent on this.

Just those specified in Distributions.jl docs

…butions.jl into bc/prototype

bencottier · 2021-03-09T13:32:58Z

One of the equality tests is failing, so there may be a problem with my struct implementation

Has it anything to do with using ==() instead of isequal(), or constructing the KeyedDistribution inside the function?

EDIT: the fix was to specify F<:VariateForm, S<:ValueSupport.

If I do kd2 = KeyedDistribution(d, keys); kd == kd2 it's the same error:

ERROR: type KeyedDistribution has no field KeyedDistribution
Stacktrace:
 [1] getproperty(::KeyedDistribution{Multivariate,Continuous,MvNormal{Float64,PDMats.PDMat{Float64,Array{Float64,2}},Array{Float64,1}}}, ::Symbol) at ./Base.jl:33
 [2] ==(::KeyedDistribution{Multivariate,Continuous,MvNormal{Float64,PDMats.PDMat{Float64,Array{Float64,2}},Array{Float64,1}}}, ::KeyedDistribution{Multivariate,Continuous,MvNormal{Float64,PDMats.PDMat{Float64,Array{Float64,2}},Array{Float64,1}}}) at /Users/bencottier/.julia/packages/KeyedDistributions/i1TQS/src/KeyedDistributions.jl:15
 [3] top-level scope at REPL[36]:1

Whereas isequal(kd, kd2) == true.

The fields and types are all equal. I assume it relates to the way @auto_hash_equals processes the struct.

Fixed by specifying the type of `F` and `S`

- Increase coverage - Separate Distribution-only methods

codecov · 2021-03-09T18:12:28Z

Codecov Report

Merging #1 (1ef92ce) into main (883d9c9) will increase coverage by 98.00%.
The diff coverage is 98.00%.

@@            Coverage Diff            @@
##           main       #1       +/-   ##
=========================================
+ Coverage      0   98.00%   +98.00%     
=========================================
  Files         0        1        +1     
  Lines         0       50       +50     
=========================================
+ Hits          0       49       +49     
- Misses        0        1        +1

Impacted Files	Coverage Δ
src/KeyedDistributions.jl	`98.00% <98.00%> (ø)`

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 883d9c9...1ef92ce. Read the comment docs.

bencottier · 2021-03-09T18:46:48Z

I think we can open issues for the remaining tasks and in this PR nail down the Distributions API at least.

Issues: #2 #3 #4 #5

Expected to fix Pkg bug

Had to move it out of the `struct`

src/KeyedDistributions.jl

test/runtests.jl

src/KeyedDistributions.jl

Consistent with AxisKeys.jl and allows Matrix-variate Distributions

Also a bit of clean-up and styling

- Docstrings - Comments - Tests

@ref

I misunderstood @ref, I guess it cannot reference external things.

bencottier · 2021-03-12T17:20:48Z

@mcabbott We're interested in what you think of this package and this PR. It applies the idea of a KeyedArray to a Distribution, and is meant to extend the AxisKeys ecosystem.

src/KeyedDistributions.jl

test/runtests.jl

src/KeyedDistributions.jl

glennmoy · 2021-03-12T17:37:20Z

Test no longer broken

- Update outdated docstring - Simplify methods

bencottier · 2021-03-15T18:04:05Z

src/KeyedDistributions.jl

+                key_lengths = map(length, keys)
+                key_lengths == _size(d) || throw(ArgumentError(
+                    "lengths of key vectors $key_lengths must match " *
+                    "size of distribution $(_size(d))"))


We may not be pirating size, but "size" in this message may be misleading. I did it for the sake of generalising to one error message.

maybe remove reference to size in the error message?

"Dimensions of key vectors $key_lengths must match the distribution $(_size(d))"))

unless you think referring to dimensions is even more confusing?

I'm not a fan of dimensions, because it's ambiguous whether it's the number of dimensions or the size on each dimension. But you know what, this is testing both of those things, so maybe it's just right.

bencottier · 2021-03-15T18:10:08Z

src/KeyedDistributions.jl

+    @eval Distributions.$f(d::KeyedDistribution{<:Univariate}) = $f(distribution(d))
+end
+
+# Needed to avoid method ambiguity


Method ambiguity errors from the tests, from old code when I put the below functions in the first @eval loop

MethodError: cdf(::KeyedDistribution{Univariate,Continuous,Normal{Float64}}, ::Int64) is ambiguous. Candidates: cdf(d::Distribution{Univariate,Continuous}, x::Real) in Distributions at /Users/bencottier/.julia/packages/Distributions/cNe2C/src/univariates.jl:367 cdf(d::KeyedDistribution{var"#s12",S,D} where D<:Distribution{var"#s12",S} where S<:ValueSupport where var"#s12"<:Univariate, args...) in KeyedDistributions at /Users/bencottier/JuliaEnvs/KeyedDistributions/src/KeyedDistributions.jl:173 Possible fix, define cdf(::KeyedDistribution{Univariate,Continuous,D} where D<:Distribution{Univariate,Continuous}, ::Real)

MethodError: insupport(::KeyedDistribution{Univariate,Continuous,Normal{Float64}}, ::Int64) is ambiguous. Candidates: insupport(d::Union{Type{D}, D}, x::Real) where D<:Distribution{Univariate,Continuous} in Distributions at /Users/bencottier/.julia/packages/Distributions/cNe2C/src/univariates.jl:127 insupport(d::KeyedDistribution{var"#s12",S,D} where D<:Distribution{var"#s12",S} where S<:ValueSupport where var"#s12"<:Univariate, b) in KeyedDistributions at /Users/bencottier/JuliaEnvs/KeyedDistributions/src/KeyedDistributions.jl:179 Possible fix, define insupport(::D, ::Real) where D<:(KeyedDistribution{Univariate,Continuous,D} where D<:Distribution{Univariate,Continuous})

I think it's because we haven't provided enough type information to the function signature to restrict it to using our method.

We've only specified that it needs to be Univariate but Distributions provides the ValueSupport. So when it sees KeyedDistribution{Univariate, Continuous} it's conflicted between which version to use.

FWIW I don't understand why it still doesn't break outside the @eval loop, but we can/should fix it with the following

Distributions.cdf(d::KeyedArray{F, S, <:Distribution{F, S}, x) where {F, S} = cdf(distribution(d), x)

same with insupport.

Yeah I don't know why it worked this way.

glennmoy · 2021-03-15T19:35:25Z

test/runtests.jl

+    end
+
+    @testset "Distributions types" begin
+        @testset "Univariate" begin


Just to note we only test on Continuous distributions. We should probably test on Discrete at some point.
#7

Implement prototype

d305bc0

glennmoy reviewed Mar 8, 2021

View reviewed changes

src/KeyedDistributions.jl Outdated Show resolved Hide resolved

test/runtests.jl Outdated Show resolved Hide resolved

test/runtests.jl Outdated Show resolved Hide resolved

Apply suggestions from code review

318c726

Co-authored-by: Glenn Moynihan <glenn.moynihan@invenialabs.co.uk>

bencottier added 2 commits March 9, 2021 11:48

Limit KeyedSampleable methods

fff2948

Just those specified in Distributions.jl docs

Merge branch 'bc/prototype' of https://github.com/invenia/KeyedDistri…

1390c89

…butions.jl into bc/prototype

glennmoy mentioned this pull request Mar 9, 2021

Introduce dataset type invenia/AxisSets.jl#6

Merged

6 tasks

bencottier added 4 commits March 9, 2021 13:51

Rearrange tests after refactor

f827bfa

Bugfix: incorrect ==

2d53943

Fixed by specifying the type of `F` and `S`

Reorganise tests

177f619

- Increase coverage - Separate Distribution-only methods

Fix main branch references

64bedc0

bencottier changed the title ~~POC: Implement KeyedDistribution and KeyedSampleable~~ Implement KeyedDistribution and KeyedSampleable Mar 9, 2021

bencottier added 5 commits March 10, 2021 10:47

Update julia compat to 1.3

5b42400

Expected to fix Pkg bug

Bugfix: Documenter cannot find docstring

5951cd9

Had to move it out of the `struct`

Remove CI on 1.0 and change julia compat to 1.5

c120037

Set main as devbranch to build docs

7ae0b91

Test _rand!

590874f

glennmoy reviewed Mar 10, 2021

View reviewed changes

morris25 reviewed Mar 10, 2021

View reviewed changes

src/KeyedDistributions.jl Outdated Show resolved Hide resolved

Remove axiskeys type piracy of Sampleable

fbd6e19

bencottier added 14 commits March 11, 2021 11:57

Remove parent and keyless methods

1ac5d27

Make keys a tuple of vectors

e97b1da

Consistent with AxisKeys.jl and allows Matrix-variate Distributions

Add subtyping tests

837af84

Support matrix-variate and univariate distributions

2d7dc6d

Also a bit of clean-up and styling

Add default keys for no-keys constructor

d00c8c1

Remove unnecessary type checks

774e083

Clarify docstrings

84da47f

Generalise cov for Matrixvariate

ec836a5

Update tests for KeyedArray/PDMats issue

104f353

Compare more properties to underlying distribution

ae1ab09

Edit comment

41cf7b7

Clear up

c001761

- Docstrings - Comments - Tests

Simplify _rand!

1210eb2

Fix docs references

00d34d3

I misunderstood @ref, I guess it cannot reference external things.