Atomic p-Orbital Electron Density

import matplotlib.pyplot as plt
import numpy as np

np.random.seed(456)

# Electron density distribution (orbital shape)
n_points = 400

# p-orbital shape
theta = np.random.uniform(0, 2*np.pi, n_points)
phi = np.random.uniform(0, np.pi, n_points)
r = np.abs(np.cos(phi)) * np.random.exponential(1, n_points)

x = r * np.sin(phi) * np.cos(theta)
y = r * np.sin(phi) * np.sin(theta)
z = r * np.cos(phi)

# Phase (positive/negative lobes)
phase = np.sign(z)
colors = ['#4927F5' if p > 0 else '#F5276C' for p in phase]

# Density for size
density = np.exp(-r) * 100 + 10

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

ax.scatter(x, y, z, c=colors, s=density, alpha=0.6, edgecolors='white', linewidths=0.2)

# Mark nucleus
ax.scatter([0], [0], [0], c='#1f2937', s=100, marker='o')

ax.set_xlabel('X (Å)', color='#1f2937', fontsize=10)
ax.set_ylabel('Y (Å)', color='#1f2937', fontsize=10)
ax.set_zlabel('Z (Å)', color='#1f2937', fontsize=10)
ax.set_title('Atomic p-Orbital Electron Density', color='#1f2937', fontsize=14, fontweight='bold', pad=20)

ax.tick_params(colors='#6b7280', labelsize=8)
ax.xaxis.pane.fill = False
ax.yaxis.pane.fill = False
ax.zaxis.pane.fill = False
ax.xaxis.pane.set_edgecolor('#e5e7eb')
ax.yaxis.pane.set_edgecolor('#e5e7eb')
ax.zaxis.pane.set_edgecolor('#e5e7eb')

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