almost done

listings/PairInference.jl

@@ -1,17 +1,20 @@
 
 @model function pair_parameter_inference(pairs)
+    N = length(pairs)
     # ______________________________________________________________________________
     # Step 1 - Sample prior parameters
 
-    # Time parameter
-    t ~ Exponential(1.0)
+    # Time parameter for each pair
+    for i ∈ 1:D
+        t[i] ~ Exponential(1.0)
+    end
     # Alignment parameters
     @submodel prefix="τ" Λ = tkf92_prior()
     # Dihedral parameters
     @submodel prefix="Θ" Ξ = jwndiff_prior()
     # Check parameter validity
-    if t ≤ 0 || any(isnan.(Ξ)) || any(isnan.(Λ))
-        Turing.@addlogprob! -Inf; return # Reject sample
+    if any(t .≤ 0) || any(isnan.(Ξ)) || any(isnan.(Λ))
+        reject_sample()
     end
 
     # ______________________________________________________________________________
@@ -27,14 +30,13 @@
     # Alignment model
     τ = TKF92([t], Λ...)
 
-    # Chain level model
-    Γ = ChainJointDistribution(ξ, τ, t)
-
     # ______________________________________________________________________________
-    # Step 3 - Observe each pair X, Y by proxy of their joint probability under Γ
-    for (X, Y) ∈ pairs
-        (X, Y) ~ Γ
+    # Step 3 - Observe each pair X, Y by proxy of their joint probability
+    for i ∈ 1:N
+        X, Y = pairs[i]
+        τ = TKF92([t[i]], Λ...)
+        (X, Y) ~ ChainJointDistribution(ξ, τ)
     end
 
-    return Γ
+    return t, Λ, Ξ
 end

listings/TrajectoryReconstruction.jl

@@ -1,6 +1,28 @@
-@model function trajectory_reconstruction(Y, Z, M_YZ, Γ)
-    @submodel X, M_XYZ = ancestor_chain_sampler(Y, Z, M_YZ)
-    left = trajectory_reconstruction(Y, X, t/2)
-    right = trajectory_reconstruction(X, Z, t/2)
 
+function trajectory_reconstruction(M_YZ::Alignment,
+                                   Y::ObservedChain, Z::ObservedChain,
+                                   ξ::MixtureProductProcess, t::Real, Λ; levels=1)
+    # ______________________________________________________________________________
+    # Base case
+    if levels == 0
+        return [Y, Z], M_YZ
+    end
+
+    # ______________________________________________________________________________
+    # Recursion
+
+    # 1. Sample X midpoint of Y and Z, as well as the
+    #    alignment M_XYZ, given the alignment M_YZ
+    X, M_XYZ = ancestor_sampling(M_YZ, Y, Z, t, ξ, Λ)
+
+    # 2. Reconstruct trajectories recursively on each branch
+    M_YX  = subalignment(M_XYZ, [2, 1])
+    traj_YX, M_Y_toX = trajectory_reconstruction(M_YX, Y, X, ξ, t/2, Λ; levels-1)
+    M_XZ = subalignment(M_XYZ, [1, 3])
+    traj_XZ, M_X_toZ = trajectory_reconstruction(M_XZ, X, Z, ξ, t/2, Λ; levels-1)
+
+    # Combine the trajectories and alignments
+    traj = [traj_YX; traj_XZ[2:end]]
+    M = glue(M_Y_toX, M_X_toZ)
+    return traj, M
 end

listings/TripleInference.jl

@@ -0,0 +1,36 @@
+
+@model function triple_alignment_sampler(Y, Z, W, Λ, ξ; max_N_X=200)
+    # _____________________________________________________________________________
+    # Step 1 - Sample prior parameters
+
+    # Time parameters
+    t_Y ~ Exponential(1.0)
+    t_Z ~ Exponential(1.0)
+    t_W ~ Exponential(1.0)
+
+    # Check parameter validity
+    if t_Y ≤ 0 || t_Z ≤ 0 || t_W ≤ 0
+        reject_sample()
+    end
+
+    # _____________________________________________________________________________
+    # Step 2 - Observe data and simultaneously construct alignment
+
+    # First, observe Y and Z and sample a triple alignment of X, Y and Z
+    τ_XYZ = TKF92([t_Y, t_Z], Λ...; known_ancestor=false)
+    (Y, Z) ~ ChainJointDistribution(ξ, τ_XYZ)
+    M_XYZ ~ ConditionedAlignmentDistribution(τ_XYZ, ξ, Y, Z)
+
+    # Construct X, the hidden ancestor chain, given alignment M_XYZ and data Y, Z
+    X = hiddenchain_from_alignment(Y, Z, t_Y, t_Z, M_XYZ, ξ)
+
+    # Finally, observe W given X and sample alignment of X and W
+    τ_XW = TKF92([t_W], Λ...; known_ancestor=true)
+    W ~ ChainTransitionDistribution(ξ, τ_XW, X)
+    M_XW ~ ConditionedAlignmentDistribution(τ_XW, ξ, X, W)
+
+    M_XYZW = combine(M_XYZ, M_XW)
+    M_YZW = subalignment(M_XYZW, [2, 3, 4])
+
+    return t_Y, t_Z, t_W, M_YZW
+end;

src/distributions/MixtureProductProcess.jl

@@ -53,6 +53,24 @@ end
     return r
 end
 
+@memoize function fulltranslogpdf!(r::AbstractMatrix{<:Real},
+                                   p::MixtureProductProcess,
+                                   t::Real,
+                                   X::ObservedChain,
+                                   Y::ObservedChain)
+    n = num_sites(X)
+    m = num_sites(Y)
+    r .= -Inf
+
+    @views jointlogpdf!(r[1:n, 1:m], p, t, X, Y)
+    @views statlogpdf!(r[1:n, m+1], p, X)
+    r[1:n, 1:m] .-= r[1:n, m+1]
+    r[1:n+1, m+1] .= 0
+    @views statlogpdf!(r[n+1, 1:m], p, Y)
+    return r
+end
+
+
 @memoize function fulljointlogpdf!(r::AbstractMatrix{<:Real},
                      p::MixtureProductProcess, t::Real,
                      X::ObservedChain,