bundle/v0/tests/test_appendix_routes.py

from __future__ import annotations

import numpy as np

from loss_aware_dro_repro.appendix_conic import (
    absolute_regression_socp,
    solve_appendix_socp,
    solve_absolute_regression,
    solve_empirical_w1_portfolio,
    solve_squared_regression,
)
from loss_aware_dro_repro.empirical_ot import empirical_wasserstein_value_gradient
from loss_aware_dro_repro.gelbrich import lower_triangular_central_difference, relative_gradient_error
from loss_aware_dro_repro.residuals import conic_residual_maximum


def _data():
    rng = np.random.default_rng(731)
    portfolio = rng.normal(size=(10, 3))
    X = rng.normal(size=(10, 1))
    y = 1.4 * X[:, 0] + rng.normal(size=10)
    return portfolio, X, y


def test_empirical_w1_portfolio_matches_closed_form_and_gradient():
    portfolio, _, _ = _data()
    epsilon = 0.2
    L = np.array([[1.1, 0.0, 0.0], [0.1, 0.9, 0.0], [-0.05, 0.08, 1.2]])
    result = solve_empirical_w1_portfolio(portfolio, epsilon, L)
    decision = result["decision"]
    expected = -portfolio.mean(axis=0) @ decision + epsilon * np.linalg.norm(np.linalg.solve(L, decision))
    numerical = lower_triangular_central_difference(lambda factor: solve_empirical_w1_portfolio(portfolio, epsilon, factor)["objective"], L)
    assert abs(result["objective"] - expected) < 1e-8
    assert abs(decision.sum() - 1.0) < 1e-9
    assert decision.min() > -1e-9
    assert conic_residual_maximum(result["residuals"]) < 1e-8
    assert relative_gradient_error(result["value_gradient"], numerical) < 1e-4


def test_absolute_regression_matches_appendix_objective_and_gradient():
    _, X, y = _data()
    epsilon = 0.2
    L = np.array([[1.1, 0.0], [0.15, 0.9]])
    result = solve_absolute_regression(X, y, epsilon, L)
    weight = result["decision"]
    coefficient = np.concatenate((-weight, [1.0]))
    expected = np.mean(np.abs(y - X @ weight)) + epsilon * np.linalg.norm(np.linalg.solve(L, coefficient))
    numerical = lower_triangular_central_difference(lambda factor: solve_absolute_regression(X, y, epsilon, factor)["objective"], L)
    assert abs(result["objective"] - expected) < 1e-8
    assert conic_residual_maximum(result["residuals"]) < 1e-8
    assert relative_gradient_error(result["value_gradient"], numerical) < 1e-4


def test_appendix_complementarity_excludes_zero_cone_equalities():
    _, X, y = _data()
    program = absolute_regression_socp(
        X,
        y,
        0.2,
        np.array([[1.1, 0.0], [0.15, 0.9]]),
    )
    result = solve_appendix_socp(program)
    expected = abs(float(result["dual"] @ result["slack"]))
    assert result["residuals"]["complementarity"] == expected


def test_squared_regression_post_solve_square_and_gradient():
    _, X, y = _data()
    epsilon = 0.2
    L = np.array([[1.1, 0.0], [0.15, 0.9]])
    result = solve_squared_regression(X, y, epsilon, L)
    weight = result["decision"]
    coefficient = np.concatenate((-weight, [1.0]))
    expected_root = np.sqrt(np.mean((y - X @ weight) ** 2)) + epsilon * np.linalg.norm(np.linalg.solve(L, coefficient))
    numerical = lower_triangular_central_difference(lambda factor: solve_squared_regression(X, y, epsilon, factor)["objective"], L)
    assert result["post_solve_square"] is True
    assert abs(result["root_objective"] - expected_root) < 1e-8
    assert abs(result["objective"] - expected_root**2) < 1e-8
    assert conic_residual_maximum(result["residuals"]) < 1e-8
    assert relative_gradient_error(result["value_gradient"], numerical) < 1e-4


def test_empirical_wasserstein_envelope_gradients_and_marginals():
    rng = np.random.default_rng(19)
    first = rng.normal(size=(8, 2))
    second = rng.normal(size=(8, 2))
    L = np.array([[1.1, 0.0], [0.2, 0.9]])
    for order in (1, 2):
        value, analytic, metadata = empirical_wasserstein_value_gradient(first, second, L, order)
        numerical = lower_triangular_central_difference(
            lambda factor: empirical_wasserstein_value_gradient(first, second, factor, order)[0], L
        )
        assert value > 0
        assert metadata["marginal_row_residual"] < 1e-12
        assert metadata["marginal_column_residual"] < 1e-12
        assert metadata["dual_feasibility_violation"] < 1e-12
        assert metadata["complementary_slackness"] < 1e-12
        assert metadata["duality_gap"] < 1e-12
        assert relative_gradient_error(analytic, numerical) < 1e-6


def test_empirical_w2_exact_match_has_zero_conservative_subgradient():
    points = np.array([[0.0, 1.0], [2.0, -1.0]])
    value, gradient, metadata = empirical_wasserstein_value_gradient(points, points, np.eye(2), 2)
    assert value == 0.0
    assert np.array_equal(gradient, np.zeros((2, 2)))
    assert metadata["transport_power_value"] == 0.0


def test_empirical_wasserstein_retains_tiny_distinct_shift():
    first = np.array([[0.0]])
    second = np.array([[1e-12]])
    for order in (1, 2):
        value, gradient, metadata = empirical_wasserstein_value_gradient(first, second, np.eye(1), order)
        np.testing.assert_allclose(value, 1e-12, rtol=1e-12, atol=0.0)
        assert gradient[0, 0] != 0.0
        assert metadata["transport_power_value"] > 0.0