import polars as pl
from aqora_cli.pyarrow import dataset
import matplotlib.pyplot as plt
# Load the unified dataset from Aqora
df = pl.scan_pyarrow_dataset(dataset("aqora/n2-cs-vqe", "v1.1.1")).collect()
print(f"Dataset shape: {df.shape}")
print(f"Geometries with experimental data: {df['has_experimental_data'].sum()}")
# Example 1: Plot the complete dissociation curve
plt.figure(figsize=(12, 6))
# Plot classical methods (all 160 points)
plt.plot(df['bond_length_angstrom'], df['energy_fci'],
label='FCI (Exact)', linewidth=2, color='black', zorder=3)
plt.plot(df['bond_length_angstrom'], df['energy_ccsd_t'],
label='CCSD(T)', linestyle='--', alpha=0.7)
plt.plot(df['bond_length_angstrom'], df['energy_cs_dd_5q'],
label='CS-DD (5q, noiseless)', linestyle='-.', color='blue')
# Add experimental quantum results (10 points only)
df_exp = df.filter(pl.col('has_experimental_data'))
plt.errorbar(df_exp['bond_length_angstrom'].to_numpy(),
df_exp['energy_cs_vqe_mean'].to_numpy(),
yerr=df_exp['energy_cs_vqe_std'].to_numpy(),
fmt='o', label='CS-VQE (Quantum Hardware)',
capsize=5, markersize=8, color='red', zorder=4)
plt.xlabel('Bond Length (Å)', fontsize=12)
plt.ylabel('Energy (Hartree)', fontsize=12)
plt.title('N₂ Dissociation Curve: Classical vs Quantum Computing', fontsize=14)
plt.legend(fontsize=10)
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
# Example 2: Error analysis for experimental geometries
fig, ax = plt.subplots(figsize=(10, 6))
# Calculate errors relative to FCI (convert to eV)
df_exp = df_exp.with_columns([
((pl.col('energy_cs_dd_5q') - pl.col('energy_fci')) * 27.2114).alias('error_cs_dd'),
((pl.col('energy_cs_vqe_mean') - pl.col('energy_fci')) * 27.2114).alias('error_cs_vqe')
])
ax.plot(df_exp['bond_length_angstrom'].to_numpy(),
df_exp['error_cs_dd'].abs().to_numpy(),
'o-', label='CS-DD (noiseless)', markersize=8)
ax.errorbar(df_exp['bond_length_angstrom'].to_numpy(),
df_exp['error_cs_vqe'].abs().to_numpy(),
yerr=df_exp['energy_cs_vqe_std'].to_numpy() * 27.2114,
fmt='o-', label='CS-VQE (hardware)', capsize=4, markersize=8)
ax.axhline(y=0.0016 * 27.2114, color='gray', linestyle='--',
label='Chemical accuracy (1 kcal/mol)', alpha=0.7)
ax.set_xlabel('Bond Length (Å)', fontsize=12)
ax.set_ylabel('Absolute Error vs FCI (eV)', fontsize=12)
ax.set_title('Quantum Computing Accuracy Across Dissociation Curve', fontsize=14)
ax.legend(fontsize=10)
ax.set_yscale('log')
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
# Example 3: Analyze entanglement entropy
plt.figure(figsize=(10, 6))
plt.plot(df['bond_length_angstrom'].to_numpy(),
df['entropy_fci'].to_numpy(),
label='FCI', linewidth=2)
plt.plot(df['bond_length_angstrom'].to_numpy(),
df['entropy_ccsd'].to_numpy(),
label='CCSD', linestyle='--')
plt.plot(df['bond_length_angstrom'].to_numpy(),
df['entropy_cs_dd_5q'].to_numpy(),
label='CS-DD', linestyle='-.', color='blue')
# Highlight experimental geometries
exp_entropy = df.filter(pl.col('has_experimental_data'))['entropy_fci'].to_numpy()
plt.scatter(df_exp['bond_length_angstrom'].to_numpy(),
exp_entropy,
color='red', s=100, zorder=5, label='Experimental geometries')
plt.xlabel('Bond Length (Å)', fontsize=12)
plt.ylabel('von Neumann Entropy', fontsize=12)
plt.title('Entanglement Entropy Along Dissociation Curve', fontsize=14)
plt.legend(fontsize=10)
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
