Files
HonorThesis/codes/plot_entropy_and_dim_3d.py
2026-02-01 13:41:12 -06:00

51 lines
1.7 KiB
Python

import numpy as np
import matplotlib.pyplot as plt
from quantum_states import sample_and_calculate
from tqdm import tqdm
from mpl_toolkits.mplot3d import Axes3D
# Define range of dimensions to test
dimensionsA = np.arange(2, 64, 2) # Test dimensions from 2 to 50 in steps of 2
dimensionsB = np.arange(2, 64, 2) # Test dimensions from 2 to 50 in steps of 2
# Create meshgrid for 3D plot
X, Y = np.meshgrid(dimensionsA, dimensionsB)
Z = np.zeros_like(X, dtype=float)
# Calculate entropies for each dimension combination
total_iterations = len(dimensionsA) * len(dimensionsB)
pbar = tqdm(total=total_iterations, desc="Calculating entropies")
for i, dim_a in enumerate(dimensionsA):
for j, dim_b in enumerate(dimensionsB):
entropies = sample_and_calculate(dim_a, dim_b, n_samples=100)
Z[j,i] = np.mean(entropies)
pbar.update(1)
pbar.close()
# Create the 3D plot
fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(111, projection='3d')
# Plot the surface
surf = ax.plot_surface(X, Y, Z, cmap='viridis')
# Add labels and title with larger font sizes
ax.set_xlabel('Dimension of Subsystem A', fontsize=12, labelpad=10)
ax.set_ylabel('Dimension of Subsystem B', fontsize=12, labelpad=10)
ax.set_zlabel('von Neumann Entropy (bits)', fontsize=12, labelpad=10)
ax.set_title('von Neumann Entropy vs. System Dimensions', fontsize=14, pad=20)
# Add colorbar
cbar = fig.colorbar(surf, ax=ax, label='Entropy')
cbar.ax.set_ylabel('Entropy', fontsize=12)
# Add tick labels with larger font size
ax.tick_params(axis='x', labelsize=10)
ax.tick_params(axis='y', labelsize=10)
ax.tick_params(axis='z', labelsize=10)
# Rotate the plot for better visibility
ax.view_init(elev=30, azim=45)
plt.show()