More careful evaluation + add tests

Author: ararslan

src/truncated/normal.jl

@@ -118,15 +118,34 @@ function entropy(d::Truncated{<:Normal{<:Real},Continuous})
     0.5 * (log2π + 1.) + log(σ * z) + (aφa - bφb) / (2.0 * z)
 end
 
-function mgf(d::Truncated{Normal{T},Continuous}, t::Real) where {T}
+function mgf(d::Truncated{<:Normal{<:Real},Continuous}, t::Real)
+    a, b = extrema(d)
+    T = promote_type(partype(d), typeof(t), typeof(a))
+    if isnan(a) || isnan(b)  # TODO: Disallow constructing `Truncated` with a `NaN` bound?
+        return T(NaN)
+    elseif isinf(a) && isinf(b) && a != b
+        # Distribution is `Truncated`-wrapped but not actually truncated
+        return T(mgf(d.untruncated, t))
+    elseif a == b
+        # Truncated to a Dirac distribution; this is `mgf(Dirac(a), t)`
+        return exp(a * t)
+    end
     d0 = d.untruncated
     μ = mean(d0)
     σ = std(d0)
     σ²t = σ^2 * t
-    a = (minimum(d) - μ) / σ - σ²t
-    b = (maximum(d) - μ) / σ - σ²t
+    a′ = (a - μ) / σ
+    b′ = (b - μ) / σ
     stdnorm = Normal{T}(zero(T), one(T))
-    return exp(t * (μ + σ²t / 2) + logdiffcdf(stdnorm, b, a) - d.logtp)
+    # log((Φ(b′ - σ²t) - Φ(a′ - σ²t)) / (Φ(b′) - Φ(a′)))
+    logratio = if isfinite(a) && isfinite(b)  # doubly truncated
+        logdiffcdf(stdnorm, b′ - σ²t, a′ - σ²t) - logdiffcdf(stdnorm, b′, a′)
+    elseif isfinite(a)  # left truncated: b = ∞, Φ(b′) = Φ(b′ - σ²t) = 1
+        logccdf(stdnorm, a′ - σ²t) - logccdf(stdnorm, a′)
+    else  # isfinite(b), right truncated: a = ∞, Φ(a′) = Φ(a′ - σ²t) = 0
+        logcdf(stdnorm, b′ - σ²t) - logcdf(stdnorm, b′)
+    end
+    return exp(t * (μ + σ²t / 2) + logratio)
 end
 
 ### sampling

test/truncated/normal.jl

@@ -69,3 +69,44 @@ end
         @test isfinite(pdf(trunc, x))
     end
 end
+
+bigly(x) = x
+bigly(x::Symbol) = x === :π || x === :ℯ ? Expr(:call, :big, x) : x
+bigly(x::Real) = Expr(:call, :big, x)
+bigly(x::Expr) = (map!(bigly, x.args, x.args); x)
+macro bigly(ex)
+    return esc(bigly(ex))
+end
+
+@testset "Truncated normal MGF" begin
+    sqrt2 = sqrt(big(2))
+    invsqrt2 = inv(sqrt2)
+    inv2sqrt2 = inv(big(2) * sqrt2)
+    twoerfsqrt2 = big(2) * erf(sqrt2)
+
+    for T in (Float32, Float64, BigFloat)
+        d = truncated(Normal{T}(zero(T), one(T)), -2, 2)
+        @test @inferred(mgf(d, 0)) == 1
+        @test @inferred(mgf(d, 1)) ≈ @bigly sqrt(ℯ) * (erf(invsqrt2) + erf(3 * invsqrt2)) / twoerfsqrt2
+        @test @inferred(mgf(d, 2.5)) ≈ @bigly exp(25//8) * (erf(9 * inv2sqrt2) - erf(inv2sqrt2)) / twoerfsqrt2
+    end
+
+    d = truncated(Normal(3, 10), 7, 8)
+    @test mgf(d, 0) == 1
+    @test mgf(d, 1) == 0
+
+    d = truncated(Normal(27, 3); lower=0)
+    @test mgf(d, 0) == 1
+    @test mgf(d, 1) ≈ @bigly 2 * exp(63//2) / (1 + erf(9 * invsqrt2))
+    @test mgf(d, 2.5) ≈ @bigly 2 * exp(765//8) / (1 + erf(9 * invsqrt2))
+
+    d = truncated(Normal(-5, 1); upper=-10)
+    @test mgf(d, 0) == 1
+    @test mgf(d, 1) ≈ @bigly erfc(3 * sqrt2) / (exp(9//2) * erfc(5 * invsqrt2))
+
+    @test isnan(mgf(truncated(Normal(); upper=NaN), 0))
+
+    @test mgf(truncated(Normal(), -Inf, Inf), 1) == mgf(Normal(), 1)
+
+    @test mgf(truncated(Normal(), 2, 2), 1) == exp(2)
+end