File size: 4,703 Bytes
84bcf84
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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


"""A Python wrapper for hmmsearch - search profile against a sequence db."""

import os
import tempfile

from absl import logging
from flax_model.alphafold3.data import parsers
from flax_model.alphafold3.data.tools import hmmbuild
from flax_model.alphafold3.data.tools import subprocess_utils


class Hmmsearch(object):
  """Python wrapper of the hmmsearch binary."""

  def __init__(
      self,
      *,
      binary_path: str,
      hmmbuild_binary_path: str,
      database_path: str,
      alphabet: str = 'amino',
      filter_f1: float | None = None,
      filter_f2: float | None = None,
      filter_f3: float | None = None,
      e_value: float | None = None,
      inc_e: float | None = None,
      dom_e: float | None = None,
      incdom_e: float | None = None,
      filter_max: bool = False,
  ):
    """Initializes the Python hmmsearch wrapper.

    Args:
      binary_path: The path to the hmmsearch executable.
      hmmbuild_binary_path: The path to the hmmbuild executable. Used to build
        an hmm from an input a3m.
      database_path: The path to the hmmsearch database (FASTA format).
      alphabet: Chain type e.g. amino, rna, dna.
      filter_f1: MSV and biased composition pre-filter, set to >1.0 to turn off.
      filter_f2: Viterbi pre-filter, set to >1.0 to turn off.
      filter_f3: Forward pre-filter, set to >1.0 to turn off.
      e_value: E-value criteria for inclusion in tblout.
      inc_e: E-value criteria for inclusion in MSA/next round.
      dom_e: Domain e-value criteria for inclusion in tblout.
      incdom_e: Domain e-value criteria for inclusion of domains in MSA/next
        round.
      filter_max: Remove all filters, will ignore all filter_f* settings.

    Raises:
      RuntimeError: If hmmsearch binary not found within the path.
    """
    self._binary_path = binary_path
    self._hmmbuild_runner = hmmbuild.Hmmbuild(
        alphabet=alphabet, binary_path=hmmbuild_binary_path
    )
    self._database_path = database_path
    flags = []
    if filter_max:
      flags.append('--max')
    else:
      if filter_f1 is not None:
        flags.extend(('--F1', filter_f1))
      if filter_f2 is not None:
        flags.extend(('--F2', filter_f2))
      if filter_f3 is not None:
        flags.extend(('--F3', filter_f3))

    if e_value is not None:
      flags.extend(('-E', e_value))
    if inc_e is not None:
      flags.extend(('--incE', inc_e))
    if dom_e is not None:
      flags.extend(('--domE', dom_e))
    if incdom_e is not None:
      flags.extend(('--incdomE', incdom_e))

    self._flags = tuple(map(str, flags))

    subprocess_utils.check_binary_exists(
        path=self._binary_path, name='hmmsearch'
    )

    if not os.path.exists(self._database_path):
      logging.error('Could not find hmmsearch database %s', database_path)
      raise ValueError(f'Could not find hmmsearch database {database_path}')

  def query_with_hmm(self, hmm: str) -> str:
    """Queries the database using hmmsearch using a given hmm."""
    with tempfile.TemporaryDirectory() as query_tmp_dir:
      hmm_input_path = os.path.join(query_tmp_dir, 'query.hmm')
      sto_out_path = os.path.join(query_tmp_dir, 'output.sto')
      with open(hmm_input_path, 'w') as f:
        f.write(hmm)

      cmd = [
          self._binary_path,
          '--noali',  # Don't include the alignment in stdout.
          *('--cpu', '8'),
      ]
      # If adding flags, we have to do so before the output and input:
      if self._flags:
        cmd.extend(self._flags)
      cmd.extend([
          *('-A', sto_out_path),
          hmm_input_path,
          self._database_path,
      ])

      subprocess_utils.run(
          cmd=cmd,
          cmd_name=f'Hmmsearch ({os.path.basename(self._database_path)})',
          log_stdout=False,
          log_stderr=True,
          log_on_process_error=True,
      )

      with open(sto_out_path) as f:
        a3m_out = parsers.convert_stockholm_to_a3m(
            f, remove_first_row_gaps=False, linewidth=60
        )

    return a3m_out

  def query_with_a3m(self, a3m_in: str) -> str:
    """Query the database using hmmsearch using a given a3m."""

    # Only the "fast" model construction makes sense with A3M, as it doesn't
    # have any way to annotate reference columns.
    hmm = self._hmmbuild_runner.build_profile_from_a3m(a3m_in)
    return self.query_with_hmm(hmm)

  def query_with_sto(
      self, msa_sto: str, model_construction: str = 'fast'
  ) -> str:
    """Queries the database using hmmsearch using a given stockholm msa."""
    hmm = self._hmmbuild_runner.build_profile_from_sto(
        msa_sto, model_construction=model_construction
    )
    return self.query_with_hmm(hmm)