3D Hertzsprung-Russell Diagram

import matplotlib.pyplot as plt
import numpy as np

np.random.seed(2023)

# Simulated HR diagram with 3D (adding distance)
n_stars = 400

# Main sequence
n_ms = 250
temp_ms = np.random.uniform(3000, 30000, n_ms)
lum_ms = (temp_ms / 5800) ** 4 * np.random.lognormal(0, 0.3, n_ms)
dist_ms = np.random.exponential(100, n_ms) + 10

# Giants
n_giants = 100
temp_giants = np.random.uniform(3500, 5500, n_giants)
lum_giants = np.random.uniform(50, 500, n_giants)
dist_giants = np.random.exponential(200, n_giants) + 50

# White dwarfs
n_wd = 50
temp_wd = np.random.uniform(8000, 40000, n_wd)
lum_wd = np.random.uniform(0.001, 0.1, n_wd)
dist_wd = np.random.exponential(50, n_wd) + 5

temp = np.concatenate([temp_ms, temp_giants, temp_wd])
lum = np.concatenate([lum_ms, lum_giants, lum_wd])
dist = np.concatenate([dist_ms, dist_giants, dist_wd])

fig = plt.figure(figsize=(10, 8), facecolor='#0a0a0f')
ax = fig.add_subplot(111, projection='3d', facecolor='#0a0a0f')

# Color by temperature (blue=hot, red=cool)
scatter = ax.scatter(temp, np.log10(lum), dist, c=temp, cmap='RdYlBu', 
                     s=20, alpha=0.7, edgecolors='none')

ax.set_xlabel('Temperature (K)', color='white', fontsize=10)
ax.set_ylabel('log₁₀(L/L☉)', color='white', fontsize=10)
ax.set_zlabel('Distance (pc)', color='white', fontsize=10)
ax.set_title('3D Hertzsprung-Russell Diagram', color='white', fontsize=14, fontweight='bold', pad=20)
ax.invert_xaxis()  # Hot stars on left

ax.tick_params(colors='#64748b', labelsize=8)
ax.xaxis.pane.fill = False
ax.yaxis.pane.fill = False
ax.zaxis.pane.fill = False
ax.xaxis.pane.set_edgecolor('#1e293b')
ax.yaxis.pane.set_edgecolor('#1e293b')
ax.zaxis.pane.set_edgecolor('#1e293b')

ax.view_init(elev=20, azim=135)
plt.tight_layout()
plt.show()