Allow inferring the alignment with the tree fixed.

haskell/MCMC.hs

@@ -50,10 +50,6 @@ foreign import bpcall "MCMC:" fnpr_unsafe_proposal :: t -> Int -> ContextIndex -
 
 walk_tree_sample_nni_unsafe tree c = sequence_ [ nni_on_branch_unsafe tree branch c | branch <- walk_tree_path tree c]
 
-foreign import bpcall "MCMC:" walk_tree_sample_alignments :: t -> ContextIndex -> IO ()
-
-foreign import bpcall "MCMC:" realign_from_tips :: t -> ContextIndex -> IO ()
-
 foreign import bpcall "MCMC:" walk_tree_sample_NNI :: t -> ContextIndex -> IO ()
 
 foreign import bpcall "MCMC:" walk_tree_sample_NNI_and_A :: t -> ContextIndex -> IO ()

haskell/Probability/Distribution/RandomAlignment.hs

@@ -6,6 +6,7 @@ import           Bio.Alignment
 import           Control.DeepSeq
 import           Data.Array
 import           Data.Maybe (mapMaybes)
+import           MCMC
 
 import qualified Data.Map as Map
 import Data.IntMap (IntMap)
@@ -123,7 +124,9 @@ annotated_alignment_prs tree hmms model alignment = do
   property "properties" (RandomAlignmentProperties pr hmms lengthp as ls length_prs)
   return $ prs
 
-
+alignment_effect (AlignmentOnTree tree n ls as) = do
+  SamplingRate 1 $ add_move $ walk_tree_sample_alignments tree as
+  SamplingRate 0.1 $ add_move $ realign_from_tips tree as
 
 data RandomAlignment t = (HasLabels t, IsTree t) => RandomAlignment (WithBranchLengths t) IModel (Map.Map Text Int) (IntMap PairHMM)
 
@@ -132,11 +135,14 @@ instance Dist (RandomAlignment t) where
     dist_name _ = "RandomAlignment"
 
 instance Sampleable (RandomAlignment t) where
-    sample dist@(RandomAlignment tree model tip_lengths hmms) = RanDistribution3 dist do_nothing triggered_modifiable_alignment (sample_alignment tree hmms tip_lengths)
+    sample dist@(RandomAlignment tree model tip_lengths hmms) = RanDistribution3 dist alignment_effect triggered_modifiable_alignment (sample_alignment tree hmms tip_lengths)
 
 instance HasAnnotatedPdf (RandomAlignment t) where
     annotated_densities dist@(RandomAlignment tree model tip_lengths hmms) = annotated_alignment_prs tree hmms model
 
 random_alignment tree model tip_lengths = RandomAlignment tree model tip_lengths (branch_hmms model tree)
 
+foreign import bpcall "MCMC:" walk_tree_sample_alignments :: t -> IntMap PairwiseAlignment -> ContextIndex -> IO ()
+
+foreign import bpcall "MCMC:" realign_from_tips :: t -> IntMap PairwiseAlignment -> ContextIndex -> IO ()
 

haskell/Probability/Distribution/Tree.hs

@@ -86,18 +86,12 @@ triggered_modifiable_tree = triggered_modifiable_structure modifiable_tree
 
 uniform_topology_effect tree = do
 --  add_move $ walk_tree_sample_NNI_unsafe tree  -- probably we should ensure that the probability of the alignment is zero if pairwise alignments don't match?
-  add_move $ walk_tree_sample_alignments tree  -- maybe this should be elsewhere?
-                                               -- if this were elsewhere then we would have to walk the whole tree for each partition... which might not be terrible...
   add_move $ walk_tree_sample_NNI tree         -- does this handle situations with no data partitions?
 
 uniform_labelled_topology taxa = do
   topology <- sample $ uniform_topology (length taxa)
   return $ add_labels (zip [0..] taxa) topology
 
-add_alignment_moves tree = do
-  SamplingRate 1 $ add_move $ walk_tree_sample_alignments tree
-  SamplingRate 0.1 $ add_move $ realign_from_tips tree
-
 add_SPR_moves tree = do
   SamplingRate 1 $ add_move $ sample_SPR_all tree
   SamplingRate 0.5 $ add_move $ sample_SPR_flat tree
@@ -113,12 +107,8 @@ add_length_moves tree = do
   SamplingRate 1 $ add_move $ walk_tree_sample_branch_lengths tree
 
 add_tree_moves tree = do
-  SamplingRate 1 $ add_topology_moves tree
-  SamplingRate 1 $ add_length_moves tree
-
-add_tree_alignment_moves tree = do
-  SamplingRate 2 $ add_tree_moves tree
-  SamplingRate 1 $ add_alignment_moves tree
+  SamplingRate 2 $ add_topology_moves tree
+  SamplingRate 2 $ add_length_moves tree
 
 uniform_labelled_tree taxa branch_lengths_dist = do
   -- These lines should be under SamplingRate 0.0 -- but then the polytomy trees won't work
@@ -127,7 +117,7 @@ uniform_labelled_tree taxa branch_lengths_dist = do
 --  branch_lengths <- sample $ independent [branch_lengths_dist topology b | b <- [0..numBranches topology-1]]
   branch_lengths <- sample $ independent $ (getUEdgesSet topology & IntMap.fromSet (branch_lengths_dist topology))
   let tree = WithBranchLengths topology branch_lengths
-  return tree `with_tk_effect` add_tree_alignment_moves
+  return tree `with_tk_effect` add_tree_moves
 
 ----
 -- choose 2 leaves, connect them to an internal node, and put that internal node on the list of leaves

src/builtins/MCMC.cc

@@ -916,13 +916,15 @@ extern "C" closure builtin_function_walk_tree_sample_alignments(OperationArgs& A
     //------------- 1a. Get argument X -----------------//
     int tree_reg = Args.reg_for_slot(0);
 
-    int c1 = Args.evaluate(1).as_int();
+    int as_reg  = Args.reg_for_slot(1);
+
+    int c1 = Args.evaluate(2).as_int();
 
     //------------ 2. Make a TreeInterface -------------//
     context_ref C1(M, c1);
 
     MCMC::MoveStats Stats;
-    owned_ptr<Model> P(claim(new Parameters(C1, tree_reg)));
+    owned_ptr<Model> P(claim(new Parameters(C1, tree_reg, {as_reg})));
     if (P.as<Parameters>()->n_data_partitions())
     {
         walk_tree_sample_alignments(P,Stats);
@@ -940,13 +942,15 @@ extern "C" closure builtin_function_realign_from_tips(OperationArgs& Args)
     //------------- 1a. Get argument X -----------------//
     int tree_reg = Args.reg_for_slot(0);
 
-    int c1 = Args.evaluate(1).as_int();
+    int as_reg  = Args.reg_for_slot(1);
+
+    int c1 = Args.evaluate(2).as_int();
 
     //------------ 2. Make a TreeInterface -------------//
     context_ref C1(M, c1);
 
     MCMC::MoveStats Stats;
-    owned_ptr<Model> P(claim(new Parameters(C1, tree_reg)));
+    owned_ptr<Model> P(claim(new Parameters(C1, tree_reg, {as_reg})));
     if (P.as<Parameters>()->n_data_partitions())
     {
         realign_from_tips(P,Stats);

src/models/A-T-prog.cc

@@ -264,7 +264,7 @@ std::string generate_atmodel_program(const variables_map& args,
         program.let(tree_var, {var("tree")});
     branch_lengths = {var("IntMap.elems"),branch_lengths};
 
-    program.perform({var("RanSamplingRate"), 1.0, {var("PerformTKEffect"), {var("add_tree_alignment_moves"), tree_var}}});
+    program.perform({var("RanSamplingRate"), 1.0, {var("PerformTKEffect"), {var("add_tree_moves"), tree_var}}});
 
 
     set<string> used_states;