diff --git a/bbconf/_version.py b/bbconf/_version.py index 350cbe9d..ea370a8e 100644 --- a/bbconf/_version.py +++ b/bbconf/_version.py @@ -1 +1 @@ -__version__ = "0.11.4" +__version__ = "0.12.0" diff --git a/bbconf/bbagent.py b/bbconf/bbagent.py index 68c318e2..59a58b6d 100644 --- a/bbconf/bbagent.py +++ b/bbconf/bbagent.py @@ -1,10 +1,12 @@ import logging from functools import cached_property from pathlib import Path -from typing import List, Union +from typing import List, Union, Dict +import numpy as np +import statistics from sqlalchemy.orm import Session -from sqlalchemy.sql import distinct, func, select +from sqlalchemy.sql import distinct, func, select, and_, or_ from bbconf.config_parser.bedbaseconfig import BedBaseConfig from bbconf.db_utils import ( @@ -16,8 +18,20 @@ UsageBedMeta, UsageFiles, UsageSearch, + GeoGsmStatus, + BedStats, + Files, +) +from bbconf.models.base_models import ( + StatsReturn, + UsageModel, + FileStats, + UsageStats, + AllFilesInfo, + FileInfo, + BinValues, + GEOStatistics, ) -from bbconf.models.base_models import StatsReturn, UsageModel, FileStats, UsageStats from bbconf.modules.bedfiles import BedAgentBedFile from bbconf.modules.bedsets import BedAgentBedSet from bbconf.modules.objects import BBObjects @@ -97,6 +111,7 @@ def get_detailed_stats(self, concise: bool = False) -> FileStats: _LOGGER.info("Getting detailed statistics for all bed files") with Session(self.config.db_engine.engine) as session: + bed_compliance = { f[0]: f[1] for f in session.execute( @@ -131,30 +146,89 @@ def get_detailed_stats(self, concise: bool = False) -> FileStats: .order_by(func.count(BedMetadata.species_name).desc()) ).all() } + file_assay = { + f[0]: f[1] + for f in session.execute( + select(BedMetadata.assay, func.count(BedMetadata.assay)) + .group_by(BedMetadata.assay) + .order_by(func.count(BedMetadata.assay).desc()) + ).all() + } slice_value = 20 + bed_comments = self._stats_comments(session) + geo_status = self._stats_geo_status(session) + + bedfiles_info = self.bed_files_info() + + number_of_regions = [bed.number_of_regions for bed in bedfiles_info.files] + list_mean_width = [bed.mean_region_width for bed in bedfiles_info.files] + list_file_size = [bed.file_size for bed in bedfiles_info.files] + + number_of_regions_bins = self._bin_number_of_regions(number_of_regions) + list_mean_width_bins = self._bin_mean_region_width(list_mean_width) + list_file_size_bins = self._bin_file_size(list_file_size) + + geo_stats = self._get_geo_stats(session) + if concise: bed_compliance_concise = dict(list(bed_compliance.items())[0:slice_value]) bed_compliance_concise["other"] = sum( list(bed_compliance.values())[slice_value:] ) + if "" in bed_compliance_concise: + bed_compliance_concise["other"] = ( + bed_compliance_concise["other"] + bed_compliance_concise[""] + ) + bed_compliance_concise.pop("") file_genomes_concise = dict(list(file_genomes.items())[0:slice_value]) file_genomes_concise["other"] = sum( list(file_genomes.values())[slice_value:] ) + file_genomes.get("other", 0) + if "" in file_genomes_concise: + file_genomes_concise["other"] = ( + file_genomes_concise["other"] + file_genomes_concise[""] + ) + file_genomes_concise.pop("") file_organism_concise = dict(list(file_organism.items())[0:slice_value]) file_organism_concise["other"] = sum( list(file_organism.values())[slice_value:] ) + file_organism.get("other", 0) + if "" in file_organism_concise: + file_organism_concise["other"] = ( + file_organism_concise["other"] + file_organism_concise[""] + ) + file_organism_concise.pop("") + file_assay_concise = dict(list(file_assay.items())[0:slice_value]) + file_assay_concise["other"] = sum( + list(file_assay.values())[slice_value:] + ) + file_assay.get("other", 0) + if "" in file_assay_concise: + file_assay_concise["other"] = ( + file_assay_concise["other"] + file_assay_concise[""] + ) + file_assay_concise.pop("") + if "OTHER" in file_assay_concise: + file_assay_concise["other"] = ( + file_assay_concise["other"] + file_assay_concise["OTHER"] + ) + file_assay_concise.pop("OTHER") return FileStats( data_format=data_format, bed_compliance=bed_compliance_concise, file_genome=file_genomes_concise, file_organism=file_organism_concise, + file_assay=file_assay_concise, + bed_comments=bed_comments, + geo_status=geo_status, + mean_region_width=list_mean_width_bins, + file_size=list_file_size_bins, + number_of_regions=number_of_regions_bins, + geo=geo_stats, ) return FileStats( @@ -162,6 +236,13 @@ def get_detailed_stats(self, concise: bool = False) -> FileStats: bed_compliance=bed_compliance, file_genome=file_genomes, file_organism=file_organism, + file_assay=file_assay, + bed_comments=bed_comments, + geo_status=geo_status, + mean_region_width=list_mean_width_bins, + file_size=list_file_size_bins, + number_of_regions=number_of_regions_bins, + geo=geo_stats, ) def get_detailed_usage(self) -> UsageStats: @@ -203,6 +284,7 @@ def get_detailed_usage(self) -> UsageStats: .order_by(func.sum(UsageSearch.count).desc()) .limit(20) ).all() + if f[0] } bedset_search_terms = { @@ -214,6 +296,17 @@ def get_detailed_usage(self) -> UsageStats: .order_by(func.sum(UsageSearch.count).desc()) .limit(20) ).all() + if f[0] + } + + bed_downloads = { + f[0].split("/")[-1].split(".")[0]: f[1] + for f in session.execute( + select(UsageFiles.file_path, func.sum(UsageFiles.count)) + .group_by(UsageFiles.file_path) + .order_by(func.sum(UsageFiles.count).desc()) + .limit(20) + ).all() } return UsageStats( @@ -221,6 +314,7 @@ def get_detailed_usage(self) -> UsageStats: bedset_metadata=bedset_metadata, bed_search_terms=bed_search_terms, bedset_search_terms=bedset_search_terms, + bed_downloads=bed_downloads, ) def get_list_genomes(self) -> List[str]: @@ -238,6 +332,22 @@ def get_list_genomes(self) -> List[str]: genomes = session.execute(statement).all() return [result[0] for result in genomes if result[0]] + def get_list_assays(self): + """ + Get list of genomes from the database + + :return: list of genomes + """ + + with Session(self.config.db_engine.engine) as session: + statement = ( + select(BedMetadata.assay) + .group_by(BedMetadata.assay) + .order_by(func.count(BedMetadata.assay).desc()) + ) + results = session.execute(statement).all() + return [result[0] for result in results if result[0]] + @cached_property def list_of_licenses(self) -> List[str]: """ @@ -364,3 +474,307 @@ def add_usage(self, stats: UsageModel) -> None: session.add(new_stats) session.commit() + + def _stats_comments(self, sa_session: Session) -> Dict[str, int]: + """ + Get statistics about comments that are present in bed files. + + :param sa_session: SQLAlchemy session + + :return: Dict[str, int] + """ + _LOGGER.info("Analyzing bed table for comments in bed files...") + + total_statement = select(func.count(Bed.id)).where(Bed.header.is_not(None)) + correct_statement = select(func.count(Bed.id)).where(Bed.header.like("#%")) + track_name_statement = select(func.count(Bed.id)).where( + Bed.header.like("track name%") + ) + track_type_statement = select(func.count(Bed.id)).where( + Bed.header.like("track type%") + ) + browser_statement = select(func.count(Bed.id)).where( + Bed.header.like("browser%") + ) + + total_bed_comments = ( + sa_session.execute(total_statement).one_or_none() or (0,) + )[0] + correct_bed_comments = ( + sa_session.execute(correct_statement).one_or_none() or (0,) + )[0] + track_name_comments = ( + sa_session.execute(track_name_statement).one_or_none() or (0,) + )[0] + track_type_comments = ( + sa_session.execute(track_type_statement).one_or_none() or (0,) + )[0] + browser_comments = ( + sa_session.execute(browser_statement).one_or_none() or (0,) + )[0] + + header_comments = ( + total_bed_comments + - correct_bed_comments + - track_name_comments + - track_type_comments + - browser_comments + ) + + return { + "correct_bed_comments": correct_bed_comments, + "track_name_comments": track_name_comments, + "track_type_comments": track_type_comments, + "browser_comments": browser_comments, + "header_comments": header_comments, + } + + def _stats_geo_status(self, sa_session: Session) -> Dict[str, int]: + """ + Get statistics about status of GEO bed file processing. + + :param sa_session: SQLAlchemy session + :return Dict[str, int] + """ + + success_statement = select( + func.count(distinct(GeoGsmStatus.sample_name)) + ).where(GeoGsmStatus.status == "SUCCESS") + + failed_count_statement = select( + func.count(distinct(GeoGsmStatus.sample_name)) + ).where(and_(GeoGsmStatus.status == "FAIL")) + mean_region_width_statement = select( + func.count(distinct(GeoGsmStatus.sample_name)) + ).where( + or_( + GeoGsmStatus.error.like("%Initial QC failed for%"), + GeoGsmStatus.error.like("%Quality control failed%"), + ) + ) + size_greater_then_statement = select( + func.count(distinct(GeoGsmStatus.sample_name)) + ).where(GeoGsmStatus.error.like("%File size is too big.%")) + + success_count = (sa_session.execute(success_statement).one_or_none() or (0,))[0] + failed_count = ( + sa_session.execute(failed_count_statement).one_or_none() or (0,) + )[0] + + mean_region_width_count = ( + sa_session.execute(mean_region_width_statement).one_or_none() or (0,) + )[0] + size_greater_than_count = ( + sa_session.execute(size_greater_then_statement).one_or_none() or (0,) + )[0] + corrupted_file_count = ( + failed_count - mean_region_width_count - size_greater_than_count + ) + + return { + "success": success_count, + "mean_region_width_lower_10": mean_region_width_count, + "size_greater_20": size_greater_than_count, + "corrupted_files": corrupted_file_count, + } + + def bed_files_info(self) -> AllFilesInfo: + """ + Get information about all bed files in bedbase. + + :param sa_session: SQLAlchemy session + :return AllFilesInfo: + { + "total": int," + "files": [ + { id: str + bed_compliance: str + data_format: str + mean_region_width: float + file_size: int + number_of_regions: int + }, + ... ] + } + """ + + all_files_statement = ( + select( + Bed.id, + Bed.bed_compliance, + Bed.data_format, + BedStats.mean_region_width, + Files.size, + BedStats.number_of_regions, + ) + .join(BedStats, BedStats.id == Bed.id) # Explicit join condition + .join(Files, Files.bedfile_id == Bed.id) # Explicit join condition + .where(Files.name == "bed_file") + ) + + results = [] + error_list = [] + with Session(self.config.db_engine.engine) as session: + sql_result = session.execute(all_files_statement).all() + + for single_bed in sql_result: + try: + results.append( + FileInfo( + id=single_bed[0], + bed_compliance=single_bed[1], + data_format=single_bed[2], + mean_region_width=single_bed[3], + file_size=single_bed[4], + number_of_regions=single_bed[5], + ) + ) + except Exception: + error_list.append(single_bed[0]) + + _LOGGER.info( + f"Number of bed records with unknown regions and region width {len(error_list)}" + ) + return AllFilesInfo( + total=len(results), + files=results, + ) + + def _bin_number_of_regions(self, number_of_regions: list) -> BinValues: + """ + Create bins for number of regions in bed files + + :param number_of_regions: list of number of regions in bed files + :return: BinValues object containing bins and values + """ + + max_value_threshold = 400_000 # set a threshold for maximum value to avoid outliers in the histogram + + filtered_number_of_regions = [ + x if x <= max_value_threshold else max_value_threshold + 1 + for x in number_of_regions + ] + + n_region_counts, n_region_bin_edges = np.histogram( + filtered_number_of_regions, bins=100 + ) + n_region_counts = n_region_counts.astype(int).tolist() + n_region_bin_edges = n_region_bin_edges.astype(int).tolist() + + return BinValues( + bins=n_region_bin_edges, + counts=n_region_counts, + mean=round(statistics.mean(number_of_regions), 2), + median=round(statistics.median(number_of_regions), 2), + ) + + def _bin_mean_region_width(self, mean_region_widths: list) -> BinValues: + """ + Create bins for number of regions in bed files + + :param mean_region_widths: list of mean region widths in bed files + :return: BinValues object containing bins and values + """ + + max_value_threshold = 5_000 # set a threshold for maximum value to avoid outliers in the histogram + + filtered_mean_region_widths = [ + x if x <= max_value_threshold else max_value_threshold + 1 + for x in mean_region_widths + ] + + mean_reg_width_counts, mean_reg_width_bin_edges = np.histogram( + filtered_mean_region_widths, bins=100 + ) + mean_reg_width_counts = mean_reg_width_counts.tolist() + mean_reg_width_bin_edges = mean_reg_width_bin_edges.tolist() + + return BinValues( + bins=mean_reg_width_bin_edges, + counts=mean_reg_width_counts, + mean=round(statistics.mean(mean_region_widths), 2), + median=round(statistics.median(mean_region_widths), 2), + ) + + def _bin_file_size(self, list_file_size: list) -> BinValues: + """ + Create bins for number of regions in bed files + + :param list_file_size: list of bed file sizes in bytes + :return: BinValues object containing bins and values + """ + + max_value_threshold = 10 * 1024 * 1024 + + filtered_list_file_size = [ + x for x in list_file_size if x <= max_value_threshold + ] + + filtered_list_file_size = [x / (1024 * 1024) for x in filtered_list_file_size] + + file_size_counts, file_size_bin_edges = np.histogram( + filtered_list_file_size, bins=100 + ) + file_size_counts = file_size_counts.astype(int).tolist() + file_size_bin_edges = file_size_bin_edges.astype(float).tolist() + + return BinValues( + bins=file_size_bin_edges, + counts=file_size_counts, + mean=round(statistics.mean(filtered_list_file_size), 2), + median=round(statistics.median(filtered_list_file_size), 2), + ) + + def _get_geo_stats(self, sa_session: Session) -> GEOStatistics: + """ + Get GEO statistics for the bedbase platform. + + :return: GEOStatistics + """ + + _LOGGER.info("Getting GEO statistics.") + + statement = select( + GeoGsmStatus.bed_id, + GeoGsmStatus.source_submission_date, + GeoGsmStatus.file_size, + ).distinct(GeoGsmStatus.sample_name) + + results = sa_session.execute(statement).all() + + years = [] + file_sizes = [] + + for result in results: + if result[1]: + years.append(result[1].year) + if result[2]: + file_sizes.append(result[2]) + + years_array = np.array([y for y in years if y is not None]) + unique_years, counts = np.unique(years_array, return_counts=True) + + cumulative_counts = np.cumsum(counts) + unique_years = unique_years.astype(str).tolist() + + years_cumulative = dict(zip(unique_years, cumulative_counts)) + years_numbers = dict(zip(unique_years, counts)) + + MAX_FILE_SIZE = 50 * 1024 * 1024 # 50 MB + file_size_filter = [ + x if x <= MAX_FILE_SIZE else MAX_FILE_SIZE + 1 for x in file_sizes + ] + + file_size_filter = [x / (1024 * 1024) for x in file_size_filter] + file_size_counts, file_size_bin_edges = np.histogram(file_size_filter, bins=100) + + return GEOStatistics( + number_of_files=years_numbers, + cumulative_number_of_files=years_cumulative, + file_sizes=BinValues( + bins=list(file_size_bin_edges), + counts=file_size_counts.astype(int).tolist(), + mean=round(statistics.mean(file_sizes), 2), + median=round(statistics.median(file_sizes), 2), + ), + ) diff --git a/bbconf/config_parser/bedbaseconfig.py b/bbconf/config_parser/bedbaseconfig.py index 879eb821..a6cedc37 100644 --- a/bbconf/config_parser/bedbaseconfig.py +++ b/bbconf/config_parser/bedbaseconfig.py @@ -6,6 +6,7 @@ import boto3 import qdrant_client +from qdrant_client import QdrantClient, models import s3fs import yacman import zarr @@ -60,6 +61,7 @@ def __init__(self, config: Union[Path, str], init_ml: bool = True): self._qdrant_engine = self._init_qdrant_backend() self._qdrant_text_engine = self._init_qdrant_text_backend() + self._qdrant_advanced_engine = self._init_qdrant_advanced_backend() if init_ml: self._b2bsi = self._init_b2bsi_object() @@ -67,7 +69,7 @@ def __init__(self, config: Union[Path, str], init_ml: bool = True): self._bivec = self._init_bivec_object() else: _LOGGER.info( - f"Skipping initialization of ML models, init_ml parameter set to False." + "Skipping initialization of ML models, init_ml parameter set to False." ) self._b2bsi = None @@ -207,7 +209,7 @@ def _init_db_engine(self) -> BaseEngine: Create database engine object using credentials provided in config file """ - _LOGGER.info(f"Initializing database engine...") + _LOGGER.info("Initializing database engine...") return BaseEngine( host=self._config.database.host, port=self._config.database.port, @@ -224,7 +226,7 @@ def _init_qdrant_backend(self) -> QdrantBackend: :return: QdrantClient """ - _LOGGER.info(f"Initializing qdrant engine...") + _LOGGER.info("Initializing qdrant engine...") try: return QdrantBackend( collection=self._config.qdrant.file_collection, @@ -247,7 +249,7 @@ def _init_qdrant_text_backend(self) -> Union[QdrantBackend, None]: :return: QdrantClient """ - _LOGGER.info(f"Initializing qdrant text engine...") + _LOGGER.info("Initializing qdrant text engine...") try: return QdrantBackend( dim=TEXT_EMBEDDING_DIMENSION, @@ -264,6 +266,65 @@ def _init_qdrant_text_backend(self) -> Union[QdrantBackend, None]: ) return None + def _init_qdrant_advanced_backend(self) -> Union[QdrantClient, None]: + """ + Create qdrant client text embedding object using credentials provided in config file + + :return: QdrantClient + """ + + COLLECTION_NAME = self.config.qdrant.search_collection + DIMENSIONS = 384 + + _LOGGER.info("Initializing qdrant text advanced engine...") + + try: + qdrant_cl = QdrantClient( + url=self.config.qdrant.host, + port=self.config.qdrant.port, + api_key=self.config.qdrant.api_key, + ) + + if not qdrant_cl.collection_exists(COLLECTION_NAME): + _LOGGER.info( + "Collection 'bedbase_query_search' does not exist, creating it." + ) + qdrant_cl.create_collection( + collection_name=COLLECTION_NAME, + vectors_config=models.VectorParams( + size=DIMENSIONS, distance=models.Distance.COSINE + ), + quantization_config=models.ScalarQuantization( + scalar=models.ScalarQuantizationConfig( + type=models.ScalarType.INT8, + quantile=0.99, + always_ram=True, + ), + ), + ) + + qdrant_cl.create_payload_index( + collection_name=COLLECTION_NAME, + field_name="assay", + field_schema="keyword", + ) + qdrant_cl.create_payload_index( + collection_name=COLLECTION_NAME, + field_name="genome_alias", + field_schema="keyword", + ) + + return qdrant_cl + + except qdrant_client.http.exceptions.ResponseHandlingException as err: + _LOGGER.error( + f"Error in Connection to qdrant! skipping... Error: {err}. Qdrant host: {self._config.qdrant.host}" + ) + warnings.warn( + f"error in Connection to qdrant! skipping... Error: {err}", UserWarning + ) + return None + def _init_bivec_object(self) -> Union[BiVectorSearchInterface, None]: """ Create BiVectorSearchInterface object using credentials provided in config file @@ -271,11 +332,11 @@ def _init_bivec_object(self) -> Union[BiVectorSearchInterface, None]: :return: BiVectorSearchInterface """ - _LOGGER.info(f"Initializing BiVectorBackend...") + _LOGGER.info("Initializing BiVectorBackend...") search_backend = BiVectorBackend( metadata_backend=self._qdrant_text_engine, bed_backend=self._qdrant_engine ) - _LOGGER.info(f"Initializing BiVectorSearchInterface...") + _LOGGER.info("Initializing BiVectorSearchInterface...") search_interface = BiVectorSearchInterface( backend=search_backend, query2vec=self.config.path.text2vec, @@ -289,7 +350,7 @@ def _init_b2bsi_object(self) -> Union[BED2BEDSearchInterface, None]: :return: Bed2BEDSearchInterface object """ try: - _LOGGER.info(f"Initializing search interfaces...") + _LOGGER.info("Initializing search interfaces...") return BED2BEDSearchInterface( backend=self.qdrant_engine, query2vec=BED2Vec(model=self._config.path.region2vec), @@ -309,13 +370,15 @@ def _init_pephubclient() -> Union[PEPHubClient, None]: :return: PephubClient """ - try: - _LOGGER.info(f"Initializing PEPHub client...") - return PEPHubClient() - except Exception as e: - _LOGGER.error(f"Error in creating PephubClient object: {e}") - warnings.warn(f"Error in creating PephubClient object: {e}", UserWarning) - return None + + # try: + # _LOGGER.info("Initializing PEPHub client...") + # return PEPHubClient() + # except Exception as e: + # _LOGGER.error(f"Error in creating PephubClient object: {e}") + # warnings.warn(f"Error in creating PephubClient object: {e}", UserWarning) + # return None + return None def _init_boto3_client( self, @@ -342,7 +405,7 @@ def _init_r2v_object(self) -> Union[Region2VecExModel, None]: Create Region2VecExModel object using credentials provided in config file """ try: - _LOGGER.info(f"Initializing R2V object...") + _LOGGER.info("Initializing R2V object...") return Region2VecExModel(self.config.path.region2vec) except Exception as e: _LOGGER.error(f"Error in creating Region2VecExModel object: {e}") diff --git a/bbconf/config_parser/const.py b/bbconf/config_parser/const.py index 1c692880..b9541f51 100644 --- a/bbconf/config_parser/const.py +++ b/bbconf/config_parser/const.py @@ -7,6 +7,7 @@ DEFAULT_QDRANT_PORT = 6333 DEFAULT_QDRANT_COLLECTION_NAME = "bedbase" DEFAULT_QDRANT_TEXT_COLLECTION_NAME = "bed_text" +DEFAULT_QDRANT_SEARCH_COLLECTION_NAME = "bedbase_query_search" DEFAULT_QDRANT_API_KEY = None DEFAULT_SERVER_PORT = 80 diff --git a/bbconf/config_parser/models.py b/bbconf/config_parser/models.py index 5eb75838..13726a62 100644 --- a/bbconf/config_parser/models.py +++ b/bbconf/config_parser/models.py @@ -21,6 +21,7 @@ DEFAULT_SERVER_HOST, DEFAULT_SERVER_PORT, DEFAULT_TEXT2VEC_MODEL, + DEFAULT_QDRANT_SEARCH_COLLECTION_NAME, ) _LOGGER = logging.getLogger(__name__) @@ -54,6 +55,7 @@ class ConfigQdrant(BaseModel): api_key: Optional[str] = None file_collection: str = DEFAULT_QDRANT_COLLECTION_NAME text_collection: Optional[str] = DEFAULT_QDRANT_TEXT_COLLECTION_NAME + search_collection: Optional[str] = DEFAULT_QDRANT_SEARCH_COLLECTION_NAME class ConfigServer(BaseModel): diff --git a/bbconf/db_utils.py b/bbconf/db_utils.py index aea952d2..d49902a8 100644 --- a/bbconf/db_utils.py +++ b/bbconf/db_utils.py @@ -1,6 +1,5 @@ import datetime import logging -import os from typing import List, Optional import pandas as pd @@ -101,6 +100,10 @@ class Bed(Base): indexed: Mapped[bool] = mapped_column( default=False, comment="Whether sample was added to qdrant" ) + file_indexed: Mapped[bool] = mapped_column( + default=False, + comment="Whether file was tokenized and added to the vector database", + ) pephub: Mapped[bool] = mapped_column( default=False, comment="Whether sample was added to pephub" ) @@ -141,7 +144,7 @@ class Bed(Base): ) license_mapping: Mapped["License"] = relationship("License", back_populates="bed") - ref_classifier: Mapped["GenomeRefStats"] = relationship( + ref_classifier: Mapped[List["GenomeRefStats"]] = relationship( "GenomeRefStats", back_populates="bed", cascade="all, delete-orphan" ) processed: Mapped[bool] = mapped_column( @@ -413,6 +416,21 @@ class License(Base): bed: Mapped[List["Bed"]] = relationship("Bed", back_populates="license_mapping") +class ReferenceGenome(Base): + __tablename__ = "reference_genomes" + + digest: Mapped[str] = mapped_column(primary_key=True, index=True) + alias: Mapped[str] = mapped_column( + nullable=False, comment="Name of the reference genome" + ) + + bed_reference: Mapped[List["GenomeRefStats"]] = relationship( + "GenomeRefStats", + back_populates="genome_object", + cascade="all, delete-orphan", + ) + + class GenomeRefStats(Base): __tablename__ = "genome_ref_stats" @@ -427,6 +445,9 @@ class GenomeRefStats(Base): compared_genome: Mapped[str] = mapped_column( nullable=False, comment="Compared Genome" ) + genome_digest: Mapped[str] = mapped_column( + ForeignKey("reference_genomes.digest", ondelete="CASCADE"), + ) xs: Mapped[float] = mapped_column(nullable=True, default=None) oobr: Mapped[float] = mapped_column(nullable=True, default=None) @@ -436,6 +457,12 @@ class GenomeRefStats(Base): bed: Mapped["Bed"] = relationship("Bed", back_populates="ref_classifier") + genome_object: Mapped["ReferenceGenome"] = relationship( + "ReferenceGenome", + back_populates="bed_reference", + lazy="joined", + ) + __table_args__ = (UniqueConstraint("bed_id", "compared_genome"),) diff --git a/bbconf/models/base_models.py b/bbconf/models/base_models.py index 69fe02ec..a7711857 100644 --- a/bbconf/models/base_models.py +++ b/bbconf/models/base_models.py @@ -26,19 +26,13 @@ class StatsReturn(BaseModel): genomes_number: int = 0 -class FileStats(BaseModel): - bed_compliance: Dict[str, int] - data_format: Dict[str, int] - file_genome: Dict[str, int] - file_organism: Dict[str, int] - - class UsageStats(BaseModel): # file_downloads: Dict[str, int] # Placeholder for tracking file download statistics in the future. bed_metadata: Dict[str, int] bedset_metadata: Dict[str, int] bed_search_terms: Dict[str, int] bedset_search_terms: Dict[str, int] + bed_downloads: Dict[str, int] class UsageModel(BaseModel): @@ -55,3 +49,57 @@ class UsageModel(BaseModel): date_from: datetime.datetime date_to: Union[datetime.datetime, None] = None + + +class FileInfo(BaseModel): + """ + Main information about a file used for BEDbase verse statistics. + """ + + id: str + bed_compliance: str + data_format: str + mean_region_width: float + file_size: int + number_of_regions: int + + +class AllFilesInfo(BaseModel): + """ + Information about all files. e.g. file sizes, mean region width, etc. + + """ + + total: int + files: List[FileInfo] + + +class BinValues(BaseModel): + bins: List[Union[int, float, str]] + counts: List[int] + mean: float + median: float + + +class GEOStatistics(BaseModel): + """ + GEO statistics for files. + """ + + number_of_files: Dict[str, int] + cumulative_number_of_files: Dict[str, int] + file_sizes: BinValues + + +class FileStats(BaseModel): + bed_compliance: Dict[str, int] + data_format: Dict[str, int] + file_genome: Dict[str, int] + file_organism: Dict[str, int] + file_assay: Dict[str, int] + geo_status: Dict[str, int] + bed_comments: Dict[str, int] + mean_region_width: BinValues + file_size: BinValues + number_of_regions: BinValues + geo: GEOStatistics diff --git a/bbconf/models/bed_models.py b/bbconf/models/bed_models.py index f101c299..a6059568 100644 --- a/bbconf/models/bed_models.py +++ b/bbconf/models/bed_models.py @@ -85,7 +85,7 @@ class BedPEPHub(BaseModel): species_id: str = "" cell_type: str = "" cell_line: str = "" - exp_protocol: str = Field("", description="Experimental protocol (e.g. ChIP-seq)") + assay: str = Field("", description="Experimental protocol (e.g. ChIP-seq)") library_source: str = Field( "", description="Library source (e.g. genomic, transcriptomic)" ) @@ -130,7 +130,8 @@ class StandardMeta(BaseModel): "", description="Library source (e.g. genomic, transcriptomic)" ) assay: str = Field( - "", description="Experimental protocol (e.g. ChIP-seq)", alias="exp_protocol" + "", + description="Experimental protocol (e.g. ChIP-seq)", ) antibody: str = Field("", description="Antibody used in the assay") target: str = Field("", description="Target of the assay (e.g. H3K4me3)") @@ -243,7 +244,8 @@ class TokenizedPathResponse(BaseModel): class RefGenValidModel(BaseModel): provided_genome: str - compared_genome: str + compared_genome: Union[str, None] + genome_digest: Union[str, None] xs: float = 0.0 oobr: Union[float, None] = None sequence_fit: Union[float, None] = None @@ -257,3 +259,22 @@ class RefGenValidReturnModel(BaseModel): id: str provided_genome: Union[str, None] = None compared_genome: List[RefGenValidModel] + + +class VectorMetadata(BaseModel): + id: str + name: str + description: str + cell_line: str + cell_type: str + tissue: str + target: str + treatment: str + assay: str + genome_alias: str + genome_digest: Union[str, None] = None + species_name: str + # summary: str + # global_sample_id: str + # original_file_name: str + # embedding: List[float] diff --git a/bbconf/modules/bedfiles.py b/bbconf/modules/bedfiles.py index 68831e51..cad3e1a0 100644 --- a/bbconf/modules/bedfiles.py +++ b/bbconf/modules/bedfiles.py @@ -12,11 +12,14 @@ from pydantic import BaseModel from qdrant_client.http.models import PointStruct from qdrant_client.models import Distance, PointIdsList, VectorParams -from sqlalchemy import and_, delete, func, or_, select +from sqlalchemy import and_, delete, func, or_, select, cast from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session, aliased from sqlalchemy.orm.attributes import flag_modified +from sqlalchemy.dialects import postgresql from tqdm import tqdm +from fastembed import TextEmbedding +from qdrant_client import models from bbconf.config_parser.bedbaseconfig import BedBaseConfig from bbconf.const import DEFAULT_LICENSE, PKG_NAME, ZARR_TOKENIZED_FOLDER @@ -59,6 +62,7 @@ TokenizedBedResponse, TokenizedPathResponse, UniverseMetadata, + VectorMetadata, ) _LOGGER = getLogger(PKG_NAME) @@ -85,6 +89,9 @@ def __init__(self, config: BedBaseConfig, bbagent_obj=None): self._config = config self.bb_agent = bbagent_obj + self._embedding_model = TextEmbedding(config.config.path.text2vec) + # self._embedding_model = TextEmbedding("BAAI/bge-large-en-v1.5") + def get(self, identifier: str, full: bool = False) -> BedMetadataAll: """ Get file metadata by identifier. @@ -469,11 +476,23 @@ def get_reference_validation(self, identifier: str) -> RefGenValidReturnModel: result_list = [] + results = [item for item in results] + + if not results: + return RefGenValidReturnModel( + id=identifier, + provided_genome=None, + compared_genome=[], + ) + for result in results: result_list.append( RefGenValidModel( provided_genome=result.provided_genome, - compared_genome=result.compared_genome, + compared_genome=( + result.genome_object.alias if result.genome_object else None + ), + genome_digest=result.genome_digest, xs=result.xs, oobr=result.oobr, sequence_fit=result.sequence_fit, @@ -527,6 +546,7 @@ def add( :param processed: true if bedfile was processed and statistics and plots were calculated :return: None """ + _LOGGER.info(f"Adding bed file to database. bed_id: {identifier}") if self.exists(identifier): @@ -662,16 +682,25 @@ def add( session.add(new_metadata) if ref_validation: - for ref_gen_check, data in ref_validation.items(): - new_gen_ref = GenomeRefStats( - **RefGenValidModel( - **data.model_dump(), - provided_genome=classification.genome_alias, - compared_genome=ref_gen_check, - ).model_dump(), - bed_id=identifier, - ) - session.add(new_gen_ref) + + new_gen_refs = self._create_ref_validation_models( + ref_validation=ref_validation, + bed_id=identifier, + provided_genome=classification.genome_alias, + ) + session.add_all(new_gen_refs) + + # for ref_gen_check, data in ref_validation.items(): + # new_gen_ref = GenomeRefStats( + # **RefGenValidModel( + # **data.model_dump(), + # provided_genome=classification.genome_alias, + # compared_genome=ref_gen_check, + # genome_digest=ref_gen_check, + # ).model_dump(), + # bed_id=identifier, + # ) + # session.add(new_gen_ref) session.commit() return None @@ -686,13 +715,13 @@ def update( classification: Union[dict, None] = None, ref_validation: Union[Dict[str, BaseModel], None] = None, license_id: str = DEFAULT_LICENSE, - upload_qdrant: bool = True, - upload_pephub: bool = True, + upload_qdrant: bool = False, + upload_pephub: bool = False, upload_s3: bool = True, local_path: str = None, overwrite: bool = False, nofail: bool = False, - processed: bool = True, + processed: bool = False, ) -> None: """ Update bed file to the database. @@ -732,7 +761,7 @@ def update( bed_metadata = StandardMeta(**metadata if metadata else {}) classification = BedClassification(**classification if classification else {}) - if upload_pephub: + if upload_pephub and metadata: metadata = BedPEPHub(**metadata) try: self.update_pephub(identifier, metadata.model_dump(), overwrite) @@ -789,12 +818,17 @@ def update( ) self._update_ref_validation( - sa_session=session, bed_object=bed_object, ref_validation=ref_validation + sa_session=session, + bed_id=identifier, + ref_validation=ref_validation, + provided_genome=bed_object.genome_alias or "", ) bed_object.processed = processed bed_object.indexed = upload_qdrant bed_object.last_update_date = datetime.datetime.now(datetime.timezone.utc) + if bed_metadata.description: + bed_object.description = bed_metadata.description session.commit() @@ -979,9 +1013,12 @@ def _update_files( f"File with name: {v.name} already exists. Updating.." ) - @staticmethod def _update_ref_validation( - sa_session: Session, bed_object: Bed, ref_validation: Dict[str, BaseModel] + self, + sa_session: Session, + bed_id: str, + ref_validation: Dict[str, BaseModel], + provided_genome: str = "", ) -> None: """ Update reference validation data @@ -989,34 +1026,57 @@ def _update_ref_validation( ! This function won't update the reference validation data, if it exists, it will skip it. :param sa_session: sqlalchemy session - :param bed_object: bed sqlalchemy object + :param bed_id: bed sqlalchemy object :param ref_validation: bed file metadata + :param provided_genome: genome reference that was provided by user """ if not ref_validation: return None - _LOGGER.info("Updating reference validation data..") + sa_session.execute( + delete(GenomeRefStats).where(GenomeRefStats.bed_id == bed_id) + ) + + new_gen_refs = self._create_ref_validation_models( + ref_validation=ref_validation, + bed_id=bed_id, + provided_genome=provided_genome, + ) + sa_session.add_all(new_gen_refs) + + # One commit for both deletes and inserts + sa_session.commit() - for ref_gen_check, data in ref_validation.items(): - new_gen_ref = GenomeRefStats( + return None + + def _create_ref_validation_models( + self, + ref_validation: Dict[str, BaseModel], + bed_id: str, + provided_genome: str = None, + ) -> list[GenomeRefStats]: + + compatibility = {} + + for k, v in ref_validation.items(): + if v.tier_ranking < 4: + compatibility[k] = v + + # Add all new GenomeRefStats objects in one go + new_gen_refs: list[GenomeRefStats] = [ + GenomeRefStats( **RefGenValidModel( **data.model_dump(), - provided_genome=bed_object.genome_alias, + provided_genome=provided_genome, compared_genome=ref_gen_check, + genome_digest=ref_gen_check, ).model_dump(), - bed_id=bed_object.id, + bed_id=bed_id, ) - try: - sa_session.add(new_gen_ref) - sa_session.commit() - except IntegrityError as _: - sa_session.rollback() - _LOGGER.info( - f"Reference validation exists for BED id: {bed_object.id} and ref_gen_check." - ) - - return None + for ref_gen_check, data in compatibility.items() + ] + return new_gen_refs def delete(self, identifier: str) -> None: """ @@ -1125,9 +1185,9 @@ def upload_file_qdrant( def _embed_file(self, bed_file: Union[str, RegionSet]) -> np.ndarray: """ - Create embeding for bed file + Create embedding for bed file - :param bed_file: bed file path or regionset + :param bed_file: bed file path or region set :param bed_file: path to the bed file, or RegionSet object :return np array of embeddings @@ -1156,14 +1216,20 @@ def _embed_file(self, bed_file: Union[str, RegionSet]) -> np.ndarray: return bed_embedding.reshape(1, vec_dim) def text_to_bed_search( - self, query: str, limit: int = 10, offset: int = 0 + self, + query: str, + limit: int = 10, + offset: int = 0, + with_metadata: bool = True, ) -> BedListSearchResult: """ Search for bed files by text query in qdrant database + This is bivec_search :param query: text query :param limit: number of results to return :param offset: offset to start from + :param with_metadata: if True, will return metadata for each result :return: list of bed file metadata """ @@ -1174,16 +1240,29 @@ def text_to_bed_search( for result in results: result_id = result["id"].replace("-", "") try: - result_meta = self.get(result_id) + if with_metadata: + result_meta = self.get(result_id) + else: + result_meta = None except BEDFileNotFoundError as e: _LOGGER.warning( f"Could not retrieve metadata for bed file: {result_id}. Error: {e}" ) continue - if result_meta: - results_list.append(QdrantSearchResult(**result, metadata=result_meta)) + results_list.append( + QdrantSearchResult( + **result, metadata=result_meta if with_metadata else None + ) + ) + + if with_metadata: + count = self._config._qdrant_advanced_engine.get_collection( + collection_name=self._config.config.qdrant.file_collection + ).points_count + else: + count = 0 return BedListSearchResult( - count=self.bb_agent.get_detailed_stats().file_genome.get("hg38", 0), + count=count, limit=limit, offset=offset, results=results_list, @@ -1195,6 +1274,15 @@ def bed_to_bed_search( limit: int = 10, offset: int = 0, ) -> BedListSearchResult: + """ + Search for bed files by using region set in qdrant database. + + :param region_set: RegionSet object to search for (bed file) + :param limit: number of results to return + :param offset: offset to start from + + :return: BedListSetResults + """ results = self._config.b2bsi.query_search( region_set, limit=limit, offset=offset ) @@ -1221,6 +1309,7 @@ def sql_search( self, query: str, genome: str = None, + assay: str = None, limit: int = 10, offset: int = 0, ) -> BedListSearchResult: @@ -1230,11 +1319,13 @@ def sql_search( :param query: text query :param genome: genome alias to filter results + :param assay: filter by assay type :param limit: number of results to return :param offset: offset to start from :return: list of bed file metadata """ + _LOGGER.debug(f"Looking for: {query}") statement = select(Bed).join(BedMetadata) @@ -1249,13 +1340,19 @@ def sql_search( BedMetadata.cell_type.ilike(sql_search_str), ) + condition_statement = or_statement + if genome_alias := genome: _LOGGER.debug(f"Filtering by genome: {genome_alias}") - condition_statement = and_(Bed.genome_alias == genome_alias, or_statement) + condition_statement = and_( + Bed.genome_alias == genome_alias, condition_statement + ) - else: - condition_statement = or_statement + if assay: + _LOGGER.debug(f"Filtering by assay: {assay}") + + condition_statement = and_(BedMetadata.assay == assay, condition_statement) statement = statement.where(condition_statement) @@ -1306,7 +1403,7 @@ def _sql_search_count(self, condition_statement) -> int: count = session.execute(statement).one() return count[0] - def reindex_qdrant(self, batch: int = 100) -> None: + def reindex_qdrant(self, batch: int = 100, purge: bool = False) -> None: """ Re-upload all files to quadrant. !Warning: only hg38 genome can be added to qdrant! @@ -1319,57 +1416,96 @@ def reindex_qdrant(self, batch: int = 100) -> None: """ bb_client = BBClient() - annotation_result = self.get_ids_list( - limit=100000, genome=QDRANT_GENOME, offset=0 - ) + with Session(self._sa_engine) as session: + if purge: + _LOGGER.info("Purging indexed files in the database ...") + session.query(Bed).update({Bed.file_indexed: False}) + session.commit() + _LOGGER.info("Purged indexed files in the database successfully!") - if not annotation_result.results: - _LOGGER.error("No bed files found.") - return None - results = annotation_result.results + statement = ( + select(Bed) + .join(BedMetadata, Bed.id == BedMetadata.id) + .where( + and_( + Bed.file_indexed == False, + Bed.genome_alias == QDRANT_GENOME, + # BedMetadata.global_experiment_id.contains(['encode']) # If we want only encode data + ) + ) + .limit(150000) + ) - with tqdm(total=len(results), position=0, leave=True) as pbar: - points_list = [] - processed_number = 0 - for record in results: - try: - bed_region_set_obj = GRegionSet(bb_client.seek(record.id)) - except FileNotFoundError: - bed_region_set_obj = bb_client.load_bed(record.id) + annotation_results = session.scalars(statement) + + results: List[Bed] = [result for result in annotation_results] + if not results: + _LOGGER.info("No files to reindex in qdrant.") + return None + + with tqdm(total=len(results), position=0, leave=True) as pbar: + points_list = [] + processed_number = 0 + for record in results: + try: + bed_region_set_obj = GRegionSet(bb_client.seek(record.id)) + except FileNotFoundError: + bed_region_set_obj = bb_client.load_bed(record.id) - pbar.set_description(f"Processing file: {record.id}") + pbar.set_description(f"Processing file: {record.id}") - file_embedding = self._embed_file(bed_region_set_obj) - points_list.append( - PointStruct( + file_embedding = self._embed_file(bed_region_set_obj) + + bed_metadata = VectorMetadata( id=record.id, - vector=file_embedding.tolist()[0], - payload=( - record.annotation.model_dump() if record.annotation else {} - ), + name=record.name, + description=record.description, + genome_alias=record.genome_alias, + genome_digest=record.genome_digest, + cell_line=record.annotations.cell_line, + cell_type=record.annotations.cell_type, + tissue=record.annotations.tissue, + target=record.annotations.target, + treatment=record.annotations.treatment, + assay=record.annotations.assay, + species_name=record.annotations.species_name, ) - ) - processed_number += 1 - if processed_number % batch == 0: - pbar.set_description(f"Uploading points to qdrant using batch...") - operation_info = self._config.qdrant_engine.qd_client.upsert( - collection_name=self._config.config.qdrant.file_collection, - points=points_list, + + record.file_indexed = True + + points_list.append( + PointStruct( + id=record.id, + vector=file_embedding.tolist()[0], + payload=(bed_metadata.model_dump()), + ) ) - pbar.write("Uploaded batch to qdrant.") - points_list = [] - assert operation_info.status == "completed" + processed_number += 1 - pbar.write(f"File: {record.id} successfully indexed.") - pbar.update(1) + if processed_number % batch == 0: + pbar.set_description( + "Uploading points to qdrant using batch..." + ) + operation_info = self._config.qdrant_engine.qd_client.upsert( + collection_name=self._config.config.qdrant.file_collection, + points=points_list, + ) + pbar.write("Uploaded batch to qdrant.") + points_list = [] + assert operation_info.status == "completed" - _LOGGER.info(f"Uploading points to qdrant using batches...") - operation_info = self._config.qdrant_engine.qd_client.upsert( - collection_name=self._config.config.qdrant.file_collection, - points=points_list, - ) - assert operation_info.status == "completed" - return None + session.commit() + + pbar.write(f"File: {record.id} successfully indexed.") + pbar.update(1) + + _LOGGER.info("Uploading points to qdrant using batches...") + operation_info = self._config.qdrant_engine.qd_client.upsert( + collection_name=self._config.config.qdrant.file_collection, + points=points_list, + ) + assert operation_info.status == "completed" + return None def delete_qdrant_point(self, identifier: str) -> None: """ @@ -1554,6 +1690,7 @@ def get_tokenized(self, bed_id: str, universe_id: str) -> TokenizedBedResponse: :return: zarr path """ + if not self.exist_tokenized(bed_id, universe_id): raise TokenizeFileNotExistError("Tokenized file not found in the database.") univers_group = self._config.zarr_root.require_group(universe_id) @@ -1853,3 +1990,257 @@ def _update_sources( flag_modified(bedmetadata_object, "global_experiment_id") session.commit() + + def reindex_semantic_search(self, batch: int = 1000, purge: bool = False) -> None: + """ + Reindex all bed files for semantic database + + :param batch: number of files to upload in one batch + :param purge: resets indexed in database for all files to False + + :return: None + """ + + # Add column that will indicate if this file is indexed or not + statement = ( + select(Bed) + .join(BedMetadata, Bed.id == BedMetadata.id) + .where(Bed.indexed == False) + .limit(150000) + ) + + with Session(self._sa_engine) as session: + + if purge: + _LOGGER.info("Purging indexed files in the database ...") + session.query(Bed).update({Bed.indexed: False}) + session.commit() + _LOGGER.info("Purged indexed files in the database successfully!") + + _LOGGER.info("Fetching data from the database ...") + results = session.scalars(statement) + + _LOGGER.info("Fetch data successfully!") + + points = [] + results = [result for result in results] + + with tqdm(total=len(results), position=0, leave=True) as pbar: + processed_number = 0 + for result in results: + text = ( + f"biosample is {result.annotations.cell_line} / {result.annotations.cell_type} / " + f"{result.annotations.tissue} with target {result.annotations.target} " + f"assay {result.annotations.assay}." + f"File name {result.name} with summary {result.description}" + ) + + embeddings_list = list(self._embedding_model.embed(text)) + # result_list.append( + data = VectorMetadata( + id=result.id, + name=result.name, + description=result.description, + genome_alias=result.genome_alias, + genome_digest=result.genome_digest, + cell_line=result.annotations.cell_line, + cell_type=result.annotations.cell_type, + tissue=result.annotations.tissue, + target=result.annotations.target, + treatment=result.annotations.treatment, + assay=result.annotations.assay, + species_name=result.annotations.species_name, + ) + + points.append( + PointStruct( + id=result.id, + vector=list(embeddings_list[0]), + payload=data.model_dump(), + ) + ) + processed_number += 1 + result.indexed = True + + if processed_number % batch == 0: + pbar.set_description( + "Uploading points to qdrant using batch..." + ) + operation_info = self._config._qdrant_advanced_engine.upsert( + collection_name=self._config.config.qdrant.search_collection, + points=points, + ) + session.commit() + pbar.write("Uploaded batch to qdrant.") + points = [] + assert operation_info.status == "completed" + + pbar.write(f"File: {result.id} successfully indexed.") + pbar.update(1) + + operation_info = self._config._qdrant_advanced_engine.upsert( + collection_name=self._config.config.qdrant.search_collection, + points=points, + ) + assert operation_info.status == "completed" + session.commit() + + return None + + def semantic_search( + self, + query: str = "liver", + genome_alias: str = "", + assay: str = "", + limit: int = 100, + offset: int = 0, + with_metadata: bool = True, + ) -> BedListSearchResult: + """ + Run semantic search for bed files using qdrant. + This is not bivec search, but usual qdrant search with embeddings. + + :param query: text query to search for + :param genome_alias: genome alias to filter results + :param assay: filter by assay type + :param limit: number of results to return + :param offset: offset to start from + :param with_metadata: if True, metadata will be returned in the results. Default is True. + + :return: list of bed file metadata + """ + + should_statement = [] + + if genome_alias: + should_statement.append( + models.FieldCondition( + key="genome_alias", + match=models.MatchValue(value=genome_alias), + ) + ) + if assay: + should_statement.append( + models.FieldCondition( + key="assay", + match=models.MatchValue(value=assay), + ) + ) + + embeddings_list = list(self._embedding_model.embed(query))[0] + + results = self._config._qdrant_advanced_engine.search( + collection_name=self._config.config.qdrant.search_collection, + query_vector=list(embeddings_list), + limit=limit, + offset=offset, + search_params=models.SearchParams( + exact=True, + ), + # query_filter=models.Filter(should=should_statement) if should_statement else None, + query_filter=( + models.Filter(must=should_statement) if should_statement else None + ), + with_payload=True, + with_vectors=True, + ) + + result_list = [] + for result in results: + result_id = result.id.replace("-", "") + + if with_metadata: + metadata = self.get(result_id, full=False) + else: + metadata = None + + result_list.append( + QdrantSearchResult( + id=result_id, + payload=result.payload, + score=result.score, + metadata=metadata, + ) + ) + + if with_metadata: + count = self.bb_agent.get_stats().bedfiles_number + else: + count = 0 + + return BedListSearchResult( + count=count, + limit=limit, + offset=offset, + results=result_list, + ) + + def search_external_file(self, source: str, accession: str) -> BedListSearchResult: + """ + Search for bed files by external source and accession number. + e.g. source='geo', accession='GSE12345' + + :param source: external source, e.g. 'geo' or 'encode' + :param accession: accession number, e.g. 'GSE12345' or 'ENCSR12345' + + :return: list of bed file metadata + """ + if source not in ["geo", "encode"]: + raise BedBaseConfError( + f"Source {source} is not supported. Supported sources are: 'geo', 'encode'." + ) + + if source == "geo" and accession.upper().startswith("GSE"): + statement = ( + select(Bed) + .join(BedMetadata, Bed.id == BedMetadata.id) + .where( + BedMetadata.global_experiment_id.contains( + cast( + [f"{source}:{accession}"], + postgresql.ARRAY(postgresql.VARCHAR), + ) + ) + ) + ) + else: + statement = ( + select(Bed) + .join(BedMetadata, Bed.id == BedMetadata.id) + .where( + BedMetadata.global_sample_id.contains( + cast( + [f"{source}:{accession}"], + postgresql.ARRAY(postgresql.VARCHAR), + ) + ) + ) + ) + + with Session(self._sa_engine) as session: + + bed_objects = session.scalars(statement) + results = [ + BedMetadataBasic( + **bedfile_obj.__dict__, + annotation=StandardMeta( + **( + bedfile_obj.annotations.__dict__ + if bedfile_obj.annotations + else {} + ) + ), + ) + for bedfile_obj in bed_objects + ] + result_list = [ + QdrantSearchResult(id=result.id, score=1, metadata=result) + for result in results + ] + + return BedListSearchResult( + count=len(result_list), + limit=len(result_list), + offset=0, + results=result_list, + ) diff --git a/bbconf/modules/bedsets.py b/bbconf/modules/bedsets.py index 74f22a7e..6179ee98 100644 --- a/bbconf/modules/bedsets.py +++ b/bbconf/modules/bedsets.py @@ -4,7 +4,7 @@ from geniml.io.utils import compute_md5sum_bedset from sqlalchemy import Float, Numeric, func, or_, select from sqlalchemy.exc import IntegrityError -from sqlalchemy.orm import Session, relationship +from sqlalchemy.orm import Session from bbconf.config_parser import BedBaseConfig from bbconf.const import PKG_NAME diff --git a/docs/changelog.md b/docs/changelog.md index 9df3a992..44bbe791 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -2,6 +2,20 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html) and [Keep a Changelog](https://keepachangelog.com/en/1.0.0/) format. +### [0.12.0] - 2025-09-11 +### Added: +- New qdrant semantic search +- Added more plots to bedbase summary page +- New reference genome compatibility that supports new refgenie +- Genomes table, with ability of automatic updates from refgenie +- Added filter in search (Assay and genome) + +### Changed: +- Improved reindexing methods + +### Fixed: +- Issues in bedfile update method + ### [0.11.4] - 2025-06-01 ### Fixed: - SQL search diff --git a/manual_testing.py b/manual_testing.py index 36884b28..2fc2205d 100644 --- a/manual_testing.py +++ b/manual_testing.py @@ -6,6 +6,9 @@ # from dotenv import load_dotenv from geniml.io import RegionSet from gtars.utils import read_tokens_from_gtok +import matplotlib.pyplot as plt +import numpy as np +import time # from gtars.tokenizers import RegionSet @@ -213,9 +216,94 @@ def compreh_stats(): from bbconf import BedBaseAgent agent = BedBaseAgent(config="/home/bnt4me/virginia/repos/bedhost/config.yaml") - results = agent.get_detailed_stats() - results = agent.get_detailed_usage() - print(results) + + time1 = time.time() + + results = agent.get_detailed_stats(concise=True) + + # results = agent.get_detailed_usage() + + time2 = time.time() + print(time2 - time1) + + # results = agent.get_detailed_stats() + # results = agent.get_detailed_usage() + + def plot_file_sizes(file_size_counts, file_size_bin_edges): + + # Plot the histogram + plt.figure(figsize=(10, 6)) + plt.bar( + file_size_bin_edges[:-1], + file_size_counts, + width=np.diff(file_size_bin_edges), + edgecolor="black", + align="edge", + ) + + # Add labels and title + plt.xlabel("File Size Bin Edges (bytes)") + plt.ylabel("Counts") + plt.title("Histogram of File Sizes") + plt.grid(axis="y", linestyle="--", alpha=0.7) + + # Show the plot + plt.show() + + def plot_region_width(mean_reg_width_counts, mean_reg_width_bin_edges): + + # Plot the histogram + plt.figure(figsize=(10, 6)) + plt.bar( + mean_reg_width_bin_edges[:-1], + mean_reg_width_counts, + width=np.diff(mean_reg_width_bin_edges), + edgecolor="black", + align="edge", + ) + + # Add labels and title + plt.xlabel("Mean Region Width Bin Edges") + plt.ylabel("Counts") + plt.title("Histogram of Mean Region Widths") + plt.grid(axis="y", linestyle="--", alpha=0.7) + + # Show the plot + plt.show() + + def plot_number_of_regions(n_region_counts, n_region_bin_edges): + + # Plot the histogram + plt.figure(figsize=(10, 6)) + plt.bar( + n_region_bin_edges[:-1], + n_region_counts, + width=np.diff(n_region_bin_edges), + edgecolor="black", + align="edge", + ) + + # Add labels and title + plt.xlabel("Region Bin Edges") + plt.ylabel("Counts") + plt.title("Histogram of Number of Regions") + plt.grid(axis="y", linestyle="--", alpha=0.7) + + # Show the plot + plt.show() + + plot_file_sizes( + results.file_size.counts, + results.file_size.bins, + ) + plot_region_width( + results.mean_region_width.counts, + results.mean_region_width.bins, + ) + plot_number_of_regions( + results.number_of_regions.counts, + results.number_of_regions.bins, + ) def get_unprocessed_files(): @@ -235,6 +323,43 @@ def get_genomes(): print(results) +def new_search(): + from bbconf import BedBaseAgent + + agent = BedBaseAgent(config="/home/bnt4me/virginia/repos/bedhost/config.yaml") + time1 = time.time() + + results = agent.bed.reindex_semantic_search() + # results = agent.bed.comp_search() + time2 = time.time() + + print(f"Time taken: {time2 - time1} seconds") + + +def get_assay_list(): + from bbconf import BedBaseAgent + + agent = BedBaseAgent(config="/home/bnt4me/virginia/repos/bedhost/config.yaml") + results = agent.get_list_assays() + print(results) + + +def external_search(): + from bbconf import BedBaseAgent + + agent = BedBaseAgent(config="/home/bnt4me/virginia/repos/bedhost/config.yaml") + + result = agent.bed.search_external_file("geo", "gsm1399546") + result + + +def reindex_files(): + from bbconf import BedBaseAgent + + agent = BedBaseAgent(config="/home/bnt4me/virginia/repos/bedhost/config.yaml") + agent.bed.reindex_qdrant(purge=True, batch=10) + + if __name__ == "__main__": # zarr_s3() # add_s3() @@ -243,8 +368,14 @@ def get_genomes(): # get_pep() # get_id_plots_missing() # neighbour_beds() - sql_search() + # sql_search() # config_t() - # compreh_stats() + compreh_stats() # get_unprocessed_files() # get_genomes() + # new_search() + + # external_search() + # get_assay_list() + + # reindex_files() diff --git a/requirements/requirements-all.txt b/requirements/requirements-all.txt index 91e2bada..ce2fca39 100644 --- a/requirements/requirements-all.txt +++ b/requirements/requirements-all.txt @@ -1,5 +1,6 @@ yacman >= 0.9.1 sqlalchemy >= 2.0.0 +gtars >= 0.4.0 geniml[ml] >= 0.7.1 psycopg >= 3.1.15 coloredlogs diff --git a/tests/test_bedfile.py b/tests/test_bedfile.py index 974e6795..c0ed1d1d 100644 --- a/tests/test_bedfile.py +++ b/tests/test_bedfile.py @@ -58,7 +58,9 @@ def test_get_all(self, bbagent_obj, mocked_phc): assert return_result is not None assert return_result.files is not None assert return_result.plots is not None - assert return_result.raw_metadata is not None + + # TODO: PEPhub is disabled + # assert return_result.raw_metadata is not None assert return_result.genome_alias == "hg38" assert return_result.stats.number_of_regions == 1 @@ -80,6 +82,9 @@ def test_get_all_not_found(self, bbagent_obj): assert return_result.genome_alias == "hg38" assert return_result.id == BED_TEST_ID + @pytest.mark.skip( + "Skipped, because PHC is disabled" + ) # TODO: should we disable PHC everywhere? def test_get_raw_metadata(self, bbagent_obj, mocked_phc): with ContextManagerDBTesting(config=bbagent_obj.config, add_data=True): return_result = bbagent_obj.bed.get_raw_metadata(BED_TEST_ID)