Architecture

This document describes the high-level architecture of the zdm package and how the various components interact.

Overview

The zdm package follows a layered architecture:

  1. Configuration Layer: Parameter management via dataclasses

  2. Data Layer: Survey definitions and FRB observations

  3. Computation Layer: Cosmology, p(DM|z), and grid calculations

  4. Analysis Layer: Likelihood computation and MCMC inference

┌─────────────────────────────────────────────────────────┐
│                    Analysis Layer                       │
│              (MCMC, iteration, analyze_cube)            │
└─────────────────────────────────────────────────────────┘
                           │
┌─────────────────────────────────────────────────────────┐
│                  Computation Layer                      │
│         (Grid, cosmology, pcosmic, energetics)          │
└─────────────────────────────────────────────────────────┘
                           │
┌─────────────────────────────────────────────────────────┐
│                     Data Layer                          │
│              (Survey, survey_data, beams)               │
└─────────────────────────────────────────────────────────┘
                           │
┌─────────────────────────────────────────────────────────┐
│                Configuration Layer                      │
│            (State, parameters, data_class)              │
└─────────────────────────────────────────────────────────┘

Core Classes

State

State is the central configuration object containing all model parameters. It aggregates several parameter dataclasses:

  • cosmo: Cosmological parameters (H0, Omega_m, etc.)

  • FRBdemo: FRB population demographics

  • energy: Luminosity function parameters

  • host: Host galaxy DM distribution

  • MW: Milky Way DM contributions

  • IGM: Intergalactic medium parameters

  • width: Intrinsic width distribution

  • scat: Scattering parameters

  • rep: Repeater population parameters

  • analysis: Analysis control flags

state = parameters.State()
# All parameter groups are accessible as attributes
state.cosmo.H0  # Hubble constant
state.energy.gamma  # Luminosity function slope

Survey

Survey represents an FRB survey with:

  • Instrument properties (beam pattern, frequency, threshold)

  • Survey metadata (observation time, field of view)

  • Detected FRBs with measured DMs and positions

  • Detection efficiency as function of DM

Surveys are initialized from ECSV files in zdm/data/Surveys/.

Grid

Grid is the computational core, building a 2D probability distribution over redshift and DM:

  1. Takes a Survey and State as input

  2. Computes detection thresholds and efficiencies

  3. Applies beam pattern weighting

  4. Calculates expected detection rates per z-DM cell

The grid represents P(detection | z, DM, survey) weighted by the intrinsic FRB rate density.

Data Flow

Typical Workflow

1. Initialize State with parameters
         │
         ▼
2. Load Survey from data file
         │
         ▼
3. Initialize cosmology distances
         │
         ▼
4. Compute p(DM|z) grid (pcosmic)
         │
         ▼
5. Build Grid for Survey
         │
         ▼
6. Compute likelihood via iteration
         │
         ▼
7. MCMC or optimization over parameters

Grid Construction

The Grid construction involves several steps:

  1. parse_grid: Store z, DM arrays and base p(DM|z) grid

  2. calc_dV: Compute comoving volume elements per z bin

  3. smear_dm: Apply DM smearing from host and halo contributions

  4. calc_thresholds: Compute energy thresholds for detection

  5. calc_pdv: Combine volume and threshold calculations

  6. set_evolution: Apply source evolution model

  7. calc_rates: Compute final detection rates

Likelihood Calculation

The get_log_likelihood() function computes:

\[\ln \mathcal{L} = \sum_i \ln p({\rm DM}_i, z_i | \theta) - N_{\rm exp}(\theta) + \ln N_{\rm obs}!\]

Where:

  • Sum is over observed FRBs

  • \(N_{\rm exp}\) is expected number of detections

  • \(\theta\) represents model parameters

Key Modules

cosmology

Implements Lambda-CDM cosmology calculations:

  • Distance measures: comoving (DM), angular diameter (DA), luminosity (DL)

  • Volume elements: differential comoving volume dV/dz

  • Source evolution functions: SFR-based and (1+z)^n power-law models

  • Energy-flux conversions for FRB observables

Uses interpolation tables for fast array operations on redshift grids.

pcosmic

Computes p(DM_cosmic | z), the probability distribution of cosmic DM given redshift. Implements the Macquart relation from Macquart et al. (2020, Nature 581, 391) with the feedback parameter F controlling variance in the cosmic baryon distribution.

energetics

Implements FRB luminosity/energy functions using the upper incomplete gamma function. Uses spline interpolation of pre-computed values for computational efficiency during grid calculations.

Supports multiple luminosity function forms:

  • Power-law: Simple truncated power-law with Emin, Emax, gamma

  • Gamma function: Schechter-like function with exponential cutoff

  • Spline interpolation: Pre-computed gamma function lookups for speed

iteration

Contains likelihood calculation routines for z-DM grids:

  • get_log_likelihood: Main likelihood function combining all components

  • calc_likelihoods_1D: For DM-only fits (non-localized FRBs)

  • calc_likelihoods_2D: For DM+z fits (localized FRBs)

Supports various likelihood components including p(DM,z), Poisson number counts, SNR distributions, and width/scattering contributions.

MCMC

Parameter estimation using emcee:

  • calc_log_posterior: Posterior evaluation

  • Supports uniform priors with configurable bounds

  • Handles multiple surveys simultaneously

File Organization

zdm/
├── __init__.py
├── parameters.py      # State and parameter dataclasses
├── data_class.py      # Base dataclass utilities
├── survey.py          # Survey class
├── grid.py            # Grid class
├── repeat_grid.py     # Repeater grid extension
├── cosmology.py       # Cosmological calculations
├── pcosmic.py         # p(DM|z) calculations
├── energetics.py      # Luminosity functions
├── iteration.py       # Likelihood calculations
├── MCMC.py            # MCMC parameter estimation
├── loading.py         # High-level loading functions
├── misc_functions.py  # Utility functions
├── figures.py         # Plotting routines
├── beams.py           # Beam pattern handling
├── data/
│   ├── Surveys/       # Survey definition files
│   ├── BeamData/      # Beam response files
│   └── Cube/          # Precomputed grids
└── tests/             # Test suite

Extension Points

Adding New Surveys

  1. Create an ECSV file in zdm/data/Surveys/ with FRB data

  2. Add beam pattern files to zdm/data/BeamData/ if needed

  3. Load via Survey class

Custom Luminosity Functions

  1. Add function to energetics.py

  2. Register in luminosity_function parameter options

  3. Update grid initialization to use new function

New Source Evolution Models

  1. Add evolution function to cosmology.py

  2. Register in source_evolution parameter options

  3. Grid will automatically use via set_evolution()