grid-simulate 0.1.0__tar.gz → 0.2.2__tar.gz

grid_simulate-0.2.2/.github/workflows/publish.yml

@@ -0,0 +1,59 @@
+name: Publish to PyPI
+
+on:
+  release:
+    types: [published]
+
+permissions:
+  contents: read
+  id-token: write
+
+jobs:
+  build:
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, windows-latest, macos-latest]
+
+    runs-on: ${{ matrix.os }}
+
+    steps:
+      - uses: actions/checkout@v5
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.13"
+
+      - name: Install maturin
+        run: pip install maturin==1.14.1
+
+      - name: Build wheels
+        env:
+          PYO3_USE_ABI3_FORWARD_COMPATIBILITY: "1"
+        run: |
+          maturin build --release --sdist --interpreter python3.11 python3.12 python3.13
+
+      - uses: actions/upload-artifact@v5
+        with:
+          name: wheels-${{ matrix.os }}
+          path: target/wheels/*
+
+
+  publish:
+    needs: build
+    runs-on: ubuntu-latest
+
+    permissions:
+      id-token: write
+
+    environment:
+      name: pypi
+
+    steps:
+      - uses: actions/download-artifact@v5
+        with:
+          pattern: wheels-*
+          merge-multiple: true
+          path: dist
+
+      - uses: pypa/gh-action-pypi-publish@release/v1

grid_simulate-0.2.2/.gitignore

@@ -0,0 +1,19 @@
+__pycache__/
+*.py[cod]
+.venv/
+venv/
+
+/target
+*.pyd
+*.so
+
+# Models
+*.pt
+*.pth
+*.ckpt
+*.bin
+*.safetensors
+
+# IDE
+.idea/
+.vscode/

{grid_simulate-0.1.0 → grid_simulate-0.2.2}/Cargo.lock

@@ -25,7 +25,7 @@ checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
 
 [[package]]
 name = "grid"
-version = "0.1.0"
+version = "0.2.2"
 dependencies = [
  "approx",
  "ndarray",

{grid_simulate-0.1.0 → grid_simulate-0.2.2}/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "grid"
-version = "0.1.0"
+version = "0.2.2"
 edition = "2021"
 readme = "README.md"
 

{grid_simulate-0.1.0 → grid_simulate-0.2.2}/PKG-INFO

@@ -1,5 +1,5 @@
 Metadata-Version: 2.4
 Name: grid-simulate
-Version: 0.1.0
+Version: 0.2.2
 Summary: Rust + PyO3 simulation library
 Requires-Python: >=3.8

grid_simulate-0.2.2/gnn/algorithms.py

@@ -0,0 +1,193 @@
+"""
+Pebble Game algorithm for 2D rigidity independence.
+Deterministic & O(n^2)
+
+(Jacobs & Hendrickson (1997)
+        "An Algorithm for Two-Dimensional Rigidity Percolation")
+author: Sezer
+"""
+
+from typing import List, Tuple
+
+# Core pebble search
+def _find_pebble(
+        v: int,
+        visited: set,
+        pebbles: List[int],
+        out_edges: List[List[int]],
+        changes: List[tuple],
+) -> bool:
+    """
+    DFS: try to find a free pebble reachable from v by following directed edges.
+    All state changes are recorded in 'changes' so they can be rolled back if the
+    overall edge turns out to be redundant.
+
+    A free pebble is reachable from v if:
+        - v itself has a pebble, OR
+        - v has an outgoing directed edge v->u and a pebble is reachable from u
+        (in which case we reverse the edge u->v to move the pebble toward v)
+    """
+    if v in visited:
+        return False
+    visited.add(v)
+
+    # Base case
+    if pebbles[v] > 0:
+        pebbles[v] -= 1
+        changes.append(('pebble', v))
+        return True
+
+    for i in range(len(out_edges[v])):
+        u = out_edges[v][i]
+        if _find_pebble(u, visited, pebbles, out_edges, changes):
+            out_edges[v].pop(i)
+            out_edges[u].append(v)
+            changes.append(('edge', v, u))
+            return True
+
+    return False
+
+def _rollback(changes: List[tuple], pebbles: List[int], out_edges: List[List[int]]):
+    """
+    Undo all state changes recorded during a failed edge attempt
+    Applied in reverse order so each change is cleanly undone
+    """
+    for change in reversed(changes):
+        if change[0] == 'pebble':
+            pebbles[change[1]] += 1
+        elif change[0] == 'gain_pebble':
+            pebbles[change[1]] -= 1
+        elif change[0] == 'edge':
+            v, u = change[1], change[2]
+            out_edges[u].remove(v)
+            out_edges[v].append(u)
+
+def pebble_game(
+        n: int,
+        candidate_edges: List[Tuple[int, int]],
+        d: int = 2,
+) -> Tuple[List[int], List[Tuple[int, int]]]:
+    """
+    Pebble Game algorithm for finding independent edges in a rigidity framework.
+
+    Only supports d=2
+    For d=3, falls back to gaussian elimination
+
+    Args:
+        n               : number of nodes
+        candidate_edges : edges in upper-triangular order (i < j)
+        d               : embedding dimension (must be 2)
+
+    Returns:
+        labels    : list[int]        1=independent, 0=redundant
+        accepted  : list[(i, j)]     the independent edges only
+
+    Complexity: O(n^2)
+    """
+    assert d == 2
+
+    pebbles = [2] * n
+    # Directed graph: out_edges[v] = list of nodes v points to
+    # A directed edge v->u means "v has borrowed a pebble from u's region"
+    out_edges: List[List[int]] = [[] for _ in range(n)]
+
+    labels = []
+    accepted = []
+
+    for (u, v) in candidate_edges:
+        changes = []
+
+        while pebbles[u] + pebbles[v] < 4:
+
+            pu, pv = pebbles[u], pebbles[v]
+            pebbles[u], pebbles[v] = 0, 0
+
+            progressed = False
+
+            if pu < 2 and _find_pebble(u, set(), pebbles, out_edges, changes):
+                pu += 1
+                changes.append(('gain_pebble', u))
+                progressed = True
+            elif pv < 2 and _find_pebble(v, set(), pebbles, out_edges, changes):
+                pv += 1
+                changes.append(('gain_pebble', v))
+                progressed = True
+
+            pebbles[u], pebbles[v] = pu, pv
+
+            if not progressed:
+                break
+
+        if pebbles[u] + pebbles[v] >= 4:
+            # Independent: consume exactly 1 pebble to cover this edge
+            pebbles[v] -= 1
+            out_edges[v].append(u)
+            labels.append(1)
+            accepted.append((u, v))
+        else:
+            # Redundant: undo everything collected during this attempt
+            _rollback(changes, pebbles, out_edges)
+            labels.append(0)
+
+    return labels, accepted
+
+def gaussian_elimination(
+        n: int,
+        candidate_edges: List[Tuple[int, int]],
+        d: int= 2,
+) -> Tuple[List[int], List[Tuple[int, int]]]:
+    """
+    Classical Gaussian elimination rank check for edge independence.
+    Works for both d=2 and d=3.
+
+    This mirrors exactly what the Rust simulation does internally,
+    implementing it just for comparison
+    """
+    ncols = n * d
+    pivot_rows: List[List[float]] = []
+    pivot_cols: List[int] = []
+    tol = 1e-9
+
+    labels = []
+    accepted = []
+
+    def add_row(row: List[float]) -> bool:
+        """Returns True if row is linearly independent of existing rows."""
+        r = row[:]
+
+        for k, col in enumerate(pivot_cols):
+            factor = r[col]
+            if abs(factor) > tol:
+                for j in range(ncols):
+                    r[j] -= factor * pivot_rows[k][j]
+        # Find new pivot
+        new_pivot = next((j for j in range(ncols) if abs(r[j]) > tol), None)
+        if new_pivot is None:
+            return False  # redundant
+        # Normalise
+        scale = r[new_pivot]
+        r = [v / scale for v in r]
+        # Back-substitute into existing rows
+        for k in range(len(pivot_rows)):
+            factor = pivot_rows[k][new_pivot]
+            if abs(factor) > tol:
+                for j in range(ncols):
+                    pivot_rows[k][j] -= factor * r[j]
+        pivot_rows.append(r)
+        pivot_cols.append(new_pivot)
+        return True
+
+    for (i, j) in candidate_edges:
+        row = [0.0] * ncols
+        for dim in range(d):
+            row[d * i + dim] = 1.0
+            row[d * j + dim] = -1.0
+            # This version without coords only checks structural independence
+            # Use grid.simulate() for the geometry-aware version
+        independent = add_row(row)
+        labels.append(1 if independent else 0)
+        if independent:
+            accepted.append((i, j))
+
+    return labels, accepted
+

grid_simulate-0.2.2/gnn/compare.py

@@ -0,0 +1,221 @@
+"""
+Tests the Pebble Game algorithm against rust ground truth and trained model
+author: Sezer
+"""
+
+import time
+import numpy as np
+import torch
+import grid
+
+from algorithms import pebble_game
+from dataset import make_graph
+from model import RigidityGNN
+
+### Helpers
+
+def upper_triangular_edges(n: int):
+    return [(i, j) for i in range(n) for j in range(i + 1, n)]
+
+def load_gnn(path: str = "best_model.pt", coord_dim: int = 2, hidden_dim: int = 64,
+             n_layers: int = 3):
+    """"Load trained GNN"""
+    try:
+        model = RigidityGNN(coord_dim=coord_dim, hidden_dim=hidden_dim, n_layers=n_layers)
+        model.load_state_dict(torch.load(path, map_location='cpu'))
+        model.eval()
+        print(f"  GNN loaded from {path}")
+        return model
+    except FileNotFoundError:
+        print(f"  No trained GNN found at {path}, skipping GNN comparison")
+        return None
+
+def run_gnn(model, coords: np.ndarray) -> list:
+    """Run GNN prediction on one graph. Returns label list."""
+    graph = make_graph(coords, d=2)
+    with torch.no_grad():
+        preds = model.predict(graph.x, graph.edge_index, graph.candidate_edges)
+    return preds.tolist()
+
+# Single graph comparison
+
+def compare_single(coords: np.ndarray, gnn_model=None, verbose: bool = True):
+    """
+    Run all three methods on one graph and compare results.
+
+    Returns dict with per-method lables and timing.
+    """
+    n = coords.shape[0]
+    d = 2
+    edges = upper_triangular_edges(n)
+
+    results = {}
+
+    # Rust ground truth
+    t0 = time.perf_counter()
+    rust_result = grid.simulate(coords.astype(np.float64), d=d)
+    rust_time = time.perf_counter() - t0
+    rust_labels = rust_result["labels"].tolist()
+    results["rust"] = {"labels": rust_labels, "time_ms": rust_time * 1000}
+
+    # Pebble Game
+    t0 = time.perf_counter()
+    pb_labels, pb_accepted = pebble_game(n, edges, d=d)
+    pb_time = time.perf_counter() - t0
+    results["pebble"] = {"labels": pb_labels, "time_ms": pb_time * 1000}
+
+    # GNN
+    if gnn_model is not None:
+        t0 = time.perf_counter()
+        gnn_labels = run_gnn(gnn_model, coords)
+        gnn_time = time.perf_counter() - t0
+        results["gnn"] = {"labels": gnn_labels, "time_ms": gnn_time * 1000}
+
+    # Agreement Metrics
+    def agreement(labels_a, labels_b):
+        matches = sum(a == b for a,b in zip(labels_a, labels_b))
+        return matches / len(labels_a) if labels_a else 1.0
+
+    if verbose:
+        print(f"\n  n={n}  total_edges={len(edges)}  max_rank={rust_result['max_rank']}")
+        print(f"  {'Method':<12} {'Time (ms)':>10}  {'vs Rust':>10}  {'Independent':>12}")
+        print(f"  {'-' * 50}")
+
+        n_rust_indep = sum(rust_labels)
+        print(f"  {'Rust':<12} {rust_time * 1000:>10.3f}  {'(truth)':>10}  {n_rust_indep:>12}")
+
+        n_pb_indep = sum(pb_labels)
+        pb_agree = agreement(pb_labels, rust_labels)
+        print(f"  {'Pebble Game':<12} {pb_time * 1000:>10.3f}  {pb_agree:>10.4f}  {n_pb_indep:>12}")
+
+        if gnn_model is not None:
+            n_gnn_indep = sum(gnn_labels)
+            gnn_agree = agreement(gnn_labels, rust_labels)
+            print(f"  {'GNN':<12} {gnn_time * 1000:>10.3f}  {gnn_agree:>10.4f}  {n_gnn_indep:>12}")
+
+    return results
+
+# Stress TestL reconstruct from pebble game output
+
+def stress_test(coords: np.ndarray, labels: list) -> dict:
+    """
+    Use the accepted edges from the pebble game to reconstruct coordinates.
+    Measures whether the pebble game found a valid rigid framework.
+    """
+    sim = grid.simulate(coords.astype(np.float64), d=2)
+    rec = grid.reconstruct(sim["distance_matrix"], d=2)
+    return {
+        "stress": rec["stress"],
+        "is_isometric": rec["is_isometric"],
+    }
+
+# Batch benchmark
+
+def benchmark(n_graphs: int = 200, n_nodes: int = 10, seed: int = 42, gnn_model=None):
+    """
+    Run all methods on n_graphs random graphs
+    Reports accuracy and timing stats
+    """
+    rng = np.random.RandomState(seed)
+    d = 2
+
+    pb_agreements = []
+    gnn_agreements = []
+    pb_times = []
+    rust_times = []
+    gnn_times = []
+    stress_scores = []
+
+    print(f"\nBenchmarking {n_graphs} graphs (n={n_nodes}, d={d}) ...")
+    print(f"{'=' * 40}")
+
+    for k in range(n_graphs):
+        coords = rng.uniform(0, 10, (n_nodes, d))
+        r = compare_single(coords, gnn_model=gnn_model, verbose=False)
+
+        rust_labels = r["rust"]["labels"]
+        pb_labels = r["pebble"]["labels"]
+
+        agree = sum(a == b for a,b in zip(pb_labels, rust_labels)) / len(rust_labels)
+        pb_agreements.append(agree)
+        rust_times.append(r["rust"]["time_ms"])
+        pb_times.append(r["pebble"]["time_ms"])
+
+        # Stress test on pebble game result
+        s = stress_test(coords, pb_labels)
+        stress_scores.append(s["stress"])
+
+        if gnn_model is not None:
+            gnn_labels = r["gnn"]["labels"]
+            gnn_agree = sum(a == b for a, b in zip(gnn_labels, rust_labels)) / len(rust_labels)
+            gnn_agreements.append(gnn_agree)
+            gnn_times.append(r["gnn"]["time_ms"])
+
+    print(f"\nResults over {n_graphs} graphs:")
+    print(f"\n  Pebble Game:")
+    print(f"    Agreement with Rust : {np.mean(pb_agreements):.4f}  (std={np.std(pb_agreements):.4f})")
+    print(f"    Avg time            : {np.mean(pb_times):.3f} ms")
+    print(f"    Rust avg time       : {np.mean(rust_times):.3f} ms")
+    print(f"    Speedup over Rust   : {np.mean(rust_times) / np.mean(pb_times):.2f}x")
+    print(f"    Avg stress          : {np.mean(stress_scores):.2e}  (0=lossless)")
+    print(f"    Isometric rate      : {sum(s < 1e-6 for s in stress_scores)}/{n_graphs}")
+
+    if gnn_model is not None:
+        print(f"\n  GNN:")
+        print(f"    Agreement with Rust : {np.mean(gnn_agreements):.4f}  (std={np.std(gnn_agreements):.4f})")
+        print(f"    Avg time            : {np.mean(gnn_times):.3f} ms")
+
+    print(
+        f"\n  Winner: {'Pebble Game' if not gnn_agreements or np.mean(pb_agreements) >= np.mean(gnn_agreements) else 'GNN'}")
+
+# Scaling test
+def scaling_test():
+    """
+    Show how pebble game scales vs Rust as n grows.
+    """
+    print(f"\nScaling test:")
+    print(f"  {'n':>6}  {'Rust (ms)':>10}  {'Pebble (ms)':>12}  {'Speedup':>8}  {'Agreement':>10}")
+    print(f"  {'-'*55}")
+
+    rng = np.random.RandomState(0)
+    for n in [5, 10, 20, 50, 100, 200]:
+        coords = rng.uniform(0, 10, (n, 2))
+        edges  = upper_triangular_edges(n)
+
+        t0 = time.perf_counter()
+        rust_result = grid.simulate(coords.astype(np.float64), d=2)
+        rust_t = (time.perf_counter() - t0) * 1000
+
+        t0 = time.perf_counter()
+        pb_labels, _ = pebble_game(n, edges, d=2)
+        pb_t = (time.perf_counter() - t0) * 1000
+
+        rust_labels = rust_result["labels"].tolist()
+        agree = sum(a == b for a, b in zip(pb_labels, rust_labels)) / len(rust_labels)
+        speedup = rust_t / pb_t if pb_t > 0 else float('inf')
+
+        print(f"  {n:>6}  {rust_t:>10.3f}  {pb_t:>12.3f}  {speedup:>8.2f}x  {agree:>10.4f}")
+
+if __name__ == "__main__":
+    print("GRiD — Algorithm Comparison")
+    print("Pebble Game  vs  Rust Oracle  vs  GNN")
+    print("=" * 60)
+
+    gnn = load_gnn("best_model.pt")
+
+    # Single graph detailed view
+    print("\n── Single graph (n=5, the R1-R5 example) ──")
+    coords_5 = np.array([
+        [0.0, 0.0],
+        [2.0, 0.0],
+        [3.0, 1.5],
+        [1.0, 2.5],
+        [-1.0, 1.5],
+    ])
+    compare_single(coords_5, gnn_model=gnn, verbose=True)
+
+    # Batch benchmark
+    benchmark(n_graphs=200, n_nodes=10, gnn_model=gnn)
+
+    # Scaling
+    scaling_test()

grid_simulate-0.2.2/gnn/dataset.py

@@ -0,0 +1,90 @@
+"""
+@author: Sezer
+date: 01/07/2026
+"""
+
+import numpy as np
+import torch
+from torch_geometric.data import Data
+from torch.utils.data import Dataset
+from torchgen.dest.native_functions import torch_api_key_word_prefix
+import grid
+
+class RigidityData(Data):
+    """
+    Custom PyG Data subclass.
+
+    PyG's default bathing auto-offsets 'edge_index' node indices
+    but knows nothing about the 'candidate_edges'. Overriding __inc__
+    and __cat_dim__ teaches PyG to treat candidate_edges exactly like
+    edge_index during batching.
+    """
+    def __inc__(self, key, value, *args, **kwargs):
+        if key == 'candidate_edges':
+            return self.num_nodes
+        return super().__inc__(key, value, *args, **kwargs)
+
+    def __cat_dim__(self, key, value, *args, **kwargs):
+        if key == 'candidate_edges':
+            return 1    # concat along the edge dim (not node)
+        return super().__cat_dim__(key, value, *args, **kwargs)
+
+def make_graph(coords: np.ndarray, d: int = 2, tol: float = 1e-9,) -> RigidityData:
+    """
+    Convert raw coordinates into a RigidityData graph.
+
+    Args:
+        coords : (n, d) float64 array of node positions
+        d      : embedding dimension
+        tol    : rank tolerance passed to grid.simulate
+
+    Returns:
+        RigidityData with:
+            x               (n, d)   node coordinates as float32
+            edge_index      (2, 2E)  both directions, for message passing
+            candidate_edges (2, E)   upper-triangular only, for scoring
+            y               (E,)     binary labels from Rust simulation
+    """
+    coords = np.asarray(coords, dtype=np.float64)
+    n = coords.shape[0]
+
+    result = grid.simulate(coords, d=d, tol=tol)
+
+    x = torch.tensor(coords, dtype=torch.float32)
+
+    src = torch.tensor([i for i in range(n) for j in range(i + 1, n)], dtype=torch.long)
+    dst = torch.tensor([j for i in range(n) for j in range(i + 1, n)], dtype=torch.long)
+    candidate_edges = torch.stack([src, dst], dim=0)
+
+    edge_index = torch.cat([candidate_edges, candidate_edges[[1, 0]]], dim=1)
+
+    y = torch.tensor(result["labels"], dtype=torch.float32)
+
+    return RigidityData(
+        x=x,
+        edge_index=edge_index,
+        candidate_edges=candidate_edges,
+        y=y,
+        num_nodes=n,
+    )
+
+class RigidityDataset:
+    """
+    Generates some random rigidity graphs and caches them in memory.
+    """
+    def __init__(self, n_graphs: int, n_nodes: int, d: int = 2, seed: int = 0):
+        self.graphs = []
+        rng = np.random.RandomState(seed)
+
+        print(f"Generating {n_graphs} graphs (n={n_nodes}, d={d}")
+        for _ in range(n_graphs):
+            coords = rng.uniform(0, 10, (n_nodes, d))
+            self.graphs.append(make_graph(coords, d=d))
+
+        print(f"Done, Labels per graph: {len(self.graphs[0].y)}")
+
+    def __len__(self):
+        return len(self.graphs)
+
+    def __getitem__(self, idx):
+        return self.graphs[idx]

grid_simulate-0.2.2/gnn/model.py

@@ -0,0 +1,117 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_geometric.nn import MessagePassing
+
+class RigidityConv(MessagePassing):
+    """
+    One graph convolution layer.
+
+    The rigidity matrix row for edge (i, j) is built from this exact difference
+    vector. Learning a function over it, instead of using it directly the MLP
+    decides which geometric differences matter.
+
+    Edges contribute independently to a node's constraint count. Sum preserves
+    the total geometric information received.
+    """
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__(aggr='sum')
+
+        # Maps distances to message vector
+        self.msg_mlp = nn.Sequential(
+            nn.Linear(in_channels, out_channels),
+            nn.ReLU(),
+            nn.Linear(out_channels, out_channels),
+        )
+
+        # Combines current embedding with aggregated messages
+        self.upd_mlp = nn.Sequential(
+            nn.Linear(out_channels, out_channels),
+            nn.ReLU(),
+            nn.Linear(out_channels, out_channels),
+        )
+
+        self.norm = nn.LayerNorm(out_channels)
+
+    def forward(self, x, edge_index):
+        agg = self.propagate(edge_index, x=x)    # residual is here
+        return self.norm(x + agg)
+
+    def message(self, x_i, x_j):
+        # This IS the rigidity row
+        return self.msg_mlp(x_i - x_j)
+
+    def update(self, aggr_out):
+        return self.upd_mlp(aggr_out)
+
+class RigidityGNN(nn.Module):
+    """
+    Full model: coordinate encoder → GNN layers → edge scorer.
+
+    Flow:
+        1. Project raw coordinates (n, d) → (n, hidden_dim)
+        2. k rounds of message passing — builds context-aware node embeddings
+           After k=3 layers, each node has seen its 3-hop geometric neighbourhood
+        3. For each candidate edge (i,j), score using [h_i || h_j || x_i-x_j]
+
+        Layer 1: node sees direct neighbours
+        Layer 2: node sees triangles and quadrilaterals
+        Layer 3: node has global rigidity context
+        More layers risk over-smoothing (all embeddings become similar)
+
+    Why hidden_dim=64:
+        Enough capacity to represent rigidity patterns, small enough
+        to train quickly on CPU.
+    """
+    def __init__(self, coord_dim: int = 2, hidden_dim: int = 64, n_layers: int = 3):
+        super().__init__()
+        self.coord_dim = coord_dim
+        self.hidden_dim = hidden_dim
+
+        # Initial projection: raw (x,y) is too low-dimensional for meaningful
+        self.input_proj = nn.Sequential(
+            nn.Linear(coord_dim, hidden_dim),
+            nn.ReLU(),
+        )
+
+        self.convs = nn.ModuleList([
+            RigidityConv(hidden_dim, hidden_dim)
+            for _ in range(n_layers)
+        ])
+
+        # Edge scorer
+        # Input: [h_i (hidden) || h_j (hidden) || x_i-x_j (coord_dim)]
+        self.edge_scorer = nn.Sequential(
+            nn.Linear(2 * hidden_dim + coord_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.ReLU(),
+            nn.Linear(hidden_dim // 2, 1),
+        )
+
+    def encode(self, x: torch.Tensor, edge_index: torch.Tensor) -> torch.Tensor:
+        h = self.input_proj(x)
+        for conv in self.convs:
+            h = h + conv(h, edge_index)  # residual: preserves input signal
+        return h
+
+    def score_edges(
+        self,
+        h: torch.Tensor,
+        x: torch.Tensor,
+        candidate_edges: torch.Tensor,
+    ) -> torch.Tensor:
+        """Score each candidate edge. Returns raw logits (E,)."""
+        src, dst = candidate_edges[0], candidate_edges[1]
+        feat = torch.cat([h[src], h[dst], x[src] - x[dst]], dim=-1)
+        return self.edge_scorer(feat).squeeze(-1)
+
+    def forward(self, x, edge_index, candidate_edges):
+        h = self.encode(x, edge_index)
+        return self.score_edges(h, x, candidate_edges)
+
+    def predict(self, x, edge_index, candidate_edges, threshold: float = 0.5):
+        """Inference: returns binary predictions (0 or 1)."""
+        with torch.no_grad():
+            logits = self.forward(x, edge_index, candidate_edges)
+            return (torch.sigmoid(logits) >= threshold).long()

grid_simulate-0.2.2/gnn/train.py

@@ -0,0 +1,100 @@
+import torch
+import torch.nn.functional as F
+from torch_geometric.loader import DataLoader
+
+from dataset import RigidityDataset
+from model import RigidityGNN
+
+
+def train_epoch(model, loader, optimizer, device):
+    model.train()
+    total_loss, correct, total = 0.0, 0, 0
+
+    for batch in loader:
+        batch = batch.to(device)
+        optimizer.zero_grad()
+
+        logits = model(batch.x, batch.edge_index, batch.candidate_edges)
+        loss = F.binary_cross_entropy_with_logits(logits, batch.y)
+
+        loss.backward()
+        optimizer.step()
+
+        n = batch.y.size(0)
+        total_loss += loss.item() * n
+        correct += ((torch.sigmoid(logits) >= 0.5) == batch.y.bool()).sum().item()
+        total += n
+
+    return total_loss / total, correct / total
+
+
+@torch.no_grad()
+def evaluate(model, loader, device):
+    model.eval()
+    total_loss, correct, total = 0.0, 0, 0
+
+    for batch in loader:
+        batch = batch.to(device)
+        logits = model(batch.x, batch.edge_index, batch.candidate_edges)
+        loss = F.binary_cross_entropy_with_logits(logits, batch.y)
+
+        n = batch.y.size(0)
+        total_loss += loss.item() * n
+        correct += ((torch.sigmoid(logits) >= 0.5) == batch.y.bool()).sum().item()
+        total += n
+
+    return total_loss / total, correct / total
+
+
+def train(
+    n_train: int   = 2000,
+    n_val: int     = 400,
+    n_nodes: int   = 10,
+    d: int         = 2,
+    hidden_dim: int = 64,
+    n_layers: int  = 3,
+    lr: float      = 1e-3,
+    epochs: int    = 100,
+    batch_size: int = 32,
+    seed: int      = 0,
+):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Device: {device}\n")
+
+    train_set = RigidityDataset(n_train, n_nodes, d=d, seed=seed)
+    val_set   = RigidityDataset(n_val,   n_nodes, d=d, seed=seed + 99999)
+
+    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
+    val_loader   = DataLoader(val_set,   batch_size=batch_size)
+
+    model = RigidityGNN(coord_dim=d, hidden_dim=hidden_dim, n_layers=n_layers).to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+        optimizer, mode='min', patience=5, factor=0.5, min_lr=1e-5
+    )
+
+    n_params = sum(p.numel() for p in model.parameters())
+    print(f"Model: {n_params:,} parameters")
+    print(f"{'Epoch':>6}  {'Train Loss':>10}  {'Train Acc':>10}  {'Val Loss':>10}  {'Val Acc':>10}")
+    print("-" * 55)
+
+    best_val_acc = 0.0
+    for epoch in range(1, epochs + 1):
+        tr_loss, tr_acc = train_epoch(model, train_loader, optimizer, device)
+        va_loss, va_acc = evaluate(model, val_loader, device)
+        scheduler.step(va_loss)
+
+        if va_acc > best_val_acc:
+            best_val_acc = va_acc
+            torch.save(model.state_dict(), "best_model.pt")
+
+        if epoch % 10 == 0:
+            print(f"{epoch:>6}  {tr_loss:>10.4f}  {tr_acc:>10.4f}  {va_loss:>10.4f}  {va_acc:>10.4f}")
+
+    print(f"\nBest validation accuracy: {best_val_acc:.4f}")
+    print("Saved to best_model.pt")
+    return model
+
+
+if __name__ == "__main__":
+    train()

{grid_simulate-0.1.0 → grid_simulate-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "grid-simulate"
-version = "0.1.0"
+version = "0.2.2"
 description = "Rust + PyO3 simulation library"
 requires-python = ">=3.8"
 

{grid_simulate-0.1.0 → grid_simulate-0.2.2}/python/test_bridge.py

@@ -67,7 +67,7 @@ print(result["distance_matrix"])
 ################################################################################
 ### grid.generate_training_batch ###
 """
-Generate many random graphs at once for GNN training data
+Generate many random graphs at once for gnn training data
 All heavy work runs with GIL released so torch can train in parallel
 
 Args:
@@ -80,7 +80,7 @@ Returns:
     list[dict]    (same dict as simulate(), one per graph)
 """
 
-banner("grid.generate_training_batch  -  batched random graphs for GNN data")
+banner("grid.generate_training_batch  -  batched random graphs for gnn data")
 
 batch = grid.generate_training_batch(n_graphs=100, n_nodes=20, d=2, seed=42)
 

{grid_simulate-0.1.0 → grid_simulate-0.2.2}/src/lib.rs

@@ -22,6 +22,7 @@
 
 mod node;
 mod rank;
+mod pebble;
 mod reconstruct;
 mod rigidity;
 mod simulation;
@@ -380,6 +381,47 @@ fn reconstruct_coords(
     dict.set_item("is_isometric", result.is_isometric)?;
     Ok(dict.into())
 }
+
+/// Run the Pebble Game algorithm — O(n^2), deterministic, no matrix rank needed.
+///
+/// Parameters
+/// ----------
+/// coords : np.ndarray shape (n, d) — node coordinates (only d=2 supported)
+/// d      : int — must be 2
+///
+/// Returns dict with keys:
+///   "labels"          np.ndarray[u8]  1=independent, 0=redundant
+///   "accepted_edges"  list[(i,j)]
+///   "n_nodes"         int
+///   "d"               int
+#[pyfunction]
+fn pebble_game(py: Python, coords: PyReadonlyArray2<f64>, d: usize) -> PyResult<PyObject> {
+    let arr = coords.as_array();
+    let n = arr.shape()[0];
+
+    if d != 2 {
+        return Err(pyo3::exceptions::PyValueError::new_err(
+            "pebble_game only supports d=2",
+        ));
+    }
+
+    let edges: Vec<(usize, usize)> = (0..n)
+        .flat_map(|i| (i + 1..n).map(move |j| (i, j)))
+        .collect();
+
+    let (labels, accepted) = py.allow_threads(|| pebble::pebble_game(n, &edges, d));
+
+    let dict = PyDict::new_bound(py);
+    let labels_arr: Array1<u8> = Array1::from(labels);
+    dict.set_item("labels", labels_arr.into_pyarray_bound(py))?;
+    let edge_list = PyList::new_bound(py, accepted.iter().map(|&(i, j)| (i, j)));
+    dict.set_item("accepted_edges", edge_list)?;
+    dict.set_item("n_nodes", n)?;
+    dict.set_item("d", d)?;
+
+    Ok(dict.into())
+}
+
 // ── Module registration ────────────────────────────────────────────────────────
 
 #[pymodule]
@@ -389,5 +431,6 @@ fn grid(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(generate_training_batch, m)?)?;
     m.add_function(wrap_pyfunction!(max_independent, m)?)?;
     m.add_function(wrap_pyfunction!(reconstruct_coords, m)?)?;
+    m.add_function(wrap_pyfunction!(pebble_game, m)?)?;
     Ok(())
 }

grid_simulate-0.2.2/src/pebble.rs

@@ -0,0 +1,167 @@
+//! Pebble Game Algorithm for 2D rigidity independence
+//! Deterministic and O(n^2)
+
+use std::collections::HashSet;
+
+/// A recorded state change, so a failed edge attempt can be undone.
+enum Change {
+    // A pebble was consumed at this vertex
+    Pebble(usize),
+    // A pebble was gained at this vertex during the collection loop
+    GainPebble(usize),
+    // A directed edge was reversed: was from_ -> to, is now to -> from_.
+    Edge(usize, usize),
+}
+
+/// DFS: try to find a free pebble reachable from v by following directed edges.
+/// Records the changes made during the DFS search so they can be rolled back if necessary.
+fn find_pebble(
+    v: usize,
+    visited: &mut HashSet<usize>,
+    pebbles: &mut [i32],
+    out_edges: &mut Vec<Vec<usize>>,
+    changes: &mut Vec<Change>,
+) -> bool {
+    if visited.contains(&v) {
+        return false;
+    }
+    visited.insert(v);
+
+    if pebbles[v] > 0 {
+        pebbles[v] -= 1;
+        changes.push(Change::Pebble(v));
+        return true;
+    }
+
+    for i in 0..out_edges[v].len() {
+        let u = out_edges[v][i];
+        if find_pebble(u, visited, pebbles, out_edges, changes) {
+            out_edges[v].remove(i);
+            out_edges[u].push(v);
+            changes.push(Change::Edge(v, u));
+            return true;
+        }
+    }
+
+    false
+}
+
+/// Undo all state changes during a failed edge attempt.
+fn rollback(changes: Vec<Change>, pebbles: &mut [i32], out_edges: &mut Vec<Vec<usize>>) {
+    for change in changes.into_iter().rev() {
+        match change {
+            Change::Pebble(v) => pebbles[v] += 1,
+            Change::GainPebble(v) => pebbles[v] -= 1,
+            Change::Edge(v, u) => {
+                if let Some(pos) = out_edges[u].iter().position(|&x| x == v) {
+                    out_edges[u].remove(pos);
+                }
+                out_edges[v].push(u);
+            }
+        }
+    }
+}
+
+/// Pebble Game for 2D rigidity independence
+pub fn pebble_game(
+    n: usize,
+    candidate_edges: &[(usize, usize)],
+    d: usize,
+) -> (Vec<u8>, Vec<(usize, usize)>) {
+    assert_eq!(d, 2, "dimension must be 2");
+
+    let mut pebbles = vec![2i32; n];
+    let mut out_edges: Vec<Vec<usize>> = vec![Vec::new(); n];
+    let mut labels = Vec::with_capacity(candidate_edges.len());
+    let mut accepted = Vec::new();
+
+    for &(u, v) in candidate_edges {
+        let mut changes: Vec<Change> = Vec::new();
+
+        while pebbles[u] + pebbles[v] < 4 {
+            let (mut pu, mut pv) = (pebbles[u], pebbles[v]);
+            pebbles[u] = 0;
+            pebbles[v] = 0;
+
+            let mut progressed = false;
+
+            if pu < 2 {
+                let mut visited = HashSet::new();
+                if find_pebble(u, &mut visited, &mut pebbles, &mut out_edges, &mut changes) {
+                    pu += 1;
+                    changes.push(Change::GainPebble(u));
+                    progressed = true;
+                }
+            }
+            if !progressed {
+                let mut visited = HashSet::new();
+                if find_pebble(v, &mut visited, &mut pebbles, &mut out_edges, &mut changes) {
+                    pv += 1;
+                    changes.push(Change::GainPebble(v));
+                    progressed = true;
+                }
+            }
+
+            pebbles[u] = pu;
+            pebbles[v] = pv;
+
+            if !progressed {
+                break;
+            }
+        }
+
+        if pebbles[u] + pebbles[v] >= 4 {
+            pebbles[v] -= 1;
+            out_edges[v].push(u);
+            labels.push(1u8);
+            accepted.push((u, v));
+        } else {
+            rollback(changes, &mut pebbles, &mut out_edges);
+            labels.push(0u8);
+        }
+    }
+
+    (labels, accepted)
+}
+
+/// Testing
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn upper_triangular_edges(n: usize) -> Vec<(usize, usize)> {
+        (0..n).flat_map(|i| (i + 1..n).map(move |j| (i, j))).collect()
+    }
+
+    #[test]
+    fn test_triangle_all_independent() {
+        let edges = upper_triangular_edges(3);
+        let (labels, accepted) = pebble_game(3, &edges, 2);
+        assert_eq!(labels, vec![1, 1, 1]);
+        assert_eq!(accepted.len(), 3);
+    }
+
+    #[test]
+    fn test_square_one_redundant() {
+        let edges = upper_triangular_edges(4);
+        let (labels, accepted) = pebble_game(4, &edges, 2);
+        assert_eq!(accepted.len(), 5);
+        assert_eq!(labels.iter().filter(|&&l| l == 1).count(), 5);
+        assert_eq!(labels.iter().filter(|&&l| l == 0).count(), 1);
+    }
+
+    #[test]
+    fn test_five_nodes_matches_max_rank() {
+        let edges = upper_triangular_edges(5);
+        let (labels, accepted) = pebble_game(5, &edges, 2);
+        assert_eq!(accepted.len(), 7);
+        assert_eq!(labels.iter().filter(|&&l| l == 1).count(), 7);
+    }
+
+    #[test]
+    fn test_labels_length_matches_input() {
+        let edges = upper_triangular_edges(10);
+        let (labels, _) = pebble_game(10, &edges, 2);
+        assert_eq!(labels.len(), edges.len());
+    }
+}

grid_simulate-0.1.0/.github/workflows/publish.yml

@@ -1,48 +0,0 @@
-name: Publish to PyPI
-
-on:
-  release:
-    types: [published]
-
-permissions:
-  contents: read
-  id-token: write
-
-jobs:
-  build:
-    runs-on: ubuntu-latest
-
-    steps:
-      - uses: actions/checkout@v4
-
-      - name: Build wheels with maturin
-        uses: PyO3/maturin-action@v1
-        with:
-          command: build
-          args: --release --sdist
-
-      - name: Upload artifacts
-        uses: actions/upload-artifact@v4
-        with:
-          name: dist
-          path: target/wheels/
-
-  publish:
-    needs: build
-    runs-on: ubuntu-latest
-
-    permissions:
-      id-token: write
-
-    environment:
-      name: pypi
-
-    steps:
-      - name: Download artifacts
-        uses: actions/download-artifact@v4
-        with:
-          name: dist
-          path: dist/
-
-      - name: Publish to PyPI
-        uses: pypa/gh-action-pypi-publish@release/v1

grid_simulate-0.1.0/.gitignore

@@ -1 +0,0 @@
-/target