normalized_heat_diffusion_sim.py

"""
Heat Diffusion Simulation
=========================
Models 2D heat spreading across a plate using finite differences.

WHY THIS FILE IS STRUCTURED THIS WAY:
We separate concerns into clear sections:
  1. Parameters  — all the numbers you might want to tune, in one place
  2. Setup       — build the initial grid state
  3. Simulation  — the core update loop
  4. Visualization — display the results

This makes it easy to change one thing without touching the rest.
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation


# =============================================================================
# SECTION 1: PARAMETERS
# All tuneable values live here. Change these to explore different behaviors.
# =============================================================================

# Grid size: how many cells wide and tall the plate is.
# WHY: More cells = more accurate, but slower. 50x50 is a good starting point.
# Note: This creates a square matrix
GRID_SIZE = 50

# Thermal diffusivity (α/alpha). Controls how fast heat spreads.
# WHY: Higher α = heat spreads faster per time step.
# Real values: Copper ~1.17e-4 m²/s, Steel ~1.2e-5 m²/s
# We use a normalized value here since our grid is unitless.
ALPHA = 0.25

# Physical size of the plate in meters (used for dx calculation).
# dx (delta x), the physical distance between two neighboring points
# WHY: We need dx to compute a stable time step (CFL condition).
# Courant-Friedriches-Lewy Condition
# "For a simulation that propagates something through a grid (heat, waves, fluid), the time step
# must be small enough that the thing being propagated cannot jump more than one cell per step."
# ^ more on that see HeatDiffusion.md
PLATE_SIZE = 1.0  # 1 meter x 1 meter plate

# Number of simulation steps to run.
# WHY: More steps = closer to steady state/0. Watch the animation to see convergence.
# Convergence = When all the cell's temperatures barely start changing
NUM_STEPS = 2000

########################################################################################################  Line of Review  ##############

# How often to capture a frame for the animation (every N steps).
# WHY: We don't want to store 2000 frames — every 20 steps is smooth enough.
FRAME_INTERVAL = 20

# Edge (boundary) temperature in degrees Celsius.
# WHY: Fixed edges simulate the plate resting on a cold surface (heat sink).
BOUNDARY_TEMP = 0.0

# Heat source temperature in degrees Celsius.
# WHY: This is the "hot spot" — like a resistor dissipating heat into the plate.
SOURCE_TEMP = 100.0

# Heat source positions as (row, col) fractions of grid size.
# WHY: Using fractions (0.0–1.0) means positions scale with GRID_SIZE.
# Try adding more sources: [(0.5, 0.5), (0.25, 0.25), (0.75, 0.75)]
HEAT_SOURCES = [
    (0.5, 0.5),   # center of the plate
]


# =============================================================================
# SECTION 2: SETUP
# Build the initial state of the simulation.
# =============================================================================

# Compute grid spacing dx (meters per cell).
# WHY: We need this for the CFL stability condition below.
dx = PLATE_SIZE / GRID_SIZE

# CFL Stability condition: dt must be ≤ dx² / (4 * alpha) for 2D.
# "The time step must be less than or equal to the cell size squared, divided by four times the thermal diffusivity"
# WHY: If dt is too large, small errors amplify each step → simulation explodes.
# We use 0.9x the maximum safe dt as a safety margin.
dt = 0.9 * (dx ** 2) / (4 * ALPHA)

print(f"Grid size:      {GRID_SIZE} x {GRID_SIZE} cells")
print(f"Cell spacing:   {dx:.4f} m")
print(f"Time step (dt): {dt:.6f} s")
print(f"Total sim time: {NUM_STEPS * dt:.4f} s")
print(f"Stability check: α·dt/dx² = {ALPHA * dt / dx**2:.4f} (must be ≤ 0.25)")
print()

# Initialize the temperature grid to boundary temperature everywhere.
# WHY: np.full creates a 2D array where every cell has the same value.
# Shape is (rows, cols) — rows = Y axis, cols = X axis.
T = np.full((GRID_SIZE, GRID_SIZE), BOUNDARY_TEMP, dtype=float)

# Apply heat sources.
# WHY: We pin these cells to SOURCE_TEMP permanently during simulation.
# Converting fractional positions to integer grid indices with int().
source_indices = []
for (row_frac, col_frac) in HEAT_SOURCES:
    r = int(row_frac * (GRID_SIZE - 1))
    c = int(col_frac * (GRID_SIZE - 1))
    source_indices.append((r, c))
    T[r, c] = SOURCE_TEMP
    print(f"Heat source placed at grid cell ({r}, {c})")

print()


# =============================================================================
# SECTION 3: SIMULATION
# Core update loop — this is where the physics happens.
# =============================================================================

# Store frames for animation.
# WHY: We collect snapshots during the sim rather than animating live,
# so the animation plays back at a controlled speed regardless of compute time.
frames = [T.copy()]  # .copy() is critical — without it, all frames point to same array

def step(T):
    """
    Perform one time step of the heat diffusion equation.

    The finite difference Laplacian:
        ∇²T[i,j] = T[i+1,j] + T[i-1,j] + T[i,j+1] + T[i,j-1] - 4*T[i,j]

    WHY SLICING INSTEAD OF LOOPS:
        A Python for-loop over every cell would be O(N²) Python operations.
        NumPy slicing shifts the entire array in one C-level operation — 
        dramatically faster. This is the key NumPy skill to internalize.

        T[1:-1, 1:-1]  → all interior cells (excludes boundary rows/cols)
        T[2:,   1:-1]  → same grid shifted UP by one row  (south neighbor)
        T[:-2,  1:-1]  → same grid shifted DOWN by one row (north neighbor)
        T[1:-1, 2:]    → same grid shifted LEFT by one col  (east neighbor)
        T[1:-1, :-2]   → same grid shifted RIGHT by one col (west neighbor)
    """
    # Compute Laplacian for all interior cells simultaneously (no Python loop!)
    laplacian = (
        T[2:,   1:-1] +   # neighbor below
        T[:-2,  1:-1] +   # neighbor above
        T[1:-1, 2:]   +   # neighbor right
        T[1:-1, :-2]  -   # neighbor left
        4 * T[1:-1, 1:-1] # minus 4x self
    )

    # Apply the update rule: T_new = T + α * dt * ∇²T
    # WHY: Only update interior cells — boundaries are fixed (Dirichlet condition).
    T_new = T.copy()
    T_new[1:-1, 1:-1] = T[1:-1, 1:-1] + ALPHA * dt * laplacian

    # Re-apply boundary condition (edges stay fixed).
    # WHY: The update above might drift boundaries slightly due to float ops.
    T_new[0, :]  = BOUNDARY_TEMP  # top edge
    T_new[-1, :] = BOUNDARY_TEMP  # bottom edge
    T_new[:, 0]  = BOUNDARY_TEMP  # left edge
    T_new[:, -1] = BOUNDARY_TEMP  # right edge

    # Re-apply heat sources (pin them to SOURCE_TEMP every step).
    # WHY: Without this, the source cell would cool down due to diffusion.
    for (r, c) in source_indices:
        T_new[r, c] = SOURCE_TEMP

    return T_new


print("Running simulation...")
for step_num in range(NUM_STEPS):
    T = step(T)

    # Save frame every FRAME_INTERVAL steps
    if (step_num + 1) % FRAME_INTERVAL == 0:
        frames.append(T.copy())

print(f"Simulation complete. {len(frames)} frames captured.")
print(f"Final max temp: {T.max():.2f}°C  |  Final min temp: {T.min():.2f}°C")
print()


# =============================================================================
# SECTION 4: VISUALIZATION
# =============================================================================

print("Rendering animation...")

fig, axes = plt.subplots(1, 2, figsize=(12, 5))
fig.suptitle("2D Heat Diffusion Simulation", fontsize=14, fontweight='bold')

# --- Left plot: animated heatmap ---
ax_heat = axes[0]
ax_heat.set_title("Temperature Field (evolving)")
ax_heat.set_xlabel("X (grid cells)")
ax_heat.set_ylabel("Y (grid cells)")

# imshow displays a 2D array as a color image.
# WHY 'hot' colormap: intuitively maps black→red→yellow→white (cold→hot).
# vmin/vmax fix the color scale so it doesn't jump around during animation.
im = ax_heat.imshow(
    frames[0],
    cmap='hot',
    vmin=BOUNDARY_TEMP,
    vmax=SOURCE_TEMP,
    origin='lower',        # WHY: 'lower' puts row 0 at the bottom (math convention)
    interpolation='bilinear'  # WHY: smooths the blocky grid for nicer visuals
)
plt.colorbar(im, ax=ax_heat, label='Temperature (°C)')

# Mark heat source locations
for (r, c) in source_indices:
    ax_heat.plot(c, r, 'b+', markersize=12, markeredgewidth=2, label='Heat source')
ax_heat.legend(loc='upper right', fontsize=8)

step_text = ax_heat.text(
    0.02, 0.95, '', transform=ax_heat.transAxes,
    color='white', fontsize=9, verticalalignment='top'
)

# --- Right plot: temperature cross-section through center ---
ax_line = axes[1]
ax_line.set_title("Temperature Profile: Horizontal Slice Through Center")
ax_line.set_xlabel("X (grid cells)")
ax_line.set_ylabel("Temperature (°C)")
ax_line.set_ylim(BOUNDARY_TEMP - 5, SOURCE_TEMP + 5)
ax_line.axhline(y=BOUNDARY_TEMP, color='gray', linestyle='--', linewidth=0.8, label='Boundary temp')

# WHY this cross-section: it lets you see the temperature gradient clearly —
# hot in the middle, dropping off toward the edges. You'll see this in real life too.
center_row = GRID_SIZE // 2
line_plot, = ax_line.plot([], [], 'r-', linewidth=2, label='T at center row')
ax_line.set_xlim(0, GRID_SIZE - 1)
ax_line.legend()
ax_line.grid(True, alpha=0.3)


def animate(frame_idx):
    """Update function called for each animation frame."""
    frame_data = frames[frame_idx]
    im.set_data(frame_data)
    step_text.set_text(f'Step: {frame_idx * FRAME_INTERVAL}')
    line_plot.set_data(np.arange(GRID_SIZE), frame_data[center_row, :])
    return im, step_text, line_plot


ani = animation.FuncAnimation(
    fig,
    animate,
    frames=len(frames),
    interval=50,    # milliseconds between frames
    blit=True       # WHY: only redraw changed artists — much faster
)

plt.tight_layout()
plt.show()

print()
print("=== WHAT TO OBSERVE ===")
print("1. Heat spreads radially from the source — symmetric because the plate is symmetric.")
print("2. The gradient steepens near the source — more heat flowing per unit distance.")
print("3. The cross-section plot converges to a fixed curve — that's steady state.")
print("4. Try adding more heat sources in HEAT_SOURCES and see how they interact!")