Research-Stack/scripts/load_complete_graph.py

#!/usr/bin/env -S uv run
# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "gremlinpython",
#   "python-dotenv",
# ]
# ///
"""
load_complete_graph.py — Load a complete directed graph K_n into mathblob

Every point touches every point: for each ordered pair of distinct vertices
(i, j) there is a directed edge i → j.  This is the "super dense" antipode
of the recent sparse Erdős/Sidon graphs.

Run with:  uv run scripts/load_complete_graph.py [--n N] [--label LABEL]
Requires:  .env.gremlin
"""

import os
import argparse
import time
from pathlib import Path
from dotenv import load_dotenv
from gremlin_python.driver import client as gremlin_client, serializer

# ── Config ────────────────────────────────────────────────────────────────────

ROOT     = Path(__file__).parent.parent
ENV_FILE = ROOT / ".env.gremlin"

load_dotenv(ENV_FILE)

ENDPOINT = os.environ["GREMLIN_ENDPOINT"]
USERNAME = os.environ["GREMLIN_USERNAME"]
PASSWORD = os.environ["GREMLIN_PASSWORD"]

BATCH_SIZE = 50

# ── Gremlin helpers ───────────────────────────────────────────────────────────

def make_client():
    return gremlin_client.Client(
        ENDPOINT, "g",
        username=USERNAME,
        password=PASSWORD,
        message_serializer=serializer.GraphSONSerializersV2d0(),
    )

def submit(c, query: str, bindings: dict = None):
    try:
        cb = c.submitAsync(query, bindings or {})
        return cb.result().all().result()
    except Exception as e:
        print(f"  ERR: {e!s:.120}")
        return None

def _safe_id(raw: str) -> str:
    """Cosmos DB element IDs cannot contain '/', '\', '|', or spaces."""
    return raw.replace("/", "_").replace("\\", "_").replace("|", "_").replace(" ", "_")

# ── Vertex / edge upserts ─────────────────────────────────────────────────────

def upsert_vertex(c, vid: str, vlabel: str, graph_label: str):
    q = (
        f"g.V().has('{vlabel}','id',vid).fold()"
        f".coalesce(unfold(),"
        f"addV('{vlabel}').property('id',vid).property('pk',vid).property('graph',graph_label))"
        f".property('graph',graph_label)"
    )
    submit(c, q, {"vid": vid, "graph_label": graph_label})

def upsert_edge(c, src: str, dst: str, elabel: str, vlabel: str, graph_label: str):
    q = (
        f"g.V().has('{vlabel}','id',src).as('a')"
        f".V().has('{vlabel}','id',dst).as('b')"
        f".coalesce(__.select('a').outE('{elabel}').where(__.inV().as('b')),"
        f"__.select('a').addE('{elabel}').to('b').property('graph',graph_label))"
        f".property('graph',graph_label)"
    )
    submit(c, q, {"src": src, "dst": dst, "graph_label": graph_label})

# ── Complete graph generation ─────────────────────────────────────────────────

def generate_complete_graph(n: int, graph_label: str):
    vertices = []
    for i in range(n):
        vid = _safe_id(f"{graph_label}:node:{i}")
        vertices.append({"id": vid, "label": "complete_node", "index": i})

    edges = []
    for i in range(n):
        for j in range(n):
            if i != j:
                src = _safe_id(f"{graph_label}:node:{i}")
                dst = _safe_id(f"{graph_label}:node:{j}")
                edges.append((src, dst, "touches"))

    return vertices, edges

# ── Load ──────────────────────────────────────────────────────────────────────

def load_complete_graph(n: int, graph_label: str):
    vertices, edges = generate_complete_graph(n, graph_label)
    print(f"K_{n}: {len(vertices)} vertices, {len(edges)} edges")

    c = make_client()
    try:
        print("Upserting vertices...")
        for idx, v in enumerate(vertices):
            upsert_vertex(c, v["id"], "complete_node", graph_label)
            if (idx + 1) % BATCH_SIZE == 0:
                print(f"  {idx + 1}/{len(vertices)}")
        print(f"  {len(vertices)} vertices done")

        print("Upserting edges...")
        for idx, (src, dst, label) in enumerate(edges):
            upsert_edge(c, src, dst, label, "complete_node", graph_label)
            if (idx + 1) % BATCH_SIZE == 0:
                print(f"  {idx + 1}/{len(edges)}")
        print(f"  {len(edges)} edges done")

        # Verification query
        vcount = submit(c, "g.V().has('graph',graph_label).count()", {"graph_label": graph_label})
        ecount = submit(c, "g.E().has('graph',graph_label).count()", {"graph_label": graph_label})
        print(f"Loaded: {vcount} vertices, {ecount} edges")
    finally:
        c.close()

# ── Main ──────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Load a complete directed graph K_n into mathblob")
    parser.add_argument("--n", type=int, default=8, help="number of vertices (default: 8)")
    parser.add_argument("--label", type=str, default="complete_graph_k8", help="graph label/tag")
    args = parser.parse_args()

    if args.n < 2:
        raise ValueError("n must be at least 2")

    start = time.time()
    load_complete_graph(args.n, args.label)
    elapsed = time.time() - start
    print(f"Finished in {elapsed:.1f}s")

if __name__ == "__main__":
    main()