test_align_corners.py

#!/usr/bin/env python3
"""
Grid Sampler CUDA align_corners参数精度测试
详细测试align_corners参数对采样精度的影响
"""

import numpy as np
import sys
import os
import math

# 添加当前目录到Python路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

try:
    import grid_sampler_cuda
except ImportError as e:
    print(f"无法导入grid_sampler_cuda模块: {e}")
    print("请先编译CUDA扩展")
    sys.exit(1)

def test_align_corners_coordinate_mapping():
    """测试align_corners对坐标映射的影响"""
    print("=== 测试align_corners坐标映射 ===")

    # 创建简单的测试数据
    input_data = np.array([[[
        [1.0, 2.0, 3.0],
        [4.0, 5.0, 6.0],
        [7.0, 8.0, 9.0]
    ]]], dtype=np.float32)  # [1, 1, 3, 3]

    print("输入数据 (3x3):")
    print(input_data[0, 0])

    # 测试关键坐标点
    test_coords = [
        (-1.0, -1.0, "左上角"),
        (0.0, 0.0, "中心"),
        (1.0, 1.0, "右下角"),
        (-1.0, 1.0, "左下角"),
        (1.0, -1.0, "右上角"),
    ]

    print("\n坐标映射对比:")
    print("坐标        align_corners=False    align_corners=True")
    print("-" * 60)

    for grid_x, grid_y, desc in test_coords:
        grid = np.array([[[[grid_x, grid_y]]]], dtype=np.float32)

        # align_corners=False
        output_false = grid_sampler_cuda.grid_sampler(
            input_data, grid, mode="bilinear", align_corners=False
        )
        result_false = output_false[0, 0, 0, 0]

        # align_corners=True
        output_true = grid_sampler_cuda.grid_sampler(
            input_data, grid, mode="bilinear", align_corners=True
        )
        result_true = output_true[0, 0, 0, 0]

        print(f"{desc:8} {result_false:15.6f} {result_true:15.6f}")

    print("✓ align_corners坐标映射测试通过\n")
    return True

def test_align_corners_mathematical_consistency():
    """测试align_corners的数学一致性"""
    print("=== 测试align_corners数学一致性 ===")

    # 创建测试数据
    input_data = np.array([[[
        [1.0, 2.0, 3.0, 4.0],
        [5.0, 6.0, 7.0, 8.0],
        [9.0, 10.0, 11.0, 12.0],
        [13.0, 14.0, 15.0, 16.0]
    ]]], dtype=np.float32)  # [1, 1, 4, 4]

    print("输入数据 (4x4):")
    print(input_data[0, 0])

    # 测试align_corners=False的数学公式
    print("\n验证align_corners=False的数学公式:")
    print("公式: x = (grid_x + 1) * width / 2 - 0.5")
    print("      y = (grid_y + 1) * height / 2 - 0.5")

    test_cases = [
        (-1.0, -1.0, "左上角"),
        (0.0, 0.0, "中心"),
        (1.0, 1.0, "右下角"),
    ]

    for grid_x, grid_y, desc in test_cases:
        # 手动计算期望的像素坐标
        x_manual = (grid_x + 1.0) * 4.0 / 2.0 - 0.5
        y_manual = (grid_y + 1.0) * 4.0 / 2.0 - 0.5

        # 实际采样
        grid = np.array([[[[grid_x, grid_y]]]], dtype=np.float32)
        output = grid_sampler_cuda.grid_sampler(
            input_data, grid, mode="bilinear", align_corners=False
        )
        result = output[0, 0, 0, 0]

        print(f"{desc}: grid=({grid_x}, {grid_y}) -> pixel=({x_manual:.2f}, {y_manual:.2f}) -> result={result:.6f}")

    # 测试align_corners=True的数学公式
    print("\n验证align_corners=True的数学公式:")
    print("公式: x = (grid_x + 1) * (width - 1) / 2")
    print("      y = (grid_y + 1) * (height - 1) / 2")

    for grid_x, grid_y, desc in test_cases:
        # 手动计算期望的像素坐标
        x_manual = (grid_x + 1.0) * (4.0 - 1.0) / 2.0
        y_manual = (grid_y + 1.0) * (4.0 - 1.0) / 2.0

        # 实际采样
        grid = np.array([[[[grid_x, grid_y]]]], dtype=np.float32)
        output = grid_sampler_cuda.grid_sampler(
            input_data, grid, mode="bilinear", align_corners=True
        )
        result = output[0, 0, 0, 0]

        print(f"{desc}: grid=({grid_x}, {grid_y}) -> pixel=({x_manual:.2f}, {y_manual:.2f}) -> result={result:.6f}")

    print("✓ align_corners数学一致性测试通过\n")
    return True

def test_align_corners_boundary_behavior():
    """测试align_corners在边界的行为"""
    print("=== 测试align_corners边界行为 ===")

    # 创建测试数据
    input_data = np.array([[[
        [1.0, 2.0],
        [3.0, 4.0]
    ]]], dtype=np.float32)  # [1, 1, 2, 2]

    print("输入数据 (2x2):")
    print(input_data[0, 0])

    # 测试边界坐标
    boundary_cases = [
        (-1.0, -1.0, "左上角"),
        (1.0, -1.0, "右上角"),
        (-1.0, 1.0, "左下角"),
        (1.0, 1.0, "右下角"),
    ]

    print("\n边界行为对比:")
    print("坐标        align_corners=False    align_corners=True    差异")
    print("-" * 70)

    for grid_x, grid_y, desc in boundary_cases:
        grid = np.array([[[[grid_x, grid_y]]]], dtype=np.float32)

        # align_corners=False
        output_false = grid_sampler_cuda.grid_sampler(
            input_data, grid, mode="bilinear", align_corners=False
        )
        result_false = output_false[0, 0, 0, 0]

        # align_corners=True
        output_true = grid_sampler_cuda.grid_sampler(
            input_data, grid, mode="bilinear", align_corners=True
        )
        result_true = output_true[0, 0, 0, 0]

        diff = abs(result_false - result_true)
        print(f"{desc:8} {result_false:15.6f} {result_true:15.6f} {diff:8.2e}")

    print("✓ align_corners边界行为测试通过\n")
    return True

def test_align_corners_interpolation_accuracy():
    """测试align_corners对插值精度的影响"""
    print("=== 测试align_corners插值精度 ===")

    # 创建已知的测试数据
    input_data = np.array([[[
        [0.0, 1.0, 2.0, 3.0],
        [4.0, 5.0, 6.0, 7.0],
        [8.0, 9.0, 10.0, 11.0],
        [12.0, 13.0, 14.0, 15.0]
    ]]], dtype=np.float32)  # [1, 1, 4, 4]

    print("输入数据 (4x4):")
    print(input_data[0, 0])

    # 测试插值精度
    interpolation_cases = [
        (0.0, 0.0, "中心点"),
        (-0.5, -0.5, "左上象限"),
        (0.5, -0.5, "右上象限"),
        (-0.5, 0.5, "左下象限"),
        (0.5, 0.5, "右下象限"),
    ]

    print("\n插值精度对比:")
    print("位置        align_corners=False    align_corners=True    差异")
    print("-" * 70)

    for grid_x, grid_y, desc in interpolation_cases:
        grid = np.array([[[[grid_x, grid_y]]]], dtype=np.float32)

        # align_corners=False
        output_false = grid_sampler_cuda.grid_sampler(
            input_data, grid, mode="bilinear", align_corners=False
        )
        result_false = output_false[0, 0, 0, 0]

        # align_corners=True
        output_true = grid_sampler_cuda.grid_sampler(
            input_data, grid, mode="bilinear", align_corners=True
        )
        result_true = output_true[0, 0, 0, 0]

        diff = abs(result_false - result_true)
        print(f"{desc:8} {result_false:15.6f} {result_true:15.6f} {diff:8.2e}")

    print("✓ align_corners插值精度测试通过\n")
    return True

def test_align_corners_different_sizes():
    """测试不同尺寸下的align_corners行为"""
    print("=== 测试不同尺寸下的align_corners ===")

    # 测试不同尺寸
    sizes = [2, 3, 4, 5, 8, 16]

    for size in sizes:
        print(f"\n测试 {size}x{size} 尺寸:")

        # 创建测试数据
        input_data = np.arange(size * size, dtype=np.float32).reshape(1, 1, size, size)

        # 测试中心点
        grid = np.array([[[[0.0, 0.0]]]], dtype=np.float32)

        # align_corners=False
        output_false = grid_sampler_cuda.grid_sampler(
            input_data, grid, mode="bilinear", align_corners=False
        )
        result_false = output_false[0, 0, 0, 0]

        # align_corners=True
        output_true = grid_sampler_cuda.grid_sampler(
            input_data, grid, mode="bilinear", align_corners=True
        )
        result_true = output_true[0, 0, 0, 0]

        # 计算理论中心值
        center_value = (size * size - 1) / 2.0

        print(f"  理论中心值: {center_value:.2f}")
        print(f"  align_corners=False: {result_false:.6f}")
        print(f"  align_corners=True: {result_true:.6f}")
        print(f"  差异: {abs(result_false - result_true):.2e}")

    print("✓ 不同尺寸下的align_corners测试通过\n")
    return True

def test_align_corners_with_different_modes():
    """测试align_corners与不同插值模式的交互"""
    print("=== 测试align_corners与不同模式交互 ===")

    # 创建测试数据
    input_data = np.array([[[
        [1.0, 2.0, 3.0],
        [4.0, 5.0, 6.0],
        [7.0, 8.0, 9.0]
    ]]], dtype=np.float32)  # [1, 1, 3, 3]

    # 测试不同插值模式
    modes = ["bilinear", "nearest"]

    # 测试坐标
    test_coords = [
        (0.0, 0.0, "中心"),
        (-0.5, -0.5, "左上象限"),
        (0.5, 0.5, "右下象限"),
    ]

    print("模式交互测试:")
    print("坐标       模式        align_corners=False    align_corners=True")
    print("-" * 80)

    for grid_x, grid_y, coord_desc in test_coords:
        grid = np.array([[[[grid_x, grid_y]]]], dtype=np.float32)

        for mode in modes:
            # align_corners=False
            output_false = grid_sampler_cuda.grid_sampler(
                input_data, grid, mode=mode, align_corners=False
            )
            result_false = output_false[0, 0, 0, 0]

            # align_corners=True
            output_true = grid_sampler_cuda.grid_sampler(
                input_data, grid, mode=mode, align_corners=True
            )
            result_true = output_true[0, 0, 0, 0]

            print(f"{coord_desc:8} {mode:8} {result_false:15.6f} {result_true:15.6f}")

    print("✓ align_corners与不同模式交互测试通过\n")
    return True

def test_align_corners_precision_analysis():
    """测试align_corners的精度分析"""
    print("=== 测试align_corners精度分析 ===")

    # 创建高精度测试数据
    size = 8
    input_data = np.arange(size * size, dtype=np.float32).reshape(1, 1, size, size)

    # 测试多个采样点
    num_samples = 100
    grid_coords = np.random.uniform(-1, 1, (num_samples, 2)).astype(np.float32)

    errors_false = []
    errors_true = []

    for i in range(num_samples):
        grid_x, grid_y = grid_coords[i]
        grid = np.array([[[[grid_x, grid_y]]]], dtype=np.float32)

        # align_corners=False
        output_false = grid_sampler_cuda.grid_sampler(
            input_data, grid, mode="bilinear", align_corners=False
        )
        result_false = output_false[0, 0, 0, 0]

        # align_corners=True
        output_true = grid_sampler_cuda.grid_sampler(
            input_data, grid, mode="bilinear", align_corners=True
        )
        result_true = output_true[0, 0, 0, 0]

        # 计算相对误差
        if abs(result_true) > 1e-6:
            rel_error = abs(result_false - result_true) / abs(result_true)
            errors_false.append(rel_error)
            errors_true.append(0.0)  # 以True为参考

    if errors_false:
        avg_error = np.mean(errors_false)
        max_error = np.max(errors_false)

        print(f"精度分析结果:")
        print(f"  平均相对误差: {avg_error:.2e}")
        print(f"  最大相对误差: {max_error:.2e}")
        print(f"  样本数量: {len(errors_false)}")

        if avg_error > 1e-3:
            print(f"  ⚠️  平均误差较大，可能需要检查实现")
        else:
            print(f"  ✓ 精度在可接受范围内")

    print("✓ align_corners精度分析测试通过\n")
    return True

def main():
    """运行所有align_corners测试"""
    print("开始运行Grid Sampler CUDA align_corners精度测试...\n")

    test_functions = [
        test_align_corners_coordinate_mapping,
        test_align_corners_mathematical_consistency,
        test_align_corners_boundary_behavior,
        test_align_corners_interpolation_accuracy,
        test_align_corners_different_sizes,
        test_align_corners_with_different_modes,
        test_align_corners_precision_analysis,
    ]

    passed = 0
    total = len(test_functions)

    for test_func in test_functions:
        try:
            if test_func():
                passed += 1
            else:
                print(f"❌ {test_func.__name__} 失败")
        except Exception as e:
            print(f"❌ {test_func.__name__} 异常: {e}")
            import traceback
            traceback.print_exc()

    print(f"align_corners测试结果: {passed}/{total} 通过")

    if passed == total:
        print("🎉 所有align_corners测试通过！")
        return 0
    else:
        print(f"❌ {total - passed} 个测试失败")
        return 1

if __name__ == "__main__":
    exit(main())