Performance: add C++ backend for `pdist_tangency_grad` (batched pairwise tangency VJP)

[koki3070]

Summary

When training with ellphi.grad.pdist_tangency_grad, the current implementation loops over all ellipse pairs in pure Python, calling tangency_grad (and solve_mu_gradients) once per pair. This makes gradient-based workflows impractical at modest point counts.

Forward-only pdist_tangency is already fast via the existing C++ backend, but the differentiable backward path remains the bottleneck.

Proposal: add a C++ kernel pdist_tangency_grad that computes, in a single call, the condensed pairwise distances plus per-pair coefficient gradients dt_dp / dt_dq using the same math as the existing Python reference. Keep the public API and VJP semantics unchanged; retain the Python implementation as a fallback when the C++ extension is unavailable.

We have a working prototype on a local branch (feature/pdist-tangency-grad-cpp, commit 9fcad0c) integrated in our downstream project TDA-ML (vendored ellphi_repo). If this direction looks good to maintainers, we are happy to push the branch to a fork and open a PR upstream.

---

Motivation / use case

In gradient-based topology optimization (e.g. learning ellipse parameters so that an anisotropic distance matrix matches a target persistence diagram), each training step needs:

1. coef_from_cov_grad — centres/covariances → conic coefficients
2. pdist_tangency_grad — all pairwise tangency distances + VJP pullback
3. downstream TDA loss (Vietoris–Rips + Wasserstein, etc.)

For a single point cloud with N = 120 ellipses, step (2) alone involves C(120, 2) = 7,140 pairs per cloud, repeated over batches and epochs.

---

Current behavior

ellphi.grad._pdist_tangency_grad_python iterates over all pairs in Python:

for k, (i, j) in enumerate(pairs):
    g = tangency_grad(coefs[i], coefs[j])
    dists[k] = g.t
    store.append((i, j, g.dt_dp, g.dt_dq))

Profiling on our workload (N=120, CPU, one cloud, forward+VJP path):

Component	Time (approx.)
Python pdist_tangency_grad (7140 × tangency_grad / solve_mu_gradients)	~2027 ms
— of which solve_mu_gradients	~82%

For the same N, forward-only C++ pdist_tangency is already ~10 ms.

With the C++ kernel, the same forward+VJP path is ~11.6 ms (~174× faster vs. the Python path above).

In our downstream training setup (30 epochs, N=120 topo loss, CPU), wall-clock time for ellphi backend training went from on the order of ~3 days with the Python VJP path (estimated from per-step profiling before the kernel existed) to ~2 hours measured with the C++ kernel (same hyperparameters; numerics unchanged).

---

Proposed solution

Add pdist_tangency_grad to the existing C++ tangency backend (_tangency_cpp_impl.cpp), exposed through _tangency_cpp.py via ctypes, mirroring the existing pdist_tangency / tangency bindings.

C++ kernel (one call):

• Input: coefs with shape (N, m)
• Output:
  • dists: condensed distances, shape (N*(N-1)//2,), same order as scipy.spatial.distance.pdist
  • dt_dp, dt_dq: shape (n_pairs, m) — ∂(pair distance)/∂(first/second ellipse coefficients)

Python layer (grad.py):

• If C++ is available and exports pdist_tangency_grad → compute distances and dt_dp/dt_dq in C++; keep the existing NumPy VJP accumulation (np.add.at) unchanged.
• Otherwise → keep _pdist_tangency_grad_python as the reference fallback.

API: no change to pdist_tangency_grad(coefs) -> (dists, vjp) signature or semantics.

---

Correctness

• Distances match pdist_tangency / Python pdist_tangency_grad (covered by existing tests in tests/test_grad.py).
• VJP matches finite-difference checks (test_pdist_tangency_grad_vjp).
• On our prototype branch, max absolute error vs. the Python reference is ≤ ~8e-12 on representative random clouds.

---

Why C++ (not “Python cannot do the math”)

Python can compute the same values; _pdist_tangency_grad_python is the reference. The issue is throughput: thousands of Python function calls per cloud per backward pass, with per-pair implicit differentiation in Python. Running the same loop in compiled code yields a large constant-factor speedup without changing the mathematical definition of tangency distance.

---

Prototype / next steps

Item	Detail
Branch	feature/pdist-tangency-grad-cpp
Commit	9fcad0c (local prototype, integrated in TDA-ML)
Size	~280 lines (C++ kernel + Python bindings)
Downstream	TDA-ML uses a path dependency on this fork until upstream merge

If maintainers are interested: we will push the branch to a public fork and open a PR, update release notes / differentiable guide, and ensure CI covers the new symbol when the C++ extension is built.

---

Questions for maintainers

1. Is a batched pdist_tangency_grad C++ entry point aligned with performance goals for ellphi.grad?
2. Any preference on naming / export surface (same as pdist_tangency, or behind a separate build flag)?
3. Any API-stability concerns before we submit a PR?

Thank you for maintaining ellphi — the grad module enabled our TDA training pipeline; this change would make ellphi-scale CPU workloads practical.

---