layout_precompute.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import random
from pathlib import Path
from typing import List, Optional, Tuple

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq


def _read_int_column(parquet_path: Path, column_name: str) -> np.ndarray:
    table = pq.read_table(parquet_path, columns=[column_name])
    return table[column_name].to_numpy()


def _build_node_index(
    source_ids: np.ndarray, target_ids: np.ndarray, node_ids: Optional[np.ndarray]
) -> np.ndarray:
    edge_node_ids = np.unique(np.concatenate((source_ids, target_ids)))
    if node_ids is None:
        return edge_node_ids
    return np.union1d(edge_node_ids, np.unique(node_ids))


def _build_edges_as_index_pairs(
    source_ids: np.ndarray, target_ids: np.ndarray, node_ids: np.ndarray
) -> List[Tuple[int, int]]:
    source_idx = np.searchsorted(node_ids, source_ids)
    target_idx = np.searchsorted(node_ids, target_ids)
    return list(zip(source_idx.tolist(), target_idx.tolist()))


def _read_layout_seed(layout_path: Path) -> Tuple[np.ndarray, np.ndarray]:
    table = pq.read_table(layout_path, columns=["id", "x", "y"])
    ids = table["id"].to_numpy()
    x = table["x"].to_numpy()
    y = table["y"].to_numpy()
    order = np.argsort(ids)
    ids = ids[order]
    coords = np.column_stack((x[order], y[order]))
    return ids, coords


def _build_incremental_seed(
    all_node_ids: np.ndarray, previous_layout_path: Path, seed: int
) -> Tuple[np.ndarray, np.ndarray]:
    prev_ids, prev_coords = _read_layout_seed(previous_layout_path)
    node_count = all_node_ids.size

    init_coords = np.zeros((node_count, 2), dtype=np.float64)
    fixed_mask = np.zeros(node_count, dtype=bool)

    prev_pos = np.searchsorted(all_node_ids, prev_ids)
    in_bounds = prev_pos < node_count
    is_match = np.zeros(prev_pos.shape, dtype=bool)
    is_match[in_bounds] = all_node_ids[prev_pos[in_bounds]] == prev_ids[in_bounds]

    matched_pos = prev_pos[is_match]
    matched_coords = prev_coords[is_match]
    init_coords[matched_pos] = matched_coords
    fixed_mask[matched_pos] = True

    new_mask = ~fixed_mask
    if np.any(new_mask):
        rng = np.random.default_rng(seed)
        if np.any(fixed_mask):
            center = init_coords[fixed_mask].mean(axis=0)
            spread = np.std(init_coords[fixed_mask], axis=0)
            spread = np.maximum(spread * 0.05, 1e-3)
            init_coords[new_mask] = rng.normal(
                loc=center, scale=spread, size=(np.count_nonzero(new_mask), 2)
            )
        else:
            init_coords[new_mask] = rng.normal(
                loc=0.0, scale=1.0, size=(np.count_nonzero(new_mask), 2)
            )

    return init_coords, fixed_mask


def _compute_layout_with_igraph(
    edges: List[Tuple[int, int]],
    node_count: int,
    seed: int,
    mode: str = "drl_default",
    init_coords: Optional[np.ndarray] = None,
    fixed_mask: Optional[np.ndarray] = None,
    drl_options: str = "default",
    fr_niter: int = 200,
    fr_start_temp: float = 2.0,
) -> np.ndarray:
    try:
        import igraph as ig
    except ModuleNotFoundError as exc:
        raise SystemExit(
            "Missing dependency: python-igraph. Install with `pip install .`"
        ) from exc

    # Make placement reproducible.
    ig.set_random_number_generator(random.Random(seed))
    graph = ig.Graph(n=node_count, edges=edges, directed=False)
    graph.simplify(multiple=True, loops=True)

    if mode == "drl_default":
        # DRL is a force-directed layout suitable for larger sparse graphs.
        layout = graph.layout_drl(options=drl_options)
    elif mode == "drl_refine":
        if init_coords is None:
            raise SystemExit("`drl_refine` requires an incremental seed layout")
        layout = graph.layout_drl(seed=init_coords, options=drl_options)
    elif mode == "fr_freeze":
        if init_coords is None or fixed_mask is None:
            raise SystemExit("`fr_freeze` requires an incremental seed layout")
        limit = 1e12
        minx = np.where(fixed_mask, init_coords[:, 0], -limit)
        maxx = np.where(fixed_mask, init_coords[:, 0], limit)
        miny = np.where(fixed_mask, init_coords[:, 1], -limit)
        maxy = np.where(fixed_mask, init_coords[:, 1], limit)
        layout = graph.layout_fruchterman_reingold(
            seed=init_coords,
            niter=fr_niter,
            start_temp=fr_start_temp,
            minx=minx,
            maxx=maxx,
            miny=miny,
            maxy=maxy,
            grid="grid",
        )
    else:
        raise SystemExit(f"Unknown layout mode: {mode}")

    return np.asarray(layout.coords, dtype=np.float64)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Precompute force-directed x/y coordinates from parquet edges."
    )
    parser.add_argument("--edges", type=Path, default=Path("model_atlas_edges.parquet"))
    parser.add_argument("--nodes", type=Path, default=Path("model_atlas_nodes.parquet"))
    parser.add_argument("--output", type=Path, default=Path("model_atlas_layout.parquet"))
    parser.add_argument("--source-col", default="source")
    parser.add_argument("--target-col", default="target")
    parser.add_argument("--node-id-col", default="id")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument(
        "--mode",
        choices=["drl_default", "drl_refine", "fr_freeze"],
        default="drl_default",
    )
    parser.add_argument("--previous-layout", type=Path, default=None)
    parser.add_argument("--drl-options", default=None)
    parser.add_argument("--fr-niter", type=int, default=200)
    parser.add_argument("--fr-start-temp", type=float, default=2.0)
    args = parser.parse_args()

    source_ids = _read_int_column(args.edges, args.source_col)
    target_ids = _read_int_column(args.edges, args.target_col)

    node_ids: Optional[np.ndarray] = None
    if args.nodes.exists():
        node_ids = _read_int_column(args.nodes, args.node_id_col)

    all_node_ids = _build_node_index(source_ids, target_ids, node_ids)
    edges_idx = _build_edges_as_index_pairs(source_ids, target_ids, all_node_ids)

    init_coords: Optional[np.ndarray] = None
    fixed_mask: Optional[np.ndarray] = None
    if args.mode in ("drl_refine", "fr_freeze"):
        if args.previous_layout is None:
            raise SystemExit(
                f"`{args.mode}` requires --previous-layout with id/x/y columns"
            )
        if not args.previous_layout.exists():
            raise SystemExit(f"Previous layout not found: {args.previous_layout}")
        init_coords, fixed_mask = _build_incremental_seed(
            all_node_ids, args.previous_layout, args.seed
        )

    drl_options = args.drl_options
    if drl_options is None:
        drl_options = "refine" if args.mode == "drl_refine" else "default"

    coords = _compute_layout_with_igraph(
        edges=edges_idx,
        node_count=all_node_ids.size,
        seed=args.seed,
        mode=args.mode,
        init_coords=init_coords,
        fixed_mask=fixed_mask,
        drl_options=drl_options,
        fr_niter=args.fr_niter,
        fr_start_temp=args.fr_start_temp,
    )

    output_table = pa.table(
        {
            "id": pa.array(all_node_ids),
            "x": pa.array(coords[:, 0]),
            "y": pa.array(coords[:, 1]),
        }
    )
    pq.write_table(output_table, args.output, compression="zstd")

    print(f"Wrote {len(all_node_ids)} coordinates to {args.output}")


if __name__ == "__main__":
    main()