Chapter 7: NumPy - The Foundation of Scientific Computing

# Question 1: What does this produce?
result = [x**2 for x in range(5) if x % 2 == 0]
# Your answer: _______

# Question 2: What's the final value of total?
total = sum([len(str(x)) for x in [10, 100, 1000]])
# Your answer: _______

# Question 3: What error (if any) occurs here?
import math
values = [1, 4, 9]
# roots = math.sqrt(values)  # Uncomment to test
# Your answer: _______

# Question 4: Can you understand this nested structure?
matrix = [[1, 2], [3, 4]]
flattened = [item for row in matrix for item in row]
# Your answer: _______

import time
import math

# Python list approach (what you know from Chapter 4)
magnitudes = [12.3, 12.5, 12.4, 12.7, 12.6] * 20000  # 100,000 measurements
fluxes = []

start = time.perf_counter()
for mag in magnitudes:
    flux = 10**(-mag/2.5)  # Convert magnitude to relative flux
    fluxes.append(flux)
list_time = time.perf_counter() - start

print(f"List approach: {list_time*1000:.2f} ms")
print(f"First 5 fluxes: {fluxes[:5]}")

import time
import numpy as np  # Standard abbreviation used universally

# NumPy array approach (what you're learning)
magnitudes = np.array([12.3, 12.5, 12.4, 12.7, 12.6] * 20000)

start = time.perf_counter()
fluxes = 10**(-magnitudes/2.5)  # Operates on entire array at once!
numpy_time = time.perf_counter() - start

print(f"NumPy approach: {numpy_time*1000:.2f} ms")
print(f"Speedup: {list_time/numpy_time:.1f}x faster")
print(f"First 5 fluxes: {fluxes[:5]}")

# Math module - works on single values only
import math
x = 2.0
math_result = math.sin(x)
print(f"math.sin({x}) = {math_result}")

# Try with a list - this fails!
# values = [0, 1, 2, 3]
# math.sin(values)  # TypeError!

# NumPy - works on single values AND arrays
x = 2.0
numpy_scalar_result = np.sin(x)
print(f"np.sin({x}) = {numpy_scalar_result}")

# NumPy shines with arrays
values = np.array([0, 1, 2, 3])
numpy_array_result = np.sin(values)
print(f"np.sin({values}) = {numpy_array_result}")

# NumPy has everything math has (and more)
print(f"\nComparison for x = {x}:")
print(f"  math.sqrt({x}) = {math.sqrt(x):.6f}")
print(f"  np.sqrt({x})   = {np.sqrt(x):.6f}")
print(f"  math.exp({x})  = {math.exp(x):.6f}")
print(f"  np.exp({x})    = {np.exp(x):.6f}")
print(f"  math.log10({x}) = {math.log10(x):.6f}")
print(f"  np.log10({x})   = {np.log10(x):.6f}")

# From Python lists (most common starting point)
measurements = [23.5, 24.1, 23.8, 24.3]
arr = np.array(measurements)
print(f"From list: {arr}, dtype: {arr.dtype}")

# Specify data type for memory efficiency
counts = np.array([1000, 2000, 1500], dtype=np.int32)
print(f"Integer array: {counts}, dtype: {counts.dtype}")

# 2D array (matrix) - like an image
image_data = np.array([[10, 20, 30],
                       [40, 50, 60],
                       [70, 80, 90]])
print(f"2D array shape: {image_data.shape}")
print(f"2D array:\n{image_data}")

# linspace: Evenly spaced values (inclusive endpoints)
# Perfect for wavelength grids, time series
wavelengths_nm = np.linspace(400, 700, 11)  # 11 points from 400 to 700 nm
print(f"Wavelengths (nm): {wavelengths_nm}")

# Convert to CGS (cm) - standard in stellar atmosphere models
wavelengths_cm = wavelengths_nm * 1e-7
print(f"Wavelengths (cm): {wavelengths_cm}")

# logspace: Logarithmically spaced values
# Essential for frequency grids, stellar masses
masses_solar = np.logspace(-1, 2, 4)  # 0.1 to 100 solar masses
masses_g = masses_solar * 1.989e33  # Convert to grams (CGS)
print("Stellar masses (g):", ", ".join(f"{x:.2e}" for x in masses_g))

# arange: Like Python's range but returns array
times_s = np.arange(0, 10, 0.1)  # 0 to 9.9 in 0.1s steps
print(f"Time points: {len(times_s)} samples from {times_s[0]} to {times_s[-1]}")

# zeros and ones: Initialize arrays
dark_frame = np.zeros((100, 100))  # 100x100 CCD dark frame
flat_field = np.ones((100, 100))   # Flat field (perfect response)
print(f"Dark frame shape: {dark_frame.shape}, sum: {dark_frame.sum()}")
print(f"Flat field shape: {flat_field.shape}, sum: {flat_field.sum()}")

# full: Create array filled with specific value
bias_level = np.full((100, 100), 500)  # CCD bias level
print(f"Bias array: all values = {bias_level[0, 0]}")

# eye: Identity matrix
identity = np.eye(3)
print(f"Identity matrix:\n{identity}")

import os  # For cleanup

# Save astronomical data in binary format
flux_data = np.random.normal(1000, 50, size=1000)
np.save('observations.npy', flux_data)  # Binary format, fast

# Load for analysis
loaded_data = np.load('observations.npy')
print(f"Loaded {len(loaded_data)} measurements")

# For text format (human-readable but slower)
np.savetxt('observations.txt', flux_data[:10], fmt='%.2f')
text_data = np.loadtxt('observations.txt')
print(f"Text file sample: {text_data[:3]}")

# Clean up files
os.remove('observations.npy')
os.remove('observations.txt')

# ALWAYS set seed for reproducibility in scientific code!
np.random.seed(42)

# Uniform distribution: random positions, phases
# Generate random sky coordinates
n_stars = 1000
ra = np.random.uniform(0, 360, n_stars)  # Right Ascension in degrees
dec = np.random.uniform(-90, 90, n_stars)  # Declination in degrees
print(f"Generated {n_stars} random sky positions")
print(f"RA range: [{ra.min():.1f}, {ra.max():.1f}]°")

# Random phases for periodic variables
phases = np.random.uniform(0, 2*np.pi, n_stars)
print(f"Phase range: [0, 2π]")

# Shorthand for uniform [0, 1)
random_fractions = np.random.rand(5)
print(f"Random fractions: {random_fractions}")

# Normal (Gaussian) distribution: measurement errors, thermal noise
# Simulate CCD readout noise
mean_counts = 1000  # electrons
read_noise = 10  # electrons RMS
n_pixels = 10000

pixel_values = np.random.normal(mean_counts, read_noise, n_pixels)
print(f"Pixel statistics:")
print(f"  Mean: {pixel_values.mean():.1f} (expected: {mean_counts})")
print(f"  Std: {pixel_values.std():.1f} (expected: {read_noise})")
print(f"  SNR: {pixel_values.mean()/pixel_values.std():.1f}")

# Shorthand for standard normal (mean=0, std=1)
standard_normal = np.random.randn(5)
print(f"Standard normal samples: {standard_normal}")

# Poisson distribution: photon counting, radioactive decay
# Simulate photon arrival statistics
mean_photons = 100  # photons per exposure
n_exposures = 1000

photon_counts = np.random.poisson(mean_photons, n_exposures)
print(f"Photon counting statistics:")
print(f"  Mean: {photon_counts.mean():.1f} (expected: {mean_photons})")
print(f"  Std: {photon_counts.std():.2f} (expected: {np.sqrt(mean_photons):.2f})")
print(f"  Variance/Mean: {photon_counts.var()/photon_counts.mean():.2f} (should be ~1)")

# Exponential distribution: time between events, decay processes
# Simulate time between supernova detections (days)
mean_interval = 30  # days
n_events = 100
intervals = np.random.exponential(mean_interval, n_events)
print(f"Supernova intervals: mean = {intervals.mean():.1f} days")

# Power-law distribution (using Pareto for x_min=1)
# Initial Mass Function approximation
alpha = 2.35  # Salpeter IMF slope
n_stars = 1000
# Generate masses from 0.1 to 100 solar masses
x = np.random.pareto(alpha, n_stars) + 1  # Pareto starts at 1
masses = 0.1 * x  # Scale to start at 0.1 solar masses
masses = masses[masses < 100]  # Truncate at 100 solar masses
print(f"Generated {len(masses)} stellar masses with power-law distribution")

# Multivariate normal: correlated parameters
# Simulate color-magnitude relationship
mean = [12, 1.0]  # Mean magnitude, mean color (B-V)
# Covariance matrix: brighter stars tend to be bluer
cov = [[1.0, -0.3],   # Variance in mag, covariance
       [-0.3, 0.25]]  # Covariance, variance in color

n_stars = 500
mag_color = np.random.multivariate_normal(mean, cov, n_stars)
magnitudes = mag_color[:, 0]
colors = mag_color[:, 1]

correlation = np.corrcoef(magnitudes, colors)[0, 1]
print(f"Generated {n_stars} stars with correlated properties")
print(f"Correlation coefficient: {correlation:.2f} (expected: ~-0.6)")

# Random choice: selecting objects from catalogs
star_types = np.array(['O', 'B', 'A', 'F', 'G', 'K', 'M'])
# Probabilities based on stellar statistics
probs = np.array([0.00003, 0.0013, 0.006, 0.03, 0.076, 0.121, 0.765])
probs = probs / probs.sum()  # Ensure normalization

n_sample = 1000
sampled_types = np.random.choice(star_types, n_sample, p=probs)

# Count occurrences
unique, counts = np.unique(sampled_types, return_counts=True)
for star_type, count in zip(unique, counts):
    print(f"Type {star_type}: {count/n_sample*100:.1f}%")

# Random permutation: bootstrap resampling
# Original dataset
data = np.array([23.5, 24.1, 23.8, 24.3, 23.9])

# Bootstrap resampling for error estimation
n_bootstrap = 1000
bootstrap_means = []

for i in range(n_bootstrap):
    # Resample with replacement
    resampled = np.random.choice(data, len(data), replace=True)
    bootstrap_means.append(resampled.mean())

bootstrap_means = np.array(bootstrap_means)
print(f"Original mean: {data.mean():.2f}")
print(f"Bootstrap mean: {bootstrap_means.mean():.2f}")
print(f"Bootstrap std error: {bootstrap_means.std():.3f}")

# Shuffle for randomization tests
shuffled = np.random.permutation(data)
print(f"Original: {data}")
print(f"Shuffled: {shuffled}")

# Basic arithmetic operates element-wise
a = np.array([1, 2, 3, 4])
b = np.array([10, 20, 30, 40])

print(f"a + b = {a + b}")      # Element-wise addition
print(f"a * b = {a * b}")      # Element-wise multiplication
print(f"a ** 2 = {a ** 2}")    # Element-wise power
print(f"b / a = {b / a}")      # Element-wise division
print(f"b // a = {b // a}")    # Element-wise floor division
print(f"b % a = {b % a}")      # Element-wise modulo

# Matrix multiplication uses @ operator
c = np.array([[1, 2], [3, 4]])
d = np.array([[5, 6], [7, 8]])
print(f"\nMatrix multiplication (c @ d):\n{c @ d}")
# Alternative: np.dot(c, d)
print(f"Using np.dot:\n{np.dot(c, d)}")

# Compare with list approach (verbose and slow)
a_list = [1, 2, 3, 4]
b_list = [10, 20, 30, 40]
result_list = []
for i in range(len(a_list)):
    result_list.append(a_list[i] + b_list[i])
print(f"\nList addition: {result_list}")  # Same result, more code!

# Trigonometric functions for coordinate transformations
angles_deg = np.array([0, 30, 45, 60, 90])
angles_rad = np.deg2rad(angles_deg)  # Convert to radians

sines = np.sin(angles_rad)
cosines = np.cos(angles_rad)

print(f"Angles (deg): {angles_deg}")
print(f"sin(θ): {sines}")
print(f"cos(θ): {cosines}")
print(f"sin²(θ) + cos²(θ): {sines**2 + cosines**2}")  # Should all be 1!

# Exponential and logarithm for magnitude scales
magnitudes = np.array([0, 1, 2, 5, 10])
flux_ratios = 10**(-magnitudes/2.5)  # Pogson's equation
print(f"\nMagnitudes: {magnitudes}")
print(f"Flux ratios: {flux_ratios}")

# Verify: magnitude difference = -2.5 * log10(flux ratio)
recovered_mags = -2.5 * np.log10(flux_ratios)
print(f"Recovered magnitudes: {recovered_mags}")

# Floating point comparison - use allclose for safety
print(f"Magnitudes match: {np.allclose(magnitudes, recovered_mags)}")

# Create sample data: Gaussian with outliers
np.random.seed(42)
data = np.random.normal(100, 15, 1000)  # Normal dist: mean=100, std=15
data[::100] = 200  # Add outliers every 100th point

# Statistical methods
print(f"Mean: {data.mean():.2f}")
print(f"Median: {np.median(data):.2f}")  # More robust to outliers
print(f"Std dev: {data.std():.2f}")
print(f"Min: {data.min():.2f}, Max: {data.max():.2f}")

# Find outliers
outliers = data > 150
n_outliers = outliers.sum()  # True counts as 1
print(f"Number of outliers (>150): {n_outliers}")

# Clean data by filtering
clean_data = data[~outliers]  # ~ means NOT
print(f"Clean mean: {clean_data.mean():.2f}")
print(f"Clean std: {clean_data.std():.2f}")

# Additional useful statistics
print(f"\nPercentiles:")
print(f"  25th: {np.percentile(data, 25):.2f}")
print(f"  75th: {np.percentile(data, 75):.2f}")
print(f"  95th: {np.percentile(data, 95):.2f}")

# 1D indexing (like lists but more powerful)
spectrum = np.array([1.0, 1.2, 1.5, 1.3, 1.1, 0.9, 0.8])
print(f"Full spectrum: {spectrum}")
print(f"First element: {spectrum[0]}")
print(f"Last element: {spectrum[-1]}")  # Negative indexing works!
print(f"Middle section: {spectrum[2:5]}")

# Fancy indexing - select multiple specific indices
important_indices = [0, 2, 4, 6]
selected = spectrum[important_indices]
print(f"Selected wavelengths: {selected}")

# 2D indexing - like accessing matrix elements
image = np.array([[10, 20, 30],
                  [40, 50, 60],
                  [70, 80, 90]])
print(f"2D array:\n{image}")
print(f"Element at row 1, col 2: {image[1, 2]}")  # Note: comma notation!
print(f"Entire row 1: {image[1, :]}")
print(f"Entire column 2: {image[:, 2]}")
print(f"Sub-image: \n{image[0:2, 1:3]}")

# Stellar catalog example
stars_mag = np.array([8.2, 12.5, 6.1, 15.3, 9.7, 11.2, 5.5])
stars_color = np.array([0.5, 1.2, 0.3, 1.8, 0.7, 1.0, 0.2])  # B-V color

# Create boolean masks
bright_mask = stars_mag < 10  # True where magnitude < 10
blue_mask = stars_color < 0.6  # True where B-V < 0.6 (blue stars)

print(f"Bright stars mask: {bright_mask}")
print(f"Bright star magnitudes: {stars_mag[bright_mask]}")

# Combine conditions with & (not 'and'), | (not 'or')
bright_and_blue = (stars_mag < 10) & (stars_color < 0.6)
print(f"Bright AND blue: {stars_mag[bright_and_blue]}")

# Count matching objects
n_bright = bright_mask.sum()  # True = 1, False = 0
print(f"Number of bright stars: {n_bright}")

# where: conditional operations and finding indices
data = np.array([1, 5, 3, 8, 2, 9, 4])
high_indices = np.where(data > 5)[0]  # Returns tuple, we want first element
print(f"Indices where data > 5: {high_indices}")
print(f"Values at those indices: {data[high_indices]}")

# Conditional replacement
clipped = np.where(data > 5, 5, data)  # Clip values above 5
print(f"Clipped data: {clipped}")

# clip: cleaner way to bound values
clipped_better = np.clip(data, 2, 8)  # Clip to range [2, 8]
print(f"Clipped with np.clip: {clipped_better}")

# unique: find unique values in catalogs
star_types = np.array(['G', 'K', 'M', 'G', 'K', 'G', 'M', 'M', 'K'])
unique_types, counts = np.unique(star_types, return_counts=True)
print(f"Unique star types: {unique_types}")
print(f"Counts: {counts}")

# histogram: bin data for analysis
magnitudes = np.random.normal(12, 2, 1000)
hist, bins = np.histogram(magnitudes, bins=20)
print(f"Histogram has {len(hist)} bins")
print(f"Bin edges from {bins[0]:.1f} to {bins[-1]:.1f}")
print(f"Peak bin has {hist.max()} stars")

# Check for NaN and Inf values
test_array = np.array([1.0, np.nan, 3.0, np.inf, 5.0])
print(f"\nNaN check: {np.isnan(test_array)}")
print(f"Inf check: {np.isinf(test_array)}")
print(f"Finite check: {np.isfinite(test_array)}")

# Simple broadcasting: scalar with array
arr = np.array([1, 2, 3, 4])
result = arr + 10  # 10 is "broadcast" to [10, 10, 10, 10]
print(f"Array + scalar: {result}")

# Row vector + column vector = matrix
row = np.array([[1, 2, 3]])      # Shape (1, 3)
col = np.array([[10], [20], [30]])  # Shape (3, 1)
matrix = row + col  # Broadcasting creates (3, 3) result
print(f"Row shape: {row.shape}")
print(f"Column shape: {col.shape}")
print(f"Result:\n{matrix}")
print(f"Result shape: {matrix.shape}")

# Practical example: Normalize each row of a matrix
data = np.array([[1, 2, 3],
                 [4, 5, 6],
                 [7, 8, 9]], dtype=float)

# Calculate row means (shape: 3,)
row_means = data.mean(axis=1)  # axis=1 means along columns
print(f"Row means: {row_means}")

# To subtract row means from each row, we need to reshape
row_means_reshaped = row_means.reshape(-1, 1)  # Shape: (3, 1)
normalized = data - row_means_reshaped  # Broadcasting!
print(f"Normalized data:\n{normalized}")
print(f"New row means: {normalized.mean(axis=1)}")  # Should be ~0

# Flat-field correction for CCD images
np.random.seed(42)

# Simulate CCD data
raw_image = np.random.poisson(1000, size=(100, 100))  # 100x100 pixels
dark_current = np.full((100, 100), 50)  # Uniform dark current
flat_field = np.ones((100, 100))
flat_field[:, :50] = 0.9  # Left half less sensitive

# Calibrate image using broadcasting
calibrated = (raw_image - dark_current) / flat_field

print(f"Raw image mean: {raw_image.mean():.1f}")
print(f"Calibrated mean: {calibrated.mean():.1f}")
print(f"Left half sensitivity: {calibrated[:, :50].mean():.1f}")
print(f"Right half sensitivity: {calibrated[:, 50:].mean():.1f}")

# Reshape: Change dimensions without changing data
data = np.arange(12)
print(f"Original: {data}")

# Reshape to 2D
matrix = data.reshape(3, 4)
print(f"As 3x4 matrix:\n{matrix}")

# Reshape to 3D
cube = data.reshape(2, 2, 3)
print(f"As 2x2x3 cube:\n{cube}")

# Use -1 to infer dimension
auto_reshape = data.reshape(3, -1)  # NumPy figures out the 4
print(f"Auto-reshape to 3x?:\n{auto_reshape}")

# Flatten back to 1D
flattened = matrix.flatten()
print(f"Flattened: {flattened}")

# Stacking arrays
a = np.array([1, 2, 3])
b = np.array([4, 5, 6])
c = np.array([7, 8, 9])

# Vertical stack (row-wise)
vstacked = np.vstack([a, b, c])
print(f"Vertical stack:\n{vstacked}")

# Horizontal stack (column-wise)
hstacked = np.hstack([a, b, c])
print(f"Horizontal stack: {hstacked}")

# Concatenate (general purpose)
concat_axis0 = np.concatenate([a, b, c])  # Default axis=0
print(f"Concatenated: {concat_axis0}")

# Stack as columns
column_stack = np.column_stack([a, b, c])
print(f"Column stack:\n{column_stack}")

# Transpose swaps axes - useful for coordinate transformations
# Example: RA/Dec coordinates to X/Y projections
ra_dec = np.array([[10.5, -25.3],   # Star 1: RA, Dec
                    [15.2, -30.1],   # Star 2: RA, Dec
                    [20.8, -22.7]])  # Star 3: RA, Dec

# Transpose to get all RAs and all Decs
coords_transposed = ra_dec.T
print(f"Original (each row is a star):\n{ra_dec}")
print(f"Transposed (row 0 = all RAs, row 1 = all Decs):\n{coords_transposed}")

all_ra = coords_transposed[0]
all_dec = coords_transposed[1]
print(f"All RA values: {all_ra}")
print(f"All Dec values: {all_dec}")

# Basic meshgrid demonstration
x = np.linspace(-2, 2, 5)  # 5 x-coordinates
y = np.linspace(-1, 1, 3)  # 3 y-coordinates
X, Y = np.meshgrid(x, y)

print(f"Original x: {x}")
print(f"Original y: {y}")
print(f"\nX coordinates (shape {X.shape}):\n{X}")
print(f"\nY coordinates (shape {Y.shape}):\n{Y}")

# Each (X[i,j], Y[i,j]) gives a coordinate pair
print(f"\nCoordinate at row 1, col 2: ({X[1,2]:.1f}, {Y[1,2]:.1f})")

# Evaluate f(x,y) = x² + y² at every grid point
Z = X**2 + Y**2  # No loops needed!
print(f"\nFunction values:\n{Z}")

# Numerical simulation example: gravitational potential field
# Simulate potential from multiple point masses
masses = [1.0, 0.5, 0.3]  # Solar masses
positions = [(0, 0), (3, 2), (-2, 1)]  # AU

# Create fine grid for field calculation
x_grid = np.linspace(-5, 5, 50)
y_grid = np.linspace(-5, 5, 50)
X_field, Y_field = np.meshgrid(x_grid, y_grid)

# Calculate gravitational potential at each grid point
G = 1  # Normalized units
potential = np.zeros_like(X_field)

for mass, (px, py) in zip(masses, positions):
    # Distance from each grid point to this mass
    R = np.sqrt((X_field - px)**2 + (Y_field - py)**2)
    # Avoid singularity at mass position
    R = np.maximum(R, 0.1)
    # Add contribution to potential (Φ = -GM/r)
    potential += -G * mass / R

print(f"Potential field shape: {potential.shape}")
print(f"Min potential: {potential.min():.2f}, Max: {potential.max():.2f}")

# Numerical gradient gives force field
Fx, Fy = np.gradient(-potential, x_grid[1]-x_grid[0], y_grid[1]-y_grid[0])
force_magnitude = np.sqrt(Fx**2 + Fy**2)
print(f"Max force magnitude: {force_magnitude.max():.2f}")

# Common use: creating synthetic star images  
x_pixels = np.linspace(0, 10, 100)
y_pixels = np.linspace(0, 10, 100)
X_img, Y_img = np.meshgrid(x_pixels, y_pixels)

# Create 2D Gaussian (like a star PSF)
sigma = 2.0
star_x, star_y = 5.0, 5.0  # Star position
psf = np.exp(-((X_img - star_x)**2 + (Y_img - star_y)**2) / (2 * sigma**2))
print(f"PSF shape: {psf.shape}, peak: {psf.max():.3f}")

# 3D meshgrid for volume simulations
x_3d = np.linspace(-1, 1, 10)
y_3d = np.linspace(-1, 1, 10) 
z_3d = np.linspace(-1, 1, 10)
X_3d, Y_3d, Z_3d = np.meshgrid(x_3d, y_3d, z_3d)
print(f"\n3D grid shape: {X_3d.shape}")  # (10, 10, 10)

# Evaluate 3D density field ρ(x,y,z) = exp(-r²)
density = np.exp(-(X_3d**2 + Y_3d**2 + Z_3d**2))
print(f"Total mass (integrated density): {density.sum():.2f}")

# gradient: numerical differentiation
# Useful for finding peaks, trends in light curves
time = np.linspace(0, 10, 100)
flux = np.sin(time) + 0.1 * np.random.randn(100)

# Calculate derivative (rate of change)
flux_gradient = np.gradient(flux, time)
print(f"Maximum rate of increase: {flux_gradient.max():.3f}")
print(f"Maximum rate of decrease: {flux_gradient.min():.3f}")

# Find turning points (where derivative ≈ 0)
turning_points = np.where(np.abs(flux_gradient) < 0.1)[0]
print(f"Found {len(turning_points)} turning points")

# Linear interpolation - crucial for spectra, light curves
wavelengths_measured = np.array([400, 500, 600, 700])  # nm
flux_measured = np.array([1.0, 1.5, 1.2, 0.8])

# Interpolate to finer grid
wavelengths_fine = np.linspace(400, 700, 31)
flux_interpolated = np.interp(wavelengths_fine, 
                              wavelengths_measured, 
                              flux_measured)

print(f"Original: {len(wavelengths_measured)} points")
print(f"Interpolated: {len(wavelengths_fine)} points")
print(f"Flux at 550 nm: {np.interp(550, wavelengths_measured, flux_measured):.3f}")

# Create a signal with multiple frequencies
t = np.linspace(0, 1, 500)  # 1 second, 500 samples
sampling_rate = 500  # Hz (samples per second)

# Build composite signal
signal = np.sin(2 * np.pi * 5 * t)  # 5 Hz component
signal += 0.5 * np.sin(2 * np.pi * 10 * t)  # 10 Hz component  
signal += 0.3 * np.sin(2 * np.pi * 50 * t)  # 50 Hz component
signal += 0.2 * np.random.normal(size=t.shape)  # Noise

# Compute FFT
fft = np.fft.fft(signal)

# Get corresponding frequencies
# fftfreq returns frequencies in cycles per unit of sample spacing
# Since our sample spacing is 1/500 seconds, frequencies are in Hz
freqs = np.fft.fftfreq(len(t), d=1/sampling_rate)

# Alternative: manually calculate frequencies
# freqs_manual = np.arange(len(t)) * (sampling_rate / len(t))

# Power spectrum (squared magnitude)
power = np.abs(fft)**2

# Due to symmetry, only look at positive frequencies
pos_mask = freqs > 0
freqs_pos = freqs[pos_mask]
power_pos = power[pos_mask]

# Find peaks (simplified - use scipy.signal.find_peaks for real work)
threshold = power_pos.max() / 10
peak_indices = np.where(power_pos > threshold)[0]
peak_freqs = freqs_pos[peak_indices]

print(f"Sampling rate: {sampling_rate} Hz")
print(f"Frequency resolution: {freqs[1]-freqs[0]:.2f} Hz")
print(f"Nyquist frequency: {sampling_rate/2} Hz")
print(f"Detected peaks at: {peak_freqs} Hz")
print(f"Expected: 5, 10, and 50 Hz")

# Practical example: Finding periodic signal in noisy data
# Simulate exoplanet transit light curve with periodic dips
time = np.linspace(0, 30, 3000)  # 30 days of observations
period = 3.456  # Planet orbital period in days
transit_duration = 0.15  # days

# Create light curve with transits
flux = np.ones_like(time)
phase = (time % period) / period
in_transit = phase < (transit_duration / period)
flux[in_transit] *= 0.99  # 1% transit depth

# Add realistic noise
np.random.seed(42)
flux += np.random.normal(0, 0.002, len(time))  # 0.2% noise

# FFT to find periodicity
fft_flux = np.fft.fft(flux - flux.mean())  # Remove DC component
freqs_flux = np.fft.fftfreq(len(time), time[1] - time[0])  # Frequencies in 1/day

# Power spectrum
power_flux = np.abs(fft_flux)**2

# Look for peak in physically reasonable range (periods from 0.5 to 10 days)
freq_mask = (freqs_flux > 0.1) & (freqs_flux < 2)
peak_freq = freqs_flux[freq_mask][np.argmax(power_flux[freq_mask])]
detected_period = 1 / peak_freq

print(f"True period: {period:.3f} days")
print(f"Detected period: {detected_period:.3f} days")
print(f"Error: {abs(detected_period - period)*24*60:.1f} minutes")

# Calculate covariance and correlation between datasets
# Simulate correlated stellar properties
np.random.seed(42)
n_stars = 100
temperature = np.random.normal(5800, 500, n_stars)  # Kelvin
# Luminosity correlates with temperature (Stefan-Boltzmann: L ∝ T⁴)
# Add noise to make it realistic
luminosity_true = (temperature/5800)**4  # Normalized to solar
luminosity = luminosity_true + np.random.normal(0, 0.1, n_stars)

# Manual calculation to understand the math
temp_mean = temperature.mean()
lum_mean = luminosity.mean()
temp_centered = temperature - temp_mean
lum_centered = luminosity - lum_mean

# Covariance: average of products of deviations
covariance_manual = (temp_centered * lum_centered).mean()
print(f"Manual covariance: {covariance_manual:.6f}")

# Standard deviations
temp_std = temperature.std()
lum_std = luminosity.std()

# Correlation coefficient
correlation_manual = covariance_manual / (temp_std * lum_std)
print(f"Manual correlation: {correlation_manual:.3f}")

# Using NumPy functions
# Covariance matrix
cov_matrix = np.cov(temperature, luminosity)
print(f"\nCovariance matrix:\n{cov_matrix}")
print(f"Temperature variance: {cov_matrix[0,0]:.1f}")
print(f"Luminosity variance: {cov_matrix[1,1]:.6f}")
print(f"Temp-Lum covariance: {cov_matrix[0,1]:.6f}")

# Correlation coefficient
corr_matrix = np.corrcoef(temperature, luminosity)
print(f"\nCorrelation matrix:\n{corr_matrix}")
print(f"Correlation coefficient: {corr_matrix[0,1]:.3f}")

# Interpretation
if abs(corr_matrix[0,1]) > 0.7:
    strength = "strong"
elif abs(corr_matrix[0,1]) > 0.3:
    strength = "moderate"
else:
    strength = "weak"
print(f"This indicates a {strength} correlation between temperature and luminosity")

# Multiple variables: complete covariance analysis
# Add more stellar properties
mass = np.random.normal(1, 0.3, n_stars)  # Solar masses
# Radius correlates with mass (mass-radius relation)
radius = mass**0.8 + np.random.normal(0, 0.05, n_stars)
# Age (uncorrelated with current properties for main sequence)
age = np.random.uniform(1, 10, n_stars)  # Gyr

# Stack all variables
data = np.vstack([temperature, luminosity, mass, radius, age])

# Full covariance matrix
full_cov = np.cov(data)
print("Full covariance matrix shape:", full_cov.shape)

# Full correlation matrix (easier to interpret)
full_corr = np.corrcoef(data)
labels = ['Temp', 'Lum', 'Mass', 'Radius', 'Age']

print("\nCorrelation matrix:")
print("       ", "  ".join(f"{l:>7}" for l in labels))
for i, label in enumerate(labels):
    print(f"{label:>7}", "  ".join(f"{full_corr[i,j]:7.3f}" for j in range(5)))

print("\nStrong correlations (|r| > 0.5):")
for i in range(5):
    for j in range(i+1, 5):
        if abs(full_corr[i,j]) > 0.5:
            print(f"  {labels[i]}-{labels[j]}: {full_corr[i,j]:.3f}")

# Cross-correlation for time series alignment
# Useful for aligning light curves or spectra
signal1 = np.sin(np.linspace(0, 10, 100))
signal2 = np.sin(np.linspace(0, 10, 100) + 1)  # Phase shifted

# Cross-correlation finds the shift
correlation = np.correlate(signal1, signal2, mode='same')
lag = np.argmax(correlation) - len(correlation)//2
print(f"Maximum correlation at lag: {lag} samples")

# This tells us signal2 is shifted relative to signal1

import time
import numpy as np

# Compare performance: loops vs vectorization
size = 100000
a = np.random.randn(size)
b = np.random.randn(size)

# Method 1: Python loop
start = time.perf_counter()
result_loop = []
for i in range(size):
    result_loop.append(a[i] * b[i] + a[i]**2)
loop_time = time.perf_counter() - start

# Method 2: NumPy vectorization
start = time.perf_counter()
result_vector = a * b + a**2
vector_time = time.perf_counter() - start

print(f"Loop time: {loop_time*1000:.2f} ms")
print(f"Vector time: {vector_time*1000:.2f} ms")
print(f"Speedup: {loop_time/vector_time:.1f}x")

# Memory usage comparison
array_float64 = np.ones(1000000, dtype=np.float64)
array_float32 = np.ones(1000000, dtype=np.float32)
array_float16 = np.ones(1000000, dtype=np.float16)

print(f"Memory usage for 1 million elements:")
print(f"float64: {array_float64.nbytes / 1e6:.1f} MB")
print(f"float32: {array_float32.nbytes / 1e6:.1f} MB")
print(f"float16: {array_float16.nbytes / 1e6:.1f} MB")

"""
Performance comparison: Lists vs NumPy
Run these on your machine - results vary by hardware!
"""
import time
import numpy as np

def benchmark_operation(name, list_func, numpy_func, size=100000):
    """Benchmark a single operation."""
    # Setup data
    list_data = list(range(size))
    numpy_data = np.arange(size, dtype=float)
    
    # Time list operation
    start = time.perf_counter()
    list_result = list_func(list_data)
    list_time = time.perf_counter() - start
    
    # Time NumPy operation
    start = time.perf_counter()
    numpy_result = numpy_func(numpy_data)
    numpy_time = time.perf_counter() - start
    
    speedup = list_time / numpy_time
    return list_time * 1000, numpy_time * 1000, speedup

# Define operations to benchmark
operations = {
    "Element-wise square": (
        lambda x: [i**2 for i in x],
        lambda x: x**2
    ),
    "Sum all elements": (
        lambda x: sum(x),
        lambda x: x.sum()
    ),
    "Element-wise sqrt": (
        lambda x: [i**0.5 for i in x],
        lambda x: np.sqrt(x)
    ),
    "Find maximum": (
        lambda x: max(x),
        lambda x: x.max()
    )
}

print("Performance Comparison (100,000 elements):")
print("-" * 60)
print(f"{'Operation':<20} {'List (ms)':<12} {'NumPy (ms)':<12} {'Speedup':<10}")
print("-" * 60)

for name, (list_op, numpy_op) in operations.items():
    list_ms, numpy_ms, speedup = benchmark_operation(name, list_op, numpy_op)
    print(f"{name:<20} {list_ms:<12.2f} {numpy_ms:<12.3f} {speedup:<10.1f}x")

import time
import numpy as np

# Matrix multiplication comparison (smaller size due to O(n³) complexity)
def matrix_multiply_lists(A, B):
    """Matrix multiplication with nested lists."""
    n = len(A)
    C = [[0]*n for _ in range(n)]
    for i in range(n):
        for j in range(n):
            for k in range(n):
                C[i][j] += A[i][k] * B[k][j]
    return C

# Create test matrices
n = 100  # 100x100 matrices
A_list = [[float(i+j) for j in range(n)] for i in range(n)]
B_list = [[float(i-j) for j in range(n)] for i in range(n)]
A_numpy = np.array(A_list)
B_numpy = np.array(B_list)

# Benchmark matrix multiplication
start = time.perf_counter()
C_list = matrix_multiply_lists(A_list, B_list)
list_matmul_time = time.perf_counter() - start

start = time.perf_counter()
C_numpy = A_numpy @ B_numpy
numpy_matmul_time = time.perf_counter() - start

print(f"\nMatrix Multiplication (100×100):")
print(f"  Nested loops: {list_matmul_time*1000:.1f} ms")
print(f"  NumPy (@):    {numpy_matmul_time*1000:.3f} ms")
print(f"  Speedup:      {list_matmul_time/numpy_matmul_time:.0f}x")