<commit_before>/*
 * Copyright 2021 Google LLC.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * 3. Neither the name of the copyright holder nor the names of its
 *    contributors may be used to endorse or promote products derived from this
 *    software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#include "deepvariant/direct_phasing.h"

#include <algorithm>
#include <string>

#include "deepvariant/protos/deepvariant.pb.h"
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/container/btree_set.h"
#include "absl/status/status.h"
#include "boost/graph/graphviz.hpp"
#include "third_party/nucleus/protos/variants.pb.h"
#include "third_party/nucleus/vendor/statusor.h"
#include "tensorflow/core/platform/logging.h"

namespace learning {
namespace genomics {
namespace deepvariant {

const int kMinRefAlleleDepth = 3;
const int kMinAllelesToPhase = 2;
const float kMinEdgeWeight = 2.0;
constexpr absl::string_view kRef = "REF";
const int kNumOfPhases = 2;

std::string ReadKey(const nucleus::genomics::v1::Read& read) {
  return absl::StrCat(read.fragment_name(), "/", read.read_number());
}

nucleus::StatusOr<std::vector<int>> DirectPhasing::PhaseReads(
    const std::vector<DeepVariantCall>& candidates,
    const std::vector<
        nucleus::ConstProtoPtr<const nucleus::genomics::v1::Read>>& reads) {
  // Build graph from candidates.
  Build(candidates, reads);
  // Iterate positions in order. Calculate the score for each combination of
  // allele pairs.
  for (int i = 0; i < positions_.size(); i++) {
    // TODO Call UpdateStartingScore if score cannot be improved from
    // position to the next position. This happens with bad data where reads
    // are mismapped. Good example is chr1:143175001-143200000.
    // In this case we can break the graph and treat each piece as separate
    // graphs. This TODO has a lower impact because it happens with "bad" data
    // and DeepVariant will reject these candidates in most of the cases.
    // The work is tracked in internal
    if (i == 0) {
      UpdateStartingScore(vertices_by_position_[positions_[i]]);
      continue;
    }

    // If any of the vertices have no incoming edges we create zero-weighted
    // edges connecting to all vertices in the previous position. This is
    // needed so that we can consider a "broken" path.
    // Example:
    // ... ------- A -------- C   ------ G ------ ...
    // ... --------C        [ T ] ------ A ------ ...
    // This is a simplified example showing how a broken path may still
    // need to be considered. In this case we will create extra edges
    // connecting T with A and T with C.
    absl::btree_set<Edge> incoming_edges;
    for (const auto& v : vertices_by_position_[positions_[i]]) {
      auto [start, end] = boost::in_edges(v, graph_);

      // If there are no incoming edges for the vertex create zero weight
      // edges to all previous vertices to connect the graph.
      if (start == end && i > 0) {
        for (const auto& prev_v : vertices_by_position_[positions_[i - 1]]) {
          incoming_edges.insert(AddEdge(prev_v, v, 0));
        }
      }
      incoming_edges.insert(start, end);
    }

    // Enumerate all edge pairs
    for (const auto& edge_1 : incoming_edges) {
      for (const auto& edge_2 : incoming_edges) {
        const Vertex& to_1 = edge_1.m_target;
        const Vertex& to_2 = edge_2.m_target;
        Score score = CalculateScore(edge_1, edge_2);
        // If the score for the given vertices already exists then we update
        // it if the new score is higher.
        auto stored = scores_[{to_1, to_2}];
        if (stored.score < score.score) {
          scores_[{to_1, to_2}] = score;
        }
      }  // for j
    }    // for i
  }
  // Backtrack from the last position. For each position where best partition is
  // not homozygous assign phases to vertices (alleles).
  AssignPhasesToVertices();

  // Phases are assigned to reads based on a set of alleles the read overlap.
  // If read overlaps more alleles of phase 1 then it is assigned a phase 1.
  // There 3 possible assignments: 0, 1, 2 where 0 is "phase unassigned".
  return AssignPhasesToReads(reads);
}

void DirectPhasing::AssignPhasesToVertices() {
  // Assigning a random valid score. The max_score should be at least no less
  // then this.
  if (scores_.empty()) {
    return;
  }
  auto max_score_it = scores_.begin();
  int max_score = 0;
  // Find the best score for the last position.
  for (const Vertex& v1 : vertices_by_position_[positions_.back()]) {
    for (const Vertex& v2 : vertices_by_position_[positions_.back()]) {
      auto scores_it = scores_.find({v1, v2});
      if (scores_it != scores_.end() &&
          scores_it->second.score > max_score) {
        max_score_it = scores_it;
        max_score = scores_it->second.score;
      }
    }
  }

  while (max_score_it != scores_.end()) {
    if (max_score_it->first.phase_1_vertex !=
        max_score_it->first.phase_2_vertex) {
      graph_[max_score_it->first.phase_1_vertex].allele_info.phase = 1;
      graph_[max_score_it->first.phase_2_vertex].allele_info.phase = 2;
    }
    // Go to the next score.
    max_score_it = scores_.find(
        {max_score_it->second.from[0], max_score_it->second.from[1]});
  }
}

std::vector<int> DirectPhasing::AssignPhasesToReads(
    const std::vector<
        nucleus::ConstProtoPtr<const nucleus::genomics::v1::Read>>& reads)
    const {
  // Assign phase reads
  // 1. For each read find all allele the read overlaps.
  // 2. Assign the phase to the read based on the majority phase of all
  // overlapped alleles.<commit_msg>Sort edges by allele bases to make the algorithm deterministic.<commit_after>/*
 * Copyright 2021 Google LLC.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * 3. Neither the name of the copyright holder nor the names of its
 *    contributors may be used to endorse or promote products derived from this
 *    software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

#include "deepvariant/direct_phasing.h"

#include <algorithm>
#include <string>

#include "deepvariant/protos/deepvariant.pb.h"
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/container/btree_map.h"
#include "absl/container/btree_set.h"
#include "absl/status/status.h"
#include "boost/graph/graphviz.hpp"
#include "third_party/nucleus/protos/variants.pb.h"
#include "third_party/nucleus/vendor/statusor.h"
#include "tensorflow/core/platform/logging.h"

namespace learning {
namespace genomics {
namespace deepvariant {

const int kMinRefAlleleDepth = 3;
const int kMinAllelesToPhase = 2;
const float kMinEdgeWeight = 2.0;
constexpr absl::string_view kRef = "REF";
const int kNumOfPhases = 2;

std::string ReadKey(const nucleus::genomics::v1::Read& read) {
  return absl::StrCat(read.fragment_name(), "/", read.read_number());
}

nucleus::StatusOr<std::vector<int>> DirectPhasing::PhaseReads(
    const std::vector<DeepVariantCall>& candidates,
    const std::vector<
        nucleus::ConstProtoPtr<const nucleus::genomics::v1::Read>>& reads) {
  // Build graph from candidates.
  Build(candidates, reads);
  // Iterate positions in order. Calculate the score for each combination of
  // allele pairs.
  for (int i = 0; i < positions_.size(); i++) {
    // TODO Call UpdateStartingScore if score cannot be improved from
    // position to the next position. This happens with bad data where reads
    // are mismapped. Good example is chr1:143175001-143200000.
    // In this case we can break the graph and treat each piece as separate
    // graphs. This TODO has a lower impact because it happens with "bad" data
    // and DeepVariant will reject these candidates in most of the cases.
    // The work is tracked in internal
    if (i == 0) {
      UpdateStartingScore(vertices_by_position_[positions_[i]]);
      continue;
    }

    // If any of the vertices have no incoming edges we create zero-weighted
    // edges connecting to all vertices in the previous position. This is
    // needed so that we can consider a "broken" path.
    // Example:
    // ... ------- A -------- C   ------ G ------ ...
    // ... --------C        [ T ] ------ A ------ ...
    // This is a simplified example showing how a broken path may still
    // need to be considered. In this case we will create extra edges
    // connecting T with A and T with C.
    absl::btree_set<Edge> incoming_edges;
    for (const auto& v : vertices_by_position_[positions_[i]]) {
      auto [start, end] = boost::in_edges(v, graph_);

      // If there are no incoming edges for the vertex create zero weight
      // edges to all previous vertices to connect the graph.
      if (start == end && i > 0) {
        for (const auto& prev_v : vertices_by_position_[positions_[i - 1]]) {
          incoming_edges.insert(AddEdge(prev_v, v, 0));
        }
      }
      incoming_edges.insert(start, end);
    }

    absl::btree_map<std::pair<std::string, std::string>, Edge> keyed_edges;
    for (const auto& edge : incoming_edges) {
      std::string edge_source = graph_[edge.m_source].allele_info.bases;
      std::string edge_target = graph_[edge.m_target].allele_info.bases;
      keyed_edges[{edge_source, edge_target}] = edge;
    }

    // Enumerate all edge pairs
    for (const auto& edge_1 : keyed_edges) {
      for (const auto& edge_2 : keyed_edges) {
        const Vertex& to_1 = edge_1.second.m_target;
        const Vertex& to_2 = edge_2.second.m_target;
        Score score = CalculateScore(edge_1.second, edge_2.second);
        // If the score for the given vertices already exists then we update
        // it if the new score is higher.
        auto stored = scores_[{to_1, to_2}];
        if (stored.score < score.score) {
          scores_[{to_1, to_2}] = score;
        }
      }  // for j
    }    // for i
  }
  // Backtrack from the last position. For each position where best partition is
  // not homozygous assign phases to vertices (alleles).
  AssignPhasesToVertices();

  // Phases are assigned to reads based on a set of alleles the read overlap.
  // If read overlaps more alleles of phase 1 then it is assigned a phase 1.
  // There 3 possible assignments: 0, 1, 2 where 0 is "phase unassigned".
  return AssignPhasesToReads(reads);
}

void DirectPhasing::AssignPhasesToVertices() {
  // Assigning a random valid score. The max_score should be at least no less
  // then this.
  if (scores_.empty()) {
    return;
  }
  auto max_score_it = scores_.begin();
  int max_score = 0;
  bool non_unique_max_score = true;
  int i = positions_.size()-1;
  while (non_unique_max_score && i >= 0) {
    max_score = 0;
    non_unique_max_score = false;
    // Iterate all scores at positions_[i] and the maximum.
    for (const Vertex& v1 : vertices_by_position_[positions_[i]]) {
      for (const Vertex& v2 : vertices_by_position_[positions_[i]]) {
        auto scores_it = scores_.find({v1, v2});
        if (scores_it != scores_.end() &&
            scores_it->second.score > max_score) {
          max_score_it = scores_it;
          max_score = scores_it->second.score;
        }
      }
    }

    // Check that max score is unique. If not, then move to previous position.
    for (const Vertex& v1 : vertices_by_position_[positions_[i]]) {
      for (const Vertex& v2 : vertices_by_position_[positions_[i]]) {
        auto scores_it = scores_.find({v1, v2});
        if (scores_it != scores_.end()
            && scores_it->second.score == max_score
            && (max_score_it->first.phase_1_vertex != v1
            || max_score_it->first.phase_2_vertex != v2)) {
          non_unique_max_score = true;
        }
      }
    }

    i--;
  }

  while (max_score_it != scores_.end()) {
    if (max_score_it->first.phase_1_vertex !=
        max_score_it->first.phase_2_vertex) {
      graph_[max_score_it->first.phase_1_vertex].allele_info.phase = 1;
      graph_[max_score_it->first.phase_2_vertex].allele_info.phase = 2;
    }
    // Go to the next score.
    max_score_it = scores_.find(
        {max_score_it->second.from[0], max_score_it->second.from[1]});
  }
}

std::vector<int> DirectPhasing::AssignPhasesToReads(
    const std::vector<
        nucleus::ConstProtoPtr<const nucleus::genomics::v1::Read>>& reads)
    const {
  // Assign phase reads
  // 1. For each read find all allele the read overlaps.
  // 2. Assign the phase to the read based on the majority phase of all
  // overlapped alleles.
{'fraction_non_alphanumeric': 0.10330358332223258, 'fraction_numerical': 0.010456906886905554, 'mean_word_length': 3.11144578313253, 'pattern_counts': {'":': 0, '<': 31, '<?xml version=': 0, '>': 51, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6619430', 'n_tokens_mistral': 4641, 'n_tokens_neox': 4363, 'n_words': 1573}