Affine (#20)

* update Affine methods

* add LinearAlgebra

* get `logjac` working properly

* bugfix

* add tests

* Try to make Julia 1.3 happy

* bump version

Author: cscherrer

Project.toml

@@ -1,13 +1,14 @@
 name = "MeasureBase"
 uuid = "fa1605e6-acd5-459c-a1e6-7e635759db14"
 authors = ["Chad Scherrer <[email protected]> and contributors"]
-version = "0.4.1"
+version = "0.4.2"
 
 [deps]
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 KeywordCalls = "4d827475-d3e4-43d6-abe3-9688362ede9f"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 MLStyle = "d8e11817-5142-5d16-987a-aa16d5891078"
 MappedArrays = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900"

src/MeasureBase.jl

@@ -62,11 +62,11 @@ include("primitives/trivial.jl")
 
 include("combinators/factoredbase.jl")
 include("combinators/weighted.jl")
-include("combinators/affine.jl")
 include("combinators/superpose.jl")
 include("combinators/product.jl")
 include("combinators/for.jl")
 include("combinators/power.jl")
+include("combinators/affine.jl")
 include("combinators/spikemixture.jl")
 include("kernel.jl")
 include("combinators/likelihood.jl")

src/combinators/affine.jl

@@ -1,6 +1,6 @@
 export Affine, AffineTransform
-
-struct AffineTransform{N,T}
+using LinearAlgebra
+@concrete terse struct AffineTransform{N,T}
     par::NamedTuple{N,T}
 end
 
@@ -22,16 +22,21 @@ Base.propertynames(d::AffineTransform{N}) where {N} = N
 (f::AffineTransform{(:μ,:σ)})(x) = f.σ * x + f.μ
 (f::AffineTransform{(:μ,:ω)})(x) = f.ω \ x + f.μ
 
+
+logjac(x::AbstractMatrix) = first(logabsdet(x))
+
+logjac(x::Number) = log(abs(x))
+
 # TODO: `log` doesn't work for the multivariate case, we need the log absolute determinant
-logjac(f::AffineTransform{(:μ,:σ)}) = log(f.σ)
-logjac(f::AffineTransform{(:μ,:ω)}) = -log(f.ω)
-logjac(f::AffineTransform{(:σ,)}) = log(f.σ)
-logjac(f::AffineTransform{(:ω,)}) = -log(f.ω)
+logjac(f::AffineTransform{(:μ,:σ)}) = logjac(f.σ)
+logjac(f::AffineTransform{(:μ,:ω)}) = -logjac(f.ω)
+logjac(f::AffineTransform{(:σ,)}) = logjac(f.σ)
+logjac(f::AffineTransform{(:ω,)}) = -logjac(f.ω)
 logjac(f::AffineTransform{(:μ,)}) = 0.0
 
 ###############################################################################
 
-struct Affine{N,M,T} <: AbstractMeasure
+@concrete terse struct Affine{N,M,T} <: AbstractMeasure
     f::AffineTransform{N,T}
     parent::M
 end
@@ -62,21 +67,40 @@ Base.propertynames(d::Affine{N}) where {N} = N ∪ (:parent,)
     end
 end
 
-# Note: We could also write
-# logdensity(d::Affine, x) = logdensity(inv(getfield(d, :f)), x)
+Base.size(d) = size(d.μ)
+Base.size(d::Affine{(:σ,)}) = (size(d.σ, 1),)
+Base.size(d::Affine{(:ω,)}) = (size(d.ω, 2),)
 
-logdensity(d::Affine{(:μ,:σ)}, x) = logdensity(d.parent, d.σ \ (x - d.μ))
-logdensity(d::Affine{(:μ,:ω)}, x) = logdensity(d.parent, d.ω * (x - d.μ))
 logdensity(d::Affine{(:σ,)}, x) = logdensity(d.parent, d.σ \ x)
 logdensity(d::Affine{(:ω,)}, x) = logdensity(d.parent, d.ω * x)
 logdensity(d::Affine{(:μ,)}, x) = logdensity(d.parent, x - d.μ) 
+logdensity(d::Affine{(:μ,:σ)}, x) = logdensity(d.parent, d.σ \ (x - d.μ)) 
+logdensity(d::Affine{(:μ,:ω)}, x) = logdensity(d.parent, d.ω * (x - d.μ)) 
+
+# logdensity(d::Affine{(:μ,:ω)}, x) = logdensity(d.parent, d.σ \ (x - d.μ))
+function logdensity(d::Affine{(:μ,:σ), Tuple{AbstractVector, AbstractMatrix}}, x)
+    z = x - d.μ
+    ldiv!(d.σ, z)
+    logdensity(d.parent, z)
+end
+    
+# logdensity(d::Affine{(:μ,:ω)}, x) = logdensity(d.parent, d.ω * (x - d.μ))
+function logdensity(d::Affine{(:μ,:ω), Tuple{AbstractVector, AbstractMatrix}}, x)
+    z = x - d.μ
+    lmul!(d.ω, z)
+    logdensity(d.parent, z)
+end
 
 basemeasure(d::Affine) = affine(getfield(d, :f), basemeasure(d.parent))
 
 # We can't do this until we know we're working with Lebesgue measure, since for
 # example it wouldn't make sense to apply a log-Jacobian to a point measure
 basemeasure(d::Affine{N,L}) where {N, L<:Lebesgue} = weightedmeasure(-logjac(d), d.parent)
 
+function basemeasure(d::Affine{N,L}) where {N, L<:PowerMeasure{typeof(identity), <:Lebesgue}}
+    weightedmeasure(-logjac(d), d.parent)
+end
+
 logjac(d::Affine) = logjac(getfield(d, :f))
 
 

test/runtests.jl

@@ -134,6 +134,19 @@ end
         @test Affine(par)(unif) == Affine(f, unif)
         @test density(Affine(f, Affine(inv(f), unif)), 0.5) == 1
     end
+
+    d = ∫exp(x -> -x^2, Lebesgue(ℝ))
+
+    μ = randn(3)
+    σ = LowerTriangular(randn(3,3))
+    ω = inv(σ)
+
+    x = randn(3)
+
+    @test logdensity(Affine((μ=μ,σ=σ), d^3), x) ≈  logdensity(Affine((μ=μ,ω=ω), d^3), x)
+    @test logdensity(Affine((σ=σ,), d^3), x) ≈  logdensity(Affine((ω=ω,), d^3), x)
+    @test logdensity(Affine((μ=μ,), d^3), x) ≈  logdensity(d^3, x-μ)
+
 end
 
 @testset "For" begin