Get table of promoter region

src/genomic_features/ensembl/ensembldb.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import ibis
+import numpy as np
 import requests
 from ibis import _
 from pandas import DataFrame, Timestamp
@@ -127,3 +128,53 @@ def genes(
     def chromosomes(self):
         """Get chromosome information."""
         return self.db.table("chromosome").execute()
+
+    def promoters(
+        self,
+        filter: _filters.AbstractFilterExpr = filters.EmptyFilter(),
+        upstream: int = 2000,
+        downstream: int = 200,
+        canonical_transcripts: bool = False,
+    ) -> DataFrame:
+        """Get promoter annotations.
+
+        Parameters
+        ----------
+        filter
+            Filter expression to apply to the genes table.
+        upstream
+            Number of base pairs upstream of the TSS (default: 2000).
+        downstream
+            Number of base pairs downstream of the TSS (default: 200).
+        canonical_transcripts
+            If True, return only canonical transcript for each gene (default: False).
+
+        Returns
+        -------
+        DataFrame
+            A table of promoter annotations.
+        """
+        # TODO: change to get transcript table with gene level columns
+        # something like:
+        # tx_table = self.transcripts(cols = set(cols + ['seq_strand', 'seq_name', 'tx_is_canonical']), filter = filter)
+        tx_table = self.genes(filter)
+
+        # Get promoter region based on strand
+        # strand = 1         |>>>>>>>>>>>>>>|
+        # strand = -1                         |<<<<<<<<<<<<<<|
+        # Tx SS:             *                               *
+        # Promoter       ------                             ------
+        tx_ss = np.where(
+            tx_table["seq_strand"] == 1,
+            tx_table["gene_seq_start"],
+            tx_table["gene_seq_end"],
+        )
+        tx_table["promoter_seq_start"] = np.where(
+            tx_table["seq_strand"] == 1, tx_ss - upstream, tx_ss - downstream
+        )
+        tx_table["promoter_seq_end"] = np.where(
+            tx_table["seq_strand"] == 1, tx_ss + downstream, tx_ss + upstream
+        )
+        # if canonical_transcripts:
+        #     tx_table = tx_table[tx_table["tx_is_canonical"] == 1]
+        return tx_table

tests/test_basic.py

@@ -16,3 +16,28 @@ def test_genes():
 def test_missing_version():
     with pytest.raises(ValueError):
         gf.ensembl.annotation("Hsapiens", 86)
+
+
+def test_promoters():
+    promoters = gf.ensembl.annotation("Hsapiens", 108).promoters()
+    assert isinstance(promoters, pd.DataFrame)
+    promoters = gf.ensembl.annotation("Hsapiens", 108).promoters(
+        upstream=100, downstream=100
+    )
+    assert ((promoters.promoter_seq_end - promoters.promoter_seq_start) == 200).all()
+    promoters = gf.ensembl.annotation("Hsapiens", 108).promoters(
+        upstream=1000, downstream=100
+    )
+    assert ((promoters.promoter_seq_end - promoters.promoter_seq_start) == 1100).all()
+    # Test strandedness
+    promoters = gf.ensembl.annotation("Hsapiens", 108).promoters(
+        upstream=1000, downstream=100
+    )
+    assert (
+        promoters[promoters.seq_strand == -1].promoter_seq_start
+        == promoters[promoters.seq_strand == -1].gene_seq_end - 100
+    ).all()
+    assert (
+        promoters[promoters.seq_strand == 1].promoter_seq_start
+        == promoters[promoters.seq_strand == 1].gene_seq_start - 1000
+    ).all()