from __future__ import annotations
from collections import defaultdict
from collections.abc import Generator
from threading import Lock
from typing import Any, cast
from intervaltree import (
Interval,
IntervalTree,
)
from gain import logging
from gain.genomic_resources.repository import (
GenomicResource,
)
from gain.genomic_resources.resource_implementation import (
ResourceConfigValidationMixin,
get_base_resource_schema,
get_required_filename_schema,
)
from gain.genomic_resources.utils import (
build_chrom_mapping,
read_resource_id_label,
)
from gain.utils.regions import (
Region,
collapse,
)
from .parsers import load_transcript_models
from .transcript_models import TranscriptModel
logger = logging.getLogger(__name__)
[docs]
class GeneModels(
ResourceConfigValidationMixin,
):
"""Manage and query gene model data from genomic resources.
This class provides access to gene models loaded from various file formats
(GTF, refFlat, refSeq, CCDS, etc.) and offers efficient querying by gene
name or genomic location.
The class maintains three internal data structures:
- transcript_models: Dict mapping transcript IDs to TranscriptModel objects
- gene_models: Dict mapping gene names to lists of TranscriptModel objects
- _tx_index: IntervalTree index for fast location-based queries
Attributes:
resource (GenomicResource): The genomic resource containing gene models.
config (dict): Validated configuration from the resource.
reference_genome_id (str | None): ID of the reference genome.
gene_models (dict[str, list[TranscriptModel]]): Gene name to
transcript models mapping.
transcript_models (dict[str, TranscriptModel]): Transcript ID to
transcript model mapping.
Example:
>>> from gain.genomic_resources.gene_models \
... .gene_models_factory import build_gene_models_from_file
>>> gene_models = build_gene_models_from_file("genes.gtf")
>>> gene_models.load()
>>> # Query by gene name
>>> tp53_transcripts = gene_models.gene_models_by_gene_name("TP53")
>>> # Query by location
>>> transcripts = gene_models.gene_models_by_location("chr17", 7676592)
Note:
The gene models must be loaded (see ``is_loaded()``) before
queries can be performed. The class is thread-safe for concurrent
access.
"""
def __init__(self, resource: GenomicResource):
self._is_loaded = False
if resource.get_type() != "gene_models":
raise ValueError(
f"wrong type of resource passed: {resource.get_type()}")
self.resource = resource
self.config = self.validate_and_normalize_schema(
resource.get_config(), resource,
)
self.reference_genome_id: str | None = \
read_resource_id_label(resource, "reference_genome")
self.gene_models: dict[str, list[TranscriptModel]] = defaultdict(list)
self._tx_index: dict[str, IntervalTree] = defaultdict(IntervalTree)
self.transcript_models: dict[str, Any] = {}
self._reset()
self.__lock = Lock()
@property
def resource_id(self) -> str:
"""The id of the gene models resource this object wraps."""
return self.resource.resource_id
[docs]
def close(self) -> None:
"""Do nothing: gene models hold no open file and stay loaded.
Here so a gene models resource closes like every other.
"""
def _reset(self) -> None:
"""Reset gene models."""
self._is_loaded = False
self.transcript_models = {}
self.gene_models = defaultdict(list)
self._tx_index = defaultdict(IntervalTree)
def _add_to_utr_index(self, tm: TranscriptModel) -> None:
self._tx_index[tm.chrom].add(Interval(tm.tx[0], tm.tx[1] + 1, tm))
def _chrom_genes(self) -> Generator[
tuple[tuple[str, str], list[TranscriptModel]], None, None]:
"""Generate chromosome and gene name keys with transcript models."""
for chrom, interval_tree in self._tx_index.items():
gene_models: dict[
tuple[str, str], list[TranscriptModel]] = defaultdict(list)
for interval in interval_tree:
tm = cast(TranscriptModel, interval.data)
assert chrom == tm.chrom
gene_models[tm.chrom, tm.gene].append(tm)
yield from gene_models.items()
def _update_indexes(self) -> None:
"""Update internal indexes."""
self.gene_models = defaultdict(list)
self._tx_index = defaultdict(IntervalTree)
for transcript in self.transcript_models.values():
self.gene_models[transcript.gene].append(transcript)
self._add_to_utr_index(transcript)
[docs]
def gene_names(self) -> list[str]:
"""Get list of all gene names in the loaded gene models.
Returns
-------
list[str]
The gene names (symbols), one for each gene. The list is
empty if the gene models hold no genes, for example when
they are not loaded yet.
Example:
>>> gene_models.load()
>>> genes = gene_models.gene_names()
>>> print(f"Loaded {len(genes)} genes")
"""
if self.gene_models is None:
logger.warning(
"gene models %s are empty", self.resource.resource_id)
return []
return list(self.gene_models.keys())
[docs]
def gene_models_by_gene_name(
self, name: str,
) -> list[TranscriptModel] | None:
"""Retrieve all transcript models for a specific gene.
Parameters
----------
name : str
The gene name (symbol) to search for.
Returns
-------
list[TranscriptModel] | None
The transcript models of the gene, or ``None`` if the gene
models hold no gene with this name. The result is never an
empty list.
Example:
>>> transcripts = gene_models.gene_models_by_gene_name("BRCA1")
>>> if transcripts:
... print(f"BRCA1 has {len(transcripts)} transcript variants")
"""
return self.gene_models.get(name, None)
[docs]
def has_chromosome(self, chrom: str) -> bool:
"""Check if a chromosome has any gene models.
Parameters
----------
chrom : str
The chromosome name to check.
Returns
-------
bool
``True`` if at least one transcript model is on the
chromosome, ``False`` otherwise.
Example:
>>> if gene_models.has_chromosome("chr1"):
... print("Chromosome 1 has gene annotations")
"""
return chrom in self._tx_index
[docs]
def gene_models_by_location(
self, chrom: str, pos_begin: int, pos_end: int | None = None,
) -> list[TranscriptModel]:
"""Retrieve transcripts overlapping a genomic position or region.
This method uses an interval tree index for efficient querying of
transcripts by genomic coordinates.
Parameters
----------
chrom : str
The chromosome name (for example "chr1" or "17").
pos_begin : int
The start position (1-based, inclusive).
pos_end : int | None, default None
The end position (1-based, inclusive). If ``None``, the
query is the single position ``pos_begin``.
Returns
-------
list[TranscriptModel]
The transcript models whose transcript region overlaps the
query position or region. The list is empty if no transcript
overlaps or if the chromosome has no gene models.
Example:
>>> # Query single position
>>> models = gene_models.gene_models_by_location("chr17", 7676592)
>>> # Query region
>>> models = gene_models.gene_models_by_location(
... "chr17", 7661779, 7687550
... )
>>> for tm in models:
... print(f"{tm.gene}: {tm.tr_id}")
Note:
Positions are swapped automatically if pos_end < pos_begin.
"""
with self.__lock:
if chrom not in self._tx_index:
return []
if pos_end is None:
pos_end = pos_begin
if pos_end < pos_begin:
pos_begin, pos_end = pos_end, pos_begin
tms_interval = self._tx_index[chrom]
result = tms_interval.overlap(pos_begin, pos_end + 1)
return [r.data for r in result]
[docs]
@staticmethod
def get_schema() -> dict[str, Any]:
"""The schema a ``gene_models`` resource's config is checked against.
The base resource schema plus ``filename``, ``format``,
``gene_mapping`` and ``chrom_mapping``. ``filename`` is required:
the loader reads it unconditionally.
Returns
-------
dict[str, Any]
The Cerberus schema of the resource config.
"""
return {
**get_base_resource_schema(),
**get_required_filename_schema(),
"format": {"type": "string"},
"gene_mapping": {"type": "string"},
"chrom_mapping": {"type": "dict", "schema": {
"filename": {
"type": "string",
"excludes": ["add_prefix", "del_prefix"],
},
"add_prefix": {"type": "string"},
"del_prefix": {"type": "string", "excludes": "add_prefix"},
}},
}
[docs]
def load(self) -> GeneModels:
"""Load gene models from the genomic resource.
This method parses the gene model file and builds internal indexes
for efficient querying. It is thread-safe and will only load once.
Returns
-------
GeneModels
This object, for method chaining.
Example:
>>> gene_models = build_gene_models_from_file("genes.gtf")
>>> gene_models.load()
>>> num_transcripts = len(gene_models.transcript_models)
>>> print(f"Loaded {num_transcripts} transcripts")
Note:
Calling load() multiple times is safe - subsequent calls return
immediately if already loaded.
"""
with self.__lock:
if self._is_loaded:
return self
self._reset()
self._install(
self._chrom_mapping(load_transcript_models(self.resource)))
return self
def _install(
self, transcript_models: dict[str, TranscriptModel],
) -> None:
"""Make these the model's transcripts, indexed and loaded.
The one transition to loaded, as ``_reset()`` is the one
transition back: transcripts, indexes and the loaded flag move
together or not at all.
"""
self.transcript_models = transcript_models
self._update_indexes()
self._is_loaded = True
def _chrom_mapping(
self, transcript_models: dict[str, TranscriptModel],
) -> dict[str, TranscriptModel]:
chrom_mapping = build_chrom_mapping(self.resource)
if chrom_mapping is None:
return transcript_models
result = {}
for tm in transcript_models.values():
chrom = chrom_mapping(tm.chrom)
if chrom is None:
logger.warning(
"transcript %s on chrom %s is removed by "
"chrom_mapping",
tm.tr_id, tm.chrom,
)
continue
tm.chrom = chrom
result[tm.tr_id] = tm
return result
[docs]
def is_loaded(self) -> bool:
"""Check whether this object holds usable gene models.
True once the model holds indexed transcripts, however it got
them: ``load()`` parsed its resource, or it was built loaded by
``join_gene_models`` or ``from_transcript_models``. ``load()``
on a loaded model is a no-op.
Returns
-------
bool
``True`` if the models are loaded and indexed, ``False`` if
nobody loaded them yet.
Example:
>>> if not gene_models.is_loaded():
... gene_models.load()
"""
with self.__lock:
return self._is_loaded
[docs]
@staticmethod
def join_gene_models(*gene_models: GeneModels) -> GeneModels:
"""Merge loaded gene models into a single, loaded GeneModels object.
The result holds every input's transcripts, answers to the first
input's resource, and is loaded (see ``is_loaded()``).
Parameters
----------
*gene_models : GeneModels
Two or more loaded ``GeneModels`` objects to merge.
Returns
-------
GeneModels
A new loaded ``GeneModels`` object that holds all
transcripts.
Raises
------
ValueError
If fewer than 2 gene models are provided, or if any of them
was never loaded: merging its empty transcripts would
silently drop everything it was meant to bring.
Example:
>>> gm1 = build_gene_models_from_file("genes1.gtf").load()
>>> gm2 = build_gene_models_from_file("genes2.gtf").load()
>>> merged = GeneModels.join_gene_models(gm1, gm2)
Note:
Transcript IDs should be unique across all input gene models;
a later input's transcript replaces an earlier one's under
the same ID.
"""
if len(gene_models) < 2:
raise ValueError("The function needs at least 2 arguments!")
for gm in gene_models:
if not gm.is_loaded():
raise ValueError(
f"gene models {gm.resource_id} were never loaded; "
f"call load() before joining them")
transcript_models: dict[str, TranscriptModel] = {}
for gm in gene_models:
transcript_models.update(gm.transcript_models)
return GeneModels.from_transcript_models(
gene_models[0].resource, transcript_models)
[docs]
@classmethod
def from_transcript_models(
cls, resource: GenomicResource,
transcript_models: dict[str, TranscriptModel],
) -> GeneModels:
"""Build a loaded, indexed gene models object over ``resource``.
The transcripts are taken as given, in the order given (the
serializers write them in that order); the resource's chrom
mapping is not applied.
Parameters
----------
resource : GenomicResource
The gene models resource the result answers to
(``resource_id``, reference genome).
transcript_models : dict[str, TranscriptModel]
Transcript ID to transcript model.
Returns
-------
GeneModels
A loaded model over ``resource``.
"""
gm = cls(resource)
gm._install(dict(transcript_models))
return gm
[docs]
def create_regions_from_genes(
gene_models: GeneModels,
genes: list[str],
regions: list[Region] | None,
gene_regions_heuristic_cutoff: int = 20,
gene_regions_heuristic_extend: int = 20000,
) -> list[Region] | None:
"""Produce a list of regions from given gene symbols.
If given a list of regions, will merge the newly-created regions
from the genes with the provided ones.
"""
assert genes is not None
assert gene_models is not None
if len(genes) == 0 or len(genes) > gene_regions_heuristic_cutoff:
return regions
gene_regions = []
for gene_name in genes:
gene_model = gene_models.gene_models_by_gene_name(gene_name)
if gene_model is None:
logger.warning("gene model for %s not found", gene_name)
continue
for gm in gene_model:
gene_regions.append( # ruff: ignore[manual-list-comprehension]
Region(
gm.chrom,
max(1, gm.tx[0] - 1 - gene_regions_heuristic_extend),
gm.tx[1] + 1 + gene_regions_heuristic_extend,
),
)
gene_regions = collapse(gene_regions)
if not regions:
regions = gene_regions or None
else:
result = []
for gene_region in gene_regions:
for region in regions:
intersection = gene_region.intersection(region)
if intersection:
result.append(intersection)
result = collapse(result)
logger.info("original regions: %s; result: %s", regions, result)
regions = result
return regions