trellis_sparse_mrp_mlx/tests/test_sparse.py

"""Correctness tests for the MLX sparse core.

spconv cannot be installed on this machine — that is the entire reason this port exists —
so there is no way to diff against upstream numerically here. Instead the vectorised MLX
implementation is checked against a deliberately naive, obviously-correct reference
written straight from the definition of submanifold convolution (dict lookup, per-voxel
Python loop). The two share no indexing code, so an off-by-one in the fast path cannot
hide in both.
"""

import sys
from pathlib import Path

import mlx.core as mx
import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from trellis_sparse_mlx.conv import SubMConv3d, build_indice_map, _kernel_offsets
from trellis_sparse_mlx.tensor import SparseTensor, subdivide


def reference_subm_conv(coords, feats, weight, bias, k):
    """Definition of submanifold conv, written for obviousness, not speed."""
    occupied = {tuple(c): i for i, c in enumerate(coords.tolist())}
    offsets = _kernel_offsets(k)
    n, out_c = coords.shape[0], weight.shape[2]
    out = np.zeros((n, out_c), dtype=np.float64)
    for i, c in enumerate(coords.tolist()):
        for oi, d in enumerate(offsets):
            nb = (c[0], c[1] + d[0], c[2] + d[1], c[3] + d[2])
            j = occupied.get(nb)
            if j is not None:
                out[i] += feats[j].astype(np.float64) @ weight[oi].astype(np.float64)
    if bias is not None:
        out += bias.astype(np.float64)
    return out


def random_sparse(n_vox, channels, batch=2, res=8, seed=0):
    rng = np.random.default_rng(seed)
    seen, coords = set(), []
    while len(coords) < n_vox:
        b = int(rng.integers(0, batch))
        z, y, x = (int(v) for v in rng.integers(0, res, 3))
        if (b, z, y, x) in seen:
            continue
        seen.add((b, z, y, x))
        coords.append((b, z, y, x))
    coords.sort()  # upstream requires batch-contiguous rows
    coords = np.array(coords, dtype=np.int32)
    feats = rng.standard_normal((n_vox, channels)).astype(np.float32)
    return coords, feats


def test_subm_conv_matches_reference(k=3, n=180, cin=12, cout=7):
    coords, feats = random_sparse(n, cin, seed=1)
    rng = np.random.default_rng(2)
    w = rng.standard_normal((k**3, cin, cout)).astype(np.float32) * 0.1
    b = rng.standard_normal((cout,)).astype(np.float32)

    conv = SubMConv3d(cin, cout, k, bias=True, indice_key="t")
    conv.weight = mx.array(w)
    conv.bias = mx.array(b)

    got = np.asarray(conv(SparseTensor(mx.array(feats), mx.array(coords))).feats)
    want = reference_subm_conv(coords, feats, w, b, k)
    err = np.abs(got - want).max()
    assert err < 2e-4, f"k={k} max abs err {err:.3g}"
    return err


def test_kernel1_is_pointwise(n=64, cin=8, cout=5):
    coords, feats = random_sparse(n, cin, seed=3)
    rng = np.random.default_rng(4)
    w = rng.standard_normal((1, cin, cout)).astype(np.float32)
    conv = SubMConv3d(cin, cout, 1, bias=False, indice_key="p")
    conv.weight = mx.array(w)
    got = np.asarray(conv(SparseTensor(mx.array(feats), mx.array(coords))).feats)
    err = np.abs(got - feats @ w[0]).max()
    assert err < 1e-4, f"k=1 err {err:.3g}"
    return err


def test_isolated_voxel_sees_only_itself():
    """A voxel with no occupied neighbours must reduce to the centre tap alone."""
    coords = np.array([[0, 0, 0, 0], [0, 50, 50, 50]], dtype=np.int32)
    feats = np.ones((2, 3), dtype=np.float32)
    w = np.zeros((27, 3, 3), dtype=np.float32)
    centre = 13  # index of (0,0,0) in centred C-order offsets
    assert tuple(_kernel_offsets(3)[centre]) == (0, 0, 0)
    w[centre] = np.eye(3)
    conv = SubMConv3d(3, 3, 3, bias=False, indice_key="iso")
    conv.weight = mx.array(w)
    got = np.asarray(conv(SparseTensor(mx.array(feats), mx.array(coords))).feats)
    assert np.abs(got - feats).max() < 1e-6, got
    return 0.0


def test_batches_do_not_leak():
    """Same spatial cell in two batch items must not become neighbours."""
    coords = np.array([[0, 1, 1, 1], [1, 1, 1, 1]], dtype=np.int32)
    feats = np.array([[1.0], [100.0]], dtype=np.float32)
    w = np.ones((27, 1, 1), dtype=np.float32)  # sum every occupied neighbour
    conv = SubMConv3d(1, 1, 3, bias=False, indice_key="b")
    conv.weight = mx.array(w)
    got = np.asarray(conv(SparseTensor(mx.array(feats), mx.array(coords))).feats)
    assert np.allclose(got, [[1.0], [100.0]]), f"batch leak: {got}"
    return 0.0


def test_indice_map_hit_rate():
    """A fully dense block: interior voxels must find all 27 neighbours."""
    coords = np.array(
        [[0, z, y, x] for z in range(4) for y in range(4) for x in range(4)],
        dtype=np.int32,
    )
    imap = build_indice_map(mx.array(coords), 3)
    lin = {tuple(c): i for i, c in enumerate(coords.tolist())}
    interior = [lin[(0, z, y, x)] for z in (1, 2) for y in (1, 2) for x in (1, 2)]
    assert (imap[:, interior] != -1).all(), "interior voxel missing a neighbour"
    corner = lin[(0, 0, 0, 0)]
    assert (imap[:, corner] != -1).sum() == 8, "corner should see exactly 8 of 27"
    return 0.0


def test_subdivide():
    coords = np.array([[0, 1, 2, 3]], dtype=np.int32)
    feats = np.array([[5.0, 6.0]], dtype=np.float32)
    out = subdivide(SparseTensor(mx.array(feats), mx.array(coords)))
    c = np.asarray(out.coords)
    assert c.shape == (8, 4) and out.feats.shape == (8, 2)
    assert set(map(tuple, c[:, 1:].tolist())) == {
        (2 + a, 4 + b, 6 + d) for a in (0, 1) for b in (0, 1) for d in (0, 1)
    }
    assert np.abs(np.asarray(out.feats) - feats).max() < 1e-6
    return 0.0


if __name__ == "__main__":
    tests = [
        ("subm k=3 vs reference", lambda: test_subm_conv_matches_reference(3)),
        ("subm k=5 vs reference", lambda: test_subm_conv_matches_reference(5, n=140)),
        ("k=1 is pointwise", test_kernel1_is_pointwise),
        ("isolated voxel", test_isolated_voxel_sees_only_itself),
        ("batch isolation", test_batches_do_not_leak),
        ("indice map hit rate", test_indice_map_hit_rate),
        ("subdivide", test_subdivide),
    ]
    failed = 0
    for name, fn in tests:
        try:
            err = fn()
            print(f"  PASS  {name:28s} (max err {err:.2e})")
        except AssertionError as e:
            print(f"  FAIL  {name:28s} {e}")
            failed += 1
        except Exception as e:  # noqa: BLE001
            print(f"  ERROR {name:28s} {type(e).__name__}: {e}")
            failed += 1
    print(f"\n{len(tests)-failed}/{len(tests)} passed")
    sys.exit(1 if failed else 0)