Add COSMOS PKN client module (cosmos subpackage)

New subpackage omnipath_client.cosmos with:
- get_pkn(organism, categories, resources, format): fetches pre-built
  PKNs from metabo.omnipathdb.org as DataFrame, dict, or AnnNet Graph
- to_annnet(df): bulk conversion to AnnNet Graph using add_vertices_bulk
  and add_edges_bulk for efficient construction
- categories(), organisms(), resources(), status(): convenience functions

Features:
- Organism normalisation via utils.ensure_ncbi_tax_id — accepts any
  form: 9606, 'human', 'Homo sapiens', 'hsapiens', 'hsa', etc.
- Uses shared _download.Downloader (dlmachine + caching)
- DataFrame output via polars (preferred) or pandas
- AnnNet conversion stores entity types as vertex attributes,
  interaction metadata as edge attributes

Usage:
  import omnipath_client as oc
  df = oc.cosmos.get_pkn('human', categories='transporters')
  g = oc.cosmos.get_pkn('mouse', format='annnet')

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: deeenes

omnipath_client/__init__.py

@@ -31,7 +31,7 @@
     'values',
 ]
 
-from omnipath_client import utils
+from omnipath_client import cosmos, utils
 from ._client import (
     OmniPath,
     params,

omnipath_client/cosmos/__init__.py

@@ -0,0 +1,37 @@
+"""COSMOS PKN client for the OmniPath Metabolomics service.
+
+Fetches pre-built COSMOS prior-knowledge networks from
+``metabo.omnipathdb.org`` as DataFrames or AnnNet Graph objects.
+
+Example::
+
+    import omnipath_client as oc
+
+    # DataFrame (polars by default)
+    df = oc.cosmos.get_pkn(organism='human', categories='transporters')
+
+    # AnnNet Graph
+    g = oc.cosmos.get_pkn(organism='human', format='annnet')
+
+    # Available categories and organisms
+    oc.cosmos.categories()   # ['transporters', 'receptors', ...]
+    oc.cosmos.organisms()    # [9606, 10090, 10116]
+"""
+
+from omnipath_client.cosmos._pkn import (
+    categories,
+    get_pkn,
+    organisms,
+    resources,
+    status,
+)
+from omnipath_client.cosmos._annnet import to_annnet
+
+__all__ = [
+    'categories',
+    'get_pkn',
+    'organisms',
+    'resources',
+    'status',
+    'to_annnet',
+]

omnipath_client/cosmos/_annnet.py

@@ -0,0 +1,115 @@
+"""Convert COSMOS PKN DataFrames to AnnNet Graph objects.
+
+Requires the ``annnet`` package (optional dependency).
+Uses bulk operations for efficient graph construction.
+"""
+
+from __future__ import annotations
+
+__all__ = ['to_annnet']
+
+import logging
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def to_annnet(df: Any) -> Any:
+    """Convert a COSMOS PKN DataFrame to an AnnNet Graph.
+
+    Uses ``add_vertices_bulk`` and ``add_edges_bulk`` for efficient
+    construction.  Entity types (protein, small_molecule) are stored
+    as vertex attributes; interaction metadata as edge attributes.
+
+    Args:
+        df: COSMOS PKN DataFrame (polars or pandas) with columns:
+            ``source``, ``target``, ``source_type``, ``target_type``,
+            ``interaction_type``, ``resource``, ``mor``.
+
+    Returns:
+        ``annnet.Graph`` instance.
+
+    Raises:
+        ImportError: If ``annnet`` is not installed.
+
+    Example::
+
+        import omnipath_client as oc
+
+        df = oc.cosmos.get_pkn('human', categories='transporters')
+        g = oc.cosmos.to_annnet(df)
+    """
+
+    try:
+        from annnet import Graph
+    except ImportError as exc:
+        raise ImportError(
+            'annnet is required for graph conversion. '
+            'Install with: pip install annnet'
+        ) from exc
+
+    # Normalise to dict-of-lists for column access
+    # (works with both polars and pandas)
+    try:
+        # Polars
+        cols = {c: df[c].to_list() for c in df.columns}
+    except AttributeError:
+        # Pandas
+        cols = {c: df[c].tolist() for c in df.columns}
+
+    n_rows = len(cols.get('source', []))
+
+    # -- Collect unique entities with types -----------------------------------
+    entity_types: dict[str, str] = {}
+
+    for i in range(n_rows):
+        src = cols['source'][i]
+        tgt = cols['target'][i]
+        src_type = cols.get('source_type', [None] * n_rows)[i]
+        tgt_type = cols.get('target_type', [None] * n_rows)[i]
+
+        if src not in entity_types and src_type:
+            entity_types[src] = src_type
+        if tgt not in entity_types and tgt_type:
+            entity_types[tgt] = tgt_type
+
+    # -- Build graph ----------------------------------------------------------
+    g = Graph()
+
+    # Bulk add vertices
+    vertices = [
+        {'vertex_id': vid, 'entity_type': etype}
+        for vid, etype in entity_types.items()
+    ]
+    g.add_vertices_bulk(vertices)
+
+    logger.info(
+        'AnnNet: added %d vertices (%d protein, %d small_molecule)',
+        len(vertices),
+        sum(1 for v in entity_types.values() if v == 'protein'),
+        sum(1 for v in entity_types.values() if v == 'small_molecule'),
+    )
+
+    # Bulk add edges
+    edges = []
+    mor_col = cols.get('mor', [0] * n_rows)
+    itype_col = cols.get('interaction_type', [''] * n_rows)
+    resource_col = cols.get('resource', [''] * n_rows)
+
+    for i in range(n_rows):
+        edges.append({
+            'source': cols['source'][i],
+            'target': cols['target'][i],
+            'weight': float(mor_col[i]) if mor_col[i] is not None else 0.0,
+            'edge_directed': True,
+            'attributes': {
+                'interaction_type': itype_col[i],
+                'resource': resource_col[i],
+            },
+        })
+
+    g.add_edges_bulk(edges)
+
+    logger.info('AnnNet: added %d edges', len(edges))
+
+    return g

omnipath_client/cosmos/_pkn.py

@@ -0,0 +1,194 @@
+"""COSMOS PKN fetch and convenience functions."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from omnipath_client._download import Downloader
+from omnipath_client._session import get_logger
+
+logger = get_logger(__name__)
+
+DEFAULT_METABO_URL = 'https://metabo.omnipathdb.org'
+
+_metabo_url: str = DEFAULT_METABO_URL
+_downloader: Downloader | None = None
+
+
+def set_url(url: str) -> None:
+    """Override the metabo service base URL."""
+
+    global _metabo_url
+    _metabo_url = url
+
+
+def _dl() -> Downloader:
+    """Lazy-initialised shared downloader."""
+
+    global _downloader
+
+    if _downloader is None:
+        _downloader = Downloader(use_cache=False)
+
+    return _downloader
+
+
+def _resolve_organism(organism: int | str) -> int:
+    """Resolve organism to NCBI taxonomy ID via the utils service."""
+
+    if isinstance(organism, int):
+        return organism
+
+    # Try parsing as int first (e.g. '9606')
+    try:
+        return int(organism)
+    except (ValueError, TypeError):
+        pass
+
+    from omnipath_client.utils import ensure_ncbi_tax_id
+
+    result = ensure_ncbi_tax_id(organism)
+
+    if result is None:
+        raise ValueError(f'Could not resolve organism: {organism!r}')
+
+    return result
+
+
+def get_pkn(
+    organism: int | str = 9606,
+    categories: str | list[str] = 'all',
+    resources: str | list[str] | None = None,
+    format: str = 'dataframe',
+) -> Any:
+    """Fetch COSMOS PKN from the metabo service.
+
+    Args:
+        organism:
+            Organism identifier — any form accepted by omnipath-utils
+            taxonomy: NCBI ID (``9606``), common name (``'human'``),
+            Latin (``'Homo sapiens'``), Ensembl (``'hsapiens'``),
+            KEGG (``'hsa'``), etc.
+        categories:
+            Category names or ``'all'``.  Available categories:
+            ``transporters``, ``receptors``, ``allosteric``,
+            ``enzyme_metabolite``, ``ppi``, ``grn``.
+        resources:
+            Optional resource filter (comma-separated string or list).
+        format:
+            Output format: ``'dataframe'`` (default), ``'parquet'``,
+            ``'dict'``, or ``'annnet'``.
+
+    Returns:
+        DataFrame (polars/pandas), raw dict, bytes (parquet), or
+        AnnNet Graph depending on *format*.
+
+    Example::
+
+        import omnipath_client as oc
+
+        # All human transporters
+        df = oc.cosmos.get_pkn('human', categories='transporters')
+
+        # Full mouse PKN as AnnNet graph
+        g = oc.cosmos.get_pkn('mouse', format='annnet')
+    """
+
+    ncbi_tax_id = _resolve_organism(organism)
+
+    if isinstance(categories, list):
+        categories = ','.join(categories)
+
+    if isinstance(resources, list):
+        resources = ','.join(resources)
+
+    params: dict[str, Any] = {
+        'organism': ncbi_tax_id,
+        'categories': categories,
+    }
+
+    if resources:
+        params['resources'] = resources
+
+    if format == 'parquet':
+        params['format'] = 'parquet'
+
+    url = f'{_metabo_url}/cosmos/pkn'
+    data = _dl().fetch_json(url, params=params)
+
+    if format == 'dict':
+        return data
+
+    if format == 'annnet':
+        df = _network_to_dataframe(data['network'])
+        from omnipath_client.cosmos._annnet import to_annnet
+        return to_annnet(df)
+
+    # Default: DataFrame
+    return _network_to_dataframe(data['network'])
+
+
+def _network_to_dataframe(records: list[dict]) -> Any:
+    """Convert network records to a DataFrame.
+
+    Uses polars if available, falls back to pandas.
+    """
+
+    if not records:
+        try:
+            import polars as pl
+            return pl.DataFrame()
+        except ImportError:
+            import pandas as pd
+            return pd.DataFrame()
+
+    try:
+        import polars as pl
+        return pl.DataFrame(records)
+    except ImportError:
+        pass
+
+    try:
+        import pandas as pd
+        return pd.DataFrame(records)
+    except ImportError:
+        pass
+
+    raise ImportError(
+        'Either polars or pandas is required for DataFrame output. '
+        'Install with: pip install polars'
+    )
+
+
+def categories() -> list[str]:
+    """List available PKN categories."""
+
+    url = f'{_metabo_url}/cosmos/categories'
+    return _dl().fetch_json(url)
+
+
+def organisms() -> list[int]:
+    """List organisms with pre-built PKNs."""
+
+    url = f'{_metabo_url}/cosmos/organisms'
+    return _dl().fetch_json(url)
+
+
+def resources(organism: int | str = 9606) -> dict[str, list[str]]:
+    """List resources available per category.
+
+    Args:
+        organism: Organism identifier (any form).
+    """
+
+    ncbi_tax_id = _resolve_organism(organism)
+    url = f'{_metabo_url}/cosmos/resources'
+    return _dl().fetch_json(url, params={'organism': ncbi_tax_id})
+
+
+def status() -> dict:
+    """Get cache status from the metabo service."""
+
+    url = f'{_metabo_url}/cosmos/status'
+    return _dl().fetch_json(url)