"""Classes for handling of gene sets and gene set collections."""
from __future__ import annotations
import abc
import gzip
import json
import os
from typing import IO, Annotated, Any, Literal, cast
from pydantic import BaseModel, Field
from gain import logging
from gain.gene_sets.gene_term import (
read_ewa_set_file,
read_gmt_file,
read_mapping_file,
)
from gain.genomic_resources.fsspec_protocol import build_local_resource
from gain.genomic_resources.histogram import (
Histogram,
load_histogram,
)
from gain.genomic_resources.memo import Memo, config_memo_key
from gain.genomic_resources.repository import (
GenomicResource,
GenomicResourceRepo,
)
from gain.genomic_resources.repository_factory import (
build_genomic_resource_repository,
)
from gain.genomic_resources.resource_types import (
GENE_SET_TYPES,
LEGACY_GENE_SET_TYPE,
PREFERRED_GENE_SET_TYPE,
)
logger = logging.getLogger(__name__)
[docs]
class BaseResourceSchema(BaseModel):
type: str | None = None
meta: MetaSchema | None = None
[docs]
class ViewRangeSchema(BaseModel):
min: float | None = None
max: float | None = None
# pylint: disable=missing-class-docstring
[docs]
class NumericHistogramSchema(BaseModel):
type: Literal["number"]
plot_function: str | None = None
number_of_bins: int | None = None
view_range: ViewRangeSchema | None = None
x_log_scale: bool | None = None
y_log_scale: bool | None = None
x_min_log: float | None = None
value_order: list[str | int] | None = None
displayed_values_count: int | None = None
# pylint: disable=missing-class-docstring
[docs]
class CategoricalHistogramSchema(BaseModel):
type: Literal["categorical"]
displayed_values_count: int | None = None
displayed_values_percent: float | None = None
value_order: list[str | int] | None = None
y_log_scale: bool | None = None
label_rotation: int | None = None
plot_function: str | None = None
enforce_type: bool | None = None
natural_order: bool | None = None
HistogramConfig = Annotated[
NumericHistogramSchema | CategoricalHistogramSchema,
Field(discriminator="type"),
]
# pylint: disable=missing-class-docstring
[docs]
class GeneSetResourceSchema(BaseModel):
resource_id: str = Field(alias="id")
filename: str | None = None
directory: str | None = None
resource_format: str | None = Field(alias="format")
web_label: str | None = None
web_format_str: str | None = None
histograms: dict[
Literal["genes_per_gene_set", "gene_sets_per_gene"],
HistogramConfig,
] | None = None
[docs]
class GeneSet:
"""A named set of genes.
Attributes:
name (str): The name of the gene set, unique in its collection.
desc (str): The description of the gene set. It can be empty.
count (int): The number of gene symbols in ``syms``.
syms (list[str]): The gene symbols of the set.
"""
# pylint: disable=too-few-public-methods
name: str
desc: str
count: int
syms: list[str]
def __init__(self, name: str, desc: str, syms: list[str]) -> None:
"""Create a gene set.
Parameters
----------
name : str
The name of the gene set.
desc : str
The description of the gene set.
syms : list[str]
The gene symbols. ``count`` is set to the length of this list.
"""
self.name = name
self.desc = desc
self.count = len(syms)
self.syms = syms
def __getitem__(self, name: str) -> Any:
# This is done so that GeneSet instances and
# denovo gene set dictionaries can be accessed in a uniform way
if name == "name":
return self.name
if name == "desc":
return self.desc
if name == "count":
return self.count
if name == "syms":
return self.syms
raise KeyError
[docs]
class BaseGeneSetCollection(abc.ABC):
"""Base class for gene set collections."""
def __init__(self, collection_id: str) -> None:
self.collection_id = collection_id
self.web_label: str = ""
self.web_format_str: str = ""
self.gene_sets: dict[str, GeneSet] = {}
[docs]
@abc.abstractmethod
def load(self) -> BaseGeneSetCollection:
"""Load the gene sets from the resource."""
[docs]
@abc.abstractmethod
def get_gene_set(self, gene_set_id: str) -> GeneSet | None:
"""Return the gene set if found; returns None if not found."""
[docs]
@abc.abstractmethod
def get_all_gene_sets(self) -> list[GeneSet]:
"""Return list of all gene sets in the collection."""
[docs]
class GeneSetCollection(
BaseGeneSetCollection,
):
"""The gene sets of a ``gene_set_collection`` resource.
Build a collection with ``build_gene_set_collection_from_resource_id``
or ``build_gene_set_collection_from_resource``. The built collection
may not be loaded: call ``load()`` before the first read.
Attributes:
resource (GenomicResource): The resource that holds the gene sets.
config (GeneSetResourceSchema): The validated config of the
resource.
collection_id (str): The ``id:`` key of the resource config.
web_label (str): The label that a web interface shows for the
collection. Empty when the config sets no ``web_label``.
web_format_str (str): The format string that a web interface uses
to show a gene set. Empty when the config sets no
``web_format_str``.
gene_sets (dict[str, GeneSet]): The loaded gene sets by name. The
dictionary is empty until ``load()`` runs.
"""
def __init__(self, resource: GenomicResource) -> None:
config = resource.get_config()
self.resource = resource
self.config = GeneSetResourceSchema.model_validate(config)
super().__init__(self.config.resource_id)
assert self.collection_id != "denovo"
if resource.get_type() not in GENE_SET_TYPES:
raise ValueError("Invalid resource type for gene set collection")
if resource.get_type() == LEGACY_GENE_SET_TYPE:
logger.warning(
"'%s' resource type is deprecated; use '%s' instead",
LEGACY_GENE_SET_TYPE, PREFERRED_GENE_SET_TYPE)
self.web_label = self.config.web_label or ""
self.web_format_str = self.config.web_format_str or ""
logger.debug("loading %s: %s", self.collection_id, config)
self.gene_sets: dict[str, GeneSet] = {}
assert self.collection_id, self.gene_sets
@property
def files(self) -> set[str]:
"""Return the names of the resource files that hold the gene sets."""
res = set()
collection_format = self.config.resource_format
if collection_format == "map":
filename = self.config.filename
assert filename is not None
res.add(filename)
names_filename = filename.removesuffix(".gz")[:-4] + "names.txt"
# The same test ``load_gene_sets`` makes. Manifest membership
# would answer it too, but reading a manifest the resource has
# not got scans and writes state across the whole root (#911).
if self.resource.file_exists(names_filename):
res.add(names_filename)
elif collection_format == "gmt":
filename = self.config.filename
assert filename is not None
res.add(filename)
elif collection_format == "directory":
directory = self.config.directory
assert directory is not None
if directory == ".":
directory = ""
for filepath, _ in self.resource.get_manifest().get_files():
if filepath.startswith(directory) and \
filepath.endswith(".txt"):
res.add(filepath)
else:
raise ValueError("Invalid collection format type")
return res
[docs]
def is_loaded(self) -> bool:
"""Return whether the collection holds any loaded gene set.
Returns
-------
bool
``True`` after ``load()`` read at least one gene set. A
resource with no gene sets always reports ``False``.
"""
return bool(self.gene_sets)
[docs]
def load(self) -> GeneSetCollection:
"""Read the gene sets from the resource.
A call on a loaded collection does nothing.
Returns
-------
GeneSetCollection
This collection, so the call chains.
Raises
------
ValueError
If the ``format:`` key of the resource is not ``directory``,
``gmt`` or ``map``.
"""
if self.is_loaded():
logger.debug(
"gene sets already loaded from %s", self.collection_id)
return self
self.gene_sets = self.load_gene_sets()
logger.debug(
"loaded %d gene sets from %s",
len(self.gene_sets), self.collection_id,
)
return self
def load_gene_sets(self) -> dict[str, GeneSet]:
"""Read the gene sets from the resource files.
:meta private:
"""
assert self.resource is not None
gene_sets = {}
collection_format = self.config.resource_format
logger.debug("loading %s", self.collection_id)
if collection_format == "map":
filename = self.config.filename
assert filename is not None
names_filename = filename.removesuffix(".gz")[:-4] + "names.txt"
names_file = None
if self.resource.file_exists(names_filename):
names_file = self.resource.open_raw_file(names_filename)
map_file: IO[Any]
if filename.endswith(".gz"):
map_file = gzip.open( # ruff: ignore[open-file-with-context-handler]
self.resource.open_raw_file(filename, "rb"), "rt")
else:
map_file = self.resource.open_raw_file(filename)
gene_terms = read_mapping_file(map_file, names_file)
elif collection_format == "gmt":
filename = self.config.filename
assert filename is not None
gene_terms = read_gmt_file(self.resource.open_raw_file(filename))
elif collection_format == "directory":
directory = self.config.directory
assert directory is not None
filepaths = []
if directory == ".":
directory = "" # Easier check with startswith
for filepath, _ in self.resource.get_manifest().get_files():
if filepath.startswith(directory) and \
filepath.endswith(".txt"):
filepaths.append(filepath)
files = [self.resource.open_raw_file(f) for f in filepaths]
gene_terms = read_ewa_set_file(files)
else:
raise ValueError("Invalid collection format type")
for key, value in gene_terms.t_desc.items():
syms = list(gene_terms.t2g[key].keys())
gene_set = GeneSet(key, value, syms)
gene_sets[gene_set.name] = gene_set
return gene_sets
[docs]
def get_gene_set(self, gene_set_id: str) -> GeneSet | None:
"""Return one gene set by its name.
Parameters
----------
gene_set_id : str
The name of the gene set.
Returns
-------
GeneSet | None
The gene set, or ``None`` when the collection has no set with
this name. The collection logs a warning in that case. An
unloaded collection has no sets, so it always returns ``None``.
"""
gene_set = self.gene_sets.get(gene_set_id)
if gene_set is None:
logger.warning(
"%s not found in %s", gene_set_id, self.gene_sets.keys(),
)
return gene_set
[docs]
def get_all_gene_sets(self) -> list[GeneSet]:
"""Return all the gene sets of the collection.
Returns
-------
list[GeneSet]
The gene sets in the order of the resource. The list is empty
when the collection is not loaded.
"""
return list(self.gene_sets.values())
def get_genes_per_gene_set_hist_image_filename(self) -> str:
"""Return the resource path of the genes per gene set plot.
:meta private:
"""
return "statistics/genes_per_gene_set_histogram.png"
def get_genes_per_gene_set_hist_filename(self) -> str:
"""Return the resource path of the genes per gene set histogram.
:meta private:
"""
return "statistics/genes_per_gene_set_histogram.json"
[docs]
def get_genes_per_gene_set_hist(self) -> Histogram | None:
"""Read the stored histogram of the number of genes in a gene set.
Returns
-------
Histogram | None
The histogram, or ``None`` when the resource has no stored
histogram because no statistics build ran.
"""
return self._load_stored_histogram(
self.get_genes_per_gene_set_hist_filename())
def get_gene_sets_per_gene_hist_image_filename(self) -> str:
"""Return the resource path of the gene sets per gene plot.
:meta private:
"""
return "statistics/gene_sets_per_gene_histogram.png"
def get_gene_sets_per_gene_hist_filename(self) -> str:
"""Return the resource path of the gene sets per gene histogram.
:meta private:
"""
return "statistics/gene_sets_per_gene_histogram.json"
[docs]
def get_gene_sets_per_gene_hist(self) -> Histogram | None:
"""Read the stored histogram of the gene sets that contain a gene.
Returns
-------
Histogram | None
The histogram, or ``None`` when the resource has no stored
histogram because no statistics build ran.
"""
return self._load_stored_histogram(
self.get_gene_sets_per_gene_hist_filename())
def _load_stored_histogram(self, filename: str) -> Histogram | None:
if not self.resource.file_exists(filename):
return None
return load_histogram(self.resource, filename)
[docs]
def get_gene_sets_list_statistics(self) -> list[dict] | None:
"""Read the stored list of gene sets with their sizes.
Returns
-------
list[dict] | None
One dictionary for each gene set, with the keys ``name``,
``count`` and ``desc``, ordered by falling ``count``. ``None``
when the resource has no stored list because no statistics
build ran.
"""
try:
with self.resource.proto.open_raw_file(
self.resource,
"statistics/gene_sets_list_statistics.json",
"rt",
) as statistics_file:
return cast(list, json.load(statistics_file))
except FileNotFoundError:
return None
[docs]
def get_gene_collection_count_statistics(self) -> dict | None:
"""Read the stored counts of the collection.
Returns
-------
dict | None
A dictionary with the keys ``number_of_gene_sets`` and
``number_of_unique_genes``. ``None`` when the resource has no
stored counts because no statistics build ran.
"""
try:
with self.resource.proto.open_raw_file(
self.resource,
"statistics/gene_collection_count_statistics.json",
"rt",
) as statistics_file:
return cast(dict, json.load(statistics_file))
except FileNotFoundError:
return None
_RESOURCE_CACHE: Memo[tuple[str, str, str], GeneSetCollection] = Memo()
_FILE_CACHE: Memo[tuple[str, str], GeneSetCollection] = Memo()
_FORMAT_BY_EXTENSION = {
".txt": "map",
".gmt": "gmt",
}
def _detect_collection_format(filename: str) -> str:
"""Return the collection format implied by ``filename``."""
if os.path.isdir(filename):
return "directory"
extension = os.path.splitext(filename)[1]
if extension not in _FORMAT_BY_EXTENSION:
raise ValueError("Cannot find collection format automatically")
return _FORMAT_BY_EXTENSION[extension]
[docs]
def build_gene_set_collection_from_file(
filename: str,
collection_id: str | None = None,
collection_format: str | None = None,
web_label: str | None = None,
web_format_str: str | None = None,
) -> GeneSetCollection:
"""Return a Gene Set Collection by adapting a file to a local resource."""
# Normalising is what puts the containing directory -- for a bare relative
# name, the working directory -- into the cache key.
filename = os.path.abspath(filename)
dirname = os.path.dirname(filename)
basename = os.path.basename(filename)
if collection_format is None:
collection_format = _detect_collection_format(filename)
if collection_id is None:
collection_id = basename
config: dict[str, Any] = {
"type": "gene_set_collection",
"id": collection_id,
"format": collection_format,
"web_label": web_label,
"web_format_str": web_format_str,
}
# A single file format is addressed by basename from the directory that
# holds it -- the dirname/basename split ADR 0010 records for these
# factories. The directory format has no such containing directory to
# fall back on: rooting it at the parent would put every unrelated
# sibling inside the resource, and reading a manifest scans -- and
# writes state beside -- everything under the root (#911).
if collection_format == "directory":
root = filename
config["directory"] = "."
else:
root = dirname
config["filename"] = basename
# Keyed on the whole config so that every config-shaping argument --
# present and future -- participates in the key. A resource id plus repo
# url identifies a resource only when it is a subdirectory of a
# repository; this one is the repository root, so which file it describes
# lives in the config alone. Hence a cache of its own, and a collection
# built directly rather than through the resource-keyed factory (#894).
return _FILE_CACHE.get_or_build(
(filename, config_memo_key(config)),
lambda: GeneSetCollection(build_local_resource(root, config)))
[docs]
def build_gene_set_collection_from_resource(
resource: GenomicResource,
) -> GeneSetCollection:
"""Build the gene set collection of a genomic resource.
The returned collection may not be loaded. Calls with the same
resource return the same object, so an earlier caller may have loaded
it.
Parameters
----------
resource : GenomicResource
A resource of type ``gene_set_collection``.
Returns
-------
GeneSetCollection
The collection of the resource. Check ``is_loaded()`` or call
``load()`` before any read.
Raises
------
ValueError
If the resource is not of type ``gene_set_collection``.
"""
return _RESOURCE_CACHE.get_or_build(
resource.get_memo_key(), lambda: GeneSetCollection(resource))
[docs]
def build_gene_set_collection_from_resource_id(
resource_id: str, grr: GenomicResourceRepo | None = None,
) -> GeneSetCollection:
"""Build the gene set collection of the resource with a given id.
The returned collection may not be loaded. Calls that find the same
resource return the same object, so an earlier caller may have loaded
it.
Parameters
----------
resource_id : str
The id of a resource of type ``gene_set_collection`` in the
repository.
grr : GenomicResourceRepo | None
The repository that resolves ``resource_id``. When it is
``None``, the function builds the default repository with
``build_genomic_resource_repository()``.
Returns
-------
GeneSetCollection
The collection of the resource. Check ``is_loaded()`` or call
``load()`` before any read.
Raises
------
FileNotFoundError
If a protocol-backed repository has no resource with this id.
ValueError
If a group repository has no resource with this id, or if the
resource is not of type ``gene_set_collection``.
"""
if grr is None:
grr = build_genomic_resource_repository()
return build_gene_set_collection_from_resource(
grr.get_resource(resource_id))