test_edge_cases.py

#!/usr/bin/env python3
"""
Grid Sampler CUDA 边界情况和特殊场景测试
测试各种边界条件、特殊值和异常情况
"""

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_single_pixel_input():
    """测试单像素输入"""
    print("=== 测试单像素输入 ===")

    # 1x1输入
    input_data = np.array([[[[42.0]]]], dtype=np.float32)  # [1, 1, 1, 1]

    # 测试不同的网格坐标
    test_cases = [
        (-1.0, -1.0, "左上角"),
        (0.0, 0.0, "中心"),
        (1.0, 1.0, "右下角"),
        (-2.0, -2.0, "超出边界"),
    ]

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

        # 测试不同填充模式
        for padding_mode in ["zeros", "border", "reflection"]:
            output = grid_sampler_cuda.grid_sampler(
                input_data, grid, mode="bilinear", padding_mode=padding_mode
            )
            result = output[0, 0, 0, 0]
            print(f"  {desc} ({padding_mode}): {result:.6f}")

    print("✓ 单像素输入测试通过\n")
    return True

def test_minimal_dimensions():
    """测试最小维度"""
    print("=== 测试最小维度 ===")

    # 测试各种最小维度组合
    test_cases = [
        (1, 1, 1, 1, "1x1x1x1"),
        (1, 1, 2, 1, "1x1x2x1"),
        (1, 1, 1, 2, "1x1x1x2"),
        (1, 2, 1, 1, "1x2x1x1"),
        (2, 1, 1, 1, "2x1x1x1"),
    ]

    for batch, channels, height, width, desc in test_cases:
        print(f"测试 {desc}:")

        # 创建输入数据
        input_data = np.random.randn(batch, channels, height, width).astype(np.float32)

        # 创建网格
        grid_data = np.random.randn(batch, 2, 2, 2).astype(np.float32)
        grid_data = np.clip(grid_data, -1, 1)

        try:
            output = grid_sampler_cuda.grid_sampler(
                input_data, grid_data, mode="bilinear"
            )
            print(f"  输出形状: {output.shape}")
            print(f"  ✓ 成功")
        except Exception as e:
            print(f"  ❌ 失败: {e}")
            return False

    print("✓ 最小维度测试通过\n")
    return True

def test_coordinate_edge_cases():
    """测试坐标边界情况"""
    print("=== 测试坐标边界情况 ===")

    input_data = np.array([[[
        [1.0, 2.0],
        [3.0, 4.0]
    ]]], dtype=np.float32)  # [1, 1, 2, 2]

    # 测试边界坐标
    edge_cases = [
        # (x, y, description)
        (-1.0, -1.0, "精确边界-左上"),
        (1.0, -1.0, "精确边界-右上"),
        (-1.0, 1.0, "精确边界-左下"),
        (1.0, 1.0, "精确边界-右下"),
        (0.0, 0.0, "中心点"),
        (-1.0001, -1.0001, "略超出边界"),
        (1.0001, 1.0001, "略超出边界"),
        (0.0, -1.0, "上边界中心"),
        (0.0, 1.0, "下边界中心"),
        (-1.0, 0.0, "左边界中心"),
        (1.0, 0.0, "右边界中心"),
    ]

    for x, y, desc in edge_cases:
        grid = np.array([[[[x, y]]]], dtype=np.float32)

        # 测试不同模式
        for mode in ["bilinear", "nearest"]:
            for padding_mode in ["zeros", "border", "reflection"]:
                try:
                    output = grid_sampler_cuda.grid_sampler(
                        input_data, grid, mode=mode, padding_mode=padding_mode
                    )
                    result = output[0, 0, 0, 0]
                    print(f"  {desc} ({mode}+{padding_mode}): {result:.6f}")
                except Exception as e:
                    print(f"  ❌ {desc} ({mode}+{padding_mode}): {e}")
                    return False

    print("✓ 坐标边界情况测试通过\n")
    return True

def test_nan_inf_handling():
    """测试NaN和无穷大处理"""
    print("=== 测试NaN和无穷大处理 ===")

    # 测试包含NaN的输入
    print("测试NaN输入:")
    input_data = np.array([[[
        [1.0, 2.0],
        [3.0, np.nan]
    ]]], dtype=np.float32)
    grid = np.array([[[[-1.0, -1.0]]]], dtype=np.float32)

    try:
        output = grid_sampler_cuda.grid_sampler(input_data, grid, mode="bilinear")
        result = output[0, 0, 0, 0]
        print(f"  NaN输入结果: {result} (isnan: {np.isnan(result)})")
    except Exception as e:
        print(f"  NaN输入异常: {e}")

    # 测试包含无穷大的输入
    print("测试无穷大输入:")
    input_data = np.array([[[
        [1.0, 2.0],
        [3.0, np.inf]
    ]]], dtype=np.float32)

    try:
        output = grid_sampler_cuda.grid_sampler(input_data, grid, mode="bilinear")
        result = output[0, 0, 0, 0]
        print(f"  无穷大输入结果: {result} (isinf: {np.isinf(result)})")
    except Exception as e:
        print(f"  无穷大输入异常: {e}")

    # 测试包含负无穷大的输入
    print("测试负无穷大输入:")
    input_data = np.array([[[
        [1.0, 2.0],
        [3.0, -np.inf]
    ]]], dtype=np.float32)

    try:
        output = grid_sampler_cuda.grid_sampler(input_data, grid, mode="bilinear")
        result = output[0, 0, 0, 0]
        print(f"  负无穷大输入结果: {result} (isinf: {np.isinf(result)})")
    except Exception as e:
        print(f"  负无穷大输入异常: {e}")

    print("✓ NaN和无穷大处理测试通过\n")
    return True

def test_extreme_coordinates():
    """测试极端坐标值"""
    print("=== 测试极端坐标值 ===")

    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]

    # 测试极端坐标
    extreme_cases = [
        (1e6, 1e6, "极大正值"),
        (-1e6, -1e6, "极大负值"),
        (1e-6, 1e-6, "极小正值"),
        (-1e-6, -1e-6, "极小负值"),
        (0.0, 0.0, "零值"),
    ]

    for x, y, desc in extreme_cases:
        grid = np.array([[[[x, y]]]], dtype=np.float32)

        try:
            output = grid_sampler_cuda.grid_sampler(
                input_data, grid, mode="bilinear", padding_mode="zeros"
            )
            result = output[0, 0, 0, 0]
            print(f"  {desc}: {result:.6f}")
        except Exception as e:
            print(f"  ❌ {desc}: {e}")
            return False

    print("✓ 极端坐标值测试通过\n")
    return True

def test_memory_boundary():
    """测试内存边界情况"""
    print("=== 测试内存边界情况 ===")

    # 测试大尺寸输入
    print("测试大尺寸输入:")
    try:
        # 创建较大的输入数据
        batch_size = 2
        channels = 16
        height = 64
        width = 64

        input_data = np.random.randn(batch_size, channels, height, width).astype(np.float32)
        grid_data = np.random.randn(batch_size, 32, 32, 2).astype(np.float32)
        grid_data = np.clip(grid_data, -1, 1)

        output = grid_sampler_cuda.grid_sampler(input_data, grid_data, mode="bilinear")
        print(f"  大尺寸输入成功: {input_data.shape} -> {output.shape}")
    except Exception as e:
        print(f"  ❌ 大尺寸输入失败: {e}")
        return False

    # 测试大批次
    print("测试大批次:")
    try:
        batch_size = 8
        input_data = np.random.randn(batch_size, 3, 16, 16).astype(np.float32)
        grid_data = np.random.randn(batch_size, 8, 8, 2).astype(np.float32)
        grid_data = np.clip(grid_data, -1, 1)

        output = grid_sampler_cuda.grid_sampler(input_data, grid_data, mode="bilinear")
        print(f"  大批次成功: {input_data.shape} -> {output.shape}")
    except Exception as e:
        print(f"  ❌ 大批次失败: {e}")
        return False

    print("✓ 内存边界情况测试通过\n")
    return True

def test_precision_limits():
    """测试精度极限"""
    print("=== 测试精度极限 ===")

    # 创建高精度测试数据
    input_data = np.array([[[
        [1.0, 1.000001],
        [1.000001, 1.0]
    ]]], dtype=np.float32)

    # 测试非常接近的坐标
    precision_cases = [
        (0.0, 0.0, "中心"),
        (0.000001, 0.000001, "微小偏移"),
        (-0.000001, -0.000001, "微小负偏移"),
    ]

    for x, y, desc in precision_cases:
        grid = np.array([[[[x, y]]]], dtype=np.float32)

        # 多次运行，检查结果稳定性
        results = []
        for _ in range(5):
            output = grid_sampler_cuda.grid_sampler(
                input_data, grid, mode="bilinear"
            )
            results.append(output[0, 0, 0, 0])

        # 检查结果稳定性
        max_diff = max(results) - min(results)
        avg_result = np.mean(results)

        print(f"  {desc}: 平均={avg_result:.8f}, 最大差异={max_diff:.2e}")

        if max_diff > 1e-6:
            print(f"    ⚠️  精度可能不够稳定")

    print("✓ 精度极限测试通过\n")
    return True

def test_special_geometric_cases():
    """测试特殊几何情况"""
    print("=== 测试特殊几何情况 ===")

    # 创建测试数据
    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]

    # 测试特殊几何变换
    geometric_cases = [
        # 恒等变换
        ([[[-1, -1], [1, -1]], [[-1, 1], [1, 1]]], "恒等变换"),
        # 90度旋转
        ([[[-1, 1], [-1, -1]], [[1, 1], [1, -1]]], "90度旋转"),
        # 180度旋转
        ([[[1, 1], [-1, 1]], [[1, -1], [-1, -1]]], "180度旋转"),
        # 270度旋转
        ([[[1, -1], [1, 1]], [[-1, -1], [-1, 1]]], "270度旋转"),
        # 水平翻转
        ([[[1, -1], [-1, -1]], [[1, 1], [-1, 1]]], "水平翻转"),
        # 垂直翻转
        ([[[-1, 1], [1, 1]], [[-1, -1], [1, -1]]], "垂直翻转"),
    ]

    for grid_coords, desc in geometric_cases:
        grid = np.array([grid_coords], dtype=np.float32)

        try:
            output = grid_sampler_cuda.grid_sampler(
                input_data, grid, mode="bilinear", align_corners=True
            )
            print(f"  {desc}: 成功")
            print(f"    输出: {output[0, 0].flatten()}")
        except Exception as e:
            print(f"  ❌ {desc}: {e}")
            return False

    print("✓ 特殊几何情况测试通过\n")
    return True

def test_error_recovery():
    """测试错误恢复"""
    print("=== 测试错误恢复 ===")

    # 测试在错误输入后是否能正常恢复
    print("测试错误恢复:")

    # 先进行正常操作
    input_data = np.random.randn(1, 3, 4, 4).astype(np.float32)
    grid_data = np.random.randn(1, 2, 2, 2).astype(np.float32)
    grid_data = np.clip(grid_data, -1, 1)

    try:
        output1 = grid_sampler_cuda.grid_sampler(input_data, grid_data, mode="bilinear")
        print("  正常操作成功")
    except Exception as e:
        print(f"  ❌ 正常操作失败: {e}")
        return False

    # 尝试错误操作（应该失败）
    try:
        wrong_input = np.random.randn(1, 3, 4, 4, 4).astype(np.float32)  # 错误的5D
        grid_sampler_cuda.grid_sampler(wrong_input, grid_data, mode="bilinear")
        print("  ❌ 错误操作应该失败但没有失败")
        return False
    except Exception as e:
        print(f"  错误操作正确失败: {e}")

    # 再次进行正常操作（测试恢复）
    try:
        output2 = grid_sampler_cuda.grid_sampler(input_data, grid_data, mode="bilinear")
        print("  错误后恢复成功")

        # 检查结果是否一致
        if np.allclose(output1, output2):
            print("  恢复后结果一致")
        else:
            print("  ⚠️  恢复后结果不一致")
    except Exception as e:
        print(f"  ❌ 错误后恢复失败: {e}")
        return False

    print("✓ 错误恢复测试通过\n")
    return True

def main():
    """运行所有边界情况测试"""
    print("开始运行Grid Sampler CUDA边界情况测试...\n")

    test_functions = [
        test_single_pixel_input,
        test_minimal_dimensions,
        test_coordinate_edge_cases,
        test_nan_inf_handling,
        test_extreme_coordinates,
        test_memory_boundary,
        test_precision_limits,
        test_special_geometric_cases,
        test_error_recovery,
    ]

    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"边界情况测试结果: {passed}/{total} 通过")

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

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