Merge pull request #93 from databio/umap

Releaase 0.13.0

Author: khoroshevskyi

bbconf/_version.py

@@ -1 +1 @@
-__version__ = "0.12.0"
+__version__ = "0.13.0"

bbconf/config_parser/bedbaseconfig.py

@@ -3,6 +3,7 @@
 import warnings
 from pathlib import Path
 from typing import List, Literal, Union
+import numpy as np
 
 import boto3
 import qdrant_client
@@ -18,6 +19,10 @@
 from geniml.search.query2vec import BED2Vec
 from pephubclient import PEPHubClient
 from zarr import Group as Z_GROUP
+from umap import UMAP
+
+import joblib
+import io
 
 from bbconf.config_parser.const import (
     S3_BEDSET_PATH_FOLDER,
@@ -67,6 +72,7 @@ def __init__(self, config: Union[Path, str], init_ml: bool = True):
             self._b2bsi = self._init_b2bsi_object()
             self._r2v = self._init_r2v_object()
             self._bivec = self._init_bivec_object()
+            self._umap_model: Union[UMAP, None] = self._init_umap_model()
         else:
             _LOGGER.info(
                 "Skipping initialization of ML models, init_ml parameter set to False."
@@ -75,6 +81,7 @@ def __init__(self, config: Union[Path, str], init_ml: bool = True):
             self._b2bsi = None
             self._r2v = None
             self._bivec = None
+            self._umap_model: Union[UMAP, None] = None
 
         self._phc = self._init_pephubclient()
         self._boto3_client = self._init_boto3_client()
@@ -414,6 +421,47 @@ def _init_r2v_object(self) -> Union[Region2VecExModel, None]:
             )
             return None
 
+    def _init_umap_model(self) -> Union[UMAP, None]:
+        """
+        Load UMAP model from the specified path, or url
+        """
+
+        if not self.config.path.umap_model:
+            _LOGGER.warning(
+                "UMAP model path is not specified in the configuration, and won't be used."
+            )
+            return None
+
+        model_path = self.config.path.umap_model
+
+        if model_path.startswith(("http://", "https://")):
+            import requests
+
+            try:
+                response = requests.get(model_path)
+                response.raise_for_status()
+                buffer = io.BytesIO(response.content)
+                umap_model = joblib.load(buffer)
+                print(f"UMAP model loaded from URL: {model_path}")
+            except requests.RequestException as e:
+                _LOGGER.error(f"Error downloading UMAP model from URL: {e}")
+                return None
+        else:
+            try:
+                with open(model_path, "rb") as file:
+                    umap_model = joblib.load(file)
+                print(f"UMAP model loaded from local path: {model_path}")
+            except FileNotFoundError as e:
+                _LOGGER.error(f"Error loading UMAP model from local path: {e}")
+                return None
+
+        if not isinstance(umap_model, UMAP):
+            _LOGGER.error(f"Loaded object is not a UMAP instance: {type(umap_model)}")
+            return None
+        # np.random.seed(42)
+        umap_model.random_state = 42
+        return umap_model
+
     def upload_s3(self, file_path: str, s3_path: Union[Path, str]) -> None:
         """
         Upload file to s3.

bbconf/config_parser/models.py

@@ -67,6 +67,7 @@ class ConfigPath(BaseModel):
     region2vec: str = DEFAULT_REGION2_VEC_MODEL
     # vec2vec: str = DEFAULT_VEC2VEC_MODEL
     text2vec: str = DEFAULT_TEXT2VEC_MODEL
+    umap_model: Union[str, None] = None  # Path or link to pre-trained UMAP model
 
 
 class AccessMethodsStruct(BaseModel):

bbconf/modules/bedfiles.py

@@ -5,7 +5,8 @@
 
 import numpy as np
 from geniml.bbclient import BBClient
-from geniml.io import RegionSet
+
+# from geniml.io import RegionSet
 from geniml.search.backends import QdrantBackend
 from gtars.models import RegionSet as GRegionSet
 from pephubclient.exceptions import ResponseError
@@ -1155,7 +1156,7 @@ def delete_pephub_sample(self, identifier: str):
     def upload_file_qdrant(
         self,
         bed_id: str,
-        bed_file: Union[str, RegionSet],
+        bed_file: Union[str, GRegionSet],
         payload: dict = None,
     ) -> None:
         """
@@ -1183,7 +1184,7 @@ def upload_file_qdrant(
         )
         return None
 
-    def _embed_file(self, bed_file: Union[str, RegionSet]) -> np.ndarray:
+    def _embed_file(self, bed_file: Union[str, GRegionSet]) -> np.ndarray:
         """
         Create embedding for bed file
 
@@ -1204,8 +1205,8 @@ def _embed_file(self, bed_file: Union[str, RegionSet]) -> np.ndarray:
             try:
                 bed_region_set = GRegionSet(bed_file)
             except RuntimeError as _:
-                bed_region_set = RegionSet(bed_file)
-        elif isinstance(bed_file, RegionSet) or isinstance(bed_file, GRegionSet):
+                bed_region_set = GRegionSet(bed_file)
+        elif isinstance(bed_file, GRegionSet) or isinstance(bed_file, GRegionSet):
             bed_region_set = bed_file
         else:
             raise BedBaseConfError(
@@ -1215,6 +1216,20 @@ def _embed_file(self, bed_file: Union[str, RegionSet]) -> np.ndarray:
         vec_dim = bed_embedding.shape[0]
         return bed_embedding.reshape(1, vec_dim)
 
+    def _get_umap_file(self, bed_file: Union[str, GRegionSet]) -> np.ndarray:
+        """
+        Create UMAP for bed file
+
+        :param bed_file: bed file path or region set
+        """
+
+        if self._config._umap_model is None:
+            raise BedBaseConfError("UMAP model is not initialized.")
+
+        bed_embedding = self._embed_file(bed_file)
+        bed_umap = self._config._umap_model.transform(bed_embedding)
+        return bed_umap
+
     def text_to_bed_search(
         self,
         query: str,
@@ -1270,7 +1285,7 @@ def text_to_bed_search(
 
     def bed_to_bed_search(
         self,
-        region_set: RegionSet,
+        region_set: GRegionSet,
         limit: int = 10,
         offset: int = 0,
     ) -> BedListSearchResult:

docs/changelog.md

@@ -2,6 +2,10 @@
 
 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.13.0] - 2025-11-24
+### Added:
+- Conversion of bedfile to umap from predefined model
+
 ### [0.12.0] - 2025-09-11
 ### Added:
 - New qdrant semantic search

requirements/requirements-all.txt

@@ -14,3 +14,4 @@ pyyaml >= 6.0.1 # for s3fs because of the errors
 s3fs >= 2024.3.1
 pandas >= 2.0.0
 pybiocfilecache == 0.6.1
+umap-learn >= 0.5.8