Add TrackedLikelihoodIntegrator which uses GriddedAbilityTracker

Author: Frankie Robertson

src/CatConfig.jl

@@ -65,6 +65,34 @@ function _find_ability_estimator_and_tracker(bits...)
     (ability_estimator, ability_tracker)
 end
 
+function collect_trackers(_)
+    return NullAbilityTracker()
+end
+
+function collect_trackers(tracker::AbilityTracker)
+    return tracker
+end
+
+function collect_trackers(config::CatConfigBase)
+    acc = NullAbilityTracker()
+    for fieldname in fieldnames(typeof(config))
+        tracker = collect_trackers(getfield(config, fieldname))
+        if !(tracker isa NullAbilityTracker)
+            acc = ConsAbilityTracker(tracker, acc)
+        end
+    end
+    return acc
+end
+
+function collect_trackers(next_item_rule::NextItemRule, ability_tracker::AbilityTracker)
+    rest = collect_trackers(next_item_rule)
+    if !(ability_tracker isa NullAbilityTracker)
+        ConsAbilityTracker(ability_tracker, rest)
+    else
+        rest
+    end
+end
+
 function CatRules(bits...)
     ability_estimator, ability_tracker = _find_ability_estimator_and_tracker(bits...)
     if ability_estimator === nothing
@@ -73,7 +101,7 @@ function CatRules(bits...)
     if ability_tracker === nothing
         error("Could not find an ability tracker in $(bits)")
     end
-    next_item = NextItemRule(bits..., ability_estimator=ability_estimator)
+    next_item = NextItemRule(bits..., ability_estimator=ability_estimator, ability_tracker=ability_tracker)
     if next_item === nothing
         error("Could not find a next item rule in $(bits)")
     end
@@ -85,7 +113,7 @@ function CatRules(bits...)
         next_item=next_item,
         termination_condition=termination_condition,
         ability_estimator=ability_estimator,
-        ability_tracker=ability_tracker
+        ability_tracker=collect_trackers(next_item, ability_tracker)
     )
 end
 

src/aggregators/Aggregators.jl

@@ -40,7 +40,7 @@ export variance, variance_given_mean, mean_1d
 # XXX: Does having a common supertype of DistributionAbilityEstimator and PointAbilityEstimator make sense?
 abstract type AbilityEstimator <: CatConfigBase end
 
-function AbilityEstimator(bits...; ability_estimator=nothing)
+function AbilityEstimator(bits...; ability_estimator=nothing, ability_tracker=nothing)
     @returnsome ability_estimator
     @returnsome find1_instance(AbilityEstimator, bits)
     item_bank = find1_type_sloppy(AbstractItemBank, bits)
@@ -73,24 +73,51 @@ end
 
 abstract type AbilityTracker <: CatConfigBase end
 
-function AbilityTracker(bits...; ability_estimator=nothing)
+function AbilityTracker(bits...; integrator=nothing, ability_estimator=nothing)
     @returnsome find1_instance(AbilityTracker, bits)
     ability_tracker = find1_type(AbilityTracker, bits)
     if (ability_tracker !== nothing)
         ability_tracker()
     end
-    NullAbilityTracker()
-    # TODO: find if ability_estimator is GriddedAbilityEstimator and then propagate stuff to GriddedAbilityTracker
+    if ability_estimator !== nothing && integrator !== nothing
+        GriddedAbilityTracker(ability_estimator, integrator)
+    else
+        NullAbilityTracker()
+    end
+end
+
+function compatible_tracker(bits...; integrator, ability_estimator, prefer_tracked)
+    ability_tracker = AbilityTracker(bits...; ability_estimator=ability_estimator)
+    if ability_tracker isa GriddedAbilityTracker && ability_tracker.integrator === integrator
+        return ability_tracker
+    end
+    if prefer_tracked
+        return AbilityTracker(bits...; integrator=integrator, ability_estimator=ability_estimator)
+    end
 end
 
 abstract type AbilityIntegrator <: CatConfigBase end
-function AbilityIntegrator(bits...; ability_estimator=nothing)
+function AbilityIntegrator(bits...; ability_estimator=nothing, prefer_tracked=false)
     @returnsome find1_instance(AbilityIntegrator, bits)
     zero_arg_intergrators = find1_type(RiemannEnumerationIntegrator, bits)
     if (zero_arg_intergrators !== nothing)
         return RiemannEnumerationIntegrator()
     end
-    @returnsome Integrator(bits...) integrator -> FunctionIntegrator(integrator)
+    integrator = Integrator(bits...)
+    if integrator === nothing
+        return nothing
+    end
+    tracker = compatible_tracker(
+        bits...;
+        integrator=integrator,
+        ability_estimator=ability_estimator,
+        prefer_tracked=prefer_tracked
+    )
+    if tracker !== nothing
+        TrackedLikelihoodIntegrator(integrator, tracker)
+    else
+        FunctionIntegrator(integrator)
+    end
 end
 
 abstract type AbilityOptimizer end

src/aggregators/ability_tracker.jl

@@ -59,6 +59,16 @@ struct NullAbilityTracker <: AbilityTracker end
 
 function track!(_, ::NullAbilityTracker) end
 
+struct ConsAbilityTracker{H <: AbilityTracker, T <: AbilityTracker} <: AbilityTracker
+    head::H
+    tail::T
+end
+
+function track!(responses, cons::ConsAbilityTracker)
+    track!(responses, cons.head)
+    track!(responses, cons.tail)
+end
+
 struct VarNormal{T <: Real}
     mean::T
     var::T

src/aggregators/ability_trackers/grid.jl

@@ -1,12 +1,14 @@
-struct GriddedAbilityTracker{AbilityEstimatorT <: DistributionAbilityEstimator, GridT <: AbstractVector{Float64}} <: AbilityTracker
+struct GriddedAbilityTracker{AbilityEstimatorT <: DistributionAbilityEstimator, IntegratorT} <: AbilityTracker
     ability_estimator::AbilityEstimatorT 
-    grid::GridT
+    integrator::IntegratorT
     cur_ability::Vector{Float64}
 end
 
-GriddedAbilityTracker(ability_estimator, grid) = GriddedAbilityTracker(ability_estimator, grid, fill(NaN, length(grid)))
+function GriddedAbilityTracker(ability_estimator::DistributionAbilityEstimator, integrator::FixedGridIntegrator)
+    GriddedAbilityTracker(ability_estimator, integrator, fill(NaN, length(integrator.grid)))
+end
 
 function track!(responses, ability_tracker::GriddedAbilityTracker)
     ability_pdf = pdf(ability_tracker.ability_estimator, responses)
-    ability_tracker.cur_ability .= ability_pdf.(ability_tracker.grid)
+    ability_tracker.cur_ability .= ability_pdf.(ability_tracker.integrator.grid)
 end

src/aggregators/integrators.jl

@@ -7,6 +7,18 @@ function (product::FunctionProduct)(x::T) where {T}
     product.f(x) * product.lh_function(x)
 end
 
+struct TrackedLikelihoodIntegrator{IntegratorT <: Integrator} <: AbilityIntegrator
+    integrator::IntegratorT
+    tracker::GriddedAbilityTracker
+end
+
+function(integrator::TrackedLikelihoodIntegrator{IntegratorT})(
+    f::F,
+    ncomp
+) where {F, IntegratorT}
+    integrator.integrator((x, y) -> f(x) * y, integrator.tracker.cur_ability, ncomp)
+end
+
 struct FunctionIntegrator{IntegratorT <: Integrator} <: AbilityIntegrator 
     integrator::IntegratorT
 end
@@ -45,7 +57,7 @@ function (integrator::RiemannEnumerationIntegrator)(
     result[]
 end
 
-function (integrator::AbilityIntegrator)(
+function (integrator::Union{RiemannEnumerationIntegrator, FunctionIntegrator})(
     f::F,
     ncomp,
     est,
@@ -54,3 +66,13 @@ function (integrator::AbilityIntegrator)(
 ) where {F}
     integrator(maybe_apply_prior(f, est), ncomp, AbilityLikelihood(tracked_responses); kwargs...)
 end
+
+function (integrator::TrackedLikelihoodIntegrator)(
+    f::F,
+    ncomp,
+    est,
+    tracked_responses::TrackedResponses;
+    kwargs...
+) where {F}
+    integrator(maybe_apply_prior(f, est), ncomp; kwargs...)
+end

src/next_item_rules/NextItemRules.jl

@@ -53,9 +53,9 @@ delegates to `ItemStrategyNextItemRule`.
 """
 abstract type NextItemRule <: CatConfigBase end
 
-function NextItemRule(bits...; ability_estimator=nothing, parallel=true)
+function NextItemRule(bits...; ability_estimator=nothing, ability_tracker=nothing, parallel=true)
     @returnsome find1_instance(NextItemRule, bits)
-    @returnsome ItemStrategyNextItemRule(bits..., ability_estimator=ability_estimator, parallel=parallel)
+    @returnsome ItemStrategyNextItemRule(bits..., ability_estimator=ability_estimator, ability_tracker=ability_tracker, parallel=parallel)
 end
 
 include("./random.jl")
@@ -135,9 +135,9 @@ struct ItemStrategyNextItemRule{NextItemStrategyT <: NextItemStrategy, ItemCrite
     criterion::ItemCriterionT
 end
 
-function ItemStrategyNextItemRule(bits...; parallel=true, ability_estimator=nothing)
+function ItemStrategyNextItemRule(bits...; parallel=true, ability_estimator=nothing, ability_tracker=nothing)
     strategy = NextItemStrategy(bits...; parallel=parallel)
-    criterion = ItemCriterion(bits...; ability_estimator=ability_estimator)
+    criterion = ItemCriterion(bits...; ability_estimator=ability_estimator, ability_tracker=ability_tracker)
     if strategy !== nothing && criterion !== nothing
         return ItemStrategyNextItemRule(strategy, criterion)
     end

src/next_item_rules/objective_function.jl

@@ -3,18 +3,18 @@ $(TYPEDEF)
 """
 abstract type ItemCriterion <: CatConfigBase end
 
-function ItemCriterion(bits...; ability_estimator=nothing)
+function ItemCriterion(bits...; ability_estimator=nothing, ability_tracker=nothing)
     @returnsome find1_instance(ItemCriterion, bits)
-    @returnsome find1_type(ItemCriterion, bits) typ -> typ(ability_estimator=ability_estimator)
-    @returnsome ExpectationBasedItemCriterion(bits...; ability_estimator=ability_estimator)
+    @returnsome find1_type(ItemCriterion, bits) typ -> typ(ability_estimator=ability_estimator, ability_tracker=ability_tracker)
+    @returnsome ExpectationBasedItemCriterion(bits...; ability_estimator=ability_estimator, ability_tracker=ability_tracker)
 end
 
 """
 $(TYPEDEF)
 """
 abstract type StateCriterion <: CatConfigBase end
 
-function StateCriterion(bits...; ability_estimator=nothing)
+function StateCriterion(bits...; ability_estimator=nothing, ability_tracker=nothing)
     @returnsome find1_instance(StateCriterion, bits)
     @returnsome find1_type(StateCriterion, bits) typ -> typ()
 end
@@ -100,12 +100,12 @@ end
 
 abstract type ExpectationBasedItemCriterion <: ItemCriterion end
 
-function ExpectationBasedItemCriterion(bits...; ability_estimator=nothing)
-    criterion = StateCriterion(bits...; ability_estimator=ability_estimator)
+function ExpectationBasedItemCriterion(bits...; ability_estimator=nothing, ability_tracker=nothing)
+    criterion = StateCriterion(bits...; ability_estimator=ability_estimator, ability_tracker=ability_tracker)
     if criterion === nothing
         return nothing
     end
-    ability_estimator = AbilityEstimator(bits..., ability_estimator=ability_estimator)
+    ability_estimator = AbilityEstimator(bits..., ability_estimator=ability_estimator, ability_tracker=ability_tracker)
     if ability_estimator === nothing
         return nothing
     end