Ancestor Callbacks are type unstable

[charlesknipp]

The title is a little misleading, but illustrates the crux of a very persistent issue. In #101, we establish that initial draws can be different than states generated from the dynamics; this implies that their types need not be the same. In cases like automatic differentiation where targeted parameters are left out of the state prior, we need disentanglement.

While #101 is merged, there are still remnants of this issue which exist uniquely in the callbacks. This was highlighted in a comment left by @THargreaves in #122 (comment), where an aesthetic change to the prior results in error due to the types.

What Gives

To collect a path of ancestors, we require pre-populating a container with roots that begin at the initial draw. As the filter progresses, subsequent states need not be the same type as the roots. For the sparse ancestry storage, this should be a brief modification to the callback:

function (c::AncestorCallback)(
    model, filter, step, state, data, ::PostPredictCallback; kwargs...
)
    particles = state.particles
    if step == 1
        N = num_particles(filter)
        c.tree = ParticleTree(getfield.(particles, :state), floor(Int64, N * log(N)))
    else
        prune!(c.tree, getfield.(particles, :ancestor))
        insert!(c.tree, getfield.(particles, :state), getfield.(particles, :ancestor))
    end
    return nothing
end

The main issue here is that we no longer store the initial draw as part of this structure. Furthermore (even though resampling shouldn't occur), we also lose a potential change in ancestry from particles at time 0 to time 1. Subsequently, we would have to change the behavior of ref_traj to something like a custom object instead of an OffsetVector.

I think this may be an easy fix, but it requires some discussion as to how we go about defining trajectories and priors.

[THargreaves]

Yeah, this is quite fiddly. At least it's a bit hidden from the user.

It's not particularly clean, but I think a reasonable solution would be:

• On step zero, we save all of the initial states in their own vector
• On step one, we create the particle tree using the (now stable) state type
• The particle tree now has two fields related to ancestry ancestor_idx and root
• If root is false, we just use the regular ancestor_idx within the tree
• If root is true the ancestor_idx points to the initial states vector
• For simplicity I suggest we don't prune the initial states since they only have a fixed memory cost

I think this is likely the best we can do without enforcing the initial states and following states have the same type.