DataDive: Python Basics for Demographic Analysis

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Rectangle, FancyArrow

# Create figure
fig = plt.figure(figsize=(15, 10), facecolor='#f8f9fa')
ax = fig.add_subplot(111)
ax.set_facecolor('#f8f9fa')
ax.set_xlim(0, 14)
ax.set_ylim(0, 10)
ax.axis('off')

# Colors
color_wide = '#4e79a7'
color_long = '#e15759'
color_arrow = '#76b7b2'

# Titles
plt.text(2.5, 9.5, 'Wide Format (Original)', ha='center', fontsize=14, weight='bold', color=color_wide)
plt.text(11, 9.5, 'Long Format (Stacked)', ha='center', fontsize=14, weight='bold', color=color_long)
plt.text(7, 3.5, 'Unstacked (Back to Wide)', ha='center', fontsize=14, weight='bold', color=color_wide)

# ========== WIDE FORMAT ==========
# Table structure
ax.add_patch(Rectangle((0.5, 6), 4, 2.5, fill=None, edgecolor=color_wide, lw=2))

# Columns
columns = ['Product', 'Q1 Sales', 'Q2 Sales']
for i, col in enumerate(columns):
    plt.text(1 + i*1.5, 8.2, col, ha='center', weight='bold', color=color_wide)

# Data
data = [
    ['Laptop', 120, 150],
    ['Tablet', 85, 95],
    ['Phone', 200, 220]
]

for r, row in enumerate(data):
    for c, val in enumerate(row):
        plt.text(1 + c*1.5, 7.5 - r*0.8, str(val), ha='center', color=color_wide)

# ========== LONG FORMAT ==========
# Table structure
ax.add_patch(Rectangle((8.5, 6), 4, 2.5, fill=None, edgecolor=color_long, lw=2))

# Columns
columns_long = ['Product', 'Quarter', 'Sales']
for i, col in enumerate(columns_long):
    plt.text(9 + i*1.5, 8.2, col, ha='center', weight='bold', color=color_long)

# Data
data_long = [
    ['Laptop', 'Q1', 120],
    ['Laptop', 'Q2', 150],
    ['Tablet', 'Q1', 85],
    ['Tablet', 'Q2', 95],
    ['Phone', 'Q1', 200],
    ['Phone', 'Q2', 220]
]

for r, row in enumerate(data_long):
    for c, val in enumerate(row):
        plt.text(9 + c*1.5, 7.7 - r*0.4, str(val), ha='center', color=color_long)

# ========== UNSTACKED FORMAT ==========
ax.add_patch(Rectangle((3, 0.5), 8, 2.5, fill=None, edgecolor=color_wide, lw=2))

# Header
plt.text(4.5, 2.7, 'Product', ha='center', weight='bold', color=color_wide)
plt.text(7, 2.7, 'Quarter', ha='center', weight='bold', color=color_wide)
plt.text(9.5, 2.7, 'Sales', ha='center', weight='bold', color=color_wide)

# Divider lines
plt.plot([6, 6], [0.5, 3], color=color_wide, linestyle='--', alpha=0.7)
plt.plot([8.5, 8.5], [0.5, 3], color=color_wide, linestyle='--', alpha=0.7)

# Unstacked data
unstacked_data = [
    ['Laptop', 'Q1', 120],
    ['', 'Q2', 150],
    ['Tablet', 'Q1', 85],
    ['', 'Q2', 95],
    ['Phone', 'Q1', 200],
    ['', 'Q2', 220]
]

for r, row in enumerate(unstacked_data):
    for c, val in enumerate(row):
        x_pos = 4.5 + c*2.5 if c < 2 else 9.5
        plt.text(x_pos, 2.2 - r*0.4, str(val), ha='center', color=color_wide)

# ========== ARROWS ==========
# Stack arrow
plt.arrow(5.5, 7.5, 2.5, 0, head_width=0.3, head_length=0.3, 
          fc=color_arrow, ec=color_arrow, width=0.05)
plt.text(6.7, 7.7, 'stack()', color=color_arrow, fontsize=12, weight='bold')

# Unstack arrow
plt.arrow(8.5, 5.5, 0, -2.2, head_width=0.3, head_length=0.2, 
          fc=color_arrow, ec=color_arrow, width=0.05)
plt.text(8.2, 4, 'unstack()', color=color_arrow, fontsize=12, 
         weight='bold', rotation=90)

# Return arrow
plt.arrow(11.5, 3.5, -2.5, 0, head_width=0.3, head_length=0.3, 
          fc=color_arrow, ec=color_arrow, width=0.05, linestyle='--')
plt.text(9.5, 3.7, 'stack()', color=color_arrow, fontsize=12, weight='bold')

# Diagram explanation
plt.text(7, 9.8, 'Pandas Stack/Unstack Transformation Diagram', 
         ha='center', fontsize=16, weight='bold', color='#2c3e50')

# Legend
legend_text = [
    "stack(): Converts columns → rows (wide to long)",
    "unstack(): Converts rows → columns (long to wide)"
]
for i, text in enumerate(legend_text):
    plt.text(0.5, 0.2 - i*0.3, text, fontsize=11, color='#34495e')

# Save and show
plt.tight_layout()
plt.savefig('stack-unstack-diagram.png', dpi=300, bbox_inches='tight', facecolor='#f8f9fa')
plt.show()