o
    ॵi                     @   s`   d Z ddlZddlZddlmZmZ ddlmZ ddlm	Z	 ddl
mZmZ G dd	 d	eZdS )
zFA Python wrapper for hmmsearch - search profile against a sequence db.    N)OptionalSequence)logging)parsers   )hmmbuildutilsc                   @   s   e Zd ZdZdddedededeee  fdd	Zed
efddZ	ed
efddZ
ded
efddZded
efddZdeded
eej fddZdS )	Hmmsearchz'Python wrapper of the hmmsearch binary.N)flagsbinary_pathhmmbuild_binary_pathdatabase_pathr
   c                C   s\   || _ tj|d| _|| _|du rg d}|| _tj| js,t	
d| td| dS )a  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).
            flags: List of flags to be used by hmmsearch.

        Raises:
            RuntimeError: If hmmsearch binary not found within the path.
        )r   N)z--F10.1z--F2r   z--F3r   z--incE100z-Er   z--domEr   z	--incdomEr   z$Could not find hmmsearch database %sz"Could not find hmmsearch database )r   r   Hmmbuildhmmbuild_runnerr   r
   ospathexistsr   error
ValueError)selfr   r   r   r
    r   i/home/ubuntu/.local/lib/python3.10/site-packages/modelscope/models/science/unifold/msa/tools/hmmsearch.py__init__   s    zHmmsearch.__init__returnc                 C      dS Nstor   r   r   r   r   output_formatO      zHmmsearch.output_formatc                 C   r   r   r   r   r   r   r   input_formatS   r!   zHmmsearch.input_formatmsa_stoc                 C   s   | j j|dd}| |S )zAQueries the database using hmmsearch using a given stockholm msa.hand)model_construction)r   build_profile_from_stoquery_with_hmm)r   r#   hmmr   r   r   queryW   s   
zHmmsearch.queryr(   c              	   C   s  t  }tj|d}tj|d}t|d}|| W d   n1 s(w   Y  | jdddg}| jr=|	| j |	d||| j
g td	| tj|tjtjd
}t dtj| j
 d | \}}	| }
W d   n1 szw   Y  |
rtd|d|	df t|}| }W d   n1 sw   Y  W d   |S W d   |S 1 sw   Y  |S )z7Queries the database using hmmsearch using a given hmm.z	query.hmmz
output.stowNz--noaliz--cpu8z-AzLaunching sub-process %s)stdoutstderrzhmmsearch (z) queryz)hmmsearch failed:
stdout:
%s

stderr:
%s
zutf-8)r   tmpdir_managerr   r   joinopenwriter   r
   extendr   r   info
subprocessPopenPIPEtimingbasenamecommunicatewaitRuntimeErrordecoderead)r   r(   query_tmp_dirhmm_input_pathout_pathfcmdprocessr,   r-   retcodeout_msar   r   r   r'   ]   sZ   





'
''zHmmsearch.query_with_hmmoutput_stringinput_sequencec                 C   s"   t j|dd}t j||dd}|S )zAGets parsed template hits from the raw string output by the tool.F)remove_first_row_gaps)query_sequence
a3m_string
skip_first)r   convert_stockholm_to_a3mparse_hmmsearch_a3m)r   rF   rG   rJ   template_hitsr   r   r   get_template_hits   s   zHmmsearch.get_template_hits)__name__
__module____qualname____doc__strr   r   r   propertyr    r"   r)   r'   r   TemplateHitrO   r   r   r   r   r	      s2    

2+r	   )rS   r   r4   typingr   r   abslr   %modelscope.models.science.unifold.msar    r   r   objectr	   r   r   r   r   <module>   s   