some refactoring (#24)

* some refactoring

* pretty printing

* some updates for resettablerng

* deepcopy on iterate

* Move ResettableRNG to MeasureTheory

* Make Pretty a const

* drop extra spaces

* formatting

* formatting

* formatting

Author: cscherrer

Project.toml

@@ -12,6 +12,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 MLStyle = "d8e11817-5142-5d16-987a-aa16d5891078"
 MappedArrays = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900"
+PrettyPrinting = "54e16d92-306c-5ea0-a30b-337be88ac337"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775"
 
@@ -23,6 +24,7 @@ KeywordCalls = "0.2"
 LogExpFunctions = "0.3"
 MLStyle = "0.4"
 MappedArrays = "0.4"
+PrettyPrinting = "0.3"
 Tricks = "0.1"
 julia = "1.3"
 

src/MeasureBase.jl

@@ -16,6 +16,10 @@ export AbstractMeasure
 
 abstract type AbstractMeasure end
 
+import PrettyPrinting
+
+const Pretty = PrettyPrinting
+
 sampletype(μ::AbstractMeasure) = typeof(testvalue(μ))
 
 # sampletype(μ::AbstractMeasure) = sampletype(basemeasure(μ))
@@ -37,7 +41,6 @@ Methods for computing density relative to other measures will be
 """
 function logdensity end
 
-
 if VERSION < v"1.7.0-beta1.0"
     @eval begin
         struct Returns{T}
@@ -57,7 +60,6 @@ include("domains.jl")
 include("utils.jl")
 include("absolutecontinuity.jl")
 include("macros.jl")
-include("resettablerng.jl")
 
 include("primitive.jl")
 include("primitives/counting.jl")

src/combinators/affine.jl

@@ -4,39 +4,42 @@ using LinearAlgebra
     par::NamedTuple{N,T}
 end
 
+quoteof(f::AffineTransform) = :(AffineTransform($(quoteof(f.par))))
+
 params(f::AffineTransform) = getfield(f, :par)
 
 @inline Base.getproperty(d::AffineTransform, s::Symbol) = getfield(getfield(d, :par), s)
 
-Base.propertynames(d::AffineTransform{N}) where {N} = N 
+Base.propertynames(d::AffineTransform{N}) where {N} = N
 
-@inline Base.inv(f::AffineTransform{(:μ,:σ)}) = AffineTransform((μ = -(f.σ \ f.μ), ω = f.σ))
-@inline Base.inv(f::AffineTransform{(:μ,:ω)}) = AffineTransform((μ = - f.ω * f.μ, σ = f.ω))
+@inline Base.inv(f::AffineTransform{(:μ, :σ)}) =
+    AffineTransform((μ = -(f.σ \ f.μ), ω = f.σ))
+@inline Base.inv(f::AffineTransform{(:μ, :ω)}) = AffineTransform((μ = -f.ω * f.μ, σ = f.ω))
 @inline Base.inv(f::AffineTransform{(:σ,)}) = AffineTransform((ω = f.σ,))
 @inline Base.inv(f::AffineTransform{(:ω,)}) = AffineTransform((σ = f.ω,))
 @inline Base.inv(f::AffineTransform{(:μ,)}) = AffineTransform((μ = -f.μ,))
 
 # `size(f) == (m,n)` means `f : ℝⁿ → ℝᵐ`  
-Base.size(f::AffineTransform{(:μ,:σ)}) = size(f.σ)
-Base.size(f::AffineTransform{(:μ,:ω)}) = size(f.ω)
-Base.size(f::AffineTransform{(:σ,)})  = size(f.σ)
-Base.size(f::AffineTransform{(:ω,)})  = size(f.ω)
+Base.size(f::AffineTransform{(:μ, :σ)}) = size(f.σ)
+Base.size(f::AffineTransform{(:μ, :ω)}) = size(f.ω)
+Base.size(f::AffineTransform{(:σ,)}) = size(f.σ)
+Base.size(f::AffineTransform{(:ω,)}) = size(f.ω)
 
 function Base.size(f::AffineTransform{(:μ,)})
     (n,) = size(f.μ)
-    return (n,n)
+    return (n, n)
 end
 
 Base.size(f::AffineTransform, n::Int) = @inbounds size(f)[n]
 
 (f::AffineTransform{(:μ,)})(x) = x + f.μ
 (f::AffineTransform{(:σ,)})(x) = f.σ * x
 (f::AffineTransform{(:ω,)})(x) = f.ω \ x
-(f::AffineTransform{(:μ,:σ)})(x) = f.σ * x + f.μ
-(f::AffineTransform{(:μ,:ω)})(x) = f.ω \ x + f.μ
+(f::AffineTransform{(:μ, :σ)})(x) = f.σ * x + f.μ
+(f::AffineTransform{(:μ, :ω)})(x) = f.ω \ x + f.μ
 
 rowsize(x) = ()
-rowsize(x::AbstractArray) = (size(x,1),)
+rowsize(x::AbstractArray) = (size(x, 1),)
 
 function rowsize(f::AffineTransform)
     size_f = size(f)
@@ -46,7 +49,7 @@ function rowsize(f::AffineTransform)
 end
 
 colsize(x) = ()
-colsize(x::AbstractArray) = (size(x,2),)
+colsize(x::AbstractArray) = (size(x, 2),)
 
 function colsize(f::AffineTransform)
     size_f = size(f)
@@ -65,7 +68,7 @@ end
     return x
 end
 
-@inline function apply!(x, f::AffineTransform{(:ω,), Tuple{F}}, z) where {F<:Factorization}
+@inline function apply!(x, f::AffineTransform{(:ω,),Tuple{F}}, z) where {F<:Factorization}
     ldiv!(x, f.ω, z)
     return x
 end
@@ -75,20 +78,20 @@ end
     return x
 end
 
-@inline function apply!(x, f::AffineTransform{(:μ,:σ)}, z)
+@inline function apply!(x, f::AffineTransform{(:μ, :σ)}, z)
     apply!(x, AffineTransform((σ = f.σ,)), z)
     apply!(x, AffineTransform((μ = f.μ,)), x)
     return x
 end
 
-@inline function apply!(x, f::AffineTransform{(:μ,:ω)}, z)
+@inline function apply!(x, f::AffineTransform{(:μ, :ω)}, z)
     apply!(x, AffineTransform((ω = f.ω,)), z)
     apply!(x, AffineTransform((μ = f.μ,)), x)
     return x
 end
 
-function logjac(x::AbstractMatrix) 
-    (m,n) = size(x)
+function logjac(x::AbstractMatrix)
+    (m, n) = size(x)
     m == n && return first(logabsdet(x))
 
     # Equivalent to sum(log, svdvals(x)), but much faster
@@ -99,8 +102,8 @@ end
 logjac(x::Number) = log(abs(x))
 
 # TODO: `log` doesn't work for the multivariate case, we need the log absolute determinant
-logjac(f::AffineTransform{(:μ,:σ)}) = logjac(f.σ)
-logjac(f::AffineTransform{(:μ,:ω)}) = -logjac(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
@@ -130,16 +133,16 @@ function params(μ::Affine)
     return merge(nt1, nt2)
 end
 
-function paramnames(::Type{A}) where {N,M, A<:Affine{N,M}}
+function paramnames(::Type{A}) where {N,M,A<:Affine{N,M}}
     tuple(union(N, paramnames(M))...)
 end
 
-Base.propertynames(d::Affine{N}) where {N} = N ∪ (:parent,:f)
+Base.propertynames(d::Affine{N}) where {N} = N ∪ (:parent, :f)
 
-@inline function Base.getproperty(d::Affine, s::Symbol) 
+@inline function Base.getproperty(d::Affine, s::Symbol)
     if s === :parent
         return getfield(d, :parent)
-    elseif s === :f 
+    elseif s === :f
         return getfield(d, :f)
     else
         return getproperty(getfield(d, :f), s)
@@ -166,18 +169,18 @@ end
 
 function logdensity(d::Affine{(:μ,)}, x)
     z = x - d.μ
-    logdensity(d.parent, z) 
+    logdensity(d.parent, z)
 end
 
-function logdensity(d::Affine{(:μ,:σ)}, x)
+function logdensity(d::Affine{(:μ, :σ)}, x)
     z = d.σ \ (x - d.μ)
-    logdensity(d.parent, z) 
+    logdensity(d.parent, z)
 end
 
-function logdensity(d::Affine{(:μ,:ω)}, x)
+function logdensity(d::Affine{(:μ, :ω)}, x)
     z = d.ω * (x - d.μ)
     logdensity(d.parent, z)
-end 
+end
 
 # # logdensity(d::Affine{(:μ,:ω)}, x) = logdensity(d.parent, d.σ \ (x - d.μ))
 # @inline function logdensity(d::Affine{(:μ,:σ), P, Tuple{V,M}}, x) where {P, V<:AbstractVector, M<:AbstractMatrix}
@@ -190,7 +193,7 @@ end
 #     end
 #     sum(zⱼ -> logdensity(d.parent, zⱼ), z)
 # end
-    
+
 # # logdensity(d::Affine{(:μ,:ω)}, x) = logdensity(d.parent, d.ω * (x - d.μ))
 # @inline function logdensity(d::Affine{(:μ,:ω), P,Tuple{V,M}}, x) where {P,V<:AbstractVector, M<:AbstractMatrix}
 #     z = x - d.μ
@@ -202,31 +205,35 @@ 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)
+basemeasure(d::Affine{N,L}) where {N,L<:Lebesgue} = weightedmeasure(-logjac(d), d.parent)
 
-function basemeasure(d::Affine{N,M}) where {N,L<:Lebesgue, M<:ProductMeasure{Returns{L}}}
+function basemeasure(d::Affine{N,M}) where {N,L<:Lebesgue,M<:ProductMeasure{Returns{L}}}
     weightedmeasure(-logjac(d), d.parent)
 end
 
 logjac(d::Affine) = logjac(getfield(d, :f))
 
-function Random.rand!(rng::Random.AbstractRNG, d::Affine, x::AbstractVector{T}, z=Vector{T}(undef, size(getfield(d,:f),2))) where {T}
+function Random.rand!(
+    rng::Random.AbstractRNG,
+    d::Affine,
+    x::AbstractVector{T},
+    z = Vector{T}(undef, size(getfield(d, :f), 2))
+) where {T}
     rand!(rng, parent(d), z)
     f = getfield(d, :f)
     apply!(x, f, z)
     return x
 end
 
-
 # function Base.rand(rng::Random.AbstractRNG, ::Type{T}, d::Affine) where {T}
 #     f = getfield(d, :f)
 #     z = rand(rng, T, parent(d))
 #     apply!(x, f, z)
 #     return z
 # end
 
-supportdim(nt::NamedTuple{(:μ,:σ)}) = colsize(nt.σ)
-supportdim(nt::NamedTuple{(:μ,:ω)}) = rowsize(nt.ω)
-supportdim(nt::NamedTuple{(:σ,)})   = colsize(nt.σ)
-supportdim(nt::NamedTuple{(:ω,)})   = rowsize(nt.ω)
-supportdim(nt::NamedTuple{(:μ,)})   = size(nt.μ)
+supportdim(nt::NamedTuple{(:μ, :σ)}) = colsize(nt.σ)
+supportdim(nt::NamedTuple{(:μ, :ω)}) = rowsize(nt.ω)
+supportdim(nt::NamedTuple{(:σ,)}) = colsize(nt.σ)
+supportdim(nt::NamedTuple{(:ω,)}) = rowsize(nt.ω)
+supportdim(nt::NamedTuple{(:μ,)}) = size(nt.μ)

src/combinators/factoredbase.jl

@@ -1,10 +1,10 @@
 export FactoredBase
 
 struct FactoredBase{R,C,V,B} <: AbstractMeasure
-    inbounds :: R
-    constℓ :: C
-    varℓ :: V
-    base :: B
+    inbounds::R
+    constℓ::C
+    varℓ::V
+    base::B
 end
 
 @inline function logdensity(d::FactoredBase, x)

src/combinators/for.jl

@@ -3,7 +3,6 @@ export For
 using Random
 import Base
 
-
 """
     For(f, base...)
 
@@ -79,10 +78,9 @@ For(f, dims...) = productmeasure(i -> f(i...), zip(dims...))
 
 For(f, inds::AbstractArray) = productmeasure(f, inds)
 
-For(f, n::Int) = productmeasure(f, 1:n)
+For(f, n::Int) = productmeasure(f, Base.OneTo(n))
 For(f, dims::Int...) = productmeasure(i -> f(Tuple(i)...), CartesianIndices(dims))
 
-
 function Base.eltype(d::ProductMeasure{F,I}) where {F,I<:AbstractArray}
     return eltype(d.f(first(d.pars)))
 end

src/combinators/likelihood.jl

@@ -106,17 +106,16 @@ struct Likelihood{F,S,X}
     x::X
 end
 
-Likelihood(μ::AbstractMeasure, x) = Likelihood(kernel(μ),x)
+Likelihood(μ::AbstractMeasure, x) = Likelihood(kernel(μ), x)
 
-Likelihood(::Type{M}, x) where {M<:AbstractMeasure} = Likelihood(kernel(M),x)
+Likelihood(::Type{M}, x) where {M<:AbstractMeasure} = Likelihood(kernel(M), x)
 
-function Base.show(io::IO, ℓ::Likelihood) 
+function Base.show(io::IO, ℓ::Likelihood)
     io = IOContext(io, :compact => true)
     k, x = ℓ.k, ℓ.x
-    print(io, "Likelihood(",k,", ", x, ")")
+    print(io, "Likelihood(", k, ", ", x, ")")
 end
 
-
-function logdensity(ℓ::Likelihood, p) 
+function logdensity(ℓ::Likelihood, p)
     return logdensity(ℓ.k(p), ℓ.x)
 end

src/combinators/mapsto.jl

@@ -5,9 +5,9 @@ end
 
 export ↦, mapsto
 
-mapsto(x,y) = x ↦ y
+mapsto(x, y) = x ↦ y
 
-↦(x::X,y::Y) where {X,Y} = MapsTo{X,Y}(x,y)
+↦(x::X, y::Y) where {X,Y} = MapsTo{X,Y}(x, y)
 
 Base.first(t::MapsTo) = t.x
 Base.last(t::MapsTo) = t.y

src/combinators/pointwise.jl

@@ -1,7 +1,7 @@
 export ⊙
 
 @concrete terse struct PointwiseProductMeasure{T} <: AbstractMeasure
-    data :: T
+    data::T
 end
 
 Base.size(μ::PointwiseProductMeasure) = size(μ.data)
@@ -31,8 +31,8 @@ Base.length(m::PointwiseProductMeasure{T}) where {T} = length(m.data)
     sum((logdensity(dⱼ, x) for dⱼ in d.data))
 end
 
-function sampletype(d::PointwiseProductMeasure) 
+function sampletype(d::PointwiseProductMeasure)
     @inbounds sampletype(first(d.data))
 end
 
-basemeasure(d::PointwiseProductMeasure) =  @inbounds basemeasure(first(d.data))
+basemeasure(d::PointwiseProductMeasure) = @inbounds basemeasure(first(d.data))

src/combinators/power.jl

@@ -34,28 +34,24 @@ const PowerMeasure{F,S,T,N,A} = ProductMeasure{F,S,Fill{T,N,A}}
 Base.:^(μ::AbstractMeasure, ::Tuple{}) = μ
 
 function Base.:^(μ::AbstractMeasure, dims::Integer...)
-    return μ ^ dims
+    return μ^dims
 end
 
 function Base.:^(μ::M, dims::NTuple{N,I}) where {M<:AbstractMeasure,N,I<:Integer}
     powermeasure(μ, dims)
 end
 
 # Same as PowerMeasure
-function Base.show(io::IO, d::ProductMeasure{Returns{T},I,C}) where {T,I,C<:CartesianIndices}
-    io = IOContext(io, :compact => true)
-    print(io, d.f.f.value, " ^ ", size(d.pars))
+function Pretty.tile(d::ProductMeasure{Returns{T},I,C}) where {T,I,C<:CartesianIndices}
+    Pretty.pair_layout(Pretty.tile(d.f.f.value), Pretty.tile(size(d.pars)); sep = " ^ ")
 end
 
-function Base.show(io::IO, d::ProductMeasure{R,I,V}) where {R<:Returns,I,V<:AbstractVector}
-    io = IOContext(io, :compact => true)
-    print(io, d.f.f.value, " ^ ", length(d.pars))
+function Pretty.tile(d::ProductMeasure{R,I,V}) where {R<:Returns,I,V<:AbstractVector}
+    Pretty.pair_layout(Pretty.tile(d.f.f.value), Pretty.tile(length(d.pars)); sep = " ^ ")
 end
 
 # sampletype(d::PowerMeasure{M,N}) where {M,N} = @inbounds Array{sampletype(first(marginals(d))), N}
 
-
-
 params(d::ProductMeasure{F,S,<:Fill}) where {F,S} = params(first(marginals(d)))
 
 params(::Type{P}) where {F,S,P<:ProductMeasure{F,S,<:Fill}} = params(D)
@@ -69,16 +65,16 @@ end
 
 # Same as PowerMeasure
 @inline function _basemeasure(d::ProductMeasure{F,S,<:Fill}, b) where {F,S}
-    b ^ size(d.pars)
+    b^size(d.pars)
 end
 
 # Same as PowerMeasure
 @inline function _basemeasure(d::ProductMeasure{F,S,<:Fill}, b::FactoredBase) where {F,S}
     n = length(d.pars)
-    inbounds(x) = all(xj -> b.inbounds(xj), x)
+    inbounds(x) = all(b.inbounds, x)
     constℓ = n * b.constℓ
     varℓ() = n * b.varℓ()
-    base = b.base ^ size(d.pars)
+    base = b.base^size(d.pars)
     FactoredBase(inbounds, constℓ, varℓ, base)
 end