app.py

import gradio as gr
import numpy as np
import scipy.sparse as sparse
import time
import os
import shutil
import math
import sys
from pathlib import Path

# Assuming flex_chunk.py and matrix_multiply.py are in the same directory
from flex_chunk import FlexChunk, save_chunk, load_chunk
from matrix_multiply import prepare_chunks, load_chunks, matrix_vector_multiply

# --- Matrix Generation (copied from test_vs_scipy.py) ---

def generate_sparse_matrix(size, density, challenging=False):
    """
    Generate a sparse test matrix with optional challenging patterns.
    
    Args:
        size: Matrix size (n x n)
        density: Target density
        challenging: Whether to include challenging patterns and extreme values
        
    Returns:
        A scipy.sparse.csr_matrix
    """
    # Calculate number of non-zeros
    nnz = int(size * size * density)
    if nnz == 0: # Ensure at least one non-zero element if density is very low
        nnz = 1 
        
    if not challenging:
        # Simple random matrix
        rows = np.random.randint(0, size, nnz)
        cols = np.random.randint(0, size, nnz)
        data = np.random.rand(nnz)
        # Ensure the matrix actually has the specified size if nnz is small
        if nnz < size:
             # Add diagonal elements to ensure size
             diag_indices = np.arange(min(nnz, size))
             rows = np.concatenate([rows, diag_indices])
             cols = np.concatenate([cols, diag_indices])
             data = np.concatenate([data, np.ones(len(diag_indices))]) # Use 1 for diagonal
        
        matrix = sparse.csr_matrix((data, (rows, cols)), shape=(size, size))
        matrix.sum_duplicates() # Consolidate duplicate entries
        return matrix
    
    # --- Challenging matrix with specific patterns ---
    # Base random matrix (80% of non-zeros)
    base_nnz = int(nnz * 0.8)
    rows = np.random.randint(0, size, base_nnz)
    cols = np.random.randint(0, size, base_nnz)
    data = np.random.rand(base_nnz)
    
    # Add diagonal elements (10% of non-zeros)
    diag_nnz = int(nnz * 0.1)
    diag_indices = np.random.choice(size, diag_nnz, replace=False)
    
    # Add extreme values (10% of non-zeros)
    extreme_nnz = max(0, nnz - base_nnz - diag_nnz) # Ensure non-negative
    extreme_rows = np.random.randint(0, size, extreme_nnz)
    extreme_cols = np.random.randint(0, size, extreme_nnz)
    
    # Mix of very large and very small values
    extreme_data = np.concatenate([
        np.random.uniform(1e6, 1e9, extreme_nnz // 2),
        np.random.uniform(1e-9, 1e-6, extreme_nnz - extreme_nnz // 2)
    ]) if extreme_nnz > 0 else np.array([])
    if extreme_nnz > 0:
        np.random.shuffle(extreme_data)
    
    # Combine all components
    all_rows = np.concatenate([rows, diag_indices, extreme_rows])
    all_cols = np.concatenate([cols, diag_indices, extreme_cols])
    all_data = np.concatenate([data, np.random.rand(diag_nnz), extreme_data])
    
    matrix = sparse.csr_matrix((all_data, (all_rows, all_cols)), shape=(size, size))
    matrix.sum_duplicates() # Consolidate duplicate entries
    return matrix

# --- Benchmark Function (Placeholder) ---

def run_benchmark(size, density, num_chunks, challenging, progress=gr.Progress()):
    # This function will contain the main logic from test_vs_scipy.py
    # Adapted for Gradio inputs and outputs
    progress(0, desc="Starting Benchmark...")
    time.sleep(1) # Placeholder
    
    # 1. Setup storage
    storage_dir = Path("./flex_chunk_temp_space")
    if storage_dir.exists():
        shutil.rmtree(storage_dir)
    storage_dir.mkdir(exist_ok=True)
    
    progress(0.1, desc="Generating Matrix...")
    # 2. Generate matrix and vector
    matrix = generate_sparse_matrix(size, density, challenging)
    vector = np.random.rand(size)
    actual_nnz = matrix.nnz
    actual_density = actual_nnz / (size * size) if size > 0 else 0
    
    matrix_info = f"Matrix: {size}x{size}, Target Density: {density:.6f}, Actual Density: {actual_density:.6f}, NNZ: {actual_nnz}"
    print(matrix_info) # For debugging in Hugging Face console
    
    # --- FlexChunk Run ---
    progress(0.2, desc="Preparing FlexChunks...")
    prepare_start = time.time()
    prepare_chunks(matrix, num_chunks, str(storage_dir), verbose=False)
    prepare_time = time.time() - prepare_start

    progress(0.4, desc="Loading FlexChunks...")
    load_start = time.time()
    chunks = load_chunks(str(storage_dir), verbose=False)
    load_time = time.time() - load_start
    
    progress(0.6, desc="Running FlexChunk SpMV...")
    flex_compute_start = time.time()
    flex_result = matrix_vector_multiply(chunks, vector, verbose=False)
    flex_compute_time = time.time() - flex_compute_start
    flex_total_time = load_time + flex_compute_time
    
    # --- SciPy Run ---
    progress(0.7, desc="Saving SciPy data...")
    scipy_temp_dir = storage_dir / "scipy_temp"
    scipy_temp_dir.mkdir(exist_ok=True)
    matrix_file = scipy_temp_dir / "matrix.npz"
    vector_file = scipy_temp_dir / "vector.npy"
    
    scipy_save_start = time.time()
    sparse.save_npz(matrix_file, matrix)
    np.save(vector_file, vector)
    scipy_save_time = time.time() - scipy_save_start
    
    progress(0.8, desc="Loading SciPy data...")
    scipy_load_start = time.time()
    loaded_matrix = sparse.load_npz(matrix_file)
    loaded_vector = np.load(vector_file)
    scipy_load_time = time.time() - scipy_load_start
    
    progress(0.9, desc="Running SciPy SpMV...")
    scipy_compute_start = time.time()
    scipy_result = loaded_matrix @ loaded_vector
    scipy_compute_time = time.time() - scipy_compute_start
    scipy_total_time = scipy_load_time + scipy_compute_time
    
    # --- Comparison ---
    progress(0.95, desc="Comparing results...")
    diff = np.abs(scipy_result - flex_result)
    max_diff = np.max(diff) if len(diff) > 0 else 0
    mean_diff = np.mean(diff) if len(diff) > 0 else 0
    is_close = np.allclose(scipy_result, flex_result, atol=1e-9) # Increased tolerance slightly
    comparison_result = f"✅ Results Match! (Max Diff: {max_diff:.2e}, Mean Diff: {mean_diff:.2e})" if is_close else f"❌ Results Differ! (Max Diff: {max_diff:.2e}, Mean Diff: {mean_diff:.2e})"

    # --- Cleanup ---
    shutil.rmtree(storage_dir)
    
    progress(1.0, desc="Benchmark Complete")
    
    # --- Format Output ---
    results_summary = f"""
{matrix_info}

**FlexChunk Performance:**
- Prepare Chunks Time: {prepare_time:.4f}s
- Load Chunks Time:    {load_time:.4f}s
- Compute Time:        {flex_compute_time:.4f}s
- **Total (Load+Compute): {flex_total_time:.4f}s**

**SciPy Performance (Out-of-Core Emulation):**
- Save Data Time:   {scipy_save_time:.4f}s (For reference)
- Load Data Time:   {scipy_load_time:.4f}s
- Compute Time:     {scipy_compute_time:.4f}s
- **Total (Load+Compute): {scipy_total_time:.4f}s**

**Comparison:**
{comparison_result}
"""
    
    return results_summary

# --- Gradio Interface ---

with gr.Blocks() as demo:
    gr.Markdown("""
    # FlexChunk: Out-of-Core Sparse Matrix-Vector Multiplication (SpMV) Demo
    
    This demo compares the performance of FlexChunk against standard SciPy for SpMV, 
    simulating an out-of-core scenario where the matrix doesn't fit entirely in memory.
    
    FlexChunk splits the matrix into smaller chunks, processing them sequentially to reduce peak memory usage. 
    SciPy performance includes the time to save and load the matrix from disk to mimic this out-of-core access.
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("**Benchmark Parameters**")
            size_input = gr.Slider(label="Matrix Size (N x N)", minimum=100, maximum=50000, value=10000, step=100)
            # Max density adjusted to prevent excessive nnz for large matrices in demo
            density_input = gr.Slider(label="Matrix Density", minimum=0.00001, maximum=0.01, value=0.0001, step=0.00001, format="%.5f") 
            chunks_input = gr.Slider(label="Number of Chunks", minimum=1, maximum=32, value=4, step=1)
            challenging_input = gr.Checkbox(label="Use Challenging Matrix (Extreme Values)", value=False)
            run_button = gr.Button("Run Benchmark", variant="primary")
            
        with gr.Column(scale=2):
            gr.Markdown("**Results**")
            output_textbox = gr.Markdown(label="Benchmark Summary")

    run_button.click(
        fn=run_benchmark, 
        inputs=[size_input, density_input, chunks_input, challenging_input], 
        outputs=[output_textbox]
    )
    
    gr.Markdown("--- Developed based on the [FlexChunk concept](https://www.lesswrong.com/posts/zpRhsdDkWygTDScxb/flexchunk-enabling-100m-100m-out-of-core-spmv-1-8-min-1-7-gb).")

# Launch the app
if __name__ == "__main__":
    demo.launch()