Chapter 8: Matplotlib - Visualizing Your Universe

import numpy as np

# Question 1: What shape will this array have?
x = np.linspace(0, 10, 100)
y = np.sin(x)
# Shape of y: _______

# Question 2: What will this boolean mask select?
data = np.array([1, 5, 3, 8, 2, 9])
mask = data > 5
selected = data[mask]
# selected contains: _______

# Question 3: What's wrong with this code?
# matrix = np.array([[1, 2], [3, 4]])
# result = matrix + [10, 20, 30]  # Will this work?
# Your answer: _______

# Question 4: What will this create?
X, Y = np.meshgrid(np.arange(5), np.arange(3))
# Shape of X and Y: _______

import numpy as np
import matplotlib.pyplot as plt

# Generate some data - a damped oscillation
time = np.linspace(0, 10, 1000)
signal = np.exp(-time/3) * np.sin(2 * np.pi * time)

# Approach 1: The Default Plot (technically correct but uninspiring)
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

ax1.plot(time, signal)
ax1.set_title('Default: Technically Correct')

# Approach 2: The Artistic Plot (same data, different choices)
ax2.plot(time, signal, color='#2E86AB', linewidth=2, alpha=0.8)
ax2.fill_between(time, signal, alpha=0.2, color='#2E86AB')
ax2.axhline(y=0, color='#A23B72', linestyle='-', linewidth=0.5, alpha=0.5)
ax2.set_xlabel('Time (s)', fontsize=11, fontweight='light')
ax2.set_ylabel('Amplitude', fontsize=11, fontweight='light')
ax2.set_title('Artistic: Same Data, Better Story', fontsize=13, fontweight='bold')
ax2.grid(True, alpha=0.15, linestyle='-', linewidth=0.5)
ax2.spines['top'].set_visible(False)
ax2.spines['right'].set_visible(False)

plt.tight_layout()
plt.show()

# Real data: stellar magnitudes with a power-law distribution
np.random.seed(42)
masses = np.random.pareto(2.35, 1000) + 0.1  # Stellar IMF (Salpeter 1955)
luminosities = masses ** 3.5  # Mass-luminosity relation
luminosities += np.random.normal(0, 0.1 * luminosities)  # Add scatter

# Try different visualizations to find what reveals the pattern best
fig, axes = plt.subplots(2, 3, figsize=(12, 8))

# Attempt 1: Simple scatter
axes[0, 0].scatter(masses, luminosities, s=1, alpha=0.5)
axes[0, 0].set_title('Linear Scale: Pattern Hidden')
axes[0, 0].set_xlabel('Mass [M_☉]')
axes[0, 0].set_ylabel('Luminosity [L_☉]')

# Attempt 2: Log-log reveals power law!
axes[0, 1].loglog(masses, luminosities, '.', markersize=2, alpha=0.5)
axes[0, 1].set_title('Log-Log: Power Law Revealed!')
axes[0, 1].set_xlabel('Mass [M_☉]')
axes[0, 1].set_ylabel('Luminosity [L_☉]')

# Attempt 3: Semilog-x (wrong choice for this data)
axes[0, 2].semilogx(masses, luminosities, '.', markersize=2, alpha=0.5)
axes[0, 2].set_title('Semilog-X: Not Helpful Here')
axes[0, 2].set_xlabel('Mass [M_☉]')
axes[0, 2].set_ylabel('Luminosity [L_☉]')

# Attempt 4: 2D histogram for density
h = axes[1, 0].hist2d(np.log10(masses), np.log10(luminosities), 
                       bins=30, cmap='YlOrRd')
axes[1, 0].set_title('2D Histogram: Shows Density')
axes[1, 0].set_xlabel('log(Mass)')
axes[1, 0].set_ylabel('log(Luminosity)')

# Attempt 5: Hexbin for large datasets
axes[1, 1].hexbin(masses, luminosities, gridsize=20, 
                  xscale='log', yscale='log', cmap='Blues')
axes[1, 1].set_title('Hexbin: Good for Large N')
axes[1, 1].set_xlabel('Mass [M_☉]')
axes[1, 1].set_ylabel('Luminosity [L_☉]')

# Attempt 6: Contours with scatter overlay
from scipy.stats import gaussian_kde
xy = np.vstack([np.log10(masses), np.log10(luminosities)])
z = gaussian_kde(xy)(xy)
idx = z.argsort()
x, y, z = np.log10(masses)[idx], np.log10(luminosities)[idx], z[idx]
axes[1, 2].scatter(x, y, c=z, s=1, cmap='viridis')
axes[1, 2].set_title('KDE-Colored: Shows Structure')
axes[1, 2].set_xlabel('log(Mass)')
axes[1, 2].set_ylabel('log(Luminosity)')

plt.suptitle('Experimentation Reveals the Best Visualization', 
             fontsize=14, fontweight='bold')
plt.tight_layout()
plt.show()

print("Key Lesson: The log-log plot immediately reveals the power-law relationship")
print("that was completely hidden in the linear plot. Always experiment!")

# The Complete Anatomy of a Matplotlib Figure
# Copyright (c) 2016 Nicolas P. Rougier - MIT License
# Adapted from: http://github.com/rougier/figure-anatomy
# This brilliant visualization shows every component of a Matplotlib figure

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator, FuncFormatter

np.random.seed(123)

# Create the figure with Rougier's specifications
X = np.linspace(0.5, 3.5, 100)
Y1 = 3 + np.cos(X)
Y2 = 1 + np.cos(1 + X/0.75)/2
Y3 = np.random.uniform(Y1, Y2, len(X))

fig = plt.figure(figsize=(8, 8), facecolor="w")
ax = fig.add_subplot(1, 1, 1, aspect=1)

# Configure the tick system
def minor_tick(x, pos):
    if not x % 1.0:
        return ""
    return "%.2f" % x

ax.xaxis.set_major_locator(MultipleLocator(1.000))
ax.xaxis.set_minor_locator(MultipleLocator(0.250))
ax.yaxis.set_major_locator(MultipleLocator(1.000))
ax.yaxis.set_minor_locator(MultipleLocator(0.250))
ax.xaxis.set_minor_formatter(FuncFormatter(minor_tick))

ax.set_xlim(0, 4)
ax.set_ylim(0, 4)

ax.tick_params(which='major', width=1.0, length=10)
ax.tick_params(which='minor', width=1.0, length=5, labelsize=10, labelcolor='0.25')

# Add the plot elements
ax.grid(linestyle="--", linewidth=0.5, color='.25', zorder=-10)
ax.plot(X, Y1, c=(0.25, 0.25, 1.00), lw=2, label="Blue signal", zorder=10)
ax.plot(X, Y2, c=(1.00, 0.25, 0.25), lw=2, label="Red signal")
ax.scatter(X, Y3, c='w', edgecolors='black', linewidth=0.5)

ax.set_title("Anatomy of a figure", fontsize=20)
ax.set_xlabel("X axis label")
ax.set_ylabel("Y axis label")
ax.legend(frameon=False)

# Add annotations for each component
def circle(x, y, radius=0.15):
    from matplotlib.patches import Circle
    from matplotlib.patheffects import withStroke
    circle = Circle((x, y), radius, clip_on=False, zorder=10, linewidth=1,
                    edgecolor='black', facecolor=(0, 0, 0, .0125),
                    path_effects=[withStroke(linewidth=5, foreground='w')])
    ax.add_artist(circle)

def text(x, y, text):
    ax.text(x, y, text, backgroundcolor="white",
            ha='center', va='top', weight='bold', color='blue')

# Label all the components
circle(0.50, -0.05)
text(0.50, -0.25, "Minor tick label")

circle(4.00, 2.00)
text(4.00, 1.80, "Major tick")

circle(1.80, -0.22)
text(1.80, -0.4, "X axis label")

circle(1.75, 2.80)
text(1.75, 2.60, "Line\n(line plot)")

circle(3.20, 1.75)
text(3.20, 1.55, "Markers\n(scatter plot)")

circle(3.00, 3.00)
text(3.00, 2.80, "Grid")

circle(3.70, 3.75)
text(3.70, 3.55, "Legend")

circle(0.5, 0.5)
text(0.5, 0.3, "Axes")

circle(-0.3, 0.65)
text(-0.3, 0.45, "Figure")

# Add the spines annotation
ax.annotate('Spines', xy=(4.0, 0.35), xycoords='data',
            xytext=(3.3, 0.5), textcoords='data',
            weight='bold', color='blue',
            arrowprops=dict(arrowstyle='->', connectionstyle="arc3", color='blue'))

plt.suptitle("Credit: Nicolas P. Rougier (http://github.com/rougier/figure-anatomy)",
             fontsize=10, family="monospace", color='.5')
plt.tight_layout()
plt.show()

# Create a figure with explicit control over components
fig = plt.figure(figsize=(10, 6))

# Add axes manually to see the structure
# [left, bottom, width, height] in figure coordinates (0-1)
ax_main = fig.add_axes([0.1, 0.3, 0.7, 0.6])   # Main plot
ax_zoom = fig.add_axes([0.65, 0.6, 0.2, 0.2])  # Inset zoom
ax_residual = fig.add_axes([0.1, 0.05, 0.7, 0.2])  # Residual panel

# Generate data
x = np.linspace(0, 10, 100)
y = np.sin(x) * np.exp(-x/10)
y_model = np.sin(x) * np.exp(-x/10.5)  # Slightly different model

# Main plot
ax_main.plot(x, y, 'ko', markersize=3, alpha=0.5, label='Data')
ax_main.plot(x, y_model, 'r-', linewidth=2, label='Model')
ax_main.set_ylabel('Signal', fontsize=12)
ax_main.legend(loc='upper right')
ax_main.grid(True, alpha=0.3)

# Zoom inset
zoom_mask = (x > 3) & (x < 5)
ax_zoom.plot(x[zoom_mask], y[zoom_mask], 'ko', markersize=2)
ax_zoom.plot(x[zoom_mask], y_model[zoom_mask], 'r-', linewidth=1)
ax_zoom.set_xlim(3, 5)
ax_zoom.grid(True, alpha=0.3)
ax_zoom.set_title('Zoom', fontsize=9)

# Residuals
residuals = y - y_model
ax_residual.scatter(x, residuals, s=5, alpha=0.5, color='blue')
ax_residual.axhline(y=0, color='red', linestyle='--', linewidth=1)
ax_residual.set_xlabel('Time', fontsize=12)
ax_residual.set_ylabel('Residuals', fontsize=10)
ax_residual.grid(True, alpha=0.3)

fig.suptitle('Explicit Control Over Figure Components', fontsize=14, fontweight='bold')
plt.show()

from matplotlib.gridspec import GridSpec

# Create figure with shared x-axes using GridSpec
fig = plt.figure(figsize=(10, 8))
gs = GridSpec(3, 1, hspace=0.05)

# Create subplots with shared x-axis
ax1 = fig.add_subplot(gs[0])
ax2 = fig.add_subplot(gs[1], sharex=ax1)
ax3 = fig.add_subplot(gs[2], sharex=ax1)

# Generate multi-wavelength data
time = np.linspace(0, 10, 200)
optical = np.sin(2 * np.pi * time / 3) + np.random.normal(0, 0.1, 200)
xray = 2 * np.sin(2 * np.pi * time / 3 + 0.5) + np.random.normal(0, 0.2, 200)
radio = 0.5 * np.sin(2 * np.pi * time / 3 - 0.3) + np.random.normal(0, 0.05, 200)

# Plot with shared axes
ax1.plot(time, optical, 'b-', alpha=0.7)
ax1.set_ylabel('Optical')
ax1.tick_params(labelbottom=False)  # Hide x labels except bottom

ax2.plot(time, xray, 'r-', alpha=0.7)
ax2.set_ylabel('X-ray')
ax2.tick_params(labelbottom=False)

ax3.plot(time, radio, 'g-', alpha=0.7)
ax3.set_ylabel('Radio')
ax3.set_xlabel('Time (days)')

fig.suptitle('Multi-wavelength Observations with Shared Time Axis')
plt.show()

# Generate different types of relationships
x = np.logspace(-1, 3, 100)  # 0.1 to 1000

# Different mathematical relationships
linear = 2 * x + 5
quadratic = 0.5 * x**2
exponential = np.exp(x/100)
power_law = 10 * x**(-1.5)
logarithmic = 50 * np.log10(x) + 10

# Create a comprehensive comparison
fig, axes = plt.subplots(5, 4, figsize=(14, 16))

datasets = [
    (linear, 'Linear: y = 2x + 5'),
    (quadratic, 'Quadratic: y = 0.5x²'),
    (exponential, 'Exponential: y = e^(x/100)'),
    (power_law, 'Power Law: y = 10x^(-1.5)'),
    (logarithmic, 'Logarithmic: y = 50log(x) + 10')
]

for i, (data, title) in enumerate(datasets):
    # Linear-linear
    axes[i, 0].plot(x, data, 'b-', linewidth=2)
    axes[i, 0].set_title(f'{title}\nLinear-Linear')
    axes[i, 0].grid(True, alpha=0.3)
    
    # Log-log
    axes[i, 1].loglog(x, np.abs(data), 'r-', linewidth=2)
    axes[i, 1].set_title('Log-Log')
    axes[i, 1].grid(True, alpha=0.3, which='both')
    
    # Semilog-x
    axes[i, 2].semilogx(x, data, 'g-', linewidth=2)
    axes[i, 2].set_title('Semilog-X')
    axes[i, 2].grid(True, alpha=0.3)
    
    # Semilog-y
    axes[i, 3].semilogy(x, np.abs(data), 'm-', linewidth=2)
    axes[i, 3].set_title('Semilog-Y')
    axes[i, 3].grid(True, alpha=0.3)

# Highlight which scale reveals linearity
axes[0, 0].set_facecolor('#E8F4F8')  # Linear shows linear
axes[2, 3].set_facecolor('#F8E8E8')  # Semilogy shows exponential
axes[3, 1].set_facecolor('#F8F8E8')  # Loglog shows power law
axes[4, 2].set_facecolor('#E8F8E8')  # Semilogx shows logarithmic

fig.suptitle('Choose the Scale that Reveals Your Data\'s Nature', 
             fontsize=16, fontweight='bold')
plt.tight_layout()
plt.show()

print("Key insight: The 'right' scale makes relationships linear!")
print("- Power laws → log-log")
print("- Exponential growth → semilog-y")
print("- Logarithmic growth → semilog-x")

# Real astronomical example: Galaxy luminosity function
np.random.seed(42)

# Schechter function parameters (Schechter 1976)
L_star = 1e10  # Characteristic luminosity
alpha = -1.25  # Faint-end slope
phi_star = 0.01  # Normalization

# Generate galaxy luminosities
L = np.logspace(7, 12, 1000)
phi = phi_star * (L/L_star)**alpha * np.exp(-L/L_star)

# Add observational scatter
observed_phi = phi * np.random.lognormal(0, 0.2, len(phi))

fig, axes = plt.subplots(2, 2, figsize=(10, 10))

# Linear scale - useless for this data
axes[0, 0].plot(L, observed_phi, '.', markersize=1, alpha=0.5)
axes[0, 0].set_xlabel('Luminosity [L_☉]')
axes[0, 0].set_ylabel('Φ (Number density)')
axes[0, 0].set_title('Linear Scale: Cannot See Structure')

# Log-log - reveals power law at faint end
axes[0, 1].loglog(L, observed_phi, '.', markersize=1, alpha=0.5, label='Data')
axes[0, 1].loglog(L, phi, 'r-', linewidth=2, label='Schechter Function')
axes[0, 1].set_xlabel('Luminosity [L_☉]')
axes[0, 1].set_ylabel('Φ (Number density)')
axes[0, 1].set_title('Log-Log: Reveals Power Law + Exponential Cutoff')
axes[0, 1].legend()

# Semilog-y - emphasizes exponential cutoff
axes[1, 0].semilogx(L, observed_phi, '.', markersize=1, alpha=0.5)
axes[1, 0].semilogx(L, phi, 'r-', linewidth=2)
axes[1, 0].set_xlabel('Luminosity [L_☉]')
axes[1, 0].set_ylabel('Φ (Number density)')
axes[1, 0].set_title('Semilog-X: Emphasizes Bright End Cutoff')

# Custom: log-log with ratio to model
ratio = observed_phi / phi
axes[1, 1].semilogx(L, ratio, '.', markersize=1, alpha=0.5)
axes[1, 1].axhline(y=1, color='red', linestyle='--', linewidth=2)
axes[1, 1].set_xlabel('Luminosity [L_☉]')
axes[1, 1].set_ylabel('Observed / Model')
axes[1, 1].set_title('Residual Plot: Shows Systematic Deviations')
axes[1, 1].set_ylim(0.1, 10)

fig.suptitle('Galaxy Luminosity Function: Scale Choice Matters', fontsize=14, fontweight='bold')
plt.tight_layout()
plt.show()

# Example: A reusable light curve plotting function
def plot_light_curve(time, flux, flux_err=None, period=None, 
                     title=None, figsize=(12, 8)):
    """
    Create a publication-quality light curve plot with optional phase folding.
    
    Parameters
    ----------
    time : array-like
        Time values (days)
    flux : array-like
        Flux or magnitude values
    flux_err : array-like, optional
        Flux uncertainties
    period : float, optional
        Period for phase folding (days)
    title : str, optional
        Plot title
    figsize : tuple, optional
        Figure size (width, height)
    
    Returns
    -------
    fig, axes : matplotlib figure and axes objects
    """
    if period is not None:
        fig, axes = plt.subplots(2, 2, figsize=figsize)
        axes = axes.flatten()
    else:
        fig, axes = plt.subplots(1, 1, figsize=(figsize[0], figsize[1]/2))
        axes = [axes]  # Make it iterable
    
    # Main light curve
    ax = axes[0]
    if flux_err is not None:
        ax.errorbar(time, flux, yerr=flux_err, fmt='k.', markersize=2,
                   alpha=0.5, elinewidth=0.5, capsize=0)
    else:
        ax.scatter(time, flux, s=1, alpha=0.5, color='black')
    
    ax.set_xlabel('Time (days)', fontsize=11)
    ax.set_ylabel('Relative Flux', fontsize=11)
    ax.set_title('Light Curve' if title is None else title, fontsize=12)
    ax.grid(True, alpha=0.3)
    
    # If period provided, add phase-folded plot
    if period is not None:
        phase = (time % period) / period
        
        # Phase folded
        ax = axes[1]
        if flux_err is not None:
            ax.errorbar(phase, flux, yerr=flux_err, fmt='b.', markersize=2,
                       alpha=0.3, elinewidth=0.5, capsize=0)
        else:
            ax.scatter(phase, flux, s=1, alpha=0.3, color='blue')
        ax.set_xlabel('Phase', fontsize=11)
        ax.set_ylabel('Relative Flux', fontsize=11)
        ax.set_title(f'Phase Folded (P = {period:.3f} days)', fontsize=12)
        ax.set_xlim(0, 1)
        ax.grid(True, alpha=0.3)
        
        # Double-plotted phase folded
        ax = axes[2]
        phase_double = np.concatenate([phase, phase + 1])
        flux_double = np.concatenate([flux, flux])
        if flux_err is not None:
            err_double = np.concatenate([flux_err, flux_err])
            ax.errorbar(phase_double, flux_double, yerr=err_double, 
                       fmt='r.', markersize=2, alpha=0.3, 
                       elinewidth=0.5, capsize=0)
        else:
            ax.scatter(phase_double, flux_double, s=1, alpha=0.3, color='red')
        ax.set_xlabel('Phase', fontsize=11)
        ax.set_ylabel('Relative Flux', fontsize=11)
        ax.set_title('Double Phase Plot', fontsize=12)
        ax.set_xlim(0, 2)
        ax.grid(True, alpha=0.3)
        
        # Binned phase curve - more efficient version
        ax = axes[3]
        n_bins = 50
        binned_flux, bin_edges = np.histogram(phase, bins=n_bins, weights=flux)
        bin_counts, _ = np.histogram(phase, bins=n_bins)
        bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2
        
        # Calculate proper error bars using standard error of the mean
        binned_mean = binned_flux / bin_counts
        
        # For error calculation, need variance
        flux_squared, _ = np.histogram(phase, bins=n_bins, weights=flux**2)
        variance = flux_squared / bin_counts - binned_mean**2
        binned_err = np.sqrt(variance / bin_counts)
        
        # Only plot bins with data
        valid_bins = bin_counts > 0
        ax.errorbar(bin_centers[valid_bins], binned_mean[valid_bins], 
                   yerr=binned_err[valid_bins],
                   fmt='go-', markersize=4, linewidth=1, capsize=3)
        ax.set_xlabel('Phase', fontsize=11)
        ax.set_ylabel('Relative Flux', fontsize=11)
        ax.set_title(f'Binned ({n_bins} bins)', fontsize=12)
        ax.set_xlim(0, 1)
        ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    return fig, axes

# Test the function with synthetic data
np.random.seed(42)
time = np.linspace(0, 30, 500)
period = 2.3456  # days
phase = 2 * np.pi * time / period
flux = 1.0 - 0.01 * np.sin(phase)**2  # Transit-like signal
flux += np.random.normal(0, 0.002, len(time))  # Add noise
flux_err = np.ones_like(flux) * 0.002

fig, axes = plot_light_curve(time, flux, flux_err, period=period, 
                            title='Exoplanet Transit Light Curve')
plt.show()

print("This reusable function encodes best practices:")
print("- Automatic phase folding when period is provided")
print("- Double phase plot to see continuity")
print("- Binned version to see average behavior")
print("- Consistent styling throughout")

# my_astro_plots.py - Your personal plotting library
import numpy as np
import matplotlib.pyplot as plt

def plot_spectrum(wavelength, flux, flux_err=None, lines=None, 
                  title=None, figsize=(12, 5)):
    """Plot a stellar spectrum with optional line identification."""
    fig, ax = plt.subplots(figsize=figsize)
    
    # Plot spectrum
    ax.plot(wavelength, flux, 'k-', linewidth=0.8, label='Spectrum')
    
    if flux_err is not None:
        ax.fill_between(wavelength, flux - flux_err, flux + flux_err,
                        alpha=0.3, color='gray', label='Uncertainty')
    
    # Mark spectral lines
    if lines is not None:
        for wave, name in lines:
            ax.axvline(x=wave, color='red', linestyle=':', alpha=0.5)
            ax.text(wave, ax.get_ylim()[1]*0.95, name, 
                   rotation=90, va='top', ha='right', fontsize=8)
    
    ax.set_xlabel('Wavelength (Å)', fontsize=11)
    ax.set_ylabel('Normalized Flux', fontsize=11)
    ax.set_title('Spectrum' if title is None else title, fontsize=12)
    ax.grid(True, alpha=0.3)
    ax.legend(loc='best')
    
    return fig, ax

def plot_cmd(color, magnitude, title=None, figsize=(8, 10)):
    """Create a color-magnitude diagram."""
    fig, ax = plt.subplots(figsize=figsize)
    
    # 2D histogram for density
    h = ax.hist2d(color, magnitude, bins=50, cmap='YlOrBr', 
                  density=True)
    
    ax.set_xlabel('Color (B-V)', fontsize=11)
    ax.set_ylabel('Absolute Magnitude', fontsize=11)
    ax.invert_yaxis()  # Astronomical convention
    ax.set_title('Color-Magnitude Diagram' if title is None else title, 
                fontsize=12)
    
    plt.colorbar(h[3], ax=ax, label='Density')
    
    return fig, ax

def plot_finding_chart(ra, dec, image=None, sources=None, 
                       figsize=(10, 10)):
    """Create a finding chart with marked sources."""
    fig, ax = plt.subplots(figsize=figsize, 
                           subplot_kw={'projection': 'rectilinear'})
    
    if image is not None:
        im = ax.imshow(image, cmap='gray_r', origin='lower',
                      extent=[ra.min(), ra.max(), dec.min(), dec.max()])
        plt.colorbar(im, ax=ax, label='Intensity')
    
    if sources is not None:
        for source in sources:
            circle = plt.Circle((source['ra'], source['dec']), 
                               source.get('radius', 0.1),
                               fill=False, color='red', linewidth=2)
            ax.add_patch(circle)
            ax.text(source['ra'], source['dec'] + 0.15, 
                   source.get('name', ''), 
                   ha='center', color='red')
    
    ax.set_xlabel('RA (degrees)', fontsize=11)
    ax.set_ylabel('Dec (degrees)', fontsize=11)
    ax.set_title('Finding Chart', fontsize=12)
    ax.invert_xaxis()  # RA increases to the left
    
    return fig, ax

# Example usage
print("Your personal plotting library is ready!")
print("Available functions:")
print("  - plot_light_curve(): For time series photometry")
print("  - plot_spectrum(): For spectroscopic data") 
print("  - plot_cmd(): For color-magnitude diagrams")
print("  - plot_finding_chart(): For sky position plots")

# Generate realistic stellar spectrum
np.random.seed(42)
wavelength = np.linspace(4000, 7000, 1000)  # Angstroms
continuum = 1 - 0.0001 * (wavelength - 5500)**2 / 1e6  # Continuum shape

# Add absorption lines (H-alpha, H-beta, etc.)
lines = [(6563, 50, 0.3), (4861, 30, 0.25), (4340, 25, 0.2)]  # center, width, depth
spectrum = continuum.copy()
for center, width, depth in lines:
    spectrum *= (1 - depth * np.exp(-0.5 * ((wavelength - center) / width)**2))

# Add noise
spectrum += np.random.normal(0, 0.01, len(wavelength))

# Create publication-quality spectrum plot
fig, ax = plt.subplots(figsize=(12, 5))

ax.plot(wavelength, spectrum, 'k-', linewidth=0.8, label='Observed')
ax.axhline(y=1, color='gray', linestyle='--', alpha=0.5, label='Continuum')

# Mark important lines
for center, width, depth in lines:
    ax.axvline(x=center, color='red', linestyle=':', alpha=0.5)
    ax.text(center, 0.65, f'{center}Å', rotation=90, 
            va='bottom', ha='right', fontsize=9)

ax.set_xlabel('Wavelength (Å)', fontsize=12)
ax.set_ylabel('Normalized Flux', fontsize=12)
ax.set_title('Stellar Spectrum with Balmer Lines', fontsize=14)
ax.set_xlim(4000, 7000)
ax.set_ylim(0.6, 1.1)
ax.grid(True, alpha=0.2)
ax.legend(loc='lower right')

plt.tight_layout()
plt.show()

# Create a color-magnitude diagram (CMD)
np.random.seed(42)

# Main sequence stars
n_ms = 500
color_ms = np.random.uniform(-0.3, 2.0, n_ms)
mag_ms = 4 * color_ms + np.random.normal(0, 0.5, n_ms) + 4

# Red giants
n_rg = 100
color_rg = np.random.uniform(0.8, 2.0, n_rg)
mag_rg = np.random.normal(0, 0.5, n_rg)

# White dwarfs
n_wd = 50
color_wd = np.random.uniform(-0.3, 0.5, n_wd)
mag_wd = np.random.normal(11, 0.5, n_wd)

# Create the CMD
fig, ax = plt.subplots(figsize=(8, 10))

ax.scatter(color_ms, mag_ms, c='navy', s=10, alpha=0.6, label='Main Sequence')
ax.scatter(color_rg, mag_rg, c='red', s=30, alpha=0.7, label='Red Giants')
ax.scatter(color_wd, mag_wd, c='lightblue', s=20, alpha=0.8, label='White Dwarfs')

ax.set_xlabel('B - V Color Index', fontsize=12)
ax.set_ylabel('V Magnitude', fontsize=12)
ax.set_title('Color-Magnitude Diagram\n(Hertzsprung-Russell Diagram)', fontsize=14)
ax.invert_yaxis()  # Astronomical convention
ax.grid(True, alpha=0.3)
ax.legend(loc='upper left')

# Add annotations
ax.annotate('Turn-off point', xy=(0.6, 6), xytext=(1.2, 7),
            arrowprops=dict(arrowstyle='->', color='black', alpha=0.5))

plt.tight_layout()
plt.show()

# Demonstrating proper error bar visualization
np.random.seed(42)

# Generate data with varying uncertainties
x = np.linspace(1, 10, 20)
y = 2 * x + 3 + np.random.normal(0, 1, 20)
# Heteroscedastic errors (varying with x)
y_err = 0.5 + 0.1 * x + np.random.uniform(0, 0.5, 20)

fig, axes = plt.subplots(1, 3, figsize=(14, 5))

# Standard error bars
axes[0].errorbar(x, y, yerr=y_err, fmt='o', capsize=3, capthick=1,
                 elinewidth=1, markersize=5, label='Data')
axes[0].set_title('Standard Error Bars')
axes[0].set_xlabel('X')
axes[0].set_ylabel('Y')
axes[0].grid(True, alpha=0.3)

# Error bars with confidence band
axes[1].errorbar(x, y, yerr=y_err, fmt='o', capsize=0,
                 elinewidth=0.5, markersize=5, alpha=0.7)
# Add model with confidence band
model_x = np.linspace(0, 11, 100)
model_y = 2 * model_x + 3
axes[1].plot(model_x, model_y, 'r-', label='Model')
axes[1].fill_between(model_x, model_y - 1, model_y + 1, 
                     alpha=0.3, color='red', label='1σ confidence')
axes[1].set_title('Model with Confidence Band')
axes[1].set_xlabel('X')
axes[1].legend()
axes[1].grid(True, alpha=0.3)

# Weighted fit visualization
weights = 1 / y_err**2
# Fit weighted linear model
coeffs = np.polyfit(x, y, 1, w=weights)
fit_y = np.polyval(coeffs, x)

axes[2].errorbar(x, y, yerr=y_err, fmt='o', capsize=3,
                 elinewidth=1, markersize=5 * weights/weights.max(),
                 alpha=0.6, label='Data (size ∝ weight)')
axes[2].plot(x, fit_y, 'g-', linewidth=2, label='Weighted fit')
axes[2].set_title('Weighted Fitting')
axes[2].set_xlabel('X')
axes[2].legend()
axes[2].grid(True, alpha=0.3)

fig.suptitle('Error Bar Visualization Best Practices', fontsize=14)
plt.tight_layout()
plt.show()

# Create synthetic galaxy image
np.random.seed(42)
x = np.linspace(-5, 5, 200)
y = np.linspace(-5, 5, 200)
X, Y = np.meshgrid(x, y)

# Exponential disk profile
R = np.sqrt(X**2 + Y**2)
disk = np.exp(-R / 1.5)

# Add spiral arms (simplified)
theta = np.arctan2(Y, X)
spiral = 1 + 0.3 * np.sin(2 * theta - R)
galaxy = disk * spiral

# Add noise and background
galaxy += np.random.normal(0, 0.02, galaxy.shape)
galaxy += 0.1  # Sky background

# Display with different scaling
fig, axes = plt.subplots(2, 3, figsize=(12, 8))

# Linear scale
im1 = axes[0, 0].imshow(galaxy, cmap='gray', origin='lower')
axes[0, 0].set_title('Linear Scale')
plt.colorbar(im1, ax=axes[0, 0], fraction=0.046)

# Log scale
from matplotlib.colors import LogNorm
galaxy_positive = galaxy - galaxy.min() + 1e-3  # Ensure positive
im2 = axes[0, 1].imshow(galaxy_positive, cmap='gray', 
                         norm=LogNorm(), origin='lower')
axes[0, 1].set_title('Log Scale')
plt.colorbar(im2, ax=axes[0, 1], fraction=0.046)

# Histogram equalization (adaptive)
from matplotlib.colors import PowerNorm
im3 = axes[0, 2].imshow(galaxy, cmap='gray', 
                         norm=PowerNorm(gamma=0.5), origin='lower')
axes[0, 2].set_title('Power Scale (γ=0.5)')
plt.colorbar(im3, ax=axes[0, 2], fraction=0.046)

# Different colormaps
im4 = axes[1, 0].imshow(galaxy, cmap='viridis', origin='lower')
axes[1, 0].set_title('Viridis Colormap')
plt.colorbar(im4, ax=axes[1, 0], fraction=0.046)

im5 = axes[1, 1].imshow(galaxy, cmap='hot', origin='lower')
axes[1, 1].set_title('Hot Colormap')
plt.colorbar(im5, ax=axes[1, 1], fraction=0.046)

im6 = axes[1, 2].imshow(galaxy, cmap='twilight', origin='lower')
axes[1, 2].set_title('Twilight Colormap')
plt.colorbar(im6, ax=axes[1, 2], fraction=0.046)

for ax in axes.flat:
    ax.set_xlabel('X (pixels)')
    ax.set_ylabel('Y (pixels)')

fig.suptitle('Galaxy Image with Different Displays', fontsize=14)
plt.tight_layout()
plt.show()

# Example of WCSAxes usage (simplified without actual FITS data)
# This would normally use astropy.wcs and astropy.io.fits

# Simulated example showing the concept
fig = plt.figure(figsize=(10, 8))

# Create synthetic data
ra_range = np.linspace(150, 151, 100)  # RA in degrees
dec_range = np.linspace(2, 3, 100)     # Dec in degrees
RA, DEC = np.meshgrid(ra_range, dec_range)

# Synthetic astronomical image
image = np.exp(-((RA - 150.5)**2 + (DEC - 2.5)**2) / 0.01)

# Standard matplotlib plot (pixel coordinates)
ax1 = fig.add_subplot(121)
im1 = ax1.imshow(image, origin='lower', cmap='viridis')
ax1.set_xlabel('X (pixels)')
ax1.set_ylabel('Y (pixels)')
ax1.set_title('Pixel Coordinates')
plt.colorbar(im1, ax=ax1)

# WCS-aware plot (would use WCS projection in practice)
ax2 = fig.add_subplot(122)
im2 = ax2.imshow(image, origin='lower', cmap='viridis',
                 extent=[ra_range[0], ra_range[-1], 
                        dec_range[0], dec_range[-1]])
ax2.set_xlabel('RA (degrees)')
ax2.set_ylabel('Dec (degrees)')
ax2.set_title('Sky Coordinates (WCS)')
ax2.invert_xaxis()  # RA increases to the left
plt.colorbar(im2, ax=ax2)

plt.suptitle('Pixel vs Sky Coordinates in Astronomical Images')
plt.tight_layout()
plt.show()

print("Note: For real WCS plotting, use:")
print("from astropy.wcs import WCS")
print("from astropy.io import fits")
print("ax = fig.add_subplot(111, projection=wcs)")

# Demonstrate perceptual uniformity
data = np.random.randn(10, 10)

fig, axes = plt.subplots(2, 4, figsize=(14, 7))

# Bad colormaps (not perceptually uniform)
bad_cmaps = ['jet', 'rainbow', 'nipy_spectral', 'gist_ncar']
for ax, cmap in zip(axes[0], bad_cmaps):
    im = ax.imshow(data, cmap=cmap)
    ax.set_title(f'{cmap} (Avoid!)')
    ax.axis('off')
    plt.colorbar(im, ax=ax, fraction=0.046)

# Good colormaps (perceptually uniform)
good_cmaps = ['viridis', 'plasma', 'cividis', 'twilight']
for ax, cmap in zip(axes[1], good_cmaps):
    im = ax.imshow(data, cmap=cmap)
    ax.set_title(f'{cmap} (Good!)')
    ax.axis('off')
    plt.colorbar(im, ax=ax, fraction=0.046)

fig.suptitle('Perceptually Uniform vs Non-Uniform Colormaps', fontsize=14)
plt.tight_layout()
plt.show()

# Example of journal-specific styling
# Note: SciencePlots package provides pre-configured styles

# Manual implementation of ApJ-like style
def apply_apj_style():
    """Apply Astrophysical Journal style settings."""
    plt.rcParams.update({
        'font.size': 10,
        'font.family': 'serif',
        'font.serif': ['Computer Modern Roman'],
        'text.usetex': False,  # Set True if LaTeX is installed
        'figure.figsize': (3.5, 3.5),  # Single column
        'figure.dpi': 300,
        'savefig.dpi': 300,
        'savefig.bbox': 'tight',
        'axes.linewidth': 0.5,
        'lines.linewidth': 1.0,
        'lines.markersize': 4,
        'xtick.major.width': 0.5,
        'ytick.major.width': 0.5,
    })

# Example plot with journal style
apply_apj_style()
fig, ax = plt.subplots()
x = np.linspace(0, 2*np.pi, 100)
ax.plot(x, np.sin(x), label='sin(x)')
ax.plot(x, np.cos(x), label='cos(x)')
ax.set_xlabel('Phase')
ax.set_ylabel('Amplitude')
ax.legend()
ax.grid(True, alpha=0.3)
plt.show()

# Reset to defaults
plt.rcParams.update(plt.rcParamsDefault)

# Conceptual example of yt usage (requires yt installation)
"""
# yt installation: pip install yt
import yt

# Load simulation data (supports many formats)
ds = yt.load("galaxy_simulation/output_0050")

# Create a slice plot through the simulation volume
slc = yt.SlicePlot(ds, 'z', 'density')
slc.set_cmap('density', 'viridis')
slc.annotate_timestamp()
slc.save()

# Volume rendering
sc = yt.create_scene(ds, 'density')
sc.save()

# Phase plots
plot = yt.PhasePlot(ds.all_data(), 
                    'density', 'temperature', 'cell_mass')
plot.save()
"""

# Demonstrate the concept with synthetic data
fig, axes = plt.subplots(2, 3, figsize=(14, 9))

# Simulate different yt-style visualizations
# 1. Density slice
x = np.linspace(-10, 10, 100)
y = np.linspace(-10, 10, 100)
X, Y = np.meshgrid(x, y)
density = np.exp(-(X**2 + Y**2)/10) + 0.5*np.exp(-((X-3)**2 + Y**2)/5)

im1 = axes[0, 0].imshow(density, cmap='magma', origin='lower')
axes[0, 0].set_title('Density Slice (z=0)')
plt.colorbar(im1, ax=axes[0, 0], label='ρ [g/cm³]')

# 2. Temperature slice
temperature = 1e4 * density**0.5 + np.random.normal(0, 100, density.shape)
im2 = axes[0, 1].imshow(temperature, cmap='hot', origin='lower')
axes[0, 1].set_title('Temperature Slice')
plt.colorbar(im2, ax=axes[0, 1], label='T [K]')

# 3. Velocity field
u = -Y / np.sqrt(X**2 + Y**2 + 0.1)
v = X / np.sqrt(X**2 + Y**2 + 0.1)
axes[0, 2].streamplot(x, y, u, v, density=1, color='blue', alpha=0.7)
axes[0, 2].set_title('Velocity Streamlines')
axes[0, 2].set_xlim(-10, 10)
axes[0, 2].set_ylim(-10, 10)

# 4. Phase plot
n_points = 1000
dens_sample = np.random.lognormal(0, 1, n_points)
temp_sample = 1e4 * dens_sample**0.5 + np.random.normal(0, 500, n_points)
h = axes[1, 0].hist2d(np.log10(dens_sample), np.log10(temp_sample), 
                       bins=30, cmap='YlOrBr')
axes[1, 0].set_xlabel('log(Density)')
axes[1, 0].set_ylabel('log(Temperature)')
axes[1, 0].set_title('Phase Diagram')
plt.colorbar(h[3], ax=axes[1, 0])

# 5. Projection (column density)
column_density = np.sum(np.stack([density]*20), axis=0)  # Fake projection
im5 = axes[1, 1].imshow(column_density, cmap='viridis', 
                         origin='lower', norm=LogNorm())
axes[1, 1].set_title('Column Density Projection')
plt.colorbar(im5, ax=axes[1, 1], label='Σ [g/cm²]')

# 6. Multi-field composite
axes[1, 2].imshow(density, cmap='Blues', alpha=0.7, origin='lower')
contour = axes[1, 2].contour(temperature, levels=5, colors='red', 
                              linewidths=2, alpha=0.8)
axes[1, 2].set_title('Multi-field Visualization')
axes[1, 2].clabel(contour, inline=True, fontsize=8)

fig.suptitle('yt-style Visualizations for Simulation Data', 
             fontsize=14, fontweight='bold')
plt.tight_layout()
plt.show()

print("\nyt advantages for simulation data:")
print("✓ Supports 40+ simulation codes natively")
print("✓ Handles AMR and particle data seamlessly")
print("✓ Unit-aware (automatic unit conversions)")
print("✓ Parallel processing for large datasets")
print("✓ Volume rendering capabilities")
print("✓ On-the-fly derived fields")
print("\nLearn more at: https://yt-project.org/")