pyobo 0.12.10__py3-none-any.whl → 0.12.12__py3-none-any.whl

pyobo/__init__.py

@@ -14,6 +14,7 @@ from .api import (
     get_filtered_relations_df,
     get_filtered_xrefs,
     get_graph,
+    get_graph_embeddings_df,
     get_hierarchy,
     get_id_definition_mapping,
     get_id_multirelations_mapping,
@@ -37,6 +38,7 @@ from .api import (
     get_obsolete,
     get_primary_curie,
     get_primary_identifier,
+    get_primary_reference,
     get_properties,
     get_properties_df,
     get_property,
@@ -44,6 +46,7 @@ from .api import (
     get_relation,
     get_relation_mapping,
     get_relations_df,
+    get_semantic_mappings,
     get_species,
     get_sssom_df,
     get_subhierarchy,
@@ -114,6 +117,7 @@ __all__ = [
     "get_filtered_relations_df",
     "get_filtered_xrefs",
     "get_graph",
+    "get_graph_embeddings_df",
     "get_grounder",
     "get_hierarchy",
     "get_id_definition_mapping",
@@ -139,6 +143,7 @@ __all__ = [
     "get_ontology",
     "get_primary_curie",
     "get_primary_identifier",
+    "get_primary_reference",
     "get_properties",
     "get_properties_df",
     "get_property",
@@ -149,6 +154,7 @@ __all__ = [
     "get_scispacy_entities",
     "get_scispacy_entity_linker",
     "get_scispacy_knowledgebase",
+    "get_semantic_mappings",
     "get_species",
     "get_sssom_df",
     "get_subhierarchy",

pyobo/api/__init__.py

@@ -5,10 +5,16 @@ from .alts import (
     get_id_to_alts,
     get_primary_curie,
     get_primary_identifier,
+    get_primary_reference,
 )
 from .combine import get_literal_mappings_subset
 from .edges import get_edges, get_edges_df, get_graph
-from .embedding import get_text_embedding, get_text_embedding_similarity, get_text_embeddings_df
+from .embedding import (
+    get_graph_embeddings_df,
+    get_text_embedding,
+    get_text_embedding_similarity,
+    get_text_embeddings_df,
+)
 from .hierarchy import (
     get_ancestors,
     get_children,
@@ -59,6 +65,7 @@ from .typedefs import get_typedef_df
 from .xrefs import (
     get_filtered_xrefs,
     get_mappings_df,
+    get_semantic_mappings,
     get_sssom_df,
     get_xref,
     get_xrefs,
@@ -80,6 +87,7 @@ __all__ = [
     "get_filtered_relations_df",
     "get_filtered_xrefs",
     "get_graph",
+    "get_graph_embeddings_df",
     "get_hierarchy",
     "get_id_definition_mapping",
     "get_id_multirelations_mapping",
@@ -105,6 +113,7 @@ __all__ = [
     "get_ontology",
     "get_primary_curie",
     "get_primary_identifier",
+    "get_primary_reference",
     "get_priority_curie",
     "get_properties",
     "get_properties_df",
@@ -113,6 +122,7 @@ __all__ = [
     "get_relation",
     "get_relation_mapping",
     "get_relations_df",
+    "get_semantic_mappings",
     "get_species",
     "get_sssom_df",
     "get_subhierarchy",

pyobo/api/alts.py

@@ -20,6 +20,7 @@ __all__ = [
     "get_id_to_alts",
     "get_primary_curie",
     "get_primary_identifier",
+    "get_primary_reference",
 ]
 
 logger = logging.getLogger(__name__)
@@ -61,13 +62,13 @@ def get_alts_to_id(prefix: str, **kwargs: Unpack[GetOntologyKwargs]) -> Mapping[
     }
 
 
-def get_primary_curie(
+def get_primary_reference(
     prefix: str | curies.Reference | curies.ReferenceTuple,
     identifier: str | None = None,
     /,
     **kwargs: Unpack[GetOntologyKwargs],
-) -> str | None:
-    """Get the primary curie for an entity."""
+) -> curies.ReferenceTuple | None:
+    """Get the primary reference for an entity."""
     reference = _get_pi(prefix, identifier)
     try:
         primary_identifier = get_primary_identifier(reference, **kwargs)
@@ -76,7 +77,20 @@ def get_primary_curie(
             raise
         # this happens on invalid prefix. maybe revise?
         return None
-    return f"{reference.prefix}:{primary_identifier}"
+    return curies.ReferenceTuple(reference.prefix, primary_identifier)
+
+
+def get_primary_curie(
+    prefix: str | curies.Reference | curies.ReferenceTuple,
+    identifier: str | None = None,
+    /,
+    **kwargs: Unpack[GetOntologyKwargs],
+) -> str | None:
+    """Get the primary curie for an entity."""
+    reference = get_primary_reference(prefix, identifier, **kwargs)
+    if reference is None:
+        return None
+    return reference.curie
 
 
 def get_primary_identifier(

pyobo/api/embedding.py

@@ -2,18 +2,29 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+import tempfile
+from pathlib import Path
+from typing import TYPE_CHECKING, Literal
 
+import bioregistry
 import curies
 import numpy as np
 import pandas as pd
+from tqdm import tqdm
+from typing_extensions import Unpack
 
-from pyobo.api.names import get_definition, get_name, get_references
+from pyobo.api.edges import get_edges_df
+from pyobo.api.names import get_definition, get_id_name_mapping, get_name
+from pyobo.api.utils import get_version_from_kwargs
+from pyobo.constants import GetOntologyKwargs, check_should_force
+from pyobo.identifier_utils import wrap_norm_prefix
+from pyobo.utils.path import CacheArtifact, get_cache_path
 
 if TYPE_CHECKING:
     import sentence_transformers
 
 __all__ = [
+    "get_graph_embeddings_df",
     "get_text_embedding",
     "get_text_embedding_model",
     "get_text_embedding_similarity",
@@ -31,38 +42,126 @@ def get_text_embedding_model() -> sentence_transformers.SentenceTransformer:
 
 def _get_text(
     reference: str | curies.Reference | curies.ReferenceTuple,
+    /,
+    *,
+    name: str | None = None,
+    **kwargs: Unpack[GetOntologyKwargs],
 ) -> str | None:
-    name = get_name(reference)
+    if name is None:
+        name = get_name(reference, **kwargs)
     if name is None:
         return None
-    description = get_definition(reference)
+    description = get_definition(reference, **kwargs)
     if description:
         name += " " + description
     return name
 
 
+def get_graph_embeddings_df(
+    prefix: str,
+    *,
+    method: Literal["pykeen", "grape"] | None = None,
+    epochs: int = 30,
+    dimension: int = 32,
+    **kwargs: Unpack[GetOntologyKwargs],
+) -> pd.DataFrame:
+    """Get graph machine learning embeddings."""
+    if method == "pykeen" or method is None:
+        from pykeen.models import PairRE
+        from pykeen.training import SLCWATrainingLoop
+        from pykeen.triples import TriplesFactory
+        from torch.optim import Adam
+
+        triples_df = get_edges_df(prefix, **kwargs)
+        training = TriplesFactory.from_labeled_triples(triples_df.values)
+        model = PairRE(triples_factory=training, embedding_dim=dimension)
+        optimizer = Adam(params=model.get_grad_params())
+        training_loop = SLCWATrainingLoop(
+            model=model, triples_factory=training, optimizer=optimizer
+        )
+        # can also set batch size here
+        training_loop.train(triples_factory=training, num_epochs=epochs)
+        embeddings = model.entity_representations[0]()
+        df = pd.DataFrame(
+            embeddings.detach().numpy(),
+            index=[training.entity_id_to_label[i] for i in range(embeddings.shape[0])],
+        )
+
+    elif method == "grape":
+        from ensmallen import Graph
+
+        edges_df = get_edges_df(prefix, **kwargs)
+        with tempfile.TemporaryDirectory() as d:
+            path = Path(d).joinpath("test.tsv")
+            edges_df[[":START_ID", ":END_ID"]].to_csv(path, header=None, sep="\t", index=False)
+            graph = Graph.from_csv(
+                edge_path=str(path),
+                edge_list_separator="\t",
+                sources_column_number=0,
+                destinations_column_number=1,
+                edge_list_numeric_node_ids=False,
+                directed=True,
+                name=bioregistry.get_name(prefix, strict=True),
+                verbose=True,
+            )
+        graph = graph.remove_disconnected_nodes()
+
+        from embiggen.embedders.ensmallen_embedders.second_order_line import (
+            SecondOrderLINEEnsmallen,
+        )
+
+        embedding = SecondOrderLINEEnsmallen(embedding_size=dimension, epochs=epochs).fit_transform(
+            graph
+        )
+        df = embedding.get_all_node_embedding()[0].sort_index()
+        # df.columns = [str(c) for c in df.columns]
+    else:
+        raise ValueError(f"invalid graph machine learning method: {method}")
+
+    df.index.name = "curie"
+    return df
+
+
+@wrap_norm_prefix
 def get_text_embeddings_df(
     prefix: str,
     *,
     model: sentence_transformers.SentenceTransformer | None = None,
+    **kwargs: Unpack[GetOntologyKwargs],
 ) -> pd.DataFrame:
     """Get embeddings for all entities in the resource.
 
     :param prefix: A reference, either as a string or Reference object
     :param model: A sentence transformer model. Defaults to ``all-MiniLM-L6-v2`` if not
         given.
+    :param kwargs: The keyword arguments to forward to ontology getter functions for
+        names, definitions, and version
+
+    :returns: A pandas dataframe with an index representing local unique identifiers and
+        columns for the values of the model returned vectors
     """
+    path = get_cache_path(
+        prefix, CacheArtifact.embeddings, version=get_version_from_kwargs(prefix, kwargs)
+    )
+    if path.is_file() and not check_should_force(kwargs):
+        df = pd.read_csv(path, sep="\t").set_index(0)
+        return df
+
+    id_to_name = get_id_name_mapping(prefix, **kwargs)
+
     luids, texts = [], []
-    for reference in get_references(prefix):
-        text = _get_text(reference)
+    for identifier, name in tqdm(id_to_name.items(), desc=f"[{prefix}] constructing text"):
+        text = _get_text(curies.ReferenceTuple(prefix, identifier), name=name, **kwargs)
         if text is None:
             continue
-        luids.append(reference.identifier)
+        luids.append(identifier)
         texts.append(text)
     if model is None:
         model = get_text_embedding_model()
-    res = model.encode(texts)
-    return pd.DataFrame(res, index=luids)
+    res = model.encode(texts, show_progress_bar=True)
+    df = pd.DataFrame(res, index=luids)
+    df.to_csv(path, sep="\t")  # index is important here!
+    return df
 
 
 def get_text_embedding(

pyobo/api/names.py

@@ -6,7 +6,7 @@ import logging
 import subprocess
 from collections.abc import Callable, Mapping
 from functools import lru_cache
-from typing import Any, TypeVar
+from typing import TypeVar
 
 import curies
 import pandas as pd
@@ -49,9 +49,15 @@ __all__ = [
 logger = logging.getLogger(__name__)
 
 
-def get_name_by_curie(curie: str, **kwargs: Any) -> str | None:
+def get_name_by_curie(
+    curie: str,
+    /,
+    *,
+    upgrade_identifier: bool | None = None,
+    **kwargs: Unpack[GetOntologyKwargs],
+) -> str | None:
     """Get the name for a CURIE, if possible."""
-    return get_name(curie, **kwargs)
+    return get_name(curie, upgrade_identifier=upgrade_identifier, **kwargs)
 
 
 X = TypeVar("X")
@@ -63,6 +69,8 @@ NO_BUILD_LOGGED: set = set()
 def _help_get(
     f: Callable[[str, Unpack[GetOntologyKwargs]], Mapping[str, X]],
     reference: Reference,
+    *,
+    upgrade_identifier: bool | None = None,
     **kwargs: Unpack[GetOntologyKwargs],
 ) -> X | None:
     """Get the result for an entity based on a mapping maker function ``f``."""
@@ -87,19 +95,32 @@ def _help_get(
             NO_BUILD_PREFIXES.add(reference.prefix)
         return None
 
-    primary_id = get_primary_identifier(reference, **kwargs)
-    return mapping.get(primary_id)
+    if upgrade_identifier is None:
+        if reference.identifier in mapping:
+            return mapping[reference.identifier]
+        else:
+            primary_id = get_primary_identifier(reference, **kwargs)
+            return mapping.get(primary_id)
+    elif upgrade_identifier is True:
+        primary_id = get_primary_identifier(reference, **kwargs)
+        return mapping.get(primary_id)
+    else:
+        return mapping.get(reference.identifier)
 
 
 def get_name(
     prefix: str | curies.Reference | curies.ReferenceTuple,
     identifier: str | None = None,
     /,
+    *,
+    upgrade_identifier: bool | None = None,
     **kwargs: Unpack[GetOntologyKwargs],
 ) -> str | None:
     """Get the name for an entity."""
     reference = _get_pi(prefix, identifier)
-    return _help_get(get_id_name_mapping, reference, **kwargs)
+    return _help_get(
+        get_id_name_mapping, reference, upgrade_identifier=upgrade_identifier, **kwargs
+    )
 
 
 @lru_cache
@@ -325,6 +346,7 @@ def get_literal_mappings(
     return rv
 
 
+@wrap_norm_prefix
 def get_literal_mappings_df(
     prefix: str,
     **kwargs: Unpack[GetOntologyKwargs],

pyobo/api/xrefs.py

@@ -5,8 +5,11 @@ import warnings
 from collections.abc import Mapping
 from functools import lru_cache
 
+import curies
 import pandas as pd
 from curies import ReferenceTuple
+from sssom_pydantic import SemanticMapping
+from sssom_pydantic.io import row_to_semantic_mapping
 from typing_extensions import Unpack
 
 from .utils import get_version_from_kwargs
@@ -19,7 +22,7 @@ from ..constants import (
     check_should_use_tqdm,
 )
 from ..getters import get_ontology
-from ..identifier_utils import wrap_norm_prefix
+from ..identifier_utils import get_converter, wrap_norm_prefix
 from ..struct import Obo
 from ..utils.cache import cached_df
 from ..utils.path import CacheArtifact, get_cache_path
@@ -27,6 +30,7 @@ from ..utils.path import CacheArtifact, get_cache_path
 __all__ = [
     "get_filtered_xrefs",
     "get_mappings_df",
+    "get_semantic_mappings",
     "get_sssom_df",
     "get_xref",
     "get_xrefs",
@@ -107,6 +111,22 @@ def get_sssom_df(
     return get_mappings_df(prefix=prefix, names=names, **kwargs)
 
 
+def get_semantic_mappings(
+    prefix: str,
+    converter: curies.Converter | None = None,
+    names: bool = True,
+    include_mapping_source_column: bool = False,
+    **kwargs: Unpack[GetOntologyKwargs],
+) -> list[SemanticMapping]:
+    """Get semantic mapping objects."""
+    df = get_mappings_df(
+        prefix, names=names, include_mapping_source_column=include_mapping_source_column, **kwargs
+    )
+    if converter is None:
+        converter = get_converter()
+    return [row_to_semantic_mapping(row, converter=converter) for _, row in df.iterrows()]
+
+
 def get_mappings_df(
     prefix: str | Obo,
     *,

pyobo/cli/cli.py

@@ -6,10 +6,7 @@ from collections.abc import Iterable
 from functools import lru_cache
 from operator import itemgetter
 
-import bioregistry
 import click
-import humanize
-from tabulate import tabulate
 
 from .database import main as database_main
 from .lookup import lookup
@@ -59,6 +56,9 @@ def clean(remove_obo: bool):
 @main.command()
 def ls():
     """List how big all of the OBO files are."""
+    import humanize
+    from tabulate import tabulate
+
     entries = [(prefix, os.path.getsize(path)) for prefix, path in _iter_cached_obo()]
     entries = [
         (prefix, humanize.naturalsize(size), "✅" if not has_nomenclature_plugin(prefix) else "❌")
@@ -69,6 +69,8 @@ def ls():
 
 def _iter_cached_obo() -> Iterable[tuple[str, str]]:
     """Iterate over cached OBO paths."""
+    import bioregistry
+
     for prefix in os.listdir(RAW_DIRECTORY):
         if prefix in GLOBAL_SKIP or _has_no_download(prefix) or bioregistry.is_deprecated(prefix):
             continue
@@ -83,6 +85,8 @@ def _iter_cached_obo() -> Iterable[tuple[str, str]]:
 
 def _has_no_download(prefix: str) -> bool:
     """Return if the prefix is not available."""
+    import bioregistry
+
     prefix_norm = bioregistry.normalize_prefix(prefix)
     return prefix_norm is not None and prefix_norm in _no_download()
 
@@ -90,6 +94,8 @@ def _has_no_download(prefix: str) -> bool:
 @lru_cache(maxsize=1)
 def _no_download() -> set[str]:
     """Get the list of prefixes not available as OBO."""
+    import bioregistry
+
     return {resource.prefix for resource in bioregistry.resources() if not resource.has_download()}