From 94cf6c0af2a533516a282048a147fa2bdf2f3dc2 Mon Sep 17 00:00:00 2001
From: "Ricardo M. S. Rosa" <rmsrosa@gmail.com>
Date: Sat, 20 Jan 2024 12:59:19 -0300
Subject: [PATCH 01/15] gradlogpdf for mixtures

---
 src/mixtures/mixturemodel.jl | 15 ++++++++++++
 test/gradlogpdf.jl           | 44 ++++++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+)

diff --git a/src/mixtures/mixturemodel.jl b/src/mixtures/mixturemodel.jl
index c3d3b1f91..e4a5a405a 100644
--- a/src/mixtures/mixturemodel.jl
+++ b/src/mixtures/mixturemodel.jl
@@ -87,6 +87,14 @@ Here, `x` can be a single sample or an array of multiple samples.
 """
 logpdf(d::AbstractMixtureModel, x::Any)
 
+"""
+    gradlogpdf(d::Union{UnivariateMixture, MultivariateMixture}, x)
+
+Evaluate the gradient of the logarithm of the (mixed) probability density function over `x`.
+Here, `x` is expected to be a single sample.
+"""
+gradlogpdf(d::AbstractMixtureModel, x::Any)
+
 """
     rand(d::Union{UnivariateMixture, MultivariateMixture})
 
@@ -359,8 +367,14 @@ function _mixlogpdf!(r::AbstractArray, d::AbstractMixtureModel, x)
     return r
 end
 
+function _mixgradlogpdf1(d::AbstractMixtureModel, x)
+    glp = insupport(d, x) ? sum(pbi * pdf(d.components[i], x) .* gradlogpdf(d.components[i], x) for (i, pbi) in enumerate(probs(d)) if pdf(d.components[i], x) > 0) / pdf(d, x) : zero(x)
+    return glp
+end
+
 pdf(d::UnivariateMixture, x::Real) = _mixpdf1(d, x)
 logpdf(d::UnivariateMixture, x::Real) = _mixlogpdf1(d, x)
+gradlogpdf(d::UnivariateMixture, x::Real) = _mixgradlogpdf1(d, x)
 
 _pdf!(r::AbstractArray{<:Real}, d::UnivariateMixture{Discrete}, x::UnitRange) = _mixpdf!(r, d, x)
 _pdf!(r::AbstractArray{<:Real}, d::UnivariateMixture, x::AbstractArray{<:Real}) = _mixpdf!(r, d, x)
@@ -368,6 +382,7 @@ _logpdf!(r::AbstractArray{<:Real}, d::UnivariateMixture, x::AbstractArray{<:Real
 
 _pdf(d::MultivariateMixture, x::AbstractVector{<:Real}) = _mixpdf1(d, x)
 _logpdf(d::MultivariateMixture, x::AbstractVector{<:Real}) = _mixlogpdf1(d, x)
+gradlogpdf(d::MultivariateMixture, x::AbstractVector{<:Real}) = _mixgradlogpdf1(d, x)
 _pdf!(r::AbstractArray{<:Real}, d::MultivariateMixture, x::AbstractMatrix{<:Real}) = _mixpdf!(r, d, x)
 _logpdf!(r::AbstractArray{<:Real}, d::MultivariateMixture, x::AbstractMatrix{<:Real}) = _mixlogpdf!(r, d, x)
 
diff --git a/test/gradlogpdf.jl b/test/gradlogpdf.jl
index f4216a67e..275e4a6ea 100644
--- a/test/gradlogpdf.jl
+++ b/test/gradlogpdf.jl
@@ -25,3 +25,47 @@ using Test
     [0.191919191919192, 1.080808080808081]   ,atol=1.0e-8)
 @test isapprox(gradlogpdf(MvTDist(5., [1., 2.], [1. 0.1; 0.1 1.]), [0.7, 0.9]),
     [0.2150711513583442, 1.2111901681759383] ,atol=1.0e-8)
+
+# Test for gradlogpdf on univariate mixture distributions
+
+x = [-0.2, 0.3, 0.8, 1.3, 10.5]
+delta = 0.001
+
+for d in (
+    MixtureModel([Normal(-4.5, 2.0)], [1.0]),
+    MixtureModel([Exponential(2.0)], [1.0]),
+    MixtureModel([Uniform(-1.0, 1.0)], [1.0]),
+    MixtureModel([Beta(1.5, 3.0), Chi(5.0), Chisq(7.0)], [0.4, 0.3, 0.3]),
+    MixtureModel([Exponential(2.0), Gamma(9.0, 0.5), Gumbel(3.5, 1.0), Laplace(7.0)], [0.3, 0.2, 0.4, 0.1]),
+    MixtureModel([Logistic(-6.0), LogNormal(5.5), TDist(8.0), Weibull(2.0)], [0.3, 0.2, 0.4, 0.1])
+)
+    xs = filter(s -> insupport(d, s), x)
+    glp1 = gradlogpdf.(d, xs)
+    glp2 = ( logpdf.(d, xs .+ delta) - logpdf.(d, xs .- delta) ) ./ 2delta
+    @test isapprox(glp1, glp2, atol = delta)
+end
+
+# Test for gradlogpdf on multivariate mixture distributions
+
+x = [[0.2, 0.3], [0.8, 1.3], [-1.0, 10.5]]
+delta = 0.001
+
+for d in (
+    MixtureModel([MvNormal([1., 2.], [1. 0.1; 0.1 1.])], [1.0]),
+    MixtureModel([MvNormal([1.0, 2.0], [0.4 0.2; 0.2 0.5]), MvNormal([2.0, 1.0], [0.3 0.1; 0.1 0.4])], [0.4, 0.6]),
+    MixtureModel([MvTDist(5., [1., 2.], [1. 0.1; 0.1 1.])], [1.0]),
+    MixtureModel([MvNormal([1.0, 2.0], [0.4 0.2; 0.2 0.5]), MvTDist(5., [1., 2.], [1. 0.1; 0.1 1.])], [0.4, 0.6])
+)
+    xs = filter(s -> insupport(d, s), x)
+    for xi in xs
+        glp = gradlogpdf(d, xi)
+        glpx = ( logpdf(d, xi .+ [delta, 0]) - logpdf(d, xi .- [delta, 0]) ) ./ 2delta
+        glpy = ( logpdf(d, xi .+ [0, delta]) - logpdf(d, xi .- [0, delta]) ) ./ 2delta
+        @test isapprox(glp, [glpx, glpy], atol = delta)
+    end
+end
+
+#@test isapprox(gradlogpdf(MixtureModel([MvNormal([1., 2.], [1. 0.1; 0.1 1.])], [1.0]), [0.7, 0.9])   , 
+#[-0.4375, 2.375]   ,atol=1.0e-8)
+#@test isapprox(gradlogpdf(MixtureModel([MvNormal([1.0, 2.0], [0.4 0.2; 0.2 0.5]), MvNormal([2.0, 1.0], [0.3 0.1; 0.1 0.4])], [0.4, 0.6]), [0.7, 0.9])   ,
+#[-0.4375, 2.375]   ,atol=1.0e-8)
\ No newline at end of file

From e7dd8ef6762b1a73158f1a577740f0de7ce252fb Mon Sep 17 00:00:00 2001
From: "Ricardo M. S. Rosa" <rmsrosa@gmail.com>
Date: Sat, 20 Jan 2024 16:37:15 -0300
Subject: [PATCH 02/15] make them work on array of multiple samples

---
 src/common.jl                | 59 ++++++++++++++++++++++++++++++++++++
 src/mixtures/mixturemodel.jl |  4 +--
 test/gradlogpdf.jl           | 22 +++++---------
 3 files changed, 69 insertions(+), 16 deletions(-)

diff --git a/src/common.jl b/src/common.jl
index 31a218e19..4dbd35adb 100644
--- a/src/common.jl
+++ b/src/common.jl
@@ -277,6 +277,46 @@ end
 # `_logpdf` should be implemented and has no default definition
 # _logpdf(d::Distribution{ArrayLikeVariate{N}}, x::AbstractArray{<:Real,N}) where {N}
 
+"""
+    gradlogpdf(d::Distribution{ArrayLikeVariate{N}}, x::AbstractArray{<:Real,N}) where {N}
+
+Evaluate the gradient of the logarithm of the probability density function of `d` at `x`.
+
+This function checks if the size of `x` is compatible with distribution `d`. This check can
+be disabled by using `@inbounds`.
+
+# Implementation
+
+Instead of `gradlogpdf` one should implement `_gradlogpdf(d, x)` which does not have to check the
+size of `x`.
+
+See also: [`pdf`](@ref).
+"""
+@inline function gradlogpdf(
+    d::Distribution{ArrayLikeVariate{N}}, x::AbstractArray{<:Real,M}
+) where {N,M}
+    if M == N
+        @boundscheck begin
+            size(x) == size(d) ||
+                throw(DimensionMismatch("inconsistent array dimensions"))
+        end
+        return _gradlogpdf(d, x)
+    else
+        @boundscheck begin
+            M > N ||
+                throw(DimensionMismatch(
+                    "number of dimensions of the variates ($M) must be greater than or equal to the dimension of the distribution ($N)"
+                ))
+            ntuple(i -> size(x, i), Val(N)) == size(d) ||
+                throw(DimensionMismatch("inconsistent array dimensions"))
+        end
+        return @inbounds map(Base.Fix1(gradlogpdf, d), eachvariate(x, variate_form(typeof(d))))
+    end
+end
+
+# `_gradlogpdf` should be implemented and has no default definition
+# _gradlogpdf(d::Distribution{ArrayLikeVariate{N}}, x::AbstractArray{<:Real,N}) where {N}
+
 # TODO: deprecate?
 """
     pdf(d::Distribution{ArrayLikeVariate{N}}, x) where {N}
@@ -315,6 +355,25 @@ Base.@propagate_inbounds function logpdf(
     return map(Base.Fix1(logpdf, d), x)
 end
 
+"""
+    gradlogpdf(d::Distribution{ArrayLikeVariate{N}}, x) where {N}
+
+Evaluate the gradient of the logarithm of the probability density function of `d` at every 
+element in a collection `x`.
+
+This function checks for every element of `x` if its size is compatible with distribution
+`d`. This check can be disabled by using `@inbounds`.
+
+Here, `x` can be
+- an array of dimension `> N` with `size(x)[1:N] == size(d)`, or
+- an array of arrays `xi` of dimension `N` with `size(xi) == size(d)`.
+"""
+Base.@propagate_inbounds function gradlogpdf(
+    d::Distribution{ArrayLikeVariate{N}}, x::AbstractArray{<:AbstractArray{<:Real,N}},
+) where {N}
+    return map(Base.Fix1(gradlogpdf, d), x)
+end
+
 """
     pdf!(out, d::Distribution{ArrayLikeVariate{N}}, x) where {N}
 
diff --git a/src/mixtures/mixturemodel.jl b/src/mixtures/mixturemodel.jl
index e4a5a405a..8bb39afea 100644
--- a/src/mixtures/mixturemodel.jl
+++ b/src/mixtures/mixturemodel.jl
@@ -91,7 +91,7 @@ logpdf(d::AbstractMixtureModel, x::Any)
     gradlogpdf(d::Union{UnivariateMixture, MultivariateMixture}, x)
 
 Evaluate the gradient of the logarithm of the (mixed) probability density function over `x`.
-Here, `x` is expected to be a single sample.
+Here, `x` can be a single sample or an array of multiple samples.
 """
 gradlogpdf(d::AbstractMixtureModel, x::Any)
 
@@ -382,7 +382,7 @@ _logpdf!(r::AbstractArray{<:Real}, d::UnivariateMixture, x::AbstractArray{<:Real
 
 _pdf(d::MultivariateMixture, x::AbstractVector{<:Real}) = _mixpdf1(d, x)
 _logpdf(d::MultivariateMixture, x::AbstractVector{<:Real}) = _mixlogpdf1(d, x)
-gradlogpdf(d::MultivariateMixture, x::AbstractVector{<:Real}) = _mixgradlogpdf1(d, x)
+_gradlogpdf(d::MultivariateMixture, x::AbstractVector{<:Real}) = _mixgradlogpdf1(d, x)
 _pdf!(r::AbstractArray{<:Real}, d::MultivariateMixture, x::AbstractMatrix{<:Real}) = _mixpdf!(r, d, x)
 _logpdf!(r::AbstractArray{<:Real}, d::MultivariateMixture, x::AbstractMatrix{<:Real}) = _mixlogpdf!(r, d, x)
 
diff --git a/test/gradlogpdf.jl b/test/gradlogpdf.jl
index 275e4a6ea..ce7843aa5 100644
--- a/test/gradlogpdf.jl
+++ b/test/gradlogpdf.jl
@@ -26,7 +26,7 @@ using Test
 @test isapprox(gradlogpdf(MvTDist(5., [1., 2.], [1. 0.1; 0.1 1.]), [0.7, 0.9]),
     [0.2150711513583442, 1.2111901681759383] ,atol=1.0e-8)
 
-# Test for gradlogpdf on univariate mixture distributions
+# Test for gradlogpdf on univariate mixture distributions against centered finite-difference on logpdf
 
 x = [-0.2, 0.3, 0.8, 1.3, 10.5]
 delta = 0.001
@@ -40,12 +40,12 @@ for d in (
     MixtureModel([Logistic(-6.0), LogNormal(5.5), TDist(8.0), Weibull(2.0)], [0.3, 0.2, 0.4, 0.1])
 )
     xs = filter(s -> insupport(d, s), x)
-    glp1 = gradlogpdf.(d, xs)
+    glp1 = gradlogpdf(d, xs)
     glp2 = ( logpdf.(d, xs .+ delta) - logpdf.(d, xs .- delta) ) ./ 2delta
     @test isapprox(glp1, glp2, atol = delta)
 end
 
-# Test for gradlogpdf on multivariate mixture distributions
+# Test for gradlogpdf on multivariate mixture distributions against centered finite-difference on logpdf
 
 x = [[0.2, 0.3], [0.8, 1.3], [-1.0, 10.5]]
 delta = 0.001
@@ -57,15 +57,9 @@ for d in (
     MixtureModel([MvNormal([1.0, 2.0], [0.4 0.2; 0.2 0.5]), MvTDist(5., [1., 2.], [1. 0.1; 0.1 1.])], [0.4, 0.6])
 )
     xs = filter(s -> insupport(d, s), x)
-    for xi in xs
-        glp = gradlogpdf(d, xi)
-        glpx = ( logpdf(d, xi .+ [delta, 0]) - logpdf(d, xi .- [delta, 0]) ) ./ 2delta
-        glpy = ( logpdf(d, xi .+ [0, delta]) - logpdf(d, xi .- [0, delta]) ) ./ 2delta
-        @test isapprox(glp, [glpx, glpy], atol = delta)
-    end
+    glp = gradlogpdf(d, xs)
+    glpx = ( logpdf(d, xs .+ [[delta, 0]]) - logpdf(d, xs .- [[delta, 0]]) ) ./ 2delta
+    glpy = ( logpdf(d, xs .+ [[0, delta]]) - logpdf(d, xs .- [[0, delta]]) ) ./ 2delta
+    @test isapprox(getindex.(glp, 1), glpx, atol=delta)
+    @test isapprox(getindex.(glp, 2), glpy, atol=delta)
 end
-
-#@test isapprox(gradlogpdf(MixtureModel([MvNormal([1., 2.], [1. 0.1; 0.1 1.])], [1.0]), [0.7, 0.9])   , 
-#[-0.4375, 2.375]   ,atol=1.0e-8)
-#@test isapprox(gradlogpdf(MixtureModel([MvNormal([1.0, 2.0], [0.4 0.2; 0.2 0.5]), MvNormal([2.0, 1.0], [0.3 0.1; 0.1 0.4])], [0.4, 0.6]), [0.7, 0.9])   ,
-#[-0.4375, 2.375]   ,atol=1.0e-8)
\ No newline at end of file

From 39309837f370028d44dc218268a623226a5e9fed Mon Sep 17 00:00:00 2001
From: "Ricardo M. S. Rosa" <rmsrosa@gmail.com>
Date: Sat, 20 Jan 2024 17:58:25 -0300
Subject: [PATCH 03/15] use iszero instead of > 0

---
 src/mixtures/mixturemodel.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/mixtures/mixturemodel.jl b/src/mixtures/mixturemodel.jl
index 8bb39afea..3c15234b0 100644
--- a/src/mixtures/mixturemodel.jl
+++ b/src/mixtures/mixturemodel.jl
@@ -368,7 +368,7 @@ function _mixlogpdf!(r::AbstractArray, d::AbstractMixtureModel, x)
 end
 
 function _mixgradlogpdf1(d::AbstractMixtureModel, x)
-    glp = insupport(d, x) ? sum(pbi * pdf(d.components[i], x) .* gradlogpdf(d.components[i], x) for (i, pbi) in enumerate(probs(d)) if pdf(d.components[i], x) > 0) / pdf(d, x) : zero(x)
+    glp = insupport(d, x) ? sum(pi * pdf(d.components[i], x) .* gradlogpdf(d.components[i], x) for (i, pi) in enumerate(probs(d)) if (!iszero(pi) && !iszero(pdf(d.components[i], x)))) / pdf(d, x) : zero(x)
     return glp
 end
 

From 65e32101374209e9f7dfa715434b29ca16c3fd97 Mon Sep 17 00:00:00 2001
From: "Ricardo M. S. Rosa" <rmsrosa@gmail.com>
Date: Sat, 20 Jan 2024 18:01:03 -0300
Subject: [PATCH 04/15] tests with a zero prior

---
 test/gradlogpdf.jl | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/test/gradlogpdf.jl b/test/gradlogpdf.jl
index ce7843aa5..14b72663c 100644
--- a/test/gradlogpdf.jl
+++ b/test/gradlogpdf.jl
@@ -35,6 +35,7 @@ for d in (
     MixtureModel([Normal(-4.5, 2.0)], [1.0]),
     MixtureModel([Exponential(2.0)], [1.0]),
     MixtureModel([Uniform(-1.0, 1.0)], [1.0]),
+    MixtureModel([Normal(-2.0, 3.5), Normal(-4.5, 2.0)], [0.0, 1.0]),
     MixtureModel([Beta(1.5, 3.0), Chi(5.0), Chisq(7.0)], [0.4, 0.3, 0.3]),
     MixtureModel([Exponential(2.0), Gamma(9.0, 0.5), Gumbel(3.5, 1.0), Laplace(7.0)], [0.3, 0.2, 0.4, 0.1]),
     MixtureModel([Logistic(-6.0), LogNormal(5.5), TDist(8.0), Weibull(2.0)], [0.3, 0.2, 0.4, 0.1])
@@ -53,6 +54,7 @@ delta = 0.001
 for d in (
     MixtureModel([MvNormal([1., 2.], [1. 0.1; 0.1 1.])], [1.0]),
     MixtureModel([MvNormal([1.0, 2.0], [0.4 0.2; 0.2 0.5]), MvNormal([2.0, 1.0], [0.3 0.1; 0.1 0.4])], [0.4, 0.6]),
+    MixtureModel([MvNormal([1.0, 2.0], [0.4 0.2; 0.2 0.5]), MvNormal([2.0, 1.0], [0.3 0.1; 0.1 0.4])], [1.0, 0.0]),
     MixtureModel([MvTDist(5., [1., 2.], [1. 0.1; 0.1 1.])], [1.0]),
     MixtureModel([MvNormal([1.0, 2.0], [0.4 0.2; 0.2 0.5]), MvTDist(5., [1., 2.], [1. 0.1; 0.1 1.])], [0.4, 0.6])
 )

From af7dde6d974b68d1aa7927c6be4f9198b6062839 Mon Sep 17 00:00:00 2001
From: "Ricardo M. S. Rosa" <rmsrosa@gmail.com>
Date: Sat, 20 Jan 2024 19:54:41 -0300
Subject: [PATCH 05/15] update docs

---
 docs/src/censored.md   | 1 +
 docs/src/mixture.md    | 1 +
 docs/src/truncate.md   | 1 +
 docs/src/univariate.md | 1 +
 4 files changed, 4 insertions(+)

diff --git a/docs/src/censored.md b/docs/src/censored.md
index 3059da015..1e11b5aa4 100644
--- a/docs/src/censored.md
+++ b/docs/src/censored.md
@@ -24,6 +24,7 @@ Many functions, including those for the evaluation of pdf and sampling, are defi
 - [`insupport(::UnivariateDistribution, x::Any)`](@ref)
 - [`pdf(::UnivariateDistribution, ::Real)`](@ref)
 - [`logpdf(::UnivariateDistribution, ::Real)`](@ref)
+- [`gradlogpdf(::UnivariateDistribution, ::Real)`](@ref)
 - [`cdf(::UnivariateDistribution, ::Real)`](@ref)
 - [`logcdf(::UnivariateDistribution, ::Real)`](@ref)
 - [`logdiffcdf(::UnivariateDistribution, ::T, ::T) where {T <: Real}`](@ref)
diff --git a/docs/src/mixture.md b/docs/src/mixture.md
index e6b24b103..5549de039 100644
--- a/docs/src/mixture.md
+++ b/docs/src/mixture.md
@@ -99,6 +99,7 @@ var(::UnivariateMixture)
 length(::MultivariateMixture)
 pdf(::AbstractMixtureModel, ::Any)
 logpdf(::AbstractMixtureModel, ::Any)
+gradlogpdf(::AbstractMixtureModel, ::Any)
 rand(::AbstractMixtureModel)
 rand!(::AbstractMixtureModel, ::AbstractArray)
 ```
diff --git a/docs/src/truncate.md b/docs/src/truncate.md
index 4d4f63f9f..2362d166b 100644
--- a/docs/src/truncate.md
+++ b/docs/src/truncate.md
@@ -26,6 +26,7 @@ are defined for all truncated univariate distributions:
 - [`insupport(::UnivariateDistribution, x::Any)`](@ref)
 - [`pdf(::UnivariateDistribution, ::Real)`](@ref)
 - [`logpdf(::UnivariateDistribution, ::Real)`](@ref)
+- [`gradlogpdf(::UnivariateDistribution, ::Real)`](@ref)
 - [`cdf(::UnivariateDistribution, ::Real)`](@ref)
 - [`logcdf(::UnivariateDistribution, ::Real)`](@ref)
 - [`logdiffcdf(::UnivariateDistribution, ::T, ::T) where {T <: Real}`](@ref)
diff --git a/docs/src/univariate.md b/docs/src/univariate.md
index 0b2c48c6e..c04da3913 100644
--- a/docs/src/univariate.md
+++ b/docs/src/univariate.md
@@ -73,6 +73,7 @@ pdfsquaredL2norm
 insupport(::UnivariateDistribution, x::Any)
 pdf(::UnivariateDistribution, ::Real)
 logpdf(::UnivariateDistribution, ::Real)
+gradlogpdf(::UnivariateDistribution, ::Real)
 loglikelihood(::UnivariateDistribution, ::AbstractArray)
 cdf(::UnivariateDistribution, ::Real)
 logcdf(::UnivariateDistribution, ::Real)

From 6ed94767c3ffabb8a472d6633e16d29f3fef0854 Mon Sep 17 00:00:00 2001
From: Ricardo Rosa <rmsrosa@gmail.com>
Date: Sat, 20 Jan 2024 21:55:08 -0300
Subject: [PATCH 06/15] no docstring

---
 docs/src/censored.md | 1 -
 1 file changed, 1 deletion(-)

diff --git a/docs/src/censored.md b/docs/src/censored.md
index 1e11b5aa4..3059da015 100644
--- a/docs/src/censored.md
+++ b/docs/src/censored.md
@@ -24,7 +24,6 @@ Many functions, including those for the evaluation of pdf and sampling, are defi
 - [`insupport(::UnivariateDistribution, x::Any)`](@ref)
 - [`pdf(::UnivariateDistribution, ::Real)`](@ref)
 - [`logpdf(::UnivariateDistribution, ::Real)`](@ref)
-- [`gradlogpdf(::UnivariateDistribution, ::Real)`](@ref)
 - [`cdf(::UnivariateDistribution, ::Real)`](@ref)
 - [`logcdf(::UnivariateDistribution, ::Real)`](@ref)
 - [`logdiffcdf(::UnivariateDistribution, ::T, ::T) where {T <: Real}`](@ref)

From fe2f778cf526aec531f85d0dfec565bcd37afb9e Mon Sep 17 00:00:00 2001
From: "Ricardo M. S. Rosa" <rmsrosa@gmail.com>
Date: Tue, 23 Jan 2024 12:51:06 +0000
Subject: [PATCH 07/15] remove implicit broadcasting of gradlogpdf

---
 src/common.jl                | 59 ------------------------------------
 src/mixtures/mixturemodel.jl |  2 +-
 test/gradlogpdf.jl           | 14 +++++----
 3 files changed, 9 insertions(+), 66 deletions(-)

diff --git a/src/common.jl b/src/common.jl
index 4dbd35adb..31a218e19 100644
--- a/src/common.jl
+++ b/src/common.jl
@@ -277,46 +277,6 @@ end
 # `_logpdf` should be implemented and has no default definition
 # _logpdf(d::Distribution{ArrayLikeVariate{N}}, x::AbstractArray{<:Real,N}) where {N}
 
-"""
-    gradlogpdf(d::Distribution{ArrayLikeVariate{N}}, x::AbstractArray{<:Real,N}) where {N}
-
-Evaluate the gradient of the logarithm of the probability density function of `d` at `x`.
-
-This function checks if the size of `x` is compatible with distribution `d`. This check can
-be disabled by using `@inbounds`.
-
-# Implementation
-
-Instead of `gradlogpdf` one should implement `_gradlogpdf(d, x)` which does not have to check the
-size of `x`.
-
-See also: [`pdf`](@ref).
-"""
-@inline function gradlogpdf(
-    d::Distribution{ArrayLikeVariate{N}}, x::AbstractArray{<:Real,M}
-) where {N,M}
-    if M == N
-        @boundscheck begin
-            size(x) == size(d) ||
-                throw(DimensionMismatch("inconsistent array dimensions"))
-        end
-        return _gradlogpdf(d, x)
-    else
-        @boundscheck begin
-            M > N ||
-                throw(DimensionMismatch(
-                    "number of dimensions of the variates ($M) must be greater than or equal to the dimension of the distribution ($N)"
-                ))
-            ntuple(i -> size(x, i), Val(N)) == size(d) ||
-                throw(DimensionMismatch("inconsistent array dimensions"))
-        end
-        return @inbounds map(Base.Fix1(gradlogpdf, d), eachvariate(x, variate_form(typeof(d))))
-    end
-end
-
-# `_gradlogpdf` should be implemented and has no default definition
-# _gradlogpdf(d::Distribution{ArrayLikeVariate{N}}, x::AbstractArray{<:Real,N}) where {N}
-
 # TODO: deprecate?
 """
     pdf(d::Distribution{ArrayLikeVariate{N}}, x) where {N}
@@ -355,25 +315,6 @@ Base.@propagate_inbounds function logpdf(
     return map(Base.Fix1(logpdf, d), x)
 end
 
-"""
-    gradlogpdf(d::Distribution{ArrayLikeVariate{N}}, x) where {N}
-
-Evaluate the gradient of the logarithm of the probability density function of `d` at every 
-element in a collection `x`.
-
-This function checks for every element of `x` if its size is compatible with distribution
-`d`. This check can be disabled by using `@inbounds`.
-
-Here, `x` can be
-- an array of dimension `> N` with `size(x)[1:N] == size(d)`, or
-- an array of arrays `xi` of dimension `N` with `size(xi) == size(d)`.
-"""
-Base.@propagate_inbounds function gradlogpdf(
-    d::Distribution{ArrayLikeVariate{N}}, x::AbstractArray{<:AbstractArray{<:Real,N}},
-) where {N}
-    return map(Base.Fix1(gradlogpdf, d), x)
-end
-
 """
     pdf!(out, d::Distribution{ArrayLikeVariate{N}}, x) where {N}
 
diff --git a/src/mixtures/mixturemodel.jl b/src/mixtures/mixturemodel.jl
index 3c15234b0..ad8d76cad 100644
--- a/src/mixtures/mixturemodel.jl
+++ b/src/mixtures/mixturemodel.jl
@@ -382,10 +382,10 @@ _logpdf!(r::AbstractArray{<:Real}, d::UnivariateMixture, x::AbstractArray{<:Real
 
 _pdf(d::MultivariateMixture, x::AbstractVector{<:Real}) = _mixpdf1(d, x)
 _logpdf(d::MultivariateMixture, x::AbstractVector{<:Real}) = _mixlogpdf1(d, x)
-_gradlogpdf(d::MultivariateMixture, x::AbstractVector{<:Real}) = _mixgradlogpdf1(d, x)
 _pdf!(r::AbstractArray{<:Real}, d::MultivariateMixture, x::AbstractMatrix{<:Real}) = _mixpdf!(r, d, x)
 _logpdf!(r::AbstractArray{<:Real}, d::MultivariateMixture, x::AbstractMatrix{<:Real}) = _mixlogpdf!(r, d, x)
 
+gradlogpdf(d::MultivariateMixture, x::AbstractVector{<:Real}) = _mixgradlogpdf1(d, x)
 
 ## component-wise pdf and logpdf
 
diff --git a/test/gradlogpdf.jl b/test/gradlogpdf.jl
index 14b72663c..3e6ab8a10 100644
--- a/test/gradlogpdf.jl
+++ b/test/gradlogpdf.jl
@@ -41,7 +41,7 @@ for d in (
     MixtureModel([Logistic(-6.0), LogNormal(5.5), TDist(8.0), Weibull(2.0)], [0.3, 0.2, 0.4, 0.1])
 )
     xs = filter(s -> insupport(d, s), x)
-    glp1 = gradlogpdf(d, xs)
+    glp1 = gradlogpdf.(d, xs)
     glp2 = ( logpdf.(d, xs .+ delta) - logpdf.(d, xs .- delta) ) ./ 2delta
     @test isapprox(glp1, glp2, atol = delta)
 end
@@ -59,9 +59,11 @@ for d in (
     MixtureModel([MvNormal([1.0, 2.0], [0.4 0.2; 0.2 0.5]), MvTDist(5., [1., 2.], [1. 0.1; 0.1 1.])], [0.4, 0.6])
 )
     xs = filter(s -> insupport(d, s), x)
-    glp = gradlogpdf(d, xs)
-    glpx = ( logpdf(d, xs .+ [[delta, 0]]) - logpdf(d, xs .- [[delta, 0]]) ) ./ 2delta
-    glpy = ( logpdf(d, xs .+ [[0, delta]]) - logpdf(d, xs .- [[0, delta]]) ) ./ 2delta
-    @test isapprox(getindex.(glp, 1), glpx, atol=delta)
-    @test isapprox(getindex.(glp, 2), glpy, atol=delta)
+    for xi in xs
+        glp = gradlogpdf(d, xi)
+        glpx = ( logpdf(d, xi .+ [delta, 0]) - logpdf(d, xi .- [delta, 0]) ) ./ 2delta
+        glpy = ( logpdf(d, xi .+ [0, delta]) - logpdf(d, xi .- [0, delta]) ) ./ 2delta
+        @test isapprox(glp[1], glpx, atol=delta)
+        @test isapprox(glp[2], glpy, atol=delta)
+    end
 end

From 56860a1956e04dafc0de1bb21b0713f19d535449 Mon Sep 17 00:00:00 2001
From: "Ricardo M. S. Rosa" <rmsrosa@gmail.com>
Date: Tue, 23 Jan 2024 12:54:01 +0000
Subject: [PATCH 08/15] make gradlogpdf type stable... hopefully

---
 src/mixtures/mixturemodel.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/mixtures/mixturemodel.jl b/src/mixtures/mixturemodel.jl
index ad8d76cad..4e6932d93 100644
--- a/src/mixtures/mixturemodel.jl
+++ b/src/mixtures/mixturemodel.jl
@@ -368,7 +368,8 @@ function _mixlogpdf!(r::AbstractArray, d::AbstractMixtureModel, x)
 end
 
 function _mixgradlogpdf1(d::AbstractMixtureModel, x)
-    glp = insupport(d, x) ? sum(pi * pdf(d.components[i], x) .* gradlogpdf(d.components[i], x) for (i, pi) in enumerate(probs(d)) if (!iszero(pi) && !iszero(pdf(d.components[i], x)))) / pdf(d, x) : zero(x)
+    pdfdx = pdf(d, x) 
+    glp = !iszero(pdfdx) ? sum(pi * pdf(d.components[i], x) .* gradlogpdf(d.components[i], x) for (i, pi) in enumerate(probs(d)) if (!iszero(pi) && !iszero(pdf(d.components[i], x)))) / pdfdx : pdfdx
     return glp
 end
 

From b50f5eb6d31ee26e378255025d14dac82124fe2f Mon Sep 17 00:00:00 2001
From: "Ricardo M. S. Rosa" <rmsrosa@gmail.com>
Date: Fri, 26 Jan 2024 11:40:05 +0000
Subject: [PATCH 09/15] make gradlogpdf type stable

---
 src/mixtures/mixturemodel.jl | 55 ++++++++++++++++++++++++++++++------
 test/gradlogpdf.jl           | 10 ++++---
 2 files changed, 53 insertions(+), 12 deletions(-)

diff --git a/src/mixtures/mixturemodel.jl b/src/mixtures/mixturemodel.jl
index 4e6932d93..31f1931fb 100644
--- a/src/mixtures/mixturemodel.jl
+++ b/src/mixtures/mixturemodel.jl
@@ -367,15 +367,31 @@ function _mixlogpdf!(r::AbstractArray, d::AbstractMixtureModel, x)
     return r
 end
 
-function _mixgradlogpdf1(d::AbstractMixtureModel, x)
-    pdfdx = pdf(d, x) 
-    glp = !iszero(pdfdx) ? sum(pi * pdf(d.components[i], x) .* gradlogpdf(d.components[i], x) for (i, pi) in enumerate(probs(d)) if (!iszero(pi) && !iszero(pdf(d.components[i], x)))) / pdfdx : pdfdx
-    return glp
-end
-
 pdf(d::UnivariateMixture, x::Real) = _mixpdf1(d, x)
 logpdf(d::UnivariateMixture, x::Real) = _mixlogpdf1(d, x)
-gradlogpdf(d::UnivariateMixture, x::Real) = _mixgradlogpdf1(d, x)
+
+function gradlogpdf(d::UnivariateMixture, x::Real)
+    cp = components(d)
+    pr = probs(d)
+    pdfx1 = pdf(cp[1], x)
+    pdfx = pr[1] * pdfx1
+    _glp = pdfx * gradlogpdf(cp[1], x)
+    glp = (!iszero(pr[1])) && (!iszero(pdfx)) ? _glp : zero(_glp)
+    @inbounds for i in Iterators.drop(eachindex(pr, cp), 1)
+        if !iszero(pr[i])
+            pdfxi = pdf(cp[i], x)
+            if !iszero(pdfxi)
+                pipdfxi = pr[i] * pdfxi
+                pdfx += pipdfxi
+                glp += pipdfxi * gradlogpdf(cp[i], x)
+            end
+        end
+    end
+    if !iszero(pdfx) # else glp is already zero
+        glp /= pdfx
+    end 
+    return glp
+end
 
 _pdf!(r::AbstractArray{<:Real}, d::UnivariateMixture{Discrete}, x::UnitRange) = _mixpdf!(r, d, x)
 _pdf!(r::AbstractArray{<:Real}, d::UnivariateMixture, x::AbstractArray{<:Real}) = _mixpdf!(r, d, x)
@@ -386,7 +402,30 @@ _logpdf(d::MultivariateMixture, x::AbstractVector{<:Real}) = _mixlogpdf1(d, x)
 _pdf!(r::AbstractArray{<:Real}, d::MultivariateMixture, x::AbstractMatrix{<:Real}) = _mixpdf!(r, d, x)
 _logpdf!(r::AbstractArray{<:Real}, d::MultivariateMixture, x::AbstractMatrix{<:Real}) = _mixlogpdf!(r, d, x)
 
-gradlogpdf(d::MultivariateMixture, x::AbstractVector{<:Real}) = _mixgradlogpdf1(d, x)
+function gradlogpdf(d::MultivariateMixture, x::AbstractVector{<:Real})
+    cp = components(d)
+    pr = probs(d)
+    pdfx1 = pdf(cp[1], x)
+    pdfx = pr[1] * pdfx1
+    glp = pdfx * gradlogpdf(cp[1], x)
+    if ( iszero(pr[1]) || iszero(pdfx) )
+        glp .= zero(eltype(glp))
+    end
+    @inbounds for i in Iterators.drop(eachindex(pr, cp), 1)
+        if !iszero(pr[i])
+            pdfxi = pdf(cp[i], x)
+            if !iszero(pdfxi)
+                pipdfxi = pr[i] * pdfxi
+                pdfx += pipdfxi
+                glp .+= pipdfxi * gradlogpdf(cp[i], x)
+            end
+        end
+    end
+    if !iszero(pdfx) # else glp is already zero
+        glp ./= pdfx
+    end
+    return glp
+end
 
 ## component-wise pdf and logpdf
 
diff --git a/test/gradlogpdf.jl b/test/gradlogpdf.jl
index 3e6ab8a10..4cdc86cb8 100644
--- a/test/gradlogpdf.jl
+++ b/test/gradlogpdf.jl
@@ -28,22 +28,24 @@ using Test
 
 # Test for gradlogpdf on univariate mixture distributions against centered finite-difference on logpdf
 
-x = [-0.2, 0.3, 0.8, 1.3, 10.5]
-delta = 0.001
+x = [-0.2, 0.3, 0.8, 1.0, 1.3, 10.5]
+delta = 0.0001
 
 for d in (
     MixtureModel([Normal(-4.5, 2.0)], [1.0]),
     MixtureModel([Exponential(2.0)], [1.0]),
     MixtureModel([Uniform(-1.0, 1.0)], [1.0]),
+    MixtureModel([Normal(1//1, 2//1), Beta(2//1, 3//1), Exponential(3//2)], [3//10, 4//10, 3//10]),
     MixtureModel([Normal(-2.0, 3.5), Normal(-4.5, 2.0)], [0.0, 1.0]),
     MixtureModel([Beta(1.5, 3.0), Chi(5.0), Chisq(7.0)], [0.4, 0.3, 0.3]),
     MixtureModel([Exponential(2.0), Gamma(9.0, 0.5), Gumbel(3.5, 1.0), Laplace(7.0)], [0.3, 0.2, 0.4, 0.1]),
     MixtureModel([Logistic(-6.0), LogNormal(5.5), TDist(8.0), Weibull(2.0)], [0.3, 0.2, 0.4, 0.1])
 )
-    xs = filter(s -> insupport(d, s), x)
+    xs = filter(s -> all(insupport.(d, [s - delta, s, s + delta])), x)
     glp1 = gradlogpdf.(d, xs)
     glp2 = ( logpdf.(d, xs .+ delta) - logpdf.(d, xs .- delta) ) ./ 2delta
-    @test isapprox(glp1, glp2, atol = delta)
+    @info "Testing `gradlogpdf` on $d"
+    @test isapprox(glp1, glp2, atol = 0.01)
 end
 
 # Test for gradlogpdf on multivariate mixture distributions against centered finite-difference on logpdf

From 1d3c281e8eadf302571e0d76592427f3a1176bd3 Mon Sep 17 00:00:00 2001
From: "Ricardo M. S. Rosa" <rmsrosa@gmail.com>
Date: Sat, 27 Jan 2024 09:41:49 -0300
Subject: [PATCH 10/15] use iterate

---
 src/mixtures/mixturemodel.jl | 68 ++++++++++++++++++++++++------------
 1 file changed, 45 insertions(+), 23 deletions(-)

diff --git a/src/mixtures/mixturemodel.jl b/src/mixtures/mixturemodel.jl
index 31f1931fb..babf66a5f 100644
--- a/src/mixtures/mixturemodel.jl
+++ b/src/mixtures/mixturemodel.jl
@@ -371,19 +371,31 @@ pdf(d::UnivariateMixture, x::Real) = _mixpdf1(d, x)
 logpdf(d::UnivariateMixture, x::Real) = _mixlogpdf1(d, x)
 
 function gradlogpdf(d::UnivariateMixture, x::Real)
-    cp = components(d)
-    pr = probs(d)
-    pdfx1 = pdf(cp[1], x)
-    pdfx = pr[1] * pdfx1
-    _glp = pdfx * gradlogpdf(cp[1], x)
-    glp = (!iszero(pr[1])) && (!iszero(pdfx)) ? _glp : zero(_glp)
-    @inbounds for i in Iterators.drop(eachindex(pr, cp), 1)
-        if !iszero(pr[i])
-            pdfxi = pdf(cp[i], x)
+    ps = probs(d)
+    cs = components(d)
+
+    # `d` is expected to have at least one distribution, otherwise this will just error
+    psi, idxps = iterate(ps)
+    csi, idxcs = iterate(cs)
+    pdfx1 = pdf(csi, x)
+    pdfx = psi * pdfx1
+    glp = pdfx * gradlogpdf(csi, x)
+    if iszero(psi) || iszero(pdfx)
+        glp = zero(glp)
+    end
+    
+    while true
+        iterps = iterate(ps, idxps)
+        itercs = iterate(cs, idxcs)
+        ( (iterps !== nothing) && (itercs !== nothing) ) || break
+        psi, idxps = iterps
+        csi, idxcs = itercs
+        if !iszero(psi)
+            pdfxi = pdf(csi, x)
             if !iszero(pdfxi)
-                pipdfxi = pr[i] * pdfxi
+                pipdfxi = psi * pdfxi
                 pdfx += pipdfxi
-                glp += pipdfxi * gradlogpdf(cp[i], x)
+                glp += pipdfxi * gradlogpdf(csi, x)
             end
         end
     end
@@ -403,27 +415,37 @@ _pdf!(r::AbstractArray{<:Real}, d::MultivariateMixture, x::AbstractMatrix{<:Real
 _logpdf!(r::AbstractArray{<:Real}, d::MultivariateMixture, x::AbstractMatrix{<:Real}) = _mixlogpdf!(r, d, x)
 
 function gradlogpdf(d::MultivariateMixture, x::AbstractVector{<:Real})
-    cp = components(d)
-    pr = probs(d)
-    pdfx1 = pdf(cp[1], x)
-    pdfx = pr[1] * pdfx1
-    glp = pdfx * gradlogpdf(cp[1], x)
-    if ( iszero(pr[1]) || iszero(pdfx) )
+    ps = probs(d)
+    cs = components(d)
+
+    # `d` is expected to have at least one distribution, otherwise this will just error
+    psi, idxps = iterate(ps)
+    csi, idxcs = iterate(cs)
+    pdfx1 = pdf(csi, x)
+    pdfx = psi * pdfx1
+    glp = pdfx * gradlogpdf(csi, x)
+    if iszero(psi) || iszero(pdfx)
         glp .= zero(eltype(glp))
     end
-    @inbounds for i in Iterators.drop(eachindex(pr, cp), 1)
-        if !iszero(pr[i])
-            pdfxi = pdf(cp[i], x)
+    
+    while true
+        iterps = iterate(ps, idxps)
+        itercs = iterate(cs, idxcs)
+        ( (iterps !== nothing) && (itercs !== nothing) ) || break
+        psi, idxps = iterps
+        csi, idxcs = itercs
+        if !iszero(psi)
+            pdfxi = pdf(csi, x)
             if !iszero(pdfxi)
-                pipdfxi = pr[i] * pdfxi
+                pipdfxi = psi * pdfxi
                 pdfx += pipdfxi
-                glp .+= pipdfxi * gradlogpdf(cp[i], x)
+                glp .+= pipdfxi * gradlogpdf(csi, x)
             end
         end
     end
     if !iszero(pdfx) # else glp is already zero
         glp ./= pdfx
-    end
+    end 
     return glp
 end
 

From 7ccc7dc6d9d1bed6fab788cd962d1779b69d9fda Mon Sep 17 00:00:00 2001
From: "Ricardo M. S. Rosa" <rmsrosa@gmail.com>
Date: Sat, 27 Jan 2024 10:05:23 -0300
Subject: [PATCH 11/15] update docstring for gradlogpdf

---
 src/mixtures/mixturemodel.jl | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/mixtures/mixturemodel.jl b/src/mixtures/mixturemodel.jl
index babf66a5f..0a53c3184 100644
--- a/src/mixtures/mixturemodel.jl
+++ b/src/mixtures/mixturemodel.jl
@@ -90,8 +90,7 @@ logpdf(d::AbstractMixtureModel, x::Any)
 """
     gradlogpdf(d::Union{UnivariateMixture, MultivariateMixture}, x)
 
-Evaluate the gradient of the logarithm of the (mixed) probability density function over `x`.
-Here, `x` can be a single sample or an array of multiple samples.
+Evaluate the gradient of the logarithm of the (mixed) probability density function over a single sample `x`.
 """
 gradlogpdf(d::AbstractMixtureModel, x::Any)
 

From 2c42066b8e6e395df87d5deb3bc112addc01acbc Mon Sep 17 00:00:00 2001
From: "Ricardo M. S. Rosa" <rmsrosa@gmail.com>
Date: Sat, 27 Jan 2024 11:26:53 -0300
Subject: [PATCH 12/15] increase coverage

---
 test/gradlogpdf.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/test/gradlogpdf.jl b/test/gradlogpdf.jl
index 4cdc86cb8..a3dcd6515 100644
--- a/test/gradlogpdf.jl
+++ b/test/gradlogpdf.jl
@@ -56,7 +56,8 @@ delta = 0.001
 for d in (
     MixtureModel([MvNormal([1., 2.], [1. 0.1; 0.1 1.])], [1.0]),
     MixtureModel([MvNormal([1.0, 2.0], [0.4 0.2; 0.2 0.5]), MvNormal([2.0, 1.0], [0.3 0.1; 0.1 0.4])], [0.4, 0.6]),
-    MixtureModel([MvNormal([1.0, 2.0], [0.4 0.2; 0.2 0.5]), MvNormal([2.0, 1.0], [0.3 0.1; 0.1 0.4])], [1.0, 0.0]),
+    MixtureModel([MvNormal([3.0, 2.0], [0.2 0.3; 0.3 0.5]), MvNormal([1.0, 2.0], [0.4 0.2; 0.2 0.5]), MvNormal([2.0, 1.0], [0.3 0.1; 0.1 0.4])], [0.0, 1.0, 0.0]),
+    MixtureModel([MvNormal([1.0, 2.0], [0.4 0.2; 0.2 0.5]), MvNormal([2.0, 1.0], [0.3 0.1; 0.1 0.4])], [0.0, 1.0]),
     MixtureModel([MvTDist(5., [1., 2.], [1. 0.1; 0.1 1.])], [1.0]),
     MixtureModel([MvNormal([1.0, 2.0], [0.4 0.2; 0.2 0.5]), MvTDist(5., [1., 2.], [1. 0.1; 0.1 1.])], [0.4, 0.6])
 )

From c87d94806906107240fd3a8e57333bafcdeccbe6 Mon Sep 17 00:00:00 2001
From: "Ricardo M. S. Rosa" <rmsrosa@gmail.com>
Date: Sat, 27 Jan 2024 11:30:37 -0300
Subject: [PATCH 13/15] remove unnecessary test

---
 test/gradlogpdf.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/test/gradlogpdf.jl b/test/gradlogpdf.jl
index a3dcd6515..6a25141b6 100644
--- a/test/gradlogpdf.jl
+++ b/test/gradlogpdf.jl
@@ -57,7 +57,6 @@ for d in (
     MixtureModel([MvNormal([1., 2.], [1. 0.1; 0.1 1.])], [1.0]),
     MixtureModel([MvNormal([1.0, 2.0], [0.4 0.2; 0.2 0.5]), MvNormal([2.0, 1.0], [0.3 0.1; 0.1 0.4])], [0.4, 0.6]),
     MixtureModel([MvNormal([3.0, 2.0], [0.2 0.3; 0.3 0.5]), MvNormal([1.0, 2.0], [0.4 0.2; 0.2 0.5]), MvNormal([2.0, 1.0], [0.3 0.1; 0.1 0.4])], [0.0, 1.0, 0.0]),
-    MixtureModel([MvNormal([1.0, 2.0], [0.4 0.2; 0.2 0.5]), MvNormal([2.0, 1.0], [0.3 0.1; 0.1 0.4])], [0.0, 1.0]),
     MixtureModel([MvTDist(5., [1., 2.], [1. 0.1; 0.1 1.])], [1.0]),
     MixtureModel([MvNormal([1.0, 2.0], [0.4 0.2; 0.2 0.5]), MvTDist(5., [1., 2.], [1. 0.1; 0.1 1.])], [0.4, 0.6])
 )

From 53e0977ac2a2b564696cb650cdcf78f77a8810f9 Mon Sep 17 00:00:00 2001
From: "Ricardo M. S. Rosa" <rmsrosa@gmail.com>
Date: Mon, 29 Jan 2024 08:00:02 -0300
Subject: [PATCH 14/15] Apply suggestions from code review

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
---
 src/mixtures/mixturemodel.jl | 14 ++++----------
 1 file changed, 4 insertions(+), 10 deletions(-)

diff --git a/src/mixtures/mixturemodel.jl b/src/mixtures/mixturemodel.jl
index 0a53c3184..ff573d822 100644
--- a/src/mixtures/mixturemodel.jl
+++ b/src/mixtures/mixturemodel.jl
@@ -383,10 +383,7 @@ function gradlogpdf(d::UnivariateMixture, x::Real)
         glp = zero(glp)
     end
     
-    while true
-        iterps = iterate(ps, idxps)
-        itercs = iterate(cs, idxcs)
-        ( (iterps !== nothing) && (itercs !== nothing) ) || break
+    while (iterps = iterate(ps, idxps)) !== nothing && (itercs = iterate(cs, idxcs)) !== nothing
         psi, idxps = iterps
         csi, idxcs = itercs
         if !iszero(psi)
@@ -424,13 +421,10 @@ function gradlogpdf(d::MultivariateMixture, x::AbstractVector{<:Real})
     pdfx = psi * pdfx1
     glp = pdfx * gradlogpdf(csi, x)
     if iszero(psi) || iszero(pdfx)
-        glp .= zero(eltype(glp))
+        fill!(glp, zero(eltype(glp)))
     end
     
-    while true
-        iterps = iterate(ps, idxps)
-        itercs = iterate(cs, idxcs)
-        ( (iterps !== nothing) && (itercs !== nothing) ) || break
+    while (iterps = iterate(ps, idxps)) !== nothing && (itercs = iterate(cs, idxcs)) !== nothing
         psi, idxps = iterps
         csi, idxcs = itercs
         if !iszero(psi)
@@ -438,7 +432,7 @@ function gradlogpdf(d::MultivariateMixture, x::AbstractVector{<:Real})
             if !iszero(pdfxi)
                 pipdfxi = psi * pdfxi
                 pdfx += pipdfxi
-                glp .+= pipdfxi * gradlogpdf(csi, x)
+                glp .+= pipdfxi .* gradlogpdf(csi, x)
             end
         end
     end

From 4d4af7e0069c7f81ed70d51049fb4b3e0c26f958 Mon Sep 17 00:00:00 2001
From: "Ricardo M. S. Rosa" <rmsrosa@gmail.com>
Date: Mon, 29 Jan 2024 13:52:48 -0300
Subject: [PATCH 15/15] add tests for single distribution "mixtures"

---
 test/gradlogpdf.jl | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/test/gradlogpdf.jl b/test/gradlogpdf.jl
index 6a25141b6..faa1ca0f9 100644
--- a/test/gradlogpdf.jl
+++ b/test/gradlogpdf.jl
@@ -26,22 +26,36 @@ using Test
 @test isapprox(gradlogpdf(MvTDist(5., [1., 2.], [1. 0.1; 0.1 1.]), [0.7, 0.9]),
     [0.2150711513583442, 1.2111901681759383] ,atol=1.0e-8)
 
-# Test for gradlogpdf on univariate mixture distributions against centered finite-difference on logpdf
+# Test for gradlogpdf on univariate mixture distributions
 
 x = [-0.2, 0.3, 0.8, 1.0, 1.3, 10.5]
 delta = 0.0001
 
+for di in (
+    Normal(-4.5, 2.0),
+    Exponential(2.0),
+    Uniform(0.0, 1.0),
+    Beta(2.0, 3.0),
+    Beta(0.5, 0.5)
+)
+    d = MixtureModel([di], [1.0])
+    glp1 = gradlogpdf.(d, x)
+    glp2 = gradlogpdf.(di, x)
+    @info "Testing `gradlogpdf` on $d"
+    @test isapprox(glp1, glp2, atol = 0.01)
+end
+
 for d in (
-    MixtureModel([Normal(-4.5, 2.0)], [1.0]),
-    MixtureModel([Exponential(2.0)], [1.0]),
-    MixtureModel([Uniform(-1.0, 1.0)], [1.0]),
     MixtureModel([Normal(1//1, 2//1), Beta(2//1, 3//1), Exponential(3//2)], [3//10, 4//10, 3//10]),
     MixtureModel([Normal(-2.0, 3.5), Normal(-4.5, 2.0)], [0.0, 1.0]),
     MixtureModel([Beta(1.5, 3.0), Chi(5.0), Chisq(7.0)], [0.4, 0.3, 0.3]),
     MixtureModel([Exponential(2.0), Gamma(9.0, 0.5), Gumbel(3.5, 1.0), Laplace(7.0)], [0.3, 0.2, 0.4, 0.1]),
     MixtureModel([Logistic(-6.0), LogNormal(5.5), TDist(8.0), Weibull(2.0)], [0.3, 0.2, 0.4, 0.1])
 )
+
+    # finite differences don't handle when not in the interior of the support
     xs = filter(s -> all(insupport.(d, [s - delta, s, s + delta])), x)
+
     glp1 = gradlogpdf.(d, xs)
     glp2 = ( logpdf.(d, xs .+ delta) - logpdf.(d, xs .- delta) ) ./ 2delta
     @info "Testing `gradlogpdf` on $d"