codex-py/codex/graph.py

"""Citation graph analytics (F-15).

Builds an in-memory NetworkX DiGraph from the ``citations`` table and
provides PageRank, bibliographic coupling, co-citation, and dangling-
citation queries.  The graph is ephemeral — rebuilt per call, < 1 s for
≤ 100 papers.

Graceful degradation
--------------------
* ``citation_pagerank`` returns uniform scores when the graph has fewer
  than 5 nodes, and logs a warning.
* All functions accept an empty graph without raising.
"""

from __future__ import annotations

import logging
from typing import Any, cast

import networkx as nx

logger = logging.getLogger(__name__)

_MIN_PAGERANK_NODES = 5

# Single source of truth for resolving ``citations.cited_id`` to a canonical
# paper id. ``cited_id`` stores OpenAlex IDs while ``papers.id`` stores DOIs /
# arXiv ids, so an in-KB reference cited by its OpenAlex id does not match
# ``papers.id`` directly. This resolves it to the canonical ``papers.id`` (via
# ``papers.openalex_id``, then the ``paper_identifiers`` alias table); dangling
# references keep their raw id. Both :func:`build_citation_graph` and
# :mod:`codex.discover` query through this so the "what counts as ingested" rule
# cannot drift between them again (audit C-1). Trusted constant — no user input.
RESOLVED_CITATIONS_SQL = """
    SELECT c.citing_id,
           COALESCE(p.id, pi.paper_id, c.cited_id) AS cited_id
    FROM citations c
    LEFT JOIN papers p ON p.openalex_id = c.cited_id
    LEFT JOIN paper_identifiers pi
           ON pi.openalex_id = c.cited_id AND p.id IS NULL
"""


def build_citation_graph(conn: Any) -> nx.DiGraph:
    """Load the ``citations`` table into a directed graph.

    Nodes are paper IDs (strings).  An edge ``citing → cited`` means
    the citing paper references the cited paper.  Both ingested papers
    and dangling targets (cited but not yet ingested) appear as nodes.

    ``cited_id`` in the citations table stores OpenAlex IDs; ``papers.id``
    stores DOIs.  A LEFT JOIN resolves cited papers that are in the KB to
    their canonical DOI key via ``papers.openalex_id``.  Dangling references
    (not in KB) keep their raw OpenAlex ID.

    Parameters
    ----------
    conn:
        Open psycopg connection (dict-row factory assumed).

    Returns
    -------
    nx.DiGraph with every (citing_id, cited_id) pair as an edge.
    """
    rows = conn.execute(RESOLVED_CITATIONS_SQL).fetchall()
    g: nx.DiGraph = nx.DiGraph()
    for row in rows:
        g.add_edge(row["citing_id"], row["cited_id"])
    logger.debug(
        "Built citation graph: %d nodes, %d edges", g.number_of_nodes(), g.number_of_edges()
    )
    return g


def citation_pagerank(graph: nx.DiGraph, *, damping: float = 0.85) -> dict[str, float]:
    """Compute PageRank over the citation graph.

    Parameters
    ----------
    graph:
        DiGraph from :func:`build_citation_graph`.
    damping:
        PageRank damping factor (default 0.85).

    Returns
    -------
    ``{paper_id: score}`` dict.  Scores sum to approximately 1.0.

    When ``graph`` has fewer than :data:`_MIN_PAGERANK_NODES` (5) nodes,
    returns a uniform distribution and logs a warning — the graph is too
    sparse for the random-walk model to converge meaningfully.  The
    ``damping`` parameter is ignored in this branch.  The user-facing
    corpus-size threshold (default 15) lives in
    :attr:`codex.config.Settings.graph_min_corpus_size` and is surfaced
    by ``codex graph report``.
    """
    n = graph.number_of_nodes()
    if n == 0:
        return {}
    if n < _MIN_PAGERANK_NODES:
        logger.warning(
            "citation_pagerank: only %d nodes — corpus too small for meaningful ranking "
            "(need ≥ %d). Returning uniform scores.",
            n,
            _MIN_PAGERANK_NODES,
        )
        uniform = 1.0 / n
        return {node: uniform for node in graph.nodes()}
    return cast(dict[str, float], nx.pagerank(graph, alpha=damping))


def find_related(paper_id: str, graph: nx.DiGraph, *, min_shared: int = 2) -> list[str]:
    """Bibliographic coupling: papers sharing ≥ ``min_shared`` references with ``paper_id``.

    Two papers are bibliographically coupled when they cite the same sources.
    The more shared references, the more likely they treat related topics.

    Parameters
    ----------
    paper_id:
        Source paper whose references form the seed set.
    graph:
        DiGraph from :func:`build_citation_graph`.
    min_shared:
        Minimum number of shared references to be included.

    Returns
    -------
    List of paper IDs ordered by shared-reference count DESC.
    """
    if paper_id not in graph:
        return []
    # papers cited by paper_id
    references: set[str] = set(graph.successors(paper_id))
    if not references:
        return []

    shared: dict[str, int] = {}
    for ref in references:
        # other papers that also cite this reference
        for co_citer in graph.predecessors(ref):
            if co_citer == paper_id:
                continue
            shared[co_citer] = shared.get(co_citer, 0) + 1

    return [
        pid for pid, count in sorted(shared.items(), key=lambda x: -x[1]) if count >= min_shared
    ]


def find_co_cited(paper_id: str, graph: nx.DiGraph) -> list[tuple[str, int]]:
    """Co-citation: papers frequently cited alongside ``paper_id``.

    A paper X is co-cited with ``paper_id`` when some third paper cites both.

    Returns
    -------
    ``[(paper_id, count)]`` ordered by count DESC.
    """
    if paper_id not in graph:
        return []
    # papers that cite paper_id
    citers: set[str] = set(graph.predecessors(paper_id))
    if not citers:
        return []

    co_cited: dict[str, int] = {}
    for citer in citers:
        for other in graph.successors(citer):
            if other == paper_id:
                continue
            co_cited[other] = co_cited.get(other, 0) + 1

    return sorted(co_cited.items(), key=lambda x: -x[1])


def dangling_citations(graph: nx.DiGraph, known_ids: set[str]) -> list[str]:
    """Return graph nodes not present in ``known_ids`` (cited but not yet ingested).

    Complements :func:`codex.discover.discovery_leads` with graph context:
    the returned IDs are reachable in the citation graph but have no paper
    record in the DB.

    Parameters
    ----------
    graph:
        DiGraph from :func:`build_citation_graph`.
    known_ids:
        Set of paper IDs that exist in the ``papers`` table.

    Returns
    -------
    List of dangling paper IDs (unordered).
    """
    return [node for node in graph.nodes() if node not in known_ids]