geographdb-core 0.5.4

Geometric graph database core - 3D spatial indexing for code analysis
Documentation
#!/usr/bin/env python3
"""Plot the FFN low-rank + residual sweep produced by low_rank_ffn.rs.

Usage:
    python examples/plot_low_rank_ffn.py --csv target/low_rank_ffn.csv --png target/low_rank_ffn.png
"""

import argparse
import csv
import os
from pathlib import Path


def main():
    parser = argparse.ArgumentParser(description="Plot FFN low-rank sweep results")
    parser.add_argument("--csv", default="target/low_rank_ffn.csv", help="input CSV path")
    parser.add_argument("--png", default="target/low_rank_ffn.png", help="output PNG path")
    args = parser.parse_args()

    try:
        import matplotlib.pyplot as plt
    except ImportError as exc:
        raise SystemExit("matplotlib is required: pip install matplotlib") from exc

    ranks = []
    low_rank_params = []
    logit_l2 = []
    logit_mse = []
    top1 = []
    residual_l2 = []
    residual_top1 = []

    with open(args.csv, newline="") as f:
        reader = csv.DictReader(f)
        for row in reader:
            ranks.append(int(row["rank"]))
            low_rank_params.append(int(row["low_rank_params"]))
            logit_l2.append(float(row["logit_l2"]))
            logit_mse.append(float(row["logit_mse"]))
            top1.append(float(row["top1_agreement"]) * 100.0)
            residual_l2.append(float(row["residual_logit_l2"]))
            residual_top1.append(float(row["residual_top1_agreement"]) * 100.0)

    fig, axes = plt.subplots(2, 2, figsize=(12, 10))
    fig.suptitle("FFN low-rank + residual prototype")

    ax = axes[0, 0]
    ax.plot(ranks, logit_l2, "o-", label="low-rank")
    ax.plot(ranks, residual_l2, "s--", label="low-rank + dense residual")
    ax.set_xlabel("rank")
    ax.set_ylabel("mean logit L2 drift")
    ax.set_xscale("log", base=2)
    ax.set_yscale("log")
    ax.legend()
    ax.grid(True, which="both", ls="--", alpha=0.5)

    ax = axes[0, 1]
    ax.plot(low_rank_params, logit_l2, "o-", label="low-rank")
    ax.plot(low_rank_params, residual_l2, "s--", label="low-rank + dense residual")
    ax.set_xlabel("low-rank parameters")
    ax.set_ylabel("mean logit L2 drift")
    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.legend()
    ax.grid(True, which="both", ls="--", alpha=0.5)

    ax = axes[1, 0]
    ax.plot(ranks, logit_mse, "o-", label="low-rank MSE")
    ax.set_xlabel("rank")
    ax.set_ylabel("mean logit MSE")
    ax.set_xscale("log", base=2)
    ax.set_yscale("log")
    ax.legend()
    ax.grid(True, which="both", ls="--", alpha=0.5)

    ax = axes[1, 1]
    ax.plot(ranks, top1, "o-", label="low-rank top-1")
    ax.plot(ranks, residual_top1, "s--", label="low-rank + residual top-1")
    ax.set_xlabel("rank")
    ax.set_ylabel("top-1 agreement (%)")
    ax.set_xscale("log", base=2)
    ax.set_ylim(-5, 105)
    ax.legend()
    ax.grid(True, which="both", ls="--", alpha=0.5)

    fig.tight_layout()
    out = Path(args.png)
    out.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out, dpi=150)
    print(f"Saved {out}")


if __name__ == "__main__":
    main()