repro_nspsg.py

#!/usr/bin/env python3
"""
Reproduction script for ICML 2026 paper #22179:
"Hard-Constrained Graph Generation with Discrete-Projection Diffusion" (NSPSG)

Claims:
  1. NSPSG achieves 99%-100% validity rates for constrained graph generation
     across heterogeneous constraints.
  2. For complex non-linear constraints, NSPSG improves data validity by up to
     43% while maintaining 99% validity.
"""

import argparse
import json
import os
import random
import time
import sys
from pathlib import Path

import networkx as nx
import numpy as np


# ---------------------------------------------------------------------------
#  SMT-based projectors (Z3) for various graph constraints
# ---------------------------------------------------------------------------

def _z3_available():
    try:
        import z3
        return True
    except ImportError:
        return False


def project_planar(graph: nx.Graph) -> nx.Graph:
    """Project onto planar graphs: greedily add edges while preserving
    planarity to obtain a maximal planar subgraph."""
    g = nx.Graph()
    g.add_nodes_from(graph.nodes())
    # Try adding edges from the original graph, keeping only those that
    # preserve planarity
    edges = list(graph.edges())
    random.shuffle(edges)
    for u, v in edges:
        g.add_edge(u, v)
        is_planar, _ = nx.check_planarity(g)
        if not is_planar:
            g.remove_edge(u, v)
    is_planar, _ = nx.check_planarity(g)
    return g, is_planar


def project_tree(graph: nx.Graph) -> nx.Graph:
    """Project onto tree / forest by removing edges until no cycles remain."""
    g = graph.copy()
    # Find a spanning forest (minimum edge removal)
    if g.number_of_edges() == 0:
        return g, True
    # Use a BFS-based cycle removal
    visited = set()
    to_remove = []
    for source in g.nodes():
        if source in visited:
            continue
        # BFS to find spanning tree edges
        parent = {source: None}
        order = [source]
        visited.add(source)
        for node in order:
            for neighbor in g.neighbors(node):
                if neighbor not in visited:
                    visited.add(neighbor)
                    parent[neighbor] = node
                    order.append(neighbor)
    # Remove non-tree edges
    tree_edges = set()
    for node, p in parent.items():
        if p is not None:
            tree_edges.add(tuple(sorted((node, p))))
    for u, v in list(g.edges()):
        if tuple(sorted((u, v))) not in tree_edges:
            g.remove_edge(u, v)
    is_forest = nx.is_forest(g)
    return g, is_forest


def project_lobster(graph: nx.Graph) -> nx.Graph:
    """Project onto the set of lobster graphs (tree where removing leaves
    yields a caterpillar, and removing those caterpillar leaves yields a path).
    We use a heuristic: first project to tree, then iteratively trim growth."""
    g, _ = project_tree(graph)
    if g.number_of_nodes() <= 2:
        return g, True
    # Check lobster property
    leaves1 = [n for n in g.nodes() if g.degree(n) <= 1]
    g2 = g.copy()
    g2.remove_nodes_from(leaves1)
    if g2.number_of_nodes() <= 1:
        return g, True
    leaves2 = [n for n in g2.nodes() if g2.degree(n) <= 1]
    g3 = g2.copy()
    g3.remove_nodes_from(leaves2)
    # g3 should be a path (max degree <= 2 and connected components are paths)
    is_lobster = all(d <= 2 for _, d in g3.degree()) if g3.number_of_nodes() > 0 else True
    if not is_lobster:
        # Trim further
        to_trim = [n for n in g3.nodes() if g3.degree(n) > 2]
        for n in to_trim:
            neighbors = list(g3.neighbors(n))
            # Keep only 2 edges for nodes with degree > 2 in g3
            keep_edges = list(g3.edges(n))[:2]
            for nbr in neighbors:
                edge = tuple(sorted((n, nbr)))
                if edge not in keep_edges:
                    if g.has_edge(n, nbr):
                        g.remove_edge(n, nbr)
    return g, True


def project_predecessor_balance(graph: nx.DiGraph, rho: float = 0.5) -> nx.DiGraph:
    """Project a directed graph to satisfy predecessor balance constraint.
    For each node: |n0 - n1| / (2 * avg) <= rho where n0 = #0-type preds,
    n1 = #1-type preds."""
    g = graph.copy()
    if _z3_available():
        return _z3_project_predecessor_balance(g, rho)
    return _heuristic_project_predecessor_balance(g, rho)


def _heuristic_project_predecessor_balance(g: nx.DiGraph, rho: float) -> nx.DiGraph:
    """Heuristic projection for predecessor balance."""
    valid = True
    for node in g.nodes():
        preds = list(g.predecessors(node))
        if not preds:
            continue
        n0 = sum(1 for p in preds if g.nodes[p].get("type", 0) == 0)
        n1 = sum(1 for p in preds if g.nodes[p].get("type", 0) == 1)
        total = len(preds)
        if total == 0:
            continue
        imbalance = abs(n0 - n1) / max(2 * total, 1)
        if imbalance > rho:
            valid = False
            # Remove edges to fix imbalance (remove oldest/most recent)
            excess = int((imbalance - rho) * total / 2)
            type1_preds = [(i, p) for i, p in enumerate(preds)
                           if g.nodes[p].get("type", 0) == 1]
            type0_preds = [(i, p) for i, p in enumerate(preds)
                           if g.nodes[p].get("type", 0) == 0]
            if n1 > n0:
                to_remove = [p for _, p in type1_preds[:excess]]
            else:
                to_remove = [p for _, p in type0_preds[:excess]]
            for p in to_remove:
                g.remove_edge(p, node)
    return g, valid


def _z3_project_predecessor_balance(g: nx.DiGraph, rho: float) -> nx.DiGraph:
    """Use Z3 SMT solver for predecessor balance projection."""
    import z3
    g = g.copy()
    opt = z3.Optimize()
    edges = {}
    for u, v in list(g.edges()):
        edges[(u, v)] = z3.Bool(f"e_{u}_{v}")

    for node in g.nodes():
        preds = list(g.predecessors(node))
        if not preds:
            continue
        n0_vars = [edges[(p, node)] for p in preds
                   if g.nodes[p].get("type", 0) == 0 and (p, node) in edges]
        n1_vars = [edges[(p, node)] for p in preds
                   if g.nodes[p].get("type", 0) == 1 and (p, node) in edges]
        n0 = z3.Sum([z3.If(v, 1, 0) for v in n0_vars])
        n1 = z3.Sum([z3.If(v, 1, 0) for v in n1_vars])
        total = z3.Sum([z3.If(v, 1, 0) for v in n0_vars + n1_vars])
        # constraint: |n0 - n1| <= rho * 2 * total
        # encoded as: n0 - n1 <= rho * 2 * total AND n1 - n0 <= rho * 2 * total
        rhs = z3.If(total > 0, rho * 2 * total, 1)
        opt.add(z3.And(z3.And(n0 - n1 <= rhs), z3.And(n1 - n0 <= rhs)))

    # Maximize number of retained edges
    edge_sum = z3.Sum([z3.If(v, 1, 0) for v in edges.values()])
    opt.maximize(edge_sum)

    if opt.check() == z3.sat:
        model = opt.model()
        to_remove = [(u, v) for (u, v), var in edges.items()
                     if not model[var]]
        for u, v in to_remove:
            g.remove_edge(u, v)

    return g, True


# ---------------------------------------------------------------------------
#  Discrete diffusion graph generator (simplified)
# ---------------------------------------------------------------------------

class DiscreteGraphDiffusion:
    """Minimal discrete diffusion for graph generation with projection."""

    def __init__(self, num_nodes: int = 32, T: int = 100, device: str = "cpu"):
        self.num_nodes = num_nodes
        self.T = T
        self.device = device

    def generate(self, num_graphs: int = 100, projection_fn=None,
                 project_every: int = 5) -> list:
        """Generate graphs using reverse diffusion with optional projection."""
        graphs = []
        for _ in range(num_graphs):
            # Start from random graph (noisy state)
            g = nx.erdos_renyi_graph(self.num_nodes, 0.3, seed=random.randint(0, 10000))
            # Simulate reverse diffusion with projection
            for step in range(self.T, 0, -1):
                # Denoise step: randomly add/remove edges toward target density
                density = 0.1 + 0.4 * (step / self.T)
                self._denoise_step(g, density)
                # Apply projection if specified at this step
                if projection_fn and step % project_every == 0:
                    g, _ = projection_fn(g)
            # Final projection
            if projection_fn:
                g, valid = projection_fn(g)
            else:
                valid = True
            graphs.append((g, valid))
        return graphs

    def _denoise_step(self, g: nx.Graph, target_density: float):
        """Single denoising step."""
        current_edges = g.number_of_edges()
        max_edges = self.num_nodes * (self.num_nodes - 1) // 2
        target_edges = int(target_density * max_edges)
        # Add edges
        while g.number_of_edges() < target_edges:
            u = random.randint(0, self.num_nodes - 1)
            v = random.randint(0, self.num_nodes - 1)
            if u != v and not g.has_edge(u, v):
                g.add_edge(u, v)
        # Remove edges
        edges_to_remove = list(g.edges())
        random.shuffle(edges_to_remove)
        while g.number_of_edges() > target_edges and edges_to_remove:
            g.remove_edge(*edges_to_remove.pop())


# ---------------------------------------------------------------------------
#  Claim 1: Structural constraint evaluation
# ---------------------------------------------------------------------------

def evaluate_claim1(output_dir: Path, num_graphs: int = 100,
                    num_nodes: int = 20, use_z3: bool = False):
    """Evaluate Claim 1: NSPSG achieves 99-100% validity for structural constraints."""
    results = {}
    generator = DiscreteGraphDiffusion(num_nodes=num_nodes, T=50, device="cpu")

    constraints = {
        "planar": project_planar,
        "tree": project_tree,
        "lobster": project_lobster,
    }

    for constraint_name, proj_fn in constraints.items():
        print(f"\n=== Evaluating {constraint_name} constraint ===")

        # Without projection
        t0 = time.time()
        graphs_no_proj = generator.generate(
            num_graphs=num_graphs, projection_fn=None
        )
        t_no_proj = time.time() - t0
        valid_no_proj = sum(1 for _, v in graphs_no_proj if v)
        valid_pct_no_proj = valid_no_proj / num_graphs * 100

        # With projection (NSPSG)
        t0 = time.time()
        graphs_proj = generator.generate(
            num_graphs=num_graphs, projection_fn=proj_fn, project_every=10
        )
        t_proj = time.time() - t0
        valid_proj = sum(1 for _, v in graphs_proj if v)

        # Re-check validity for projected graphs
        rechecked = 0
        for g, _ in graphs_proj:
            if constraint_name == "planar":
                is_valid = nx.check_planarity(g)[0]
            elif constraint_name == "tree":
                is_valid = nx.is_forest(g)
            elif constraint_name == "lobster":
                is_valid = _check_lobster(g)
            else:
                is_valid = True
            if is_valid:
                rechecked += 1
        valid_pct_proj = rechecked / num_graphs * 100

        print(f"  Without projection: {valid_pct_no_proj:.1f}% valid "
              f"({t_no_proj:.2f}s)")
        print(f"  With NSPSG projection: {valid_pct_proj:.1f}% valid "
              f"({t_proj:.2f}s)")

        results[constraint_name] = {
            "no_projection_valid_pct": round(valid_pct_no_proj, 1),
            "nspsg_valid_pct": round(valid_pct_proj, 1),
            "no_projection_time_s": round(t_no_proj, 2),
            "nspsg_time_s": round(t_proj, 2),
            "num_graphs": num_graphs,
            "num_nodes": num_nodes,
        }

    # Save results
    results_path = output_dir / "claim1_results.json"
    with open(results_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nClaim 1 results saved to {results_path}")

    # Summary
    print("\n=== Claim 1 Summary ===")
    all_above_99 = all(r["nspsg_valid_pct"] >= 99.0 for r in results.values())
    print(f"NSPSG achieves >=99% validity across all constraints: {all_above_99}")
    for name, r in results.items():
        print(f"  {name}: {r['nspsg_valid_pct']}% valid "
              f"(baseline: {r['no_projection_valid_pct']}%)")

    return results


def _check_lobster(g: nx.Graph) -> bool:
    """Check if a graph is a lobster graph."""
    if not nx.is_forest(g):
        return False
    if g.number_of_nodes() <= 2:
        return True
    leaves1 = [n for n in g.nodes() if g.degree(n) <= 1]
    g2 = g.copy()
    g2.remove_nodes_from(leaves1)
    if g2.number_of_nodes() <= 1:
        return True
    leaves2 = [n for n in g2.nodes() if g2.degree(n) <= 1]
    g3 = g2.copy()
    g3.remove_nodes_from(leaves2)
    return all(d <= 2 for _, d in g3.degree()) if g3.number_of_nodes() > 0 else True


# ---------------------------------------------------------------------------
#  Claim 2: Non-linear constraint evaluation
# ---------------------------------------------------------------------------

def evaluate_claim2(output_dir: Path, num_graphs: int = 100,
                    num_nodes: int = 15, use_z3: bool = False):
    """Evaluate Claim 2: NSPSG improves validity by up to 43% on non-linear
    constraints (predecessor balance)."""
    results = {}

    rhos = [0.0, 0.5, 1.0]

    for rho in rhos:
        print(f"\n=== Evaluating predecessor balance constraint (rho={rho}) ===")

        # Generate random DAGs
        valid_before = 0
        valid_after = 0

        for _ in range(num_graphs):
            # Create DAG with random topological order
            g = nx.DiGraph()
            g.add_nodes_from(range(num_nodes))
            # Random node types
            for n in g.nodes():
                g.nodes[n]["type"] = random.randint(0, 1)

            # Add edges respecting topological order
            order = list(range(num_nodes))
            random.shuffle(order)
            for i in range(num_nodes):
                for j in range(i + 1, num_nodes):
                    if random.random() < 0.3:
                        g.add_edge(order[j], order[i])  # j -> i (j after i in order)

            # Check original validity
            _, valid_orig = _heuristic_project_predecessor_balance(
                g.copy(), rho
            )
            if valid_orig:
                valid_before += 1

            # Apply SMT projection
            g_proj, _ = project_predecessor_balance(g, rho)
            # Check projected validity (re-check)
            _, valid_proj = _heuristic_project_predecessor_balance(
                g_proj, rho
            )
            if valid_proj:
                valid_after += 1

        valid_before_pct = valid_before / num_graphs * 100
        valid_after_pct = valid_after / num_graphs * 100
        improvement = valid_after_pct - valid_before_pct

        print(f"  Before projection: {valid_before_pct:.1f}% valid")
        print(f"  After NSPSG projection: {valid_after_pct:.1f}% valid")
        print(f"  Improvement: {improvement:.1f} pp")

        results[f"rho_{rho}"] = {
            "rho": rho,
            "before_valid_pct": round(valid_before_pct, 1),
            "nspsg_valid_pct": round(valid_after_pct, 1),
            "improvement_pp": round(improvement, 1),
            "num_graphs": num_graphs,
        }

    results_path = output_dir / "claim2_results.json"
    with open(results_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nClaim 2 results saved to {results_path}")

    print("\n=== Claim 2 Summary ===")
    max_improvement = max(r["improvement_pp"] for r in results.values())
    print(f"Max improvement: {max_improvement} percentage points")
    for name, r in results.items():
        print(f"  {name}: {r['nspsg_valid_pct']}% valid "
              f"(improvement: {r['improvement_pp']} pp)")

    return results


# ---------------------------------------------------------------------------
#  Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Reproduce ICML 2026 NSPSG paper claims"
    )
    parser.add_argument("--output-dir", default="./repro_output",
                        help="Output directory for results")
    parser.add_argument("--num-graphs", type=int, default=100,
                        help="Number of graphs to generate per experiment")
    parser.add_argument("--num-nodes", type=int, default=20,
                        help="Number of nodes per graph")
    parser.add_argument("--claim", type=int, choices=[1, 2, 0], default=0,
                        help="Which claim to evaluate (0=both)")
    parser.add_argument("--use-z3", action="store_true",
                        help="Use Z3 SMT solver instead of heuristics")
    args = parser.parse_args()

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Save environment info
    env_info = {
        "python_version": sys.version,
        "networkx_version": nx.__version__,
        "z3_available": _z3_available(),
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S UTC", time.gmtime()),
        "args": vars(args),
    }
    with open(output_dir / "env_info.json", "w") as f:
        json.dump(env_info, f, indent=2)

    all_results = {}

    if args.claim in (0, 1):
        print("=" * 60)
        print("EVALUATING CLAIM 1: Structural Constraints")
        print("=" * 60)
        c1 = evaluate_claim1(output_dir, num_graphs=args.num_graphs,
                             num_nodes=args.num_nodes, use_z3=args.use_z3)
        all_results["claim1"] = c1

    if args.claim in (0, 2):
        print("\n" + "=" * 60)
        print("EVALUATING CLAIM 2: Non-Linear Arithmetic Constraints")
        print("=" * 60)
        c2 = evaluate_claim2(output_dir, num_graphs=args.num_graphs,
                             num_nodes=args.num_nodes, use_z3=args.use_z3)
        all_results["claim2"] = c2

    # Final report
    print("\n" + "=" * 60)
    print("FINAL REPRODUCTION REPORT")
    print("=" * 60)

    if "claim1" in all_results:
        c1 = all_results["claim1"]
        all_valid = all(r["nspsg_valid_pct"] >= 99.0 for r in c1.values())
        print(f"\nClaim 1: {all_valid}")
        print(f"  Reported: NSPSG achieves 99%-100% validity")
        for name, r in c1.items():
            print(f"  {name}: {r['nspsg_valid_pct']}% (baseline: {r['no_projection_valid_pct']}%)")

    if "claim2" in all_results:
        c2 = all_results["claim2"]
        max_imp = max(r["improvement_pp"] for r in c2.values())
        print(f"\nClaim 2: Improvement up to {max_imp} pp")
        print(f"  Reported: up to 43% improvement while maintaining 99% validity")
        for name, r in c2.items():
            print(f"  {name}: {r['nspsg_valid_pct']}% valid (improvement: {r['improvement_pp']} pp)")

    print("\nResults written to:", output_dir.resolve())


if __name__ == "__main__":
    main()