source/gconvex_authors/tests/test_FCOTSeparable.py

import torch
import pytest
from optimal_transport.ot_fc_sep_map import FCOTSeparable
from models import FiniteSeparableModel


################################################################################
# Helper: Simple 1D quadratic kernel
################################################################################

def quad_kernel(x, y):
    return (x - y) ** 2


################################################################################
# Fixture: tiny model that runs fast
################################################################################

@pytest.fixture
def small_model():
    model = FiniteSeparableModel(
        kernel=quad_kernel,
        num_dims=1,
        radius=2.0,
        y_accuracy=0.5,
        x_accuracy=0.5,
        mode="concave",
        temp=5.0,
        epsilon=1e-6,
        cache_gradients=False,
    )
    return model


@pytest.fixture
def small_solver(small_model):
    solver = FCOTSeparable(
        input_dim=1,
        model=small_model,
        inverse_kx=lambda x, p: x - 0.5 * p,
        outer_lr=1e-2,
        warmup_lr=1e-1,
        full_refresh_every=5,               # frequency must be small for testing
        reactivate_every=3,
        reactivate_eps=1e-3,
    )
    return solver


################################################################################
# 1. INTERCEPT REFRESH TEST
################################################################################

def test_refresh_updates_intercepts(small_model):
    # Start from a nontrivial intercept configuration so refresh has work to do
    torch.manual_seed(0)
    with torch.no_grad():
        small_model.intercepts.copy_(torch.randn_like(small_model.intercepts))

    old = small_model.intercepts.clone()

    changed = small_model.refresh_intercepts_via_transform()
    new = small_model.intercepts

    # Function must run without error and produce a valid intercept tensor.
    # Depending on initialization and kernel, refresh may be idempotent.
    assert isinstance(changed, int)
    assert new.shape == old.shape


################################################################################
# 2. MOMENTUM RESET TEST (Adam or TwoStageOptimizer)
################################################################################

def test_momentum_reset(small_solver):
    solver = small_solver
    opt = solver.optimizer

    p = solver.model.intercepts

    # Simulate momentum state
    for group in opt.param_groups:
        for param in group["params"]:
            if param is p:
                state = opt.state.setdefault(param, {})
                state["momentum_buffer"] = torch.ones_like(p) * 3.14

    # Apply reset
    solver._reset_intercept_momentum()

    # Check cleared:
    state = opt.state[p]
    assert len(state) == 0, "Momentum/state was not fully cleared."


################################################################################
# 3. JITTER TEST
################################################################################

def test_jitter_changes_intercepts(small_solver):
    model = small_solver.model
    before = model.intercepts.clone()
    small_solver._jitter_intercepts(strength=1e-2)
    after = model.intercepts

    assert not torch.allclose(before, after), "Jitter must modify intercepts."
    diff = (after - before).abs().mean().item()
    assert diff > 0, "Jitter must have nonzero magnitude."


################################################################################
# 4. FULL REFRESH (with reset + jitter)
################################################################################

def test_full_refresh_triggers_all(small_solver):
    solver = small_solver
    model = solver.model

    # Overwrite refresh to count calls
    calls = {"refresh": 0, "reset": 0, "jitter": 0}

    orig_refresh = model.refresh_intercepts_via_transform
    def counted_refresh():
        calls["refresh"] += 1
        return orig_refresh()

    solver.model.refresh_intercepts_via_transform = counted_refresh

    orig_reset = solver._reset_intercept_momentum
    def counted_reset():
        calls["reset"] += 1
        return orig_reset()

    solver._reset_intercept_momentum = counted_reset

    orig_jitter = solver._jitter_intercepts
    def counted_jitter(strength=None):
        calls["jitter"] += 1
        return orig_jitter(strength)

    solver._jitter_intercepts = counted_jitter

    # Force a step_count where refresh should trigger
    solver._step_count = solver._full_refresh_every
    solver._maybe_full_refresh()

    assert calls["refresh"] == 1, "Refresh should be called exactly once."
    assert calls["reset"] == 1, "Momentum reset should be called exactly once."
    assert calls["jitter"] == 1, "Jitter should be called exactly once."


################################################################################
# 5. FORWARD AND STE TEST
################################################################################

def test_forward_shapes(small_model):
    X = torch.tensor([[0.2], [-0.5], [1.0]])
    choice, fx = small_model.forward(X, selection_mode="hard")

    assert choice.shape == (3, 1)
    assert fx.shape == (3,)
    assert not fx.isnan().any()


def test_ste_gradient_correctness(small_model):
    # Hard selection with STE: gradient should propagate through approx derivative.
    X = torch.tensor([[0.1], [0.5]], requires_grad=True)
    _, fx = small_model.forward(X, selection_mode="hard")
    loss = fx.sum()
    loss.backward()

    assert X.grad is not None, "Gradient must flow through STE wrapper."
    assert not torch.isnan(X.grad).any()
    # gradient magnitude must be finite / nonzero for interior x
    assert X.grad.abs().max() > 0


################################################################################
# 6. TRANSFORM CONSISTENCY TEST
################################################################################

def test_transform_consistency(small_model):
    Z = torch.linspace(-1, 1, 5).reshape(-1, 1)
    Xopt, vals, conv = small_model.inf_transform(Z)

    assert Xopt.shape == Z.shape, "Transform output must match input shape."
    assert vals.shape == (5,), "Value shape mismatch."
    assert conv is True, "Separable transform must always converge."
    assert not torch.isnan(vals).any()


################################################################################
# 7. SMALL TRAINING LOOP TEST
################################################################################

def test_training_makes_progress():
    """
    Non-monotone-safe training test.

    Checks that:
      - objective stays finite,
      - gradients stay reasonable,
      - progress occurs over a window,
    without assuming monotonic ascent (refresh/jitter/STE break monotonicity).
    """

    import math
    import torch
    from optimal_transport.ot_fc_sep_map import FCOTSeparable

    torch.manual_seed(0)

    # 1D synthetic example
    X = torch.randn(200, 1) * 0.5
    Y = torch.randn(200, 1) + 1.0

    # Build a small separable solver
    solver = FCOTSeparable.initialize_right_architecture(
        dim=1,
        radius=3.0,
        n_params=40,                   # modest grid size: 40 intercepts in 1D
        x_accuracy=0.1,
        kernel_1d=lambda x, y: (x - y) ** 2,
        inverse_kx=lambda x, p: x - 0.5 * p,
        outer_lr=5e-3,
        warmup_lr=1e-2,
        warmup_grad_threshold=0.1,
        warmup_max_steps=20,
        full_refresh_every=10,         # refresh allowed
        reactivate_every=15,
    )

    history = []
    for _ in range(50):
        res = solver.step(X, Y)
        obj = res["dual"]
        grad_norm = res["grad_norm"]

        # finite values
        assert not math.isnan(obj)
        assert not math.isinf(obj)
        assert grad_norm < 1e6

        history.append(obj)

    # progress over windows
    early_best = max(history[:5])
    late_best = max(history[-10:])

    assert late_best > early_best, (
        f"OT objective failed to improve over time: "
        f"early={early_best:.4f}, late={late_best:.4f}"
    )

################################################################################
# 8. TRANSPORT MAP TEST
################################################################################

def test_transport_map(small_solver):
    solver = small_solver
    X = torch.tensor([[-0.5], [0.0], [0.5]], dtype=torch.float32)

    # Compute transport map explicitly from potential and inverse_kx
    X_requires_grad = X.to(solver.device).requires_grad_(True)
    _, u_X = solver.model.forward(X_requires_grad, selection_mode="hard")
    grad_u = torch.autograd.grad(u_X.sum(), X_requires_grad, create_graph=False)[0]
    Ypred = solver.inverse_kx(X_requires_grad.detach(), grad_u.detach())

    assert Ypred.shape == X.shape
    assert not torch.isnan(Ypred).any()
    # In 1D quadratic, T(x) = x − 0.5 * ∇u(x), so should not be extreme.
    assert Ypred.abs().max() < 10, "Transport values look unreasonable."