Dev (#51)

cscherrer · web-flow · commit d9fdb6a3d3b1 · 2022-05-24T06:41:23.000-07:00
* 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 * license * affero * copyright notice * merge * Drop Create Commons * cleanup after merge * update support computations * insupport(d::SuperpositionMeasure, x) * dorp ParamWeighted * insupport(d::FactoredBase, x) * export unsafe_logdensityof * call promote_type instead of promote_rule * logdensity_def for named tuple product measures * type annotation for now * debugging * drop shows * speed up mapped arrays * throw an error for `Union{}` types * MT tests passing * updates * get tests passing * MIT license for MeasureBase * bump version * cleanup * spacing * Move ConditionalMeasure to MeasureBase * add LogarithmicNumbers * export basemeasure_sequence * update superpose * fix logdensity_rel * remove FIXME (it's fixed!!) * logdensityof(d::Density, x) * simplify insupport(::Lebesgue, ::Real) * clean up * assume insupport yields Bool * change logdensity_rel fall-through to warning and return NaN * update logdensity_rel * drop old code * fix warning * export logdensity_rel * logdensity_def(μ::Dirac, ν::Dirac, x) * logdensity_def methods * drop `static` * ]add StatsFuns * Fixing up superposition * [compat] entries * trying to speed things up * bugfixes * logdensity_rel tests * logdensity_rel tests * drop qualifier, and add a test * more tests * type constraint in "logdensityof(μ::AbstractMeasure, x)" (was piracy, oops) * add some docs * docs * docs * typo * moar speed * don't export Test * some more updates * logdensity_rel for products * `kleisli` docs * update instance_type * instance_type => Core.Typeof * `powermeasure` bug fix * fix logdensity_rel bug * get `commonbase` to take x type into account * test powers * commonbase docstring * deprecate instance_type * avoid breakage * switch || terms * @ifelse macro * simplify logdensity_rel * give up on this @ifelse business * bump version * working on likelihoods * update likelihood * powerweightedmeasure * powerweighted update * more powerweighted methods * bugfix * dropFactoredBase * drop FactoredBase * (::ProductMeasure) | constraint * update conditional measure * update Dirac * move conditional.jl down in the `include`s * Kleisli => TransitionKernel * simplify logdensity_def(::PowerMeasure, x) * rename kleisli.jl to kernel.jl * update Dirac tests * update Half * get tests passing * update kernel * Update Project.toml * no call-site inlining * restrict single-arg `kernel` to <:ParameterizedMeasure * export log_likelihood_ratio * Drop DensityKind(::Likelihood), at least for now * isfinite(x) instead of x>-Inf * add `condition` constructor * EOF newline * simplify logdensity_def for power measures * finishing up * updates * kernel stuff * kernel stuff * update showe methods * ass a TODO * use `dot` instead of `sum` * drop old code * typo * formatting * cleanup * kernel updates * uncomment * bugfix * drop old code * pretty printing * exports, cleanup * drop old for.jl * Make DensityKind(::AbstractLikelihood) = IsDensity() * update Compat version
diff --git a/Project.toml b/Project.toml
@@ -23,7 +23,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775"
 
 [compat]
-Compat = "3.35"
+Compat = "3.35, 4"
 ConstructionBase = "1.3"
 DensityInterface = "0.4"
 FillArrays = "0.12, 0.13"
diff --git a/src/combinators/for.jl b/src/combinators/for.jl
diff --git a/src/combinators/likelihood.jl b/src/combinators/likelihood.jl
@@ -121,6 +121,8 @@ struct Likelihood{K,X} <: AbstractLikelihood
     Likelihood(μ, x) = Likelihood(kernel(μ), x)
 end
 
+DensityInterface.DensityKind(::AbstractLikelihood) = IsDensity()
+
 function Pretty.quoteof(ℓ::Likelihood)
     k = Pretty.quoteof(ℓ.k)
     x = Pretty.quoteof(ℓ.x)
@@ -132,11 +134,17 @@ function Base.show(io::IO, ℓ::Likelihood)
     Pretty.pprint(io, ℓ)
 end
 
-# @inline function logdensity_def(ℓ::Likelihood, p)
-#     return logdensity_def(ℓ.k(p), ℓ.x)
-# end
+insupport(ℓ::AbstractLikelihood, p) = insupport(ℓ.k(p), ℓ.x)
+
+@inline function logdensityof(ℓ::AbstractLikelihood, p)
+    result = dynamic(unsafe_logdensityof(ℓ, p))
+    ifelse(insupport(ℓ, p) == true, result, oftype(result, -Inf))
+end
+
+@inline function unsafe_logdensityof(ℓ::AbstractLikelihood, p)
+    return unsafe_logdensityof(ℓ.k(p), ℓ.x)
+end
 
-# basemeasure(ℓ::Likelihood, p) = basemeasure(ℓ.k(p), ℓ.x)
 
 # basemeasure(ℓ::Likelihood) = @error "Likelihood requires local base measure"
 
diff --git a/src/combinators/pointwise.jl b/src/combinators/pointwise.jl
@@ -11,22 +11,20 @@ function Pretty.tile(d::PointwiseProductMeasure)
     Pretty.pair_layout(Pretty.tile(d.prior), Pretty.tile(d.likelihood), sep = " ⊙ ")
 end
 
-⊙(μ, ℓ) = pointwiseproduct(μ, ℓ)
+⊙(prior, ℓ) = pointwiseproduct(prior, ℓ)
+
+@inbounds function insupport(d::PointwiseProductMeasure, p)
+    prior, ℓ = d
+    insupport(prior, p) && insupport(ℓ, p)
+end
 
 @inline function logdensity_def(d::PointwiseProductMeasure, p)
-    μ, ℓ = d
-    logdensityof(ℓ.k(p), ℓ.x)
+    prior, ℓ = d
+    unsafe_logdensityof(ℓ, p)
 end
 
+basemeasure(d::PointwiseProductMeasure) = d.prior
+
 function gentype(d::PointwiseProductMeasure)
     gentype(d.prior)
 end
-
-@inbounds function insupport(d::PointwiseProductMeasure, p)
-    μ, ℓ = d
-    insupport(μ, p) && insupport(ℓ.k(p), ℓ.x)
-end
-
-basemeasure(d::PointwiseProductMeasure, x) = d.prior
-
-basemeasure(d::PointwiseProductMeasure) = basemeasure(d.prior)
diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl
@@ -54,7 +54,7 @@ productmeasure(tup::Tuple) = ProductMeasure(tup)
 productmeasure(f, param_maps, pars) = ProductMeasure(kernel(f, param_maps), pars)
 
 function productmeasure(k::ParameterizedTransitionKernel, pars)
-    productmeasure(k.f, k.param_maps, pars)
+    productmeasure(k.suff, k.param_maps, pars)
 end
 
 function productmeasure(f::Returns{W}, ::typeof(identity), pars) where {W<:WeightedMeasure}
@@ -117,8 +117,6 @@ end
 ###############################################################################
 # TransitionKernel
 
-kernel(f, pars::NamedTuple) = ParameterizedTransitionKernel(f, pars)
-
 # kernel(Normal(μ=2))
 function kernel(μ::M) where {M<:ParameterizedMeasure}
     kernel(M)
@@ -128,16 +126,57 @@ function kernel(d::PowerMeasure)
     Base.Fix2(powermeasure, d.axes) ∘ kernel(d.parent)
 end
 
-# kernel(Normal{(:μ,), Tuple{Int64}})
-function kernel(::Type{M}) where {M<:AbstractMeasure}
-    constructorof(M)
+function kernel(f)
+    T = Core.Compiler.return_type(f, Tuple{Any} )
+    _kernel(f, T)
+end
+
+function _kernel(f, ::Type{T}) where {T}
+    GenericTransitionKernel(f)
+end
+
+function _kernel(f, ::Type{P}) where {N,P<:ParameterizedMeasure{N}}
+    k = length(N)
+    C = constructorof(P)
+    maps = ntuple(Val(k)) do i
+        x -> @inbounds x[i]
+    end
+
+    kernel(params ∘ f, C, NamedTuple{N}(maps))
 end
 
-# kernel(::Type{P}, op::O) where {O, N, P<:ParameterizedMeasure{N}} = kernel{constructorof(P),O}(op)
+kernel(f::F, ::Type{M}; kwargs...) where {F<:Function,M} = kernel(f, M, NamedTuple(kwargs))
 
-function kernel(::Type{M}; param_maps...) where {M}
-    nt = NamedTuple(param_maps)
-    kernel(M, nt)
+function kernel(f::F, ::Type{M}, nt::NamedTuple) where {F<:Function,M}
+    ParameterizedTransitionKernel(M, f, nt)
 end
 
-kernel(k::ParameterizedTransitionKernel) = k
+function kernel(f::F, ::Type{M}, ::NamedTuple{()}) where {F<:Function,M}
+    T = Core.Compiler.return_type(f, Tuple{Any})
+    _kernel(f, M, T)
+end
+
+kernel(::Type{P}, nt::NamedTuple) where {P<:ParameterizedMeasure} = kernel(identity, P, nt)
+
+kernel(::Type{T}; kwargs...) where {T} = kernel(T, NamedTuple(kwargs))
+
+function kernel(::Type{M}, ::NamedTuple{()}) where {M}
+    C = constructorof(M)
+    TypedTransitionKernel(C, identity)
+end
+
+function _kernel(f::F, ::Type{M}, ::Type{NT}) where {M,F,N,NT<:NamedTuple{N}}
+    k = length(N)
+    maps = ntuple(Val(k)) do i
+        x -> @inbounds x[i]
+    end
+
+    ParameterizedTransitionKernel(M, values ∘ f, NamedTuple{N}(maps))
+end
+
+kernel(f::F; kwargs...) where {F<:Function} = kernel(f, NamedTuple(kwargs))
+
+function kernel(f::F, nt::NamedTuple{()}) where {F<:Function}
+    T = Core.Compiler.return_type(f, Tuple{Any})
+    _kernel(f, T)
+end
diff --git a/src/combinators/weighted.jl b/src/combinators/weighted.jl
@@ -31,7 +31,7 @@ _logweight(μ::WeightedMeasure) = μ.logweight
 basemeasure(μ::AbstractWeightedMeasure) = μ.base
 
 function Pretty.tile(d::WeightedMeasure)
-    weight = round(exp(d.logweight), sigdigits = 4)
+    weight = round(dynamic(exp(d.logweight)), sigdigits = 4)
     Pretty.pair_layout(Pretty.tile(weight), Pretty.tile(d.base), sep = " * ")
 end
 
diff --git a/src/density.jl b/src/density.jl
@@ -100,7 +100,6 @@ Define a new measure in terms of a log-density `f` over some measure `base`.
 """
 ∫exp(f::Function, μ) = ∫(logfuncdensity(f), μ)
 
-# TODO: `density` and `logdensity` functions for `DensityMeasure`
 
 """
     logdensityof(m::AbstractMeasure, x) 
diff --git a/src/domains.jl b/src/domains.jl
@@ -133,7 +133,7 @@ projectto!(x, ::Simplex) = normalize!(x, 1)
 struct Sphere <: CodimOne end
 
 function zeroset(::Sphere)
-    f(x::AbstractArray{T}) where {T} = sum(xⱼ -> xⱼ^2, x) - one(T)
+    f(x::AbstractArray{T}) where {T} = dot(x, x) - one(T)
     ∇f(x::AbstractArray{T}) where {T} = x
     ZeroSet(f, ∇f)
 end
diff --git a/src/kernel.jl b/src/kernel.jl
@@ -1,31 +1,59 @@
-# TODO: Dangerous to export this - let's not
+export AbstractTransitionKernel,
+    GenericTransitionKernel, TypedTransitionKernel, ParameterizedTransitionKernel
+
 abstract type AbstractTransitionKernel <: AbstractMeasure end
 
-struct ParameterizedTransitionKernel{F,N,T} <: AbstractTransitionKernel
+struct GenericTransitionKernel{F} <: AbstractTransitionKernel
+    f::F
+end
+
+(k::GenericTransitionKernel)(x) = k.f(x)
+
+struct TypedTransitionKernel{M,F} <: AbstractTransitionKernel
+    m::M
     f::F
+end
+
+(k::TypedTransitionKernel)(x) = (k.m ∘ k.f)(x)
+struct ParameterizedTransitionKernel{M,S,N,T} <: AbstractTransitionKernel
+    m::M
+    suff::S
     param_maps::NamedTuple{N,T}
 
     function ParameterizedTransitionKernel(
-        ::Type{F},
+        ::Type{M},
+        suff::S,
         param_maps::NamedTuple{N,T},
-    ) where {F,N,T}
-        new{Type{F},N,T}(F, param_maps)
+    ) where {M,S,N,T}
+        new{Type{M},S,N,T}(M, suff, param_maps)
     end
-    function ParameterizedTransitionKernel(f::F, param_maps::NamedTuple{N,T}) where {F,N,T}
-        new{F,N,T}(f, param_maps)
+    function ParameterizedTransitionKernel(
+        m::M,
+        suff::S,
+        param_maps::NamedTuple{N,T},
+    ) where {M,S,N,T}
+        new{M,S,N,T}(m, suff, param_maps)
     end
 end
 
 """
-    kernel(f, M)
-    kernel((f1, f2, ...), M)
+A *kernel* is a function that returns a measure.
+
+    k1 = kernel() do x
+        Normal(x, x^2)
+    end
 
-A kernel `κ = kernel(f, m)` returns a wrapper around a function `f` giving the
-parameters for a measure of type `M`, such that `κ(x) = M(f(x)...)` respective
-`κ(x) = M(f1(x), f2(x), ...)`
+    k2 = kernel(Normal) do x
+        (μ = x, σ = x^2)
+    end
+
+    k3 = kernel(Normal; μ = identity, σ = abs2)
+
+    k4 = kernel(Normal; μ = first, σ = last) do x
+        (x, x^2)
+    end
 
-If the argument is a named tuple `(;a=f1, b=f1)`, `κ(x)` is defined as
-`M(;a=f(x),b=g(x))`.
+    x = randn(); k1(x) == k2(x) == k3(x) == k4(x)
 
 This function is not exported, because "kernel" can have so many other meanings.
 See for example https://github.com/JuliaGaussianProcesses/KernelFunctions.jl for
@@ -37,23 +65,17 @@ another common use of this term.
 """
 function kernel end
 
-# kernel(Normal) do x
-#     (μ=x,σ=x^2)
-# end
-
-kernel(f, ::Type{M}) where {M} = kernel(M, f)
-
 mapcall(t, x) = map(func -> func(x), t)
 
-# (k::TransitionKernel{Type{P},<:Tuple})(x) where {P<:ParameterizedMeasure} = k.f(mapcall(k.param_maps, x)...)
+function (k::ParameterizedTransitionKernel)(x)
+    s = k.suff(x)
+    k.m(; mapcall(k.param_maps, s)...)
+end
 
-(k::ParameterizedTransitionKernel)(x) = k.f(; mapcall(k.param_maps, x)...)
+(k::AbstractTransitionKernel)(x1, x2, xs...) = k((x1, x2, xs...))
 
-(k::ParameterizedTransitionKernel)(x...) = k(x)
+(k::AbstractTransitionKernel)(; kwargs...) = k(NamedTuple(kwargs))
 
-function (k::ParameterizedTransitionKernel)(x::Tuple)
-    k.f(NamedTuple{k.param_maps}(x))
-end
 
 """
 For any `k::TransitionKernel`, `basekernel` is expected to satisfy
@@ -63,29 +85,31 @@ basekernel(k)(p) == (basemeasure ∘ k)(p)
 
 The main purpose of `basekernel` is to make it efficient to compute
 ```
-basemeasure(d::ProductMeasure) = productmeasure(basekernel(d.f), d.xs)
+basemeasure(d::ProductMeasure) == productmeasure(basekernel(d.f), d.xs)
 ```
 """
 function basekernel end
 
 # TODO: Find a way to do better than this
 basekernel(f) = basemeasure ∘ f
 
-basekernel(k::ParameterizedTransitionKernel) = kernel(basekernel(k.f), k.param_maps)
-
 basekernel(f::Returns) = Returns(basemeasure(f.value))
 
 function Base.show(io::IO, μ::AbstractTransitionKernel)
     io = IOContext(io, :compact => true)
     Pretty.pprint(io, μ)
 end
 
-function Pretty.quoteof(k::ParameterizedTransitionKernel)
-    qf = Pretty.quoteof(k.f)
-    qg = Pretty.quoteof(k.param_maps)
-    :(ParameterizedTransitionKernel($qf, $qg))
+function Pretty.tile(k::K) where {K<:AbstractTransitionKernel}
+    Pretty.list_layout(
+        Pretty.tile.([getproperty(k, p) for p in propertynames(k)]),
+        prefix = nameof(constructorof(K)),
+    )
 end
 
+
 const kleisli = kernel
 
 export kleisli
+
+kernel(k::AbstractTransitionKernel) = k
diff --git a/src/parameterized.jl b/src/parameterized.jl
@@ -10,23 +10,24 @@ function Base.propertynames(μ::ParameterizedMeasure{N}) where {N}
     return N
 end
 
-function Base.show(io::IO, μ::ParameterizedMeasure{()})
-    io = IOContext(io, :compact => true)
-    print(io, nameof(typeof(μ)), "()")
+function Pretty.tile(d::ParameterizedMeasure)
+    result = Pretty.literal(nameof(typeof(d)))
+    par = getfield(d, :par)
+    result *= Pretty.literal(sprint(show, par; context = :compact => true))
+    result
 end
 
-function Base.show(io::IO, μ::ParameterizedMeasure{N}) where {N}
-    io = IOContext(io, :compact => true)
-    print(io, nameof(typeof(μ)))
-    print(io, getfield(μ, :par))
+function Pretty.tile(d::ParameterizedMeasure{()})
+    result = Pretty.literal(nameof(typeof(d)))
+    par = getfield(d, :par)
+    result *= Pretty.literal("()")
+    result
 end
 
 # Allow things like
 #
 # julia> Normal{(:μ,)}(2)
 # Normal(μ = 2,)
-#
-
 function kernel(::Type{P}) where {N,P<:ParameterizedMeasure{N}}
     C = constructorof(P)
     _kernel(C, Val(N))
@@ -41,13 +42,17 @@ end
         C(NamedTuple{N,Tuple{T}}((arg,)))::C{N,Tuple{T}}
     end
 
-    f
+    kernel(f)
 end
 
 function (::Type{P})(nt::NamedTuple) where {N,P<:ParameterizedMeasure{N}}
     C = constructorof(P)
     arg = NamedTuple{N}(nt)
-    return C(arg)
+    return _parameterized(C, arg)
+end
+
+function _parameterized(::Type{C}, arg::NamedTuple{N,T}) where {C,N,T}
+    return C(arg)::C{N,T}
 end
 
 function (::Type{P})(args...) where {N,P<:ParameterizedMeasure{N}}
