test_extreme_values.py

#!/usr/bin/env python3
"""
Grid Sampler CUDA 极端值和累积误差测试
测试极小值、极大值、NaN、无穷大等极端情况以及累积误差
"""

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_extreme_small_values():
    """测试极小值处理"""
    print("=== 测试极小值处理 ===")

    # 测试不同的极小值
    extreme_values = [
        1e-10,  # 极小正值
        -1e-10, # 极小负值
        1e-20,  # 极极小正值
        -1e-20, # 极极小负值
        0.0,    # 零值
    ]

    for val in extreme_values:
        print(f"\n测试值: {val}")

        # 创建包含极小值的输入
        input_data = np.array([[[
            [val, val * 2],
            [val * 3, val * 4]
        ]]], dtype=np.float32)

        # 测试不同的网格坐标
        test_coords = [
            (-1.0, -1.0, "左上角"),
            (0.0, 0.0, "中心"),
            (1.0, 1.0, "右下角"),
        ]

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

            try:
                output = grid_sampler_cuda.grid_sampler(
                    input_data, grid, mode="bilinear"
                )
                result = output[0, 0, 0, 0]

                # 检查结果是否合理
                if np.isnan(result):
                    print(f"  ❌ {desc}: 结果为NaN")
                    return False
                elif np.isinf(result):
                    print(f"  ❌ {desc}: 结果为无穷大")
                    return False
                else:
                    print(f"  ✓ {desc}: {result:.2e}")

            except Exception as e:
                print(f"  ❌ {desc}: 异常 {e}")
                return False

    print("✓ 极小值处理测试通过\n")
    return True

def test_extreme_large_values():
    """测试极大值处理"""
    print("=== 测试极大值处理 ===")

    # 测试不同的极大值
    extreme_values = [
        1e6,   # 大正值
        -1e6,  # 大负值
        1e10,  # 极大正值
        -1e10, # 极大负值
        1e20,  # 极极大正值
        -1e20, # 极极大负值
    ]

    for val in extreme_values:
        print(f"\n测试值: {val}")

        # 创建包含极大值的输入
        input_data = np.array([[[
            [val, val * 0.5],
            [val * 1.5, val * 2.0]
        ]]], dtype=np.float32)

        # 检查输入是否溢出
        if np.any(np.isinf(input_data)) or np.any(np.isnan(input_data)):
            print(f"  输入值溢出，跳过测试")
            continue

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

        try:
            output = grid_sampler_cuda.grid_sampler(
                input_data, grid, mode="bilinear"
            )
            result = output[0, 0, 0, 0]

            # 检查结果是否合理
            if np.isnan(result):
                print(f"  ❌ 结果为NaN")
                return False
            elif np.isinf(result):
                print(f"  ⚠️  结果为无穷大（可能是预期的）")
            else:
                print(f"  ✓ 结果: {result:.2e}")

        except Exception as e:
            print(f"  ❌ 异常: {e}")
            return False

    print("✓ 极大值处理测试通过\n")
    return True

def test_nan_handling():
    """测试NaN处理"""
    print("=== 测试NaN处理 ===")

    # 测试包含NaN的输入
    nan_cases = [
        (np.nan, 1.0, "左上角NaN"),
        (1.0, np.nan, "右上角NaN"),
        (np.nan, np.nan, "全部NaN"),
        (1.0, 2.0, "无NaN"),
    ]

    for val1, val2, desc in nan_cases:
        print(f"\n测试 {desc}:")

        input_data = np.array([[[
            [val1, val2],
            [3.0, 4.0]
        ]]], dtype=np.float32)

        # 测试不同的网格坐标
        test_coords = [
            (-1.0, -1.0, "左上角"),
            (0.0, 0.0, "中心"),
            (1.0, 1.0, "右下角"),
        ]

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

            try:
                output = grid_sampler_cuda.grid_sampler(
                    input_data, grid, mode="bilinear"
                )
                result = output[0, 0, 0, 0]

                print(f"  {coord_desc}: {result} (isnan: {np.isnan(result)})")

                # 如果输入包含NaN，输出可能也是NaN（这是合理的）
                if np.isnan(val1) or np.isnan(val2):
                    if not np.isnan(result):
                        print(f"    ⚠️  输入有NaN但输出不是NaN")
                else:
                    if np.isnan(result):
                        print(f"    ❌ 输入无NaN但输出是NaN")
                        return False

            except Exception as e:
                print(f"  ❌ {coord_desc}: 异常 {e}")
                return False

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

def test_infinity_handling():
    """测试无穷大处理"""
    print("=== 测试无穷大处理 ===")

    # 测试包含无穷大的输入
    inf_cases = [
        (np.inf, 1.0, "正无穷大"),
        (-np.inf, 1.0, "负无穷大"),
        (np.inf, np.inf, "双无穷大"),
        (1.0, 2.0, "无无穷大"),
    ]

    for val1, val2, desc in inf_cases:
        print(f"\n测试 {desc}:")

        input_data = np.array([[[
            [val1, val2],
            [3.0, 4.0]
        ]]], dtype=np.float32)

        # 测试中心点
        grid = np.array([[[[0.0, 0.0]]]], 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)}, isnan: {np.isnan(result)})")

            # 检查结果是否合理
            if np.isnan(result) and not (np.isnan(val1) or np.isnan(val2)):
                print(f"    ❌ 输入无NaN但输出是NaN")
                return False

        except Exception as e:
            print(f"  ❌ 异常: {e}")
            return False

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

def test_accumulated_error():
    """测试累积误差"""
    print("=== 测试累积误差 ===")

    # 创建初始数据
    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("初始数据:")
    print(input_data[0, 0])

    # 进行多次恒等变换
    current_data = input_data.copy()
    num_iterations = 20

    # 创建恒等变换网格
    grid = np.array([[[
        [[-1, -1], [1, -1]],
        [[-1, 1], [1, 1]]
    ]]], dtype=np.float32)  # 2x2输出

    print(f"\n进行{num_iterations}次恒等变换...")

    errors = []

    for i in range(num_iterations):
        # 将输出重新采样回原始大小
        # 创建从2x2到4x4的网格
        upscale_grid = np.array([[[
            [[-1, -1], [-1, 1], [1, -1], [1, 1]],
            [[-1, -1], [-1, 1], [1, -1], [1, 1]],
            [[-1, -1], [-1, 1], [1, -1], [1, 1]],
            [[-1, -1], [-1, 1], [1, -1], [1, 1]]
        ]]], dtype=np.float32)

        # 先下采样
        current_data = grid_sampler_cuda.grid_sampler(
            current_data, grid, mode="bilinear", align_corners=True
        )

        # 再上采样回原始大小
        current_data = grid_sampler_cuda.grid_sampler(
            current_data, upscale_grid, mode="bilinear", align_corners=True
        )

        # 计算与原始数据的误差
        error = np.mean(np.abs(current_data - input_data))
        errors.append(error)

        if i % 5 == 0:  # 每5次检查一次
            print(f"  第{i+1}次变换后误差: {error:.2e}")

    # 分析累积误差
    print(f"\n累积误差分析:")
    print(f"  初始误差: {errors[0]:.2e}")
    print(f"  最终误差: {errors[-1]:.2e}")
    print(f"  最大误差: {max(errors):.2e}")
    print(f"  误差增长率: {(errors[-1] - errors[0]) / errors[0] * 100:.2f}%")

    # 检查误差是否在可接受范围内
    if errors[-1] > 1e-3:
        print(f"  ❌ 累积误差过大: {errors[-1]:.2e}")
        return False
    else:
        print(f"  ✓ 累积误差在可接受范围内")

    print("✓ 累积误差测试通过\n")
    return True

def test_precision_degradation():
    """测试精度退化"""
    print("=== 测试精度退化 ===")

    # 创建高精度测试数据
    input_data = np.array([[[
        [1.0, 1.000001, 1.000002],
        [1.000003, 1.000004, 1.000005],
        [1.000006, 1.000007, 1.000008]
    ]]], dtype=np.float32)  # [1, 1, 3, 3]

    print("初始数据（高精度）:")
    print(input_data[0, 0])

    # 进行多次变换
    current_data = input_data.copy()
    num_iterations = 10

    # 创建变换网格
    grid = np.array([[[
        [[-1, -1], [1, -1]],
        [[-1, 1], [1, 1]]
    ]]], dtype=np.float32)  # 2x2输出

    print(f"\n进行{num_iterations}次变换...")

    precision_losses = []

    for i in range(num_iterations):
        # 下采样
        current_data = grid_sampler_cuda.grid_sampler(
            current_data, grid, mode="bilinear", align_corners=True
        )

        # 上采样回原始大小
        upscale_grid = np.array([[[
            [[-1, -1], [-1, 1], [1, -1], [1, 1]],
            [[-1, -1], [-1, 1], [1, -1], [1, 1]],
            [[-1, -1], [-1, 1], [1, -1], [1, 1]]
        ]]], dtype=np.float32)

        current_data = grid_sampler_cuda.grid_sampler(
            current_data, upscale_grid, mode="bilinear", align_corners=True
        )

        # 计算精度损失
        precision_loss = np.mean(np.abs(current_data - input_data))
        precision_losses.append(precision_loss)

        if i % 2 == 0:  # 每2次检查一次
            print(f"  第{i+1}次变换后精度损失: {precision_loss:.2e}")

    # 分析精度退化
    print(f"\n精度退化分析:")
    print(f"  初始精度损失: {precision_losses[0]:.2e}")
    print(f"  最终精度损失: {precision_losses[-1]:.2e}")
    print(f"  最大精度损失: {max(precision_losses):.2e}")

    # 检查精度退化是否在可接受范围内
    if precision_losses[-1] > 1e-4:
        print(f"  ⚠️  精度退化较大: {precision_losses[-1]:.2e}")
    else:
        print(f"  ✓ 精度退化在可接受范围内")

    print("✓ 精度退化测试通过\n")
    return True

def test_numerical_stability():
    """测试数值稳定性"""
    print("=== 测试数值稳定性 ===")

    # 测试不同数据范围的稳定性
    test_ranges = [
        (1e-10, 1e-5, "极小值范围"),
        (1e-5, 1e-2, "小值范围"),
        (1e-2, 1e2, "中等值范围"),
        (1e2, 1e5, "大值范围"),
    ]

    for min_val, max_val, desc in test_ranges:
        print(f"\n测试{desc}: [{min_val:.0e}, {max_val:.0e}]")

        # 创建测试数据
        input_data = np.random.uniform(min_val, max_val, (1, 1, 4, 4)).astype(np.float32)
        grid_data = np.random.uniform(-1, 1, (1, 2, 2, 2)).astype(np.float32)

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

        # 检查结果一致性
        max_variation = 0
        for i in range(1, len(results)):
            variation = np.max(np.abs(results[i] - results[0]))
            max_variation = max(max_variation, variation)

        print(f"  最大变化: {max_variation:.2e}")

        if max_variation > 1e-6:
            print(f"  ❌ 数值稳定性不足")
            return False
        else:
            print(f"  ✓ 数值稳定")

    print("✓ 数值稳定性测试通过\n")
    return True

def test_edge_case_combinations():
    """测试边界情况组合"""
    print("=== 测试边界情况组合 ===")

    # 测试各种边界情况的组合
    test_cases = [
        # (input_data, grid_data, description)
        (np.array([[[[np.nan, 1.0], [2.0, 3.0]]]], dtype=np.float32),
         np.array([[[[-1.0, -1.0]]]], dtype=np.float32),
         "NaN + 边界坐标"),

        (np.array([[[[np.inf, 1.0], [2.0, 3.0]]]], dtype=np.float32),
         np.array([[[[0.0, 0.0]]]], dtype=np.float32),
         "无穷大 + 中心坐标"),

        (np.array([[[[1e-10, 1e-10], [1e-10, 1e-10]]]], dtype=np.float32),
         np.array([[[[1.0, 1.0]]]], dtype=np.float32),
         "极小值 + 边界坐标"),

        (np.array([[[[1e10, 1e10], [1e10, 1e10]]]], dtype=np.float32),
         np.array([[[[-1.0, -1.0]]]], dtype=np.float32),
         "极大值 + 边界坐标"),
    ]

    for input_data, grid_data, desc in test_cases:
        print(f"\n测试 {desc}:")

        try:
            output = grid_sampler_cuda.grid_sampler(
                input_data, grid_data, mode="bilinear"
            )
            result = output[0, 0, 0, 0]

            print(f"  结果: {result} (isnan: {np.isnan(result)}, isinf: {np.isinf(result)})")

            # 检查结果是否合理
            if np.isnan(result) and not np.any(np.isnan(input_data)):
                print(f"  ❌ 输入无NaN但输出是NaN")
                return False

        except Exception as e:
            print(f"  ❌ 异常: {e}")
            return False

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

def main():
    """运行所有极端值和累积误差测试"""
    print("开始运行Grid Sampler CUDA极端值和累积误差测试...\n")

    test_functions = [
        test_extreme_small_values,
        test_extreme_large_values,
        test_nan_handling,
        test_infinity_handling,
        test_accumulated_error,
        test_precision_degradation,
        test_numerical_stability,
        test_edge_case_combinations,
    ]

    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())