Compare two numpy ndarrays

compare_two_numpy_darrays.py

import pickle

import numpy as np


def load_outputs(file1, file2):
    """Load both outputs"""
    with open(file1, "rb") as f:
        original = pickle.load(f)
    with open(file2, "rb") as f:
        refactored = pickle.load(f)
    return original, refactored


def compare_arrays(arr1, arr2, name, tolerance=1e-10):
    """Comprehensive array comparison"""
    print(f"\n{'='*80}")
    print(f"Comparing: {name}")
    print(f"{'='*80}")

    # Basic shape check
    if arr1.shape != arr2.shape:
        print("❌ SHAPE MISMATCH!")
        print(f"   Original: {arr1.shape}")
        print(f"   Refactored: {arr2.shape}")
        return

    print(f"✓ Shape matches: {arr1.shape}")

    # Check if arrays are identical
    if np.array_equal(arr1, arr2):
        print("✓ Arrays are IDENTICAL")
        return

    # Check with tolerance
    close = np.allclose(arr1, arr2, rtol=tolerance, atol=tolerance)
    if close:
        print(f"✓ Arrays are equal within tolerance {tolerance}")
        max_diff = np.max(np.abs(arr1 - arr2))
        print(f"   Max absolute difference: {max_diff:.2e}")
        return

    # Find differences
    diff = np.abs(arr1 - arr2)
    different_mask = diff > tolerance
    num_different = np.sum(different_mask)
    total_elements = arr1.size

    print("❌ DIFFERENCES FOUND!")
    print(
        f"   Different elements: {num_different} / {total_elements} ({100*num_different/total_elements:.2f}%)"
    )
    print(f"   Max absolute difference: {np.max(diff):.6e}")
    print(f"   Mean absolute difference: {np.mean(diff):.6e}")
    print(f"   Median absolute difference: {np.median(diff):.6e}")

    # Find where differences start (first divergence point)
    if len(arr1.shape) == 3:  # Color array: (T, L, D)
        T, L, D = np.where(different_mask)
        if len(T) > 0:
            first_idx = np.argmin(T * 1000 + L * 100 + D)  # Sort by T, then L, then D
            print("\n   First divergence at:")
            print(f"      T={T[first_idx]}, L={L[first_idx]}, D={D[first_idx]}")
            print(
                f"      Original value: {arr1[T[first_idx], L[first_idx], D[first_idx]]:.10f}"
            )
            print(
                f"      Refactored value: {arr2[T[first_idx], L[first_idx], D[first_idx]]:.10f}"
            )
            print(
                f"      Difference: {diff[T[first_idx], L[first_idx], D[first_idx]]:.10e}"
            )

            # Show first 10 divergence points
            print("\n   First 10 divergence points:")
            for i in range(min(10, len(T))):
                idx = np.argmin(T * 1000 + L * 100 + D)
                t, l, d = T[idx], L[idx], D[idx]
                print(
                    f"      T={t:3d}, L={l}, D={d} | "
                    f"Original: {arr1[t,l,d]:10.6f} | "
                    f"Refactored: {arr2[t,l,d]:10.6f} | "
                    f"Diff: {diff[t,l,d]:.6e}"
                )
                # Remove this point for next iteration
                T = np.delete(T, idx)
                L = np.delete(L, idx)
                D = np.delete(D, idx)

    elif len(arr1.shape) == 2:  # LeafColor array: (T, L)
        T, L = np.where(different_mask)
        if len(T) > 0:
            first_idx = np.argmin(T * 100 + L)
            print("\n   First divergence at:")
            print(f"      T={T[first_idx]}, L={L[first_idx]}")
            print(f"      Original value: {arr1[T[first_idx], L[first_idx]]:.10f}")
            print(f"      Refactored value: {arr2[T[first_idx], L[first_idx]]:.10f}")
            print(f"      Difference: {diff[T[first_idx], L[first_idx]]:.10e}")

            # Show first 10 divergence points
            print("\n   First 10 divergence points:")
            for i in range(min(10, len(T))):
                idx = np.argmin(T * 100 + L)
                t, l = T[idx], L[idx]
                print(
                    f"      T={t:3d}, L={l} | "
                    f"Original: {arr1[t,l]:10.6f} | "
                    f"Refactored: {arr2[t,l]:10.6f} | "
                    f"Diff: {diff[t,l]:.6e}"
                )
                T = np.delete(T, idx)
                L = np.delete(L, idx)

    return {
        "num_different": num_different,
        "max_diff": np.max(diff),
        "mean_diff": np.mean(diff),
        "different_mask": different_mask,
    }


def generate_report(original, refactored):
    """Generate comprehensive comparison report"""
    results = {}

    # Compare Color array
    results["Color"] = compare_arrays(
        original["Color"], refactored["Color"], "Color (LDColor)"
    )

    # Compare LeafColor array
    results["LeafColor"] = compare_arrays(
        original["LeafColor"], refactored["LeafColor"], "LeafColor (LColor)"
    )

    return results


if __name__ == "__main__":
    # Load outputs
    original, refactored = load_outputs("original_output.pkl", "refactored_output.pkl")

    # Generate comparison report
    results = generate_report(original, refactored)

    print("\n" + "=" * 80)
    print("SUMMARY")
    print("=" * 80)

    if results["Color"] is None and results["LeafColor"] is None:
        print("✓ ALL OUTPUTS MATCH PERFECTLY!")
    else:
        print("❌ Differences detected - see details above")