Repository: bioinf-mcb/Metagenomic-DeepFRI (language: python)

Issue:
Error with Foldcomp databases (AttributeError: 'pyopal.lib.Database' object has no attribute 'search')
[log.txt](https://github.com/bioinf-mcb/Metagenomic-DeepFRI/files/14037344/log.txt)



Discussion / hints:
I downloaded afdb_swissprot, afdb_swissprot_v4 and highquality_clust30, none of them worked due to the error as given in log
Hello,

Thank you for bringing this up! 
What version of `pyopal` do you use? The app was tested when `0.4.1` was the most recent release. 
I see that the new release introduced a lot of API changes, I will need to look for them in detail. 



Fix patch:
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 4f392cb..971961c 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -13,13 +13,7 @@ jobs:
     strategy:
       matrix:
         python-tag:
-        - cp38-manylinux_aarch64
-        - cp39-manylinux_aarch64
-        - cp310-manylinux_aarch64
         - cp311-manylinux_aarch64
-        - cp312-manylinux_aarch64
-        - pp38-manylinux_aarch64
-        - pp39-manylinux_aarch64
     steps:
     - uses: actions/checkout@v4
     - name: Set up QEMU
@@ -34,6 +28,7 @@ jobs:
         CIBW_BUILD: ${{ matrix.python-tag }}
         CIBW_BEFORE_BUILD: pip install cython
         CIBW_BUILD_VERBOSITY: 2
+        CIBW_TEST_COMMAND: python -m unittest mDeepFRI.tests -vv
         CIBW_TEST_REQUIRES: importlib-resources
       with:
         output-dir: dist
@@ -48,13 +43,7 @@ jobs:
     strategy:
       matrix:
         python-tag:
-        - cp38-manylinux_x86_64
-        - cp39-manylinux_x86_64
-        - cp310-manylinux_x86_64
         - cp311-manylinux_x86_64
-        - cp312-manylinux_x86_64
-        - pp38-manylinux_x86_64
-        - pp39-manylinux_x86_64
     steps:
     - uses: actions/checkout@v4
     - name: Build manylinux wheels
@@ -64,6 +53,7 @@ jobs:
         CIBW_BUILD: ${{ matrix.python-tag }}
         CIBW_BEFORE_BUILD: pip install cython
         CIBW_BUILD_VERBOSITY: 2
+        CIBW_TEST_COMMAND: python -m unittest mDeepFRI.tests -vv
         CIBW_TEST_REQUIRES: importlib-resources
       with:
         output-dir: dist
@@ -78,13 +68,7 @@ jobs:
     strategy:
       matrix:
         python-tag:
-        - cp38-macosx_x86_64
-        - cp39-macosx_x86_64
-        - cp310-macosx_x86_64
         - cp311-macosx_x86_64
-        - cp312-macosx_x86_64
-        - pp38-macosx_x86_64
-        - pp39-macosx_x86_64
     steps:
     - uses: actions/checkout@v4
     - name: Build manylinux wheels
@@ -94,6 +78,7 @@ jobs:
         CIBW_BUILD: ${{ matrix.python-tag }}
         CIBW_BEFORE_BUILD: pip install cython
         CIBW_BUILD_VERBOSITY: 2
+        CIBW_TEST_COMMAND: python -m unittest mDeepFRI.tests -vv
         CIBW_TEST_REQUIRES: importlib-resources
       with:
         output-dir: dist
@@ -122,6 +107,7 @@ jobs:
         CIBW_BUILD: ${{ matrix.python-tag }}
         CIBW_BEFORE_BUILD: pip install cython
         CIBW_BUILD_VERBOSITY: 2
+        CIBW_TEST_COMMAND: python -m unittest mDeepFRI.tests -vv
         CIBW_TEST_REQUIRES: importlib-resources
       with:
         output-dir: dist
@@ -168,6 +154,8 @@ jobs:
       run: python -m pip install -U pip setuptools wheel
     - name: Install built wheel
       run: python -m pip install --no-binary mDeepFRI --find-links=dist mDeepFRI
+    - name: Run tests without coverage
+      run: python -m unittest mDeepFRI.tests -vv
 
   upload:
     environment: PyPI
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f4f1af9..047acea 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,8 +5,23 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+[Unreleased]: https://github.com/bioinf-MCB/Metagenomic-DeepFRI/compare/v1.1.2...HEAD
+
+## [v1.1.2] - 2024-02-22
+[v1.1.2]: https://github.com/bioinf-MCB/Metagenomic-DeepFRI/compare/v1.1.1...v1.1.2
+
+### Added
+- chunked loading of MMSeqs2 preselected targets - improves runtime
+
+### Fixed
+- C++17 lowered to C++14 for compatibility
+- pyopal API v0.5 integrated
+- removed `scipy` distance calculation with function
+- replaced `biopython` with `biotite` for structure manipulation - better runtime and easier deploy
+- tests fixed & added to CI/CD
 
 ## [v1.1.1] - 2024-03-01
+[v1.1.1]: https://github.com/bioinf-MCB/Metagenomic-DeepFRI/compare/v1.0.0...v1.1.1
 
 ### Added
 
diff --git a/README.md b/README.md
index 422ba99..c558ca5 100644
--- a/README.md
+++ b/README.md
@@ -4,7 +4,7 @@
 ## 🔍 Overview
 Proteins perform most of the work of living cells. Amino acid sequence and structural features of proteins determine a wide range of functions: from binding specificity and conferring mechanical stability, to catalysis of biochemical reactions, transport, and signal transduction.
 DeepFRI is a neural network designed to predict protein function within the framework of the Gene Ontology (GO). The exponential growth in the number of available protein sequences, driven by advancements in low-cost sequencing technologies and computational methods (e.g., gene prediction), has resulted in a pressing need for efficient software to facilitate the annotation of protein databases.
-Metagenomic-DeepFRI addresses such need, building upon efficient libraries. It incorporates novel databases of predicted structures (AlphaFold, ESM-Fold, MIP, etc.) and improves runtimes of DeepFRI by [2-12 times](https://github.com/bioinf-mcb/Metagenomic-DeepFRI/blob/main/weight_convert/onnx_vs_tf2.png)!
+Metagenomic-DeepFRI addresses such need, building upon efficient libraries. It incorporates novel databases of predicted structures (AlphaFold, ESMFold, MIP, etc.) and improves runtimes of DeepFRI by [2-12 times](https://github.com/bioinf-mcb/Metagenomic-DeepFRI/blob/main/weight_convert/onnx_vs_tf2.png)!
 
 ### 📋 Pipeline stages
 
@@ -33,9 +33,9 @@ cd Metagenomic-DeepFRI
 conda env create --name deepfri --file environment.yml
 conda activate deepfri
 ```
-3. Install `mDeepFRI`
+3. Show help message
 ```{code-block} bash
-pip install .
+mDeepFRI --help
 ```
 
 ## 💡 Usage
diff --git a/docs/source/usage.md b/docs/source/usage.md
index 11c185a..0a14381 100644
--- a/docs/source/usage.md
+++ b/docs/source/usage.md
@@ -54,7 +54,7 @@ Other available parameters can be found upon command `deepfri --help`.
    * - `--generate_contacts`
      - Gap fill during contact map alignment.
      - 2
-   * - `--mmseqs_min_bit_score`
+   * - `--mmseqs_min_bitscore`
      - Minimum bit score for MMSeqs2 search.
      - None
    * - `--mmseqs_max_evalue`
@@ -63,7 +63,7 @@ Other available parameters can be found upon command `deepfri --help`.
    * - `--mmseqs_min_identity`
      - Minimum sequence identity for MMSeqs2 search.
      - 0.5
-   * - `--mmseqs_min_bit_score`
+   * - `--mmseqs_min_bitscore`
      - Minimum bit score for MMSeqs2 search.
      - None
    * - `--threads`
diff --git a/environment.yml b/environment.yml
index 50ed5d3..7d1d369 100644
--- a/environment.yml
+++ b/environment.yml
@@ -5,32 +5,16 @@ channels:
 dependencies:
   - click
   - numpy>=1.19
-  - scipy
-  - biopython
+  - biotite
   - mmseqs2=14.7e284
   - pip
-  - pytest
   - python==3.11
   - libdeflate
   - pip:
-    - pyopal
-    - pysam>=0.21.0
+    - pysam==0.21.0
+    - pyopal>=0.5
     - foldcomp
-    - cython
-    - requests
     # onnx stuff
-    - coloredlogs
-    - flatbuffers
-    - packaging
-    - protobuf
-    - sympy
-    - --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
-    - ort_nightly_gpu
-    - --extra-index-url=https://pypi.ngc.nvidia.com
-    - nvidia-cublas==11.5.1.101
-    - nvidia-cublas-cu117==11.10.1.25
-    - nvidia-cuda-runtime-cu114==11.4.148
-    - nvidia-cudnn==8.2.0.51
-    - nvidia-cudnn-cu11==8.5.0.96
-    - nvidia-cufft==10.4.2.58
-    - nvidia-curand==10.2.4.58
+    - onnxruntime-gpu
+    # mDeepFRI
+    - mdeepfri
diff --git a/mDeepFRI/alignment.py b/mDeepFRI/alignment.py
index 2104d86..af907ca 100644
--- a/mDeepFRI/alignment.py
+++ b/mDeepFRI/alignment.py
@@ -1,103 +1,125 @@
 import logging
+from collections import defaultdict
 from functools import partial
-from multiprocessing.pool import ThreadPool
+from multiprocessing.pool import Pool
 
 import numpy as np
 import pyopal
 
-from mDeepFRI.bio_utils import (AlignmentResult, load_query_sequences,
+from mDeepFRI.bio_utils import (AlignmentResult, load_fasta_as_dict,
                                 retrieve_fasta_entries_as_dict)
 from mDeepFRI.mmseqs import (filter_mmseqs_results, run_mmseqs_search,
                              validate_mmseqs_database)
 
 logger = logging.getLogger(__name__)
 
+##TODO: do not collect large database
+##TODO: align 1 query against top-k databases
 
-def align_best_score(query_item: str, database: pyopal.Database,
-                     target_names: list, gap_open: int, gap_extend: int,
-                     identity_threshold: float):
-
-    query_name, query_sequence = query_item
-
-    logger.debug("Aligning %s", query_name)
-    # current bug in pyOpal
-    # https://github.com/althonos/pyopal/issues/3
-    try:
-        results = database.search(query_sequence,
-                                  mode="score",
-                                  algorithm="nw",
-                                  gap_open=gap_open,
-                                  gap_extend=gap_extend)
-
-        best = max(results, key=lambda x: x.score)
-        best_sequence = database[best.target_index]
-        best_name = target_names[best.target_index]
-        single_seq = pyopal.Database([best_sequence])
-        alignment = single_seq.search(query_sequence,
-                                      mode="full",
-                                      algorithm="nw",
-                                      gap_open=gap_open,
-                                      gap_extend=gap_extend)[0].alignment
-
-        result = AlignmentResult(query_name, query_sequence, best_name,
-                                 best_sequence, alignment)
-        result.insert_gaps().calculate_identity()
-        logger.debug("Alignment of %s finished.", query_name)
-        if result.identity >= identity_threshold:
-            pass
-        else:
-            logger.debug("No alignment for  %s.", query_name)
-            result = None
-    except RuntimeError:
-        logger.debug("Alignment of %s failed.", query_name)
-        result = None
-
-    return result
-
-
-def align_query(query_seqs: dict,
-                target_seqs: dict,
-                alignment_gap_open: float,
-                alignment_gap_extend: float,
-                identity_threshold: str = 0.3,
-                threads: int = 1):
-
-    logger.info("Pairwise alignment started.")
-    logger.info("Aligning %i queries against %i database sequences.",
-                len(query_seqs), len(target_seqs))
-    # Default substitution matrix BLOSUM50
-    database = pyopal.Database(list(target_seqs.values()))
-    align_against_db = partial(align_best_score,
-                               database=database,
-                               target_names=list(target_seqs.keys()),
-                               gap_open=alignment_gap_open,
-                               gap_extend=alignment_gap_extend,
-                               identity_threshold=identity_threshold)
-
-    with ThreadPool(threads) as pool:
-        all_alignments = pool.map(align_against_db, query_seqs.items())
-
-    filtered_alignments = list(filter(None, all_alignments))
-    logger.info("Found %i alignments.", len(filtered_alignments))
-    logger.info("Pairwise alignment finished.")
-
-    return filtered_alignments
+
+def best_hit_database(query,
+                      target_sequences,
+                      gap_open: int = 10,
+                      gap_extend: int = 1):
+    """
+    Find the best hit in the database and return index.
+
+    Args:
+        query (str): The query sequence.
+        target_sequences (dict): The target sequences.
+
+    Returns:
+        int: The index of the best hit in the database.
+    """
+
+    aligner = pyopal.Aligner(gap_open=gap_open, gap_extend=gap_extend)
+    target_database = pyopal.Database(target_sequences.values())
+
+    # Retrieve the best hit
+    results = aligner.align(query,
+                            target_database,
+                            mode="score",
+                            overflow="buckets",
+                            algorithm="nw")
+    best_hit = max(results, key=lambda x: x.score)
+    best_index = best_hit.target_index
+    best_idx = list(target_sequences.keys())[best_index]
+    best_seq = target_sequences[best_idx]
+
+    return best_idx, best_seq
+
+
+def align_pairwise(query, target, gap_open: int = 10, gap_extend: int = 1):
+    """
+    Aligns the query against the target and returns the alignment.
+
+    Args:
+        query (str): The query sequence.
+        target (str): The target sequence.
+        aligner (pyopal.Aligner): The aligner object.
+
+    Returns:
+        pyopal.Alignment: The alignment of the query against the target.
+    """
+
+    aligner = pyopal.Aligner(gap_open=gap_open, gap_extend=gap_extend)
+    database = pyopal.Database([target])
+    # Align the sequences
+    alignment = aligner.align(query, database, algorithm="nw", mode="full")
+
+    return alignment[0].alignment
+
+
+def pairwise_against_database(query_id,
+                              query_sequence,
+                              target_sequences,
+                              gap_open: int = 10,
+                              gap_extend: int = 1):
+    """
+    Finds the best alignment of the query against the target.
+    """
+
+    # find best hits in the database
+    best_idx, best_target = best_hit_database(query_sequence, target_sequences,
+                                              gap_open, gap_extend)
+    # align the query against the best hit
+    alignment = align_pairwise(query_sequence, best_target, gap_open,
+                               gap_extend)
+    # create an alignment object
+    alignment_result = AlignmentResult(query_id, query_sequence, best_idx,
+                                       best_target, alignment)
+
+    return alignment_result
+
+
+def create_partial_database(top_matches, target_sequences):
+    """
+    Create a partial database from the target sequences.
+
+    Args:
+        query_id (str): The
+    """
+
+    partial_db = {k: target_sequences[k] for k in top_matches}
+
+    return partial_db
 
 
 def run_alignment(query_file: str,
                   mmseqs_db: str,
                   sequence_db: str,
                   output_path: str,
-                  mmseqs_min_bit_score: int = None,
+                  mmseqs_min_bitscore: int = None,
                   mmseqs_max_eval: float = 10e-5,
                   mmseqs_min_identity: float = 0.3,
-                  top_k: int = 30,
+                  top_k: int = 5,
                   alignment_gap_open: int = 10,
-                  alignment_gap_continuation: int = 1,
-                  identity_threshold: float = 0.5,
+                  alignment_gap_extend: int = 1,
                   threads: int = 1):
 
-    query_seqs = load_query_sequences(query_file, output_path)
+    align_partial = partial(pairwise_against_database,
+                            gap_open=alignment_gap_open,
+                            gap_extend=alignment_gap_extend)
     mmseqs_valid = validate_mmseqs_database(mmseqs_db)
     if not mmseqs_valid:
         raise ValueError(f"MMSeqs2 database not found in {mmseqs_db}.")
@@ -105,21 +127,35 @@ def run_alignment(query_file: str,
     mmseqs_results = run_mmseqs_search(query_file, mmseqs_db, output_path,
                                        threads)
     filtered_mmseqs_results = filter_mmseqs_results(mmseqs_results,
-                                                    mmseqs_min_bit_score,
+                                                    mmseqs_min_bitscore,
                                                     mmseqs_max_eval,
                                                     mmseqs_min_identity,
                                                     k_best_hits=top_k,
                                                     threads=threads)
 
+    query_seqs = load_fasta_as_dict(query_file)
+
+    unique_queries = defaultdict(list)
+    for result in filtered_mmseqs_results:
+        unique_queries[result[0]].append(result[1])
+
     # if MMSeqs2 alignment is empty
     if filtered_mmseqs_results is None:
         pass
     else:
         target_ids = np.unique(filtered_mmseqs_results["target"]).tolist()
         target_seqs = retrieve_fasta_entries_as_dict(sequence_db, target_ids)
-
-        alignments = align_query(query_seqs, target_seqs, alignment_gap_open,
-                                 alignment_gap_continuation,
-                                 identity_threshold, threads)
+        # create partial databases
+        partial_databases = {
+            k: create_partial_database(v, target_seqs)
+            for k, v in unique_queries.items()
+        }
+
+        # align the query against the best hit
+        with Pool(threads) as pool:
+            alignments = pool.starmap(
+                align_partial,
+                [(query_id, query_seqs[query_id], partial_databases[query_id])
+                 for query_id in unique_queries])
 
     return alignments
diff --git a/mDeepFRI/alignment_utils.pyx b/mDeepFRI/alignment_utils.pyx
index 92bf488..184dc61 100644
--- a/mDeepFRI/alignment_utils.pyx
+++ b/mDeepFRI/alignment_utils.pyx
@@ -1,4 +1,9 @@
 cimport cython
+
+import numpy as np
+
+cimport numpy as np
+from libc.math cimport sqrt
 from libc.stdlib cimport free, malloc
 from libc.string cimport strlen
 
@@ -35,3 +40,27 @@ def alignment_identity(str query, str target):
         float: Sequence identity between two sequences in an alignment.
     """
     return alignment_sequences_identity(query.encode(), target.encode())
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+@cython.cdivision(True)
+cpdef pairwise_sqeuclidean(float[:, ::1] X):
+
+    cdef int n = X.shape[0]
+    cdef int m = X.shape[1]
+    cdef int i, j, k
+    cdef float d, diff
+    cdef float[:, ::1] D = np.zeros((n, n), dtype=np.float32)
+
+    with nogil:
+        for i in range(n):
+            for j in range(i + 1, n):
+                d = 0
+                for k in range(m):
+                    diff = X[i, k] - X[j, k]
+                    d += diff * diff
+                D[i, j] = d
+                D[j, i] = d
+
+    return np.asarray(D)
diff --git a/mDeepFRI/bio_utils.py b/mDeepFRI/bio_utils.py
index a5da191..1a5cb22 100644
--- a/mDeepFRI/bio_utils.py
+++ b/mDeepFRI/bio_utils.py
@@ -1,117 +1,15 @@
-import json
-import logging
 from io import StringIO
-from pathlib import Path
 from typing import Dict, List, Tuple
 
 import foldcomp
 import numpy as np
 import requests
-from Bio.PDB import MMCIFParser, PDBParser
+from biotite.sequence import ProteinSequence
+from biotite.structure.io.pdb import PDBFile
+from biotite.structure.io.pdbx import PDBxFile, get_structure
 from pysam import FastaFile, FastxFile
-from scipy.spatial.distance import pdist, squareform
-
-from mDeepFRI.alignment_utils import alignment_identity
-
-# copied from Biopython to remove dependency
-protein_letters_1to3 = {
-    "A": "Ala",
-    "C": "Cys",
-    "D": "Asp",
-    "E": "Glu",
-    "F": "Phe",
-    "G": "Gly",
-    "H": "His",
-    "I": "Ile",
-    "K": "Lys",
-    "L": "Leu",
-    "M": "Met",
-    "N": "Asn",
-    "P": "Pro",
-    "Q": "Gln",
-    "R": "Arg",
-    "S": "Ser",
-    "T": "Thr",
-    "V": "Val",
-    "W": "Trp",
-    "Y": "Tyr",
-}
-protein_letters_1to3_extended = {
-    **protein_letters_1to3,
-    **{
-        "B": "Asx",
-        "X": "Xaa",
-        "Z": "Glx",
-        "J": "Xle",
-        "U": "Sec",
-        "O": "Pyl"
-    },
-}
-
-protein_letters_3to1_extended = {
-    value: key
-    for key, value in protein_letters_1to3_extended.items()
-}
-
-PROTEIN_LETTERS = dict()
-for k, v in protein_letters_3to1_extended.items():
-    PROTEIN_LETTERS[str.upper(k)] = v
-PROTEIN_LETTERS["UNK"] = "X"
 
-
-def seq2onehot(seq: str) -> np.ndarray:
-    """Create 26-dim one-hot encoding of a protein sequence.
-
-    Args:
-        seq (str): Protein sequence.
-
-    Returns:
-        np.ndarray: One-hot encoding of protein sequence.
-    """
-
-    chars = [
-        '-', 'D', 'G', 'U', 'L', 'N', 'T', 'K', 'H', 'Y', 'W', 'C', 'P', 'V',
-        'S', 'O', 'I', 'E', 'F', 'X', 'Q', 'A', 'B', 'Z', 'R', 'M'
-    ]
-    vocab_size = len(chars)
-    vocab_embed = dict(zip(chars, range(vocab_size)))
-
-    # Convert vocab to one-hot
-    vocab_one_hot = np.zeros((vocab_size, vocab_size), int)
-    for _, val in vocab_embed.items():
-        vocab_one_hot[val, val] = 1
-
-    embed_x = [vocab_embed[v] for v in seq]
-    seqs_x = np.array([vocab_one_hot[j, :] for j in embed_x])
-
-    return seqs_x
-
-
-def insert_gaps(sequence: str, reference: str,
-                alignment_string: str) -> Tuple[str, str]:
-    """
-    Inserts gaps into query and target sequences.
-
-    Args:
-        sequence (str): Query sequence.
-        reference (str): Target sequence.
-        alignment_string (str): Alignment string.
-
-    Returns:
-        gapped_sequence (str): Query sequence with gaps.
-        gapped_target (str): Target sequence with gaps.
-    """
-
-    sequence = list(sequence)
-    reference = list(reference)
-    alignment_string = list(alignment_string)
-
-    for i, a in enumerate(alignment_string):
-        if a == "I":
-            sequence.insert(i, "-")
-        elif a == "D":
-            reference.insert(i, "-")
-    return "".join(sequence), "".join(reference)
+from mDeepFRI.alignment_utils import alignment_identity, pairwise_sqeuclidean
 
 
 class AlignmentResult:
@@ -132,24 +30,16 @@ class AlignmentResult:
         insert_gaps: Inserts gaps into query and target sequences.
         calculate_identity: Calculates identity between query and target.
     """
-    def __init__(self,
-                 query_name,
-                 query_sequence,
-                 target_name,
-                 target_sequence,
-                 alignment,
-                 gapped_sequence="",
-                 gapped_target="",
-                 identity=None):
+    def __init__(self, query_name, query_sequence, target_name,
+                 target_sequence, alignment):
 
         self.query_name = query_name
         self.query_sequence = query_sequence
         self.target_name = target_name
         self.target_sequence = target_sequence
         self.alignment = alignment
-        self.gapped_sequence = gapped_sequence
-        self.gapped_target = gapped_target
-        self.identity = identity
+        self.insert_gaps()
+        self.calculate_identity()
         self.db_name = None
 
     def __str__(self):
@@ -163,23 +53,10 @@ def __repr__(self):
     def insert_gaps(self):
         """
         Inserts gaps into query and target sequences.
-
-        Args:
-            None
-        Returns:
-            self
         """
-        seq = list(self.query_sequence)
-        ref = list(self.target_sequence)
-        aln = list(self.alignment)
-
-        for i, a in enumerate(aln):
-            if a == "I":
-                seq.insert(i, "-")
-            elif a == "D":
-                ref.insert(i, "-")
-        self.gapped_sequence = "".join(seq)
-        self.gapped_target = "".join(ref)
+
+        self.gapped_sequence, self.gapped_target = insert_gaps(
+            self.query_sequence, self.target_sequence, self.alignment)
 
         return self
 
@@ -199,6 +76,61 @@ def calculate_identity(self):
         return self
 
 
+def seq2onehot(seq: str) -> np.ndarray:
+    """Create 26-dim one-hot encoding of a protein sequence.
+
+    Args:
+        seq (str): Protein sequence.
+
+    Returns:
+        np.ndarray: One-hot encoding of protein sequence.
+    """
+
+    chars = [
+        '-', 'D', 'G', 'U', 'L', 'N', 'T', 'K', 'H', 'Y', 'W', 'C', 'P', 'V',
+        'S', 'O', 'I', 'E', 'F', 'X', 'Q', 'A', 'B', 'Z', 'R', 'M'
+    ]
+    vocab_size = len(chars)
+    vocab_embed = dict(zip(chars, range(vocab_size)))
+
+    # Convert vocab to one-hot
+    vocab_one_hot = np.zeros((vocab_size, vocab_size), int)
+    for _, val in vocab_embed.items():
+        vocab_one_hot[val, val] = 1
+
+    embed_x = [vocab_embed[v] for v in seq]
+    seqs_x = np.array([vocab_one_hot[j, :] for j in embed_x])
+
+    return seqs_x
+
+
+def insert_gaps(sequence: str, reference: str,
+                alignment_string: str) -> Tuple[str, str]:
+    """
+    Inserts gaps into query and target sequences.
+
+    Args:
+        sequence (str): Query sequence.
+        reference (str): Target sequence.
+        alignment_string (str): Alignment string.
+
+    Returns:
+        gapped_sequence (str): Query sequence with gaps.
+        gapped_target (str): Target sequence with gaps.
+    """
+
+    sequence = list(sequence)
+    reference = list(reference)
+    alignment_string = list(alignment_string)
+
+    for i, a in enumerate(alignment_string):
+        if a == "I":
+            sequence.insert(i, "-")
+        elif a == "D":
+            reference.insert(i, "-")
+    return "".join(sequence), "".join(reference)
+
+
 def load_fasta_as_dict(fasta_file: str) -> Dict[str, str]:
     """
     Load FASTA file as dict of headers to sequences
@@ -207,11 +139,12 @@ def load_fasta_as_dict(fasta_file: str) -> Dict[str, str]:
         fasta_file (str): Path to FASTA file. Can be compressed.
 
     Returns:
-        Dict[str, str]: Dictionary of FASTA entries.
+        Dict[str, str]: Dictionary of FASTA entries sorted by length.
     """
 
     with FastxFile(fasta_file) as fasta:
         fasta_dict = {entry.name: entry.sequence for entry in fasta}
+
     return fasta_dict
 
 
@@ -235,65 +168,6 @@ def retrieve_fasta_entries_as_dict(fasta_file: str,
     return fasta_dict
 
 
-def load_query_sequences(query_file: str, output_path: str) -> Dict[str, str]:
-    """
-    Loads query protein sequences from FASTA file. Filters
-    out sequences that are too short or too long.
-
-    Args:
-        query_file (str): Path to FASTA file with query protein sequences.
-        output_path (str): Path to output folder.
-
-    Returns:
-        query_seqs (dict): Dictionary with query protein headers to sequences.
-    """
-
-    # By DeepFRI design (60, 1000)
-    MIN_PROTEIN_LENGTH = 60
-    MAX_PROTEIN_LENGTH = 1_000
-
-    query_file = Path(query_file)
-    output_path = Path(output_path)
-
-    query_seqs = load_fasta_as_dict(query_file)
-
-    if len(query_seqs) == 0:
-        raise ValueError(
-            f"{query_file} does not contain parsable protein sequences.")
-
-    logging.info("Found total of %i protein sequences in %s", len(query_seqs),
-                 query_file)
-
-    # filter out sequences that are too short or too long
-    prot_len_outliers = {}
-    for prot_id, sequence in query_seqs.items():
-        prot_len = len(sequence)
-        if prot_len > MAX_PROTEIN_LENGTH or prot_len < MIN_PROTEIN_LENGTH:
-            prot_len_outliers[prot_id] = prot_len
-
-    for outlier in prot_len_outliers.keys():
-        query_seqs.pop(outlier)
-
-    # write skipped protein ids to file
-    if len(prot_len_outliers) > 0:
-        logging.info(
-            "Skipping %i proteins due to sequence length outside range %i-%i aa.",
-            len(prot_len_outliers), MIN_PROTEIN_LENGTH, MAX_PROTEIN_LENGTH)
-        logging.info("Skipped protein ids will be saved in " \
-                     "metadata_skipped_ids_length.json.")
-        json.dump(prot_len_outliers,
-                  open(output_path / 'metadata_skipped_ids_due_to_length.json',
-                       "w",
-                       encoding="utf-8"),
-                  indent=4,
-                  sort_keys=True)
-
-    assert len(query_seqs
-               ) > 0, "All proteins were filtered out due to sequence length."
-
-    return query_seqs
-
-
 def calculate_contact_map(coordinates: np.ndarray,
                           threshold=6.0,
                           mode="matrix") -> np.ndarray:
@@ -309,8 +183,10 @@ def calculate_contact_map(coordinates: np.ndarray,
     Returns:
         np.ndarray: Contact map.
     """
+    # squared euclidean is used for efficiency
+    threshold = threshold**2
 
-    distances = squareform(pdist(coordinates))
+    distances = pairwise_sqeuclidean(coordinates)
     cmap = (distances < threshold).astype(np.int32)
 
     if mode == "sparse":
@@ -321,10 +197,8 @@ def calculate_contact_map(coordinates: np.ndarray,
     return cmap
 
 
-def extract_residues_coordinates(
-        structure_string: str,
-        chain: str = None,
-        max_seq_len: int = 1_000) -> Tuple[List, np.ndarray]:
+def extract_residues_coordinates(structure_string: str,
+                                 chain: str = None) -> Tuple[List, np.ndarray]:
     """
     Extracts residues and coordinates from structural file.
     Automatically processes PDB and mmCIF files.
@@ -336,44 +210,25 @@ def extract_residues_coordinates(
     Returns:
         Tuple[List, np.ndarray]: Tuple of residues and coordinates.
     """
-
     try:
-        parser = MMCIFParser()
-        structure = parser.get_structure("", StringIO(structure_string))
-        if chain:
-            structure = structure[0][chain]
-
-        residues = list(structure.get_residues())
-        # get CA atom coordinates for mmCIF
-        coords = np.array([
-            atom.get_coord() for atom in structure.get_atoms()
-            if atom.get_name() == "CA"
-        ])
-
-    except KeyError:
-        parser = PDBParser()
-        structure = parser.get_structure("", StringIO(structure_string))
-        if chain:
-            structure = structure[0][chain]
-
-        residues = [x for x in structure.get_residues()][:max_seq_len]
-        coords = np.array([residue["CA"].get_coord() for residue in residues])
-
-    #     parser = MMCIFParser()
-    #     structure = parser.get_structure("", StringIO(structure_string))
-    #     if chain:
-    #         structure = structure[0][chain]
-
-    #     residues = list(structure.get_residues())
-    #     # get CA atom coordinates for mmCIF
-    #     coords = np.array([atom.get_coord() for atom in structure.get_atoms() if atom.get_name() == "CA"])
+        mmcif = PDBxFile.read(StringIO(structure_string))
+        structure = get_structure(mmcif, model=1)
+    except UnboundLocalError:
+        pdb = PDBFile(StringIO(structure_string))
+        structure = pdb.get_structure()[0]
+
+    protein_chain = structure[structure.chain_id == chain]
+    # extract CA atoms coordinates
+    ca_atoms = protein_chain[protein_chain.atom_name == "CA"]
+    residues = str(ProteinSequence(ca_atoms.res_name))
+    coords = ca_atoms.coord
+
     return (residues, coords)
 
 
-def retrieve_structure_features(
-        idx: str,
-        database_path: str = "pdb100",
-        max_seq_len: int = 1_000) -> Tuple[List, np.ndarray]:
+def retrieve_structure_features(idx: str,
+                                database_path: str = "pdb100"
+                                ) -> Tuple[List, np.ndarray]:
     """
     Retrieves structure either from PDB or supplied FoldComp database.
     Extracts sequence and coordinate infromaton.
@@ -408,9 +263,7 @@ def retrieve_structure_features(
                 for _, pdb in db:
                     structure = pdb
 
-    residues, coords = extract_residues_coordinates(structure,
-                                                    chain=chain,
-                                                    max_seq_len=max_seq_len)
+    residues, coords = extract_residues_coordinates(structure, chain=chain)
 
     return (residues, coords)
 
@@ -549,9 +402,7 @@ def retrieve_align_contact_map(
         database = "pdb100"
 
     idx = alignment.target_name.rsplit(".", 1)[0]
-    coordinates = retrieve_structure_features(idx,
-                                              database_path=database,
-                                              max_seq_len=max_seq_len)[1]
+    coordinates = retrieve_structure_features(idx, database_path=database)[1]
 
     cmap = calculate_contact_map(coordinates,
                                  threshold=threshold,
diff --git a/mDeepFRI/cli.py b/mDeepFRI/cli.py
index db1602d..9e2d1a9 100644
--- a/mDeepFRI/cli.py
+++ b/mDeepFRI/cli.py
@@ -106,10 +106,10 @@ def get_models(ctx, output, version):
     help="Gap fill threshold during contact map alignment.",
 )
 @click.option(
-    "--mmseqs-min-bit-score",
+    "--mmseqs-min-bitscore",
     default=None,
     type=float,
-    help="Minimum bit score for MMseqs2 alignment.",
+    help="Minimum bitscore for MMseqs2 alignment.",
 )
 @click.option(
     "--mmseqs-max-evalue",
@@ -125,7 +125,7 @@ def get_models(ctx, output, version):
 )
 @click.option(
     "--top-k",
-    default=30,
+    default=5,
     type=int,
     help="Number of top MMSeqs2 alignment for"
     "precise pairwise alignment check. Default is 30.",
@@ -172,7 +172,7 @@ def get_models(ctx, output, version):
 @click.pass_context
 def predict_function(ctx, input, db_path, weights, output, processing_modes,
                      angstrom_contact_thresh, generate_contacts,
-                     mmseqs_min_bit_score, mmseqs_max_evalue,
+                     mmseqs_min_bitscore, mmseqs_max_evalue,
                      mmseqs_min_identity, top_k, alignment_gap_open,
                      alignment_gap_extend, alignment_min_identity,
                      remove_intermediate, overwrite, threads):
@@ -190,7 +190,7 @@ def predict_function(ctx, input, db_path, weights, output, processing_modes,
         deepfri_processing_modes=processing_modes,
         angstrom_contact_threshold=angstrom_contact_thresh,
         generate_contacts=generate_contacts,
-        mmseqs_min_bit_score=mmseqs_min_bit_score,
+        mmseqs_min_bitscore=mmseqs_min_bitscore,
         mmseqs_max_eval=mmseqs_max_evalue,
         mmseqs_min_identity=mmseqs_min_identity,
         top_k=top_k,
diff --git a/mDeepFRI/mmseqs.py b/mDeepFRI/mmseqs.py
index b4be161..7cca9bf 100644
--- a/mDeepFRI/mmseqs.py
+++ b/mDeepFRI/mmseqs.py
@@ -163,7 +163,7 @@ def filter_mmseqs_results(results_file: str,
                           min_bit_score: float = None,
                           max_evalue: float = None,
                           min_identity: float = None,
-                          k_best_hits: int = 30,
+                          k_best_hits: int = 5,
                           threads: int = 1) -> np.recarray:
     """
     Filters MMSeqs results retrieving only k best hits based on identity
@@ -182,33 +182,31 @@ def filter_mmseqs_results(results_file: str,
         output (numpy.recarray): Filtered results.
     """
     def select_top_k(query, db, k=30):
-        return db[db["query"] == query][-k:]
+        return db[db["query"] == query][:k]
 
     output = np.recfromcsv(results_file,
                            delimiter="\t",
                            encoding="utf-8",
                            names=MMSEQS_COLUMN_NAMES)
 
-    logger.info("%i MMSeqs2 hits in the database.", output.size)
-
     # check if output is not empty
     if output.size == 0:
-        logger.info("No hits found in MMSeqs2 database.")
         final_database = None
+
     else:
         # MMSeqs2 alginment filters
         if min_identity:
-            filtered = output[output['identity'] >= min_identity]
+            output = output[output['identity'] >= min_identity]
         if min_bit_score:
-            filtered = output[output['bit_score'] >= min_bit_score]
+            output = output[output['bit_score'] >= min_bit_score]
         if max_evalue:
-            filtered = output[output['e_value'] <= max_evalue]
+            output = output[output['e_value'] <= max_evalue]
 
         # Get k best hits
-        filtered.sort(order=["query", "identity", "e_value"], kind="quciksort")
-        top_k_db = partial(select_top_k, db=filtered, k=k_best_hits)
+        output.sort(order=["query", "identity", "e_value"], kind="quciksort")
+        top_k_db = partial(select_top_k, db=output[::-1], k=k_best_hits)
         with ThreadPool(threads) as pool:
-            top_k_chunks = pool.map(top_k_db, np.unique(filtered["query"]))
+            top_k_chunks = pool.map(top_k_db, np.unique(output["query"]))
 
         final_database = np.concatenate(top_k_chunks)
 
diff --git a/mDeepFRI/pipeline.py b/mDeepFRI/pipeline.py
index 36f2287..d4af94e 100644
--- a/mDeepFRI/pipeline.py
+++ b/mDeepFRI/pipeline.py
@@ -8,7 +8,7 @@
 import numpy as np
 
 from mDeepFRI.alignment import run_alignment
-from mDeepFRI.bio_utils import load_query_sequences, retrieve_align_contact_map
+from mDeepFRI.bio_utils import load_fasta_as_dict, retrieve_align_contact_map
 from mDeepFRI.database import build_database
 from mDeepFRI.pdb import create_pdb_mmseqs
 from mDeepFRI.predict import Predictor
@@ -31,10 +31,10 @@ def predict_protein_function(
         deepfri_processing_modes: List[str] = ["ec", "bp", "mf", "cc"],
         angstrom_contact_threshold: float = 6,
         generate_contacts: int = 2,
-        mmseqs_min_bit_score: float = None,
+        mmseqs_min_bitscore: float = None,
         mmseqs_max_eval: float = 10e-5,
-        mmseqs_min_identity: float = 0.3,
-        top_k: int = 30,
+        mmseqs_min_identity: float = 0.5,
+        top_k: int = 5,
         alignment_gap_open: float = 10,
         alignment_gap_continuation: float = 1,
         identity_threshold: float = 0.5,
@@ -42,7 +42,11 @@ def predict_protein_function(
         overwrite=False,
         threads: int = 1):
 
+    MIN_SEQ_LEN = 60
     MAX_SEQ_LEN = 1000
+    logger.info("DeepFRI protein sequence limit: %i-%i", MIN_SEQ_LEN,
+                MAX_SEQ_LEN)
+
     query_file = pathlib.Path(query_file)
     weights = pathlib.Path(weights)
     output_path = pathlib.Path(output_path)
@@ -65,16 +69,10 @@ def predict_protein_function(
     # PDB100 database
     logger.info(
         "Creating PDB100 database. This may take a bit during a first run.")
-    # try:
     pdb100 = create_pdb_mmseqs()
     deepfri_dbs.append(pdb100)
     logger.info("PDB100 database created.")
 
-    # except RuntimeError:
-    #     logger.warning(
-    #         "PDB100 server is not responding. Skipping PDB100 database creation."
-    #     )
-
     # design solution
     # database is built in the same directory
     # where the structure database is stored
@@ -90,18 +88,22 @@ def predict_protein_function(
 
     # SEQUENCE ALIGNMENT
     aligned_queries = {}
-    query_seqs = load_query_sequences(query_file, output_path)
+    query_seqs = load_fasta_as_dict(query_file)
 
     for db in deepfri_dbs:
         aligned = len(aligned_queries)
         logger.info("Aligning %s sequences against %s", len(query_seqs),
                     db.name)
         alignments = run_alignment(query_file, db.mmseqs_db, db.sequence_db,
-                                   output_path, mmseqs_min_bit_score,
+                                   output_path, mmseqs_min_bitscore,
                                    mmseqs_max_eval, mmseqs_min_identity, top_k,
                                    alignment_gap_open,
-                                   alignment_gap_continuation,
-                                   identity_threshold, threads)
+                                   alignment_gap_continuation, threads)
+        # filter alignments by identity
+        alignments = [
+            aln for aln in alignments if aln.identity > identity_threshold
+        ]
+
         # set a db name for alignments
         for aln in alignments:
             aln.db_name = db.name
@@ -142,7 +144,6 @@ def predict_protein_function(
                         len(db_alignments), db.name)
             partial_align = partial(retrieve_align_contact_map,
                                     database=db.foldcomp_db,
-                                    max_seq_len=MAX_SEQ_LEN,
                                     threshold=angstrom_contact_threshold,
                                     generated_contacts=generate_contacts)
 
@@ -151,6 +152,10 @@ def predict_protein_function(
 
             aligned_cmaps.extend(partial_cmaps)
 
+    # sort cmaps by length of query sequence
+    aligned_cmaps = sorted(aligned_cmaps,
+                           key=lambda x: len(x[0].query_sequence))
+
     output_file_name = output_path / "results.tsv"
     output_buffer = open(output_file_name, "w", encoding="utf-8")
     csv_writer = csv.writer(output_buffer, delimiter="\t")
@@ -180,13 +185,23 @@ def predict_protein_function(
             gcn = Predictor(gcn_path, threads=threads)
 
             for i, (aln, aligned_cmap) in enumerate(aligned_cmaps):
+                if len(aln.query_sequence) > MAX_SEQ_LEN:
+                    logger.info("Skipping %s; sequence too long %i",
+                                aln.query_name, len(aln.query_sequence))
+                    continue
+
+                elif len(aln.query_sequence) < MIN_SEQ_LEN:
+                    logger.info("Skipping %s; sequence too short %i",
+                                aln.query_name, len(aln.query_sequence))
+                    continue
+
                 logger.info("Predicting %s; %i/%i", aln.query_name, i + 1,
                             gcn_prots)
                 # running the actual prediction
                 prediction_rows = gcn.predict_function(
                     seqres=aln.query_sequence,
                     cmap=aligned_cmap,
-                    chain=aln.query_name)
+                    chain=str(aln.query_name))
                 # writing the results to the output file
                 for row in prediction_rows:
                     row.extend([net_type, mode])
@@ -202,14 +217,24 @@ def predict_protein_function(
         # CNN for queries without satisfying alignments
         if cnn_prots > 0:
             net_type = "cnn"
-            logging.info("Predicting with CNN: %i proteins", cnn_prots)
+            logger.info("Predicting with CNN: %i proteins", cnn_prots)
             cnn_path = deepfri_models_config[net_type][mode]
             cnn = Predictor(cnn_path, threads=threads)
             for i, query_id in enumerate(unaligned_queries):
-                logging.info("Predicting %s; %i/%i", query_id, i + 1,
-                             cnn_prots)
+                seq = query_seqs[query_id]
+                if len(seq) > MAX_SEQ_LEN:
+                    logger.info("Skipping %s; sequence too long %i", query_id,
+                                len(seq))
+                    continue
+
+                elif len(seq) < MIN_SEQ_LEN:
+                    logger.info("Skipping %s; sequence too short %i", query_id,
+                                len(seq))
+                    continue
+
+                logger.info("Predicting %s; %i/%i", query_id, i + 1, cnn_prots)
                 prediction_rows = cnn.predict_function(
-                    seqres=query_seqs[query_id], chain=query_id)
+                    seqres=query_seqs[query_id], chain=str(query_id))
                 for row in prediction_rows:
                     row.extend([net_type, mode])
                     row.extend([np.nan, np.nan, np.nan])
diff --git a/setup.py b/setup.py
index ea23084..a1689f3 100644
--- a/setup.py
+++ b/setup.py
@@ -148,7 +148,7 @@ def run(self):
 
 install_requires = ["cython", "numpy", "requests"]
 setup_requires = ["cython", "requests", "numpy"]
-extra_compile_args = ["-std=c++17", "-O3"]
+extra_compile_args = ["-std=c++14", "-O3"]
 
 EXTENSIONS = [
     Extension("mDeepFRI.predict",
@@ -160,7 +160,7 @@ def run(self):
               sources=[SRC_DIR + "/alignment_utils.pyx"],
               language="c++",
               libraries=["stdc++"],
-              extra_compile_args=extra_compile_args),
+              extra_compile_args=extra_compile_args.extend("-w")),
 ]
 
 extras = {}


Test patch:
diff --git a/tests/data/structures/corrupted.pdb b/mDeepFRI/tests/__init__.py
similarity index 100%
rename from tests/data/structures/corrupted.pdb
rename to mDeepFRI/tests/__init__.py
diff --git a/tests/data/GCA_000731455.1.proteins.fa.gz b/mDeepFRI/tests/data/GCA_000731455.1.proteins.fa.gz
similarity index 100%
rename from tests/data/GCA_000731455.1.proteins.fa.gz
rename to mDeepFRI/tests/data/GCA_000731455.1.proteins.fa.gz
diff --git a/tests/data/small_query.faa b/mDeepFRI/tests/data/small_query.faa
similarity index 100%
rename from tests/data/small_query.faa
rename to mDeepFRI/tests/data/small_query.faa
diff --git a/tests/data/structures/1S3P-A.pdb b/mDeepFRI/tests/data/structures/1S3P-A.pdb
similarity index 100%
rename from tests/data/structures/1S3P-A.pdb
rename to mDeepFRI/tests/data/structures/1S3P-A.pdb
diff --git a/tests/data/structures/1gqv-sf.cif.gz b/mDeepFRI/tests/data/structures/1gqv-sf.cif.gz
similarity index 100%
rename from tests/data/structures/1gqv-sf.cif.gz
rename to mDeepFRI/tests/data/structures/1gqv-sf.cif.gz
diff --git a/tests/data/structures/1gvh.pdb.gz b/mDeepFRI/tests/data/structures/1gvh.pdb.gz
similarity index 100%
rename from tests/data/structures/1gvh.pdb.gz
rename to mDeepFRI/tests/data/structures/1gvh.pdb.gz
diff --git a/tests/data/structures/1nzo.pdb.gz b/mDeepFRI/tests/data/structures/1nzo.pdb.gz
similarity index 100%
rename from tests/data/structures/1nzo.pdb.gz
rename to mDeepFRI/tests/data/structures/1nzo.pdb.gz
diff --git a/tests/data/structures/1yt3.cif.gz b/mDeepFRI/tests/data/structures/1yt3.cif.gz
similarity index 100%
rename from tests/data/structures/1yt3.cif.gz
rename to mDeepFRI/tests/data/structures/1yt3.cif.gz
diff --git a/tests/data/structures/6a0j.cif.gz b/mDeepFRI/tests/data/structures/6a0j.cif.gz
similarity index 100%
rename from tests/data/structures/6a0j.cif.gz
rename to mDeepFRI/tests/data/structures/6a0j.cif.gz
diff --git a/tests/data/structures/AF-A0A2Z5TJB0-F1-model_v4.pdb.gz b/mDeepFRI/tests/data/structures/AF-A0A2Z5TJB0-F1-model_v4.pdb.gz
similarity index 100%
rename from tests/data/structures/AF-A0A2Z5TJB0-F1-model_v4.pdb.gz
rename to mDeepFRI/tests/data/structures/AF-A0A2Z5TJB0-F1-model_v4.pdb.gz
diff --git a/tests/data/structures/AF-A7YWM6-F1-model_v4.pdb.gz b/mDeepFRI/tests/data/structures/AF-A7YWM6-F1-model_v4.pdb.gz
similarity index 100%
rename from tests/data/structures/AF-A7YWM6-F1-model_v4.pdb.gz
rename to mDeepFRI/tests/data/structures/AF-A7YWM6-F1-model_v4.pdb.gz
diff --git a/tests/test_build_database.py b/mDeepFRI/tests/data/structures/corrupted.pdb
similarity index 100%
rename from tests/test_build_database.py
rename to mDeepFRI/tests/data/structures/corrupted.pdb
diff --git a/mDeepFRI/tests/test_alignment.py b/mDeepFRI/tests/test_alignment.py
new file mode 100644
index 0000000..adeb407
--- /dev/null
+++ b/mDeepFRI/tests/test_alignment.py
@@ -0,0 +1,32 @@
+import unittest
+
+import pyopal
+
+from mDeepFRI.alignment import align_pairwise, best_hit_database
+
+
+class TestAlignment(unittest.TestCase):
+    def setUp(self) -> None:
+        self.queries = dict(
+            query_seq="MAGFLKVVQLLAKYGSKAVQWAWANKGKILDWLNAGQAIDWVVS")
+        self.targets = dict(
+            seq1="MESILDLQELETSEEESALMAASTVSNNC",
+            seq2="MKKAVIVENKGCATCSIGAACLVDGPIPDFEIAGATGLFGLWG",
+            seq3="MAGFLKVVQILAKYGSKAVQWAWANKGKILDWINAGQAIDWVVE",
+            seq4="MAGFLKVVQILAKYGSKAVQWAWANKGKILDWINAGQAIDWVVE",
+        )
+        self.aligner = pyopal.Aligner()
+
+        return super().setUp()
+
+    def test_best_hit_database(self):
+
+        best_hit, _ = best_hit_database(self.queries["query_seq"],
+                                        self.targets)
+        self.assertEqual(best_hit, "seq3")
+
+    def test_align_pairwise(self):
+        alignment = align_pairwise(self.queries["query_seq"],
+                                   self.targets["seq3"])
+        self.assertEqual(alignment,
+                         "MMMMMMMMMXMMMMMMMMMMMMMMMMMMMMMMXMMMMMMMMMMX")
diff --git a/mDeepFRI/tests/test_bio_utils.py b/mDeepFRI/tests/test_bio_utils.py
new file mode 100644
index 0000000..248bc29
--- /dev/null
+++ b/mDeepFRI/tests/test_bio_utils.py
@@ -0,0 +1,30 @@
+import unittest
+
+import numpy as np
+
+from mDeepFRI.alignment_utils import alignment_identity, pairwise_sqeuclidean
+from mDeepFRI.bio_utils import insert_gaps
+
+
+class TestBioUtils(unittest.TestCase):
+    def test_insert_gaps(self):
+        self.assertEqual(insert_gaps('AACT', 'AAT', 'MMDM'), ('AACT', 'AA-T'))
+        self.assertEqual(insert_gaps('AAT', 'AATC', 'MMMI'), ('AAT-', 'AATC'))
+        self.assertEqual(insert_gaps('AAT', 'FGTC', 'XXMI'), ('AAT-', 'FGTC'))
+
+    def test_identity(self):
+        self.assertEqual(round(alignment_identity('AASDS', 'AASDS'), 2), 1)
+        self.assertEqual(round(alignment_identity('AASDS', 'ADSDS'), 2), 0.8)
+        self.assertEqual(round(alignment_identity('AASSS', 'A-S-S'), 2), 0.6)
+        self.assertEqual(round(alignment_identity('A--SS', 'A-S--'), 2), 0.2)
+
+    def test_pairwise_sqeuclidean(self):
+        np.random.seed(42)
+        expected = np.array(
+            [[0, 1.01354558, 0.12442072], [1.01354558, 0, 0.99467713],
+             [0.12442072, 0.99467713, 0]],
+            dtype=np.float32)
+
+        matrix = np.random.rand(3, 3).astype(np.float32)
+        result = pairwise_sqeuclidean(matrix)
+        np.allclose(result, expected)
diff --git a/mDeepFRI/tests/test_build_database.py b/mDeepFRI/tests/test_build_database.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/test_install.py b/mDeepFRI/tests/test_install.py
similarity index 100%
rename from tests/test_install.py
rename to mDeepFRI/tests/test_install.py
diff --git a/tests/test_pdb.py b/mDeepFRI/tests/test_pdb.py
similarity index 100%
rename from tests/test_pdb.py
rename to mDeepFRI/tests/test_pdb.py
diff --git a/mDeepFRI/tests/test_utils.py b/mDeepFRI/tests/test_utils.py
new file mode 100644
index 0000000..c23217c
--- /dev/null
+++ b/mDeepFRI/tests/test_utils.py
@@ -0,0 +1,15 @@
+import unittest
+
+from mDeepFRI import utils
+
+
+class TestUtils(unittest.TestCase):
+    def test_run_command(self):
+        command = "echo 'Hello World!'"
+        result = utils.run_command(command).strip()
+        self.assertEqual(result, "Hello World!")
+
+    def test_mmseqs(self):
+        command = "mmseqs"
+        result = utils.run_command(command).strip()
+        self.assertEqual(result.split()[0], "MMseqs2")
diff --git a/tests/test_alignment.py b/tests/test_alignment.py
deleted file mode 100644
index ec7697e..0000000
--- a/tests/test_alignment.py
+++ /dev/null
@@ -1,26 +0,0 @@
-import pytest
-
-from mDeepFRI.alignment import insert_gaps
-from mDeepFRI.alignment_utils import alignment_identity
-
-
-@pytest.mark.parametrize('query,target,expected', [('AASDS', 'AASDS', 1),
-                                                   ('AASDS', 'ADSDS', 0.8),
-                                                   ('AASSS', 'A-S-S', 0.6),
-                                                   ('A--SS', 'A-S--', 0.2)])
-def test_alignment_identity(query, target, expected):
-    seq_id = alignment_identity(query, target)
-    assert round(seq_id, 6) == expected
-
-
-@pytest.mark.parametrize(
-    'query,reference,alignment,gapped_query,gapped_reference',
-    [('AACT', 'AAT', 'MMDM', 'AACT', 'AA-T'),
-     ('AAT', 'AATC', 'MMMI', 'AAT-', 'AATC'),
-     ('AAT', 'FGTC', 'XXMI', 'AAT-', 'FGTC')])
-def test_insert_gaps_full_length(query, reference, alignment, gapped_query,
-                                 gapped_reference):
-    gapped_seq, gapped_ref = insert_gaps(query, reference, alignment)
-    assert gapped_seq == gapped_query
-    assert gapped_ref == gapped_reference
-    assert len(gapped_seq) == len(gapped_ref)
diff --git a/tests/test_cli.py b/tests/test_cli.py
deleted file mode 100644
index 84d1e94..0000000
--- a/tests/test_cli.py
+++ /dev/null
@@ -1,42 +0,0 @@
-import pytest
-from click.testing import CliRunner
-
-from mDeepFRI import cli
-
-
-@pytest.fixture
-def runner():
-    return CliRunner()
-
-
-# test click cli
-def test_cli_help(runner):
-    result = runner.invoke(cli.main, ["--help"])
-    assert result.exit_code == 0
-    assert "--help" in result.output
-
-
-def test_cli_build_db(runner):
-    result = runner.invoke(cli.main, ["build-db", "--help"])
-    assert result.exit_code == 0
-    assert "--help" in result.output
-    assert "-t" in result.output
-
-
-def test_cli_get_models(runner):
-    result = runner.invoke(cli.main, ["get-models", "--help"])
-    assert result.exit_code == 0
-    assert "--help" in result.output
-    assert "-o" in result.output
-
-
-def test_cli_predict_function(runner):
-    result = runner.invoke(cli.main, ["predict-function", "--help"])
-    assert result.exit_code == 0
-    assert "--help" in result.output
-    assert "-i" in result.output
-    assert "-o" in result.output
-    assert "-m" in result.output
-    assert "-t" in result.output
-    assert "-d" in result.output
-    assert "-w" in result.output
diff --git a/tests/test_deepfri.py b/tests/test_deepfri.py
deleted file mode 100644
index 66b2780..0000000
--- a/tests/test_deepfri.py
+++ /dev/null
@@ -1,50 +0,0 @@
-import pathlib
-import subprocess
-
-import pytest
-
-INPUT_STRUCTURES = 'tests/data'
-OUTPUT_DATABASE = 'test_database'
-QUERY_FILE = 'tests/data/small_query.faa'
-RESULTS = 'test_results'
-
-
-def execute_command(command):
-    process = subprocess.run(command.split(),
-                             check=False,
-                             capture_output=True,
-                             text=True)
-    if process.returncode != 0:
-        raise RuntimeError(
-            f'Command {command} failed with exit code {process.returncode}\n{process.stderr}'
-        )
-    return process.stdout, process.stderr
-
-
-@pytest.fixture
-def deepfri_database():
-    stdout, stderr = execute_command(
-        f'mDeepFRI build-db -i {INPUT_STRUCTURES} -o {OUTPUT_DATABASE}')
-    print(stdout)
-
-
-def test_database(deepfri_database):
-    return
-
-
-@pytest.fixture
-def deepri_results(deepfri_database):
-    stdout, stderr = execute_command(
-        f"mDeepFRI -i {QUERY_FILE} -db {OUTPUT_DATABASE} -o {RESULTS}")
-    print(stdout)
-
-
-def test_results(deepri_results):
-    results_path = pathlib.Path(RESULTS)
-    assert results_path.exists()
-    assert results_path.is_dir()
-    assert len(list(results_path.glob('*.tsv'))) == 8
-
-
-def test_clean(deepri_results):
-    execute_command(f'rm -r {OUTPUT_DATABASE} {RESULTS}')
diff --git a/tests/test_main_pipeline.py b/tests/test_main_pipeline.py
deleted file mode 100644
index 03a2d60..0000000
--- a/tests/test_main_pipeline.py
+++ /dev/null
@@ -1,8 +0,0 @@
-import pytest
-
-from mDeepFRI.utils import utils
-
-
-def test_deepfri_errors():
-    with pytest.raises(RuntimeError):
-        utils.run_command("mDeepFri").strip()
diff --git a/tests/test_parsers.py b/tests/test_parsers.py
deleted file mode 100644
index a6d13f7..0000000
--- a/tests/test_parsers.py
+++ /dev/null
@@ -1,67 +0,0 @@
-import gzip
-import pathlib
-
-import numpy as np
-import pytest
-
-from mDeepFRI.parsers import parse_mmcif, parse_pdb
-
-
-@pytest.fixture
-def pdb_file():
-    test_file = pathlib.Path(
-        __file__
-    ).parent / "data" / "structures" / "AF-A0A2Z5TJB0-F1-model_v4.pdb.gz"
-    return gzip.open(test_file, "rb").read().split(b"\n")
-
-
-def test_parse_pdb_sequence(pdb_file):
-    sequence, _, _ = parse_pdb(pdb_file)
-
-    assert sequence[:4] == [b'MET', b'MET', b'MET', b'MET']
-
-
-def test_parse_pdb_positions(pdb_file):
-    _, positions, _ = parse_pdb(pdb_file)
-    expected = np.array([-47.025, 49.762, -23.074], dtype=np.float32)
-    assert np.array_equal(positions[:3], expected)
-
-
-def test_parse_pdb_groups(pdb_file):
-    _, _, groups = parse_pdb(pdb_file)
-    assert groups[:1] == [b'A   1']
-
-
-def test_parse_pdb_lengths(pdb_file):
-    sequence, positions, groups = parse_pdb(pdb_file)
-    assert len(sequence) > 0
-    assert len(sequence) == len(positions) / 3 == len(groups)
-
-
-@pytest.fixture
-def mmcif_file():
-    test_file = pathlib.Path(
-        __file__).parent / "data" / "structures" / "6a0j.cif.gz"
-    return gzip.open(test_file, "rb").read().split(b"\n")
-
-
-def test_parse_mmcif_sequence(mmcif_file):
-    sequence, _, _ = parse_mmcif(mmcif_file)
-    assert sequence[:4] == [b'VAL', b'VAL', b'VAL', b'VAL']
-
-
-def test_parse_mmcif_positions(mmcif_file):
-    _, positions, _ = parse_mmcif(mmcif_file)
-    expected = np.array([-6.123, 31.928, 92.243], dtype=np.float32)
-    assert np.array_equal(positions[:3], expected)
-
-
-def test_parse_mmcif_groups(mmcif_file):
-    _, _, groups = parse_mmcif(mmcif_file)
-    assert groups[:1] == [b'A22']
-
-
-def test_parse_mmcif_lengths(mmcif_file):
-    sequence, positions, groups = parse_mmcif(mmcif_file)
-    assert len(sequence) > 0
-    assert len(sequence) == len(positions) / 3 == len(groups)
diff --git a/tests/test_utils.py b/tests/test_utils.py
deleted file mode 100644
index b56f6e7..0000000
--- a/tests/test_utils.py
+++ /dev/null
@@ -1,40 +0,0 @@
-from mDeepFRI import bio_utils, utils
-
-
-def test_protein_letters():
-    expected = {
-        'ALA': 'A',
-        'CYS': 'C',
-        'ASP': 'D',
-        'GLU': 'E',
-        'PHE': 'F',
-        'GLY': 'G',
-        'HIS': 'H',
-        'ILE': 'I',
-        'LYS': 'K',
-        'LEU': 'L',
-        'MET': 'M',
-        'ASN': 'N',
-        'PRO': 'P',
-        'GLN': 'Q',
-        'ARG': 'R',
-        'SER': 'S',
-        'THR': 'T',
-        'VAL': 'V',
-        'TRP': 'W',
-        'TYR': 'Y',
-        'ASX': 'B',
-        'XAA': 'X',
-        'GLX': 'Z',
-        'XLE': 'J',
-        'SEC': 'U',
-        'PYL': 'O',
-        'UNK': 'X'
-    }
-    assert bio_utils.PROTEIN_LETTERS == expected
-
-
-def test_run_command():
-    command = "echo 'Hello World!'"
-    result = utils.run_command(command).strip()
-    assert result == "Hello World!"