STASH

src/combinators/hierarchical.jl

@@ -1,61 +1,67 @@
 export HierarchicalMeasure
 
+"""
+ struct HierarchicalMeasure{F,M<:AbstractMeasure,G} <: AbstractMeasure
 
-# TODO: Document and use FlattenMode
-abstract type FlattenMode end
-struct NoFlatten <: FlattenMode end
-struct AutoFlatten <: FlattenMode end
+Represents a hierarchical measure.
 
-
-struct HierarchicalMeasure{F,M<:AbstractMeasure,FM<:FlattenMode} <: AbstractMeasure
+User code should not instantiate `HierarchicalMeasure` directly, use
+[`hierarchical_measure`](@ref) instead.
+"""
+struct HierarchicalMeasure{F,M<:AbstractMeasure,G} <: AbstractMeasure
  f::F
  m::M
- flatten_mode::FM
+ flatten::G
 end
 
-# TODO: Document
-const HierarchicalProductMeasure{F,M<:AbstractMeasure} = HierarchicalMeasure{F,M,NoFlatten}
-export HierarchicalProductMeasure
-
-HierarchicalProductMeasure(f, m::AbstractMeasure) = HierarchicalMeasure(f, m, NoFlatten())
+"""
+ hierarchical_measure(f, m::AbstractMeasure, flatten)
 
-# TODO: Document
-const FlatHierarchicalMeasure{F,M<:AbstractMeasure} = HierarchicalMeasure{F,M,AutoFlatten}
-export FlatHierarchicalMeasure
+Construct a hierarchical measure from a function `f`, measure `m` and
+"""
+@inline function hierarchical_measure(f, m::AbstractMeasure, flatten)
+ F, M, G = Core.Typeof(f), Core.Typeof(m), Core.Typeof(flatten)
+ HierarchicalProductMeasure{F,M,G}(f, m, flatten)
+end
 
-FlatHierarchicalMeasure(f, m::AbstractMeasure) = HierarchicalMeasure(f, m, AutoFlatten())
 
-HierarchicalMeasure(f, m::AbstractMeasure) = FlatHierarchicalMeasure(f, m)
+#!!!!!!
+const HierarchicalProductMeasure{F,M<:AbstractMeasure} = HierarchicalMeasure{F,M,::typeof(=>)}
+const FlatHierarchicalMeasure{F,M<:AbstractMeasure} = HierarchicalMeasure{F,M,::typeof(vcat)}
 
 
 
-function _split_variate_after(::NoFlatten, μ::AbstractMeasure, x::Tuple{2})
- @assert x isa Tuple{2}
- return x[1], x[2]
+function _split_variate(::typeof(=>), ::AbstractMeasure, x::Pair)
+ return x.first, x.second
 end
 
+function _split_variate(flatten::F, μ_primary::AbstractMeasure, x) where F
+ test_primary = testvalue(μ_primary)
+ return _split_variate_byvalue(flatten, test_primary, x)
+end
 
-function _split_variate_after(::AutoFlatten, μ::AbstractMeasure, x)
- a_test = testvalue(μ)
- return _autosplit_variate_after_testvalue(a_test, x)
+function _split_variate(::Type{F}, μ::AbstractMeasure, x) where F
+ test_primary = testvalue(μ)
+ return _split_variate_byvalue(F, test_primary, x)
 end
 
-function _autosplit_variate_after_testvalue(::Any, x)
+
+function _split_variate_byvalue(::Any, x)
  @assert x isa Tuple{2}
  return x[1], x[2]
 end
 
-function _autosplit_variate_after_testvalue(a_test::AbstractVector, x::AbstractVector)
- n, m = length(eachindex(a_test)), length(eachindex(x))
+function _split_variate_byvalue(test_primary::AbstractVector, x::AbstractVector)
+ n, m = length(eachindex(test_primary)), length(eachindex(x))
  # TODO: Use getindex or view?
  return x[begin:n], x[begin+n:m]
 end
 
-function _autosplit_variate_after_testvalue(::Tuple{N}, x::Tuple{M}) where {N,M}
+function _split_variate_byvalue(::Tuple{N}, x::Tuple{M}) where {N,M}
  return ntuple(i -> x[i], Val(1:N)), ntuple(i -> x[i], Val(N+1:M))
 end
 
-@generated function _autosplit_variate_after_testvalue(::NamedTuple{names_a}, x::NamedTuple{names}) where {names_a,names}
+@generated function _split_variate_byvalue(::NamedTuple{names_a}, x::NamedTuple{names}) where {names_a,names}
  # TODO: implement
  @assert false
 end
@@ -147,7 +153,7 @@ end
 
 function _to_std_with_rest(flatten_mode::FlattenMode, ν_inner::StdMeasure, μ::AbstractMeasure, x)
  dof_μ = getdof(μ)
- x_μ, x_rest = _split_variate_after(flatten_mode, μ, x)
+ x_μ, x_rest = _split_variate(flatten_mode, μ, x)
  y = transport_to(ν_inner^dof_μ, μ, x_μ)
  return y, x_rest
 end