File size: 5,452 Bytes
35cdf53
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
// Copyright 2024 DeepMind Technologies Limited
//
// AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
// this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
//
// To request access to the AlphaFold 3 model parameters, follow the process set
// out at https://github.com/google-deepmind/alphafold3. You may only use these
// if received directly from Google. Use is subject to terms of use available at
// https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md

#include <algorithm>
#include <cstddef>
#include <stdexcept>
#include <string>
#include <vector>

#include "absl/strings/ascii.h"
#include "absl/strings/str_format.h"
#include "absl/strings/string_view.h"
#include "pybind11/pybind11.h"
#include "pybind11/stl.h"

namespace {

namespace py = pybind11;

std::vector<std::string> ConvertA3MToStockholm(
    std::vector<absl::string_view> a3m_sequences) {
  std::vector<std::string> stockholm_sequences(a3m_sequences.size());
  auto max_length_element =
      std::max_element(a3m_sequences.begin(), a3m_sequences.end(),
                       [](absl::string_view lhs, absl::string_view rhs) {
                         return lhs.size() < rhs.size();
                       });

  for (auto& out : stockholm_sequences) {
    out.reserve(max_length_element->size());
  }

  // While any sequence has remaining columns.
  while (std::any_of(a3m_sequences.begin(), a3m_sequences.end(),
                     [](absl::string_view in) { return !in.empty(); })) {
    if (std::any_of(a3m_sequences.begin(), a3m_sequences.end(),
                    [](absl::string_view in) {
                      return !in.empty() && absl::ascii_islower(in.front());
                    })) {
      // Insertion(s) found at column.
      for (std::size_t i = 0; i < a3m_sequences.size(); ++i) {
        absl::string_view& in = a3m_sequences[i];
        std::string& out = stockholm_sequences[i];
        if (!in.empty() && absl::ascii_islower(in.front())) {
          // Consume insertion.
          out.push_back(absl::ascii_toupper(in.front()));
          in.remove_prefix(1);
        } else {
          // Row requires padding.
          out.push_back('-');
        }
      }
    } else {
      // No insertions found.
      for (std::size_t i = 0; i < a3m_sequences.size(); ++i) {
        absl::string_view& in = a3m_sequences[i];
        std::string& out = stockholm_sequences[i];
        if (!in.empty()) {
          // Consume entire column.
          out.push_back(in.front());
          in.remove_prefix(1);
        } else {
          // One alignment is shorter than the others. Should not happen with
          // valid A3M input.
          throw std::invalid_argument(absl::StrFormat(
              "a3m rows have inconsistent lengths; row %d has no columns left "
              "but not all rows are exhausted",
              i));
        }
      }
    }
  }
  return stockholm_sequences;
}

std::string AlignSequenceToGaplessQuery(absl::string_view sequence,
                                        absl::string_view query_sequence) {
  if (sequence.size() != query_sequence.size()) {
    throw py::value_error(
        absl::StrFormat("The sequence (%d) and the query sequence (%d) don't "
                        "have the same length.",
                        sequence.size(), query_sequence.size()));
  }
  std::string output;
  for (std::size_t residue_index = 0, sequence_length = sequence.size();
       residue_index < sequence_length; ++residue_index) {
    const char query_residue = query_sequence[residue_index];
    const char residue = sequence[residue_index];
    if (query_residue != '-') {
      // No gap in the query, so the residue is aligned.
      output += residue;
    } else if (residue == '-') {
      // Gap in both sequence and query, simply skip.
      continue;
    } else {
      // Gap only in the query, so this must be an inserted residue.
      output += absl::ascii_tolower(residue);
    }
  }
  return output;
}

constexpr char kConvertA3mToStockholm[] = R"(
Converts a list of sequences in a3m format to stockholm format sequences.

As an example if the input is:
abCD
CgD
fCDa

Then the output will be:
ABC-D-
--CGD-
F-C-DA

Args:
  a3m_sequences: A list of strings in a3m format.

Returns
  A list of strings converted to stockholm format.
)";

constexpr char kAlignSequenceToGaplessQuery[] = R"(
Aligns a sequence to a gapless query sequence.

This is useful when converting Stockholm MSA to A3M MSA. Example:
Seq  : AB--E
Query: A--DE
Output: Ab-E.

Args:
  sequence: A string containing to be aligned.
  query_sequence: A string containing the reference sequence to align to.

Returns
  The input sequence with gaps dropped where both the `sequence` and
  `query_sequence` have gaps, and sequence elements non-capitalized where the
  `query_sequence` has a gap, but the `sequence` does not.
)";

}  // namespace

namespace alphafold3 {

void RegisterModuleMsaConversion(pybind11::module m) {
  m.def("convert_a3m_to_stockholm", &ConvertA3MToStockholm,
        py::arg("a3m_sequences"), py::call_guard<py::gil_scoped_release>(),
        py::doc(kConvertA3mToStockholm + 1));
  m.def("align_sequence_to_gapless_query", &AlignSequenceToGaplessQuery,
        py::arg("sequence"), py::arg("query_sequence"),
        py::call_guard<py::gil_scoped_release>(),
        py::doc(kAlignSequenceToGaplessQuery + 1));
}

}  // namespace alphafold3