tests/test_rex_solver.py

import json
import math
from pathlib import Path

import numpy as np

from rex_repro.claims import (
    ATTEMPT_ID,
    CHALLENGE_CLAIM_HASHES,
    PAPER_ID,
    SNAPSHOT_ID,
    UPSTREAM_REVISION,
    build_evidence_bundle,
)
from rex_repro.solvers import (
    adaptive_rex_integrate,
    convergence_slope,
    integrate_rex,
    solver_registry,
    stable_step_interval,
)


def test_attempt_metadata_matches_claimed_rex_attempt():
    assert PAPER_ID == "7pQIzVNctu"
    assert ATTEMPT_ID == "11b90d4c-61f2-4d93-949e-8d4618aca972"
    assert SNAPSHOT_ID == "38427a3596a3ab91e5e2a3e94a0e7b110f6497ad1b2c40dfb611ff48c585d83f"
    assert UPSTREAM_REVISION == (
        "arxiv:2502.08834+github:zblasingame/Rex-solver"
        "@e39b57415d5608b18d7c5631595f1d38f06813b8"
    )
    assert CHALLENGE_CLAIM_HASHES == [
        "06ee77e870a2c0447848e1f6159454496f17d144d02bc08fe44441f6b7ad332f",
        "69eedf49ae10686f77613801c126d0825e1a2ea7198e4d9f31c945e00670b8e0",
        "be5532066024dda765f5b69ee4444b86c339c6adc9beeedeb4c995b2e61d0f13",
        "311e73b22c834fd47107a52f11e90aa005067fa136e9b477108effe648d04cb2",
    ]


def test_rex_forward_then_backward_recovers_initial_state():
    def rhs(_t, y):
        return np.array([-0.4 * y[0] + 0.1 * y[1], -0.2 * y[1]])

    y0 = np.array([1.25, -0.75])
    result = integrate_rex(rhs, y0, t0=0.0, t1=1.0, steps=32, method="rk4")

    assert np.max(np.abs(result.recovered - y0)) < 1e-11
    assert result.round_trip_error < 1e-11


def test_rk4_rex_converges_faster_than_midpoint_and_euler():
    def rhs(t, y):
        return y

    y0 = np.array([1.0])
    exact = np.array([math.e])
    euler_slope = convergence_slope(rhs, y0, exact, method="euler", step_counts=[8, 16, 32, 64])
    midpoint_slope = convergence_slope(
        rhs, y0, exact, method="midpoint", step_counts=[8, 16, 32, 64]
    )
    rk4_slope = convergence_slope(rhs, y0, exact, method="rk4", step_counts=[8, 16, 32, 64])

    assert 0.85 < euler_slope < 1.25
    assert 1.75 < midpoint_slope < 2.30
    assert rk4_slope > 3.5
    assert rk4_slope > midpoint_slope + 1.0


def test_adaptive_rex_changes_step_count_and_preserves_reversibility():
    def rhs(t, y):
        return np.array([math.sin(4.0 * t) - 0.3 * y[0]])

    y0 = np.array([0.2])
    loose = adaptive_rex_integrate(rhs, y0, 0.0, 1.0, method="heun_euler", rtol=1e-3)
    tight = adaptive_rex_integrate(rhs, y0, 0.0, 1.0, method="heun_euler", rtol=1e-6)

    assert loose.accepted_steps != tight.accepted_steps
    assert loose.round_trip_error < 1e-10
    assert tight.round_trip_error < 1e-10


def test_registry_records_solver_subsumption_and_stability_probe():
    registry = solver_registry()

    assert {"ddim", "dpm_solver_1", "dpm_solver_2", "seeds_1"}.issubset(
        registry["diffusion_specializations"]
    )
    assert {"euler", "midpoint", "rk4", "heun_euler"}.issubset(registry["tableaus"])
    assert registry["supports_sde"] is True
    assert registry["supports_adaptive"] is True
    assert stable_step_interval(method="euler", lambda_real=-1.0) > 0.0


def test_evidence_bundle_is_machine_readable_and_scoped_to_target_claims(tmp_path):
    out = tmp_path / "bundle.json"
    bundle = build_evidence_bundle(out)

    loaded = json.loads(out.read_text())
    assert loaded == bundle
    assert loaded["paper_id"] == PAPER_ID
    assert loaded["attempt_id"] == ATTEMPT_ID
    assert loaded["snapshot_id"] == SNAPSHOT_ID
    assert [claim["challenge_claim_sha256"] for claim in loaded["claims"]] == CHALLENGE_CLAIM_HASHES
    assert {claim["status"] for claim in loaded["claims"]} <= {"verified", "toy", "unavailable"}
    assert loaded["non_target_claims"]["gpu_image_generation"]["status"] == "unavailable"