Refactoring (#31)

* 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

* add DensityInterface dependency

* faster `rootmeasure`

* updates

* primitive measure docs

* get tests passing

* cleaning up

* formatting

* drop some unneeded methods

* update for DensityInterface

* updating densities to DensityInterface approach

* update domains

* moving things around

* updates to IntegerBounds

* more domain mucking

* fix some exports

* update deps

* update `using`

* working on tests

* Working toward tests passing

* some refactoring

* working on tests

* cleaning up

* get tests to pass

* tests passing

* update logdensity_def for pointwise product

* PrettyPrinting + tests

* speed up `rootmeasure`

* oops didn't mean to include that

* tile(::FactoredBase)

* Update Half and FactoredBase

* drop exp.jl

* simplify basemeasure

* drop old `include`

* drop redundant method

* update compat Returns

* add AbstractDensity

* Move Affine to MeasureTheory

* drop some old `For` code

* update counting measure

* add testvalue(::Type{T})

* bugfix

* some dispatch adjustments

* simplify show

* formatting

* transformed measures

* ZeroSets

* using LinearAlgebra, Statistics

* working on MeasureTheory tests

* updates

* updates

* typo

* update default to mimic Base

* fix tile(::Lebesgue)

* Add test_interface function

* small doc update

* DensityKind(::Likelihood)

* fixing show(::Likelihood)

* adding some docs

* update `rand` method

* drop old integration code

* add `rebase`

* export rebase

* law for ⊗

* typo

* Make Likelihood more flexible

* update kernel methods

* add Likelihood method (avoid stack overflow)

* refactoring

* compat

* Maybe Comat just works?

* refactoring

* some new stuff

* working on tests

* update powermeasure combinator

* bugfix

* comment out debugging lines

* more refactoring

* test @inferred basemeasure_depth(μ)

* drop `constructor` (just use ConstructionBase.constructorof`)

* debugging

* update help

* update interface

* make tests harder

* fixes

* Dirac bugfix

* formatting

* improve type inference

* working on type inference

* update interface

* udpates

* get test passing

* @test !isabstracttype(typejoin(...))

* work on show methods

* update CI

* remove old code

* update productmeasure

* prettyprinting stuff

* Drop te @constprop :aggressive stuff (maybe don't need it?)

* nerline

* dropping some old code

* update tbasemeasure_type(::PowerMeasure)

* moar tests

* update SpikeMixture

* update superpose type parameter name

* drop old tests

* func_string

* more updates

* getting closer

* almost there!

* generated function for type stability

* tests passing!

* newline

* more fixes

* exports and bugfix

* insupport(μ::Counting{T}, x) where {T<:Type}

* working on MeasureTheory tests

* MeasureTheory tests passing

* drop some old code

* inlining

* improve inference

* update `tile(::For)`

* tighten down infrerence

* update basemeasure(::For) for generators

* loosen type bound on instance_type

* drop debugging code

* small update for Likelihood, and a test

* fixing up likelihoods

* improve `basemeasure_depth` dispatch

* still some trouble with inferred basemeasure_depth

* clean up `For` dispatch

* simplify _logdensityof

* optimize for Returns{True} case

* rework basemeasure_depth

* aggressive tests passing!!

* drop type-level stuff

* drop help

* merge

* partial-static

* cleanup

* fix merge bug

* fix typo

* instance_type(T::DataType)

* fine-tuning some logdensity_def methods

* Constructor method for Counting

* some adjustments to partialstatic

* update Static.dynamic(x::PartialStatic)

* back to actual tests

* update logdensity_def(d::PowerMeasure{M}, x)

* solve!

* license

* Move `solve` to MeasureTheory

* speed up likelihoods and pointwise products

* promote_rule for PartialStatic

* some updates

* logdensityof(ℓ::AbstractLikelihood, par)

* fix test bug

Author: cscherrer

src/MeasureBase.jl

@@ -64,6 +64,7 @@ function logdensity_def end
 
 using Compat
 
+include("partial-static.jl")
 include("proxies.jl")
 include("kleisli.jl")
 include("parameterized.jl")

src/combinators/for.jl

@@ -17,12 +17,42 @@ end
 
 # For(f, gen::Base.Generator) = ProductMeasure(Base.Generator(f ∘ gen.f, gen.iter))
 
+@inline function logdensity_def(d::For{T,F,I}, x::AbstractVector{X}) where {X,T,F,I<:Tuple{<:AbstractVector}}
+    ℓ = zero(float(Core.Compiler.return_type(logdensity_def, Tuple{T,X})))
+    @inbounds for j in eachindex(x)
+        ℓ += logdensity_def(d.f(j), x[j])
+    end
+    ℓ
+end
+
+function logdensity_def(d::For, x::AbstractVector) 
+    sum(eachindex(x)) do i
+        @inbounds logdensity_def(d.f(getindex.(d.inds,i)...), x[i])
+    end
+end
+
+function logdensity_def(d::For{T,F,I}, x::AbstractArray{X}) where {T,F,I,X}
+    ℓ = zero(float(Core.Compiler.return_type(logdensity_def, Tuple{T,X})))
+
+    @inbounds for j in CartesianIndices(x)
+        i = (getindex(ind, j) for ind in d.inds)
+        ℓ += logdensity_def(d.f(i...), x[j])
+    end
+    ℓ
+end
+
 function logdensity_def(d::For{T,F,I}, x) where {N,T,F,I<:NTuple{N,<:Base.Generator}}
     sum(zip(x, d.inds...)) do (xⱼ, dⱼ...)
         logdensity_def(d.f(dⱼ...), xⱼ)
     end
 end
 
+function logdensity_def(d::For{T,F,I}, x::AbstractVector) where {N,T,F,I<:NTuple{N,<:Base.Generator}}
+    sum(zip(x, d.inds...)) do (xⱼ, dⱼ...)
+        logdensity_def(d.f(dⱼ...), xⱼ)
+    end
+end
+
 function marginals(d::For{T,F,I}) where {T,F,I}
     f(x...) = d.f(x...)::T
     mappedarray(f, d.inds...)
@@ -43,9 +73,8 @@ end
 end
 
 @inline function _basemeasure(d::For{T,F,I}, ::Type{B}, ::False) where {T,F,I,B<:AbstractMeasure}
-    new_f = basekleisli(d.f)
-    new_F = typeof(new_f)
-    For{B,new_F, I}(new_f, d.inds)
+    f = basekleisli(d.f)
+    _For(B, f, d.inds)
 end
 
 @inline function _basemeasure(d::For{T,F,I}, ::Type{B}, ::False) where {T,F,I,B}
@@ -58,8 +87,7 @@ end
 
 function _basemeasure(d::For{T,F,I}, ::Type{B}, ::False) where {N,T<:AbstractMeasure,F,I<:NTuple{N,<:Base.Generator},B}
     f = basekleisli(d.f)
-    newF = typeof(f)
-    For{B,newF,I}(f, d.inds)
+    _For(B, f, d.inds)
 end
 
 function Pretty.tile(d::For{T}) where {T}
@@ -74,6 +102,10 @@ function Pretty.tile(d::For{T}) where {T}
     )
 end
 
+function _For(::Type{T}, f::F, inds::I) where {T,F,I}
+    For{T,F,I}(f,inds)
+end
+
 
 """
     For(f, base...)

src/combinators/likelihood.jl

@@ -1,4 +1,8 @@
-export Likelihood
+export AbstractLikelihood, Likelihood
+
+abstract type AbstractLikelihood end
+
+@inline logdensityof(ℓ::AbstractLikelihood, par) = logdensity_def(ℓ, par)
 
 @doc raw"""
     Likelihood(k::AbstractKleisli, x)
@@ -101,7 +105,7 @@ and we observe `x=3`. We can compute the posterior measure on `μ` as
     julia> logdensity_def(post, 2)
     -2.5
 """
-struct Likelihood{K,X}
+struct Likelihood{K,X} <: AbstractLikelihood
     k::K
     x::X
 
@@ -124,10 +128,10 @@ function Base.show(io::IO, ℓ::Likelihood)
     Pretty.pprint(io, ℓ)
 end
 
-@inline function logdensity_def(ℓ::Likelihood, p)
+@inline function logdensity_def(ℓ::Likelihood, p::Tuple)
     return logdensity_def(ℓ.k(p), ℓ.x)
 end
 
-@inline function logdensityof(ℓ::Likelihood, p)
-    return logdensityof(ℓ.k(p), ℓ.x)
+@inline function logdensity_def(ℓ::Likelihood, p)
+    return logdensity_def(ℓ.k((p,)), ℓ.x)
 end

src/combinators/pointwise.jl

@@ -1,14 +1,20 @@
 export ⊙
 
-struct PointwiseProductMeasure{T} <: AbstractMeasure
-    data::T
+struct PointwiseProductMeasure{M,L} <: AbstractMeasure
+    measure::M
+    likelihood::L
+
+    function PointwiseProductMeasure(μ::M, ℓ::L) where {M,L}
+        @assert static_hasmethod(logdensity_def, Tuple{L, gentype(μ)})
+        return new{M,L}(μ, ℓ)
+    end
 end
 
 Base.size(μ::PointwiseProductMeasure) = size(μ.data)
 
 function Base.show(io::IO, μ::PointwiseProductMeasure)
     io = IOContext(io, :compact => true)
-    print(io, join(string.(μ.data), " ⊙ "))
+    print(io, μ.measure, " ⊙ ", μ.likelihood)
 end
 
 function Base.show_unquoted(io::IO, μ::PointwiseProductMeasure, indent::Int, prec::Int)
@@ -28,13 +34,11 @@ Base.length(m::PointwiseProductMeasure{T}) where {T} = length(m.data)
 ⊙(args...) = pointwiseproduct(args...)
 
 @inline function logdensity_def(d::PointwiseProductMeasure, x)
-    sum(d.data) do dⱼ
-        logdensity_def(dⱼ, x)
-    end
+    logdensity_def(d.measure, x) + logdensity_def(d.likelihood, x)
 end
 
 function gentype(d::PointwiseProductMeasure)
-    @inbounds gentype(first(d.data))
+    @inbounds gentype(d.measure)
 end
 
-basemeasure(d::PointwiseProductMeasure) = @inbounds basemeasure(first(d.data))
+basemeasure(d::PointwiseProductMeasure) = @inbounds basemeasure(d.measure)

src/combinators/power.jl

@@ -59,8 +59,12 @@ params(d::PowerMeasure) = params(first(marginals(d)))
     basemeasure(d.parent) ^ d.axes
 end
 
-@inline function logdensity_def(d::PowerMeasure, x)
-    sum(x) do xj
-        logdensity_def(d.parent, xj)
+@inline function logdensity_def(d::PowerMeasure{M}, x) where {M}
+    T = eltype(x)
+    ℓ = zero(float(Core.Compiler.return_type(logdensity_def, Tuple{M,T})))
+    parent = d.parent
+    @inbounds for xj in x
+        ℓ += logdensity_def(parent, xj)
     end
+    ℓ
 end

src/combinators/product.jl

@@ -21,6 +21,7 @@ function Base.:(==)(a::AbstractProductMeasure, b::AbstractProductMeasure)
 end
 Base.length(μ::AbstractProductMeasure) = length(marginals(μ))
 Base.size(μ::AbstractProductMeasure) = size(marginals(μ))
+
 basemeasure(d::AbstractProductMeasure) = productmeasure(map(basemeasure, marginals(d)))
 
 function Base.rand(rng::AbstractRNG, ::Type{T}, d::AbstractProductMeasure) where {T}

src/combinators/smart-constructors.jl

@@ -7,11 +7,8 @@ half(μ::AbstractMeasure) = Half(μ)
 ###############################################################################
 # PointwiseProductMeasure
 
-pointwiseproduct(μ::AbstractMeasure...) = PointwiseProductMeasure(μ)
-
 function pointwiseproduct(μ::AbstractMeasure, ℓ::Likelihood)
-    data = (μ, ℓ)
-    return PointwiseProductMeasure(data)
+    return PointwiseProductMeasure(μ, ℓ)
 end
 
 ###############################################################################

src/density.jl

@@ -101,36 +101,35 @@ Define a new measure in terms of a log-density `f` over some measure `base`.
 
 # TODO: `density` and `logdensity` functions for `DensityMeasure`
 
-@inline logdensityof(μ, x) = _logdensityof(μ, x)
+@inline logdensityof(μ, x) = dynamic(_logdensityof(μ, x))
 
 @inline _logdensityof(μ, x) = _logdensityof(μ, basemeasure(μ, x), x)
 
 @inline function  _logdensityof(μ, α, x)
-    ℓ = dynamic(logdensity_def(μ, x))
-    L = typeof(ℓ)
-    _logdensityof(μ, α, x, ℓ)::L
+    _logdensityof(μ, α, x, partialstatic(logdensity_def(μ, x)))
 end
 
 @inline function _logdensityof(μ::M, β::M, x, ℓ) where {M}
     return ℓ
 end
 
 @inline function _logdensityof(μ::M, β, x, ℓ) where {M}
-    n = basemeasure_depth(μ) - static(1)
+    n = static(basemeasure_depth(β))
     _logdensityof(β, basemeasure(β,x), x, ℓ, n)
 end
 
-@generated function _logdensityof(μ, β, x, ℓ::T, ::StaticInt{n}) where {n,T}
+@generated function _logdensityof(μ, β, x, ℓ, ::StaticInt{n}) where {n}
     nsteps = max(n, 0)
     quote
         $(Expr(:meta, :inline))
+        # @show ℓ
         Base.Cartesian.@nexprs $nsteps i -> begin
-            Δℓ = oftype(ℓ, logdensity_def(μ, x))
+            Δℓ = logdensity_def(μ, x)
             # @show μ
             # @show Δℓ
             # println()
             μ,β = β, basemeasure(β, x)
-            ℓ += Δℓ
+            ℓ += partialstatic(Δℓ)
         end
         return ℓ
     end
@@ -164,7 +163,6 @@ export logdensityof
 export density_def
 
 density_def(μ, ν::AbstractMeasure, x) = exp(logdensity_def(μ, ν, x))
-
 density_def(μ, x) = exp(logdensity_def(μ, x))
 
 """

src/interface.jl

@@ -15,7 +15,6 @@ function test_interface(μ::M) where {M}
         μ = $μ
         @testset "$μ" begin
             μ = $μ
-            M = $M
 
             ###########################################################################
             # basemeasure_depth
@@ -31,14 +30,13 @@ function test_interface(μ::M) where {M}
             ###########################################################################
             # testvalue, logdensityof
 
-            x = testvalue(μ)
+            x = @inferred testvalue(μ)
             β = @inferred basemeasure(μ, x)
 
             ℓμ = @inferred logdensityof(μ, x)
             ℓβ = @inferred logdensityof(β, x)
 
             @test ℓμ ≈ logdensity_def(μ, x) + ℓβ
-
         end
     end
 end

src/parameterized.jl

@@ -29,13 +29,19 @@ end
 export kleisli
 
 function kleisli(::Type{P}) where {N,P<:ParameterizedMeasure{N}}
-    C = constructorof(P)
-    function(args...) C(NamedTuple{N}(args...)) end
+    C = constructorof(P) 
+    _kleisli(C, Val(N))
+end
+
+@inline function _kleisli(::Type{C}, ::Val{N}) where {C,N}
+    @inline function(args::T) where {T<:Tuple}
+        C(NamedTuple{N,T}(args))::C{N,T}
+    end
 end
 
 function (::Type{P})(args...) where {N,P<:ParameterizedMeasure{N}}
     C = constructorof(P)
-    return C(NamedTuple{N}(args...))
+    return C(NamedTuple{N}(args))::C{N,typeof(args)}
 end
 
 (::Type{P})(; kwargs...) where {P<:ParameterizedMeasure} = P(NamedTuple(kwargs))