app.py

"""
NVIDIA CUDA-Q VQE Demo
=====================

This Gradio application demonstrates the Variational Quantum Eigensolver (VQE) 
algorithm using NVIDIA's CUDA Quantum library. The demo allows users to:

1. Select from available molecules
2. Adjust bond length parameters
3. Run VQE simulations with real-time visualization
4. View energy convergence and molecular structure in 3D

Running the Demo
---------------
1. Ensure you have all dependencies installed:
   ```bash
   pip install -r requirements.txt
   ```

2. Launch the application:
   ```bash
   python app.py
   ```

3. Open your browser to the local URL displayed in the terminal
   (typically http://localhost:7860)

Using the Interface
-----------------
- Select a molecule from the dropdown
- Adjust bond length using the slider (0.5 to 2.5 Å)
- Click "Run VQE Simulation" or watch real-time updates as you adjust parameters
- View results in three panels:
  * Simulation Results: Shows final energy and optimization status
  * VQE Convergence: Plots energy vs. iteration
  * Molecule Visualization: 3D representation of the molecule

GPU Acceleration
--------------
The app automatically detects if NVIDIA GPUs are available:
- With GPU: Uses CUDA-Q's nvidia backend for acceleration
- Without GPU: Falls back to CPU-based quantum simulation

References
---------
- Installation: See Cuda-Q_install.md
- VQE Implementation: See VQE_Example.md
- Project Structure: See quantum-demo-blueprint.md
"""

import spaces
from gradio_client import Client
from gradio.routes import Request as gr_Request
import subprocess
import sys
import json
import os
import logging
from logging.handlers import RotatingFileHandler
import gradio as gr
import numpy as np
from quantum_utils import run_vqe_simulation
from visualization import plot_convergence, create_molecule_viewer, format_molecule_params
import plotly.graph_objects as go

# Configure logging
os.makedirs('logs', exist_ok=True)

# Create formatters
file_formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
console_formatter = logging.Formatter('%(levelname)s - %(message)s')

# Set up file handler with rotation
file_handler = RotatingFileHandler(
    'logs/vqe_simulation.log',
    maxBytes=1024 * 1024,  # 1MB per file
    backupCount=3  # Keep 3 backup files
)
file_handler.setFormatter(file_formatter)
file_handler.setLevel(logging.DEBUG)

# Set up console handler with less verbose output
console_handler = logging.StreamHandler()
console_handler.setFormatter(console_formatter)
console_handler.setLevel(logging.INFO)  # Only show INFO and above in console

# Configure root logger
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)

# Remove any existing handlers
for handler in logger.handlers[:]:
    logger.removeHandler(handler)

# Add our handlers
logger.addHandler(file_handler)
logger.addHandler(console_handler)

# Disable other loggers that might be too verbose
logging.getLogger('httpcore.http11').setLevel(logging.WARNING)
logging.getLogger('httpx').setLevel(logging.WARNING)
logging.getLogger('gradio').setLevel(logging.WARNING)

# Log startup message
logger.info("VQE Demo Application Starting")

def install_dependencies():
    print("""Install required packages from requirements.txt.""")
    try:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-r", "requirements.txt"])
        print("Dependencies installed successfully")
    except subprocess.CalledProcessError as e:
        print(f"Error installing dependencies: {e}")
        sys.exit(1)

# Test if cudaq is already installed
try:
    import cudaq
    print("CUDA-Q already installed, skipping dependency installation")
except ImportError:
    print("CUDA-Q not found, installing dependencies...")
    # Run installation before any other imports
    install_dependencies()


# Load molecule definitions
def load_molecules():
    """Load molecule definitions from JSON file."""
    try:
        with open('molecules.json', 'r') as f:
            all_molecules = json.load(f)
            # Filter to only include enabled molecules
            enabled_molecules = {k: v for k, v in all_molecules.items() if v.get('enabled', 0) == 1}
            if not enabled_molecules:
                print("Warning: No enabled molecules found in molecules.json", file=sys.stderr)
            return enabled_molecules
    except Exception as e:
        print(f"Error loading molecules.json: {e}", file=sys.stderr)
        return {}

# Now import the rest of the modules
print("about to import numpy and quantum_utils", file=sys.stderr, flush=True)
import numpy as np
from quantum_utils import run_vqe_simulation
from visualization import plot_convergence, create_molecule_viewer, format_molecule_params

# Add immediate logging
print("Imported all modules successfully", file=sys.stderr, flush=True)

def update_simulation(molecule_choice: str, scale_factor: float) -> tuple:
    """
    Run VQE simulation and update visualizations.
    
    Args:
        molecule_choice: Selected molecule from available options
        scale_factor: Factor to scale the molecule geometry by (1.0 = original size)
    Returns:
        Tuple of (energy text, convergence plot, molecule HTML)
    """
    print("UPDATE_SIMULATION CALLED", file=sys.stderr, flush=True)
    try:
        # Get molecule data
        molecules = load_molecules()
        if molecule_choice not in molecules:
            raise ValueError(f"Unknown molecule: {molecule_choice}")
        
        molecule_data = molecules[molecule_choice]
        print(f"Starting simulation with molecule: {molecule_choice}, scale: {scale_factor}", 
              file=sys.stderr, flush=True)
        
        # Run VQE simulation
        results = run_vqe_simulation(
            molecule_data=molecule_data,
            scale_factor=scale_factor
        )
        
        print(f"VQE simulation completed with results: {results}", 
              file=sys.stderr, flush=True)
        
        # Create plots
        print("Creating convergence plot...", file=sys.stderr)
        convergence_plot = plot_convergence(results)
        print("Convergence plot created", file=sys.stderr)
        
        print("Creating molecule visualization...", file=sys.stderr)
        molecule_html = create_molecule_viewer(molecule_choice, scale_factor)
        print("Molecule visualization created", file=sys.stderr)
        
        # Format energy output
        energy_text = f"""
        Final Energy: {results['final_energy']:.6f} Hartree
        Optimization Status: {'Success' if results['success'] else 'Failed'}
        Message: {results['message']}
        Number of Iterations: {len(results['history'])}
        Scale Factor: {scale_factor:.2f}
        """
        
        print("Returning results...", file=sys.stderr)
        return energy_text, convergence_plot, molecule_html
    
    except Exception as e:
        import traceback
        error_msg = f"Error in simulation:\n{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
        print(error_msg, file=sys.stderr)  # This will show in the server logs
        return error_msg, None, None

def create_interface():
    """Create the Gradio interface for the VQE demo."""
    
    def set_client_for_session(request: gr.Request):
        """Initialize client with user's IP token to handle rate limiting."""
        logger.info(f"Setting client for session with request: {request}")
        try:
            # Try to initialize client with HF token from environment variable
            client = Client("A19grey/Cuda-Quantum-Molecular-Playground", hf_token=os.getenv("HF_TOKEN"))
            logger.info(f"Client initialized with A19grey HF token")
            return client
        except Exception as e:
            logger.error(f"Error setting client with HF token: {e}")
            try:
                # Fallback to using x-ip-token from request headers
                x_ip_token = request.headers['x-ip-token']
                logger.info(f"Client initialized with x-ip-token from request headers")
                return Client("A19grey/Cuda-Quantum-Molecular-Playground", headers={"x-ip-token": x_ip_token})
            except Exception as e:
                logger.error(f"Error setting client with x-ip-token: {e}")
                return None

    # Load available molecules
    molecules = load_molecules()
    if not molecules:
        raise ValueError("No molecules found in molecules.json")
    
    with gr.Blocks(title="CUDA-Q VQE Demo") as demo:
        gr.Markdown("""
        # 🔬 NVIDIA CUDA-Q VQE Demo

        This demo harnesses the power of **NVIDIA's GPU-accelerated CUDA Quantum** library to simulate quantum systems 
        on classical hardware. The process works in two steps:

        1. 📊 Generate the [Hamiltonian](https://quantumfrontiers.com/2017/01/31/hamiltonian-an-american-musical-without-americana-or-music/) - *the mathematical description of the system's energy*
        2. ⚛️ Use the **Variational Quantum Eigensolver (VQE)** algorithm to calculate the ground state energy

        ## How to Use
        > 🔍 Select a molecule below, adjust its parameters, and watch as the simulation computes the lowest 
        > possible energy state using quantum-inspired algorithms accelerated by NVIDIA GPUs.

        ## Important Note
        ⚠️ Even for small molecules, the quantum circuits can become quite complex. In this demo, we show only the first terms. 
        
        > 💡 **Pro Tip:** In real-world applications, even with GPU acceleration, scientists use more advanced methods 
        > that focus on simulating only outer shell electrons for practical quantum chemistry calculations.

        ---
        *Inspired by the [NVIDIA CUDA-Q VQE Example](https://nvidia.github.io/cuda-quantum/latest/applications/python/vqe_advanced.html)*
        """)
        
        client = gr.State()
        state = gr.State({})
        
        # Top section: 3 columns
        with gr.Row():
            with gr.Column(scale=1):
                molecule_choice = gr.Dropdown(
                    choices=list(molecules.keys()),
                    value=None,
                    label="Select Molecule",
                    interactive=True,
                    allow_custom_value=False
                )
                
                default_molecule = molecules[list(molecules.keys())[0]]
                scale_factor = gr.Slider(
                    minimum=default_molecule['scale_range']['min'],
                    maximum=default_molecule['scale_range']['max'],
                    value=default_molecule['default_scale'],
                    step=default_molecule['scale_range']['step'],
                    label="Scale Factor (1.0 = original size)",
                    interactive=True
                )
            
            with gr.Column(scale=1):
                params_display = gr.HTML(
                    label="Molecule Parameters",
                    value="<div>Select a molecule to view its parameters</div>"
                )
            
            with gr.Column(scale=1):
                molecule_viz = gr.Plot(
                    label="3D Molecule Viewer",
                    show_label=True
                )
        
        # Middle section: Buttons and results in 2 columns
        with gr.Row():
            with gr.Column(scale=1):
                # Buttons side by side
                with gr.Row():
                    generate_ham_button = gr.Button("Generate Hamiltonian")
                    run_vqe_button = gr.Button("Find that ground state!", interactive=False)
                
                # Hamiltonian Information
                simulation_results = gr.HTML(
                    label="Hamiltonian Information",
                    value="<div style='padding: 10px; background-color: #f5f5f5; border-radius: 5px; font-size: 0.8em;'>" +
                          "Click 'Generate Hamiltonian' to start.</div>"
                )
            
            with gr.Column(scale=1):
                # VQE Convergence Plot
                convergence_plot = gr.Plot(
                    label="VQE Convergence",
                    show_label=True
                )
        
        # Bottom section: Full width circuit visualization
        circuit_latex_display = gr.Markdown(
            label="Quantum Circuit Visualization",
            value="Circuit will be displayed after generation.",
            elem_classes="circuit-output"
        )
        
        # Add CSS styling
        gr.Markdown("""
        <style>
            .circuit-output {
                font-family: monospace;
                font-size: 0.6em;
                white-space: pre;
                overflow-x: auto;
                padding: 10px;
                background-color: #f8f8f8;
                width: 100%;
                max-width: 100%;
                box-sizing: border-box;
                display: block;
            }
            .circuit-output pre {
                margin: 0 auto;
                background: none !important;
                border: none !important;
                padding: 0 !important;
                width: 100%;
                overflow-x: auto;
            }
            .circuit-output code {
                background: none !important;
                border: none !important;
                padding: 0 !important;
            }
            .gradio-container .prose {
                max-width: 100% !important;
                width: 100% !important;
            }
            .markdown-text {
                width: 100% !important;
                max-width: none !important;
            }
        </style>
        """)
        
        def generate_hamiltonian(molecule_choice: str, scale_factor: float) -> tuple:
            """Generate Hamiltonian for the selected molecule."""
            logger.info(f"Starting Hamiltonian generation for molecule: {molecule_choice}, scale: {scale_factor}")
            try:
                # Get molecule data
                logger.debug("Loading molecule data from molecules dictionary")
                if molecule_choice not in molecules:
                    logger.error(f"Molecule {molecule_choice} not found in molecules dictionary")
                    return "<div>Error: Invalid molecule selection</div>", "", gr.update(interactive=False), None
                    
                molecule_data = molecules[molecule_choice]
                logger.debug(f"Loaded molecule data: {molecule_data}")
                
                # Generate Hamiltonian only
                logger.info("Calling run_vqe_simulation with hamiltonian_only=True")
                results = run_vqe_simulation(
                    molecule_data=molecule_data,
                    scale_factor=scale_factor,
                    hamiltonian_only=True
                )
                logger.debug(f"VQE simulation results: {results}")
                
                # Format Hamiltonian results
                logger.debug("Formatting Hamiltonian results")
                hamiltonian_html = f"""
                <div style='padding: 10px; background-color: #f5f5f5; border-radius: 5px; font-size: 0.8em;'>
                    <h3>Hamiltonian Generated:</h3>
                    <ul style='list-style-type: none; padding-left: 0;'>
                        <li><b>Number of Qubits:</b> {results['qubit_count']}</li>
                        <li><b>Number of Electrons:</b> {results['electron_count']}</li>
                        <li><b>Hamiltonian Terms:</b> {results['hamiltonian_terms']}</li>
                        <li style='margin-top: 8px; color: #2a6f97;'><b>Quantum Circuit Parameters:</b> {results['parameter_count']} <i>(parameters needed to represent the system in quantum circuit form)</i></li>
                        <li><b>Scale Factor:</b> {scale_factor:.2f}</li>
                    </ul>
                    <p><i>Ready to run VQE optimization.</i></p>
                </div>
                """
                
                # Format circuit output
                logger.debug("Formatting circuit LaTeX output")
                circuit_latex = f"""<div class='circuit-output'><pre>{results.get('circuit_latex', 'Circuit visualization not available').replace('`', '&#96;')}</pre></div>"""
                
                logger.info("Successfully generated Hamiltonian and circuit visualization")
                return hamiltonian_html, circuit_latex, gr.update(interactive=True), None
                
            except Exception as e:
                import traceback
                logger.error(f"Error in generate_hamiltonian: {str(e)}\n{traceback.format_exc()}")
                error_msg = f"<div style='color: red;'>Error generating Hamiltonian:<br>{str(e)}</div>"
                return error_msg, "Error generating circuit visualization", gr.update(interactive=False), None

        def run_vqe_optimization(molecule_choice: str, scale_factor: float) -> tuple:
            """Run VQE optimization after Hamiltonian is generated."""
            try:
                # Get molecule data
                molecule_data = molecules[molecule_choice]
                
                # Run full VQE simulation
                results = run_vqe_simulation(
                    molecule_data=molecule_data,
                    scale_factor=scale_factor,
                    hamiltonian_only=False
                )
                
                # Create plots
                convergence_plot = plot_convergence(results)
                
                # Format simulation results without the optimization status
                simulation_html = f"""
                <div style='padding: 10px; background-color: #f5f5f5; border-radius: 5px;'>
                    <h3>VQE Optimization Results:</h3>
                    <ul style='list-style-type: none; padding-left: 0;'>
                        <li><b>Final Energy:</b> {results['final_energy']:.6f} Hartree</li>
                        <li><b>Parameter Count:</b> {results['parameter_count']}</li>
                        <li><b>Hamiltonian Terms:</b> {results['hamiltonian_terms']}</li>
                        <li><b>Iterations:</b> {len(results['history'])}</li>
                        <li><b>Scale Factor:</b> {scale_factor:.2f}</li>
                    </ul>
                    <p><i>{results['message']}</i></p>
                </div>
                """
                
                return simulation_html, convergence_plot
            
            except Exception as e:
                error_msg = f"<div style='color: red;'>Error in VQE optimization:<br>{str(e)}</div>"
                return error_msg, None

        def update_molecule_info(molecule_choice):
            """Update description and scale range when molecule is selected."""
            try:
                logger.info(f"Updating molecule info for: {molecule_choice}")
                mol_data = molecules[molecule_choice]
                
                # Format parameters display
                params_html = format_molecule_params(mol_data)
                logger.debug(f"Generated parameter HTML for {molecule_choice}")
                
                # Get scale range values with defaults
                min_scale = mol_data.get('scale_range', {}).get('min', 0.1)
                max_scale = mol_data.get('scale_range', {}).get('max', 3.0)
                default_scale = mol_data.get('default_scale', 1.0)
                step_size = mol_data.get('scale_range', {}).get('step', 0.05)
                logger.debug(f"Scale parameters - min: {min_scale}, max: {max_scale}, default: {default_scale}")
                
                # Create molecule visualization using Plotly
                logger.info(f"Creating Plotly molecule visualization for {molecule_choice}")
                mol_plot = create_molecule_viewer(molecule_choice, default_scale)
                if mol_plot is None:
                    logger.error(f"No molecule visualization created for {molecule_choice}")
                    mol_plot = go.Figure()  # Empty figure as fallback
                
                return [
                    params_html,
                    gr.update(
                        minimum=min_scale,
                        maximum=max_scale,
                        value=default_scale,
                        step=step_size
                    ),
                    mol_plot
                ]
            except Exception as e:
                logger.error(f"Error updating molecule info: {str(e)}", exc_info=True)
                return [
                    "<div>Error loading molecule parameters</div>",
                    gr.update(
                        minimum=0.1,
                        maximum=3.0,
                        value=1.0,
                        step=0.05
                    ),
                    go.Figure()  # Empty figure as fallback
                ]
        
        # Update parameters and 3D view when molecule changes
        molecule_choice.change(
            fn=update_molecule_info,
            inputs=[molecule_choice],
            outputs=[params_display, scale_factor, molecule_viz]
        )
        # Generate Hamiltonian when button is clicked
        generate_ham_button.click(
            fn=generate_hamiltonian,
            inputs=[molecule_choice, scale_factor],
            outputs=[simulation_results, circuit_latex_display, run_vqe_button, convergence_plot]
        )
        
        # Run VQE optimization when button is clicked
        run_vqe_button.click(
            fn=run_vqe_optimization,
            inputs=[molecule_choice, scale_factor],
            outputs=[simulation_results, convergence_plot]
        )

        # Add load event to set client
        logger.info("Adding load event to set client")
        demo.load(set_client_for_session, None, client)

    return demo

if __name__ == "__main__":
    print("Starting VQE Demo Application")
    demo = create_interface()
    print("Launching demo with custom headers support")
    demo.launch(
        server_name="0.0.0.0", 
        server_port=7860,
        ssr_mode=False,
        show_api=True  # Enable API endpoints
    )