Optimise cell component responsibility with expanded likelihood matrix

Author: Daniel Cooke

src/core/models/genotype/variational_bayes_mixture_mixture_model.cpp

@@ -30,8 +30,9 @@ VariationalBayesMixtureMixtureModel::evaluate(const LogProbabilityVector& genoty
                                               std::vector<LogProbabilityVector> seeds) const
 {
     const auto group_log_priors = to_logs(group_priors);
+    const auto expanded_log_likelihoods = expand(log_likelihoods);
     const auto evaluate_seed = [&] (auto&& seed) {
-        return this->evaluate(genotype_log_priors, log_likelihoods, group_log_priors, group_concentrations, mixture_concentrations, std::move(seed)); };
+        return this->evaluate(genotype_log_priors, log_likelihoods, expanded_log_likelihoods, group_log_priors, group_concentrations, mixture_concentrations, std::move(seed)); };
     std::vector<PointInferences> seed_inferences(seeds.size());
     if (options_.parallel_execution) {
         parallel_transform(std::make_move_iterator(std::begin(seeds)), std::make_move_iterator(std::end(seeds)), std::begin(seed_inferences), evaluate_seed);
@@ -163,6 +164,30 @@ VariationalBayesMixtureMixtureModel::to_logs(const GroupOptionalPriorArray& prio
     return result;
 }
 
+VariationalBayesMixtureMixtureModel::ExpandedHaplotypeLikelihoodMatrix 
+VariationalBayesMixtureMixtureModel::expand(const HaplotypeLikelihoodMatrix& likelihoods) const
+{
+    const auto S = likelihoods.size();
+    const auto G = likelihoods[0].size();
+    const auto T = likelihoods[0][0].size();
+    ExpandedHaplotypeLikelihoodMatrix result(S, ExpandedGroupLikelihoodVector(T));
+    for (std::size_t s {0}; s < S; ++s) {
+        const auto N = likelihoods[s][0][0][0].size();
+        for (std::size_t t {0}; t < T; ++t) {
+            const auto K = likelihoods[0][0][t].size();
+            result[s][t].assign(K, ExpandedGenotype(N, ExpandedLikelihood(G)));
+            for (std::size_t k {0}; k < K; ++k) {
+                for (std::size_t g {0}; g < G; ++g) {
+                    for (std::size_t n {0}; n < N; ++n) {
+                        result[s][t][k][n][g] = likelihoods[s][g][t][k][n];
+                    }
+                }
+            }
+        }
+    }
+    return result;
+}
+
 namespace {
 
 VariationalBayesMixtureMixtureModel::ProbabilityVector&
@@ -229,6 +254,7 @@ bool all_equal_sizes(const VariationalBayesMixtureMixtureModel::MixtureConcentra
 VariationalBayesMixtureMixtureModel::PointInferences
 VariationalBayesMixtureMixtureModel::evaluate(const LogProbabilityVector& genotype_log_priors,
                                               const HaplotypeLikelihoodMatrix& log_likelihoods,
+                                              const ExpandedHaplotypeLikelihoodMatrix& expanded_log_likelihoods,
                                               const GroupOptionalLogPriorArray& group_log_priors,
                                               const GroupConcentrationVector& prior_group_concentrations,
                                               const MixtureConcentrationArray& prior_mixture_concentrations,
@@ -242,7 +268,7 @@ VariationalBayesMixtureMixtureModel::evaluate(const LogProbabilityVector& genoty
     latents.group_responsibilities = init_responsibilities(group_log_priors, prior_group_concentrations, prior_mixture_concentrations,
                                                            latents.genotype_posteriors, log_likelihoods);
     latents.component_responsibilities = init_responsibilities(prior_group_concentrations, prior_mixture_concentrations,
-                                                               latents.genotype_posteriors, latents.group_responsibilities, log_likelihoods);
+                                                               latents.genotype_posteriors, latents.group_responsibilities, expanded_log_likelihoods);
     auto prev_evidence = std::numeric_limits<double>::lowest();
     for (unsigned i {0}; i < options_.max_iterations; ++i) {
         update_genotype_log_posteriors(latents.genotype_log_posteriors, genotype_log_priors,
@@ -266,7 +292,7 @@ VariationalBayesMixtureMixtureModel::evaluate(const LogProbabilityVector& genoty
                                 latents.mixture_concentrations, latents.genotype_posteriors,
                                 latents.component_responsibilities, log_likelihoods);
         update_responsibilities(latents.component_responsibilities, latents.group_concentrations, latents.mixture_concentrations,
-                                latents.genotype_posteriors, latents.group_responsibilities, log_likelihoods);
+                                latents.genotype_posteriors, latents.group_responsibilities, expanded_log_likelihoods);
     }
     return {std::move(latents), prev_evidence};
 }
@@ -377,7 +403,7 @@ VariationalBayesMixtureMixtureModel::init_responsibilities(const GroupConcentrat
                                                            const MixtureConcentrationArray& mixture_concentrations,
                                                            const ProbabilityVector& genotype_priors,
                                                            const GroupResponsibilityVector& group_responsibilities,
-                                                           const HaplotypeLikelihoodMatrix& log_likelihoods) const
+                                                           const ExpandedHaplotypeLikelihoodMatrix& log_likelihoods) const
 {
     const auto S = log_likelihoods.size();
     const auto T = group_concentrations.size();
@@ -389,7 +415,7 @@ VariationalBayesMixtureMixtureModel::init_responsibilities(const GroupConcentrat
     }
     ComponentResponsibilityMatrix result(S, ComponentResponsibilityVectorArray(T, ComponentResponsibilityVector(K)));
     for (std::size_t s {0}; s < S; ++s) {
-        const auto N = log_likelihoods[s][0][0][0].size();
+        const auto N = log_likelihoods[s][0][0].size();
         for (std::size_t t {0}; t < T; ++t) {
             for (std::size_t k {0}; k < K; ++k) {
                 result[s][t][k].resize(N);
@@ -423,21 +449,23 @@ VariationalBayesMixtureMixtureModel::update_responsibilities(ComponentResponsibi
                                                              const MixtureConcentrationArray& mixture_concentrations,
                                                              const ProbabilityVector& genotype_posteriors,
                                                              const GroupResponsibilityVector& group_responsibilities,
-                                                             const HaplotypeLikelihoodMatrix& log_likelihoods) const
+                                                             const ExpandedHaplotypeLikelihoodMatrix& log_likelihoods) const
 {
     const auto T = group_concentrations.size();
     const auto S = log_likelihoods.size();
     const auto max_K = result[0][0].size();
+    std::vector<float> float_genotype_posteriors {std::cbegin(genotype_posteriors), std::cend(genotype_posteriors)};
     for (std::size_t s {0}; s < S; ++s) {
-        const auto N = log_likelihoods[s][0][0][0].size();
+        const auto N = log_likelihoods[s][0][0].size();
         for (std::size_t t {0}; t < T; ++t) {
             const auto ln_exp_pi = dirichlet_expectation_log(mixture_concentrations[s][t]);
             const auto K = ln_exp_pi.size();
             for (std::size_t k {0}; k < max_K; ++k) {
                 for (std::size_t n {0}; n < N; ++n) {
                     if (t == 0) result[s][t][k][n] = 0;
                     if (k < K) {
-                        result[s][t][k][n] += ln_exp_pi[k] + inner_product(genotype_posteriors, log_likelihoods[s], t, k, n);
+                        //result[s][t][k][n] += ln_exp_pi[k] + inner_product(genotype_posteriors, log_likelihoods[s], t, k, n);
+                        result[s][t][k][n] += ln_exp_pi[k] + inner_product(float_genotype_posteriors, log_likelihoods[s][t][k][n]);
                     } else {
                         result[s][t][k][n] = options_.null_log_probability;
                     }

src/core/models/genotype/variational_bayes_mixture_mixture_model.hpp

@@ -135,13 +135,21 @@ class VariationalBayesMixtureMixtureModel
     using GroupOptionalLogPriorVector = boost::optional<LogProbabilityVector>; // One element per group
     using GroupOptionalLogPriorArray = std::vector<GroupOptionalLogPriorVector>; // One element per sample
 
+    using ExpandedLikelihood = std::vector<float>; // One element per genotype combination
+    using ExpandedGenotype = std::vector<ExpandedLikelihood>; // One element per read
+    using ExpandedHaplotypeLikelihoodVector = std::vector<ExpandedGenotype>; // One element per haplotype in genotype
+    using ExpandedGroupLikelihoodVector = std::vector<ExpandedHaplotypeLikelihoodVector>; // One element per group
+    using ExpandedHaplotypeLikelihoodMatrix = std::vector<ExpandedGroupLikelihoodVector>; // One element per sample
+
     Options options_;
 
     GroupOptionalLogPriorVector to_logs(const GroupOptionalPriorVector& prior) const;
     GroupOptionalLogPriorArray to_logs(const GroupOptionalPriorArray& priors) const;
+    ExpandedHaplotypeLikelihoodMatrix expand(const HaplotypeLikelihoodMatrix& likelihoods) const;
     PointInferences
     evaluate(const LogProbabilityVector& genotype_log_priors,
              const HaplotypeLikelihoodMatrix& log_likelihoods,
+             const ExpandedHaplotypeLikelihoodMatrix& expanded_log_likelihoods,
              const GroupOptionalLogPriorArray& group_log_priors,
              const GroupConcentrationVector& group_concentrations,
              const MixtureConcentrationArray& mixture_concentrations,
@@ -165,14 +173,14 @@ class VariationalBayesMixtureMixtureModel
                           const MixtureConcentrationArray& mixture_concentrations,
                           const ProbabilityVector& genotype_priors,
                           const GroupResponsibilityVector& group_responsibilities,
-                          const HaplotypeLikelihoodMatrix& log_likelihoods) const;
+                          const ExpandedHaplotypeLikelihoodMatrix& log_likelihoods) const;
     void
     update_responsibilities(ComponentResponsibilityMatrix& result,
                             const GroupConcentrationVector& group_concentrations,
                             const MixtureConcentrationArray& mixture_concentrations,
                             const ProbabilityVector& genotype_posteriors,
                             const GroupResponsibilityVector& group_responsibilities,
-                            const HaplotypeLikelihoodMatrix& log_likelihoods) const;
+                            const ExpandedHaplotypeLikelihoodMatrix& log_likelihoods) const;
     void
     update_genotype_log_posteriors(LogProbabilityVector& result,
                                    const LogProbabilityVector& genotype_log_priors,