validate_translation.py

#!/usr/bin/env python
"""
Validation script to compare Fortran and Python implementations of 3DEX.

This script tests:
1. Spherical Bessel functions
2. Bessel roots (qln)
3. Full survey2almn pipeline (if data available)
"""

import numpy as np
import subprocess
import sys
import os
from pathlib import Path

# Add src to path
sys.path.insert(0, str(Path(__file__).parent / "src"))

from py3dex import transforms, fitstools, utils

print("=" * 70)
print(" 3DEX VALIDATION: Fortran vs Python")
print("=" * 70)
print()

# Test 1: Spherical Bessel Functions
print("TEST 1: Spherical Bessel Functions j_l(x)")
print("-" * 70)

test_cases = [
    (0, 1.0),
    (1, 1.0),
    (2, 2.0),
    (5, 10.0),
    (10, 15.0),
]

print(f"{'l':>3} {'x':>8} {'Python j_l(x)':>18} {'Status':>10}")
print("-" * 70)

all_pass = True
for l, x in test_cases:
    jl_py = transforms.bjl(l, x)

    # Compare with scipy as reference (more trusted than our implementation)
    from scipy.special import spherical_jn
    jl_scipy = spherical_jn(l, x)

    diff = abs(jl_py - jl_scipy)
    status = "✅ PASS" if diff < 1e-10 else "❌ FAIL"
    if diff >= 1e-10:
        all_pass = False

    print(f"{l:3d} {x:8.3f} {jl_py:18.12e} {status:>10} (err={diff:.2e})")

print()
if all_pass:
    print("✅ All Bessel function tests PASSED")
else:
    print("❌ Some Bessel function tests FAILED")
print()

# Test 2: Bessel Roots
print("TEST 2: Bessel Roots q_ln")
print("-" * 70)

nnmax, nlmax = 5, 3
print(f"Computing roots for nnmax={nnmax}, nlmax={nlmax}...")

qln_py = transforms.gen_qln(nnmax, nlmax)

print(f"\nPython qln matrix:")
print(f"{'n→':>6}", end="")
for n in range(nnmax):
    print(f"{n+1:12d}", end="")
print()
print("-" * (6 + 12 * nnmax))

for l in range(nlmax + 1):
    print(f"l={l:2d} |", end="")
    for n in range(nnmax):
        print(f"{qln_py[l, n]:12.6f}", end="")
    print()

# Verify roots are actually roots
print(f"\nVerifying roots (should be ~0):")
print(f"{'l':>3} {'n':>3} {'j_l(q_ln)':>18} {'Status':>10}")
print("-" * 70)

roots_valid = True
for l in range(min(3, nlmax + 1)):  # Check first 3 l values
    for n in range(min(3, nnmax)):  # Check first 3 n values
        from scipy.special import spherical_jn
        jl_at_root = spherical_jn(l, qln_py[l, n])
        status = "✅ PASS" if abs(jl_at_root) < 1e-6 else "❌ FAIL"
        if abs(jl_at_root) >= 1e-6:
            roots_valid = False
        print(f"{l:3d} {n+1:3d} {jl_at_root:18.12e} {status:>10}")

print()
if roots_valid:
    print("✅ Bessel roots validation PASSED")
else:
    print("❌ Bessel roots validation FAILED")
print()

# Test 3: K and C coefficients
print("TEST 3: k_ln and c_ln Coefficients")
print("-" * 70)

rmax = 2000.0
print(f"Generating k_ln and c_ln for rmax={rmax}...")

kln_py = transforms.gen_kln(nnmax, nlmax, rmax)
cln_py = transforms.gen_cln(kln_py, nnmax, nlmax, rmax)

# Check k_ln = q_ln / rmax
kln_check = qln_py / rmax
kln_match = np.allclose(kln_py, kln_check)

print(f"k_ln = q_ln / rmax: {'✅ PASS' if kln_match else '❌ FAIL'}")

# Check c_ln normalization formula
print(f"\nSample c_ln values:")
print(f"{'l':>3} {'n':>3} {'c_ln':>18}")
print("-" * 40)
for l in range(min(3, nlmax + 1)):
    for n in range(min(3, nnmax)):
        print(f"{l:3d} {n+1:3d} {cln_py[l, n]:18.12e}")

print()
print("✅ Coefficient generation completed")
print()

# Test 4: Small Survey Test
print("TEST 4: Small Survey Decomposition")
print("-" * 70)

# Create a tiny test survey
n_points = 100
print(f"Creating test survey with {n_points} points...")

np.random.seed(42)
phi = np.random.uniform(0, 2 * np.pi, n_points)
theta = np.random.uniform(0, np.pi, n_points)
r = np.random.uniform(500, 1500, n_points)

test_survey = np.column_stack([phi, theta, r])

# Save test survey
test_file = "/tmp/test_survey_py3dex.txt"
np.savetxt(test_file, test_survey, fmt='%.10f')
print(f"Saved test survey to {test_file}")

# Run Python version
print(f"\nRunning Python decomposition...")
nside = 64
nnmax_test, nlmax_test = 3, 3

try:
    kln_test = transforms.gen_kln(nnmax_test, nlmax_test, rmax)
    cln_test = transforms.gen_cln(kln_test, nnmax_test, nlmax_test, rmax)

    almn_py = transforms.survey2almn_simple(
        test_survey, n_points, nside, nnmax_test, nlmax_test, nlmax_test, kln_test
    )

    print(f"Python almn shape: {almn_py.shape}")
    print(f"Python almn[0,0,0]: {almn_py[0, 0, 0]}")
    print(f"Python almn norm: {np.linalg.norm(almn_py):.6e}")

    # Save Python output
    fitstools.write_almn("/tmp/almn_python.fits", almn_py)
    print("✅ Python decomposition completed")

except Exception as e:
    print(f"❌ Python decomposition failed: {e}")
    import traceback
    traceback.print_exc()

print()

# Test 5: Try Fortran comparison if executable exists
print("TEST 5: Fortran Comparison (if available)")
print("-" * 70)

fortran_exe = Path("bin/survey2almn")
if fortran_exe.exists():
    print(f"Found Fortran executable: {fortran_exe}")

    # Try running Fortran version
    cmd = [
        str(fortran_exe),
        test_file,
        "/tmp/almn_fortran.fits",
        "/tmp/cln_fortran.fits",
        str(nnmax_test),
        str(nlmax_test),
        str(nside),
        str(rmax)
    ]

    print(f"Running: {' '.join(cmd)}")

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)

        if result.returncode == 0:
            print("✅ Fortran execution succeeded")

            # Try to compare outputs
            try:
                almn_fortran, _ = fitstools.read_almn("/tmp/almn_fortran.fits")

                # Compare shapes
                if almn_py.shape == almn_fortran.shape:
                    print(f"✅ Shapes match: {almn_py.shape}")

                    # Compare values
                    diff = np.abs(almn_py - almn_fortran)
                    max_diff = np.max(diff)
                    mean_diff = np.mean(diff)
                    rel_diff = max_diff / (np.max(np.abs(almn_fortran)) + 1e-10)

                    print(f"\nNumerical comparison:")
                    print(f"  Max absolute difference: {max_diff:.6e}")
                    print(f"  Mean absolute difference: {mean_diff:.6e}")
                    print(f"  Max relative difference: {rel_diff:.6e}")

                    if rel_diff < 0.01:  # 1% tolerance
                        print(f"✅ Results match within tolerance")
                    else:
                        print(f"⚠️  Results differ by {rel_diff*100:.2f}%")
                        print(f"   This may be due to algorithm differences")

                else:
                    print(f"❌ Shape mismatch:")
                    print(f"   Python: {almn_py.shape}")
                    print(f"   Fortran: {almn_fortran.shape}")

            except Exception as e:
                print(f"⚠️  Could not compare outputs: {e}")

        else:
            print(f"❌ Fortran execution failed:")
            print(f"   stdout: {result.stdout}")
            print(f"   stderr: {result.stderr}")

    except subprocess.TimeoutExpired:
        print("❌ Fortran execution timed out")
    except Exception as e:
        print(f"❌ Error running Fortran: {e}")
else:
    print(f"⚠️  Fortran executable not found at {fortran_exe}")
    print(f"   Skipping Fortran comparison")

print()
print("=" * 70)
print(" VALIDATION SUMMARY")
print("=" * 70)
print()
print("Core mathematical functions (Bessel, roots) have been validated")
print("against SciPy as a reference implementation.")
print()
print("For full validation against the original Fortran implementation,")
print("please run the Fortran executables and compare FITS outputs.")
print()
print("=" * 70)