Merge pull request #143 from MuhammedHasan/custom_fasta

custom fasta

Author: avantikalal

src/grelu/io/fasta.py

@@ -18,7 +18,7 @@ def check_fasta(fasta_file: str) -> bool:
     Returns:
         True if the file path has a valid FASTA extension and exists, False otherwise.
     """
-    fasta_extensions = (".fa", ".fasta", ".fa.gz", ".fasta.gz")
+    fasta_extensions = (".fa", ".fasta", ".fa.gz", ".fasta.gz", ".fa.bgz", ".fasta.bgz")
     return (
         isinstance(fasta_file, str)
         and fasta_file.endswith(fasta_extensions)
@@ -36,14 +36,18 @@ def read_fasta(fasta_file: str) -> List[str]:
     Returns:
         A list of DNA sequences as strings.
     """
-    from Bio import SeqIO
+    from Bio import SeqIO, bgzf
 
     assert check_fasta(fasta_file), "Input is not a valid FASTA file."
 
     if fasta_file.endswith(".gz"):
         # Read sequences from a gzipped FASTA file
         with gzip.open(fasta_file, "rt") as handle:
             return [str(record.seq) for record in SeqIO.parse(handle, "fasta")]
+    elif fasta_file.endswith(".bgz"):
+        # Read sequences from a bgzipped FASTA file
+        with bgzf.BgzfReader(fasta_file, "rt") as handle:
+            return [str(record.seq) for record in SeqIO.parse(handle, "fasta")]
     else:
         # Read sequences from a FASTA file
         with open(fasta_file, "rt") as handle:

src/grelu/io/genome.py

@@ -6,8 +6,45 @@
 import os
 from typing import List, Optional, Union
 
-import genomepy
 import pandas as pd
+import pyfaidx
+import genomepy
+
+
+class CustomGenome:
+    """
+    A custom genome object that can be used to load a genome from a file.
+
+    Args:
+        genome: Path to the genome file.
+    """
+    def __init__(self, genome: str):
+        self.genome = genome
+        self._genome = pyfaidx.Fasta(genome, rebuild=False)
+        fai_file = genome + ".fai"
+        if not os.path.isfile(fai_file):
+            raise FileNotFoundError(
+                f"Genome file {fai_file} not found. "
+                "Please provide a genome name or a path to a chromosome sizes file. "
+                f"Or generate one with: `samtools faidx {genome}`."
+            )
+        self._sizes_file = genome + ".sizes"
+
+    def get_seq(self, chrom: str, start: int, end: int, rc: bool = False) -> str:
+        """
+        Get the sequence for a given chromosome and interval.
+        """
+        return self._genome.get_seq(chrom, start, end, rc=rc)
+
+    @property
+    def sizes_file(self) -> str:
+        if not os.path.isfile(self._sizes_file):
+            raise FileNotFoundError(
+                f"Genome file {self._sizes_file} not found. "
+                "Please provide a genome name or a path to a chromosome sizes file. "
+                f"Or generate one with: `faidx -i chromsizes {self.genome} > {self._sizes_file}`."
+            )
+        return self._sizes_file
 
 
 def read_sizes(genome: str = "hg38") -> pd.DataFrame:
@@ -24,16 +61,13 @@ def read_sizes(genome: str = "hg38") -> pd.DataFrame:
         and "size" (chromosome size).
     """
     # Get file path
-    if not os.path.isfile(genome):
-        genome = get_genome(genome).sizes_file
-
-    # Read file
+    genome = get_genome(genome).sizes_file
     return pd.read_table(
         genome, header=None, names=["chrom", "size"], dtype={"chrom": str, "size": int}
     )
 
 
-def get_genome(genome: str, **kwargs) -> genomepy.Genome:
+def get_genome(genome: str, **kwargs) -> Union[CustomGenome, genomepy.Genome]:
     """
     Install a genome from genomepy and load it as a Genome object
 
@@ -44,11 +78,13 @@ def get_genome(genome: str, **kwargs) -> genomepy.Genome:
     Returns:
         Genome object
     """
-    if genome not in genomepy.list_installed_genomes():
-        return genomepy.install_genome(genome, annotation=False, **kwargs)
+    if os.path.isfile(genome):
+        return CustomGenome(genome, **kwargs)
     else:
-        return genomepy.Genome(genome)
-
+        if genome not in genomepy.list_installed_genomes():
+            return genomepy.install_genome(genome, annotation=False, **kwargs)
+        else:
+            return genomepy.Genome(genome, **kwargs)
 
 def read_gtf(
     genome: str, features: Optional[Union[str, List[str]]] = None

tests/files/test.fa

@@ -0,0 +1,5 @@
+>seq1
+AAC
+
+>seq2
+ATG

tests/files/test.fa.bgz

@@ -0,0 +1,2 @@
+seq1	3	6	3	4
+seq2	3	3145728	3	3

tests/files/test.fa.bgz.fai

@@ -0,0 +1,2 @@
+seq1	3	6	3	4
+seq2	3	17	3	3

tests/files/test.fa.fai

@@ -0,0 +1,2 @@
+seq1	3
+seq2	3

tests/files/test.fa.gz

@@ -1,14 +1,16 @@
 import os
 
+import pytest
 import numpy as np
 import pandas as pd
+import genomepy
 from pandas.testing import assert_frame_equal
 
 from grelu.io import read_tomtom
 from grelu.io.bed import read_bed
 from grelu.io.bigwig import read_bigwig
 from grelu.io.fasta import read_fasta
-from grelu.io.genome import read_sizes
+from grelu.io.genome import read_sizes, CustomGenome
 from grelu.io.motifs import read_meme_file, read_modisco_report
 from grelu.sequence.utils import resize
 
@@ -20,6 +22,16 @@ def test_read_sizes():
     assert df.shape == (194, 2)
     assert df.iloc[0].to_dict() == {"chrom": "chr1", "size": 248956422}
 
+    with pytest.raises(FileNotFoundError):
+        CustomGenome("tests/files/test.fa.bgz").sizes_file
+
+    chrom_sizes = CustomGenome("tests/files/test.fa").sizes_file
+    assert chrom_sizes == "tests/files/test.fa.sizes"
+    df = read_sizes("tests/files/test.fa")
+    assert df.shape == (2, 2)
+    assert df.iloc[0].to_dict() == {'chrom': 'seq1', 'size': 3}
+    assert df.iloc[1].to_dict() == {'chrom': 'seq2', 'size': 3}
+
 
 def test_read_tomtom():
     # No q-value threshold
@@ -34,9 +46,15 @@ def test_read_tomtom():
 
 
 def test_read_fasta():
+    fa_file = os.path.join(cwd, "files", "test.fa")
+    assert np.all(read_fasta(fa_file) == ["AAC", "ATG"])
+
     fa_file = os.path.join(cwd, "files", "test.fa.gz")
     assert np.all(read_fasta(fa_file) == ["AAC", "ATG"])
 
+    fa_file = os.path.join(cwd, "files", "test.fa.bgz")
+    assert np.all(read_fasta(fa_file) == ["AAC", "ATG"])
+
 
 expected_intervals = pd.DataFrame(
     {"chrom": ["chr1", "chr1", "chr2"], "start": [1, 3, 3], "end": [4, 6, 6]}

tests/files/test.fa.sizes

@@ -129,7 +129,20 @@ def test_seq_formatting():
         convert_input_type(indices[0], "one_hot", add_batch_axis=True), batch[[0]]
     )
 
+    # Test custom genome
+    intervals = pd.DataFrame(
+        {"chrom": ["seq1", "seq2"], "start": [1, 1], "end": [3, 3]}
+    )
+    assert convert_input_type(intervals, "strings", genome="tests/files/test.fa") == ["AC", "TG"]
 
+    intervals = pd.DataFrame(
+        {"chrom": ["seq1", "seq2"], "start": [1, 1], "end": [3, 3]}
+    )
+    assert convert_input_type(intervals, "strings", genome="tests/files/test.fa.bgz") == ["AC", "TG"]
+
+    intervals = pd.DataFrame(
+        {"chrom": ["seq1", "seq2"], "start": [1, 1], "end": [3, 3]}
+    )
 # Test Metrics functions