pixal3d_mrp_mlx/tests/test_sampler.py

"""Flow Euler sampler tests — schedule, integration, and guidance call counts."""
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 pixal3d_mlx.sampler import FlowEulerSampler  # noqa: E402


def test_schedule_matches_upstream():
    s = FlowEulerSampler()
    worst = 0.0
    for r in (1.0, 3.0, 0.5):
        got = s.timesteps(7, r)
        seq = np.linspace(1, 0, 8)
        seq = r * seq / (1 + (r - 1) * seq)
        want = [(seq[i], seq[i + 1]) for i in range(7)]
        worst = max(worst, max(abs(a - c) + abs(b - d) for (a, b), (c, d) in zip(got, want)))
    assert worst < 1e-12, f"schedule err {worst:.3g}"
    return worst


def test_constant_velocity_integrates_exactly():
    """v constant => x_end = x_start - v*(1-0), regardless of step count."""
    s = FlowEulerSampler()
    class Const:
        def __call__(self, x, t, cond):
            return mx.ones_like(x) * 2.0
    for steps in (1, 10, 37):
        out = np.asarray(s.sample(Const(), mx.zeros((1, 4)), cond=None, steps=steps))
        err = abs(out[0, 0] - (-2.0))
        assert err < 1e-5, f"steps={steps} gave {out[0,0]}, want -2.0"
    return 0.0


def test_guidance_interval_skips_negative_pass():
    """Outside the interval only the conditional branch runs — half the model calls."""
    s = FlowEulerSampler()
    calls = []
    class Counting:
        def __call__(self, x, t, cond):
            calls.append(float(t[0]) / 1000.0)
            return mx.zeros_like(x)
    s.sample(Counting(), mx.zeros((1, 4)), cond="p", neg_cond="n", steps=4,
             guidance_strength=3.0, guidance_interval=(0.0, 0.5))
    assert len(calls) == 6, f"expected 6 calls (2 in-window + 1 out), got {len(calls)}"
    from collections import Counter
    c = Counter(round(t, 3) for t in calls)
    assert c[1.0] == 1 and c[0.5] == 2, dict(c)
    return 0.0


def test_guidance_is_a_lerp():
    """upstream uses g*pos + (1-g)*neg, not neg + g*(pos-neg)."""
    s = FlowEulerSampler()
    class Two:
        def __call__(self, x, t, cond):
            return mx.ones_like(x) * (1.0 if cond == "p" else 3.0)
    out = np.asarray(s.sample(Two(), mx.zeros((1, 1)), cond="p", neg_cond="n",
                              steps=1, guidance_strength=2.0))
    # g*1 + (1-g)*3 = 2 - 3 = -1  -> x = 0 - 1*(-1) = +1
    assert abs(out[0, 0] - 1.0) < 1e-5, f"got {out[0,0]}, want +1.0 for a LERP"
    return 0.0


def test_guidance_rescale_matches_upstream():
    """CFG rescale, diffed against upstream's ClassifierFreeGuidanceSamplerMixin.

    The shipped ss config sets guidance_rescale=0.7 (shape_slat 0.5), so this path
    runs at the pipeline's own defaults — it is not an exotic option. Upstream's
    torch code is imported and run directly rather than reimplemented in the test.
    """
    import torch
    sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "upstream" / "Pixal3D"))

    rng = np.random.default_rng(7)
    x_t = rng.standard_normal((2, 8, 4, 4)).astype(np.float32)
    pos = rng.standard_normal((2, 8, 4, 4)).astype(np.float32)
    neg = rng.standard_normal((2, 8, 4, 4)).astype(np.float32)
    t, g, gr = 0.6, 7.5, 0.7

    s = FlowEulerSampler()
    class Branch:
        def __call__(self, x, tt, cond):
            return mx.array(pos if cond == "p" else neg)
    mine = np.asarray(s._inference(Branch(), mx.array(x_t), t, "p", "n", g, None, gr))

    # upstream, verbatim
    sm = 1e-5
    xt_, p_, n_ = torch.from_numpy(x_t), torch.from_numpy(pos), torch.from_numpy(neg)
    pred = g * p_ + (1 - g) * n_
    to_x0 = lambda pr: (1 - sm) * xt_ - (sm + (1 - sm) * t) * pr
    x0_pos, x0_cfg = to_x0(p_), to_x0(pred)
    std_pos = x0_pos.std(dim=[1, 2, 3], keepdim=True)
    std_cfg = x0_cfg.std(dim=[1, 2, 3], keepdim=True)
    x0 = gr * (x0_cfg * (std_pos / std_cfg)) + (1 - gr) * x0_cfg
    ref = (((1 - sm) * xt_ - x0) / (sm + (1 - sm) * t)).numpy()

    err = float(np.abs(mine - ref).max())
    assert err < 2e-4, f"rescale diverges from upstream by {err:.3e}"

    # and it must be a genuine no-op at 0, or the default path silently changes
    plain = np.asarray(s._inference(Branch(), mx.array(x_t), t, "p", "n", g, None, 0.0))
    assert abs(float(np.abs(plain - (g * pos + (1 - g) * neg)).max())) < 1e-5
    return err


if __name__ == "__main__":
    tests = [("schedule vs upstream", test_schedule_matches_upstream),
             ("constant v integrates", test_constant_velocity_integrates_exactly),
             ("guidance interval", test_guidance_interval_skips_negative_pass),
             ("guidance is a lerp", test_guidance_is_a_lerp),
             ("guidance rescale vs upstream", test_guidance_rescale_matches_upstream)]
    failed = 0
    for n, fn in tests:
        try:
            fn(); print(f"  PASS  {n}")
        except AssertionError as e:
            print(f"  FAIL  {n}: {e}"); failed += 1
    print(f"\n{len(tests)-failed}/{len(tests)} passed")
    sys.exit(1 if failed else 0)