Part III: Scientific Applications | Chapter 12

ML in Cosmology & Astrophysics

Galaxy classification, photometric redshifts, gravitational lensing, simulation-based inference, and neural emulators for cosmological simulations

The Data Deluge in Cosmology

Modern cosmological surveys produce datasets of unprecedented scale: the Vera C. Rubin Observatory (LSST) will catalogue ~20 billion galaxies; Euclid will map the 3D distribution of galaxies across 10 billion light-years; and the SKA will produce exabytes of radio data. Classical analysis pipelines cannot keep pace. Machine learning has become essential for extracting cosmological information from these massive, complex datasets.

This chapter covers the major ML applications in cosmology: galaxy morphology classification, photometric redshift estimation, gravitational lens detection, simulation-based inference for parameter estimation, and neural emulators that replace expensive N-body simulations.

1. Galaxy Morphology Classification

Galaxy morphology (spiral, elliptical, irregular, merging) correlates with physical properties such as star formation rate, stellar mass, and environment. The Hubble classification has been extended by citizen-science projects (Galaxy Zoo) and now by deep learning.

CNN Classification Pipeline

A convolutional neural network takes a multi-band galaxy image$I \in \mathbb{R}^{H \times W \times C}$ (where $C$ = number of photometric bands) and outputs class probabilities:

$$p(y = k | I) = \frac{e^{z_k}}{\sum_{j=1}^{K} e^{z_j}}, \quad z = f_\theta(I) \in \mathbb{R}^K$$

The cross-entropy loss for $N$ galaxies with soft labels (vote fractions from Galaxy Zoo) is:

$$\mathcal{L} = -\frac{1}{N}\sum_{i=1}^{N}\sum_{k=1}^{K} p_k^{(i)} \log \hat{p}_k^{(i)}$$

Dieleman et al. (2015) achieved near-human accuracy using rotationally augmented CNNs. Modern approaches use equivariant CNNs that are exactly invariant under rotations, avoiding the need for augmentation.

Python

script.py73 lines

import numpy as np

# Simulate galaxy classification with a simple CNN-like approach
np.random.seed(42)

N = 200
n_features = 10  # simulated image features
K = 3  # classes: spiral, elliptical, irregular

# Generate synthetic galaxy features
# Spirals: high asymmetry, blue colour
# Ellipticals: smooth, red
# Irregulars: high clumpiness

features = np.zeros((N, n_features))
labels = np.zeros((N, K))

for i in range(N):
    cls = np.random.choice(K)
    if cls == 0:  # spiral
        features[i] = np.array([0.8, 0.2, 0.7, 0.1, -0.5, 0.3, 0.6, -0.2, 0.4, 0.1])
        labels[i] = [0.85, 0.10, 0.05]
    elif cls == 1:  # elliptical
        features[i] = np.array([-0.3, 0.9, -0.2, 0.8, 0.6, -0.4, -0.3, 0.5, -0.1, 0.7])
        labels[i] = [0.05, 0.90, 0.05]
    else:  # irregular
        features[i] = np.array([0.5, -0.1, 0.3, -0.3, -0.2, 0.8, 0.2, -0.5, 0.7, -0.3])
        labels[i] = [0.10, 0.05, 0.85]
    features[i] += np.random.randn(n_features) * 0.3

# Simple softmax classifier
W = np.random.randn(n_features, K) * 0.1
b = np.zeros(K)

def softmax(z):
    e = np.exp(z - np.max(z, axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

lr = 0.05
print("Galaxy Morphology Classification:")
for epoch in range(200):
    # Forward
    logits = features @ W + b
    probs = softmax(logits)

# Cross-entropy loss with soft labels
    loss = -np.mean(np.sum(labels * np.log(probs + 1e-10), axis=1))

# Gradient
    dlogits = (probs - labels) / N
    dW = features.T @ dlogits
    db = dlogits.sum(axis=0)

W -= lr * dW
    b -= lr * db

if epoch % 40 == 0 or epoch == 199:
        pred_class = np.argmax(probs, axis=1)
        true_class = np.argmax(labels, axis=1)
        acc = np.mean(pred_class == true_class)
        print(f"  Epoch {epoch:3d}: loss = {loss:.4f}, accuracy = {acc:.3f}")

# Confusion matrix
print("\nConfusion matrix:")
classes = ["Spiral", "Elliptical", "Irregular"]
for i in range(K):
    row = []
    for j in range(K):
        count = np.sum((true_class == i) & (pred_class == j))
        row.append(f"{count:4d}")
    print(f"  True {classes[i]:>10}: {' '.join(row)}")
print(f"  {'':>16} {'  '.join(f'{c[:4]:>4}' for c in classes)} <- Predicted")

Click Run to execute the Python code

Code will be executed with Python 3 on the server

2. Photometric Redshift Estimation

Spectroscopic redshifts are accurate but expensive. Photometric redshifts (photo-z) estimate$z$ from broadband photometry (typically 5-10 filters), enabling redshift estimation for billions of galaxies.

Photo-z as a Regression Problem

Given magnitudes $\mathbf{m} = (m_u, m_g, m_r, m_i, m_z) \in \mathbb{R}^5$ (or equivalently, colours $c_{ij} = m_i - m_j$), we want to predict the redshift$z_{\text{spec}}$. The standard metrics are:

$$\text{bias} = \langle \Delta z \rangle, \quad \Delta z = \frac{z_{\text{phot}} - z_{\text{spec}}}{1 + z_{\text{spec}}}$$

$$\sigma_{\text{NMAD}} = 1.4826 \cdot \text{median}(|\Delta z - \text{median}(\Delta z)|)$$

The outlier fraction $\eta$ is the percentage of galaxies with$|\Delta z| > 0.15$. State-of-the-art photo-z methods achieve$\sigma_{\text{NMAD}} \sim 0.01\text{-}0.03$ with $\eta < 5\%$.

Probabilistic Photo-z: Mixture Density Networks

Rather than predicting a single redshift, a mixture density network (MDN) predicts the full posterior $p(z | \mathbf{m})$ as a Gaussian mixture:

$$p(z | \mathbf{m}) = \sum_{k=1}^{K} \pi_k(\mathbf{m}) \cdot \mathcal{N}(z \,|\, \mu_k(\mathbf{m}), \sigma_k^2(\mathbf{m}))$$

The network outputs $3K$ parameters: mixing coefficients$\pi_k$ (softmax), means $\mu_k$, and variances$\sigma_k^2$ (softplus). The loss is the negative log-likelihood:$\mathcal{L} = -\sum_i \log p(z_i^{\text{spec}} | \mathbf{m}_i)$. This captures multi-modal posteriors caused by colour-redshift degeneracies.

Python

script.py81 lines

import numpy as np

# Photometric redshift estimation
np.random.seed(42)

N = 500

# Simulate galaxy photometry
# Redshift determines colours via k-correction
z_true = np.random.uniform(0.0, 2.0, N)

# Simplified colour-redshift relation
# u-g increases with z, g-r has a bump at z~0.4 (4000A break), etc.
ug = 1.0 + 0.8 * z_true + 0.3 * np.random.randn(N)
gr = 0.5 + 0.4 * np.sin(2 * z_true) + 0.2 * np.random.randn(N)
ri = 0.3 + 0.3 * z_true - 0.1 * z_true**2 + 0.15 * np.random.randn(N)
iz = 0.2 + 0.2 * z_true + 0.1 * np.random.randn(N)
mag_r = 18.0 + 2.0 * z_true + 0.5 * np.random.randn(N)  # r-band magnitude

features = np.column_stack([ug, gr, ri, iz, mag_r])
features -= features.mean(axis=0)
features /= features.std(axis=0)

# Train/test split
N_train = 400
X_train, z_train = features[:N_train], z_true[:N_train]
X_test, z_test = features[N_train:], z_true[N_train:]

# Neural network: 5 -> 32 -> 16 -> 1
def relu(x): return np.maximum(0, x)
def relu_grad(x): return (x > 0).astype(float)

W1 = np.random.randn(5, 32) * np.sqrt(2.0/5)
b1 = np.zeros(32)
W2 = np.random.randn(32, 16) * np.sqrt(2.0/32)
b2 = np.zeros(16)
W3 = np.random.randn(16, 1) * np.sqrt(2.0/16)
b3 = np.zeros(1)

lr = 0.001

print("Photometric Redshift Estimation (Neural Network):")
for epoch in range(500):
    # Forward
    h1 = relu(X_train @ W1 + b1)
    h2 = relu(h1 @ W2 + b2)
    z_pred = (h2 @ W3 + b3).flatten()

loss = np.mean((z_pred - z_train)**2)

# Backward
    dz = 2 * (z_pred - z_train).reshape(-1, 1) / N_train
    dW3 = h2.T @ dz; db3 = dz.sum(0)
    dh2 = dz @ W3.T; dh2 *= relu_grad(h1 @ W2 + b2)
    dW2 = h1.T @ dh2; db2 = dh2.sum(0)
    dh1 = dh2 @ W2.T; dh1 *= relu_grad(X_train @ W1 + b1)
    dW1 = X_train.T @ dh1; db1 = dh1.sum(0)

W1 -= lr*dW1; b1 -= lr*db1
    W2 -= lr*dW2; b2 -= lr*db2
    W3 -= lr*dW3; b3 -= lr*db3

if epoch % 100 == 0 or epoch == 499:
        print(f"  Epoch {epoch:3d}: MSE = {loss:.6f}")

# Evaluate on test set
h1_t = relu(X_test @ W1 + b1)
h2_t = relu(h1_t @ W2 + b2)
z_pred_test = (h2_t @ W3 + b3).flatten()

delta_z = (z_pred_test - z_test) / (1 + z_test)
bias = np.mean(delta_z)
sigma_nmad = 1.4826 * np.median(np.abs(delta_z - np.median(delta_z)))
outlier_frac = np.mean(np.abs(delta_z) > 0.15)

print(f"\nTest set photo-z metrics:")
print(f"  Bias: {bias:.4f}")
print(f"  sigma_NMAD: {sigma_nmad:.4f}")
print(f"  Outlier fraction (|dz|>0.15): {outlier_frac*100:.1f}%")
print(f"  Pearson r: {np.corrcoef(z_test, z_pred_test)[0,1]:.4f}")

Click Run to execute the Python code

Code will be executed with Python 3 on the server

3. Gravitational Lens Finding

Strong gravitational lensing produces spectacular arcs and multiple images of background galaxies. These rare events ($\sim 1$ per $10^4$ galaxies) contain information about dark matter substructure and the Hubble constant.

The Lens Equation

For a point mass $M$, the lens equation relates the source position$\beta$ to the image position $\theta$:

$$\beta = \theta - \frac{\theta_E^2}{\theta}, \quad \theta_E = \sqrt{\frac{4GM}{c^2}\frac{D_{ls}}{D_l D_s}}$$

where $\theta_E$ is the Einstein radius. For extended mass distributions, the convergence $\kappa(\boldsymbol{\theta})$ and shear$\gamma(\boldsymbol{\theta})$ characterise the local lensing effect:

$$\kappa(\boldsymbol{\theta}) = \frac{\Sigma(\boldsymbol{\theta})}{\Sigma_{\text{cr}}}, \quad \Sigma_{\text{cr}} = \frac{c^2}{4\pi G}\frac{D_s}{D_l D_{ls}}$$

CNNs trained on simulated lens images can find lenses orders of magnitude faster than human inspection, critical for upcoming surveys that will image billions of galaxies.

Python

script.py105 lines

import numpy as np

# Simulate gravitational lens detection
np.random.seed(42)

def simulate_lens_image(has_lens=True, img_size=32):
    """Generate a simplified 1D 'image' of a galaxy with/without lensing features"""
    x = np.linspace(-3, 3, img_size)

# Background galaxy: Gaussian profile
    galaxy = np.exp(-x**2 / (2 * 0.5**2))

if has_lens:
        # Einstein ring features: two arc-like peaks
        theta_E = 1.0 + 0.3 * np.random.randn()
        arc1 = 0.5 * np.exp(-(x - theta_E)**2 / (2 * 0.15**2))
        arc2 = 0.5 * np.exp(-(x + theta_E)**2 / (2 * 0.15**2))
        # Counter-image
        counter = 0.2 * np.exp(-(x - 0.3)**2 / (2 * 0.1**2))
        image = galaxy + arc1 + arc2 + counter
    else:
        image = galaxy

# Add noise
    image += np.random.randn(img_size) * 0.1
    return image

# Generate training data
N = 400
img_size = 32
X = np.zeros((N, img_size))
y = np.zeros(N)

for i in range(N):
    is_lens = i < N // 2
    X[i] = simulate_lens_image(has_lens=is_lens, img_size=img_size)
    y[i] = 1.0 if is_lens else 0.0

# Shuffle
perm = np.random.permutation(N)
X, y = X[perm], y[perm]
X_train, y_train = X[:300], y[:300]
X_test, y_test = X[300:], y[300:]

# Normalize
mu, sigma = X_train.mean(), X_train.std()
X_train = (X_train - mu) / sigma
X_test = (X_test - mu) / sigma

# Binary classifier
def sigmoid(z): return 1.0 / (1.0 + np.exp(-np.clip(z, -500, 500)))

W1 = np.random.randn(img_size, 32) * np.sqrt(2.0/img_size)
b1 = np.zeros(32)
W2 = np.random.randn(32, 1) * np.sqrt(2.0/32)
b2 = np.zeros(1)

def relu(x): return np.maximum(0, x)
def relu_grad(x): return (x > 0).astype(float)

lr = 0.01
print("Gravitational Lens Finder (Binary Classifier):")
for epoch in range(300):
    h1 = relu(X_train @ W1 + b1)
    logits = (h1 @ W2 + b2).flatten()
    probs = sigmoid(logits)

# Binary cross-entropy
    loss = -np.mean(y_train * np.log(probs + 1e-10) + (1 - y_train) * np.log(1 - probs + 1e-10))

# Backward
    dlogits = (probs - y_train).reshape(-1, 1) / len(y_train)
    dW2 = h1.T @ dlogits; db2 = dlogits.sum(0)
    dh1 = dlogits @ W2.T; dh1 *= relu_grad(X_train @ W1 + b1)
    dW1 = X_train.T @ dh1; db1 = dh1.sum(0)

W1 -= lr*dW1; b1 -= lr*db1
    W2 -= lr*dW2; b2 -= lr*db2

if epoch % 60 == 0 or epoch == 299:
        pred = (probs > 0.5).astype(float)
        acc = np.mean(pred == y_train)
        print(f"  Epoch {epoch:3d}: loss = {loss:.4f}, accuracy = {acc:.3f}")

# Test evaluation
h1_t = relu(X_test @ W1 + b1)
probs_test = sigmoid((h1_t @ W2 + b2).flatten())
pred_test = (probs_test > 0.5).astype(float)

tp = np.sum((pred_test == 1) & (y_test == 1))
fp = np.sum((pred_test == 1) & (y_test == 0))
fn = np.sum((pred_test == 0) & (y_test == 1))
tn = np.sum((pred_test == 0) & (y_test == 0))

precision = tp / (tp + fp) if (tp + fp) > 0 else 0
recall = tp / (tp + fn) if (tp + fn) > 0 else 0
f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0

print(f"\nTest Results:")
print(f"  Accuracy:  {np.mean(pred_test == y_test):.3f}")
print(f"  Precision: {precision:.3f}")
print(f"  Recall:    {recall:.3f}")
print(f"  F1 Score:  {f1:.3f}")
print(f"  TP={int(tp)}, FP={int(fp)}, FN={int(fn)}, TN={int(tn)}")

Click Run to execute the Python code

Code will be executed with Python 3 on the server

4. Simulation-Based Inference (SBI)

In cosmology, the likelihood $p(D|\theta)$ is often intractable — we cansimulate data from the model but cannot write down the likelihood in closed form. Simulation-based inference (also called likelihood-free inference) uses simulations to approximate the posterior $p(\theta|D)$ directly.

Neural Posterior Estimation (NPE)

NPE trains a conditional density estimator (e.g., normalising flow)$q_\phi(\theta | D)$ to approximate the posterior:

Step 1: Sample parameters from the prior:$\theta_i \sim p(\theta)$

Step 2: Run simulator:$D_i \sim p(D | \theta_i)$

Step 3: Train the network by minimising:

$$\mathcal{L}(\phi) = -\frac{1}{N}\sum_{i=1}^{N} \log q_\phi(\theta_i | D_i)$$

At inference time, given observed data $D_{\text{obs}}$, the trained network directly outputs the approximate posterior$q_\phi(\theta | D_{\text{obs}}) \approx p(\theta | D_{\text{obs}})$. This amortises inference: a single forward pass replaces expensive MCMC sampling.

Neural Likelihood Estimation (NLE)

Alternatively, one can learn the likelihood $q_\phi(D|\theta) \approx p(D|\theta)$and then use standard MCMC with the learned likelihood. The advantage is that classical statistical tests (e.g., likelihood ratio) remain available.

A third approach, Neural Ratio Estimation (NRE), directly learns the likelihood-to-evidence ratio$r(\theta, D) = p(D|\theta)/p(D)$ via binary classification between$(\theta, D)$ pairs drawn jointly vs. marginally.

Python

script.py100 lines

import numpy as np

# Simulation-based inference: learn posterior for a simple cosmological model
np.random.seed(42)

# Toy model: "observe" the Hubble constant from distance-redshift data
# True model: d = (c/H0) * z * (1 + (1-q0)*z/2) (2nd order expansion)
c = 3e5  # km/s

def simulate(H0, q0, n_galaxies=20):
    """Simulate distance-redshift observations"""
    z = np.random.uniform(0.01, 0.3, n_galaxies)
    d_true = (c / H0) * z * (1 + (1 - q0) * z / 2)
    d_obs = d_true + np.random.randn(n_galaxies) * 50  # Mpc noise
    # Summary statistics
    slope = np.polyfit(z, d_obs, 1)[0]
    curvature = np.polyfit(z, d_obs, 2)[0]
    scatter = np.std(d_obs - np.polyval(np.polyfit(z, d_obs, 1), z))
    return np.array([slope, curvature, scatter])

# Generate training set: prior -> simulate -> (theta, summary)
N_sims = 2000
thetas = np.zeros((N_sims, 2))
summaries = np.zeros((N_sims, 3))

for i in range(N_sims):
    H0 = np.random.uniform(50, 100)  # prior on H0
    q0 = np.random.uniform(-1.5, 1.0)  # prior on q0
    thetas[i] = [H0, q0]
    summaries[i] = simulate(H0, q0)

# Normalize
theta_mean, theta_std = thetas.mean(0), thetas.std(0)
sum_mean, sum_std = summaries.mean(0), summaries.std(0)
thetas_n = (thetas - theta_mean) / theta_std
summaries_n = (summaries - sum_mean) / sum_std

# Train a simple MDN to approximate p(theta | summary)
# Input: summary (3D), Output: Gaussian parameters for theta (2D)
def relu(x): return np.maximum(0, x)
def relu_grad(x): return (x > 0).astype(float)

W1 = np.random.randn(3, 32) * np.sqrt(2.0/3)
b1 = np.zeros(32)
W2 = np.random.randn(32, 16) * np.sqrt(2.0/32)
b2 = np.zeros(16)
W3 = np.random.randn(16, 4) * 0.1  # output: mu_H0, mu_q0, log_sig_H0, log_sig_q0
b3 = np.zeros(4)

lr = 0.001
N_train = 1500

for epoch in range(600):
    h1 = relu(summaries_n[:N_train] @ W1 + b1)
    h2 = relu(h1 @ W2 + b2)
    out = h2 @ W3 + b3  # (N, 4)

mu = out[:, :2]
    log_sig = out[:, 2:]
    sig = np.exp(np.clip(log_sig, -5, 5))

# Gaussian NLL
    nll = 0.5 * np.sum(((thetas_n[:N_train] - mu) / sig)**2, axis=1) + np.sum(log_sig, axis=1)
    loss = np.mean(nll)

# Backward
    d_mu = (mu - thetas_n[:N_train]) / sig**2 / N_train
    d_log_sig = (1.0 - ((thetas_n[:N_train] - mu) / sig)**2) / N_train
    d_out = np.concatenate([d_mu, d_log_sig], axis=1)

dW3 = h2.T @ d_out; db3 = d_out.sum(0)
    dh2 = d_out @ W3.T; dh2 *= relu_grad(h1 @ W2 + b2)
    dW2 = h1.T @ dh2; db2 = dh2.sum(0)
    dh1 = dh2 @ W2.T; dh1 *= relu_grad(summaries_n[:N_train] @ W1 + b1)
    dW1 = summaries_n[:N_train].T @ dh1; db1 = dh1.sum(0)

W1-=lr*dW1; b1-=lr*db1; W2-=lr*dW2; b2-=lr*db2; W3-=lr*dW3; b3-=lr*db3

if epoch % 120 == 0:
        print(f"  Epoch {epoch:3d}: NLL = {loss:.4f}")

# Test: "observe" with known parameters
H0_true, q0_true = 70.0, -0.55
obs_summary = simulate(H0_true, q0_true)
obs_n = (obs_summary - sum_mean) / sum_std

h1_t = relu(obs_n.reshape(1, -1) @ W1 + b1)
h2_t = relu(h1_t @ W2 + b2)
out_t = (h2_t @ W3 + b3).flatten()

mu_pred = out_t[:2] * theta_std + theta_mean
sig_pred = np.exp(np.clip(out_t[2:], -5, 5)) * theta_std

print(f"\nSimulation-Based Inference Results:")
print(f"  True H0 = {H0_true:.1f}, True q0 = {q0_true:.2f}")
print(f"  Posterior H0 = {mu_pred[0]:.1f} +/- {sig_pred[0]:.1f}")
print(f"  Posterior q0 = {mu_pred[1]:.2f} +/- {sig_pred[1]:.2f}")
print(f"  H0 within 2-sigma: {abs(H0_true - mu_pred[0]) < 2*sig_pred[0]}")
print(f"  q0 within 2-sigma: {abs(q0_true - mu_pred[1]) < 2*sig_pred[1]}")

Click Run to execute the Python code

Code will be executed with Python 3 on the server

5. Neural Emulators for Cosmological Simulations

Full N-body cosmological simulations (e.g., Illustris, EAGLE) take millions of CPU-hours. Neural emulators learn to predict simulation outputs as a function of cosmological parameters, enabling rapid exploration of parameter space.

Power Spectrum Emulation

The matter power spectrum $P(k)$ encodes the clustering of matter as a function of wavenumber $k$. A neural emulator learns:

$$\hat{P}(k; \theta_{\text{cosmo}}) = \text{NN}_\phi(k, \Omega_m, \Omega_b, h, n_s, \sigma_8)$$

Training data consists of power spectra computed from N-body simulations at different cosmologies (e.g., a Latin hypercube sampling of the parameter space). The emulator interpolates between training cosmologies.

Examples include EuclidEmulator, CosmicEmu, and CAMELS emulators. Accuracies of$< 1\%$ are achieved for $P(k)$ over the relevant range of scales and cosmologies, with speed-ups of $10^5\text{-}10^6$ over full simulations.

Python

script.py96 lines

import numpy as np

# Power spectrum emulator: learn P(k) as function of cosmological parameters
np.random.seed(42)

def true_power_spectrum(k, Omega_m, sigma8, ns):
    """Simplified CDM-like power spectrum (toy model)"""
    # Transfer function approximation
    keq = 0.01 * Omega_m  # matter-radiation equality scale
    T_k = np.log(1 + 2.34 * k / keq) / (2.34 * k / keq) * (
        1 + 3.89 * (k / keq) + (16.1 * k / keq)**2 + (5.46 * k / keq)**3 + (6.71 * k / keq)**4
    )**(-0.25)
    # Primordial + transfer + normalisation
    P = sigma8**2 * (k / 0.05)**ns * T_k**2 * k
    return P

# Generate training data: Latin hypercube over cosmological parameters
N_cosmo = 100
N_k = 30

# Parameter ranges
Om_range = (0.2, 0.4)
s8_range = (0.6, 1.0)
ns_range = (0.9, 1.05)

k_values = np.logspace(-3, 0, N_k)  # h/Mpc

# Training set
params_train = np.column_stack([
    np.random.uniform(*Om_range, N_cosmo),
    np.random.uniform(*s8_range, N_cosmo),
    np.random.uniform(*ns_range, N_cosmo),
])
Pk_train = np.zeros((N_cosmo, N_k))
for i in range(N_cosmo):
    Pk_train[i] = true_power_spectrum(k_values, *params_train[i])

# Work in log space for better conditioning
log_Pk_train = np.log10(Pk_train + 1e-30)
log_k = np.log10(k_values)

# Normalise
param_mean, param_std = params_train.mean(0), params_train.std(0)
logPk_mean, logPk_std = log_Pk_train.mean(), log_Pk_train.std()

X_train = (params_train - param_mean) / param_std
Y_train = (log_Pk_train - logPk_mean) / logPk_std

# Simple MLP emulator: 3 params -> N_k P(k) values
def relu(x): return np.maximum(0, x)
def relu_grad(x): return (x > 0).astype(float)

W1 = np.random.randn(3, 64) * np.sqrt(2.0/3)
b1 = np.zeros(64)
W2 = np.random.randn(64, 32) * np.sqrt(2.0/64)
b2 = np.zeros(32)
W3 = np.random.randn(32, N_k) * np.sqrt(2.0/32)
b3 = np.zeros(N_k)

lr = 0.001
for epoch in range(1000):
    h1 = relu(X_train @ W1 + b1)
    h2 = relu(h1 @ W2 + b2)
    Y_pred = h2 @ W3 + b3

loss = np.mean((Y_pred - Y_train)**2)

dY = 2 * (Y_pred - Y_train) / (N_cosmo * N_k)
    dW3 = h2.T @ dY; db3 = dY.sum(0)
    dh2 = dY @ W3.T; dh2 *= relu_grad(h1 @ W2 + b2)
    dW2 = h1.T @ dh2; db2 = dh2.sum(0)
    dh1 = dh2 @ W2.T; dh1 *= relu_grad(X_train @ W1 + b1)
    dW1 = X_train.T @ dh1; db1 = dh1.sum(0)

W1-=lr*dW1; b1-=lr*db1; W2-=lr*dW2; b2-=lr*db2; W3-=lr*dW3; b3-=lr*db3

# Test on new cosmology
Om_test, s8_test, ns_test = 0.3, 0.8, 0.97
x_test = ((np.array([Om_test, s8_test, ns_test]) - param_mean) / param_std).reshape(1, -1)

h1_t = relu(x_test @ W1 + b1)
h2_t = relu(h1_t @ W2 + b2)
logPk_pred = (h2_t @ W3 + b3).flatten() * logPk_std + logPk_mean
Pk_pred = 10**logPk_pred

Pk_true = true_power_spectrum(k_values, Om_test, s8_test, ns_test)

rel_error = np.abs(Pk_pred - Pk_true) / Pk_true * 100
print("Power Spectrum Emulator Test:")
print(f"  Test cosmology: Om={Om_test}, s8={s8_test}, ns={ns_test}")
print(f"  Mean relative error: {np.mean(rel_error):.2f}%")
print(f"  Max relative error:  {np.max(rel_error):.2f}%")
print(f"\n  k [h/Mpc]    P(k) true    P(k) emul    error%")
for i in range(0, N_k, 6):
    print(f"  {k_values[i]:.4f}    {Pk_true[i]:12.4f}  {Pk_pred[i]:12.4f}    {rel_error[i]:5.1f}%")

Click Run to execute the Python code

Code will be executed with Python 3 on the server

6. Further Applications

Weak Lensing Mass Mapping

CNNs and U-Nets reconstruct dark matter mass maps from weak lensing shear measurements, outperforming traditional Kaiser-Squires inversion by capturing non-Gaussian information in the convergence field.

Transient Classification

RNNs and transformers classify astronomical transients (supernovae, kilonovae, tidal disruption events) from their multi-band light curves in real time, essential for triggering follow-up observations.

21cm Cosmology

Neural networks extract cosmological information from noisy 21cm intensity maps, separating the faint cosmological signal from foreground contamination that is$10^4\text{-}10^5$ times brighter.

Gravitational Wave Detection

Deep learning detects and characterises gravitational wave signals in LIGO/Virgo data, performing matched filtering equivalent to searching $\sim 10^5$ templates in milliseconds rather than hours.

7. Challenges & Open Problems

Distribution Shift

ML models trained on simulations must generalise to real observations. Systematic differences between simulated and real data (e.g., imperfect PSF models, selection effects, foreground contamination) cause distribution shift. Domain adaptation and calibration techniques are essential.

Uncertainty Quantification

Cosmological parameter constraints require rigorous uncertainty estimates. Standard neural networks produce point predictions without well-calibrated uncertainties. Bayesian neural networks, ensemble methods, and conformal prediction are being explored to provide coverage guarantees for ML-derived constraints.

Interpretability

When an ML model discovers an anomaly or makes a surprising classification, physicists need to understand why. Gradient-based attribution (saliency maps), attention visualisation, and symbolic distillation help extract physical insight from black-box models.

Scalability

Future surveys will produce petabytes of data. ML pipelines must scale to billions of objects while maintaining latency requirements for real-time transient classification and alert brokering. Distributed training, model compression, and edge deployment are active research areas.

Chapter Summary

• Galaxy classification uses CNNs to morphologically classify billions of galaxies, matching human accuracy.
• Photometric redshifts are estimated from broadband photometry using neural networks; mixture density networks capture multi-modal posteriors.
• Gravitational lens finding is automated by CNNs trained on simulated lensing images, critical for upcoming surveys.
• Simulation-based inference (NPE, NLE, NRE) enables Bayesian parameter estimation when the likelihood is intractable, amortising inference via neural density estimators.
• Neural emulators replace expensive N-body simulations, predicting power spectra and summary statistics as functions of cosmological parameters with $< 1\%$ error and $10^5\times$ speed-up.

← ML for Molecular Dynamics Course Overview →