refactor(solvers): consolidate local solver API, drop numba, prune dead code

.github/workflows/ci.yml

@@ -0,0 +1,29 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.11", "3.12"]
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: pip
+
+      - name: Install package
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e .[dev]
+
+      - name: Run tests
+        run: pytest -q

.gitignore

@@ -137,5 +137,8 @@ dmypy.json
 # Local-only planning/design notes
 docs/superpowers/
 
+# Local Claude/Codex-adjacent session state
+.claude/
+
 # Per-environment solver baselines (regenerate with tests/baselines/generate.py)
 tests/baselines/*.json

README.md

@@ -23,14 +23,6 @@ pip install -r requirements-macos-arm64.txt
 pip install -e . --no-deps
 ```
 
-**Optional: numba AOT compilation.** GraphIK's Riemannian solver has an AOT-compiled cost-gradient kernel for speed. The default solver path runs without it (`use_jit=False`); for the JIT path, build the extension once after install:
-
-```bash
-cd graphik/solvers && python costs.py
-```
-
-This produces a per-platform `.so` that's gitignored and env-local — regenerate it after each environment recreation.
-
 ## SDP solvers (Mosek recommended)
 
 GraphIK's SDP-relaxation solvers (`solve_with_cidgik` in `graphik/solvers/convex_iteration.py`, the SDP formulations in `graphik/solvers/sdp_*.py`) run out of the box on a free solver (Clarabel by default, falling back to SCS or CVXOPT) bundled with `cvxpy`. `experiments/cidgik_example.py` works without any extra setup.
@@ -48,7 +40,7 @@ To enable the Mosek path:
 
 3. Place the license file (`mosek.lic`) at the path Mosek expects (typically `~/mosek/mosek.lic`).
 
-The SDP tests in `tests/test_sdp_snl*.py` are tuned to Mosek's tolerances and skip cleanly when it isn't installed.
+The SDP tests in `tests/test_sdp_snl*.py` run without Mosek: constraint-construction tests are solver-free, and the end-to-end solves use whichever free solver cvxpy picks, with tolerances loose enough for it.
 
 ## Usage
 Use of GraphIK can be summarized by four key steps, which we'll walk through below (see the scripts in [experiments/](https://github.com/utiasSTARS/GraphIK/tree/main/experiments) for more details).
@@ -66,7 +58,7 @@ GraphIK's interface between robot models and IK solvers is the [`ProblemGraph`](
 If you are considering an environment with spherical obstacles, you can include constraints that prevent collisions. In this example, we will use a set of spheres that approximate a table: 
 
 ```python
-from graphik.utils.utils import table_environment
+from graphik.utils.environments import table_environment
 obstacles = table_environment()
 # This loop is not needed if you are not using obstacle avoidance constraints 
 for idx, obs in enumerate(obstacles):
@@ -82,11 +74,14 @@ T_goal = robot.pose(q_goal, f"p{robot.n}")
 ```
 
 ### 4. Solve the IK Problem
-The main purpose of our graphical interpretation of robot kinematics is to develop distance-geometric IK solvers. One example is the [Riemannian optimization-based solver](https://arxiv.org/abs/2011.04850) implemented in [`RiemannianSolver`](https://github.com/utiasSTARS/GraphIK/blob/main/graphik/solvers/riemannian_solver.py). 
+The main purpose of our graphical interpretation of robot kinematics is to develop distance-geometric IK solvers. One example is the [Riemannian optimization-based solver](https://arxiv.org/abs/2011.04850) implemented in [`RiemannianSolver`](https://github.com/utiasSTARS/GraphIK/blob/main/graphik/solvers/riemannian.py). 
 
 ```python
-from graphik.solvers.riemannian_solver import solve_with_riemannian
-q_sol, solution_points = solve_with_riemannian(graph, T_goal, use_jit=False)  # Returns None if infeasible or didn't solve
+from graphik.solvers import RiemannianSolver
+
+solver = RiemannianSolver(graph)
+result = solver.solve(T_goal)
+q_sol = result.q if result.feasible else None  # feasible == within joint limits
 ```
 
 For a similar example using [`CIDGIK`](https://arxiv.org/abs/2109.03374), a convex optimization-based approach, please see [experiments/cidgik_example.py](https://github.com/utiasSTARS/GraphIK/blob/main/experiments/cidgik_example.py).
@@ -176,7 +171,7 @@ arXiv: [Inverse Kinematics for Serial Kinematic Chains via Sum of Squares Optimi
 
 ```bibtex
 @misc{marić2022convex_arxiv,
-  author={Filip Marić and {Matthew Giamou and Soroush Khoubyarian and Ivan Petrović and Jonathan Kelly},
+  author={Filip Marić and Matthew Giamou and Soroush Khoubyarian and Ivan Petrović and Jonathan Kelly},
   title={Inverse Kinematics for Serial Kinematic Chains via Sum of Squares Optimization}, 
   year={2020},
   eprint={1909.09318},

experiments/cidgik_example.py

@@ -1,4 +1,4 @@
-from graphik.utils.utils import table_environment
+from graphik.utils.environments import table_environment
 from graphik.solvers.convex_iteration import solve_with_cidgik
 
 # Multiple robot models to try out, or you can implement your own

experiments/riemannian_example.py

@@ -1,7 +1,7 @@
 from time import perf_counter
 
-from graphik.utils.utils import table_environment
-from graphik.solvers.riemannian_solver import solve_with_riemannian
+from graphik.utils.environments import table_environment
+from graphik.solvers import RiemannianSolver
 
 # Multiple robot models to try out, or you can implement your own
 from graphik.utils.roboturdf import load_ur10
@@ -22,9 +22,10 @@
     q_goal = robot.random_configuration()
     T_goal = robot.pose(q_goal, f"p{robot.n}")  # Can be any desired pose, this is just a simple example
 
-    # Run Riemannian solver without jit compiled cost function and gradient computations
+    # Run the Riemannian solver
+    solver = RiemannianSolver(graph)
     t0 = perf_counter()
-    q_sol, solution_points = solve_with_riemannian(graph, T_goal, use_jit=False)  # Returns None if infeasible or didn't solve
+    result = solver.solve(T_goal)
     solve_time = perf_counter() - t0
 
     # Compare the solution's end effector pose to the goal.
@@ -36,10 +37,10 @@
     print(q_goal)
     print("--------------------------------------------")
     print(f"Solve time: {solve_time:.4f} s")
-    if q_sol:
+    if result.feasible:
         print("Riemannian solution's pose: ")
-        print(robot.pose(q_sol, f"p{robot.n}"))
+        print(robot.pose(result.q, f"p{robot.n}"))
         print("Riemannian configuration: ")
-        print(q_sol)
+        print(result.q)
     else:
         print("Riemannian did not return a feasible solution.")

experiments/rtr_preconditioner_study.py

@@ -0,0 +1,70 @@
+"""Preconditioner study for the Riemannian solver's truncated CG."""
+import time
+
+import numpy as np
+
+from graphik.solvers import RiemannianSolver
+import graphik.solvers.costs as costs_mod
+from graphik.utils.roboturdf import load_ur10
+
+
+def run_case(graph, T_goal, precon):
+    solver = RiemannianSolver(graph, precon=precon)
+
+    n_hvp = [0]
+    saved = costs_mod.for_riemannian
+
+    def counting_for_riemannian(*args, **kwargs):
+        cost, egrad, ehvp = saved(*args, **kwargs)
+
+        def counted_ehvp(Y, Z):
+            n_hvp[0] += 1
+            return ehvp(Y, Z)
+
+        return cost, egrad, counted_ehvp
+
+    costs_mod.for_riemannian = counting_for_riemannian
+    try:
+        t0 = time.perf_counter()
+        res = solver.solve(T_goal)
+        wall = time.perf_counter() - t0
+    finally:
+        costs_mod.for_riemannian = saved
+
+    T_sol = graph.robot.pose(res.q, f"p{graph.robot.n}")
+    trans_err = np.linalg.norm(T_sol[:3, 3] - T_goal[:3, 3])
+    rot_err = np.linalg.norm(T_sol[:3, :3] - T_goal[:3, :3])
+    return {
+        "iters": res.iterations,
+        "hvps": n_hvp[0],
+        "time": wall,
+        "trans_err": trans_err,
+        "rot_err": rot_err,
+        "success": trans_err < 1e-2 and rot_err < 1e-2 and res.feasible,
+    }
+
+
+def main(n_goals=20, seed=0):
+    np.random.seed(seed)
+    robot, graph = load_ur10()
+    goals = []
+    for _ in range(n_goals):
+        q = robot.random_configuration()
+        goals.append(robot.pose(q, f"p{robot.n}"))
+
+    for precon in (None, "jacobi", "gn"):
+        rows = [run_case(graph, T, precon) for T in goals]
+        succ = sum(r["success"] for r in rows)
+        med = lambda k: np.median([r[k] for r in rows])
+        tot = lambda k: np.sum([r[k] for r in rows])
+        print(
+            f"{str(precon):>8}: success {succ}/{n_goals} | "
+            f"median outer {med('iters'):.0f} | median HVPs {med('hvps'):.0f} | "
+            f"HVPs/outer {tot('hvps')/tot('iters'):.1f} | "
+            f"median time {med('time')*1e3:.0f} ms | "
+            f"median trans_err {med('trans_err'):.2e}"
+        )
+
+
+if __name__ == "__main__":
+    main()

experiments/rtr_profile.py

@@ -1,13 +1,7 @@
 """Profile the in-house RTR solver on a single robot.
 
-Runs N IK problems on ur10 (or another robot via --robot), reports per-
-pose wall-clock and convergence stats, and dumps a cProfile of the
-whole run sorted by cumulative time.
-
-Usage:
-    python experiments/rtr_profile.py [--n-poses 20] [--init spectral|bsmooth]
-                                       [--robot ur10] [--seed 0] [--top 20]
-                                       [--jit]
+Runs N IK problems, reports per-pose wall-clock/convergence stats, and dumps
+a cProfile of the whole run sorted by cumulative time.
 """
 from __future__ import annotations
 
@@ -19,20 +13,15 @@
 
 import numpy as np
 
-from graphik.utils import (
-    adjacency_matrix_from_graph,
-    bound_smoothing,
-    distance_matrix_from_graph,
-    graph_from_pos,
-)
+from graphik.solvers import RiemannianSolver
+from graphik.solvers.initializations import INIT_STRATEGIES
 from graphik.utils.roboturdf import (
     load_kuka,
     load_panda,
     load_schunk_lwa4d,
     load_schunk_lwa4p,
     load_ur10,
 )
-from graphik.solvers.riemannian_solver import RiemannianSolver
 
 
 ROBOTS = {
@@ -44,33 +33,32 @@
 }
 
 
-def _run_one(graph, T_goal, init: str, jit: bool, params: dict) -> dict:
-    G_partial = graph.from_pose(T_goal)
-    D_goal = distance_matrix_from_graph(G_partial)
-    omega = adjacency_matrix_from_graph(G_partial)
-    bounds = bound_smoothing(G_partial) if init in ("bsmooth", "bspectral") else None
-
+def _run_one(graph, T_goal, init: str, params: dict) -> dict:
+    solver = RiemannianSolver(graph, init=init, rtr_params=params)
     t0 = time.perf_counter()
-    solver = RiemannianSolver(graph, jit=jit, init=init)
-    solver.params.update(params)
-    sol = solver.solve(D_goal, omega, use_limits=True, bounds=bounds)
-    wall = time.perf_counter() - t0
+    try:
+        res = solver.solve(T_goal)
+    except Exception:
+        return {
+            "wall_s": time.perf_counter() - t0,
+            "iterations": 0,
+            "stopping_criterion": "solve/decode raised",
+            "fx": float("nan"),
+            "q_sol": None,
+            "solver_feasible": False,
+        }
     return {
-        "wall_s": wall,
-        "rtr_time_s": sol["time"],
-        "iterations": sol["iterations"],
-        "stopping_criterion": sol["stopping_criterion"],
-        "gradnorm": sol["gradnorm"],
-        "fx": sol["f(x)"],
-        "Y_sol": sol["x"],
+        "wall_s": time.perf_counter() - t0,
+        "iterations": res.iterations,
+        "stopping_criterion": res.status,
+        "fx": res.cost,
+        "q_sol": res.q,
+        "solver_feasible": res.feasible,
     }
 
 
-def _pose_error(graph, robot, ee: str, T_goal: np.ndarray, Y_sol: np.ndarray) -> tuple[bool, float, float]:
-    G_sol = graph_from_pos(Y_sol, graph.node_ids)
-    try:
-        q_sol = graph.joint_variables(G_sol, {ee: T_goal})
-    except Exception:
+def _pose_error(robot, ee: str, T_goal: np.ndarray, q_sol) -> tuple[bool, float, float]:
+    if q_sol is None:
         return False, float("nan"), float("nan")
     T_got = np.asarray(robot.pose(q_sol, ee))
     pos = float(np.linalg.norm(T_goal[:3, 3] - T_got[:3, 3]))
@@ -82,20 +70,23 @@ def _pose_error(graph, robot, ee: str, T_goal: np.ndarray, Y_sol: np.ndarray) ->
 
 def _print_summary(results: list[dict], total_wall: float, args) -> None:
     walls = np.array([r["wall_s"] for r in results])
-    rtr_times = np.array([r["rtr_time_s"] for r in results])
     iters = np.array([r["iterations"] for r in results])
     feas = np.array([r["feasible"] for r in results], dtype=bool)
 
     print(f"=== RTR profile on {args.robot} ===")
     print(f"  init:                    {args.init}")
-    print(f"  jit:                     {args.jit}")
     print(f"  n_poses:                 {args.n_poses}")
     print(f"  seed:                    {args.seed}")
     print(f"  total wall (incl prof):  {total_wall*1000:.0f}ms")
     print(f"  feasible:                {feas.sum()}/{args.n_poses}  ({feas.mean()*100:.1f}%)")
-    print(f"  iterations:              mean {iters.mean():.1f}  median {np.median(iters):.0f}  max {iters.max()}")
-    print(f"  per-pose wall:           mean {walls.mean()*1000:.1f}ms  median {np.median(walls)*1000:.1f}ms  max {walls.max()*1000:.1f}ms")
-    print(f"  rtr-internal time:       mean {rtr_times.mean()*1000:.1f}ms  (vs {walls.mean()*1000:.1f}ms wall — overhead {(walls.mean()-rtr_times.mean())*1000:.1f}ms)")
+    print(
+        f"  iterations:              mean {iters.mean():.1f}  "
+        f"median {np.median(iters):.0f}  max {iters.max()}"
+    )
+    print(
+        f"  per-pose wall:           mean {walls.mean()*1000:.1f}ms  "
+        f"median {np.median(walls)*1000:.1f}ms  max {walls.max()*1000:.1f}ms"
+    )
     if feas.any():
         pos = np.array([r["pos_err"] for r in results])[feas]
         rot = np.array([r["rot_err"] for r in results])[feas]
@@ -116,13 +107,11 @@ def _print_profile(profiler: cProfile.Profile, top: int) -> None:
 
 def main() -> None:
     p = argparse.ArgumentParser(description=__doc__)
-    p.add_argument("--n-poses", type=int, default=20, help="number of IK problems (default: 20)")
+    p.add_argument("--n-poses", type=int, default=20, help="number of IK problems")
     p.add_argument("--robot", choices=tuple(ROBOTS), default="ur10")
-    p.add_argument("--init", choices=("spectral", "bsmooth"), default="spectral")
+    p.add_argument("--init", choices=INIT_STRATEGIES, default="spectral")
     p.add_argument("--seed", type=int, default=0)
-    p.add_argument("--top", type=int, default=20, help="cProfile rows to print (default: 20)")
-    p.add_argument("--jit", action="store_true", help="use AOT-compiled costgrd kernels")
-    # Solver hyperparameter overrides (forwarded to RiemannianSolver.params).
+    p.add_argument("--top", type=int, default=20)
     p.add_argument("--mingradnorm", type=float, default=None)
     p.add_argument("--maxiter", type=int, default=None)
     p.add_argument("--theta", type=float, default=None)
@@ -146,31 +135,33 @@ def main() -> None:
         ("delta_bar", "Delta_bar"),
         ("delta0", "Delta0"),
     ]:
-        v = getattr(args, cli_name)
-        if v is not None:
-            params[key] = v
+        value = getattr(args, cli_name)
+        if value is not None:
+            params[key] = value
 
     np.random.seed(args.seed)
     robot, graph = ROBOTS[args.robot]()
     ee = f"p{robot.n}"
-    poses = [np.asarray(robot.pose(robot.random_configuration(), ee)) for _ in range(args.n_poses)]
+    poses = [
+        np.asarray(robot.pose(robot.random_configuration(), ee))
+        for _ in range(args.n_poses)
+    ]
 
     profiler = cProfile.Profile()
     results: list[dict] = []
     t0 = time.perf_counter()
     profiler.enable()
     try:
         for T_goal in poses:
-            results.append(_run_one(graph, T_goal, args.init, args.jit, params))
+            results.append(_run_one(graph, T_goal, args.init, params))
     finally:
         profiler.disable()
     total_wall = time.perf_counter() - t0
 
-    # Feasibility check is outside the profile so it doesn't pollute solver
-    # function timings; it's a tiny pose-FK roundtrip per pose.
+    # Decode is inside solve; this loop only computes pose-FK error.
     for r, T_goal in zip(results, poses):
-        ok, pos_err, rot_err = _pose_error(graph, robot, ee, T_goal, r["Y_sol"])
-        r["feasible"] = ok
+        ok, pos_err, rot_err = _pose_error(robot, ee, T_goal, r["q_sol"])
+        r["feasible"] = r["solver_feasible"] and ok
         r["pos_err"] = pos_err
         r["rot_err"] = rot_err