# Technical Roadmap: High-Performance Sparse Covariance Engine (Rust + Sparse Kernels)

**Goal:** Eliminate Python-side loop overhead and memory scaling by offloading the General Law of Error Propagation (GLOP) to a highly optimized, sparse-matrix Rust backend.

**Mathematical Foundation:**
The engine must compute the propagation of the covariance matrix under a transformation characterized by the Jacobian .

Where:

- is the global sparse covariance matrix (state).
- is the sparse Jacobian of the current operation (inputs outputs).

---

## Phase 1: The Sparse Algebra Kernel (Math Layer)

**Objective:** Implement the "Sandwich Product" () efficiently using Compressed Sparse Row/Column (CSR/CSC) formats. The current `sprs` dependency in `Cargo.toml` must be leveraged or upgraded.

### 1.1. Sparse Sandwich Product Implementation

The core bottleneck is multiplying three matrices where the middle one () is large and sparse, and the outer ones () are extremely sparse (often diagonal or block-diagonal).

- **Action:** Implement a trait `SandwichProduct` in `measurekit_core/src/math.rs`.
- **Algorithm:**
  Do not perform naive . Instead, use a row-iterator approach that exploits the structure of .
  For a generic row of the result:

Since is sparse, iterate only over non-zero and .

- **Technical Requirement:**
- Inputs: `J` (Sparse Matrix or List of Sparse Blocks), `Sigma_slice` (View of the Global Store).
- Output: `Sigma_block` (The update delta).

- **Deliverable:** A unit-tested Rust function `sparse_sandwich(J: &CsMat, Sigma: &CsMat) -> CsMat`.

### 1.2. Handling Cross-Correlations (The Dependency Problem)

The engine must explicitly handle binary operations between variables that _already_ share a history.
If , the variance depends on .

- **Action:** Redesign the `update_covariance` signature to accept a **Block-Map**.
- **Logic:**

1. Identify indices for and for in the global store.
2. Extract the submatrix (cross-covariance).
3. Construct the full local covariance block:

4. Construct the concatenated Jacobian .
5. Compute .

---

## Phase 2: Memory Architecture & The Global Store

**Objective:** Refactor `CovarianceStore` in `measurekit_core/src/covariance.rs` to store **values**, not just update timestamps.

### 2.1. The Block-Sparse Storage

A single monolithic `CsMat` is inefficient for updates (insertion requires shifting data).

- **Action:** Replace `matrix: CsMatI<f64, usize>` with a **Block-Map Architecture**.
- **Structure Definition:**

```rust
// measurekit_core/src/covariance.rs

type VariableID = u64; // Hash of the variable lineage or UUID

struct CovarianceStore {
    // Maps (Row_Var, Col_Var) -> Covariance Block
    blocks: HashMap<(VariableID, VariableID), CsMat<f64>>,

    // Metadata for pruning
    access_ledger: HashMap<VariableID, u64>,
}

```

- **Benefit:** Updates to independent variables become hash map insertions. Reading correlations between and only locks their specific block, not the whole matrix (enabling concurrency).

### 2.2. Zero-Copy FFI (PyO3 Integration)

Pass Jacobian data from Python (NumPy/Torch/JAX) to Rust without copying memory.

- **Action:** Update `lib.rs` and `covariance.rs` to use `numpy::PyReadonlyArrayDyn`.
- **Signature:**

```rust
#[pyo3(signature = (out_id, input_ids, jacobians))]
fn propagate(
    &mut self,
    out_id: VariableID,
    input_ids: Vec<VariableID>,
    jacobians: Vec<PyReadonlyArrayDyn<f64>>
) -> PyResult<()>

```

- **Implementation Details:**
- Convert dense NumPy Jacobians (if small) to internal Sparse format on the fly.
- If inputs are JAX/Torch tensors on GPU, this roadmap assumes a `cpu().numpy()` sync initially. (Phase 4 would address GPU-direct pointer passing).

---

## Phase 3: Pruning & Numerical Stability

**Objective:** Prevent the global matrix from growing infinitely (the "Memory Leak" of physical history).

### 3.1. Relevance-Based Pruning (The Metric)

The current "Age" pruning is physically invalid (old correlations are still real).

- **Action:** Implement `prune_weak_correlations` in Rust.
- **Algorithm:**
  Iterate over off-diagonal blocks in the `HashMap`.
  Compute the Pearson Correlation Coefficient for the block norm:

If (e.g., ), drop the key from the HashMap.

- **Thread Safety:** This garbage collection can run on a background thread using `RwLock` on the Store.

---

## Phase 4: Python-Side Integration Strategy

**Objective:** Connect the Python `AutogradPropagator` to the new Rust engine.

### 4.1. The Vectorized Backend Dispatcher

Modify `measurekit/domain/measurement/vectorized_uncertainty.py`.

- **Logic Change:**
  Currently, it calculates `out_cov` using high-level Python ops.
  **New Flow:**

1. `AutogradPropagator` computes Jacobians .
2. `CovarianceModel` extracts `input_ids` (pointers to Rust store).
3. Call `rust_store.propagate(new_id, [id1, id2], [J1, J2])`.
4. Return a light `CovarianceModel` containing only the `new_id`. (The actual data stays in Rust).

### 4.2. Lazy Evaluation (Optional Optimization)

Don't compute immediately. Just record the operation node in Rust.
Only compute the matrix multiplication when the user requests `.std_dev` or `.value`.

- **Mechanism:** Build a Directed Acyclic Graph (DAG) in Rust.
- **Trigger:** `store.resolve(id)`.

---

## Summary of Implementation Steps (Priority Order)

1. **Rust Structs:** Rewrite `CovarianceStore` to use `HashMap<(u64, u64), CsMat>`.
2. **Rust Math:** Implement `sparse_sandwich_product` using `sprs`.
3. **Rust API:** Expose `register_variable(data)` and `propagate(jacs, ids)`.
4. **Python Bindings:** Update `measurekit_core` PyO3 module.
5. **Python Logic:** Rip out the Python-side matrix math in `vectorized_uncertainty.py` and replace it with calls to the Rust backend.
6. **Tests:** Verify vs (uncorrelated) vs (perfectly correlated).

This roadmap provides the most direct path to a production-grade, physically accurate engine capable of handling millions of parameters with minimal overhead.
