Inbounds Error Fixes for some Univariates (#1742)

* Inbounds Error Fixes for some Univariates

Addresses #1265

- Fixed univariate normal
- Fixed univariate gamma
- Fixed poisson

* Fixes some type instabilities

* Update src/univariate/continuous/normal.jl

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>

* Update src/univariate/continuous/normal.jl

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>

* Update src/univariate/continuous/normal.jl

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>

* Update src/univariate/continuous/normal.jl

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>

* Update src/univariate/continuous/normal.jl

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>

* Fixes some type instabilities and adds tests

- Fixes more type instabilities in normal
- Adds tests for normal
- Adds tests for gamma
- Adds tests for poisson

* Slight modifications in tests

- Edited tests for normal
- Edited tests for gamma
- Edited tests for poisson

* Update tests

---------

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>

src/univariate/continuous/gamma.jl

@@ -152,7 +152,7 @@ function suffstats(::Type{<:Gamma}, x::AbstractArray{T}, w::AbstractArray{Float6
     sx = zero(T)
     slogx = zero(T)
     tw = zero(T)
-    for i = 1:n
+    for i in eachindex(x, w)
         @inbounds xi = x[i]
         @inbounds wi = w[i]
         sx += wi * xi

src/univariate/continuous/normal.jl

@@ -131,15 +131,15 @@ function suffstats(::Type{<:Normal}, x::AbstractArray{T}) where T<:Real
     n = length(x)
 
     # compute s
-    s = x[1]
-    for i = 2:n
+    s = zero(T) + zero(T)
+    for i in eachindex(x)
         @inbounds s += x[i]
     end
     m = s / n
 
     # compute s2
-    s2 = abs2(x[1] - m)
-    for i = 2:n
+    s2 = zero(m)
+    for i in eachindex(x)
         @inbounds s2 += abs2(x[i] - m)
     end
 
@@ -150,18 +150,18 @@ function suffstats(::Type{<:Normal}, x::AbstractArray{T}, w::AbstractArray{Float
     n = length(x)
 
     # compute s
-    tw = w[1]
-    s = w[1] * x[1]
-    for i = 2:n
+    tw = 0.0
+    s = 0.0 * zero(T)
+    for i in eachindex(x, w)
         @inbounds wi = w[i]
         @inbounds s += wi * x[i]
         tw += wi
     end
     m = s / tw
 
     # compute s2
-    s2 = w[1] * abs2(x[1] - m)
-    for i = 2:n
+    s2 = zero(m)
+    for i in eachindex(x, w)
         @inbounds s2 += w[i] * abs2(x[i] - m)
     end
 
@@ -182,18 +182,18 @@ end
 
 function suffstats(g::NormalKnownMu, x::AbstractArray{T}) where T<:Real
     μ = g.μ
-    s2 = abs2(x[1] - μ)
-    for i = 2:length(x)
+    s2 = zero(T) + zero(μ)
+    for i in eachindex(x)
         @inbounds s2 += abs2(x[i] - μ)
     end
     NormalKnownMuStats(g.μ, s2, length(x))
 end
 
 function suffstats(g::NormalKnownMu, x::AbstractArray{T}, w::AbstractArray{Float64}) where T<:Real
     μ = g.μ
-    s2 = abs2(x[1] - μ) * w[1]
-    tw = w[1]
-    for i = 2:length(x)
+    s2 = 0.0 * abs2(zero(T) - zero(μ))
+    tw = 0.0
+    for i in eachindex(x, w)
         @inbounds wi = w[i]
         @inbounds s2 += abs2(x[i] - μ) * wi
         tw += wi

src/univariate/discrete/poisson.jl

@@ -123,7 +123,7 @@ function suffstats(::Type{<:Poisson}, x::AbstractArray{T}, w::AbstractArray{Floa
     n == length(w) || throw(DimensionMismatch("Inconsistent array lengths."))
     sx = 0.
     tw = 0.
-    for i = 1 : n
+    for i in eachindex(x, w)
         @inbounds wi = w[i]
         @inbounds sx += x[i] * wi
         tw += wi

test/univariate/continuous/gamma.jl

@@ -1,3 +1,25 @@
+using Test, Distributions, OffsetArrays
+
 test_cgf(Gamma(1  ,1  ), (0.9, -1, -100f0, -1e6))
 test_cgf(Gamma(10 ,1  ), (0.9, -1, -100f0, -1e6))
 test_cgf(Gamma(0.2, 10), (0.08, -1, -100f0, -1e6))
+
+@testset "Gamma suffstats and OffsetArrays" begin
+    a = rand(Gamma(), 11)
+    wa = 1.0:11.0
+
+    resulta = @inferred(suffstats(Gamma, a))
+
+    resultwa = @inferred(suffstats(Gamma, a, wa))
+
+    b = OffsetArray(a, -5:5)
+    wb = OffsetArray(wa, -5:5)
+
+    resultb = @inferred(suffstats(Gamma, b))
+    @test resulta == resultb
+
+    resultwb = @inferred(suffstats(Gamma, b, wb))
+    @test resultwa == resultwb
+
+    @test_throws DimensionMismatch suffstats(Gamma, a, wb)
+end

test/univariate/continuous/normal.jl

@@ -1,4 +1,4 @@
-using Test, Distributions, StatsFuns, ForwardDiff
+using Test, Distributions, StatsFuns, ForwardDiff, OffsetArrays
 
 isnan_type(::Type{T}, v) where {T} = isnan(v) && v isa T
 
@@ -203,3 +203,23 @@ end
 # affine transformations
 test_affine_transformations(Normal, randn(), randn()^2)
 test_affine_transformations(NormalCanon, randn()^2, randn()^2)
+
+@testset "Normal suffstats and OffsetArrays" begin
+    a = rand(Normal(), 11)
+    wa = 1.0:11.0
+
+    resulta = @inferred(suffstats(Normal, a))
+
+    resultwa = @inferred(suffstats(Normal, a, wa))
+
+    b = OffsetArray(a, -5:5)
+    wb = OffsetArray(wa, -5:5)
+
+    resultb = @inferred(suffstats(Normal, b))
+    @test resulta == resultb
+
+    resultwb = @inferred(suffstats(Normal, b, wb))
+    @test resultwa == resultwb
+
+    @test_throws DimensionMismatch suffstats(Normal, b, wa)
+end

test/univariate/discrete/poisson.jl

@@ -1,4 +1,25 @@
+using Test, Distributions, OffsetArrays
 
 test_cgf(Poisson(1   ), (1f0,2f0,10.0,50.0))
 test_cgf(Poisson(10  ), (1f0,2f0,10.0,50.0))
 test_cgf(Poisson(1e-3), (1f0,2f0,10.0,50.0))
+
+@testset "Poisson suffstats and OffsetArrays" begin
+    a = rand(Poisson(), 11)
+    wa = 1.0:11.0
+
+    resulta = @inferred(suffstats(Poisson, a))
+
+    resultwa = @inferred(suffstats(Poisson, a, wa))
+
+    b = OffsetArray(a, -5:5)
+    wb = OffsetArray(wa, -5:5)
+
+    resultb = @inferred(suffstats(Poisson, b))
+    @test resulta == resultb
+
+    resultwb = @inferred(suffstats(Poisson, b, wb))
+    @test resultwa == resultwb
+
+    @test_throws DimensionMismatch suffstats(Poisson, a, wb)
+end