diff --git a/src/univariate/continuous/gamma.jl b/src/univariate/continuous/gamma.jl index 3c4e3c44c4..866255fb7d 100644 --- a/src/univariate/continuous/gamma.jl +++ b/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 diff --git a/src/univariate/continuous/normal.jl b/src/univariate/continuous/normal.jl index 1dcadadac3..4231522674 100644 --- a/src/univariate/continuous/normal.jl +++ b/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,9 +150,9 @@ 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 @@ -160,8 +160,8 @@ function suffstats(::Type{<:Normal}, x::AbstractArray{T}, w::AbstractArray{Float 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,8 +182,8 @@ 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)) @@ -191,9 +191,9 @@ 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 diff --git a/src/univariate/discrete/poisson.jl b/src/univariate/discrete/poisson.jl index b7f64dba7f..b794fe7a3b 100644 --- a/src/univariate/discrete/poisson.jl +++ b/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 diff --git a/test/univariate/continuous/gamma.jl b/test/univariate/continuous/gamma.jl index b7b866a746..3937af191e 100644 --- a/test/univariate/continuous/gamma.jl +++ b/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 diff --git a/test/univariate/continuous/normal.jl b/test/univariate/continuous/normal.jl index ccc6266dbb..0378e05486 100644 --- a/test/univariate/continuous/normal.jl +++ b/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 diff --git a/test/univariate/discrete/poisson.jl b/test/univariate/discrete/poisson.jl index 657acb3c1f..0861de4a7f 100644 --- a/test/univariate/discrete/poisson.jl +++ b/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 \ No newline at end of file