feat(adr-185): P1 AETHER bindings (wifi_densepose.aether) + parity harness

Bind the ADR-024 contrastive CSI embedding surface into the wheel behind
a gated [aether] extra / Cargo `aether` feature, per ADR-185 section 3.2.

Surface (bound against the REAL code at HEAD, not the ADR wishlist):
- AetherConfig       -> EmbeddingConfig {d_model,d_proj,temperature,normalize}
- CsiAugmenter       -> augment_pair(window, seed)
- EmbeddingExtractor -> embed(csi): 128-dim L2-normed, GIL-released
- info_nce_loss, cosine_similarity (module fns)

ADR-185 section 3.2 also names aether_loss/VICReg components,
alignment_metric, uniformity_metric, forward_dual, and vicreg_* config
fields. None exist in embedding.rs at HEAD, so they are intentionally NOT
bound (a Rust-side gap, not a binding gap) rather than fabricated.

Parity (section 4.1, release-blocking): committed fixture
aether_input.json -> native Rust reference (tests/aether_parity.rs, calls
sensing-server EmbeddingExtractor directly) locks
tests/golden/aether_embedding.sha256; pytest (tests/test_aether.py) runs
the same fixture through the binding and asserts the identical SHA-256 of
the LE f32 bytes. Both pass.

Verified: cargo test --features aether --test aether_parity -> 2/2 pass;
maturin develop --features aether + pytest tests/test_aether.py -> 9/9
pass; default cargo build clean with 0 sensing-server refs in the dep
graph (gate keeps the base wheel lean).

HONEST WHEEL-SIZE FINDING (ADR-185 section 9 / Open Q 11.1, unresolved):
wifi-densepose-sensing-server is an Axum/tokio crate whose tokio/axum/
worldgraph/ruvector deps are NON-optional (only mqtt/matter are
features), so default-features=false does NOT drop them. An [aether]
wheel therefore links the full server tree and BREAKS the ADR-117 section
5.4 <=5 MB budget. The fix is the ADR-sanctioned hoist of embedding.rs
(+ its graph_transformer/sona siblings) into a leaf crate so the wheel
links pure compute only -- a change INSIDE wifi-densepose-sensing-server,
owned by another agent this session, so deferred as a required
pre-release follow-up. P1 binds real code and proves parity today; the
hoist is a wheel-size optimization, not a functional blocker.

Note: building required initializing the worldgraph/rufield/ruvector
submodules (absent in the fresh worktree).
This commit is contained in:
ruv
2026-07-21 16:28:14 -07:00
parent 5426fe830c
commit d060998e3b
11 changed files with 4276 additions and 26 deletions
+242
View File
@@ -0,0 +1,242 @@
//! ADR-185 P1 — PyO3 bindings for AETHER contrastive CSI embeddings.
//!
//! Surfaces the **pure-sync** contrastive-embedding compute from
//! `wifi-densepose-sensing-server::embedding` (ADR-024) into
//! `wifi_densepose.aether`:
//!
//! - `AetherConfig` — wraps `EmbeddingConfig` (d_model / d_proj /
//! temperature / normalize)
//! - `CsiAugmenter` — SimCLR-style augmentation pair generator
//! - `EmbeddingExtractor`— backbone + projection → 128-dim L2-normed embedding
//! - `info_nce_loss` — NT-Xent contrastive loss (module function)
//! - `cosine_similarity` — re-ID similarity helper (module function)
//!
//! ## Honest scope vs ADR-185 §3.2
//!
//! ADR-185 §3.2 names an aspirational surface (`aether_loss` returning
//! VICReg components, `alignment_metric`, `uniformity_metric`,
//! `forward_dual`, an `AetherConfig` with `vicreg_*` fields). Those do
//! **not** exist in the backing crate at HEAD — `embedding.rs` exposes
//! `EmbeddingConfig { d_model, d_proj, temperature, normalize }`,
//! `info_nce_loss` (plain `f32`), `CsiAugmenter::augment_pair`, and
//! `EmbeddingExtractor::extract`. This binding surfaces **what actually
//! exists** rather than fabricating the ADR's wished-for API. The
//! VICReg loss / metric surface is a Rust-side gap, not a binding gap.
//!
//! ## GIL release strategy (per ADR-117 §7, matching bindings/vitals.rs)
//!
//! `extract`, `augment_pair`, and `info_nce_loss` are pure-sync matrix
//! ops touching no Python objects, so they run inside
//! `py.allow_threads(|| ...)`.
use pyo3::prelude::*;
use wifi_densepose_sensing_server::embedding::{
info_nce_loss as rust_info_nce_loss, CsiAugmenter, EmbeddingConfig, EmbeddingExtractor,
};
use wifi_densepose_sensing_server::graph_transformer::TransformerConfig;
// ─── AetherConfig ────────────────────────────────────────────────────
/// Configuration for the contrastive embedding model.
///
/// Python:
/// ```python
/// from wifi_densepose.aether import AetherConfig
/// cfg = AetherConfig(d_model=64, d_proj=128, temperature=0.07, normalize=True)
/// ```
#[pyclass(frozen, name = "AetherConfig")]
#[derive(Clone)]
pub struct PyAetherConfig {
inner: EmbeddingConfig,
}
#[pymethods]
impl PyAetherConfig {
#[new]
#[pyo3(signature = (d_model=64, d_proj=128, temperature=0.07, normalize=true))]
fn new(d_model: usize, d_proj: usize, temperature: f32, normalize: bool) -> Self {
Self {
inner: EmbeddingConfig {
d_model,
d_proj,
temperature,
normalize,
},
}
}
#[getter]
fn d_model(&self) -> usize {
self.inner.d_model
}
#[getter]
fn d_proj(&self) -> usize {
self.inner.d_proj
}
#[getter]
fn temperature(&self) -> f32 {
self.inner.temperature
}
#[getter]
fn normalize(&self) -> bool {
self.inner.normalize
}
fn __repr__(&self) -> String {
format!(
"AetherConfig(d_model={}, d_proj={}, temperature={}, normalize={})",
self.inner.d_model, self.inner.d_proj, self.inner.temperature, self.inner.normalize,
)
}
}
// ─── CsiAugmenter ────────────────────────────────────────────────────
/// SimCLR-style CSI augmentation. `augment_pair` returns two distinct
/// augmented views of the same CSI window for contrastive pretraining.
///
/// Python:
/// ```python
/// from wifi_densepose.aether import CsiAugmenter
/// aug = CsiAugmenter()
/// view_a, view_b = aug.augment_pair(window, seed=42)
/// ```
#[pyclass(name = "CsiAugmenter")]
pub struct PyCsiAugmenter {
inner: CsiAugmenter,
}
#[pymethods]
impl PyCsiAugmenter {
#[new]
fn new() -> Self {
Self {
inner: CsiAugmenter::new(),
}
}
/// Produce two augmented views `(view_a, view_b)` of `window`
/// (frames × subcarriers) using the deterministic `seed`. GIL is
/// released during augmentation.
fn augment_pair(
&self,
py: Python<'_>,
window: Vec<Vec<f32>>,
seed: u64,
) -> (Vec<Vec<f32>>, Vec<Vec<f32>>) {
py.allow_threads(|| self.inner.augment_pair(&window, seed))
}
fn __repr__(&self) -> String {
"CsiAugmenter(SimCLR-style CSI augmentation)".to_string()
}
}
// ─── EmbeddingExtractor ──────────────────────────────────────────────
/// Full AETHER embedding extractor: CSI→pose transformer backbone +
/// projection head → a `d_proj`-dim (default 128) L2-normalized
/// embedding. Weights are deterministically seeded, so `embed` is a
/// pure function of its input for a fixed config.
///
/// Python:
/// ```python
/// from wifi_densepose.aether import AetherConfig, EmbeddingExtractor
/// ext = EmbeddingExtractor(n_subcarriers=56, config=AetherConfig())
/// emb = ext.embed(window) # list[float], len == config.d_proj
/// ```
#[pyclass(name = "EmbeddingExtractor")]
pub struct PyEmbeddingExtractor {
inner: EmbeddingExtractor,
embedding_dim: usize,
}
#[pymethods]
impl PyEmbeddingExtractor {
/// Construct an extractor. The transformer backbone is sized from
/// `n_subcarriers` and `config.d_model`; `config.d_proj` sets the
/// embedding dimension.
#[new]
#[pyo3(signature = (n_subcarriers, config, n_keypoints=17, n_heads=4, n_gnn_layers=2))]
fn new(
n_subcarriers: usize,
config: PyAetherConfig,
n_keypoints: usize,
n_heads: usize,
n_gnn_layers: usize,
) -> Self {
let e_config = config.inner.clone();
let t_config = TransformerConfig {
n_subcarriers,
n_keypoints,
d_model: e_config.d_model,
n_heads,
n_gnn_layers,
};
let embedding_dim = e_config.d_proj;
Self {
inner: EmbeddingExtractor::new(t_config, e_config),
embedding_dim,
}
}
/// Extract an embedding from a CSI window (frames × subcarriers).
/// Returns a `d_proj`-length vector (L2-normed when the config's
/// `normalize` is set). GIL released during the forward pass.
fn embed(&mut self, py: Python<'_>, csi_features: Vec<Vec<f32>>) -> Vec<f32> {
py.allow_threads(|| self.inner.extract(&csi_features))
}
#[getter]
fn embedding_dim(&self) -> usize {
self.embedding_dim
}
fn __repr__(&self) -> String {
format!("EmbeddingExtractor(embedding_dim={})", self.embedding_dim)
}
}
// ─── Module functions ────────────────────────────────────────────────
/// InfoNCE (NT-Xent) contrastive loss between two batches of embeddings.
/// Delegates to the identical Rust implementation. GIL released.
#[pyfunction]
#[pyo3(signature = (embeddings_a, embeddings_b, temperature=0.07))]
fn info_nce_loss(
py: Python<'_>,
embeddings_a: Vec<Vec<f32>>,
embeddings_b: Vec<Vec<f32>>,
temperature: f32,
) -> f32 {
py.allow_threads(|| rust_info_nce_loss(&embeddings_a, &embeddings_b, temperature))
}
/// Cosine similarity between two embeddings — the re-ID scoring
/// primitive. Byte-identical to the private `cosine_similarity` in the
/// backing crate (same dot-product / norm formula, `f32`).
#[pyfunction]
fn cosine_similarity(a: Vec<f32>, b: Vec<f32>) -> f32 {
let n = a.len().min(b.len());
let dot: f32 = (0..n).map(|i| a[i] * b[i]).sum();
let na = (0..n).map(|i| a[i] * a[i]).sum::<f32>().sqrt();
let nb = (0..n).map(|i| b[i] * b[i]).sum::<f32>().sqrt();
if na > 1e-10 && nb > 1e-10 {
dot / (na * nb)
} else {
0.0
}
}
pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PyAetherConfig>()?;
m.add_class::<PyCsiAugmenter>()?;
m.add_class::<PyEmbeddingExtractor>()?;
m.add_function(wrap_pyfunction!(info_nce_loss, m)?)?;
m.add_function(wrap_pyfunction!(cosine_similarity, m)?)?;
Ok(())
}
+11
View File
@@ -17,6 +17,8 @@
use pyo3::prelude::*;
mod bindings {
#[cfg(feature = "aether")]
pub mod aether;
pub mod bfld;
pub mod keypoint;
pub mod pose;
@@ -43,6 +45,8 @@ fn build_features() -> Vec<&'static str> {
feats.push("p2-pose-bindings"); // BoundingBox + PersonPose + PoseEstimate
feats.push("p3-vitals-bindings"); // BreathingExtractor + HeartRateExtractor + VitalEstimate
feats.push("p3.5-bfld-bindings"); // BfldFrame + BfldReport + BfldKind (stub Rust)
#[cfg(feature = "aether")]
feats.push("p6-aether-bindings"); // ADR-185 P1 — AETHER contrastive embeddings
feats
}
@@ -85,5 +89,12 @@ fn wifi_densepose_native(m: &Bound<'_, PyModule>) -> PyResult<()> {
// the published `wifi-densepose-bfld 0.3.0` crate, not the Python port).
// Closes ADR-125 §2.1.d at the binding boundary.
bindings::privacy_gate::register(m)?;
// ADR-185 P1 — AETHER contrastive CSI embedding bindings, compiled
// and registered only under the `aether` feature so the default
// wheel links none of the sensing-server dependency tree.
#[cfg(feature = "aether")]
bindings::aether::register(m)?;
Ok(())
}