address #16

Author: OkonSamuel

src/arrays.jl

@@ -243,6 +243,12 @@ function Base.Broadcast.broadcasted(::typeof(mode),
     mode_flat = map(1:length(u)) do i
         max_prob = maximum(dic[ref][i] for ref in keys(dic))
         m = zero(R)
+        
+        # `maximum` of any iterable containing `NaN` would return `NaN` 
+        # For this case the index `m` won't be updated in the loop as relations
+        # involving NaN as one of it's argument always returns false 
+        # (e.g `==(NaN, NaN)` returns false)
+        throw_nan_error_if_needed(max_prob)
         for ref in keys(dic)
             if dic[ref][i] == max_prob
                 m = ref

src/methods.jl

@@ -196,6 +196,12 @@ function Dist.mode(d::UnivariateFinite)
     p = values(dic)
     max_prob = maximum(p)
     m = first(first(dic)) # mode, just some ref for now
+
+    # `maximum` of any iterable containing `NaN` would return `NaN` 
+    # For this case the index `m` won't be updated in the loop below as relations
+    # involving NaN as one of it's arguments always returns false
+    # (e.g `==(NaN, NaN)` returns false)
+    throw_nan_error_if_needed(max_prob)
     for (x, prob) in dic
         if prob == max_prob
             m = x
@@ -205,6 +211,18 @@ function Dist.mode(d::UnivariateFinite)
     return d.decoder(m)
 end
 
+function throw_nan_error_if_needed(x)
+    if isnan(x)
+        throw(
+            DomainError(
+                NaN,
+                "`mode` is invalid for `UnivariateFininite` distribution "*
+                "with `pdf` containing `NaN`s"
+            )
+        )
+    end
+end
+
 # mode(v::Vector{UnivariateFinite}) = mode.(v)
 # mode(u::UnivariateFiniteVector{2}) =
 #     [u.support[ifelse(s > 0.5, 2, 1)] for s in u.scores]

test/arrays.jl

@@ -198,6 +198,17 @@ end
     u   = UnivariateFinite(P, pool=missing)
     expected = mode.([u...])
     @test all(mode.(u) .== expected)
+
+    # `mode` broadcasting of `Univariate` objects containing `NaN` in probs.
+    unf_arr = UnivariateFinite(
+        [
+            0.1 0.2 NaN 0.1 NaN;
+            0.2 0.1 0.1 0.4 0.2;
+            0.3 NaN 0.2 NaN 0.3
+        ],
+        pool=missing
+    )
+    @test_throws DomainError mode.(unf_arr)
 end
 
 @testset "cat for UnivariateFiniteArray" begin

test/methods.jl

@@ -175,6 +175,10 @@ end
     p[42] = 0.5
     d = UnivariateFinite(v, p)
     @test mode(d) == 42
+
+    # `mode` of `Univariate` objects containing `NaN` in probs.
+    unf = UnivariateFinite([0.1, 0.2, NaN, 0.1, NaN], pool=missing)
+    @test_throws DomainError mode(unf)
 end
 
 @testset "UnivariateFinite methods" begin