Refactor Mixture as direct Density and standardize sampling on FluxModelsDistributions to resemble <: Sampleable (#1882)

* Refactor Mixture as direct Density and standardize sampling on FluxModelsDistributions to resemble <: Sambleable

* rm getFactorMechanics

---------

Co-authored-by: Johannes Terblanche <Affie@users.noreply.github.com>

Author: Affie

IncrementalInference/ext/IncrInfrFluxFactorsExt.jl

@@ -35,7 +35,7 @@ function Random.rand(nfb::FluxModelsDistribution, N::Integer = 1)
   numModels = length(nfb.models)
   allPreds = 1:numModels |> collect
   # TODO -- compensate when there arent enough prediction models
-  if !(N isa Nothing) && numModels < N
+  if numModels < N
     reps = (N ÷ numModels) + 1
     allPreds = repeat(allPreds, reps)
     resize!(allPreds, N)
@@ -44,20 +44,16 @@ function Random.rand(nfb::FluxModelsDistribution, N::Integer = 1)
   # can suppress shuffle for NN training purposes
   selPred = 1 < numModels && nfb.shuffle[] ? rand(allPreds, N) : view(allPreds, 1:N)
 
-  # dev function, TODO simplify to direct call 
-  _sample() = map(pred -> (nfb.models[pred])(nfb.data), selPred)
-
-  return _sample()
-  # return [_sample() for _ in 1:N]
+  # FIXME not type stable, maybe use Univariate and Multivariate from Distributions.jl
+  samples = map(pred -> (nfb.models[pred])(nfb.data), selPred)
+  dim = length(samples[1]) # dim 1 is Univariate
+  if N == 1
+    return dim == 1 ? samples[1][1] : samples[1]
+  else
+    return dim == 1 ? reduce(vcat, samples) : reduce(hcat, samples)
+  end
 end
 
-sampleTangent(M::AbstractManifold, fmd::FluxModelsDistribution, p = 0) = rand(fmd, 1)[1]
-sampleTangent(M::AbstractLieGroup, fmd::FluxModelsDistribution, p = 0) = rand(fmd, 1)[1]
-
-samplePoint(M::AbstractManifold, fmd::FluxModelsDistribution, p = 0) = rand(fmd, 1)[1]
-function samplePoint(M::AbstractLieGroup, fmd::FluxModelsDistribution, p = 0)
-  return rand(fmd, 1)[1]
-end
 
 function FluxModelsDistribution(
   inDim::NTuple{ID, Int},
@@ -138,7 +134,7 @@ Related
 Mixture, FluxModelsDistribution
 """
 function MixtureFluxModels(
-  F_::AbstractObservation,
+  F_::Type{<:AbstractObservation},
   nnModels::Vector{P},
   inDim::NTuple{ID, Int},
   data::D,
@@ -174,10 +170,6 @@ function MixtureFluxModels(
   return Mixture(F_, ntup, diversity)
 end
 
-function MixtureFluxModels(::Type{F}, w...; kw...) where {F <: AbstractObservation}
-  return MixtureFluxModels(F(LinearAlgebra.I), w...; kw...)
-end
-
 #
 
 include("FluxModelsSerialization.jl")

IncrementalInference/src/Deprecated.jl

@@ -5,32 +5,32 @@
 #TODO this looks like dead code, should be removed
 # TODO deprecate testshuffle
 function _checkErrorCCWNumerics(
-  ccwl::Union{CommonConvWrapper{F}, CommonConvWrapper{Mixture{N_, F, S, T}}},
+  ccwl::CommonConvWrapper{F},
   testshuffle::Bool = false,
-) where {N_, F <: AbstractRelativeMinimize, S, T}
+) where {F <: AbstractRelativeMinimize}
   return nothing
 end
 function _checkErrorCCWNumerics(
-  ccwl::Union{CommonConvWrapper{F}, CommonConvWrapper{Mixture{N_, F, S, T}}},
+  ccwl::CommonConvWrapper{F},
   testshuffle::Bool = false,
-) where {N_, F <: AbstractManifoldMinimize, S, T}
+) where {F <: AbstractManifoldMinimize}
   return nothing
 end
 
 #TODO this looks like dead code, should be removed
 function _perturbIfNecessary(
-  fcttype::Union{F, <:Mixture{N_, F, S, T}},
+  fcttype::AbstractRelativeMinimize,
   len::Int = 1,
   perturbation::Real = 1e-10,
-) where {N_, F <: AbstractRelativeMinimize, S, T}
+)
   return 0
 end
 
 function _perturbIfNecessary(
-  fcttype::Union{F, <:Mixture{N_, F, S, T}},
+  fcttype::AbstractManifoldMinimize,
   len::Int = 1,
   perturbation::Real = 1e-10,
-) where {N_, F <: AbstractManifoldMinimize, S, T}
+)
   return 0
 end
 #
@@ -151,6 +151,10 @@ function resetData!(vdata::DFG.FunctionNodeData)
   error("resetData!(vdata::FunctionNodeData) is deprecated, use resetData!(state::FactorState) instead")
 end
 
+function sampleTangent(x::ManifoldKernelDensity, p = mean(x))
+  error("sampleTangent(x::ManifoldKernelDensity, p) should be replaced by sampleTangent(M<:AbstractManifold, x::ManifoldKernelDensity, p)")
+end
+
 ## ================================================================================================
 ## ================================================================================================
 

IncrementalInference/src/Factors/LinearRelative.jl

@@ -32,6 +32,7 @@ LinearRelative(nm::MvNormal) = LinearRelative{length(nm.μ), typeof(nm)}(nm)
 function LinearRelative(nm::Union{<:BallTreeDensity, <:ManifoldKernelDensity})
   return LinearRelative{Ndim(nm), typeof(nm)}(nm)
 end
+LinearRelative(z) = LinearRelative{length(z)}(z)
 
 getManifold(::InstanceType{LinearRelative{N}}) where {N} = getManifold(ContinuousEuclid{N})
 

IncrementalInference/src/Factors/Mixture.jl

@@ -1,6 +1,3 @@
-
-_defaultNamesMixtures(N::Int) = ((Symbol[Symbol("c$i") for i = 1:N])...,)
-
 """
 $(TYPEDEF)
 
@@ -12,195 +9,77 @@ Notes
 - `N` is the number of components used to make the mixture, so two bumps from two Normal components means `N=2`.
 
 DevNotes
-- FIXME swap API order so Mixture of distibutions works like a distribtion, see Caesar.jl #808
-  - Should not have field mechanics.
 - TODO on sampling see #1099 and #1094 and #1069 
 
-
 Example
-```juila
+```julia
 # prior factor
-msp = Mixture(Prior, 
-              [Normal(0,0.1), Uniform(-pi/1,pi/2)],
-              [0.5;0.5])
+msp = Prior(Mixture([Normal(0,0.1), Uniform(-pi/1,pi/2)], [0.5;0.5])
 
 addFactor!(fg, [:head], msp, tags=[:MAGNETOMETER;])
 
 # Or relative
-mlr = Mixture(LinearRelative, 
+mlr = LinearRelative(Mixture(
               (correlator=AliasingScalarSampler(...), naive=Normal(0.5,5), lucky=Uniform(0,10)),
-              [0.5;0.4;0.1])
+              [0.5;0.4;0.1]
+      )
 
 addFactor!(fg, [:x0;:x1], mlr)
 ```
 """
-struct Mixture{N, F <: AbstractObservation, S, T <: Tuple} <: AbstractObservation
-  """ factor mechanics """
-  mechanics::F
-  components::NamedTuple{S, T}
-  diversity::Distributions.Categorical
-  """ dimension of factor, so range measurement would be dims=1 """
-  dims::Int
-  labels::Vector{Int}
+# we can do <: Sampleable{VF, VS}, but then all components must implement Sampleable
+struct Mixture{C <: NamedTuple, P <: Categorical}
+    components::C
+    prior::P
 end
 
+Base.length(d::Mixture) = length(d.components[1])
+#FIXME next not working
+# Base.rand(rng::AbstractRNG, d::Mixture) = rand(rng, d.components[rand(rng, d.prior)])
+Base.rand(d::Mixture) = rand(d.components[rand(d.prior)])
+getDimension(p::Mixture) = length(p)
+
+Mixture(z::NamedTuple) = Mixture(z, Categorical(length(z)))
+Mixture(z::NamedTuple, c::Union{<:Tuple, <:AbstractVector}) = Mixture(z, Categorical(c))
+
 function Mixture(
-  f::Type{F},
-  z::NamedTuple{S, T},
-  c::Distributions.DiscreteNonParametric,
-) where {F <: AbstractObservation, S, T}
-  return Mixture{length(z), F, S, T}(
-    f(LinearAlgebra.I),
-    z,
-    c,
-    size(rand(z[1], 1), 1),
-    zeros(Int, 0),
-  )
-end
-function Mixture(
-  f::F,
-  z::NamedTuple{S, T},
-  c::Distributions.DiscreteNonParametric,
-) where {F <: AbstractObservation, S, T}
-  return Mixture{length(z), F, S, T}(f, z, c, size(rand(z[1], 1), 1), zeros(Int, 0))
-end
-function Mixture(
-  f::Union{F, Type{F}},
-  z::NamedTuple{S, T},
-  c::AbstractVector{<:Real},
-) where {F <: AbstractObservation, S, T}
-  return Mixture(f, z, Categorical([c...]))
-end
-function Mixture(
-  f::Union{F, Type{F}},
-  z::NamedTuple{S, T},
-  c::NTuple{N, <:Real},
-) where {N, F <: AbstractObservation, S, T}
-  return Mixture(f, z, [c...])
-end
-function Mixture(
-  f::Union{F, Type{F}},
-  z::Tuple,
-  c::Union{
-    <:Distributions.DiscreteNonParametric,
-    <:AbstractVector{<:Real},
-    <:NTuple{N, <:Real},
-  },
-) where {F <: AbstractObservation, N}
-  return Mixture(f, NamedTuple{_defaultNamesMixtures(length(z))}(z), c)
-end
-function Mixture(
-  f::Union{F, Type{F}},
-  z::AbstractVector{<:SamplableBelief},
-  c::Union{
-    <:Distributions.DiscreteNonParametric,
-    <:AbstractVector{<:Real},
-    <:NTuple{N, <:Real},
-  },
-) where {F <: AbstractObservation, N}
-  return Mixture(f, (z...,), c)
+    z::Union{<:Tuple, <:AbstractVector},
+    args...,
+)
+    cnames = tuple((Symbol("c", i) for i = 1:length(z))...)
+    return Mixture(NamedTuple{cnames}(z), args...)
 end
 
-function Base.resize!(mp::Mixture, s::Int)
-  return resize!(mp.labels, s)
+function Mixture(f::Type{<:AbstractObservation}, args...)
+    return f(Mixture(args...))
 end
 
-_lengthOrNothing(val) = length(val)
-_lengthOrNothing(val::Nothing) = 0
-
-getManifold(m::Mixture) = getManifold(m.mechanics)
-
-# TODO make in-place memory version
-function sampleFactor(cf::CalcFactor{<:Mixture}, N::Int = 1)
-  #
-  # TODO consolidate #927, case if mechanics has a special sampler
-  # TODO slight bit of waste in computation, but easiest way to ensure special tricks in s.mechanics::F are included
-  ## example case is old FluxModelsPose2Pose2 requiring velocity
-  # FIXME better consolidation of when to pass down .mechanics, also see #1099 and #1094 and #1069
-
-  cf_ = CalcFactorNormSq(
-    cf.factor.mechanics,
-    0,
-    cf._legacyParams,
-    cf._allowThreads,
-    cf.cache,
-    cf.fullvariables,
-    cf.solvefor,
-    cf.manifold,
-    cf.measurement,
-    nothing,
-  )
-  smpls = [getSample(cf_) for _ = 1:N]
-  # smpls = Array{Float64,2}(undef,s.dims,N)
-  #out memory should be right size first
-  length(cf.factor.labels) != N ? resize!(cf.factor.labels, N) : nothing
-  cf.factor.labels .= rand(cf.factor.diversity, N)
-  M = cf.manifold
-  
-  # mixture needs to be refactored so let's make it worse :-)
-  if cf.factor.mechanics isa AbstractPriorObservation
-    samplef = samplePoint
-  elseif cf.factor.mechanics isa AbstractRelativeObservation
-    samplef = sampleTangent
-  end
-
-  for i = 1:N
-    mixComponent = cf.factor.components[cf.factor.labels[i]]
-    # measurements relate to the factor's manifold (either tangent vector or manifold point)
-    setPointsMani!(smpls,  samplef(M, mixComponent), i)
-  end
-
-  # TODO only does first element of meas::Tuple at this stage, see #1099
-  return smpls
+Base.@kwdef struct PackedMixture{C} <: PackedBelief
+  _type::String = "IncrementalInference.PackedMixture"
+  components::C
+  prior::PackedCategorical
 end
 
-function DistributedFactorGraphs.isPrior(::Type{Mixture{N, F, S, T}}) where {N, F, S, T}
-  return F <: AbstractPriorObservation
+# FIXME 🦨 use JSON properly (in JSONv1 upgrade)
+function PackedMixture(_type::String, comp_obj::JSON3.Object, prior_obj::JSON3.Object) 
+  PackedMixture(
+    _type,
+    NamedTuple(map(c -> c[1]=>convert(PackedBelief, c[2]), collect(pairs(comp_obj)))),
+    PackedCategorical(; prior_obj...),
+  )
 end
 
-"""
-$(TYPEDEF)
 
-Serialization type for `Mixture`.
-"""
-Base.@kwdef mutable struct PackedMixture <: AbstractPackedObservation
-  N::Int
-  # store the packed type for later unpacking
-  F_::String
-  S::Vector{String}
-  components::Vector{PackedBelief}
-  diversity::PackedBelief
-end
-
-function convert(::Type{<:PackedMixture}, obj::Mixture{N, F, S, T}) where {N, F, S, T}
-  allcomp = PackedBelief[]
-  for val in obj.components
-    dtr_ = convert(PackedBelief, val)
-    # FIXME ON FIRE, likely to be difficult for non-standard "Samplable" types -- e.g. Flux models in RoME
-    push!(allcomp, dtr_)
-  end
-  if hasmethod(pack, (typeof(obj.mechanics),))
-    pm = pack(obj.mechanics)
-  else
-    @warn("No pack method for mechanics type $(typeof(obj.mechanics)), using deprecated convert instead.")
-    pm = convert(DFG.convertPackedType(obj.mechanics), obj.mechanics)
-  end
-  sT = string(typeof(pm))
-  dvst = convert(PackedBelief, obj.diversity)
-  return PackedMixture(N, sT, string.(collect(S)), allcomp, dvst)
+function DFG.packDistribution(m::Mixture)
+  PackedMixture(;
+    components = map(packDistribution, m.components),
+    prior = PackedCategorical(; p = m.prior.p)
+  )
 end
 
-function convert(::Type{<:Mixture}, obj::PackedMixture)
-  N = obj.N
-  F1 = getfield(Main, Symbol(obj.F_))
-  S = (Symbol.(obj.S)...,)
-  F2 = DFG.convertStructType(F1)
-
-  components = map(c -> convert(SamplableBelief, c), obj.components)
-  diversity = convert(SamplableBelief, obj.diversity)
-  # tupcomp = (components...,)
-  ntup = NamedTuple{S}(components) # ,typeof(tupcomp)
-  return Mixture(F2, ntup, diversity)
+function DFG.unpackDistribution(pm::PackedMixture)
+  Mixture(
+    map(unpackDistribution, pm.components),
+    unpackDistribution(pm.prior)
+  )
 end
-
-#

IncrementalInference/src/Factors/PartialPrior.jl

@@ -15,6 +15,7 @@ struct PartialPrior{T <: SamplableBelief, P <: Tuple} <: AbstractPriorObservatio
 end
 
 # TODO, standardize, but shows error on testPartialNH.jl
+#FIXME should call samplePoint(M, cf.factor.Z)
 getSample(cf::CalcFactor{<:PartialPrior}) = samplePoint(cf.factor.Z)   # remove in favor of ManifoldSampling.jl
 # getManifold(pp::PartialPrior) = TranslationGroup(length(pp.partial)) # uncomment
 

IncrementalInference/src/IncrementalInference.jl

@@ -169,6 +169,7 @@ include("Factors/GenericMarginal.jl")
 include("entities/AliasScalarSampling.jl")
 include("entities/ExtDensities.jl") # used in BeliefTypes.jl::SamplableBeliefs
 include("entities/ExtFactors.jl")
+include("Factors/Mixture.jl")
 include("entities/BeliefTypes.jl")
 
 include("services/HypoRecipe.jl")
@@ -218,7 +219,6 @@ include("Variables/DefaultVariables.jl")
 
 # included factors, see RoME.jl for more examples
 include("Factors/GenericFunctions.jl")
-include("Factors/Mixture.jl")
 include("Factors/DefaultPrior.jl")
 include("Factors/LinearRelative.jl")
 include("Factors/EuclidDistance.jl")