def make_forward_tree_defn( subst_model, tree, bin_names, with_indel_params=True, kn=True ): """Pairwise Fwd""" indel = make_indel_model_defn(with_indel_params, kn) subst = subst_model.make_fundamental_param_controller_defns(bin_names) leaf = NonParamDefn("leaf", dimensions=("edge",)) if len(bin_names) > 1: switch = ProbabilityParamDefn("bin_switch", dimensions=["locus"]) bprobs = PartitionDefn( [1.0 / len(bin_names) for bin in bin_names], name="bprobs", dimensions=["locus"], dimension=("bin", bin_names), ) edge_args = [switch, bprobs] edge_defn_constructor = EdgeSumAndAlignDefnWithBins else: edge_args = [] edge_defn_constructor = EdgeSumAndAlignDefn mprobs = subst["word_probs"] bin_data = CalcDefn(BinData)(mprobs, indel, subst["Qd"]) bin_data = bin_data.across_dimension("bin", bin_names) edge_args.extend(bin_data) (top, scores) = _recursive_defns( tree, subst, leaf, edge_defn_constructor, edge_args ) defn = FwdDefn(top) # defn = SumDefn(*scores) return AnnotateFloatDefn(defn, top)
def make_indel_model_defn(with_indel_params=True, kn=True): if kn: klass = indel_model.KnudsenMiyamotoIndelModel else: klass = indel_model.SimpleIndelModel if with_indel_params: a = IndelParameterDefn("indel_length") # P(extend indel) r = IndelParameterDefn("indel_rate") # indels per substitution return CalcDefn(klass, name="indels")(r, a) else: # not optimisable parameter, a constant. Another example is the # alignment in an LikFunc return NonParamDefn("indel_model")
def make_total_loglikelihood_defn(tree, leaves, psubs, mprobs, bprobs, bin_names, locus_names, sites_independent): fixed_motifs = NonParamDefn("fixed_motif", ["edge"]) lht = LikelihoodTreeDefn(leaves, tree=tree) plh = make_partial_likelihood_defns(tree, lht, psubs, fixed_motifs) # After the root partial likelihoods have been calculated it remains to # sum over the motifs, local sites, other sites (ie: cpus), bins and loci. # The motifs are always done first, but after that it gets complicated. # If a bin HMM is being used then the sites from the different CPUs must # be interleaved first, otherwise summing over the CPUs is done last to # minimise inter-CPU communicaton. root_mprobs = mprobs.select_from_dimension("edge", "root") lh = CalcDefn(numpy.inner, name="lh")(plh, root_mprobs) if len(bin_names) > 1: if sites_independent: site_pattern = CalcDefn(BinnedSiteDistribution, name="bdist")(bprobs) else: switch = ProbabilityParamDefn("bin_switch", dimensions=["locus"]) site_pattern = CalcDefn(PatchSiteDistribution, name="bdist")(switch, bprobs) blh = CallDefn(site_pattern, lht, name="bindex") tll = CallDefn(blh, *lh.across_dimension("bin", bin_names), **dict(name="tll")) else: lh = lh.select_from_dimension("bin", bin_names[0]) tll = CalcDefn(log_sum_across_sites, name="logsum")(lht, lh) if len(locus_names) > 1: # currently has no .make_likelihood_function() method. tll = SumDefn(*tll.across_dimension("locus", locus_names)) else: tll = tll.select_from_dimension("locus", locus_names[0]) return tll