From ce145476913f8edd060e37be190cd806f7882d77 Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Wed, 15 May 2024 10:11:18 -0700 Subject: [PATCH 1/4] DEV: add documentation pages --- .github/workflows/pages.yml | 66 ++++ .gitignore | 4 +- README.rst | 51 ++- doc/Makefile | 20 + doc/index.rst | 8 + doc/make.bat | 35 ++ doc/pages_requirements.txt | 9 + .../_templates/custom-class-template.rst | 32 ++ .../_templates/custom-module-template.rst | 68 ++++ doc/source/conf.py | 37 ++ doc/source/index.rst | 25 ++ pyproject.toml | 4 +- wcosmo/__init__.py | 2 +- wcosmo/models.py | 191 --------- wcosmo/utils.py | 71 ++++ wcosmo/wcosmo.py | 374 +++++++++--------- 16 files changed, 596 insertions(+), 401 deletions(-) create mode 100644 .github/workflows/pages.yml create mode 100644 doc/Makefile create mode 100644 doc/index.rst create mode 100644 doc/make.bat create mode 100644 doc/pages_requirements.txt create mode 100644 doc/source/_templates/custom-class-template.rst create mode 100644 doc/source/_templates/custom-module-template.rst create mode 100644 doc/source/conf.py create mode 100644 doc/source/index.rst delete mode 100644 wcosmo/models.py diff --git a/.github/workflows/pages.yml b/.github/workflows/pages.yml new file mode 100644 index 0000000..97a9d5f --- /dev/null +++ b/.github/workflows/pages.yml @@ -0,0 +1,66 @@ +name: GitHub Pages + +permissions: + contents: read + pages: write + id-token: write + +on: + push: + branches: + - main + pull_request: + branches: + - main + schedule: + - cron: "0 0 * * 0" + +jobs: + build: + runs-on: ubuntu-20.04 + concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - uses: s-weigand/setup-conda@v1 + with: + update-conda: true + python-version: "3.11" + conda-channels: anaconda, conda-forge + + - name: Install dependencies + run: | + python -m pip install -r requirements.txt + python -m pip install -r doc/pages_requirements.txt + python -m pip install . + + - name: Build documentation + run: | + # cp examples/*.ipynb docs + cd doc + make clean + make html + - name: Upload artifact + uses: actions/upload-pages-artifact@v3 + with: + path: 'doc/build/html' + + deploy: + environment: + name: github-pages + url: ${{ steps.deployment.outputs.page_url }} + runs-on: ubuntu-latest + if: github.ref == 'refs/heads/main' + needs: build + steps: + - name: Deploy to GitHub Pages + id: deployment + uses: actions/deploy-pages@v4 diff --git a/.gitignore b/.gitignore index 067ef89..5d1d6fa 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ *.pyc *.egg-info -build/ \ No newline at end of file +build/ +doc/build/ +doc/source/api/ diff --git a/README.rst b/README.rst index 275e749..981a57a 100644 --- a/README.rst +++ b/README.rst @@ -4,6 +4,34 @@ Backend agnostic astropy-like cosmology The primary intention for this package is for use with :code:`GWPopulation` but the main functionality can be used externally. +Installation and contribution +----------------------------- + +Currently installation is only available from source + +.. code-block:: console + + $ pip install git+https://github.com/ColmTalbot/wcosmo.git + +for development you should follow a standard fork-and-pull workflow. + +- First create a new fork at :code:`github.com/UserName/wcosmo`. +- Clone your fork + + .. code-block:: console + + $ git clone git@github.com:UserName/wcosmo.git + + or use a GitHub codespace. +- Install the local version with + + .. code-block:: console + + $ python -m pip install . + +- Make any desired edits and push to your fork. +- Open a pull request into :code:`git@github.com:ColmTalbot/wcosmo.git`. + Basic usage ----------- @@ -14,21 +42,9 @@ To import an astropy-like cosmology (without units) from wcosmo import FlatwCDM cosmology = FlatwCDM(H0=70, Om0=0.3, w0=-1) -To use this in :code:`GWPopulation` - -.. code-block:: python - - import gwpopulation as gwpop - from wcosmo.models import CosmoModel, PowerLawRedshift - - model = CosmoModel( - model_functions=[ - gwpop.models.mass.two_component_primary_mass_ratio, - gwpop.models.spin.iid_spin, - PowerLawRedshift(cosmo_model="FlatwCDM"), - ], - cosmo_model="FlatwCDM", - ) +This code is automatically used in :code:`GWPopulation` when using either +:code:`gwpopulation.experimental.cosmo_models.CosmoModel` and/or +:code:`PowerLawRedshift` Changing backend ---------------- @@ -47,6 +63,7 @@ Manual backend setting can be done as follows: import jax.numpy as jnp from jax.scipy.linalg.toeplitz import toeplitz - from wcosmo import wcosmo + from wcosmo import wcosmo, utils wcosmo.xp = jnp - wcosmo.toeplitz = toeplitz + utils.xp = jnp + utils.toeplitz = toeplitz diff --git a/doc/Makefile b/doc/Makefile new file mode 100644 index 0000000..d0c3cbf --- /dev/null +++ b/doc/Makefile @@ -0,0 +1,20 @@ +# Minimal makefile for Sphinx documentation +# + +# You can set these variables from the command line, and also +# from the environment for the first two. +SPHINXOPTS ?= +SPHINXBUILD ?= sphinx-build +SOURCEDIR = source +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/doc/index.rst b/doc/index.rst new file mode 100644 index 0000000..62d4936 --- /dev/null +++ b/doc/index.rst @@ -0,0 +1,8 @@ +.. currentmodule:: wcosmo + +.. include:: ../README.rst + +.. autosummary:: + + wcosmo.wcosmo + wcosmo.utils diff --git a/doc/make.bat b/doc/make.bat new file mode 100644 index 0000000..747ffb7 --- /dev/null +++ b/doc/make.bat @@ -0,0 +1,35 @@ +@ECHO OFF + +pushd %~dp0 + +REM Command file for Sphinx documentation + +if "%SPHINXBUILD%" == "" ( + set SPHINXBUILD=sphinx-build +) +set SOURCEDIR=source +set BUILDDIR=build + +%SPHINXBUILD% >NUL 2>NUL +if errorlevel 9009 ( + echo. + echo.The 'sphinx-build' command was not found. Make sure you have Sphinx + echo.installed, then set the SPHINXBUILD environment variable to point + echo.to the full path of the 'sphinx-build' executable. Alternatively you + echo.may add the Sphinx directory to PATH. + echo. + echo.If you don't have Sphinx installed, grab it from + echo.https://www.sphinx-doc.org/ + exit /b 1 +) + +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/doc/pages_requirements.txt b/doc/pages_requirements.txt new file mode 100644 index 0000000..703f1a2 --- /dev/null +++ b/doc/pages_requirements.txt @@ -0,0 +1,9 @@ +ipython_genutils +jinja2 +nbsphinx +numpydoc +pandoc +pygments +sphinx +sphinx_rtd_theme +sphinx-nefertiti diff --git a/doc/source/_templates/custom-class-template.rst b/doc/source/_templates/custom-class-template.rst new file mode 100644 index 0000000..16ebb2f --- /dev/null +++ b/doc/source/_templates/custom-class-template.rst @@ -0,0 +1,32 @@ +{{ fullname | escape | underline}} + +.. currentmodule:: {{ module }} + +.. autoclass:: {{ objname }} + :members: + :show-inheritance: + :inherited-members: + + {% block methods %} + .. automethod:: __init__ + + {% if methods %} + .. rubric:: {{ _('Methods') }} + + .. autosummary:: + {% for item in methods %} + ~{{ name }}.{{ item }} + {%- endfor %} + {% endif %} + {% endblock %} + + {% block attributes %} + {% if attributes %} + .. rubric:: {{ _('Attributes') }} + + .. autosummary:: + {% for item in attributes %} + ~{{ name }}.{{ item }} + {%- endfor %} + {% endif %} + {% endblock %} \ No newline at end of file diff --git a/doc/source/_templates/custom-module-template.rst b/doc/source/_templates/custom-module-template.rst new file mode 100644 index 0000000..0fbc6da --- /dev/null +++ b/doc/source/_templates/custom-module-template.rst @@ -0,0 +1,68 @@ +{{ fullname | escape | underline}} + +.. currentmodule:: {{ fullname }} + +.. automodule:: {{ fullname }} + + {% block attributes %} + {% if attributes %} + .. rubric:: Module Attributes + + .. autosummary:: + :toctree: + {% for item in attributes %} + {{ item }} + {%- endfor %} + {% endif %} + {% endblock %} + + {% block functions %} + {% if functions %} + .. rubric:: {{ _('Functions') }} + + .. autosummary:: + :toctree: + {% for item in functions %} + {{ item }} + {%- endfor %} + {% endif %} + {% endblock %} + + {% block classes %} + {% if classes %} + .. rubric:: {{ _('Classes') }} + + .. autosummary:: + :toctree: + :template: custom-class-template.rst + {% for item in classes %} + {{ item }} + {%- endfor %} + {% endif %} + {% endblock %} + + {% block exceptions %} + {% if exceptions %} + .. rubric:: {{ _('Exceptions') }} + + .. autosummary:: + :toctree: + {% for item in exceptions %} + {{ item }} + {%- endfor %} + {% endif %} + {% endblock %} + +{% block modules %} +{% if modules %} +.. rubric:: Modules + +.. autosummary:: + :toctree: + :template: custom-module-template.rst + :recursive: +{% for item in modules %} + {{ item }} +{%- endfor %} +{% endif %} +{% endblock %} \ No newline at end of file diff --git a/doc/source/conf.py b/doc/source/conf.py new file mode 100644 index 0000000..823058f --- /dev/null +++ b/doc/source/conf.py @@ -0,0 +1,37 @@ +# 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 + +# -- Project information ----------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information + +project = 'wcosmo' +copyright = '2024, Colm Talbot, Amanda Farah' +author = 'Colm Talbot, Amanda Farah' +release = '0.0.0' + +# -- General configuration --------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration + +extensions = [ + 'sphinx.ext.autodoc', + 'sphinx.ext.mathjax', + 'sphinx.ext.githubpages', + 'sphinx.ext.autosummary', + 'numpydoc', + 'nbsphinx', + 'sphinx.ext.viewcode', +] + +templates_path = ['_templates'] +exclude_patterns = [] + +root_doc = 'index' + + +# -- Options for HTML output ------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output + +html_theme = 'sphinx_nefertiti' +html_static_path = ['_static'] diff --git a/doc/source/index.rst b/doc/source/index.rst new file mode 100644 index 0000000..cd2c2c5 --- /dev/null +++ b/doc/source/index.rst @@ -0,0 +1,25 @@ +.. currentmodule:: wcosmo + +.. include:: ../../README.rst + +.. toctree:: + :maxdepth: 2 + :caption: Contents: + + +.. autosummary:: + :toctree: api + :template: custom-module-template.rst + :caption: API + :recursive: + + wcosmo + utils + + +Indices and tables +================== + +* :ref:`genindex` +* :ref:`modindex` +* :ref:`search` diff --git a/pyproject.toml b/pyproject.toml index 876d774..d7470b6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ packages = ["wcosmo"] [project.entry-points."gwpopulation.xp"] wcosmo = "wcosmo.wcosmo" -wcosmo-models = "wcosmo.models" +wcosmo-utils = "wcosmo.utils" [project.entry-points."gwpopulation.other"] -wcosmo = "wcosmo.wcosmo:scipy.linalg.toeplitz" +wcosmo-utils = "wcosmo.utils:scipy.linalg.toeplitz" diff --git a/wcosmo/__init__.py b/wcosmo/__init__.py index 99b5bcd..710533c 100644 --- a/wcosmo/__init__.py +++ b/wcosmo/__init__.py @@ -1 +1 @@ -from .wcosmo import * \ No newline at end of file +from .wcosmo import * diff --git a/wcosmo/models.py b/wcosmo/models.py deleted file mode 100644 index e6dc71a..0000000 --- a/wcosmo/models.py +++ /dev/null @@ -1,191 +0,0 @@ -import numpy as xp -from gwpopulation.experimental.jax import NonCachingModel -from gwpopulation.models.redshift import _Redshift - -from .wcosmo import FlatwCDM, available, z_at_value - -__all__ = ["CosmoModel", "PowerLawRedshift"] - - -class CosmologyMixin: - - def cosmology_variables(self, parameters): - return {key: parameters[key] for key in self.cosmology_names} - - def cosmology(self, parameters): - if isinstance(self._cosmo, FlatwCDM): - return self._cosmo - else: - return self._cosmo(**self.cosmology_variables(parameters)) - - def __init__(self, cosmo_model="Planck15"): - - self.cosmo_model = cosmo_model - if self.cosmo_model == "FlatwCDM": - self.cosmology_names = ["H0", "Om0", "w0"] - elif self.cosmo_model == "FlatLambdaCDM": - self.cosmology_names = ["H0", "Om0"] - else: - self.cosmology_names = [] - self._cosmo = available[cosmo_model] - - def detector_frame_to_source_frame(self, data): - - cosmo = self.cosmology(self.parameters) - - samples = dict() - samples["redshift"] = z_at_value( - cosmo.luminosity_distance, - data["luminosity_distance"], - ) - jacobian = cosmo.dDLdz(samples["redshift"]) - for key in data: - if key.endswith("_detector"): - samples[key.strip("_detector")] = data[key] / (1 + samples["redshift"]) - jacobian *= 1 + samples["redshift"] - elif key != "luminosity_distance": - samples[key] = data[key] - return samples, jacobian - - def luminosity_distance_to_redshift_jacobian(self, redshift): - - """ - Calculates the luminosity distance to redshift jacobian - - Parameters - ========== - redshift: array-like - - Returns - ======= - luminosity_distance: array-like - """ - return self.cosmology(self.parameters).dDLdz(redshift) - - -class CosmoModel(NonCachingModel, CosmologyMixin): - """ - Modified version of bilby.hyper.model.Model that disables caching for jax. - """ - - def __init__(self, model_functions, cosmo_model="Planck15"): - super().__init__(model_functions=model_functions) - CosmologyMixin.__init__(self, cosmo_model=cosmo_model) - - def prob(self, data, **kwargs): - """ - Compute the total population probability for the provided data given - the keyword arguments. - - Parameters - ========== - data: dict - Dictionary containing the points at which to evaluate the - population model. - kwargs: dict - The population parameters. These cannot include any of - :code:`["dataset", "data", "self", "cls"]` unless the - :code:`variable_names` attribute is available for the relevant - model. - """ - - samples_in_source, jacobian = self.detector_frame_to_source_frame(data) - probability = super().prob(samples_in_source, **kwargs) - return probability / jacobian - - -class _Redshift(CosmologyMixin, _Redshift): - """ - Base class for models which include a term like dVc/dz / (1 + z) - """ - - _variable_names = None - - @property - def variable_names(self): - vars = self.cosmology_names.copy() - if self._variable_names is not None: - vars += self._variable_names - return vars - - def __init__(self, cosmo_model="Planck15", z_max=2.3): - - super().__init__(cosmo_model=cosmo_model) - self.z_max = z_max - self.zs = xp.linspace(1e-6, z_max, 2500) - - def __call__(self, dataset, **kwargs): - return self.probability(dataset=dataset, **kwargs) - - def normalisation(self, parameters): - normalisation_data = self.differential_spacetime_volume( - dict(redshift=self.zs), bounds=True, **parameters - ) - norm = xp.trapz(normalisation_data, self.zs) - return norm - - def probability(self, dataset, **parameters): - differential_volume = self.differential_spacetime_volume( - dataset, bounds=True, **parameters - ) - norm = self.normalisation(parameters) - - return differential_volume / norm - - def psi_of_z(self, redshift, **parameters): - raise NotImplementedError - - def dvc_dz(self, redshift, **parameters): - return ( - 4 - * xp.pi - * self.cosmology(parameters).differential_comoving_volume(redshift) - ) - - def differential_spacetime_volume(self, dataset, bounds=False, **parameters): - r""" - Compute the differential spacetime volume. - - .. math:: - d\mathcal{V} = \frac{1}{1+z} \frac{dVc}{dz} \psi(z|\Lambda) - - Parameters - ---------- - dataset: dict - Dictionary containing entry "redshift" - parameters: dict - Dictionary of parameters - Returns - ------- - differential_volume: (float, array-like) - Differential spacetime volume - """ - differential_volume = ( - self.psi_of_z(redshift=dataset["redshift"], **parameters) - / (1 + dataset["redshift"]) - * self.dvc_dz(redshift=dataset["redshift"], **parameters) - ) - if bounds: - differential_volume *= dataset["redshift"] <= self.z_max - return differential_volume - - -class PowerLawRedshift(_Redshift): - r""" - Redshift model from Fishbach+ https://arxiv.org/abs/1805.10270 - - .. math:: - p(z|\gamma, \kappa, z_p) &\propto \frac{1}{1 + z}\frac{dV_c}{dz} \psi(z|\gamma, \kappa, z_p) - - \psi(z|\gamma, \kappa, z_p) &= (1 + z)^\lambda - - Parameters - ---------- - lamb: float - The spectral index. - """ - - _variable_names = ["lamb"] - - def psi_of_z(self, redshift, **parameters): - return (1 + redshift) ** parameters["lamb"] diff --git a/wcosmo/utils.py b/wcosmo/utils.py index 2510c93..80444c6 100644 --- a/wcosmo/utils.py +++ b/wcosmo/utils.py @@ -1,3 +1,7 @@ +import numpy as xp +from scipy.linalg import toeplitz + + def maybe_jit(func, *args, **kwargs): """ A decorator to jit the function if using jax. @@ -9,8 +13,75 @@ def maybe_jit(func, *args, **kwargs): `gwpopulation` regardless of cosmology. """ from .wcosmo import xp + if "jax" in xp.__name__: from jax import jit return jit(func, *args, **kwargs) return func + + +def pade(an, m, n=None): + """ + Return Pade approximation to a polynomial as the ratio of two polynomials. + + Parameters + ---------- + an : (N,) array_like + Taylor series coefficients. + m : int + The order of the returned approximating polynomial `q`. + n : int, optional + The order of the returned approximating polynomial `p`. By default, + the order is ``len(an)-1-m``. + + Returns + ------- + p, q : Polynomial class + The Pade approximation of the polynomial defined by `an` is + ``p(x)/q(x)``. + + Examples + -------- + >>> import numpy as np + >>> from scipy.interpolate import pade + >>> e_exp = [1.0, 1.0, 1.0/2.0, 1.0/6.0, 1.0/24.0, 1.0/120.0] + >>> p, q = pade(e_exp, 2) + + >>> e_exp.reverse() + >>> e_poly = np.poly1d(e_exp) + + Compare ``e_poly(x)`` and the Pade approximation ``p(x)/q(x)`` + + >>> e_poly(1) + 2.7166666666666668 + + >>> p(1)/q(1) + 2.7179487179487181 + + Notes + ----- + This code has been slightly edited from the numpy implementation to: + + - Use xp instead of np to support multiple backends + - Direclty use the fact that part of the matrix is Toeplitz + + """ + an = xp.asarray(an) + if n is None: + n = len(an) - 1 - m + if n < 0: + raise ValueError("Order of q must be smaller than len(an)-1.") + if n < 0: + raise ValueError("Order of p must be greater than 0.") + N = m + n + if N > len(an) - 1: + raise ValueError("Order of q+p must be smaller than len(an).") + an = an[: N + 1] + Akj = xp.eye(N + 1, n + 1, dtype=an.dtype) + Bkj = toeplitz(xp.r_[0.0, -an[:-1]], xp.zeros(m)) + Ckj = xp.hstack((Akj, Bkj)) + pq = xp.linalg.solve(Ckj, an) + p = pq[: n + 1] + q = xp.r_[1.0, pq[n + 1 :]] + return p[::-1], q[::-1] diff --git a/wcosmo/wcosmo.py b/wcosmo/wcosmo.py index 7c86ae9..dc57444 100644 --- a/wcosmo/wcosmo.py +++ b/wcosmo/wcosmo.py @@ -1,9 +1,8 @@ from functools import partial import numpy as xp -from scipy.linalg import toeplitz -from .utils import maybe_jit +from .utils import maybe_jit, pade __all__ = [ @@ -34,21 +33,37 @@ _cosmology_docstrings_ = dict( - z=""" - z : array_like + z="""z: array_like Redshift""", - Om0=""" - Om0 : array_like + Om0="""Om0: array_like The matter density fraction""", - w0=""" - w0: array_like + w0="""w0: array_like The (constant) equation of state parameter for dark energy""", - H0=""" - H0 : float + H0="""H0: float The Hubble constant in km/s/Mpc""", + zmin="""zmin: float + The minimum redshift used in the conversion from distance to redshift, + default=1e-4""", + zmax="""zmax: float + The maximum redshift used in the conversion from distance to redshift, + default=100""", + name="""name: str + The name for the cosmology, mostly used for fixed instances""", + meta="""meta: dict + Additional metadata describing the cosmology, e.g., citation + information""", ) +def _autodoc(func): + """ + Simple decorator to mark that a docstring needs formatting + """ + func.__doc__ = func.__doc__.format(**_cosmology_docstrings_) + return func + + +@_autodoc def flat_wcdm_taylor_expansion(w0): r""" Taylor coefficients expansion of E(z) as as a function @@ -56,120 +71,56 @@ def flat_wcdm_taylor_expansion(w0): .. math:: - F(x; w) = \sum_{n=0}^{{\infty} (-1/2 C n) x^n / (1/2 - 3wn) - = 2 \sum_{n=0}^{\infty} (-1/2 C n) x^n / (1 - 6wn) + F(x; w) = \sum_{{n=0}}^{{\infty}} \binom{{-1/2}}{{n}} x^n / (1/2 - 3wn) + = 2 \sum_{{n=0}}^{{\infty}} \binom{{-1/2}}{{n}} x^n / (1 - 6wn) Parameters ---------- - w0: array_like - The (constant) equation of state parameter for dark energy + {w0} Returns ------- xp.ndarray The Taylor expansion coefficients. """ - return xp.array([ - xp.ones_like(w0), - -1 / (2 * (1 - 6 * w0)), - 3 / (8 * (1 - 12 * w0)), - -5 / (16 * (1 - 18 * w0)), - 35 / (128 * (1 - 24 * w0)), - -63 / (256 * (1 - 30 * w0)), - 231 / (1024 * (1 - 36 * w0)), - -429 / (2048 * (1 - 42 * w0)), - 6435 / (32768 * (1 - 48 * w0)), - -12155 / (65536 * (1 - 54 * w0)), - 46189 / (262144 * (1 - 60 * w0)), - -88179 / (524288 * (1 - 66 * w0)), - 676039 / (4194304 * (1 - 72 * w0)), - -1300075 / (8388608 * (1 - 78 * w0)), - 5014575 / (33554432 * (1 - 84 * w0)), - -9694845 / (67108864 * (1 - 90 * w0)), - 300540195 / (268435456 * (1 - 96 * w0)), - ]) - - -def pade(an, m, n=None): - """ - Return Pade approximation to a polynomial as the ratio of two polynomials. - - Parameters - ---------- - an : (N,) array_like - Taylor series coefficients. - m : int - The order of the returned approximating polynomial `q`. - n : int, optional - The order of the returned approximating polynomial `p`. By default, - the order is ``len(an)-1-m``. - - Returns - ------- - p, q : Polynomial class - The Pade approximation of the polynomial defined by `an` is - ``p(x)/q(x)``. - - Examples - -------- - >>> import numpy as np - >>> from scipy.interpolate import pade - >>> e_exp = [1.0, 1.0, 1.0/2.0, 1.0/6.0, 1.0/24.0, 1.0/120.0] - >>> p, q = pade(e_exp, 2) - - >>> e_exp.reverse() - >>> e_poly = np.poly1d(e_exp) - - Compare ``e_poly(x)`` and the Pade approximation ``p(x)/q(x)`` - - >>> e_poly(1) - 2.7166666666666668 - - >>> p(1)/q(1) - 2.7179487179487181 - - Notes - ----- - This code has been slightly edited from the numpy implementation to: - - - Use xp instead of np to support multiple backends - - Direclty use the fact that part of the matrix is Toeplitz - - """ - an = xp.asarray(an) - if n is None: - n = len(an) - 1 - m - if n < 0: - raise ValueError("Order of q must be smaller than len(an)-1.") - if n < 0: - raise ValueError("Order of p must be greater than 0.") - N = m + n - if N > len(an) - 1: - raise ValueError("Order of q+p must be smaller than len(an).") - an = an[: N + 1] - Akj = xp.eye(N + 1, n + 1, dtype=an.dtype) - Bkj = toeplitz(xp.r_[0.0, -an[:-1]], xp.zeros(m)) - Ckj = xp.hstack((Akj, Bkj)) - pq = xp.linalg.solve(Ckj, an) - p = pq[: n + 1] - q = xp.r_[1.0, pq[n + 1 :]] - return p[::-1], q[::-1] + return xp.array( + [ + w0**0, + -1 / (2 * (1 - 6 * w0)), + 3 / (8 * (1 - 12 * w0)), + -5 / (16 * (1 - 18 * w0)), + 35 / (128 * (1 - 24 * w0)), + -63 / (256 * (1 - 30 * w0)), + 231 / (1024 * (1 - 36 * w0)), + -429 / (2048 * (1 - 42 * w0)), + 6435 / (32768 * (1 - 48 * w0)), + -12155 / (65536 * (1 - 54 * w0)), + 46189 / (262144 * (1 - 60 * w0)), + -88179 / (524288 * (1 - 66 * w0)), + 676039 / (4194304 * (1 - 72 * w0)), + -1300075 / (8388608 * (1 - 78 * w0)), + 5014575 / (33554432 * (1 - 84 * w0)), + -9694845 / (67108864 * (1 - 90 * w0)), + 300540195 / (268435456 * (1 - 96 * w0)), + ] + ) @maybe_jit +@_autodoc def efunc(z, Om0, w0=-1): - f""" + """ Compute the E(z) function for a flat Lambda CDM cosmology. Parameters ---------- - {_cosmology_docstrings_["z"]} - {_cosmology_docstrings_["Om0"]} - {_cosmology_docstrings_["w0"]} + {z} + {Om0} + {w0} Returns ------- - E(z) : array_like + E(z): array_like The E(z) function """ zp1 = 1 + z @@ -177,35 +128,37 @@ def efunc(z, Om0, w0=-1): @maybe_jit +@_autodoc def inv_efunc(z, Om0, w0=-1): - f""" + """ Compute the inverse of the E(z) function for a flat Lambda CDM cosmology. Parameters ---------- - {_cosmology_docstrings_["z"]} - {_cosmology_docstrings_["Om0"]} - {_cosmology_docstrings_["w0"]} + {z} + {Om0} + {w0} Returns ------- - inv_efunc : array_like + inv_efunc: array_like The inverse of the E(z) function """ return 1 / efunc(z, Om0, w0) +@_autodoc def hubble_distance(H0): - f""" + """ Compute the Hubble distance. Parameters ---------- - {_cosmology_docstrings_["H0"]} + {H0} Returns ------- - D_H : float + D_H: float The Hubble distance in Mpc Notes @@ -216,21 +169,22 @@ def hubble_distance(H0): return speed_of_light / H0 +@_autodoc def Phi(z, Om0, w0=-1): - f""" + """ Compute the Pade approximation to 1 / E(z) as described in arXiv:1111.6396. We extend it to include a variable dark energy equation of state and include more terms in the expansion. Parameters ---------- - {_cosmology_docstrings_["z"]} - {_cosmology_docstrings_["Om0"]} - {_cosmology_docstrings_["w0"]} + {z} + {Om0} + {w0} Returns ------- - Phi : array_like + Phi: array_like The Pade approximation to 1 / E(z) """ x = (1 - Om0) / Om0 * (1 + z) ** (-3) @@ -238,19 +192,20 @@ def Phi(z, Om0, w0=-1): return xp.polyval(p, x) / xp.polyval(q, x) +@_autodoc def flat_wcdm_pade_coefficients(w0=-1): - f""" + """ Compute the Pade coefficients as described in arXiv:1111.6396. I expand the series to a bunch more terms to get a better fit. Parameters ---------- - {_cosmology_docstrings_["w0"]} + {w0} Returns ------- - p, q : xp.ndarray + p, q: xp.ndarray The Pade coefficients """ coeffs = flat_wcdm_taylor_expansion(w0=w0) @@ -258,41 +213,43 @@ def flat_wcdm_pade_coefficients(w0=-1): return p, q +@_autodoc def analytic_integral(z, Om0, w0=-1): - f""" + """ Evaluate the analytic integral of 1 / E(z) from infty to z assuming the Pade approximation to 1 / E(z). Parameters ---------- - {_cosmology_docstrings_["z"]} - {_cosmology_docstrings_["Om0"]} - {_cosmology_docstrings_["w0"]} + {z} + {Om0} + {w0} Returns ------- - integral : array_like + integral: array_like The integral of 1 / E(z) from infty to z """ return -2 / Om0**0.5 * Phi(z, Om0, w0) / (1 + z) ** 0.5 @maybe_jit +@_autodoc def comoving_distance(z, H0, Om0, w0=-1): - f""" + """ Compute the comoving distance using an analytic integral of the Pade approximation. Parameters ---------- - {_cosmology_docstrings_["z"]} - {_cosmology_docstrings_["H0"]} - {_cosmology_docstrings_["Om0"]} - {_cosmology_docstrings_["w0"]} - + {z} + {H0} + {Om0} + {w0} + Returns ------- - comoving_distance : array_like + comoving_distance: array_like The comoving distance in Mpc """ integral = analytic_integral(z, Om0=Om0) - analytic_integral(0, Om0=Om0, w0=w0) @@ -300,41 +257,43 @@ def comoving_distance(z, H0, Om0, w0=-1): @maybe_jit +@_autodoc def luminosity_distance(z, H0, Om0, w0=-1): - f""" + """ Compute the luminosity distance using an analytic integral of the Pade approximation. Parameters ---------- - {_cosmology_docstrings_["z"]} - {_cosmology_docstrings_["H0"]} - {_cosmology_docstrings_["Om0"]} - {_cosmology_docstrings_["w0"]} - + {z} + {H0} + {Om0} + {w0} + Returns ------- - luminosity_distance : array_like + luminosity_distance: array_like The luminosity distance in Mpc """ return (1 + z) * comoving_distance(z, H0, Om0, w0) @maybe_jit +@_autodoc def dDLdz(z, H0, Om0, w0=-1): - f""" + """ The Jacobian for the conversion of redshift to luminosity distance. Parameters ---------- - {_cosmology_docstrings_["z"]} - {_cosmology_docstrings_["H0"]} - {_cosmology_docstrings_["Om0"]} - {_cosmology_docstrings_["w0"]} + {z} + {H0} + {Om0} + {w0} Returns ------- - dDLdz : array_like + dDLdz: array_like The derivative of the luminosity distance with respect to redshift Notes @@ -348,6 +307,7 @@ def dDLdz(z, H0, Om0, w0=-1): return dC + (1 + z) * D_H * Ez_i +@_autodoc def z_at_value(func, fval, zmin=1e-4, zmax=100): """ Compute the redshift at which a function equals a given value. @@ -357,18 +317,16 @@ def z_at_value(func, fval, zmin=1e-4, zmax=100): Parameters ---------- - func : callable + func: callable The function to evaluate, e.g., luminosity_distance - fval : float + fval: float The value of the function at the desired redshift - zmin : float, optional - The minimum redshift to consider, default=1e-4 - zmax : float, optional - The maximum redshift to consider, default=100 + {zmin} + {zmax} Returns ------- - z : float + z: float The redshift at which the function equals the desired value """ zs = xp.logspace(xp.log10(zmin), xp.log10(zmax), 1000) @@ -376,20 +334,21 @@ def z_at_value(func, fval, zmin=1e-4, zmax=100): @maybe_jit +@_autodoc def differential_comoving_volume(z, H0, Om0, w0=-1): - f""" + """ Compute the differential comoving volume element. Parameters ---------- - {_cosmology_docstrings_["z"]} - {_cosmology_docstrings_["H0"]} - {_cosmology_docstrings_["Om0"]} - {_cosmology_docstrings_["w0"]} + {z} + {H0} + {Om0} + {w0} Returns ------- - dVc : array_like + dVc: array_like The differential comoving volume element in Gpc^3 """ dC = comoving_distance(z, H0, Om0, w0=w0) @@ -399,10 +358,9 @@ def differential_comoving_volume(z, H0, Om0, w0=-1): @maybe_jit -def detector_to_source_frame( - m1z, m2z, dL, H0, Om0, w0=-1, zmin=1e-4, zmax=100 -): - f""" +@_autodoc +def detector_to_source_frame(m1z, m2z, dL, H0, Om0, w0=-1, zmin=1e-4, zmax=100): + """ Convert masses and luminosity distance from the detector frame to source frame masses and redshift. @@ -411,23 +369,21 @@ def detector_to_source_frame( Parameters ---------- - m1z : array_like + m1z: array_like The primary mass in the detector frame - m2z : array_like + m2z: array_like The secondary mass in the detector frame - dL : array_like + dL: array_like The luminosity distance in Mpc - {_cosmology_docstrings_["H0"]} - {_cosmology_docstrings_["Om0"]} - {_cosmology_docstrings_["w0"]} - zmin : float, optional - The minimum redshift to consider, default=1e-4 - zmax : float, optional - The maximum redshift to consider, default=100 + {H0} + {Om0} + {w0} + {zmin} + {zmax} Returns ------- - m1, m2, z : array_like + m1, m2, z: array_like The primary and secondary masses in the source frame and the redshift """ distance_func = partial(luminosity_distance, H0=H0, Om0=Om0, w0=w0) @@ -438,24 +394,25 @@ def detector_to_source_frame( @maybe_jit +@_autodoc def source_to_detector_frame(m1, m2, z, H0, Om0, w0=-1): - f""" + """ Convert masses and redshift from the source frame to the detector frame. Parameters ---------- - m1 : array_like + m1: array_like The primary mass in the source frame - m2 : array_like + m2: array_like The secondary mass in the source frame - {_cosmology_docstrings_["z"]} - {_cosmology_docstrings_["H0"]} - {_cosmology_docstrings_["Om0"]} - {_cosmology_docstrings_["w0"]} - + {z} + {H0} + {Om0} + {w0} + Returns ------- - m1z, m2z, dL : array_like + m1z, m2z, dL: array_like The primary and secondary masses in the detector frame and the luminosity distance """ @@ -464,26 +421,47 @@ def source_to_detector_frame(m1, m2, z, H0, Om0, w0=-1): @maybe_jit +@_autodoc def comoving_volume(z, H0, Om0, w0=-1): - f""" + """ Compute the comoving volume out to redshift z. Parameters ---------- - {_cosmology_docstrings_["z"]} - {_cosmology_docstrings_["H0"]} - {_cosmology_docstrings_["Om0"]} - {_cosmology_docstrings_["w0"]} - + {z} + {H0} + {Om0} + {w0} + Returns ------- - Vc : array_like + Vc: array_like The comoving volume in Gpc^3 """ return 4 / 3 * xp.pi * comoving_distance(z, H0, Om0, w0=w0) ** 3 +@_autodoc class FlatwCDM: + r""" + Implementation of FlatwCDM cosmology to (approximately) match the astropy + API. + + .. math:: + + E(z) = \sqrt{{\Omega_{{m,0}} (1 + z)^3 + (1 - \Omega_{{m,0}}) (1 + z)^{{3(1 - w0)}}}} + + Parameters + ---------- + {H0} + {Om0} + {w0} + {zmin} + {zmax} + {name} + {meta} + """ + def __init__( self, H0, @@ -560,7 +538,25 @@ def comoving_volume(self, z): return comoving_volume(z, **self._kwargs) +@_autodoc class FlatLambdaCDM(FlatwCDM): + r""" + Implementation of FlatLambdaCDM cosmology to (approximately) match the + astropy API. This is the same as the :code:`FlatwCDM` with :code:`w0=-1`. + + .. math:: + + E(z) = \sqrt{{\Omega_{{m,0}} (1 + z)^3 + (1 - \Omega_{{m,0}})}} + + Parameters + ---------- + {H0} + {Om0} + {zmin} + {zmax} + {name} + {meta} + """ def __init__(self, H0, Om0, *, zmin=0.0001, zmax=100, name=None, meta=None): super().__init__(H0, Om0, w0=-1, zmin=zmin, zmax=zmax, name=name, meta=meta) From 71cccf0f78ea8b205006a3ec4c18d1f2e39b6dc0 Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Wed, 15 May 2024 10:14:09 -0700 Subject: [PATCH 2/4] CI: fix installation --- .github/workflows/pages.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/pages.yml b/.github/workflows/pages.yml index 97a9d5f..9c6530e 100644 --- a/.github/workflows/pages.yml +++ b/.github/workflows/pages.yml @@ -38,7 +38,6 @@ jobs: - name: Install dependencies run: | - python -m pip install -r requirements.txt python -m pip install -r doc/pages_requirements.txt python -m pip install . From 11da9454b2aea57dd2c9f04a96ed5f9cdb514752 Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Wed, 15 May 2024 14:00:03 -0700 Subject: [PATCH 3/4] DOC: update docstrings and fix autosummary --- .../_templates/custom-class-template.rst | 6 +- doc/source/conf.py | 2 + wcosmo/utils.py | 73 ++++- wcosmo/wcosmo.py | 285 ++++++++++++------ 4 files changed, 265 insertions(+), 101 deletions(-) diff --git a/doc/source/_templates/custom-class-template.rst b/doc/source/_templates/custom-class-template.rst index 16ebb2f..5e252d6 100644 --- a/doc/source/_templates/custom-class-template.rst +++ b/doc/source/_templates/custom-class-template.rst @@ -11,7 +11,7 @@ .. automethod:: __init__ {% if methods %} - .. rubric:: {{ _('Methods') }} + .. rubric:: Methods .. autosummary:: {% for item in methods %} @@ -22,11 +22,11 @@ {% block attributes %} {% if attributes %} - .. rubric:: {{ _('Attributes') }} + .. rubric:: Attributes .. autosummary:: {% for item in attributes %} ~{{ name }}.{{ item }} {%- endfor %} {% endif %} - {% endblock %} \ No newline at end of file + {% endblock %} diff --git a/doc/source/conf.py b/doc/source/conf.py index 823058f..fe1ca29 100644 --- a/doc/source/conf.py +++ b/doc/source/conf.py @@ -35,3 +35,5 @@ html_theme = 'sphinx_nefertiti' html_static_path = ['_static'] + +autosummary_generate = True diff --git a/wcosmo/utils.py b/wcosmo/utils.py index 80444c6..d21c6d8 100644 --- a/wcosmo/utils.py +++ b/wcosmo/utils.py @@ -1,7 +1,78 @@ +""" +Helper functions that are not directly relevant to cosmology. +""" + import numpy as xp from scipy.linalg import toeplitz +_cosmology_docstrings_ = dict( + z="""z: array_like + Redshift""", + Om0="""Om0: array_like + The matter density fraction""", + w0="""w0: array_like + The (constant) equation of state parameter for dark energy""", + H0="""H0: array_like + The Hubble constant in km/s/Mpc""", + zmin="""zmin: float + The minimum redshift used in the conversion from distance to redshift, + default=1e-4""", + zmax="""zmax: float + The maximum redshift used in the conversion from distance to redshift, + default=100""", + name="""name: str + The name for the cosmology, mostly used for fixed instances""", + meta="""meta: dict + Additional metadata describing the cosmology, e.g., citation + information""", + m1="""m1: array_like + The primary mass in the source frame""", + m2="""m2: array_like + The secondary mass in the source frame""", + m1z="""m1z: array_like + The primary mass in the detector frame""", + m2z="""m2z: array_like + The secondary mass in the detector frame""", + dL="""dL: array_like + The luminosity distance in Mpc""", +) + + +def _autodoc(func, strip_wcdm=False, alt=None): + """ + Simple decorator to mark that a docstring needs formatting + """ + func.__doc__ = func.__doc__.format(**_cosmology_docstrings_) + return func + + +def _method_autodoc(strip_wcdm=False, alt=None): + """ + Simple decorator to mark that a docstring needs formatting. + This will strip the class level attributes of :code:`FlatwCDM` + from the dosctring and allow a docstring to be taken from + another function. + """ + + def new_wrapper(func): + def _strip_wcdm_parameters(doc): + for key in ["H0", "Om0", "w0"]: + doc = doc.replace(_cosmology_docstrings_[key], "") + return doc + + if alt is not None: + doc = alt.__doc__ + else: + doc = func.__doc__ + if strip_wcdm: + doc = _strip_wcdm_parameters(doc) + func.__doc__ = doc + return func + + return new_wrapper + + def maybe_jit(func, *args, **kwargs): """ A decorator to jit the function if using jax. @@ -10,7 +81,7 @@ def maybe_jit(func, *args, **kwargs): e.g., to specify static arguments. This function is pretty useful and so might make it into - `gwpopulation` regardless of cosmology. + :code:`gwpopulation` regardless of cosmology. """ from .wcosmo import xp diff --git a/wcosmo/wcosmo.py b/wcosmo/wcosmo.py index dc57444..d1da365 100644 --- a/wcosmo/wcosmo.py +++ b/wcosmo/wcosmo.py @@ -1,8 +1,12 @@ +""" +Core implementation of cosmology functionality. +""" + from functools import partial import numpy as xp -from .utils import maybe_jit, pade +from .utils import maybe_jit, pade, _autodoc, _method_autodoc __all__ = [ @@ -32,47 +36,18 @@ ] -_cosmology_docstrings_ = dict( - z="""z: array_like - Redshift""", - Om0="""Om0: array_like - The matter density fraction""", - w0="""w0: array_like - The (constant) equation of state parameter for dark energy""", - H0="""H0: float - The Hubble constant in km/s/Mpc""", - zmin="""zmin: float - The minimum redshift used in the conversion from distance to redshift, - default=1e-4""", - zmax="""zmax: float - The maximum redshift used in the conversion from distance to redshift, - default=100""", - name="""name: str - The name for the cosmology, mostly used for fixed instances""", - meta="""meta: dict - Additional metadata describing the cosmology, e.g., citation - information""", -) - - -def _autodoc(func): - """ - Simple decorator to mark that a docstring needs formatting - """ - func.__doc__ = func.__doc__.format(**_cosmology_docstrings_) - return func - - @_autodoc def flat_wcdm_taylor_expansion(w0): r""" - Taylor coefficients expansion of E(z) as as a function - of w. + Taylor coefficients expansion of :math:`E(z)` as as a function + of :math:`w_0`. .. math:: - F(x; w) = \sum_{{n=0}}^{{\infty}} \binom{{-1/2}}{{n}} x^n / (1/2 - 3wn) - = 2 \sum_{{n=0}}^{{\infty}} \binom{{-1/2}}{{n}} x^n / (1 - 6wn) + F(x; w_0) = 2\sum_{{n=0}}^{{\infty}} \binom{{-\frac{{1}}{{2}}}}{{n}} + \frac{{x^n}}{{\left(1 - 6nw_0\right)}} + + We include terms up to :math:`n=16`. Parameters ---------- @@ -80,7 +55,7 @@ def flat_wcdm_taylor_expansion(w0): Returns ------- - xp.ndarray + array_like The Taylor expansion coefficients. """ return xp.array( @@ -109,8 +84,13 @@ def flat_wcdm_taylor_expansion(w0): @maybe_jit @_autodoc def efunc(z, Om0, w0=-1): - """ - Compute the E(z) function for a flat Lambda CDM cosmology. + r""" + Compute the :math:`E(z)` function for a flat wCDM cosmology. + + .. math:: + + E(z; \Omega_{{m,0}}, w_0) = \sqrt{{\Omega_{{m,0}} (1 + z)^3 + + (1 - \Omega_{{m,0}}) (1 + z)^{{3(1 - w_0)}}}} Parameters ---------- @@ -131,7 +111,7 @@ def efunc(z, Om0, w0=-1): @_autodoc def inv_efunc(z, Om0, w0=-1): """ - Compute the inverse of the E(z) function for a flat Lambda CDM cosmology. + Compute the inverse of the E(z) function for a flat wCDM cosmology. Parameters ---------- @@ -149,8 +129,8 @@ def inv_efunc(z, Om0, w0=-1): @_autodoc def hubble_distance(H0): - """ - Compute the Hubble distance. + r""" + Compute the Hubble distance :math:`D_H = c H_0^{{-1}}`. Parameters ---------- @@ -160,21 +140,54 @@ def hubble_distance(H0): ------- D_H: float The Hubble distance in Mpc - - Notes - ----- - I hard code the speed of light in km/s """ speed_of_light = 299792.458 return speed_of_light / H0 @_autodoc -def Phi(z, Om0, w0=-1): +def hubble_parameter(z, H0, Om0, w0=-1): + r""" + Compute the Hubble parameter :math:`H(z)` for a flat wCDM cosmology. + + .. math:: + + H(z; H_0, \Omega_{{m,0}}, w_0) = \frac{{d_H(H_0)}}{{E(z; \Omega_{{m,0}}, w_0)}} + + Parameters + ---------- + {z} + {H0} + {Om0} + {w0} + + Returns + ------- + H(z): array_like + The Hubble parameter """ - Compute the Pade approximation to 1 / E(z) as described in arXiv:1111.6396. - We extend it to include a variable dark energy equation of state and - include more terms in the expansion. + return hubble_distance(H0=H0) * inv_efunc(z=z, H0=H0, Om0=Om0, w0=w0) + + +@_autodoc +def Phi(z, Om0, w0=-1): + r""" + Compute the Pade approximation to :math:`1 / E(z)` as described in + arXiv:1111.6396. We extend it to include a variable dark energy + equation of state and include more terms in the expansion + + .. math:: + + \Phi(z; \Omega_{{m, 0}}, w_0) = + \frac{{\sum_i^n \alpha_i x^i}}{{1 + \sum_{{j=1}}^m \beta_j x^j}}. + + Here the expansion is in terms of + + .. math:: + + x = \left(\frac{{1 - \Omega_{{m, 0}}}}{{\Omega_{{m, 0}}}}\right) (1 + z)^{{3 w_0}}. + + In practice we use :math:`m=n=7` whereas Adachi and Kasai use :math:`m=n=3`. Parameters ---------- @@ -185,9 +198,9 @@ def Phi(z, Om0, w0=-1): Returns ------- Phi: array_like - The Pade approximation to 1 / E(z) + The Pade approximation to :math:`1 / E(z)` """ - x = (1 - Om0) / Om0 * (1 + z) ** (-3) + x = (1 - Om0) / Om0 * (1 + z) ** (3 * w0) p, q = flat_wcdm_pade_coefficients(w0=w0) return xp.polyval(p, x) / xp.polyval(q, x) @@ -196,8 +209,11 @@ def Phi(z, Om0, w0=-1): def flat_wcdm_pade_coefficients(w0=-1): """ Compute the Pade coefficients as described in arXiv:1111.6396. - I expand the series to a bunch more terms to get a better fit. + We make two primary changes: + - allow a variable dark energy equation of state :math:`w_0` by changing + the definition of :math:`x`. + - include more terms (17) in the Taylor expansion. Parameters ---------- @@ -205,7 +221,7 @@ def flat_wcdm_pade_coefficients(w0=-1): Returns ------- - p, q: xp.ndarray + p, q: array_like The Pade coefficients """ coeffs = flat_wcdm_taylor_expansion(w0=w0) @@ -215,9 +231,15 @@ def flat_wcdm_pade_coefficients(w0=-1): @_autodoc def analytic_integral(z, Om0, w0=-1): - """ - Evaluate the analytic integral of 1 / E(z) from infty to z - assuming the Pade approximation to 1 / E(z). + r""" + .. math:: + + f(z; \Omega_{{m, 0}}, w_0) = \int_{{\infty}}^{{z}} + \frac{{dz'}}{{E(z'; \Omega_{{m, 0}}, w_0)}} + = \frac{{-2\Phi(z; \Omega_{{m, 0}}, w_0)}}{{\sqrt{{\Omega_{{m, 0}}(1 + z)}}}}. + + The integral is approximated using the Pade approximation and is up + to a factor the term in the braces in (1.1) of Adachi and Kasai. Parameters ---------- @@ -228,7 +250,7 @@ def analytic_integral(z, Om0, w0=-1): Returns ------- integral: array_like - The integral of 1 / E(z) from infty to z + The integral of :math:`1 / E(z)` from :math:`\infty` to :math:`z` """ return -2 / Om0**0.5 * Phi(z, Om0, w0) / (1 + z) ** 0.5 @@ -236,10 +258,16 @@ def analytic_integral(z, Om0, w0=-1): @maybe_jit @_autodoc def comoving_distance(z, H0, Om0, w0=-1): - """ + r""" Compute the comoving distance using an analytic integral of the Pade approximation. + .. math:: + + d_{{C}} = \frac{{c}}{{H_0}} \frac{{2}}{{\sqrt{{\Omega_{{m,0}}}}}} + \left( \Phi(0; \Omega_{{m, 0}}, w_0) + - \frac{{\Phi(z; \Omega_{{m, 0}}, w_0)}}{{\sqrt{{(1 + z)}}}}\right) + Parameters ---------- {z} @@ -259,10 +287,16 @@ def comoving_distance(z, H0, Om0, w0=-1): @maybe_jit @_autodoc def luminosity_distance(z, H0, Om0, w0=-1): - """ + r""" Compute the luminosity distance using an analytic integral of the Pade approximation. + .. math:: + + d_L = \frac{{c}}{{H_0}} \frac{{2(1 + z)}}{{\sqrt{{\Omega_{{m,0}}}}}} + \left( \Phi(0; \Omega_{{m, 0}}, w_0) + - \frac{{\Phi(z; \Omega_{{m, 0}}, w_0)}}{{\sqrt{{(1 + z)}}}}\right) + Parameters ---------- {z} @@ -281,9 +315,17 @@ def luminosity_distance(z, H0, Om0, w0=-1): @maybe_jit @_autodoc def dDLdz(z, H0, Om0, w0=-1): - """ + r""" The Jacobian for the conversion of redshift to luminosity distance. + .. math:: + + \frac{{dd_{{L}}}}{{z}} = d_C(z; H_0, \Omega_{{m,0}}, w_0) + + (1 + z) d_{{H}} E(z; \Omega_{{m, 0}}, w0) + + Here :math:`d_{{C}}` is comoving distance and :math:`d_{{H}}` is the Hubble + distance. + Parameters ---------- {z} @@ -295,11 +337,13 @@ def dDLdz(z, H0, Om0, w0=-1): ------- dDLdz: array_like The derivative of the luminosity distance with respect to redshift + in Mpc Notes ----- - This function does not have a direct analog in the `astropy` cosmology - objects, but is needed for accounting for fitting in luminosity distance + This function does not have a direct analog in the :code:`astropy` + cosmology objects, but is needed for accounting for expressing + distributions of redshift as distributions over luminosity distance. """ dC = comoving_distance(z, H0=H0, Om0=Om0, w0=w0) Ez_i = inv_efunc(z, Om0=Om0, w0=w0) @@ -312,13 +356,14 @@ def z_at_value(func, fval, zmin=1e-4, zmax=100): """ Compute the redshift at which a function equals a given value. - This follows the approach in `astropy`'s `z_at_value` function - closely, but uses linear interpolation instead of a root finder. + This follows the approach in :code:`astropy`'s :code:`z_at_value` + function closely, but uses linear interpolation instead of a root finder. Parameters ---------- func: callable - The function to evaluate, e.g., luminosity_distance + The function to evaluate, e.g., :code:`Planck15.luminosity_distance`, + this should take :code:`fval` as the only input. fval: float The value of the function at the desired redshift {zmin} @@ -336,9 +381,14 @@ def z_at_value(func, fval, zmin=1e-4, zmax=100): @maybe_jit @_autodoc def differential_comoving_volume(z, H0, Om0, w0=-1): - """ + r""" Compute the differential comoving volume element. + .. math:: + + \frac{{dV_{{C}}}}{{dz}} = d_C^2(z; H_0, \Omega_{{m,0}}, w_0) d_H + E(z; \Omega_{{m, 0}}, w_0) + Parameters ---------- {z} @@ -349,7 +399,7 @@ def differential_comoving_volume(z, H0, Om0, w0=-1): Returns ------- dVc: array_like - The differential comoving volume element in Gpc^3 + The differential comoving volume element in :math:`\rm{{Gpc}}^3` """ dC = comoving_distance(z, H0, Om0, w0=w0) Ez_i = inv_efunc(z, Om0, w0=w0) @@ -369,12 +419,9 @@ def detector_to_source_frame(m1z, m2z, dL, H0, Om0, w0=-1, zmin=1e-4, zmax=100): Parameters ---------- - m1z: array_like - The primary mass in the detector frame - m2z: array_like - The secondary mass in the detector frame - dL: array_like - The luminosity distance in Mpc + {m1z} + {m2z} + {dL} {H0} {Om0} {w0} @@ -401,10 +448,8 @@ def source_to_detector_frame(m1, m2, z, H0, Om0, w0=-1): Parameters ---------- - m1: array_like - The primary mass in the source frame - m2: array_like - The secondary mass in the source frame + {m1} + {m2} {z} {H0} {Om0} @@ -423,9 +468,13 @@ def source_to_detector_frame(m1, m2, z, H0, Om0, w0=-1): @maybe_jit @_autodoc def comoving_volume(z, H0, Om0, w0=-1): - """ + r""" Compute the comoving volume out to redshift z. + .. math:: + + V_C = \frac{{4\pi}}{{3}} d^3_C(z; H_0, \Omega_{{m,0}}, w_0) + Parameters ---------- {z} @@ -436,7 +485,7 @@ def comoving_volume(z, H0, Om0, w0=-1): Returns ------- Vc: array_like - The comoving volume in Gpc^3 + The comoving volume in :math:`\rm{{Gpc}}^3` """ return 4 / 3 * xp.pi * comoving_distance(z, H0, Om0, w0=w0) ** 3 @@ -444,12 +493,13 @@ def comoving_volume(z, H0, Om0, w0=-1): @_autodoc class FlatwCDM: r""" - Implementation of FlatwCDM cosmology to (approximately) match the astropy - API. + Implementation of flat wCDM cosmology to (approximately) match the + :code:`astropy` API. .. math:: - E(z) = \sqrt{{\Omega_{{m,0}} (1 + z)^3 + (1 - \Omega_{{m,0}}) (1 + z)^{{3(1 - w0)}}}} + E(z) = \sqrt{{\Omega_{{m,0}} (1 + z)^3 + + (1 - \Omega_{{m,0}}) (1 + z)^{{3(1 - w_0)}}}} Parameters ---------- @@ -481,8 +531,16 @@ def __init__( self.name = name self.meta = meta + @property + def _kwargs(self): + return {"H0": self.H0, "Om0": self.Om0, "w0": self.w0} + @property def meta(self): + """ + Meta data for the cosmology to hold additional information, e.g., + citation information + """ return self._meta @meta.setter @@ -492,48 +550,69 @@ def meta(self, meta): self._meta = meta @property - def _kwargs(self): - return {"H0": self.H0, "Om0": self.Om0, "w0": self.w0} - - @property + @_method_autodoc(strip_wcdm=True, alt=hubble_distance) def hubble_distance(self): return hubble_distance(self.H0) + @_method_autodoc(strip_wcdm=True, alt=luminosity_distance) def luminosity_distance(self, z): return luminosity_distance(z, **self._kwargs) + @_autodoc def dLdH(self, z): + r""" + Derivative of the luminosity distance w.r.t. the Hubble distance. + + .. math:: + + \frac{{dd_L}}{{dd_H}} = \frac{{d_L}}{{d_H}} + + Parameters + ---------- + {z} + + Returns + ------- + array_like: + The derivative of the luminosity distance w.r.t., the Hubble distance + """ return self.luminosity_distance(z) / self.hubble_distance + @_method_autodoc(strip_wcdm=True, alt=dDLdz) def dDLdz(self, z): return dDLdz(z, **self._kwargs) + @_method_autodoc(strip_wcdm=True, alt=differential_comoving_volume) def differential_comoving_volume(self, z): return differential_comoving_volume(z, **self._kwargs) + @_method_autodoc(strip_wcdm=True, alt=detector_to_source_frame) def detector_to_source_frame(self, m1z, m2z, dL): return detector_to_source_frame( m1z, m2z, dL, **self._kwargs, zmin=self.zmin, zmax=self.zmax ) + @_method_autodoc(strip_wcdm=True, alt=source_to_detector_frame) def source_to_detector_frame(self, m1, m2, z): return source_to_detector_frame(m1, m2, z, **self._kwargs) - def log_differential_comoving_volume(self, z): - return xp.log(self.differential_comoving_volume(z)) - + @_method_autodoc(strip_wcdm=True, alt=efunc) def efunc(self, z): - return efunc(z, self.Om0) + return efunc(z, self.Om0, self.w0) + @_method_autodoc(strip_wcdm=True, alt=inv_efunc) def inv_efunc(self, z): - return inv_efunc(z, self.Om0) + return inv_efunc(z, self.Om0, self.w0) + @_method_autodoc(strip_wcdm=True, alt=hubble_parameter) def H(self, z): - return self.H0 * self.efunc(z) + return hubble_parameter(z, **self._kwargs) + @_method_autodoc(strip_wcdm=True, alt=comoving_distance) def comoving_distance(self, z): return comoving_distance(z, **self._kwargs) + @_method_autodoc(strip_wcdm=True, alt=comoving_volume) def comoving_volume(self, z): return comoving_volume(z, **self._kwargs) @@ -541,8 +620,9 @@ def comoving_volume(self, z): @_autodoc class FlatLambdaCDM(FlatwCDM): r""" - Implementation of FlatLambdaCDM cosmology to (approximately) match the - astropy API. This is the same as the :code:`FlatwCDM` with :code:`w0=-1`. + Implementation of a flat :math:`\Lambda\rm{{CDM}}` cosmology to + (approximately) match the :code:`astropy` API. This is the same as + the :code:`FlatwCDM` with :math:`w_0=-1`. .. math:: @@ -558,8 +638,19 @@ class FlatLambdaCDM(FlatwCDM): {meta} """ - def __init__(self, H0, Om0, *, zmin=0.0001, zmax=100, name=None, meta=None): - super().__init__(H0, Om0, w0=-1, zmin=zmin, zmax=zmax, name=name, meta=meta) + def __init__( + self, + H0, + Om0, + *, + zmin=1e-4, + zmax=100, + name=None, + meta=None, + ): + super().__init__( + H0=H0, Om0=Om0, w0=-1, zmin=zmin, zmax=zmax, name=name, meta=meta + ) Planck13 = FlatLambdaCDM(H0=67.77, Om0=0.30712, name="Planck13") From 9d3e8a0aabd7ce4bfe426ab2d8daf065f1994f69 Mon Sep 17 00:00:00 2001 From: Colm Talbot Date: Wed, 15 May 2024 14:06:57 -0700 Subject: [PATCH 4/4] DOC: typo fixes in documentation --- wcosmo/wcosmo.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/wcosmo/wcosmo.py b/wcosmo/wcosmo.py index d1da365..40bcbc1 100644 --- a/wcosmo/wcosmo.py +++ b/wcosmo/wcosmo.py @@ -90,7 +90,7 @@ def efunc(z, Om0, w0=-1): .. math:: E(z; \Omega_{{m,0}}, w_0) = \sqrt{{\Omega_{{m,0}} (1 + z)^3 - + (1 - \Omega_{{m,0}}) (1 + z)^{{3(1 - w_0)}}}} + + (1 - \Omega_{{m,0}}) (1 + z)^{{3(1 + w_0)}}}} Parameters ---------- @@ -499,7 +499,7 @@ class FlatwCDM: .. math:: E(z) = \sqrt{{\Omega_{{m,0}} (1 + z)^3 - + (1 - \Omega_{{m,0}}) (1 + z)^{{3(1 - w_0)}}}} + + (1 - \Omega_{{m,0}}) (1 + z)^{{3(1 + w_0)}}}} Parameters ----------