debug_fftconvolve.py

#!/usr/bin/env python3
"""
Debug script to identify the source of NaN values in fftconvolve calls.
This will help track down the RuntimeWarning about invalid values encountered in multiply.
"""

import numpy as np
from scipy import signal
import warnings

def debug_stress_to_current(fine_time, fine_stress, tau_arr, k_arr, g=0.4, h=1):
    """
    Debug version of stress_to_current function to identify NaN sources.
    """
    print("=== DEBUGGING stress_to_current ===")
    
    # Check input arrays for NaN/inf values
    print(f"fine_time shape: {fine_time.shape}")
    print(f"fine_stress shape: {fine_stress.shape}")
    print(f"tau_arr shape: {tau_arr.shape}")
    print(f"k_arr shape: {k_arr.shape}")
    
    print(f"fine_time has NaN: {np.any(np.isnan(fine_time))}")
    print(f"fine_stress has NaN: {np.any(np.isnan(fine_stress))}")
    print(f"tau_arr has NaN: {np.any(np.isnan(tau_arr))}")
    print(f"k_arr has NaN: {np.any(np.isnan(k_arr))}")
    
    print(f"fine_time has inf: {np.any(np.isinf(fine_time))}")
    print(f"fine_stress has inf: {np.any(np.isinf(fine_stress))}")
    print(f"tau_arr has inf: {np.any(np.isinf(tau_arr))}")
    print(f"k_arr has inf: {np.any(np.isinf(k_arr))}")
    
    # Check for zeros in tau_arr (which would cause division by zero)
    print(f"tau_arr has zeros: {np.any(tau_arr == 0)}")
    print(f"tau_arr values: {tau_arr}")
    
    # Compute derivatives
    first_derivative = np.gradient(fine_stress)
    second_derivative = np.gradient(np.gradient(fine_stress))
    
    print(f"first_derivative has NaN: {np.any(np.isnan(first_derivative))}")
    print(f"second_derivative has NaN: {np.any(np.isnan(second_derivative))}")
    print(f"first_derivative has inf: {np.any(np.isinf(first_derivative))}")
    print(f"second_derivative has inf: {np.any(np.isinf(second_derivative))}")
    
    # Compute ds
    ds = np.r_[g * first_derivative + h * second_derivative].reshape(-1, 1)
    print(f"ds has NaN: {np.any(np.isnan(ds))}")
    print(f"ds has inf: {np.any(np.isinf(ds))}")
    print(f"ds shape: {ds.shape}")
    
    # Compute kernels
    k1 = k_arr[0] * np.exp(-fine_time[None].T / tau_arr[0])
    print(f"k1 has NaN: {np.any(np.isnan(k1))}")
    print(f"k1 has inf: {np.any(np.isinf(k1))}")
    print(f"k1 shape: {k1.shape}")
    
    K_peak, K_steady = k_arr[3], tau_arr[3]
    k2 = k_arr[1] * (K_peak * np.exp(-fine_time[None].T / tau_arr[1]) + K_steady)
    print(f"k2 has NaN: {np.any(np.isnan(k2))}")
    print(f"k2 has inf: {np.any(np.isinf(k2))}")
    print(f"k2 shape: {k2.shape}")
    
    k3 = k_arr[2] * np.exp(-fine_time[None].T / tau_arr[2])
    print(f"k3 has NaN: {np.any(np.isnan(k3))}")
    print(f"k3 has inf: {np.any(np.isinf(k3))}")
    print(f"k3 shape: {k3.shape}")
    
    # Test fftconvolve with error handling
    print("\n=== Testing fftconvolve calls ===")
    
    try:
        uRI = signal.fftconvolve(k1, ds, mode='full')[:fine_time.size]
        print("✓ uRI computation successful")
    except Exception as e:
        print(f"✗ uRI computation failed: {e}")
        return None
    
    try:
        uSI = signal.fftconvolve(k2, ds, mode='full')[:fine_time.size]
        print("✓ uSI computation successful")
    except Exception as e:
        print(f"✗ uSI computation failed: {e}")
        return None
    
    try:
        uUSI = signal.fftconvolve(k3, ds, mode='full')[:fine_time.size]
        print("✓ uUSI computation successful")
    except Exception as e:
        print(f"✗ uUSI computation failed: {e}")
        return None
    
    print("=== All convolutions successful ===")
    return np.column_stack([uRI, uSI, uUSI])

def test_with_sample_data():
    """
    Test the debug function with sample data to see if we can reproduce the issue.
    """
    print("=== Testing with sample data ===")
    
    # Create sample data
    fine_time = np.linspace(0, 100, 1000)
    fine_stress = np.sin(fine_time * 0.1) + np.random.normal(0, 0.1, 1000)
    tau_arr = np.array([10.0, 50.0, 100.0, 200.0])
    k_arr = np.array([1.0, 0.5, 0.3, 0.1])
    
    result = debug_stress_to_current(fine_time, fine_stress, tau_arr, k_arr)
    
    if result is not None:
        print("✓ Sample test successful")
    else:
        print("✗ Sample test failed")

if __name__ == "__main__":
    # Enable warnings to see the RuntimeWarning
    warnings.filterwarnings('always')
    
    test_with_sample_data()