Tutorials

This section contains tutorials demonstrating common use cases for zdm.

Available Notebooks

The docs/nb/ directory contains Jupyter notebooks with worked examples:

CHIME_pzDM.ipynb: Working with CHIME survey data and p(z|DM) calculations
Exploring_H0.ipynb: Exploring the effect of varying H0
Exploring_Emax.ipynb: Exploring maximum energy parameter effects
Exploring_alpha.ipynb: Exploring spectral index variations
Grid_sz.ipynb: Understanding grid size and resolution
Max_Like.ipynb: Maximum likelihood estimation
Median_DMcosmic.ipynb: Computing median cosmic DM
Omegab_H0.ipynb: Relationship between Omega_b and H0
Speedup_IGamma.ipynb: Performance optimization with gamma functions

Tutorial: Basic Likelihood Calculation

This tutorial walks through computing likelihoods for a simple case.

Step 1: Setup

import numpy as np
from zdm import parameters, survey, cosmology as cos
from zdm import misc_functions as mf
from zdm import grid as zdm_grid
from zdm import iteration as it

# Create state with default parameters
state = parameters.State()

Step 2: Initialize Cosmology

# Set cosmological parameters
cos.set_cosmology(state.params)
cos.init_dist_measures()

Step 3: Define Grid Dimensions

# DM range (pc/cm^3)
dmvals = np.linspace(0, 3000, 500)

# Redshift range
zvals = np.linspace(0.01, 3.0, 400)

Step 4: Load Survey

# Load ASKAP ICS survey
s = survey.Survey(
    state,
    survey_name='CRAFT_ICS',
    filename='CRAFT_ICS.ecsv',
    dmvals=dmvals,
    zvals=zvals
)

Step 5: Build z-DM Grid

# Get base p(DM|z) grid
zDMgrid, zvals, dmvals, smear = mf.get_zdm_grid(
    state, new=True, plot=False
)

# Build survey-specific grid
g = zdm_grid.Grid(
    survey=s,
    state=state,
    zDMgrid=zDMgrid,
    zvals=zvals,
    dmvals=dmvals,
    smear_mask=smear
)

Step 6: Compute Likelihood

# Calculate log-likelihood
ll = it.get_log_likelihood(g, s)
print(f"Log-likelihood: {ll:.2f}")

Tutorial: Parameter Exploration

Explore how likelihood varies with a parameter.

import numpy as np
import matplotlib.pyplot as plt
from zdm import loading, iteration as it

# Load CHIME data
dmvals, zvals, grids, surveys = loading.load_CHIME(Nbin=6)

# Range of H0 values to test
H0_values = np.linspace(60, 80, 21)
likelihoods = []

for H0 in H0_values:
    # Update H0 in state
    state = grids[0].state
    state.update_param('H0', H0)

    # Rebuild grids and compute likelihood
    # (simplified - full implementation needs grid rebuild)
    total_ll = sum(it.get_log_likelihood(g, s)
                   for g, s in zip(grids, surveys))
    likelihoods.append(total_ll)

# Plot results
plt.figure()
plt.plot(H0_values, likelihoods)
plt.xlabel('H0 (km/s/Mpc)')
plt.ylabel('Log-likelihood')
plt.title('Likelihood vs H0')
plt.show()

Tutorial: MCMC Parameter Estimation

Run MCMC to constrain parameters.

from zdm import MCMC
from zdm import loading
import emcee

# Load surveys
dmvals, zvals, grids, surveys = loading.load_CHIME(Nbin=6)
state = grids[0].state

# Define parameters to vary with bounds
params = {
    'H0': {'min': 50, 'max': 100},
    'gamma': {'min': -2.5, 'max': 0},
    'lEmax': {'min': 40, 'max': 43},
}

# Initial walker positions
nwalkers = 32
ndim = len(params)
p0 = np.random.uniform(
    low=[v['min'] for v in params.values()],
    high=[v['max'] for v in params.values()],
    size=(nwalkers, ndim)
)

# Setup sampler
surveys_sep = [surveys, []]  # Non-repeaters, repeaters
sampler = emcee.EnsembleSampler(
    nwalkers, ndim, MCMC.calc_log_posterior,
    args=[state, params, surveys_sep]
)

# Run MCMC
nsteps = 1000
sampler.run_mcmc(p0, nsteps, progress=True)

# Get results
samples = sampler.get_chain(discard=200, flat=True)

Tutorial: Working with Repeaters

Analyze repeating FRB populations.

from zdm import repeat_grid

# Load repeater survey data
# (requires appropriate survey file with repeater info)

# Create repeater grid
rg = repeat_grid.repeat_Grid(
    survey=s,
    state=state,
    zDMgrid=zDMgrid,
    zvals=zvals,
    dmvals=dmvals,
    smear_mask=smear
)

# Compute likelihood including repeater terms
ll = it.get_log_likelihood(rg, s, pNreps=True)

Further Resources

See the API Reference for complete function documentation
See the Parameters for all available parameters
Check the papers/ directory for publication-specific analysis scripts