""" quantum_utils.py Helper library for running a simplified VQE or quantum simulation that references CUDA-Q for computations. References: - @Cuda-Q_install.md for installation details - @VQE_Example.md for example VQE code using nvidia or nvidia-mqpu backend """ import cudaq import cudaq_solvers as solvers # Import solvers from cudaq_solvers package import numpy as np from scipy.optimize import minimize import spaces from typing import Dict, List, Union, Any, Tuple import sys import logging import os from logging.handlers import RotatingFileHandler import openfermion import openfermionpyscf from openfermion.transforms import jordan_wigner, get_fermion_operator # Create logs directory if it doesn't exist 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=10*1024*1024, # 10MB backupCount=5 ) 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 logger logger = logging.getLogger(__name__) 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) # Log startup message logger.info("VQE Simulation module initialized") print("quantum_utils imported", file=sys.stderr, flush=True) def setup_target(): """Set up CUDA-Q target based on available hardware.""" try: print("Setting up target...", file=sys.stderr, flush=True) gpu_count = cudaq.num_available_gpus() print(f"Number of available GPUs: {gpu_count}", file=sys.stderr, flush=True) if gpu_count > 0: print("Attempting to set NVIDIA GPU target...", file=sys.stderr, flush=True) cudaq.set_target("nvidia") print("Successfully set NVIDIA GPU target", file=sys.stderr, flush=True) else: print("No GPU found, attempting to set CPU target...", file=sys.stderr, flush=True) cudaq.set_target("qpp-cpu") print("Successfully set CPU target", file=sys.stderr, flush=True) except Exception as e: print(f"Error setting up quantum target: {str(e)}", file=sys.stderr, flush=True) import traceback print(f"Traceback:\n{traceback.format_exc()}", file=sys.stderr, flush=True) raise def vqe_callback(xk): """Callback function for VQE optimization to track progress.""" try: logging.info(f"VQE iteration - Current parameters: {xk}") return True except Exception as e: logging.error(f"Error in VQE callback: {str(e)}") return False def validate_ansatz(kernel_generator, init_params, qubit_num, electron_num): """Validate the ansatz kernel before VQE.""" try: logging.debug("Validating ansatz kernel...") # Set up target before validation setup_target() # Generate the kernel with the correct number of qubits and electrons kernel = kernel_generator(qubit_num, electron_num) # Log the types and values of parameters and qubits logging.debug(f"Kernel parameters type: {type(init_params)}, values: {init_params}") logging.debug(f"Kernel qubits type: {type(qubit_num)}, value: {qubit_num}") logging.debug(f"Kernel electrons type: {type(electron_num)}, value: {electron_num}") # Try running the kernel with initial parameters result = cudaq.sample(kernel, init_params.tolist()) # Convert to list for validation logging.debug(f"Ansatz validation successful. Sample result: {result}") return True except Exception as e: logging.error(f"Ansatz validation failed: {str(e)}") return False # ---------------------------------------------------------- # New or updated utility function for generating a Hamiltonian # ---------------------------------------------------------- def create_molecular_hamiltonian(geometry: list[tuple[str, tuple[float, float, float]]], basis: str, multiplicity: int, charge: int) -> cudaq.SpinOperator: """ Create a SpinOperator for a given molecule using PySCF + OpenFermion + CUDA-Q. Parameters ---------- geometry : list of (str, (float, float, float))) List describing each atom, e.g. [('H', (0.,0.,0.)), ('H',(0.,0.,bond_length))] basis : str Basis set, e.g. "sto-3g" multiplicity : int Spin multiplicity charge : int Net charge of the molecule Returns ------- cudaq.SpinOperator The qubit Hamiltonian in CUDA-Q spin-operator form. """ try: # Set up the CUDA-Q target for GPU acceleration setup_target() # Use PySCF + OpenFermion to build a molecular object molecule = openfermionpyscf.run_pyscf( openfermion.MolecularData(geometry, basis, multiplicity, charge) ) # Construct the fermionic Hamiltonian molecular_hamiltonian = molecule.get_molecular_hamiltonian() fermion_hamiltonian = get_fermion_operator(molecular_hamiltonian) qubit_hamiltonian = jordan_wigner(fermion_hamiltonian) # Convert to CUDA-Q SpinOperator spin_op = cudaq.SpinOperator(qubit_hamiltonian) return spin_op except Exception as e: logging.error(f"Failed to create molecular Hamiltonian: {str(e)}") raise # ---------------------------------------------------------- # New or updated function for generating a UCCSD ansatz kernel # ---------------------------------------------------------- @cudaq.kernel #required to create the kernel object for Nvidia def kernel(qubit_num: int, electron_num: int, thetas: list[float]): """ Generate a UCCSD ansatz kernel given the number of qubits and electrons. Parameters ---------- qubit_num : int Number of qubits in the system electron_num : int Number of electrons (which is the same as number of X gates for HF ref.) Returns ------- A CUDA-Q kernel that accepts parameters for the UCCSD ansatz. """ qubits = cudaq.qvector(qubit_num) # Prepare the Hartree-Fock reference: X on each occupied orbital for i in range(electron_num): x(qubits[i]) # Apply X gate to each occupied orbital it is defined in the cuda kernel decorator space # Add the UCCSD terms cudaq.kernels.uccsd(qubits, thetas, electron_num, qubit_num) # define the cost function as the observation of the spin_hamiltonian w.r.t. the ansatz def cost_function(kernel, spin_hamiltonian, qubit_count, electron_count, thetas): # The "cost" is the expectation of the spin_operator w.r.t. the ansatz, for these params exp_val = cudaq.observe(kernel, spin_hamiltonian, qubit_count, electron_count, thetas).expectation() logging.info(f"Cost function evaluation: {exp_val:.6f}") return exp_val def expand_geometry(geometry_template: List[List[Any]], scale_factor: float) -> List[Tuple[str, Tuple[float, float, float]]]: """ Expand or contract a molecule's geometry by a scale factor. Args: geometry_template: List of [atom_symbol, [x, y, z]] coordinates scale_factor: Factor to scale the geometry by (1.0 = no change) Returns: List of (atom_symbol, (x, y, z)) tuples with scaled coordinates """ scaled_geometry = [] logging.debug("expand_geometry: Starting with scale_factor=%f", scale_factor) logging.debug("expand_geometry: Input geometry_template=%s", geometry_template) # Find the center of mass (assuming equal weights for simplicity) coords = np.array([coord for _, coord in geometry_template]) logging.debug("expand_geometry: Computed coordinates array: %s", coords) center = np.mean(coords, axis=0) logging.debug("expand_geometry: Computed center of mass: %s", center) for atom_symbol, coord in geometry_template: # Convert coord to numpy array for vector operations coord = np.array(coord) logging.debug("expand_geometry: Processing atom: %s", atom_symbol) logging.debug("expand_geometry: Original coordinate: %s", coord) # Calculate vector from center vec = coord - center logging.debug("expand_geometry: Computed vector from center: %s", vec) # Scale the vector and add back to center scaled_coord = center + vec * scale_factor logging.debug("expand_geometry: Scaled coordinate: %s", scaled_coord) # Convert back to tuple scaled_geometry.append((atom_symbol, tuple(scaled_coord))) logging.debug("expand_geometry: Final scaled geometry: %s", scaled_geometry) return scaled_geometry def generate_hamiltonian(molecule_data: Dict[str, Any], scale_factor: float) -> Dict[str, Any]: """ Generate the Hamiltonian and its parameters for a given molecule without running VQE optimization. Parameters ---------- molecule_data : Dict[str, Any] Dictionary containing all molecule metadata and parameters scale_factor : float Factor to scale the molecule geometry by (1.0 = original size) Returns ------- Dict[str, Any] Dictionary containing Hamiltonian information: - 'hamiltonian': The CUDA-Q SpinOperator Hamiltonian - 'qubit_count': Number of qubits needed - 'electron_count': Number of electrons in the system - 'parameter_count': Number of UCCSD parameters needed - 'hamiltonian_terms': Number of terms in the Hamiltonian - 'circuit_latex': LaTeX representation of the quantum circuit """ # Get GPU time from molecule data or use default gpu_time = molecule_data.get('GPU_time', 60) logger.info(f"Generating Hamiltonian with GPU time: {gpu_time}") def _generate_hamiltonian_inner(): logging.info(f"Generating Hamiltonian for {molecule_data['name']} with scale factor {scale_factor}") # Set up the CUDA-Q target setup_target() # Create scaled geometry geometry = expand_geometry(molecule_data['geometry_template'], scale_factor) logging.info(f"Created scaled geometry with factor = {scale_factor}") # Generate the Hamiltonian spin_hamiltonian = create_molecular_hamiltonian( geometry=geometry, basis=molecule_data['basis'], multiplicity=molecule_data['multiplicity'], charge=molecule_data['charge'] ) # Get Hamiltonian info term_count = spin_hamiltonian.get_term_count() logging.debug(f"Hamiltonian has {term_count} terms.") # Commented out because it's too much info to show in the logs # logging.debug(f"Hamiltonian details:\n{spin_hamiltonian}") # Get system parameters qubit_count = 2 * molecule_data['spatial_orbitals'] electron_count = molecule_data['electron_count'] # Compute UCCSD parameters needed parameter_count = cudaq.kernels.uccsd_num_parameters(electron_count, qubit_count) logging.info(f"Number of UCCSD parameters needed = {parameter_count}") # Generate LaTeX representation of the circuit try: logging.info("Starting circuit LaTeX generation") thetas_draw = np.random.normal(0, 1, parameter_count) # Draw the circuit in fancy nice stuff circuit_latex = cudaq.draw(kernel, qubit_count, electron_count, thetas_draw) # Limit circuit output size MAX_CIRCUIT_LENGTH = 15000 if len(circuit_latex) > MAX_CIRCUIT_LENGTH: circuit_latex = circuit_latex[:MAX_CIRCUIT_LENGTH] + "\n... (circuit visualization clipped for size)" logging.info("Successfully generated circuit LaTeX representation") logging.debug(f"Circuit LaTeX length: {len(circuit_latex)} characters") except Exception as e: logging.error(f"Failed to generate circuit LaTeX: {str(e)}") circuit_latex = "Error generating circuit visualization" return { 'hamiltonian': spin_hamiltonian, 'qubit_count': qubit_count, 'electron_count': electron_count, 'parameter_count': parameter_count, 'hamiltonian_terms': term_count, 'circuit_latex': circuit_latex } return _generate_hamiltonian_inner() def run_vqe_simulation(molecule_data: Dict[str, Any], scale_factor: float, hamiltonian_only: bool = False) -> Dict[str, Any]: """ Run a VQE simulation using CUDA-Q. Parameters ---------- molecule_data : Dict[str, Any] Dictionary containing all molecule metadata and parameters scale_factor : float Factor to scale the molecule geometry by (1.0 = original size) hamiltonian_only : bool If True, only generate the Hamiltonian without running VQE optimization Returns ------- Dict[str, Any] The dictionary containing either just Hamiltonian info or full VQE results """ # Get GPU time from molecule data or use default gpu_time = molecule_data.get('GPU_time', 60) logger.info(f"Running VQE simulation with GPU time: {gpu_time}") @spaces.GPU(duration=gpu_time) def _run_vqe_simulation_inner(): setup_target() # Generate Hamiltonian and get system parameters logger.info("Generating Hamiltonian") ham_info = generate_hamiltonian(molecule_data, scale_factor) # If only Hamiltonian generation is requested, return early if hamiltonian_only: return { 'parameter_count': ham_info['parameter_count'], 'hamiltonian_terms': ham_info['hamiltonian_terms'], 'qubit_count': ham_info['qubit_count'], 'electron_count': ham_info['electron_count'], 'message': "Hamiltonian generated successfully", 'circuit_latex': ham_info['circuit_latex'] } # Get max iterations from molecule data or use default optimizer_max_iterations = molecule_data.get('iterations', 25) # Initialize the parameters for the UCCSD ansatz thetas0 = np.random.normal(0, 1, ham_info['parameter_count']) logging.info(f"Initial parameters: {thetas0}") # Create a wrapper function that only takes the parameters to optimize def objective_wrapper(thetas): return cost_function(kernel, ham_info['hamiltonian'], ham_info['qubit_count'], ham_info['electron_count'], thetas) # Define the callback function that uses the wrapper exp_vals = [] def callback(xk): val = objective_wrapper(xk) exp_vals.append(val) return True optimization_success = False optimization_message = "" final_energy = float("inf") final_parameters = [] try: # Test the wrapper function before optimization test_val = objective_wrapper(thetas0) logging.info(f"Debug cost check (initial params): {test_val:.6f}") if not np.isfinite(test_val): logging.warning("Debug cost check returned non-finite value. The optimizer may fail.") # Perform the VQE optimization using the wrapper function result = minimize(objective_wrapper, thetas0, method='COBYLA', callback=callback, options={'maxiter': optimizer_max_iterations}) # Store optimization results final_energy = result.fun final_parameters = result.x optimization_success = result.success optimization_message = ( "VQE optimization completed successfully." if result.success else f"VQE optimization completed with status: {result.message}" ) except Exception as e: logging.error(f"VQE optimization error: {str(e)}") optimization_message = f"VQE optimization error: {str(e)}" optimization_success = False logging.info(optimization_message) logging.debug(f"Final energy: {final_energy}") logging.debug(f"Optimized parameters: {final_parameters}") # Build the results dictionary results = { "final_energy": float(final_energy), "parameters": list(final_parameters), "success": optimization_success, "iterations": len(exp_vals), "history": exp_vals, "message": optimization_message, "parameter_count": ham_info['parameter_count'], "hamiltonian_terms": ham_info['hamiltonian_terms'] } return results return _run_vqe_simulation_inner()