diff --git a/.readthedocs.yaml b/.readthedocs.yaml new file mode 100644 index 00000000..8805a95a --- /dev/null +++ b/.readthedocs.yaml @@ -0,0 +1,30 @@ +# Read the Docs configuration file +# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details + +version: 2 + +# Set the OS, Python version, and other tools +build: + os: ubuntu-22.04 + tools: + python: "3.11" + +# Build documentation in the docs/ directory with Sphinx +sphinx: + configuration: docs/conf.py + fail_on_warning: false + +# Optional: include PDF and ePub outputs +formats: + - pdf + +# Install dependencies +python: + install: + # Install docs requirements + - requirements: docs/requirements.txt + # Install the package itself (without external FRB deps for docs) + - method: pip + path: . + extra_requirements: + - docs diff --git a/docs/Makefile b/docs/Makefile new file mode 100644 index 00000000..5a4746e4 --- /dev/null +++ b/docs/Makefile @@ -0,0 +1,18 @@ +# Minimal makefile for Sphinx documentation + +# You can set these variables from the command line. +SPHINXOPTS = +SPHINXBUILD = sphinx-build +SOURCEDIR = . +BUILDDIR = _build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +.PHONY: help Makefile + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/docs/api.rst b/docs/api.rst new file mode 100644 index 00000000..374ced77 --- /dev/null +++ b/docs/api.rst @@ -0,0 +1,238 @@ +.. _api: + +============= +API Reference +============= + +This section provides detailed API documentation for all public modules +and classes in the ``zdm`` package. + +Core Modules +============ + +These modules form the backbone of the zdm package, providing parameter +management, survey handling, and z-DM grid computation. + +parameters +---------- + +Central configuration system for FRB z-DM analysis. All model parameters +are organized into dataclasses grouped by category (cosmology, FRB demographics, +host galaxy, IGM, etc.). The ``State`` class aggregates all parameter groups. + +.. automodapi:: zdm.parameters + :no-inheritance-diagram: + +survey +------ + +FRB Survey class for modeling telescope observations. Encapsulates instrument +characteristics, beam patterns, detection efficiencies, and detected FRB data. +Survey files are loaded from ECSV format with header metadata and FRB tables. + +.. automodapi:: zdm.survey + :no-inheritance-diagram: + +grid +---- + +Core z-DM grid class for FRB population modeling. Computes 2D probability +distributions of FRB detection rates as a function of redshift and DM by +combining cosmological volumes, p(DM|z), telescope efficiency, and luminosity +functions. + +.. automodapi:: zdm.grid + :no-inheritance-diagram: + +repeat_grid +----------- + +Grid extension for repeating FRB population modeling. Models expected numbers +of single and repeated detections given a distribution of repeater rates +following dN/dR ~ R^Rgamma. + +.. automodapi:: zdm.repeat_grid + :no-inheritance-diagram: + +Computation Modules +=================== + +Physics and statistics modules for computing cosmological quantities, +DM probability distributions, and energy functions. + +cosmology +--------- + +Lambda CDM cosmology calculations including distance measures (comoving, +angular diameter, luminosity), volume elements, and source evolution functions. +Uses interpolation tables for fast array operations. + +.. automodapi:: zdm.cosmology + :no-inheritance-diagram: + +pcosmic +------- + +Probability distribution of cosmic dispersion measure given redshift, p(DM|z). +Implements the Macquart relation from Macquart et al. (2020) with the feedback +parameter F controlling variance. Also provides host galaxy DM convolution kernels. + +.. automodapi:: zdm.pcosmic + :no-inheritance-diagram: + +energetics +---------- + +FRB luminosity/energy function implementations including power-law and +gamma-function distributions. Uses spline interpolation of the upper incomplete +gamma function for computational efficiency during grid calculations. + +.. automodapi:: zdm.energetics + :no-inheritance-diagram: + +iteration +--------- + +Likelihood calculation routines for z-DM grids. Computes log-likelihoods of +FRB survey data given model predictions, including components for p(DM,z), +Poisson number counts, SNR distributions, and width/scattering. + +.. automodapi:: zdm.iteration + :no-inheritance-diagram: + +Utility Modules +=============== + +Helper functions for data loading, plotting, and common operations. + +loading +------- + +High-level functions for loading surveys and initializing analysis state. +Provides convenience functions like ``set_state()`` for creating properly +configured State objects with default or best-fit parameters. + +.. automodapi:: zdm.loading + :no-inheritance-diagram: + +misc_functions +-------------- + +Miscellaneous utility functions including grid initialization, parameter +updates, probability calculations, and other common operations used +throughout the package. + +.. automodapi:: zdm.misc_functions + :no-inheritance-diagram: + +figures +------- + +Plotting functions for visualizing z-DM grids and FRB data. Provides +publication-quality plots of probability distributions and analysis results. + +.. automodapi:: zdm.figures + :no-inheritance-diagram: + +beams +----- + +Telescope beam pattern modeling utilities. Provides functions for generating +and loading beam patterns (Gaussian, Airy, measured) that affect solid angle +and sensitivity variations across the field of view. + +.. automodapi:: zdm.beams + :no-inheritance-diagram: + +Analysis Modules +================ + +Parameter estimation and analysis tools. + +MCMC +---- + +MCMC parameter estimation using the emcee package. Provides functions for +running samplers, computing log-posteriors, and exploring parameter space +with support for multiple surveys and uniform priors. + +.. automodapi:: zdm.MCMC + :no-inheritance-diagram: + +MCMC_analysis +------------- + +Analysis utilities for MCMC chains including convergence diagnostics, +corner plots, and posterior summaries. + +.. automodapi:: zdm.MCMC_analysis + :no-inheritance-diagram: + +analyze_cube +------------ + +Tools for analyzing pre-computed parameter cubes. + +.. automodapi:: zdm.analyze_cube + :no-inheritance-diagram: + +vvmax +----- + +V/Vmax statistical tests for FRB population analysis. + +.. automodapi:: zdm.vvmax + :no-inheritance-diagram: + +Data Classes +============ + +Base classes and data structures. + +data_class +---------- + +Base dataclass utilities for parameter management. Provides common functionality +for serialization, dictionary access, and parameter metadata handling used by +all parameter dataclasses. + +.. automodapi:: zdm.data_class + :no-inheritance-diagram: + +survey_data +----------- + +Data structures for survey metadata and FRB observations. + +.. automodapi:: zdm.survey_data + :no-inheritance-diagram: + +Specialized Modules +=================== + +Additional modules for specific use cases. + +galactic_dm_models +------------------ + +Models for Galactic DM contributions including different halo models. + +.. automodapi:: zdm.galactic_dm_models + :no-inheritance-diagram: + +optical +------- + +Optical counterpart and host galaxy association utilities. + +.. automodapi:: zdm.optical + :no-inheritance-diagram: + +io +-- + +Input/output utilities for file handling including JSON I/O and +grid data persistence. + +.. automodapi:: zdm.io + :no-inheritance-diagram: diff --git a/docs/architecture.rst b/docs/architecture.rst new file mode 100644 index 00000000..084fa9ea --- /dev/null +++ b/docs/architecture.rst @@ -0,0 +1,257 @@ +.. _architecture: + +============ +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 + +.. code-block:: text + + ┌─────────────────────────────────────────────────────────┐ + │ 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 +----- + +:class:`~zdm.parameters.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 + +.. code-block:: python + + state = parameters.State() + # All parameter groups are accessible as attributes + state.cosmo.H0 # Hubble constant + state.energy.gamma # Luminosity function slope + +Survey +------ + +:class:`~zdm.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 +---- + +:class:`~zdm.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 +---------------- + +.. code-block:: text + + 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 :func:`~zdm.iteration.get_log_likelihood` function computes: + +.. math:: + + \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 +- :math:`N_{\rm exp}` is expected number of detections +- :math:`\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 +================= + +.. code-block:: text + + 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 :class:`~zdm.survey.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()`` diff --git a/docs/conf.py b/docs/conf.py new file mode 100644 index 00000000..3a9b2bd1 --- /dev/null +++ b/docs/conf.py @@ -0,0 +1,113 @@ +# Configuration file for the Sphinx documentation builder. +# +# For the full list of built-in configuration values, see the documentation: +# https://www.sphinx-doc.org/en/master/usage/configuration.html + +import os +import sys + +# Add the project root to the path for autodoc +sys.path.insert(0, os.path.abspath('..')) + +# -- Project information ----------------------------------------------------- +project = 'zdm' +copyright = '2024, Clancy James and contributors' +author = 'Clancy James' + +# The full version, including alpha/beta/rc tags +release = '0.1' +version = '0.1' + +# -- General configuration --------------------------------------------------- +extensions = [ + 'sphinx.ext.autodoc', + 'sphinx.ext.autosummary', + 'sphinx.ext.napoleon', + 'sphinx.ext.viewcode', + 'sphinx.ext.intersphinx', + 'sphinx.ext.mathjax', + 'sphinx_rtd_theme', + 'sphinx_automodapi.automodapi', + 'sphinx_automodapi.smart_resolver', +] + +# Autosummary settings +autosummary_generate = True +numpydoc_show_class_members = False + +# Napoleon settings for Google/NumPy style docstrings +napoleon_google_docstring = True +napoleon_numpy_docstring = True +napoleon_include_init_with_doc = True +napoleon_include_private_with_doc = False +napoleon_include_special_with_doc = True +napoleon_use_admonition_for_examples = False +napoleon_use_admonition_for_notes = False +napoleon_use_admonition_for_references = False +napoleon_use_ivar = False +napoleon_use_param = True +napoleon_use_rtype = True +napoleon_type_aliases = None + +# Intersphinx mapping +intersphinx_mapping = { + 'python': ('https://docs.python.org/3/', None), + 'numpy': ('https://numpy.org/doc/stable/', None), + 'scipy': ('https://docs.scipy.org/doc/scipy/', None), + 'astropy': ('https://docs.astropy.org/en/stable/', None), + 'matplotlib': ('https://matplotlib.org/stable/', None), +} + +# Templates path +templates_path = ['_templates'] + +# List of patterns to exclude +exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store', 'nb/.ipynb_checkpoints'] + +# The suffix of source filenames +source_suffix = '.rst' + +# The master toctree document +master_doc = 'index' + +# -- Options for HTML output ------------------------------------------------- +html_theme = 'sphinx_rtd_theme' + +html_theme_options = { + 'canonical_url': '', + 'analytics_id': '', + 'logo_only': False, + 'display_version': True, + 'prev_next_buttons_location': 'bottom', + 'style_external_links': False, + 'style_nav_header_background': '#2980B9', + # Toc options + 'collapse_navigation': True, + 'sticky_navigation': True, + 'navigation_depth': 4, + 'includehidden': True, + 'titles_only': False +} + +html_static_path = ['_static'] + +# -- Options for LaTeX output ------------------------------------------------ +latex_elements = { + 'papersize': 'letterpaper', + 'pointsize': '10pt', +} + +# Autodoc settings +autodoc_default_options = { + 'members': True, + 'member-order': 'bysource', + 'special-members': '__init__', + 'undoc-members': True, + 'exclude-members': '__weakref__' +} + +# Mock imports for modules that may not be available during doc build +autodoc_mock_imports = [ + 'ne2001', + 'frb', +] diff --git a/docs/contributing.rst b/docs/contributing.rst new file mode 100644 index 00000000..658dd2fa --- /dev/null +++ b/docs/contributing.rst @@ -0,0 +1,77 @@ +.. _contributing: + +============ +Contributing +============ + +We welcome contributions to ``zdm``! This document provides guidelines +for contributing to the project. + +Getting Started +=============== + +1. Fork the repository on GitHub +2. Clone your fork locally: + + .. code-block:: bash + + git clone https://github.com/YOUR_USERNAME/zdm.git + cd zdm + +3. Install in development mode: + + .. code-block:: bash + + pip install -e .[dev] + +4. Create a branch for your changes: + + .. code-block:: bash + + git checkout -b my-feature-branch + +Running Tests +============= + +Before submitting changes, ensure all tests pass: + +.. code-block:: bash + + # Run full test suite + pytest + + # Run with coverage + pytest --cov=zdm + + # Run specific test file + pytest zdm/tests/test_energetics.py + +Code Style +========== + +We follow PEP 8 guidelines. Check your code with: + +.. code-block:: bash + + tox -e codestyle + +Submitting Changes +================== + +1. Commit your changes with descriptive messages +2. Push to your fork +3. Open a Pull Request against the main repository +4. Ensure CI tests pass + +Reporting Issues +================ + +Please report issues on the GitHub issue tracker: +https://github.com/FRBs/zdm/issues + +Include: + +- A clear description of the problem +- Steps to reproduce +- Expected vs actual behavior +- Python and package versions diff --git a/docs/index.rst b/docs/index.rst new file mode 100644 index 00000000..08c536ad --- /dev/null +++ b/docs/index.rst @@ -0,0 +1,78 @@ +.. zdm documentation master file + +=== +zdm +=== + +**zdm** is a Python package for Fast Radio Burst (FRB) redshift-dispersion measure +(z-DM) calculations. It provides tools for statistical modeling of FRB populations, +computing likelihoods, and constraining cosmological and FRB population parameters. + +The package implements the theoretical framework for relating observed FRB dispersion +measures to their redshifts, accounting for contributions from the Milky Way, host +galaxies, and the intergalactic medium. + +.. image:: https://github.com/FRBs/zdm/workflows/CI%20Tests/badge.svg + :target: https://github.com/FRBs/zdm/actions + +Key Features +============ + +- **Cosmological Modeling**: Lambda-CDM cosmology with configurable parameters +- **z-DM Grids**: Efficient computation of probability distributions over redshift and dispersion measure +- **Survey Support**: Built-in support for CHIME, ASKAP, Parkes and other FRB surveys +- **Likelihood Analysis**: Compute likelihoods for FRB populations given survey data +- **MCMC Parameter Estimation**: Integration with emcee for Bayesian parameter inference +- **Repeater Modeling**: Support for repeating FRB population models + +Getting Started +=============== + +.. toctree:: + :maxdepth: 2 + :caption: User Guide + + installation + quickstart + tutorials + +.. toctree:: + :maxdepth: 2 + :caption: Reference + + architecture + parameters + api + +.. toctree:: + :maxdepth: 1 + :caption: Development + + contributing + +Citation +======== + +If you use ``zdm`` in your research, please cite: + +.. code-block:: bibtex + + @article{james2022, + author = {James, Clancy W. and others}, + title = {A measurement of Hubble's Constant using Fast Radio Bursts}, + journal = {MNRAS}, + year = {2022} + } + +License +======= + +``zdm`` is licensed under a 3-clause BSD style license. See the LICENSE file +in the repository for details. + +Indices and tables +================== + +* :ref:`genindex` +* :ref:`modindex` +* :ref:`search` diff --git a/docs/installation.rst b/docs/installation.rst new file mode 100644 index 00000000..4b893c37 --- /dev/null +++ b/docs/installation.rst @@ -0,0 +1,129 @@ +.. _installation: + +============ +Installation +============ + +Requirements +============ + +``zdm`` requires Python 3.10 or later and depends on several scientific Python packages: + +- numpy >= 1.18 +- scipy >= 1.4 +- astropy >= 5.2.1 +- matplotlib >= 3.3 +- pandas >= 1.3 +- emcee >= 3.1.4 +- h5py >= 3.10.0 +- mpmath >= 1.2.0 + +Additionally, ``zdm`` requires two external FRB-related packages: + +- `ne2001 `_: Galactic electron density model +- `frb `_: FRB utilities library + +Installing from Source +====================== + +Clone the repository and install in development mode: + +.. code-block:: bash + + git clone https://github.com/FRBs/zdm.git + cd zdm + pip install -e . + +For development with testing tools: + +.. code-block:: bash + + pip install -e .[dev] + +Installing Dependencies +======================= + +The external FRB packages must be installed from GitHub: + +.. code-block:: bash + + pip install git+https://github.com/FRBs/ne2001.git#egg=ne2001 + pip install git+https://github.com/FRBs/FRB.git#egg=frb + +Full Installation Script +======================== + +For a complete installation from scratch: + +.. code-block:: bash + + # Clone and enter the repository + git clone https://github.com/FRBs/zdm.git + cd zdm + + # Install external dependencies + pip install git+https://github.com/FRBs/ne2001.git#egg=ne2001 + pip install git+https://github.com/FRBs/FRB.git#egg=frb + + # Install zdm in development mode + pip install -e .[dev] + +Verifying Installation +====================== + +To verify your installation, run the test suite: + +.. code-block:: bash + + pytest + +Or run a quick import test: + +.. code-block:: python + + import zdm + from zdm import parameters, survey, grid + print("zdm installed successfully!") + +Using tox +========= + +For isolated testing environments, use tox: + +.. code-block:: bash + + pip install tox + tox -e test-alldeps + +Available tox environments: + +- ``test``: Basic test suite +- ``test-alldeps``: Tests with all optional dependencies +- ``test-astropydev``: Tests with development version of astropy +- ``codestyle``: Check code style with pycodestyle + +Troubleshooting +=============== + +ne2001 Installation Issues +-------------------------- + +If you encounter issues installing ``ne2001``, ensure you have a C compiler available: + +.. code-block:: bash + + # On Ubuntu/Debian + sudo apt-get install build-essential + + # On macOS with Homebrew + xcode-select --install + +Missing frb Package +------------------- + +The ``frb`` package is required for DM calculations. If import errors occur, reinstall: + +.. code-block:: bash + + pip uninstall frb + pip install git+https://github.com/FRBs/FRB.git#egg=frb diff --git a/docs/make.bat b/docs/make.bat new file mode 100644 index 00000000..4555c171 --- /dev/null +++ b/docs/make.bat @@ -0,0 +1,22 @@ +@ECHO OFF + +pushd %~dp0 + +REM Command file for Sphinx documentation + +if "%SPHINXBUILD%" == "" ( + set SPHINXBUILD=sphinx-build +) +set SOURCEDIR=. +set BUILDDIR=_build + +if "%1" == "" goto help + +%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% +goto end + +:help +%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% + +:end +popd diff --git a/docs/parameters.rst b/docs/parameters.rst new file mode 100644 index 00000000..5cd923ce --- /dev/null +++ b/docs/parameters.rst @@ -0,0 +1,368 @@ +.. _parameters: + +========== +Parameters +========== + +The ``zdm`` package uses a hierarchical parameter system organized into +dataclasses. All parameters are accessible through the :class:`~zdm.parameters.State` +object. + +State Object +============ + +The :class:`~zdm.parameters.State` object is the central configuration container: + +.. code-block:: python + + from zdm import parameters + + state = parameters.State() + + # Access parameter groups + print(state.cosmo.H0) # Cosmology parameters + print(state.energy.gamma) # Energy parameters + print(state.host.lmean) # Host galaxy parameters + +Updating Parameters +------------------- + +Single parameter update: + +.. code-block:: python + + state.update_param('H0', 70.0) + +Multiple parameter update via dictionary: + +.. code-block:: python + + vparams = { + 'cosmo': {'H0': 70.0}, + 'energy': {'gamma': -1.1} + } + state.update_param_dict(vparams) + +Cosmology Parameters (cosmo) +============================ + +Parameters for Lambda-CDM cosmology. Default values from Planck18. + +.. list-table:: + :widths: 20 15 15 50 + :header-rows: 1 + + * - Parameter + - Default + - Unit + - Description + * - ``H0`` + - 67.66 + - km/s/Mpc + - Hubble constant + * - ``Omega_k`` + - 0.0 + - -- + - Curvature density parameter + * - ``Omega_lambda`` + - 0.6889 + - -- + - Dark energy density parameter + * - ``Omega_m`` + - 0.3111 + - -- + - Matter density parameter + * - ``Omega_b`` + - 0.0490 + - -- + - Baryon density parameter + * - ``Omega_b_h2`` + - 0.0224 + - -- + - Baryon density times h^2 + * - ``fix_Omega_b_h2`` + - True + - -- + - Keep Omega_b_h2 fixed when varying H0 + +FRB Demographics Parameters (FRBdemo) +===================================== + +Parameters controlling FRB source population evolution. + +.. list-table:: + :widths: 20 15 15 50 + :header-rows: 1 + + * - Parameter + - Default + - Unit + - Description + * - ``source_evolution`` + - 0 + - -- + - Evolution function: 0=SFR^n, 1=(1+z)^(2.7n) + * - ``alpha_method`` + - 1 + - -- + - Scaling method: 0=k-correction, 1=rate interpretation + * - ``sfr_n`` + - 1.77 + - -- + - Scaling of FRB rate with star-formation rate + * - ``lC`` + - 3.3249 + - -- + - log10 rate constant (Gpc^-3 day^-1 at z=0) + +Energy Parameters (energy) +========================== + +Parameters for the FRB luminosity/energy function. + +.. list-table:: + :widths: 20 15 15 50 + :header-rows: 1 + + * - Parameter + - Default + - Unit + - Description + * - ``lEmin`` + - 30.0 + - log10(erg) + - Minimum FRB energy + * - ``lEmax`` + - 41.84 + - log10(erg) + - Maximum FRB energy + * - ``alpha`` + - 1.54 + - -- + - Spectral index (rate ~ nu^alpha) + * - ``gamma`` + - -1.16 + - -- + - Luminosity function slope + * - ``luminosity_function`` + - 2 + - -- + - LF type: 0=power-law, 1=gamma, 2=spline+gamma, 3=gamma+linear+log10 + +Host Galaxy Parameters (host) +============================= + +Parameters for the host galaxy DM contribution. + +.. list-table:: + :widths: 20 15 15 50 + :header-rows: 1 + + * - Parameter + - Default + - Unit + - Description + * - ``lmean`` + - 2.16 + - log10(pc/cm^3) + - Log-mean of host DM distribution + * - ``lsigma`` + - 0.51 + - log10(pc/cm^3) + - Log-sigma of host DM distribution + +Milky Way Parameters (MW) +========================= + +Parameters for Galactic DM contributions. + +.. list-table:: + :widths: 20 15 15 50 + :header-rows: 1 + + * - Parameter + - Default + - Unit + - Description + * - ``ISM`` + - 35.0 + - pc/cm^3 + - Assumed DM for Galactic ISM + * - ``DMhalo`` + - 50.0 + - pc/cm^3 + - DM contribution from Galactic halo + * - ``halo_method`` + - 0 + - -- + - Halo model: 0=uniform, 1=Yamasaki & Totani, 2=Das+ + * - ``sigmaDMG`` + - 0.0 + - -- + - Fractional uncertainty in Galactic ISM DM + * - ``sigmaHalo`` + - 15.0 + - pc/cm^3 + - Uncertainty in halo DM + * - ``logu`` + - False + - -- + - Use log-normal for DMG distributions + +IGM Parameters (IGM) +==================== + +Parameters for the intergalactic medium DM distribution. + +.. list-table:: + :widths: 20 15 15 50 + :header-rows: 1 + + * - Parameter + - Default + - Unit + - Description + * - ``logF`` + - -0.49 + - -- + - log10(F), cosmic web fluctuation parameter + +Width Parameters (width) +======================== + +Parameters for intrinsic FRB width distribution. + +.. list-table:: + :widths: 20 15 15 50 + :header-rows: 1 + + * - Parameter + - Default + - Unit + - Description + * - ``Wlogmean`` + - -0.29 + - log10(ms) + - Log-mean of intrinsic width distribution + * - ``Wlogsigma`` + - 0.65 + - log10(ms) + - Log-sigma of intrinsic width distribution + * - ``Wmethod`` + - 2 + - -- + - Width method: 0=ignore, 1=intrinsic, 2=+scattering, 3=+z-dep, 4=specific + * - ``WidthFunction`` + - 2 + - -- + - Width function: 0=log-constant, 1=log-normal, 2=half-lognormal + * - ``WNbins`` + - 12 + - -- + - Number of width bins + * - ``Wthresh`` + - 0.5 + - -- + - Starting fraction for width histogramming + +Scattering Parameters (scat) +============================ + +Parameters for FRB scattering distribution. + +.. list-table:: + :widths: 20 15 15 50 + :header-rows: 1 + + * - Parameter + - Default + - Unit + - Description + * - ``Slogmean`` + - -1.3 + - log10(ms) + - Mean of log-scattering distribution at 600 MHz + * - ``Slogsigma`` + - 0.2 + - log10(ms) + - Sigma of log-scattering distribution + * - ``Sfnorm`` + - 1000 + - MHz + - Reference frequency for scattering + * - ``Sfpower`` + - -4.0 + - -- + - Frequency scaling power-law index + * - ``ScatFunction`` + - 2 + - -- + - Scattering function: 0=log-constant, 1=lognormal, 2=half-lognormal + * - ``Sbackproject`` + - False + - -- + - Calculate p(tau|w,DM,z) arrays + +Repeater Parameters (rep) +========================= + +Parameters for repeating FRB population models. + +.. list-table:: + :widths: 20 15 15 50 + :header-rows: 1 + + * - Parameter + - Default + - Unit + - Description + * - ``lRmin`` + - -3.0 + - log10(day^-1) + - Minimum repeater rate + * - ``lRmax`` + - 1.0 + - log10(day^-1) + - Maximum repeater rate + * - ``Rgamma`` + - -2.375 + - -- + - Differential index of repeater density + * - ``RC`` + - 0.01 + - -- + - Constant repeater density + * - ``RE0`` + - 1e39 + - erg + - Energy at which rates are defined + +Analysis Parameters (analysis) +============================== + +Parameters controlling analysis behavior. + +.. list-table:: + :widths: 20 15 15 50 + :header-rows: 1 + + * - Parameter + - Default + - Unit + - Description + * - ``NewGrids`` + - True + - -- + - Generate new z-DM grids + * - ``sprefix`` + - "Std" + - -- + - Calculation detail: "Full" or "Std" + * - ``min_lat`` + - -1.0 + - degrees + - Minimum absolute Galactic latitude + * - ``DMG_cut`` + - None + - pc/cm^3 + - Maximum DMG to include diff --git a/docs/quickstart.rst b/docs/quickstart.rst new file mode 100644 index 00000000..fb03d72b --- /dev/null +++ b/docs/quickstart.rst @@ -0,0 +1,184 @@ +.. _quickstart: + +========== +Quickstart +========== + +This guide walks you through the basic workflow of using ``zdm`` to analyze +FRB populations and compute likelihoods. + +Basic Concepts +============== + +The ``zdm`` package is built around several key concepts: + +1. **State**: A configuration object holding all model parameters +2. **Survey**: Represents an FRB survey with instrument properties and detected FRBs +3. **Grid**: A 2D probability grid over redshift (z) and dispersion measure (DM) +4. **Likelihood**: Statistical comparison between model predictions and observations + +Setting Up a State +================== + +The :class:`~zdm.parameters.State` object holds all model parameters: + +.. code-block:: python + + from zdm import parameters + from zdm import loading + + # Create a state with default parameters + state = loading.set_state() + + # Or create a blank state and customize + state = parameters.State() + + # Modify parameters + vparams = { + 'cosmo': {'H0': 70.0}, + 'energy': {'gamma': -1.1, 'lEmax': 41.5}, + 'host': {'lmean': 2.2, 'lsigma': 0.5} + } + state.update_param_dict(vparams) + +Initializing Cosmology +====================== + +Before computing grids, initialize the cosmological distance measures: + +.. code-block:: python + + from zdm import cosmology as cos + + # Initialize with state parameters + cos.set_cosmology(state.params) + cos.init_dist_measures() + +Loading a Survey +================ + +Surveys are loaded from data files in the ``zdm/data/Surveys/`` directory: + +.. code-block:: python + + import numpy as np + from zdm import survey + + # Define DM and z grids + dmvals = np.linspace(0, 3000, 1000) + zvals = np.linspace(0, 3, 500) + + # Load an ASKAP survey + s = survey.Survey( + state, + survey_name='CRAFT_ICS', + filename='CRAFT_ICS.ecsv', + dmvals=dmvals, + zvals=zvals + ) + +Building a Grid +=============== + +The :class:`~zdm.grid.Grid` computes detection probabilities across z-DM space: + +.. code-block:: python + + from zdm import misc_functions as mf + + # Get the z-DM probability grid (p(DM|z)) + zDMgrid, zvals, dmvals, smear = mf.get_zdm_grid( + state, + new=True, + plot=False + ) + + # Create the survey grid + from zdm import grid as zdm_grid + + g = zdm_grid.Grid( + survey=s, + state=state, + zDMgrid=zDMgrid, + zvals=zvals, + dmvals=dmvals, + smear_mask=smear + ) + +Computing Likelihoods +===================== + +Compute the log-likelihood of the model given observed FRBs: + +.. code-block:: python + + from zdm import iteration as it + + # Compute log-likelihood + ll = it.get_log_likelihood(g, s) + print(f"Log-likelihood: {ll:.2f}") + +Complete Example +================ + +Here's a complete example loading CHIME data and computing likelihoods: + +.. code-block:: python + + from zdm import loading + from zdm import cosmology as cos + from zdm import iteration as it + + # Load CHIME survey and grids + dmvals, zvals, grids, surveys = loading.load_CHIME( + Nbin=6, + make_plots=False + ) + + # Compute total likelihood across all declination bins + total_ll = 0 + for g, s in zip(grids, surveys): + ll = it.get_log_likelihood(g, s) + total_ll += ll + print(f"Survey {s.name}: ll = {ll:.2f}") + + print(f"Total log-likelihood: {total_ll:.2f}") + +Working with Parameters +======================= + +Individual parameters can be updated and the grid recalculated: + +.. code-block:: python + + # Update a single parameter + state.update_param('H0', 73.0) + + # Update multiple parameters + vparams = { + 'energy': {'gamma': -1.2}, + 'host': {'lmean': 2.0} + } + state.update_param_dict(vparams) + + # Rebuild grid with new parameters + # (Note: some parameters require full grid rebuild) + +Visualization +============= + +The ``figures`` module provides plotting utilities: + +.. code-block:: python + + from zdm import figures + + # Plot the z-DM grid + figures.plot_grid(g, s, show=True) + +Next Steps +========== + +- See :ref:`parameters` for a complete list of model parameters +- See :ref:`architecture` for details on the code structure +- See :ref:`tutorials` for more detailed examples diff --git a/docs/requirements.txt b/docs/requirements.txt new file mode 100644 index 00000000..19437192 --- /dev/null +++ b/docs/requirements.txt @@ -0,0 +1,21 @@ +# Documentation build requirements for ReadTheDocs +sphinx>=4.0 +sphinx_rtd_theme>=1.0 +sphinx-automodapi>=0.14 +numpydoc>=1.1 + +# Package dependencies needed for autodoc +numpy>=1.18 +scipy>=1.4 +astropy>=5.2.1 +matplotlib>=3.3 +pandas>=1.3 +emcee>=3.1.4 +h5py>=3.10.0 +mpmath>=1.2.0 +PyYAML>=5.1 +IPython>=7.2.0 +corner>=2.2.3 +tqdm>=4.67.1 +importlib_resources>=6.0.0 +cmasher>=1.9 diff --git a/docs/tutorials.rst b/docs/tutorials.rst new file mode 100644 index 00000000..22286e57 --- /dev/null +++ b/docs/tutorials.rst @@ -0,0 +1,218 @@ +.. _tutorials: + +========= +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 +------------- + +.. code-block:: python + + 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 +---------------------------- + +.. code-block:: python + + # Set cosmological parameters + cos.set_cosmology(state.params) + cos.init_dist_measures() + +Step 3: Define Grid Dimensions +------------------------------ + +.. code-block:: python + + # 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 +------------------- + +.. code-block:: python + + # 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 +----------------------- + +.. code-block:: python + + # 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 +-------------------------- + +.. code-block:: python + + # 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. + +.. code-block:: python + + 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. + +.. code-block:: python + + 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. + +.. code-block:: python + + 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 :ref:`api` for complete function documentation +- See the :ref:`parameters` for all available parameters +- Check the ``papers/`` directory for publication-specific analysis scripts diff --git a/zdm/MCMC.py b/zdm/MCMC.py index 42d00fc5..5cf95581 100644 --- a/zdm/MCMC.py +++ b/zdm/MCMC.py @@ -1,13 +1,33 @@ """ -File: MCMC.py +MCMC parameter estimation for FRB z-DM analysis. + +This module provides functions for running Markov Chain Monte Carlo (MCMC) +parameter estimation using the emcee package. It interfaces with the zdm +likelihood calculations to explore the parameter space and constrain +FRB population and cosmological parameters. + +Main Functions +-------------- +- `calc_log_posterior`: Compute log-posterior for a parameter vector +- `run_mcmc`: Execute MCMC sampling with emcee +- `get_initial_walkers`: Initialize walker positions + +Features +-------- +- Uniform priors with configurable bounds +- Optional log/linear priors for specific parameters (DMhalo, host DM) +- Support for multiple surveys and repeater populations +- Grid re-initialization on each evaluation for parameter exploration + +Example +------- +>>> from zdm import MCMC +>>> params = {'gamma': {'min': -2.5, 'max': -0.5}, ...} +>>> sampler = MCMC.run_mcmc(state, params, surveys, nwalkers=32, nsteps=1000) +>>> samples = sampler.get_chain(flat=True) + Author: Jordan Hoffmann Date: 06/12/23 -Purpose: - Contains functions used for MCMC runs of the zdm code. MCMC_wrap.py is the - main function which does the calling and this holds functions which do the - MCMC analysis. - - This function re-initialises the grids on every run """ import numpy as np @@ -33,28 +53,43 @@ def calc_log_posterior(param_vals, state, params, surveys_sep, Pn=False, pNreps=True, ptauw=False, log_halo=False, lin_host=False, ind_surveys=False, g0info=None): - """ - Calculates the log posterior for a given set of parameters. Assumes uniform - priors between the minimum and maximum values provided in 'params'. + """Calculate log-posterior probability for a parameter vector. - Inputs: - param_vals (np.array) = Array of the parameter values for this step - state (params.state) = State object to modify - params (dictionary) = Parameter names, min and max values - surveys_sep (list) = surveys_sep[0] : list of non-repeater surveys - surveys_sep[1] : list of repeater surveys - grid_params (dictionary) = nz, ndm, dmmax - Pn (bool) = Include Pn or not - pNreps (bool) = Include p(N repeaters) or not - ptauw (bool) = Include p(tau,w) or not - log_halo (bool) = Use a log uniform prior on DMhalo - lin_host (bool) = Use a linear uniform prior on host mean - ind_surveys (bool) = Return likelihoods for each survey - g0info (list) = List of [zDMgrid, zvals, DMvals] Passed to use as speedup if needed - - Outputs: - llsum (double) = Total log likelihood for param_vals which is equivalent - to log posterior (un-normalised) due to uniform priors + This is the main function called by emcee samplers. It evaluates the + log-posterior (proportional to log-likelihood for uniform priors) by + building grids and computing likelihoods for all surveys. + + Parameters + ---------- + param_vals : ndarray + Array of parameter values for this MCMC step. + state : parameters.State + State object to be updated with new parameter values. + params : dict + Dictionary defining parameters to vary. Each key is a parameter name, + with value dict containing 'min' and 'max' for prior bounds. + surveys_sep : list + Two-element list: [non_repeater_surveys, repeater_surveys]. + Pn : bool, optional + Include Poisson likelihood for total number of FRBs. Default False. + pNreps : bool, optional + Include likelihood for number of repeaters. Default True. + ptauw : bool, optional + Include p(tau, width) likelihood. Default False. + log_halo : bool, optional + Use log-uniform prior on DMhalo. Default False. + lin_host : bool, optional + Use linear-uniform prior on host DM mean. Default False. + ind_surveys : bool, optional + If True, return list of individual survey likelihoods. Default False. + g0info : list, optional + Pre-computed [zDMgrid, zvals, DMvals] for speedup. + + Returns + ------- + float or tuple + Log-posterior value. Returns -inf if parameters outside prior bounds. + If ind_surveys=True, returns (llsum, ll_list) with individual likelihoods. """ # t0 = time.time() diff --git a/zdm/beams.py b/zdm/beams.py index f46a02b3..33e8d8de 100644 --- a/zdm/beams.py +++ b/zdm/beams.py @@ -1,4 +1,24 @@ -# collection of functions to handle telescope beam effects +""" +Telescope beam pattern modeling utilities. + +This module provides functions for modeling and loading telescope beam patterns, +which are essential for computing FRB detection efficiencies. Different telescopes +have different beam shapes (Gaussian, Airy, measured) that affect the solid angle +and sensitivity variations across the field of view. + +Main Functions +-------------- +- `gauss_beam`: Generate Gaussian beam pattern +- `load_beam`: Load measured beam pattern from file +- `Airy_beam`: Generate Airy disk beam pattern + +The beam is typically represented as a histogram of response values (b) and +corresponding solid angle fractions (omega), allowing efficient integration +over the beam when computing detection rates. + +Author: C.W. James +""" + from importlib.resources import files import os import numpy as np diff --git a/zdm/cosmology.py b/zdm/cosmology.py index 292d8555..5cf7b35d 100644 --- a/zdm/cosmology.py +++ b/zdm/cosmology.py @@ -1,13 +1,39 @@ -################ COSMOLOGY.PY ############### - -# Author: Clancy W. James -# clancy.w.james@gmail.com - -# This file contains functions relating to -# cosmological evaluations. It simply -# implements standard Lambda CDM cosmology -# written with alpha such that F_nu ~ \nu**-alpha -############################################## +""" +Cosmology module for the zdm package. + +This module implements standard Lambda CDM cosmology calculations for +Fast Radio Burst (FRB) redshift-dispersion measure analysis. It provides +functions for computing cosmological distance measures, volume elements, +and energy-fluence conversions. + +The module uses a convention where the spectral index alpha is defined such +that F_nu ~ nu^(-alpha), i.e., positive alpha corresponds to steeper spectra +at higher frequencies. + +Key Features +------------ +- Hubble parameter and scale factor calculations +- Distance measures: comoving (DM), angular diameter (DA), luminosity (DL) +- Cosmological volume elements (dV) and time-volume elements (dVdtau) +- Interpolated lookup tables for fast array operations +- Energy-fluence conversions for FRB analysis +- Source evolution functions (SFR, power-law) + +Usage +----- +Before using interpolated functions (dm, da, dl, dv, dvdtau), call +`init_dist_measures()` to populate the lookup tables. The cosmology +can be configured via `set_cosmology()`. + +Example +------- +>>> from zdm import cosmology as cos +>>> cos.init_dist_measures() +>>> z = 0.5 +>>> d_L = cos.dl(z) # Luminosity distance in Mpc + +Author: Clancy W. James (clancy.w.james@gmail.com) +""" import scipy.constants as constants @@ -67,14 +93,29 @@ def print_cosmology(params): - """ Print cosmological parameters - - The current values of cosmological parameters are printed. + """Print the current cosmological parameters. + + Parameters + ---------- + params : dict or State + Parameter dictionary or State object containing 'cosmo' key + with CosmoParams object. """ print("Hubble constant in default cosmology, H0: ",params['cosmo'].H0," [km/s/Mpc]") #print("Hubble constant in current epoch, H0: ",params['cosmo'].current_H0," [km/s/Mpc]") def set_cosmology(params): + """Set the global cosmology parameters for this module. + + This function must be called before using distance measure functions + if non-default cosmology is desired. + + Parameters + ---------- + params : dict or State + Parameter dictionary or State object containing 'cosmo' key + with CosmoParams object specifying H0, Omega_m, Omega_lambda, Omega_k. + """ global cosmo cosmo = params['cosmo'] @@ -110,50 +151,147 @@ def stupid_trick(a,b,c,d): def H(z): - """Hubble parameter (km/s/Mpc)""" + """Hubble parameter (km/s/Mpc) + + Args: + z (float): redshift + + Returns: + float: Hubble parameter [km/s/Mpc] + """ return E(z)*cosmo.H0 def E(z): - """scale factor, assuming a simplified cosmology.""" + """Dimensionless Hubble parameter E(z) = H(z)/H0. + + Computes the normalized expansion rate assuming flat Lambda CDM cosmology. + + Parameters + ---------- + z : float or array_like + Redshift(s) at which to evaluate E(z). + + Returns + ------- + float or ndarray + Dimensionless Hubble parameter E(z) = sqrt(Omega_m*(1+z)^3 + Omega_k*(1+z)^2 + Omega_lambda). + """ a=1.+z #inverse scale factor return (cosmo.Omega_m*a**3+cosmo.Omega_k*a**2+cosmo.Omega_lambda)**0.5 def inv_E(z): - """ inverse of scale factor""" + """Inverse of dimensionless Hubble parameter, 1/E(z). + + Used as integrand for computing comoving distance. + + Parameters + ---------- + z : float + Redshift. + + Returns + ------- + float + 1/E(z), where E(z) is the dimensionless Hubble parameter. + """ return E(z)**-1 def DM(z): - """comoving distance [Mpc]""" + """Comoving distance (line-of-sight) via numerical integration. + + This is the slow but accurate version. For array operations, + use the interpolated version `dm()` after calling `init_dist_measures()`. + + Parameters + ---------- + z : float + Redshift. + + Returns + ------- + float + Comoving distance in Mpc. + """ res,err=integrate.quad(inv_E,0,z) DH=c_light_kms/cosmo.H0 return DH*res def DA(z): - """angular diameter distance [Mpc]""" + """Angular diameter distance via numerical integration. + + This is the slow but accurate version. For array operations, + use the interpolated version `da()` after calling `init_dist_measures()`. + + Parameters + ---------- + z : float + Redshift. + + Returns + ------- + float + Angular diameter distance in Mpc. + """ return DM(z)/(1+z) def DL(z): - """luminosity distance [Mpc]""" + """Luminosity distance via numerical integration. + + This is the slow but accurate version. For array operations, + use the interpolated version `dl()` after calling `init_dist_measures()`. + + Parameters + ---------- + z : float + Redshift. + + Returns + ------- + float + Luminosity distance in Mpc. + """ return DM(z)*(1+z) -#cosmological volume element calculation -# comoving Mpc^3 per dz dOmega -# note this is simply a volume - no rate adjustment -# is taken into account. def dV(z): - """ cosmological volume element [Mpc^3 /redshift /sr]""" + """Comoving volume element per unit redshift per steradian. + + Computes dV/dz/dOmega, the differential comoving volume element. + This does not include any time dilation factor for rate conversions. + + Parameters + ---------- + z : float + Redshift. + + Returns + ------- + float + Volume element in Mpc^3 per unit redshift per steradian. + """ DH=c_light_kms/cosmo.H0 return DH*(1+z)**2*DA(z)**2/E(z) def dVdtau(z): - """ cosmological time-volume element [Mpc^3 /redshift /sr] - it is weighted by an extra (1+z) factor to reflect the rate - in the rest frame vs the observer frame + """Time-weighted comoving volume element per unit redshift per steradian. + + Similar to dV(z) but includes an extra (1+z)^-1 factor to account for + cosmological time dilation when converting between source-frame and + observer-frame event rates. + + Parameters + ---------- + z : float + Redshift. + + Returns + ------- + float + Time-weighted volume element in Mpc^3 per unit redshift per steradian. """ DH=c_light_kms/cosmo.H0 return DH*(1+z)*DA(z)**2/E(z) #changed (1+z)**2 to (1+z) @@ -164,12 +302,26 @@ def dVdtau(z): def init_dist_measures(this_ZMIN=DEF_ZMIN,this_ZMAX=DEF_ZMAX,this_NZ=DEF_NZ, verbose=False): - """ Initialises cosmological distance measures. - - Fills in look-up tables that can operate on input numpy arrays. - For speed. - It will use default values of cosmological parameters if - nothing has yet been specified. + """Initialize interpolation tables for fast distance measure calculations. + + This function populates lookup tables for comoving distance, angular + diameter distance, luminosity distance, and volume elements. After + initialization, the lowercase functions (dm, da, dl, dv, dvdtau) can + be used for fast array operations. + + Must be called before using the interpolated distance functions. + Uses current cosmological parameters set via `set_cosmology()`. + + Parameters + ---------- + this_ZMIN : float, optional + Minimum redshift for interpolation grid. Default is 0. + this_ZMAX : float, optional + Maximum redshift for interpolation grid. Default is 10. + this_NZ : int, optional + Number of redshift points in interpolation grid. Default is 1000. + verbose : bool, optional + If True, print confirmation message. Default is False. """ #sets up initial grid of DMZ @@ -204,8 +356,19 @@ def init_dist_measures(this_ZMIN=DEF_ZMIN,this_ZMAX=DEF_ZMAX,this_NZ=DEF_NZ, def dm(z): - """ Comoving distance [Mpc]. - Must have initialised w. init_dist_measures + """Comoving distance via linear interpolation (fast, array-compatible). + + Requires prior call to `init_dist_measures()` to populate lookup tables. + + Parameters + ---------- + z : float or array_like + Redshift(s). Must be within [ZMIN, ZMAX] set during initialization. + + Returns + ------- + float or ndarray + Comoving distance in Mpc. """ iz=np.array(np.floor(z/DZ)).astype('int') kz=z/DZ-iz @@ -213,8 +376,19 @@ def dm(z): def dl(z): - """ Luminosity distance [Mpc]. - Must have initialised w. init_dist_measures + """Luminosity distance via linear interpolation (fast, array-compatible). + + Requires prior call to `init_dist_measures()` to populate lookup tables. + + Parameters + ---------- + z : float or array_like + Redshift(s). Must be within [ZMIN, ZMAX] set during initialization. + + Returns + ------- + float or ndarray + Luminosity distance in Mpc. """ iz=np.array(np.floor(z/DZ)).astype('int') kz=z/DZ-iz @@ -222,8 +396,19 @@ def dl(z): def da(z): - """ Angular diameter distance [Mpc]. - Must have initialised w. init_dist_measures + """Angular diameter distance via linear interpolation (fast, array-compatible). + + Requires prior call to `init_dist_measures()` to populate lookup tables. + + Parameters + ---------- + z : float or array_like + Redshift(s). Must be within [ZMIN, ZMAX] set during initialization. + + Returns + ------- + float or ndarray + Angular diameter distance in Mpc. """ iz=np.array(np.floor(z/DZ)).astype('int') kz=z/DZ-iz @@ -231,19 +416,41 @@ def da(z): def dv(z): - """ Comoving volume element [Mpc^3 dz sr^-1]. - Must have initialised w. init_dist_measures - """ + """Comoving volume element via linear interpolation (fast, array-compatible). + + Requires prior call to `init_dist_measures()` to populate lookup tables. + + Parameters + ---------- + z : float or array_like + Redshift(s). Must be within [ZMIN, ZMAX] set during initialization. + + Returns + ------- + float or ndarray + Volume element in Mpc^3 per unit redshift per steradian. + """ iz=np.array(np.floor(z/DZ)).astype('int') kz=z/DZ-iz return dvs[iz]*(1.-kz)+dvs[iz+1]*kz def dvdtau(z): - """ Comoving volume element [Mpc^3 dz sr^-1]. - normalsied by proper time (1+z)^-1 to convert - rates. - """ + """Time-weighted volume element via linear interpolation (fast, array-compatible). + + Includes (1+z)^-1 factor for source-frame to observer-frame rate conversion. + Requires prior call to `init_dist_measures()` to populate lookup tables. + + Parameters + ---------- + z : float or array_like + Redshift(s). Must be within [ZMIN, ZMAX] set during initialization. + + Returns + ------- + float or ndarray + Time-weighted volume element in Mpc^3 per unit redshift per steradian. + """ iz=np.array(np.floor(z/DZ)).astype('int') kz=z/DZ-iz return (dvdtaus[iz]*(1.-kz)+dvdtaus[iz+1]*kz) #removed the 1/(1+z) dependency @@ -253,40 +460,66 @@ def dvdtau(z): def E_to_F(E,z,alpha=0, bandwidth=1e9): - """ Converts an energy to a fluence - Formula from Macquart & Ekers 2018 - Energy is assumed to be in ergs - Fluence returned in Jy ms + """Convert isotropic-equivalent energy to observed fluence. + + Uses the formula from Macquart & Ekers (2018) to convert burst energy + at the source to observed fluence at the telescope. + + Parameters + ---------- + E : float or array_like + Isotropic-equivalent burst energy in erg. + z : float or array_like + Redshift of the source. + alpha : float, optional + Spectral index where F_nu ~ nu^(-alpha). Default is 0 (flat spectrum). + bandwidth : float, optional + Observation bandwidth in Hz. Default is 1e9 (1 GHz). + + Returns + ------- + float or ndarray + Fluence in Jy ms. """ F=E/(4*np.pi*(dl(z))**2/(1.+z)**(2.-alpha)) F /= 9.523396e22*bandwidth # see below for constant calculation return F -# inverse of above def F_to_E(F,z,alpha=0, bandwidth=1e9, Fobs=1.3e9, Fref=1.3e9): - """ Converts a fluence in Jy ms to an energy in erg - Formula from Macquart & Ekers 2018 - Works with an array of z. - - Arguments are: - Fluence: of an FRB [Jy ms] - - Redshift: assumed redshift of an FRB producing the fluence F. - Standard cosmological definition [unitless] - - alpha: F(\nu)~\nu^-\alpha. Note that this is an internal definition. - The paper uses ^alpha, not ^-alpha. [unitless] - - Bandwidth: over which to integrate fluence [Hz] - - Fobs: the observation frequency [Hz] - - Fref: reference frequency at which FRB energies E are normalised. - It defaults to 1.3 GHz (ASKAP lat50, Parkes). - - Return value: energy [erg] - + """Convert observed fluence to isotropic-equivalent energy. + + Uses the formula from Macquart & Ekers (2018) to convert observed + fluence to source-frame burst energy. Supports array inputs for z. + + Parameters + ---------- + F : float or array_like + Observed fluence in Jy ms. + z : float or array_like + Redshift of the source. + alpha : float, optional + Spectral index where F_nu ~ nu^(-alpha). Default is 0 (flat spectrum). + Note: This uses the convention F ~ nu^(-alpha), opposite to some papers. + bandwidth : float, optional + Observation bandwidth in Hz. Default is 1e9 (1 GHz). + Fobs : float, optional + Observation frequency in Hz. Default is 1.3e9 (1.3 GHz, ASKAP/Parkes). + Fref : float, optional + Reference frequency for energy normalization in Hz. Default is 1.3e9. + + Returns + ------- + float or ndarray + Isotropic-equivalent burst energy in erg. + + Notes + ----- + Unit conversion factor 9.523396e22 accounts for: + - 10^-26 from Jy to W/m^2/Hz + - 1e-3 from ms to s + - (3.086e22 m/Mpc)^2 for distance conversion + - 1e7 from J to erg """ E=F*4*np.pi*(dl(z))**2/(1.+z)**(2.-alpha) # now convert from dl in MPc and F in Jy ms @@ -296,24 +529,36 @@ def F_to_E(F,z,alpha=0, bandwidth=1e9, Fobs=1.3e9, Fref=1.3e9): # 1e7 from J to erg # total factor is 9.523396e22 E *= 9.523396e22*bandwidth - + # now corrects for reference frequency # according to value of alpha # effectively: if fluence was X at F0, it was X*(F0/Fref)**alpha at Fref # i.e. if alpha is positive (stronger at low frequencies), we reduce E # This acts to reduce the telescope threshold at higher frequencies E *= (Fobs/Fref)**alpha - + return E -# calculates the 'bin size scaling factor' -# essentially differentiates the above expression -# for observed dF to emitted dF -# does this include the rate factor??? NO! def dFnu_to_dEnu(z,alpha=0,bandwidth=1.e9): - """ Converts differential dF to differential dE - Good for probing "per fluence" stats to "per energy" + """Compute the Jacobian dE/dF for fluence-to-energy transformations. + + Useful for converting "per fluence" statistics to "per energy" statistics. + This is the derivative of energy with respect to fluence at fixed z. + + Parameters + ---------- + z : float or array_like + Redshift. + alpha : float, optional + Spectral index where F_nu ~ nu^(-alpha). Default is 0. + bandwidth : float, optional + Observation bandwidth in Hz. Default is 1e9. + + Returns + ------- + float or ndarray + Jacobian dE/dF for the transformation. """ #Fnu is Jy #Jy: 1e-26 J/s /m2 /Hz @@ -324,20 +569,32 @@ def dFnu_to_dEnu(z,alpha=0,bandwidth=1.e9): ######### possible source evolution functions go here ########## def choose_source_evolution_function(which=0): - """ - Selects which source evolution function to use - These are now generalised to take multiple parameters - Could implement arbitrarily many of these - - Arguments: - which (int). Selects which pre-defined model - to use for FRB source evolution. - Currently implemented values are: - 0: star-formation rate from Madau - & Dickenson, to the power n - 1: (1+z)^2.7n, i.e. 0 but without the - denominator - + """Select a source evolution function for FRB population modeling. + + Returns a function that computes the relative FRB source density + as a function of redshift. These functions are parameterized to + allow for different evolution scenarios. + + Parameters + ---------- + which : int, optional + Model selection: + - 0: Star formation rate from Madau & Dickinson (2014) raised to power n. + SFR(z)^n where SFR follows the cosmic star formation history. + - 1: Simple power law (1+z)^(2.7*n), without the high-z turnover. + Useful for comparison with SFR model. + Default is 0. + + Returns + ------- + callable + Source evolution function with signature f(z, n) where z is redshift + and n is the evolution power parameter. + + Raises + ------ + ValueError + If `which` is not 0 or 1. """ if which==0: source_evolution=sfr_evolution @@ -348,30 +605,64 @@ def choose_source_evolution_function(which=0): return source_evolution def sfr_evolution(z,*params): - """ - Madau & dickenson 2014 - Arguments: - z (float): redshift - params: n (float) Scaling parameter. + """Star formation rate evolution model from Madau & Dickinson (2014). + + Computes the cosmic star formation rate raised to a power n, normalized + to unity at z=0. + + Parameters + ---------- + z : float or array_like + Redshift. + *params : float + First parameter is n, the power to which the SFR is raised. + Typically n=1 for direct SFR tracking. + + Returns + ------- + float or ndarray + Relative source density at redshift z, normalized to 1 at z=0. """ return (1.0025738*(1+z)**2.7 / (1 + ((1+z)/2.9)**5.6))**params[0] - + def opz_evolution(z,*params): - """ - Same as SFR, but without denominator, i.e. just (1+z)**2.7 - Factor of 2.7 is kept so that resulting n-values are comparable - Arguments: - z:(float, numpy array) redshift - params: n (float) Scaling parameter. + """Simple power-law source evolution model. + + Computes (1+z)^(2.7*n), which matches the low-z behavior of the SFR + model but lacks the high-z turnover. The factor 2.7 ensures that + n-values are comparable between models. + + Parameters + ---------- + z : float or array_like + Redshift. + *params : float + First parameter is n, the evolution power parameter. + + Returns + ------- + float or ndarray + Relative source density at redshift z. """ return (1+z)**(2.7*params[0]) - -# outdated code -# returns a population density proportional to the star-formation rate to the power n -#Madau & dickenson 2014 -# units: solar masses per year per cubic Mpc def sfr(z): + """Star formation rate density from Madau & Dickinson (2014). + + .. deprecated:: + Use `sfr_evolution(z, 1)` instead for parameterized models. + + Parameters + ---------- + z : float or array_like + Redshift. + + Returns + ------- + float or ndarray + Star formation rate in solar masses per year per cubic Mpc, + normalized to approximately 1 at z=0. + """ return 1.0025738*(1+z)**2.7 / (1 + ((1+z)/2.9)**5.6) diff --git a/zdm/data_class.py b/zdm/data_class.py index efabb409..657f8180 100644 --- a/zdm/data_class.py +++ b/zdm/data_class.py @@ -1,4 +1,18 @@ -""" Items related to dataclasses """ +""" +Base dataclass utilities for parameter management. + +This module provides base classes used by the parameter dataclasses +in the zdm package. These provide common functionality for serialization, +dictionary access, and parameter metadata handling. + +Classes +------- +myDataClass + Base class for parameter dataclasses with metadata access methods. +myData + Base class for composite data objects containing multiple dataclasses. +""" + import numpy as np import json import pandas diff --git a/zdm/energetics.py b/zdm/energetics.py index b6e06ab2..f47b0b8a 100644 --- a/zdm/energetics.py +++ b/zdm/energetics.py @@ -1,31 +1,80 @@ +""" +FRB luminosity/energy function implementations. + +This module provides functions for computing FRB luminosity (energy) distributions, +including both cumulative and differential forms. The main luminosity functions +implemented are: + +1. **Power Law**: Simple power-law distribution dN/dE ~ E^gamma between Emin and Emax +2. **Gamma Function**: Upper incomplete gamma function distribution with exponential cutoff + +The gamma function implementation uses spline interpolation for efficiency, as +direct evaluation of the incomplete gamma function is computationally expensive +when called many times during grid calculations. + +Key Functions +------------- +- `vector_cum_power_law`: Cumulative power-law luminosity function +- `vector_cum_gamma_spline`: Cumulative gamma function with spline interpolation +- `array_cum_gamma_spline`: N-dimensional array wrapper for gamma function +- `init_igamma_splines`: Initialize spline lookup tables for gamma functions + +Module Variables +---------------- +igamma_splines : dict + Cache of spline interpolators keyed by gamma value. +SplineMin, SplineMax : float + Log10 range for spline interpolation. +SplineLog : bool + If True, perform interpolation in log-log space (recommended). + +Example +------- +>>> from zdm import energetics +>>> Eth = np.logspace(38, 42, 100) # Energy thresholds in erg +>>> params = (1e38, 1e42, -1.5) # Emin, Emax, gamma +>>> fraction = energetics.vector_cum_power_law(Eth, *params) +""" + import numpy as np from scipy import interpolate import mpmath from IPython import embed +# Global cache for spline interpolators igamma_splines = {} igamma_linear = {} igamma_linear_log10 = {} -SplineMin = -6 -SplineMax = 6 -NSpline = 1000 -SplineLog = True - -############## this section defines different luminosity functions ########## -def init_igamma_splines(gammas, reinit=False,k=3): - """ - gammas [list of floats]: list of values of gamma at which splines - must be created - reinit [bool]: if True, will re-initialise even if a spline for - that gamma has already been created - k [int]: degree of spline to use. 3 by default (cubic splines). - Formal range: integers 1 <= k <= 5. Do NOT use 2 or 4. - - If SplineLog is set, interpolations are performed in log-space, - i.e. the results is a spline interpolation of the log10 of the - answer in terms of the log10 of the input +# Spline interpolation settings +SplineMin = -6 # Log10 of minimum argument for incomplete gamma +SplineMax = 6 # Log10 of maximum argument +NSpline = 1000 # Number of spline points +SplineLog = True # Use log-space interpolation (more accurate) + +def init_igamma_splines(gammas, reinit=False, k=3): + """Initialize spline interpolators for the upper incomplete gamma function. + + Pre-computes spline representations of the upper incomplete gamma function + Gamma(gamma, x) for fast evaluation during grid calculations. Splines are + cached globally and reused across calls. + + Parameters + ---------- + gammas : list of float + Values of gamma (the shape parameter) for which to create splines. + reinit : bool, optional + If True, reinitialize splines even if they already exist. Default False. + k : int, optional + Degree of spline interpolation. Default is 3 (cubic). Valid range: 1-5. + Note: k=2 and k=4 not recommended due to numerical issues. + + Notes + ----- + If module variable `SplineLog` is True (default), interpolation is performed + in log-log space, which provides better accuracy over the wide dynamic range + of the incomplete gamma function. """ global SplineMin,SplineMax,NSpline,SplineLog for gamma in gammas: @@ -47,14 +96,20 @@ def init_igamma_splines(gammas, reinit=False,k=3): igamma_splines[gamma] = interpolate.splrep(avals, numer,k=k) -def init_igamma_linear(gammas:list, reinit:bool=False, - log:bool=False): - """ Setup the linear interpolator for gamma - - Args: - gammas (list): values of gamma - reinit (bool, optional): If True, redo the calculation. - log (bool, optional): Perform in log10 space +def init_igamma_linear(gammas: list, reinit: bool = False, + log: bool = False): + """Initialize linear interpolators for the upper incomplete gamma function. + + Alternative to spline interpolation using scipy's interp1d. + + Parameters + ---------- + gammas : list of float + Values of gamma for which to create interpolators. + reinit : bool, optional + If True, reinitialize even if interpolator exists. Default False. + log : bool, optional + If True, perform interpolation in log10 space. Default False. """ for gamma in gammas: @@ -125,9 +180,24 @@ def array_cum_power_law(Eth,*params): ########### simple power law functions ############# -def vector_cum_power_law(Eth,*params): - """ Calculates the fraction of bursts above a certain power law - for a given Eth. +def vector_cum_power_law(Eth, *params): + """Cumulative power-law luminosity function. + + Computes the fraction of bursts with energy above threshold Eth for + a power-law distribution dN/dE ~ E^gamma between Emin and Emax. + + Parameters + ---------- + Eth : ndarray + Energy threshold values in erg. + *params : tuple + (Emin, Emax, gamma) - minimum energy, maximum energy, power-law index. + Gamma is typically negative (e.g., -1.5). + + Returns + ------- + ndarray + Fraction of bursts with E > Eth. Returns 1 for Eth < Emin, 0 for Eth > Emax. """ params=np.array(params) Emin=params[0] @@ -195,9 +265,28 @@ def vector_diff_power_law(Eth,*params): ########### gamma functions ############# -def vector_cum_gamma(Eth,*params): - """ Calculates the fraction of bursts above a certain gamma function - for a given Eth. +def vector_cum_gamma(Eth, *params): + """Cumulative gamma-function luminosity function (slow, exact version). + + Computes the fraction of bursts with energy above threshold Eth using + an upper incomplete gamma function distribution. This version evaluates + the gamma function directly using mpmath - accurate but slow. + + Parameters + ---------- + Eth : ndarray + Energy threshold values in erg. + *params : tuple + (Emin, Emax, gamma) - minimum energy, characteristic energy, shape parameter. + + Returns + ------- + ndarray + Fraction of bursts with E > Eth. Returns 1 for Eth < Emin. + + See Also + -------- + vector_cum_gamma_spline : Fast spline-interpolated version (recommended). """ params=np.array(params) Emin=params[0] @@ -217,14 +306,27 @@ def vector_cum_gamma(Eth,*params): result[low]=1. return result -def vector_cum_gamma_spline(Eth:np.ndarray, *params): - """ Calculate cumulative Gamma function using a spline +def vector_cum_gamma_spline(Eth: np.ndarray, *params): + """Cumulative gamma-function luminosity function using spline interpolation. - Args: - Eth (np.ndarray): [description] + Fast version of `vector_cum_gamma` that uses pre-computed spline + interpolators. This is the recommended function for grid calculations. - Returns: - np.ndarray: [description] + Parameters + ---------- + Eth : ndarray + Energy threshold values in erg. + *params : tuple + (Emin, Emax, gamma) - minimum energy, characteristic energy, shape parameter. + + Returns + ------- + ndarray + Fraction of bursts with E > Eth. Returns 1 for Eth < Emin. + + Notes + ----- + Automatically initializes splines for new gamma values if needed. """ global SplineLog @@ -335,8 +437,23 @@ def array_cum_gamma_linear(Eth,*params): result=result.reshape(dims) return result -def vector_diff_gamma(Eth,*params): - """ Calculates the differential fraction of bursts for a gamma function +def vector_diff_gamma(Eth, *params): + """Differential gamma-function luminosity function. + + Computes dN/dE normalized such that the integral from Emin to infinity + equals 1. This is the probability density of burst energies. + + Parameters + ---------- + Eth : ndarray + Energy values in erg at which to evaluate. + *params : tuple + (Emin, Emax, gamma) - minimum energy, characteristic energy, shape parameter. + + Returns + ------- + ndarray + Probability density dN/dE at each energy. Returns 0 for E < Emin. """ Emin=params[0] Emax=params[1] diff --git a/zdm/figures.py b/zdm/figures.py index bb09e6d4..8dc48792 100644 --- a/zdm/figures.py +++ b/zdm/figures.py @@ -1,6 +1,16 @@ """ -This file contains functions designed to produce plots, -and associated helper functions. +Plotting functions for visualizing z-DM grids and FRB data. + +This module provides functions for creating publication-quality plots of +z-DM probability distributions, FRB populations, and analysis results. + +Main Functions +-------------- +- `plot_grid`: Plot 2D z-DM probability grid with optional FRB overlays +- `plot_zdm_basic`: Basic z-DM grid visualization +- `plot_1d_projection`: Project grid onto z or DM axis + +Author: C.W. James """ import numpy as np diff --git a/zdm/grid.py b/zdm/grid.py index fa6eefd2..4b74e36e 100644 --- a/zdm/grid.py +++ b/zdm/grid.py @@ -1,3 +1,31 @@ +""" +Core z-DM grid class for FRB population modeling. + +This module provides the Grid class, which computes 2D probability distributions +of FRB detection rates as a function of redshift (z) and dispersion measure (DM). + +The Grid combines: +- Cosmological volume elements and source evolution +- p(DM|z) from the Macquart relation (cosmic + host) +- Telescope detection efficiency (fluence threshold, beam pattern) +- FRB luminosity/energy function + +Key Features +------------ +- Builds normalized 2D probability grids for expected FRB rates +- Handles beam response and detection efficiency +- Supports multiple luminosity functions (power-law, gamma) +- Efficient updating for MCMC parameter exploration + +Example +------- +>>> from zdm import grid +>>> g = grid.Grid(survey, state, zDMgrid, zvals, dmvals, smear_mask) +>>> expected_rate = g.rates # Expected detection rate per (z, DM) bin + +Author: C.W. James +""" + from IPython.terminal.embed import embed import numpy as np import datetime @@ -10,38 +38,56 @@ import time import warnings + class Grid: - """A class to hold a grid of z-dm plots - - Fundamental assumption: each z point represents FRBs created *at* that redshift - Interpolation is performed under this assumption. - - It also assumes a linear uniform grid. + """2D grid for computing FRB detection rates as a function of z and DM. + + The Grid class is the core computational object in zdm. It builds a 2D + probability distribution representing expected FRB detection rates across + the redshift-DM plane for a given survey and parameter set. + + Assumptions: + - Each z bin represents FRBs originating at that redshift (not integrated) + - Linear uniform spacing in both z and DM + - DM includes cosmic + host contributions convolved together + + Attributes + ---------- + rates : ndarray + 2D array of expected FRB rates per (z, DM) bin. + zvals : ndarray + Redshift bin centers. + dmvals : ndarray + DM bin centers in pc/cm^3. + state : parameters.State + Parameter state used for grid calculation. + survey : survey.Survey + Associated survey object. """ def __init__(self, survey, state, zDMgrid, zvals, dmvals, smear_mask, wdist=None, prev_grid=None): - """ - Class constructor. - - Args: - survey (survey.Survey): - state (parameters.State): - Defines the parameters of the analysis - Note, each grid holds the *same* copy so modifying - it in one place affects them all. - zvals (np.1darray, float): - redshift values of the grid. These are "bin centres", - representing ranges from +- dz/2. - dmvals (np.1darray, float: - DM values of the grid. These are These are "bin centres", - representing ranges from +- dDM/2. - smear_mask (np.1darray, float): - 1D array - wdist (bool): - If True, allow for a distribution of widths - prev_grid (grid.Grid): - Another grid with the same parameters just - corresponding to a different survey + """Initialize the Grid for a survey and parameter state. + + Parameters + ---------- + survey : survey.Survey + Survey object with telescope properties and FRB data. + state : parameters.State + Parameter state defining the model. Note: grids share the same + State object, so modifications affect all grids. + zDMgrid : ndarray + 2D array of p(DM|z) probabilities, shape (nz, ndm). + zvals : ndarray + Redshift bin centers. Bins span [z - dz/2, z + dz/2]. + dmvals : ndarray + DM bin centers in pc/cm^3. Bins span [DM - dDM/2, DM + dDM/2]. + smear_mask : ndarray + 1D convolution kernel for host DM smearing. + wdist : bool, optional + If True, include width distribution effects. + prev_grid : Grid, optional + Another Grid with same z/DM values but different survey. + Allows reusing pre-computed cosmological quantities. """ self.grid = None self.survey = survey diff --git a/zdm/io.py b/zdm/io.py index df99f331..a75f8b60 100644 --- a/zdm/io.py +++ b/zdm/io.py @@ -1,9 +1,24 @@ +""" +Input/output utilities for the zdm package. + +This module provides helper functions for file I/O operations including +JSON file handling and grid data persistence. + +Functions +--------- +- `process_jfile`: Load a JSON file +- `savejson`: Save a Python object to JSON +- `save_grid`: Save grid data to compressed file +- `load_grid`: Load grid data from compressed file +""" + import os import gzip import json import numpy as np -def process_jfile(jfile:str): + +def process_jfile(jfile: str): """Load up a JSON file Args: diff --git a/zdm/iteration.py b/zdm/iteration.py index 39465dda..1333ab77 100644 --- a/zdm/iteration.py +++ b/zdm/iteration.py @@ -1,5 +1,31 @@ -### this file includes various routines to iterate /maximise / minimise -# values on a zdm grid +""" +Likelihood calculation routines for z-DM grids. + +This module provides functions for computing log-likelihoods of FRB survey +data given model predictions from a Grid object. These likelihoods are used +for parameter estimation via MCMC or maximum likelihood methods. + +The main likelihood components include: +- p(DM, z): Probability of observed DM and redshift (if localized) +- p(N): Poisson probability of total number of FRBs detected +- p(SNR): Signal-to-noise ratio distribution +- p(w, tau): Width and scattering time distributions + +Main Functions +-------------- +- `get_log_likelihood`: Compute total log-likelihood for a grid/survey pair +- `calc_likelihoods_1D`: Likelihood for 1D (DM only) FRB data +- `calc_likelihoods_2D`: Likelihood for 2D (DM + z) localized FRBs + +Example +------- +>>> from zdm import iteration +>>> ll = iteration.get_log_likelihood(grid, survey) +>>> print(f"Log-likelihood: {ll}") + +Author: C.W. James +""" + import os import time from IPython.terminal.embed import embed @@ -15,19 +41,35 @@ def get_log_likelihood(grid, s, norm=True, psnr=True, Pn=False, pNreps=True, ptauw=False, pwb=False): - """ - Returns the likelihood for the grid given the survey. + """Compute total log-likelihood for a grid given survey data. - Inputs: - grid = Grid used - s = Survey to compare with the grid - norm = Normalise - psnr = Include psnr in likelihood - Pn = Include Pn in likelihood - ptauw = Include p(tau,width) - - Outputs: - llsum = Total loglikelihood for the grid + This is the main likelihood function used for parameter estimation. + It combines multiple likelihood components depending on what data + is available (DM only, localized z, SNR, etc.). + + Parameters + ---------- + grid : Grid or repeat_Grid + Grid object containing model predictions. + s : Survey + Survey object containing FRB observations. + norm : bool, optional + If True, include normalization terms. Default True. + psnr : bool, optional + If True, include SNR distribution likelihood. Default True. + Pn : bool, optional + If True, include Poisson likelihood for total number. Default False. + pNreps : bool, optional + If True, include number of repeaters likelihood. Default True. + ptauw : bool, optional + If True, include p(tau, width) likelihood. Default False. + pwb : bool, optional + If True, include width/beam likelihood. Default False. + + Returns + ------- + float + Total log-likelihood value. """ if isinstance(grid, zdm_repeat_grid.repeat_Grid): diff --git a/zdm/loading.py b/zdm/loading.py index c641046c..4b92056d 100644 --- a/zdm/loading.py +++ b/zdm/loading.py @@ -1,4 +1,21 @@ -""" Load up the Real data """ +""" +High-level functions for loading surveys and initializing analysis state. + +This module provides convenience functions for setting up zdm analysis, +including loading survey data and creating properly configured State objects. + +Main Functions +-------------- +- `set_state`: Create a State object with best-fit or default parameters +- `load_survey`: Load a single survey from file +- `load_CHIME`, `load_Parkes`, etc.: Survey-specific loaders + +Example +------- +>>> from zdm import loading +>>> state = loading.set_state() +>>> surveys, grids = loading.surveys_and_grids() +""" # It should be possible to remove all the matplotlib calls from this # but in the current implementation it is not removed. @@ -15,7 +32,27 @@ from IPython import embed + def set_state(alpha_method=1, cosmo=Planck18): + """Create a State object with default or best-fit parameters. + + Initializes a State with parameters appropriate for FRB z-DM analysis. + Parameter values depend on the chosen alpha_method. + + Parameters + ---------- + alpha_method : int, optional + Method for handling spectral index alpha. + 0: Spectral index interpretation with k-correction. + 1: Rate interpretation with (1+z)^alpha evolution (default). + cosmo : astropy.cosmology.Cosmology, optional + Astropy cosmology to use. Default is Planck18. + + Returns + ------- + parameters.State + Initialized State object with all parameters set. + """ ############## Initialise parameters ############## state = parameters.State() diff --git a/zdm/misc_functions.py b/zdm/misc_functions.py index c0042e59..c15228f1 100644 --- a/zdm/misc_functions.py +++ b/zdm/misc_functions.py @@ -1,3 +1,18 @@ +""" +Miscellaneous utility functions for the zdm package. + +This module contains various helper functions used throughout the zdm package +including grid initialization, parameter updates, probability calculations, +and other common operations. + +Main Functions +-------------- +- `make_zDMgrid`: Initialize z-DM probability grids +- `get_zdm_grids`: Create grids for multiple surveys +- `update_grid`: Update grid with new parameter values +- `interpolate_grid`: Interpolate grid values +""" + import os import sys import numpy as np diff --git a/zdm/parameters.py b/zdm/parameters.py index 73fcfdae..289e40cd 100644 --- a/zdm/parameters.py +++ b/zdm/parameters.py @@ -1,3 +1,52 @@ +""" +Parameter dataclasses for the zdm package. + +This module defines the central configuration system for FRB z-DM analysis. +All model parameters are organized into dataclasses grouped by category: + +Parameter Classes +----------------- +AnalysisParams + Analysis-level settings (grid generation, FRB cuts). +CosmoParams + Cosmological parameters (H0, Omega_m, Omega_b, etc.). +FRBDemoParams + FRB population demographics (source evolution, rates). +RepeatParams + Repeating FRB parameters (rate distribution). +MWParams + Milky Way DM contributions (ISM, halo). +HostParams + Host galaxy DM distribution parameters. +IGMParams + Intergalactic medium parameters (cosmic DM variance). +WidthParams + Intrinsic FRB width distribution parameters. +ScatParams + Scattering timescale distribution parameters. +EnergeticsParams + FRB energy/luminosity function parameters. + +State Class +----------- +The `State` class aggregates all parameter dataclasses into a single object +that can be passed to grid and survey calculations. It provides methods for +updating individual parameters and importing astropy cosmologies. + +Example +------- +>>> from zdm.parameters import State +>>> state = State() +>>> state.cosmo.H0 # Access Hubble constant +67.66 +>>> state.update_param('gamma', -1.5) # Update energy distribution slope + +Notes +----- +All dataclasses inherit from `data_class.myDataClass` which provides +serialization and parameter metadata handling. +""" + from IPython import embed import numpy as np from dataclasses import dataclass, field @@ -8,9 +57,9 @@ from zdm import data_class -# Analysis parameters @dataclass class AnalysisParams(data_class.myDataClass): + """Analysis-level configuration parameters.""" NewGrids: bool = field(default=True, metadata={"help": "Generate new z, DM grids?"}) sprefix: str = field( default="Std", @@ -34,9 +83,9 @@ class AnalysisParams(data_class.myDataClass): } ) -# Cosmology parameters @dataclass class CosmoParams(data_class.myDataClass): + """Cosmological parameters for Lambda CDM universe.""" H0: float = field( default=Planck18.H0.value, metadata={ @@ -86,9 +135,9 @@ class CosmoParams(data_class.myDataClass): ) -# FRB Demographics -- FRBdemo @dataclass class FRBDemoParams(data_class.myDataClass): + """FRB population demographics parameters.""" source_evolution: int = field( default=0, metadata={ @@ -121,9 +170,9 @@ class FRBDemoParams(data_class.myDataClass): -# FRB Demographics -- repeaters @dataclass class RepeatParams(data_class.myDataClass): + """Parameters for repeating FRB population modeling.""" lRmin: float = field( default=-3., metadata={'help': 'Minimum repeater rate', @@ -155,9 +204,9 @@ class RepeatParams(data_class.myDataClass): 'Notation': '$E_R$', }) -# Galactic parameters @dataclass class MWParams(data_class.myDataClass): + """Milky Way DM contribution parameters (ISM and halo).""" ISM: float = field( default=35., metadata={'help': 'Assumed DM for the Galactic ISM', @@ -190,9 +239,9 @@ class MWParams(data_class.myDataClass): 'Notation': '{\\rm DM}_{\\rm halo}', }) -# Host parameters -- host @dataclass class HostParams(data_class.myDataClass): + """Host galaxy DM contribution distribution parameters.""" lmean: float = field( default=2.16, metadata={ @@ -211,9 +260,9 @@ class HostParams(data_class.myDataClass): ) -# IGM parameters @dataclass class IGMParams(data_class.myDataClass): + """Intergalactic medium parameters for cosmic DM distribution.""" logF: float = field( default=np.log10(0.32), metadata={ @@ -224,9 +273,9 @@ class IGMParams(data_class.myDataClass): ) -# FRB intrinsic width parameters @dataclass class WidthParams(data_class.myDataClass): + """FRB intrinsic width distribution parameters.""" WidthFunction: int = field( default=2, metadata={ @@ -287,9 +336,9 @@ class WidthParams(data_class.myDataClass): ) -# FRB intrinsic scattering parameters @dataclass class ScatParams(data_class.myDataClass): + """Scattering timescale distribution parameters.""" ScatFunction: int = field( default=2, metadata={ @@ -356,9 +405,9 @@ class ScatParams(data_class.myDataClass): ) -# FRB Energetics -- energy @dataclass class EnergeticsParams(data_class.myDataClass): + """FRB energy/luminosity function parameters.""" lEmin: float = field( default=30.0, metadata={ @@ -400,9 +449,43 @@ class EnergeticsParams(data_class.myDataClass): class State(data_class.myData): - """Initialize the full state for the analysis - with the default parameters + """Central configuration object containing all model parameters. + The State class aggregates all parameter dataclasses into a single object + that is passed to Grid, Survey, and other components. It provides dictionary-like + access to parameter groups and methods for updating parameters. + + Attributes + ---------- + analysis : AnalysisParams + Analysis-level settings. + cosmo : CosmoParams + Cosmological parameters. + FRBdemo : FRBDemoParams + FRB population parameters. + rep : RepeatParams + Repeater parameters. + MW : MWParams + Milky Way parameters. + host : HostParams + Host galaxy parameters. + IGM : IGMParams + IGM parameters. + width : WidthParams + Width distribution parameters. + scat : ScatParams + Scattering parameters. + energy : EnergeticsParams + Energy function parameters. + + Example + ------- + >>> state = State() + >>> state.cosmo.H0 + 67.66 + >>> state['cosmo'].H0 # Dictionary-style access + 67.66 + >>> state.update_param('gamma', -1.5) """ def __init__(self): @@ -410,6 +493,7 @@ def __init__(self): self.set_params() def set_dataclasses(self): + """Initialize all parameter dataclass instances with defaults.""" self.scat = ScatParams() self.width = WidthParams() self.MW = MWParams() @@ -421,7 +505,19 @@ def set_dataclasses(self): self.energy = EnergeticsParams() self.rep = RepeatParams() - def update_param(self, param:str, value): + def update_param(self, param: str, value): + """Update a single parameter by name. + + Automatically finds which dataclass contains the parameter and + updates it. Handles special cases like H0 updating Omega_b. + + Parameters + ---------- + param : str + Parameter name (e.g., 'H0', 'gamma', 'lEmax'). + value : any + New value for the parameter. + """ # print(self.params) DC = self.params[param] setattr(self[DC], param, value) @@ -433,11 +529,12 @@ def update_param(self, param:str, value): ) def set_astropy_cosmo(self, cosmo): - """Slurp the values from an astropy Cosmology object - into our format + """Import cosmological parameters from an astropy Cosmology object. - Args: - cosmo (astropy.cosmology): [description] + Parameters + ---------- + cosmo : astropy.cosmology.Cosmology + Astropy cosmology object (e.g., Planck18, WMAP9). """ self.cosmo.H0 = cosmo.H0.value self.cosmo.Omega_lambda = cosmo.Ode0 diff --git a/zdm/pcosmic.py b/zdm/pcosmic.py index 61da886b..7f52dd56 100644 --- a/zdm/pcosmic.py +++ b/zdm/pcosmic.py @@ -1,12 +1,36 @@ -########### P(DM|z) ############ - -# this library implements Eq 33 from -# Macquart et al (Nature, submitted) -# It also includes other functions and so on and so forth related to that work -# Imported by C.W. James - - -############################# +""" +Probability distribution of cosmic dispersion measure given redshift, p(DM|z). + +This module implements the Macquart relation - the probability distribution +of cosmic (extragalactic) dispersion measure as a function of redshift. It is +based on the formalism described in Macquart et al. (2020, Nature) and +implements Equation 4 from that work. + +The cosmic DM arises from free electrons in the intergalactic medium (IGM). +The distribution is characterized by a feedback parameter F that controls +the variance, and normalization constant C0 that ensures = DM_cosmic. + +Key Features +------------ +- p(DM|z) calculation using the modified lognormal distribution +- Mean cosmic DM from the Macquart relation using astropy/frb cosmology +- Lognormal host galaxy DM contribution convolution +- Interpolation utilities for fast grid-based calculations + +Main Functions +-------------- +- `pcosmic`: Core probability distribution function p(delta|z,F) +- `get_mean_DM`: Compute mean cosmic DM at given redshifts +- `get_pDM_grid`: Generate 2D grid of p(DM|z) probabilities +- `get_dm_mask`: Generate host galaxy DM convolution kernel + +References +---------- +Macquart et al. (2020), Nature, 581, 391 - "A census of baryons in the +Universe from localized fast radio bursts" + +Author: C.W. James +""" # from fcntl import F_ADD_SEALS import sys @@ -43,20 +67,38 @@ def p(DM, z, F): def pcosmic(delta, z, logF, C0): - """ Equation 4 page 33 - - delta: = DM_cosmic/ - i.e. fractional cosmic DM - - z: redshift (sigma depends on this) - - logF: log10 of the fluctuation constant, F - - C0: constant to be optimised - - alpha, beta: these are fitted parameters to be optimised - - constraints: std dev must be F*z^0.5 + """Probability density of fractional cosmic DM, p(delta|z). + + Implements Equation 4 from Macquart et al. (2020) for the probability + distribution of cosmic DM fluctuations. The distribution is parameterized + such that delta = DM_cosmic / follows a modified log-normal + with redshift-dependent width. + + Parameters + ---------- + delta : float or array_like + Fractional cosmic DM, defined as DM_cosmic / . + Must be positive. + z : float + Redshift. Used to compute the width parameter sigma = F * z^(-0.5). + logF : float + Log10 of the feedback/fluctuation parameter F. Controls the width + of the distribution. Typical values are logF ~ -0.5 to 0. + C0 : float + Normalization constant that ensures = 1. Should be computed + via `iterate_C0()` for each (z, F) pair. + + Returns + ------- + float or ndarray + Probability density p(delta|z). Not normalized; integrate over delta + for normalization. + + Notes + ----- + Uses module-level constants alpha=3, beta=3 as shape parameters. + The standard deviation scales as F * z^(-0.5), reflecting the central + limit theorem averaging of IGM fluctuations along the line of sight. """ ### logF compensation @@ -71,9 +113,34 @@ def pcosmic(delta, z, logF, C0): def p_delta_DM(z, F, C0, deltas=None, dmin=1e-3, dmax=10, ndelta=10000): - """ - Calculates probability distribution of delta DM as function of feedback - redshift and the constant C0 + """Calculate the probability distribution of fractional DM. + + Wrapper around `pcosmic` that generates a grid of delta values + and computes the probability at each point. + + Parameters + ---------- + z : float + Redshift. + F : float + Log10 of the feedback parameter. + C0 : float + Normalization constant. + deltas : array_like, optional + Pre-defined delta values. If None, generates a linear grid. + dmin : float, optional + Minimum delta value if generating grid. Default is 1e-3. + dmax : float, optional + Maximum delta value if generating grid. Default is 10. + ndelta : int, optional + Number of delta points if generating grid. Default is 10000. + + Returns + ------- + deltas : ndarray + Array of fractional DM values. + pdeltas : ndarray + Probability density at each delta value. """ if not deltas: deltas = np.linspace(dmin, dmax, ndelta) @@ -82,12 +149,27 @@ def p_delta_DM(z, F, C0, deltas=None, dmin=1e-3, dmax=10, ndelta=10000): def iterate_C0(z, F, C0=1, Niter=10): - """ - Iteratively solves for C_0 as a function of z and F - - C0 goes through 10 iterations, where each iteration - uses the prior value of C0 to calculate C0. - + """Iteratively solve for the normalization constant C0. + + The constant C0 ensures that the mean of the p(delta) distribution + equals unity, i.e., = 1. This is required for the distribution + to correctly represent DM_cosmic / . + + Parameters + ---------- + z : float + Redshift. + F : float + Log10 of the feedback parameter. + C0 : float, optional + Initial guess for C0. Default is 1. + Niter : int, optional + Number of iterations. Default is 10. + + Returns + ------- + float + Converged value of C0. """ dmin = 1e-3 dmax = 10 @@ -107,8 +189,19 @@ def iterate_C0(z, F, C0=1, Niter=10): def make_C0_grid(zeds, F): - """ Pre-generates normalisation constants for C0 - Does this from a grid of z and a given F + """Pre-compute C0 normalization constants for a grid of redshifts. + + Parameters + ---------- + zeds : ndarray + Array of redshift values. + F : float + Log10 of the feedback parameter. + + Returns + ------- + ndarray + Array of C0 values, one per redshift. """ C0s = np.zeros([zeds.size]) for i, z in enumerate(zeds): @@ -116,18 +209,32 @@ def make_C0_grid(zeds, F): return C0s -def get_mean_DM(zeds: np.ndarray, state: parameters.State,Plot=False): - """ Gets mean average z to which can be applied deltas - - Args: - zeds (np.ndarray): redshifts (must be linearly spaced) - These zeds are assumed to represent mid-points of - bins, i.e. from 0.5, 1.5, 2.5 etc dz - state (parameters.State): - Plot (bool): create a test plot of DM vs z - - Returns: - np.ndarray: DM_cosmic +def get_mean_DM(zeds: np.ndarray, state: parameters.State, Plot=False): + """Compute the mean cosmic DM at given redshifts (Macquart relation). + + Calculates (z) using the IGM electron density model from + the frb package. Assumes a FlatLambdaCDM cosmology with parameters + from the state object. + + Parameters + ---------- + zeds : ndarray + Redshifts at which to compute mean DM. Must be linearly spaced + and represent bin centers (e.g., 0.5*dz, 1.5*dz, 2.5*dz, ...). + state : parameters.State + State object containing cosmological parameters (H0, Omega_b, Omega_m). + Plot : bool, optional + If True, generate diagnostic plots of DM vs z. Default is False. + + Returns + ------- + ndarray + Mean cosmic DM values in pc/cm^3 at each redshift. + + Notes + ----- + Uses `frb.dm.igm.average_DM` for the underlying calculation. + The redshifts must be evenly spaced for correct bin assignment. """ # Generate the cosmology cosmo = FlatLambdaCDM( @@ -189,15 +296,22 @@ def get_mean_DM(zeds: np.ndarray, state: parameters.State,Plot=False): def get_log_mean_DM(zeds: np.ndarray, state: parameters.State): - """ Gets mean average z to which can be applied deltas - Does NOT assume that the zeds are linearly spaced. - - Args: - zeds (np.ndarray): redshifts (any order/spacing). - state (parameters.State): - - Returns: - np.ndarray: DM_cosmic + """Compute mean cosmic DM for arbitrarily-spaced redshifts. + + Similar to `get_mean_DM` but does not require linearly-spaced redshifts. + Slower because it evaluates each redshift independently. + + Parameters + ---------- + zeds : ndarray + Redshifts at which to compute mean DM. Can have any spacing. + state : parameters.State + State object containing cosmological parameters. + + Returns + ------- + ndarray + Mean cosmic DM values in pc/cm^3 at each redshift. """ # Generate the cosmology cosmo = FlatLambdaCDM( @@ -215,8 +329,29 @@ def get_log_mean_DM(zeds: np.ndarray, state: parameters.State): def get_C0(z, F, zgrid, Fgrid, C0grid): - """ Takes a pre-generated table of C0 values, - and calculates the p(DM) distribution based on this """ + """Interpolate C0 from a pre-computed grid. + + Performs bilinear interpolation on a pre-computed grid of C0 values + to obtain C0 for arbitrary (z, F) pairs. + + Parameters + ---------- + z : float + Redshift at which to interpolate. + F : float + Log10 of feedback parameter. + zgrid : ndarray + Redshift grid points used to build C0grid. + Fgrid : ndarray + Log10(F) grid points used to build C0grid. + C0grid : ndarray + 2D array of pre-computed C0 values, shape (len(zgrid), len(Fgrid)). + + Returns + ------- + float + Interpolated C0 value. + """ F = 10 ** F Fgrid = 10 ** Fgrid @@ -248,10 +383,29 @@ def get_C0(z, F, zgrid, Fgrid, C0grid): def get_pDM(z, F, DMgrid, zgrid, Fgrid, C0grid, zlog=False): - """ Gets pDM for an arbitrary z value - - zlog (bool): True if zs are log-spaced - False if linearly spaced + """Compute p(DM|z) for a single redshift using pre-computed tables. + + Parameters + ---------- + z : float + Redshift. + F : float + Log10 of feedback parameter. + DMgrid : ndarray + DM values at which to evaluate the distribution. + zgrid : ndarray + Redshift grid for C0 interpolation. + Fgrid : ndarray + Feedback parameter grid for C0 interpolation. + C0grid : ndarray + Pre-computed C0 values. + zlog : bool, optional + If True, use log-spaced z grid. If False (default), use linear spacing. + + Returns + ------- + ndarray + Probability density p(DM|z) at each DM grid point. """ C0 = get_C0(z, F, zgrid, Fgrid, C0grid) if zlog: @@ -266,17 +420,31 @@ def get_pDM(z, F, DMgrid, zgrid, Fgrid, C0grid, zlog=False): def get_pDM_grid( state: parameters.State, DMgrid, zgrid, C0s, verbose=False, zlog=False ): - """ Gets pDM when the zvals are the same as the zgrid - state - C0grid: C0 values obtained by convergence - DMgrid: range of DMs for which we are generating a histogram - This represent bin centres - zgrid: redshifts. These do not have to be in any particular - order or spacing. We just iterature through these - zlog (bool): True if zs are log-spaced - False if linearly spaced - - + """Generate a 2D grid of p(DM|z) probabilities. + + Computes the probability distribution of cosmic DM for each redshift + in the grid. This is the main function for building z-DM grids. + + Parameters + ---------- + state : parameters.State + State object containing cosmological and IGM parameters. + DMgrid : ndarray + DM values (bin centers) for the output grid, in pc/cm^3. + zgrid : ndarray + Redshift values for the output grid. + C0s : ndarray + Pre-computed C0 normalization constants for each redshift. + verbose : bool, optional + If True, print diagnostic information. Default is False. + zlog : bool, optional + If True, zgrid is log-spaced. If False (default), linearly spaced. + + Returns + ------- + ndarray + 2D array of shape (len(zgrid), len(DMgrid)) containing normalized + p(DM|z) values. Each row sums to 1. """ # added H0 dependency if zlog: @@ -296,15 +464,27 @@ def get_pDM_grid( return pDMgrid -#### defines DMx (excess contribution) #### -# family of lognormal curves. These tae either linear or logarithmic arguments, -# and either do or do not include the 1/x fatcor when converting log to dlin. -# the DM part is because integral p(logDM)dlogDM = 1 -# DM is normal, logmean and logsigma are natural logs of these parameters +#### Host galaxy DM contribution (lognormal distribution) #### + def linlognormal_dlin(DM, *args): - """ x values are in linear space, - args in logspace, - returns p dx """ + """Lognormal probability density in linear DM space. + + Computes p(DM) where DM follows a lognormal distribution with + parameters given in natural log space. Includes the 1/DM Jacobian + for conversion from log to linear probability density. + + Parameters + ---------- + DM : float or array_like + Dispersion measure values (linear scale). + *args : tuple + (logmean, logsigma) - mean and standard deviation in natural log space. + + Returns + ------- + float or ndarray + Probability density p(DM) such that integral p(DM) dDM = 1. + """ logmean = args[0] logsigma = args[1] logDM = np.log(DM) @@ -313,10 +493,22 @@ def linlognormal_dlin(DM, *args): def loglognormal_dlog(logDM, *args): - """x values, mean and sigma are already in logspace - returns p dlogx - That is, this function is simply a Gaussian, - and the arguments happen to be in log space. + """Gaussian probability density in log DM space. + + This is simply a Gaussian distribution where the variable happens to + be log(DM). Does not include Jacobian for log-to-linear conversion. + + Parameters + ---------- + logDM : float or array_like + Natural log of dispersion measure. + *args : tuple + (logmean, logsigma, norm) - Gaussian parameters in log space. + + Returns + ------- + float or ndarray + Probability density p(log DM) such that integral p(log DM) d(log DM) = 1. """ logmean = args[0] logsigma = args[1] @@ -339,31 +531,38 @@ def plot_mean(zvals, saveas, title="Mean DM"): def get_dm_mask(dmvals, params, zvals=None, plot=False): - """ - Generates a mask over which to integrate the lognormal - distribution of FRM host galaxy DM contributions. It's - essentially just a probability distribution of p(DMhost), - such that p(DMeg = DMhost + DMcosmic) can be quickly - calculated, as DM[i] = DM[set[i]]*mask[i] - - DMvals (np.ndarray): DMs over which to calculate the mask. - These represent local probabilities of p(DM), i.e. - the probability of getting a DM between - DMval - dDM/2. and DMval + dDM/2. - - params [vector, 2]: mean and sigma of the lognormal (log10) - host galaxy DM distribution - - zvals [np.ndarray]: redshift values at which to calculate this. - If None: return a single, redshift-independent vector. - If not None: return a mask for each value of z, with - DMhost reduced by the (1+z) value. - In future: add a parameter to scale this as (1+z)^xi. - - We simply assign lognormal values at the midpoints - The renormalisation constants then give some idea of the - error in this procedure - + """Generate a convolution kernel for host galaxy DM contribution. + + Creates a probability distribution p(DM_host) following a lognormal + distribution. This kernel can be convolved with p(DM_cosmic|z) to + obtain the full p(DM_extragalactic|z) = p(DM_cosmic + DM_host|z). + + The host DM contribution is redshift-corrected: observed DM_host is + reduced by (1+z) relative to the rest-frame value. + + Parameters + ---------- + dmvals : ndarray + DM grid values (bin centers) in pc/cm^3. + params : array_like + [log10_mean, log10_sigma] - lognormal parameters for host DM + distribution in log10 space. + zvals : ndarray, optional + If provided, compute a redshift-dependent mask where host DM is + scaled by 1/(1+z). If None, return a single z-independent mask. + plot : bool, optional + If True, generate diagnostic plots. Default is False. + + Returns + ------- + ndarray + If zvals is None: 1D array of shape (len(dmvals),) containing p(DM_host). + If zvals is provided: 2D array of shape (len(zvals), len(dmvals)). + Each row/vector is normalized to sum to 1. + + Notes + ----- + The mask is designed for discrete convolution: p(DM_EG)[i] = sum_j p(DM_cosmic)[i-j] * mask[j] """ if len(params) != 2: @@ -440,31 +639,34 @@ def get_dm_mask(dmvals, params, zvals=None, plot=False): def integrate_pdm(ddm, ndm, logmean, logsigma, quick=True, plot=False): - """ - Assigns probabilities of DM smearing (e.g. due to the host galaxy contribution) - to a histogram in dm space. - - Here, the resulting mask assumes DM values of 0 (0 to 0.5), 1( 0.5 to 1.5) etc. - - Two methods: quick (use central values of DM bins), and slow (integrate bins) - - Arguments: - - ddm (float) [pc/cm3]: spacing of dm bins - - ndm (int): number of dm bins. Bins assumed to start at 0 - - logmean: natural logarithm of the mean of the DM distribution - - logsigma: sigma in natural log space of DM distribution - - quick (bool): True uses the speedup, False takes things slowly - - plot (bool): If True, compares quick and slow methods, then exits - to avoid generating infinite plots. - - Returns: - mask (np.ndarray) + """Discretize a lognormal DM distribution onto histogram bins. + + Converts a continuous lognormal probability distribution into discrete + bin probabilities for use in convolution operations. Two methods are + available: fast (evaluate at bin centers) and slow (integrate each bin). + + Parameters + ---------- + ddm : float + Bin spacing in pc/cm^3. + ndm : int + Number of DM bins. Bins span [0, ndm*ddm] with first bin [0, 0.5*ddm]. + logmean : float + Mean of the lognormal distribution in natural log space. + logsigma : float + Standard deviation in natural log space. + quick : bool, optional + If True (default), use fast method evaluating at bin centers. + If False, integrate over each bin (slower but more accurate). + plot : bool, optional + If True, generate comparison plot of quick vs slow methods and exit. + Default is False. + + Returns + ------- + ndarray + Array of shape (ndm,) containing probability mass in each bin. + Not normalized; caller should normalize if needed. """ # do this for the z=0 case (handling of z>0 can be performed at # when calling the routine by multiplying dDM values) diff --git a/zdm/repeat_grid.py b/zdm/repeat_grid.py index 06d66c47..ad4c87df 100644 --- a/zdm/repeat_grid.py +++ b/zdm/repeat_grid.py @@ -1,41 +1,42 @@ - """ -Class definition for repeating FRBs. +Grid extension for repeating FRB population modeling. -Input: - A grid, as well as a time-per-field. - We have two calculation methods: exact and MC. - MCs take a very long time to converge on the average (guess:10^4 iterations?) - Exact only calculates and - Repetition parameters Rmin, Rmax, and Rgamma are stored in grid.state - - Rmin: Minimum repetition rate of repeaters - - Rmax: Maximum repetition rate of repeaters - - Rgamma: *Differential* number of repeaters: dN/dR ~ R^Rgamma +This module provides the repeat_Grid class, which extends the base Grid class +to handle repeating FRB populations. It models the expected number of single +detections and repeated detections given a distribution of repeater rates. -Some notes regarding time dilation: - - #1: dVdtau - grid.py includes the time-dilation effect, dVdtau, in grid.dV. - We need to undo this: the number of repeaters does not change - with dtau, just the rate per repeater. Hence, all instances - of dV need to have an extra multiple of (1+z) attached, to undo - the time-dilation effect. - - #2: Rmult. When alpha method == 1, we use "rate scaling". This means - that the rate *per repeater* must change with frequency, - since otherwise the number of progenitors is frequency- - dependent, which is nonsense. This should be handled - under "Rmult", and implies two corrections: - - correct from nominal frequency to central frequency - - correct for (1+z) factor - - #3: sfr factor - grid.sfr, if alpha_method=0, includes the rate scaling with alpha - However, here we use this to calculate the number of progenitors, - thus we must calculate sfr from first principles if alpha_method=1. - This is now handled by assigning self.use_sfr to be self.sfr - when alpha_method=0, or it is recalculated if alpha_method=1 - +Key Concepts +------------ +- Repeater rate distribution: dN/dR ~ R^Rgamma for R in [Rmin, Rmax] +- Single bursts: First detections from repeaters (and possibly non-repeaters) +- Repeat bursts: Subsequent detections from the same source + +Time Dilation Effects +--------------------- +1. **dVdtau**: The base grid uses time-dilated volume elements. For repeater + counts, we need the undilated volume (number of sources doesn't change). + +2. **Rmult**: For alpha_method=1 (rate scaling), the per-source rate changes + with frequency, requiring corrections for central frequency and redshift. + +3. **SFR factor**: Source evolution must be recalculated for alpha_method=1 + since we're counting progenitors, not burst rates. + +Parameters +---------- +Repetition parameters are stored in state.rep: +- Rmin: Minimum repeater rate (bursts/day) +- Rmax: Maximum repeater rate (bursts/day) +- Rgamma: Power-law index of rate distribution + +Example +------- +>>> from zdm import repeat_grid +>>> rg = repeat_grid.repeat_Grid(survey, state, ...) +>>> n_singles = rg.exact_singles # Expected single detections +>>> n_repeats = rg.exact_reps # Expected repeat detections + +Author: C.W. James """ from zdm import grid diff --git a/zdm/survey.py b/zdm/survey.py index 6a4d6f99..0379f753 100644 --- a/zdm/survey.py +++ b/zdm/survey.py @@ -1,8 +1,37 @@ -# ############################################### -# This file defines a class to hold an FRB survey -# Essentially, this is relevant when multiple -# FRBs are discovered by the same instrument -# ############################################## +""" +FRB Survey class for modeling telescope observations. + +This module provides the Survey class, which encapsulates all properties of an +FRB survey including instrument characteristics, beam patterns, detection +efficiencies, and detected FRB data. + +The Survey class handles: +- Loading survey definition files (ECSV format) with FRB metadata +- Beam pattern modeling and solid angle calculations +- Detection efficiency as a function of DM and width +- DM budget calculations (Galactic, halo, host contributions) +- Threshold calculations for FRB detection + +Survey Definition Files +----------------------- +Survey files are stored in `zdm/data/Surveys/` and contain: +- Header metadata: telescope parameters, beam settings, frequency, etc. +- FRB table: TNS name, DM, position, width, fluence for each detection + +Example +------- +>>> from zdm import survey +>>> from zdm.parameters import State +>>> state = State() +>>> dmvals = np.linspace(0, 3000, 1000) +>>> s = survey.Survey(state, 'CRAFT/ICS', 'CRAFT_ICS.ecsv', dmvals) +>>> s.NFRB # Number of FRBs +10 +>>> s.DMEGs # Extragalactic DM values +array([...]) + +Author: C.W. James +""" import numpy as np import os @@ -33,30 +62,61 @@ MIN_THRESH = 1e-10 class Survey: - def __init__(self, state, survey_name:str, - filename:str, - dmvals:np.ndarray, - zvals:np.ndarray=None, - NFRB:int=None, - iFRB:int=0, + """Represents an FRB survey with instrument properties and detected FRBs. + + The Survey class is the primary interface for defining telescope + characteristics and loading FRB detections. It computes detection + efficiencies across the DM-width-fluence parameter space. + + Attributes + ---------- + name : str + Survey identifier (e.g., 'CRAFT/ICS', 'Parkes/Mb'). + NFRB : int + Number of FRBs in the survey. + DMEGs : ndarray + Extragalactic DM values for detected FRBs. + Zs : ndarray or None + Redshifts for localized FRBs (None if unknown). + meta : dict + Survey metadata from file header. + efficiencies : ndarray + Detection efficiency grid as function of DM and width. + """ + + def __init__(self, state, survey_name: str, + filename: str, + dmvals: np.ndarray, + zvals: np.ndarray = None, + NFRB: int = None, + iFRB: int = 0, edir=None, rand_DMG=False, survey_dict=None): - """ Init an FRB Survey class + """Initialize an FRB Survey. - Args: - state (_type_): _description_ - survey_name (str): - Name of the survey - filename (str): _description_ - dmvals (np.ndarray): Extragalactic DM values at which to calculate efficiencies - zvals (np.ndarray): Redshift values at which to calculate efficiencies - NFRB (int, optional): _description_. Defaults to None. - iFRB (int, optional): _description_. Defaults to 0. - edir (str, optional): Location of efficiency files - rand_DMG (bool): If true randomise the DMG values within uncertainty - survey_dict (dict, optional): Dict of survey data meta-parameters to over-ride - values in the survey file + Parameters + ---------- + state : parameters.State + State object containing model parameters. + survey_name : str + Identifier for the survey (e.g., 'CRAFT/ICS'). + filename : str + Path to survey definition file (ECSV format). + dmvals : ndarray + Extragalactic DM grid values for efficiency calculations. + zvals : ndarray, optional + Redshift grid values. Required for some width methods. + NFRB : int, optional + Limit number of FRBs loaded from file. + iFRB : int, optional + Starting index for FRB selection. Default is 0. + edir : str, optional + Directory containing pre-computed efficiency files. + rand_DMG : bool, optional + If True, randomize Galactic DM within uncertainty. Default False. + survey_dict : dict, optional + Override survey metadata parameters. """ # Proceed self.state = state