Merge pull request #48 from JuliaAI/indexing

Optimize broadcasting on the pdf function

Author: OkonSamuel

src/arrays.jl

@@ -151,10 +151,10 @@ end
 
 # dummy function
 # returns `x[i]` for `Array` inputs `x`
-# For non-Array inputs returns `zero(dtype)`
+# For non-Array inputs returns the input
 #This avoids using an if statement
-_getindex(x::Array, i, dtype)=x[i]
-_getindex(::Nothing, i, dtype) = zero(dtype)
+_getindex(x::Array, i) = x[i]
+_getindex(x, i) = x
 
 # pdf.(u, cv)
 function Base.Broadcast.broadcasted(
@@ -164,15 +164,16 @@ function Base.Broadcast.broadcasted(
 
     # we assume that we compare categorical values by their unwrapped value
     # and pick the index of this value from classes(u)
-    cv_loc = findfirst(==(cv), classes(u))
-    cv_loc == 0 && throw(err_missing_class(cv))
+    _classes = classes(u)
+    cv_loc = get(CategoricalArrays.pool(_classes), cv, zero(R))
+    isequal(cv_loc, 0) && throw(err_missing_class(cv))
 
     f() = zeros(P, size(u)) #default caller function
 
     return Base.Broadcast.Broadcasted(
         identity,
         (get(f, u.prob_given_ref, cv_loc),)
-        )
+    )
 end
 
 Base.Broadcast.broadcasted(
@@ -191,15 +192,23 @@ function Base.Broadcast.broadcasted(
     length(u) == length(v) ||throw(DimensionMismatch(
         "Arrays could not be broadcast to a common size; "*
         "got a dimension with lengths $(length(u)) and $(length(v))"))
+    
+    _classes = classes(u)
+    _classes_pool = CategoricalArrays.pool(_classes)
+    T = eltype(v) >: Missing ? Missing : Union{}
+    v_loc_flat = Vector{Tuple{Union{R, T}, Int}}(undef, length(v))
+    
+    
+    for (i, x) in enumerate(v)
+        cv_ref = ismissing(x) ? missing : get(_classes_pool, x, zero(R))
+        isequal(cv_ref, 0) && throw(err_missing_class(x))
+        v_loc_flat[i] = (cv_ref, i)
+    end
 
-    v_loc_flat = [(ismissing(x) ? missing : findfirst(==(x), classes(u)), i)
-                  for (i, x) in enumerate(v)]
-    any(isequal(0), v_loc_flat) && throw(err_missing_class(cv))
-
-    getter((cv_loc, i), dtype) =
-        _getindex(get(u.prob_given_ref, cv_loc, nothing), i, dtype)
-    getter(::Tuple{Missing,Any}, dtype) = missing
-    ret_flat = getter.(v_loc_flat, P)
+    getter((cv_ref, i)) =
+        _getindex(get(u.prob_given_ref, cv_ref, zero(P)), i)
+    getter(::Tuple{Missing,Any}) = missing
+    ret_flat = getter.(v_loc_flat)
     return reshape(ret_flat, size(u))
 end
 

test/arrays.jl

@@ -149,6 +149,9 @@ end
     u2 = UnivariateFinite(v[1:2], probs, augment=true)
     @test pdf.(u2, v[3]) == zeros(3)
     @test isequal(logpdf.(u2, v[3]), log.(zeros(3)))
+
+    ## Check that the appropriate errors are thrown
+    @test_throws DomainError pdf.(u,"strange_level")
 end
 
 _skip(v) = collect(skipmissing(v))
@@ -166,18 +169,27 @@ _skip(v) = collect(skipmissing(v))
         @test _skip(broadcast(logpdf, u, unwrap.(v))) ==
                       _skip([logpdf(u[i], v[i]) for i in 1:length(u)])
     end
+
+    ## Check that the appropriate errors are thrown
+    v1 = categorical([v0[1:end-1]...;"strange_level"])
+    v2 = [v0...;rand(rng, v0)] #length(u) !== length(v2)
+    v3 = categorical([vm[end:-1:begin+1]...;"strange_level"])
+    @test_throws DimensionMismatch broadcast(pdf, u, v2)
+    @test_throws DomainError broadcast(pdf, u, v1)
+    @test_throws DomainError broadcast(pdf, u, v3)
+
 end
 
-@testset "broadcasting: check indexing in `getter((cv, i), dtype)` see PR#375" begin
+@testset "broadcasting: check indexing in `getter((cv_ref, i))` see PR#375 from MLJBase" begin
     c  = categorical([0,1,1])
     d = UnivariateFinite(c[1:1], [1 1 1]')
     v = categorical([0,1,1,1])
     @test broadcast(pdf, d, v[2:end]) == [0,0,0]
 end
 
 @testset "_getindex" begin
-   @test CategoricalDistributions._getindex(collect(1:4), 2, Int64) == 2
-   @test CategoricalDistributions._getindex(nothing, 2, Int64) == zero(Int64)
+   @test CategoricalDistributions._getindex(collect(1:4), 2) == 2
+   @test CategoricalDistributions._getindex(0, 2) === 0
 end
 
 @testset "broadcasting mode" begin