![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\"){kind=icon}
"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "source": [
+ "# Spectral Siren Population Inference on GWTC-3\n",
+ "\n",
+ "The thrid gravitational-wave transient catalog [GWTC-3](https://arxiv.org/abs/1811.12907) includes all compact binary coalescences observed during Advanced LIGO/Virgo's first three oberving runs.\n",
+ "\n",
+ "\"Spectral siren\" cosmology (see, e.g., [Ezquiaga & Holz](https://arxiv.org/abs/2202.08240)) uses features in the black hole mass spectrum to constrain the distance-redshift relation.\n",
+ "\n",
+ "`GWPopulation` builds upon [Bilby](https://git.ligo.org/lscsoft/bilby) ([arXiv:1811.02042](https://arxiv.org/abs/1811.02042)) to provide simple, modular, user-friendly, population inference.\n",
+ "`wcosmo` provides optimized cosmology functionality for `GWPopulation`.\n",
+ "\n",
+ "There are many [implemented models](https://colmtalbot.github.io/gwpopulation/_autosummary/gwpopulation.models.html#module-gwpopulation.models) and an example of defining custom models is included below.\n",
+ "In this example we use:\n",
+ "\n",
+ "- A mass distribution in primary mass and mass ratio from Talbot & Thrane (2018) ([arXiv:1801:02699](https://arxiv.org/abs/1801.02699)). This is equivalent to the `PowerLaw + Peak` model used in LVK analyses without the low-mass smoothing for computational efficiency.\n",
+ "- Half-Gaussian + isotropic spin tilt distribution from Talbot & Thrane (2017) ([arXiv:1704.08370](https://arxiv.org/abs/1704.08370)).\n",
+ "- Beta spin magnitude distribution from Wysocki+ (2018) ([arXiv:1805:06442](https://arxiv.org/abs/1805.06442)).\n",
+ "- Each of these are also available with independent but identically distributed spins.\n",
+ "- Redshift evolution model as in Fishbach+ (2018) ([arXiv:1805.10270](https://arxiv.org/abs/1805.10270)).\n",
+ "- A variable Flat Lambda CDM cosmology using a modified version of the approximation in [arXiv:1111.6396](https://arxiv.org/abs/1111.6396) as implemented in `wcosmo`.\n",
+ "\n",
+ "For more information on `GWPopulation` see the [git repository](https://github.com/ColmTalbot/gwpopulation), [documentation](https://colmtalbot.github.io/gwpopulation/).\n",
+ "\n",
+ "\n",
+ "## Setup\n",
+ "\n",
+ "If you're using colab.research.google.com you will want to choose a GPU-accelerated runtime (I'm going to use a T4 GPU).\n",
+ "\n",
+ "\"runtime\"->\"change runtime type\"->\"Hardware accelerator = GPU\""
+ ],
+ "metadata": {
+ "id": "aeUxtTpEfzbn"
+ }
+ },
+ {
+ "cell_type": "markdown",
+ "source": [
+ "## Install some needed packages\n",
+ "\n",
+ "All of the dependencies for this are integrated into `GWPopulation`.\n",
+ "These include `wcosmo` for cosmology, and `Bilby` and `dynesty` for sampling."
+ ],
+ "metadata": {
+ "id": "Dc8zj3wngiYu"
+ }
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {
+ "id": "Tsv8JYinDZv0",
+ "collapsed": true
+ },
+ "outputs": [],
+ "source": [
+ "!pip install gwpopulation --quiet --progress-bar off"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "source": [
+ "## Download data\n",
+ "\n",
+ "We need to download the data for the events and simmulated \"injections\" used to characterize the detection sensitivity.\n",
+ "\n",
+ "### Event posteriors\n",
+ "\n",
+ "We're using the posteriors from the GWTC-3 data release in a pre-processed format.\n",
+ "\n",
+ "The file was produced by [gwpopulation-pipe](https://docs.ligo.org/ratesAndPopulations/gwpopulation_pipe) to reduce the many GB of posterior sample files to a single ~30Mb file.\n",
+ "\n",
+ "The choice of events in this file was not very careful and should only be considered qualitatively correct.\n",
+ "\n",
+ "The data file can be found [here](https://drive.google.com/drive/folders/1wyfR6sYvYVdBefF9_vrVTp0Btu03OlzL?usp=drive_link).\n",
+ "The original data can be found at [zenodo:5546663](https://zenodo.org/records/5546663) and [zenodo:6513631](https://zenodo.org/records/6513631) along with citation information.\n",
+ "\n",
+ "### Sensitivity injections\n",
+ "\n",
+ "Again I have pre-processed the full injection set using `gwpopulation-pipe` to reduce the filesize.\n",
+ "The original data is available at [zenodo:7890398](https://zenodo.org/records/7890398) along with citation information."
+ ],
+ "metadata": {
+ "id": "TdiqeGZtHwRw"
+ }
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "!gdown https://drive.google.com/uc?id=16gStLIjt65gWBkw-gNOVUqNbZ89q8CLF\n",
+ "!gdown https://drive.google.com/uc?id=10pevUCM3V2-D-bROFEMAcTJsX_9RzeM6"
+ ],
+ "metadata": {
+ "collapsed": true,
+ "id": "s8bqvo26HABd",
+ "outputId": "efbf2542-dfb3-40d2-bda8-4e9047f5f8d1",
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ }
+ },
+ "execution_count": 2,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "Downloading...\n",
+ "From: https://drive.google.com/uc?id=16gStLIjt65gWBkw-gNOVUqNbZ89q8CLF\n",
+ "To: /content/gwtc-3-injections.pkl\n",
+ "100% 2.69M/2.69M [00:00<00:00, 60.0MB/s]\n",
+ "Downloading...\n",
+ "From: https://drive.google.com/uc?id=10pevUCM3V2-D-bROFEMAcTJsX_9RzeM6\n",
+ "To: /content/gwtc-3-samples.pkl\n",
+ "100% 36.4M/36.4M [00:00<00:00, 91.6MB/s]\n"
+ ]
+ }
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "source": [
+ "## Imports\n",
+ "\n",
+ "Import the packages required for the script.\n",
+ "We also set the backend for array operations to `jax` which allows us to take advantage of just-in-time (jit) compilation in addition to GPU-parallelisation when available."
+ ],
+ "metadata": {
+ "id": "I1MQIoc_gvR-"
+ }
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "import bilby as bb\n",
+ "import gwpopulation as gwpop\n",
+ "import jax\n",
+ "import matplotlib.pyplot as plt\n",
+ "import numpy as np\n",
+ "import pandas as pd\n",
+ "from bilby.core.prior import PriorDict, Uniform\n",
+ "from gwpopulation.experimental.cosmo_models import CosmoModel\n",
+ "from gwpopulation.experimental.jax import JittedLikelihood\n",
+ "from wcosmo import Planck15\n",
+ "\n",
+ "gwpop.set_backend(\"jax\")\n",
+ "\n",
+ "xp = gwpop.utils.xp\n",
+ "\n",
+ "%matplotlib inline"
+ ],
+ "metadata": {
+ "id": "mKC8rhSUDmtg"
+ },
+ "execution_count": 3,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "source": [
+ "## Load posteriors\n",
+ "\n",
+ "We remove two events from the file that shouldn't be there that have NS-like secondaries as we are just interested in BBHs for this demonstration.\n",
+ "\n",
+ "We also need to modify some of the stored parameters as we are going to include a cosmological fit so we will fit detector frame masses and luminosity distance."
+ ],
+ "metadata": {
+ "id": "6MIE6_uMg4vw"
+ }
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "posteriors = pd.read_pickle(\"gwtc-3-samples.pkl\")\n",
+ "del posteriors[15]\n",
+ "del posteriors[38]\n",
+ "\n",
+ "for post in posteriors:\n",
+ " zs = post.pop(\"redshift\").values\n",
+ " post[\"mass_1_detector\"] = post.pop(\"mass_1\") * (1 + zs)\n",
+ " post[\"luminosity_distance\"] = np.asarray(Planck15.luminosity_distance(zs))\n",
+ " post[\"prior\"] /= np.asarray(Planck15.dDLdz(zs) * (1 + zs))"
+ ],
+ "metadata": {
+ "id": "1pAakjJ1DwrK"
+ },
+ "execution_count": 4,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "source": [
+ "## Load injections\n",
+ "\n",
+ "Load the injections used to characterize the sensitivity of the gravitaitonal-wave survey.\n",
+ "\n",
+ "Again, we need to make some modifications for fitting in detector-frame quantities."
+ ],
+ "metadata": {
+ "id": "82GMaSlJhKm1"
+ }
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "import dill\n",
+ "\n",
+ "with open(\"gwtc-3-injections.pkl\", \"rb\") as ff:\n",
+ " injections = dill.load(ff)\n",
+ "\n",
+ "zs = np.asarray(injections.pop(\"redshift\"))\n",
+ "injections[\"mass_1_detector\"] = injections.pop(\"mass_1\") * (1 + zs)\n",
+ "injections[\"luminosity_distance\"] = np.asarray(Planck15.luminosity_distance(zs))\n",
+ "injections[\"prior\"] /= np.asarray(Planck15.dDLdz(zs) * (1 + zs))"
+ ],
+ "metadata": {
+ "id": "IuHqDXFQE9Mw"
+ },
+ "execution_count": 5,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "source": [
+ "## Define some models and the likelihood\n",
+ "\n",
+ "We need to define `Bilby` `Model` objects for the numerator and denominator independently as these cache some computations interally.\n",
+ "\n",
+ "Note that we are using a `CosmoModel`, this in an experimental feature and so the specific API may change in future, but the basic funcionality should be stable.\n",
+ "We create a model that can infer the three parameters of `wCDM`.\n",
+ "\n",
+ "The `HyperparameterLikelihood` marginalises over the local merger rate, with a uniform-in-log prior.\n",
+ "The posterior for the merger rate can be recovered in post-processing.\n",
+ "\n",
+ "We provide:\n",
+ "- `posteriors`: a list of `pandas` DataFrames.\n",
+ "- `hyper_prior`: our population model, as defined above.\n",
+ "- `selection_function`: anything which evaluates the selection function.\n",
+ "\n",
+ "We can also provide:\n",
+ "- `conversion_function`: this converts between the parameters we sample in and those needed by the model, e.g., for sampling in the mean and variance of the beta distribution.\n",
+ "- `max_samples`: the maximum number of samples to use from each posterior, this defaults to the length of the shortest posterior."
+ ],
+ "metadata": {
+ "id": "I5FzBbnniQfl"
+ }
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "model = CosmoModel(\n",
+ " model_functions=[\n",
+ " gwpop.models.mass.two_component_primary_mass_ratio,\n",
+ " gwpop.models.spin.iid_spin,\n",
+ " gwpop.models.redshift.PowerLawRedshift(cosmo_model=\"FlatwCDM\"),\n",
+ " ],\n",
+ " cosmo_model=\"FlatwCDM\",\n",
+ ")\n",
+ "\n",
+ "vt = gwpop.vt.ResamplingVT(model=model, data=injections, n_events=len(posteriors))\n",
+ "\n",
+ "likelihood = gwpop.hyperpe.HyperparameterLikelihood(\n",
+ " posteriors=posteriors,\n",
+ " hyper_prior=model,\n",
+ " selection_function=vt,\n",
+ ")"
+ ],
+ "metadata": {
+ "id": "7IrVTnK5GbaZ"
+ },
+ "execution_count": 6,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "source": [
+ "## Define our prior\n",
+ "\n",
+ "The mass model has eight parameters that we vary that are described in arXiv:1801:02699. This model is sometimes referred to as \"PowerLaw+Peak\"\n",
+ "\n",
+ "The spin magnitude model is a `Beta` distribution with the usual parameterization, and the spin orientation model is a mixure of a uniform component and a truncated Gaussian that peaks at aligned spin. This combination is sometimes referred to as \"Default\".\n",
+ "\n",
+ "For redshift we use a model that looks like\n",
+ "\n",
+ "$$p(z) \\propto \\frac{d V_{c}}{dz} (1 + z)^{λ - 1}$$\n",
+ "\n",
+ "Finally, we set priors on the three parameters of the `wCDM` model, $H_0$ (the Hubble constant), $\\Omega_{m,0}$ (the matter fraction of the Universe at current time), $w_0$ (the constant dark energy equation of state)."
+ ],
+ "metadata": {
+ "id": "lQW_m3bNioih"
+ }
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "priors = PriorDict()\n",
+ "\n",
+ "# mass\n",
+ "priors[\"alpha\"] = Uniform(minimum=-2, maximum=4, latex_label=\"$\\\\alpha$\")\n",
+ "priors[\"beta\"] = Uniform(minimum=-4, maximum=12, latex_label=\"$\\\\beta$\")\n",
+ "priors[\"mmin\"] = Uniform(minimum=2, maximum=2.5, latex_label=\"$m_{\\\\min}$\")\n",
+ "priors[\"mmax\"] = Uniform(minimum=80, maximum=100, latex_label=\"$m_{\\\\max}$\")\n",
+ "priors[\"lam\"] = Uniform(minimum=0, maximum=1, latex_label=\"$\\\\lambda_{m}$\")\n",
+ "priors[\"mpp\"] = Uniform(minimum=10, maximum=50, latex_label=\"$\\\\mu_{m}$\")\n",
+ "priors[\"sigpp\"] = Uniform(minimum=1, maximum=10, latex_label=\"$\\\\sigma_{m}$\")\n",
+ "priors[\"gaussian_mass_maximum\"] = 100\n",
+ "# spin\n",
+ "priors[\"amax\"] = 1\n",
+ "priors[\"alpha_chi\"] = Uniform(\n",
+ " minimum=1, maximum=6, latex_label=\"$\\\\alpha_{\\\\chi}$\"\n",
+ ")\n",
+ "priors[\"beta_chi\"] = Uniform(minimum=1, maximum=6, latex_label=\"$\\\\beta_{\\\\chi}$\")\n",
+ "priors[\"xi_spin\"] = Uniform(minimum=0, maximum=1, latex_label=\"$\\\\xi$\")\n",
+ "priors[\"sigma_spin\"] = Uniform(minimum=0.3, maximum=4, latex_label=\"$\\\\sigma$\")\n",
+ "\n",
+ "priors[\"H0\"] = Uniform(minimum=20, maximum=200, latex_label=\"$H_0$\")\n",
+ "priors[\"Om0\"] = Uniform(minimum=0, maximum=1, latex_label=\"$\\\\Omega_{m,0}$\")\n",
+ "priors[\"w0\"] = Uniform(minimum=-1.5, maximum=-0.5, latex_label=\"$w_{0}$\")\n",
+ "priors[\"lamb\"] = Uniform(minimum=-1, maximum=10, latex_label=\"$\\\\lambda_{z}$\")"
+ ],
+ "metadata": {
+ "id": "IRf7FFQDGda-"
+ },
+ "execution_count": 7,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "source": [
+ "## Just-in-time compile\n",
+ "\n",
+ "We JIT compile the likelihood object before starting the sampler.\n",
+ "This is done using the `gwpopulation.experimental.jax.JittedLikelihood` class.\n",
+ "\n",
+ "We then time the original likelihood object and the JIT-ed version.\n",
+ "Note that we do two evaluations for each object as the first evaluation must compile the likelihood and so takes longer. (In addition to the JIT compilation, `JAX` compiles GPU functionality at the first evaluation, but this is less extreme than the full JIT compilation.)"
+ ],
+ "metadata": {
+ "id": "ncci_peFisbm"
+ }
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "parameters = priors.sample()\n",
+ "likelihood.parameters.update(parameters)\n",
+ "likelihood.log_likelihood_ratio()\n",
+ "%time print(likelihood.log_likelihood_ratio())\n",
+ "jit_likelihood = JittedLikelihood(likelihood)\n",
+ "jit_likelihood.parameters.update(parameters)\n",
+ "%time print(jit_likelihood.log_likelihood_ratio())\n",
+ "%time print(jit_likelihood.log_likelihood_ratio())"
+ ],
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "fgkofuRrGfX7",
+ "outputId": "990ef6f2-e36f-4fe7-9b7d-95b07ce39441"
+ },
+ "execution_count": 8,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "-350.4763359944699\n",
+ "CPU times: user 266 ms, sys: 70.4 ms, total: 337 ms\n",
+ "Wall time: 302 ms\n",
+ "-350.4763359944699\n",
+ "CPU times: user 5.73 s, sys: 311 ms, total: 6.05 s\n",
+ "Wall time: 5.83 s\n",
+ "-350.4763359944699\n",
+ "CPU times: user 4.86 ms, sys: 84 µs, total: 4.95 ms\n",
+ "Wall time: 19 ms\n"
+ ]
+ }
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "source": [
+ "## Run the sampler\n",
+ "\n",
+ "We'll use the sampler `dynesty` and use a small number of live points to reduce the runtime (total runtime should be approximately 5 minutes on T4 GPUs via Google colab).\n",
+ "The settings here may not give publication quality results, a convergence test should be performed before making strong quantitative statements.\n",
+ "\n",
+ "`bilby` times a single likelihood evaluation before beginning the run, however, this isn't well defined with JAX.\n",
+ "\n",
+ "**Note:** sometimes this finds a high likelihood mode, likely due to [breakdowns in the approximation](https://arxiv.org/abs/2304.06138) used to estimate the likelihood. If you see `dlogz > -80`, you should interrupt the execution and restart."
+ ],
+ "metadata": {
+ "id": "fX0FoqX0ixo0"
+ }
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "result = bb.run_sampler(\n",
+ " likelihood=jit_likelihood,\n",
+ " priors=priors,\n",
+ " sampler=\"dynesty\",\n",
+ " nlive=100,\n",
+ " label=\"cosmo\",\n",
+ " sample=\"acceptance-walk\",\n",
+ " naccept=5,\n",
+ " save=\"hdf5\",\n",
+ ")"
+ ],
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/",
+ "height": 763,
+ "referenced_widgets": [
+ "f301c06678c64a8baa6c68471f3c1b0b",
+ "75b24e49be914fdc95b18d1ab02ffbbb",
+ "a29a8eebf19a4e8fb6fea7237fc36205",
+ "edbc0b93226a466dae89cb9161df791a",
+ "1a3afc9c38cd414d80412fc4631eefb9",
+ "499aadfa6bb943608e8e813c36e1028e",
+ "ddaa2c6a652e40a0b6fef99f2b35c85c",
+ "3334f0e9517c41fead74795f5c19531e",
+ "cfa1cbe417564744ad94bf86d6d44f66",
+ "be29267ae6024e2c8e82f3c0f7ea4043",
+ "1707cba22e1e4ddc967a4d8a9e41075a"
+ ]
+ },
+ "id": "lHB2loW6Ghx-",
+ "outputId": "ceba1e1e-3be2-4195-8789-350ce0de9250"
+ },
+ "execution_count": 9,
+ "outputs": [
+ {
+ "metadata": {
+ "tags": null
+ },
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "16:01 bilby INFO : Running for label 'cosmo', output will be saved to 'outdir'\n",
+ "/usr/local/lib/python3.10/dist-packages/_distutils_hack/__init__.py:33: UserWarning: Setuptools is replacing distutils.\n",
+ " warnings.warn(\"Setuptools is replacing distutils.\")\n",
+ "16:01 bilby INFO : Analysis priors:\n",
+ "16:01 bilby INFO : alpha=Uniform(minimum=-2, maximum=4, name=None, latex_label='$\\\\alpha$', unit=None, boundary=None)\n",
+ "16:01 bilby INFO : beta=Uniform(minimum=-4, maximum=12, name=None, latex_label='$\\\\beta$', unit=None, boundary=None)\n",
+ "16:01 bilby INFO : mmin=Uniform(minimum=2, maximum=2.5, name=None, latex_label='$m_{\\\\min}$', unit=None, boundary=None)\n",
+ "16:01 bilby INFO : mmax=Uniform(minimum=80, maximum=100, name=None, latex_label='$m_{\\\\max}$', unit=None, boundary=None)\n",
+ "16:01 bilby INFO : lam=Uniform(minimum=0, maximum=1, name=None, latex_label='$\\\\lambda_{m}$', unit=None, boundary=None)\n",
+ "16:01 bilby INFO : mpp=Uniform(minimum=10, maximum=50, name=None, latex_label='$\\\\mu_{m}$', unit=None, boundary=None)\n",
+ "16:01 bilby INFO : sigpp=Uniform(minimum=1, maximum=10, name=None, latex_label='$\\\\sigma_{m}$', unit=None, boundary=None)\n",
+ "16:01 bilby INFO : alpha_chi=Uniform(minimum=1, maximum=6, name=None, latex_label='$\\\\alpha_{\\\\chi}$', unit=None, boundary=None)\n",
+ "16:01 bilby INFO : beta_chi=Uniform(minimum=1, maximum=6, name=None, latex_label='$\\\\beta_{\\\\chi}$', unit=None, boundary=None)\n",
+ "16:01 bilby INFO : xi_spin=Uniform(minimum=0, maximum=1, name=None, latex_label='$\\\\xi$', unit=None, boundary=None)\n",
+ "16:01 bilby INFO : sigma_spin=Uniform(minimum=0.3, maximum=4, name=None, latex_label='$\\\\sigma$', unit=None, boundary=None)\n",
+ "16:01 bilby INFO : H0=Uniform(minimum=20, maximum=200, name=None, latex_label='$H_0$', unit=None, boundary=None)\n",
+ "16:01 bilby INFO : Om0=Uniform(minimum=0, maximum=1, name=None, latex_label='$\\\\Omega_{m,0}$', unit=None, boundary=None)\n",
+ "16:01 bilby INFO : w0=Uniform(minimum=-1.5, maximum=-0.5, name=None, latex_label='$w_{0}$', unit=None, boundary=None)\n",
+ "16:01 bilby INFO : lamb=Uniform(minimum=-1, maximum=10, name=None, latex_label='$\\\\lambda_{z}$', unit=None, boundary=None)\n",
+ "16:01 bilby INFO : gaussian_mass_maximum=100\n",
+ "16:01 bilby INFO : amax=1\n",
+ "16:01 bilby INFO : Analysis likelihood class: | \n", + " | H0 | \n", + "Om0 | \n", + "alpha | \n", + "alpha_chi | \n", + "amax | \n", + "beta | \n", + "beta_chi | \n", + "gaussian_mass_maximum | \n", + "lam | \n", + "lamb | \n", + "... | \n", + "var_67 | \n", + "var_68 | \n", + "var_69 | \n", + "var_7 | \n", + "var_70 | \n", + "var_8 | \n", + "var_9 | \n", + "variance | \n", + "w0 | \n", + "xi_spin | \n", + "
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| count | \n", + "780.000000 | \n", + "780.000000 | \n", + "780.000000 | \n", + "780.000000 | \n", + "780.0 | \n", + "780.000000 | \n", + "780.000000 | \n", + "780.0 | \n", + "780.000000 | \n", + "780.000000 | \n", + "... | \n", + "780.000000 | \n", + "780.000000 | \n", + "780.000000 | \n", + "780.000000 | \n", + "780.000000 | \n", + "780.000000 | \n", + "780.000000 | \n", + "780.000000 | \n", + "780.000000 | \n", + "780.000000 | \n", + "
| mean | \n", + "87.540933 | \n", + "0.496407 | \n", + "2.804120 | \n", + "1.708842 | \n", + "1.0 | \n", + "1.920281 | \n", + "4.800247 | \n", + "100.0 | \n", + "0.069283 | \n", + "2.410931 | \n", + "... | \n", + "0.000886 | \n", + "0.001200 | \n", + "0.000478 | \n", + "0.000330 | \n", + "0.000755 | \n", + "0.000817 | \n", + "0.000934 | \n", + "1.447951 | \n", + "-0.975746 | \n", + "0.652503 | \n", + "
| std | \n", + "38.448334 | \n", + "0.257284 | \n", + "0.275604 | \n", + "0.387138 | \n", + "0.0 | \n", + "0.998621 | \n", + "0.816641 | \n", + "0.0 | \n", + "0.109628 | \n", + "0.937311 | \n", + "... | \n", + "0.000400 | \n", + "0.000631 | \n", + "0.000199 | \n", + "0.000128 | \n", + "0.000287 | \n", + "0.000335 | \n", + "0.000388 | \n", + "0.667548 | \n", + "0.272346 | \n", + "0.236770 | \n", + "
| min | \n", + "24.685531 | \n", + "0.013414 | \n", + "0.874584 | \n", + "1.007838 | \n", + "1.0 | \n", + "-0.526036 | \n", + "2.481648 | \n", + "100.0 | \n", + "0.005521 | \n", + "-0.615947 | \n", + "... | \n", + "0.000194 | \n", + "0.000254 | \n", + "0.000162 | \n", + "0.000127 | \n", + "0.000155 | \n", + "0.000277 | \n", + "0.000268 | \n", + "0.463470 | \n", + "-1.497994 | \n", + "0.000616 | \n", + "
| 25% | \n", + "59.168513 | \n", + "0.284552 | \n", + "2.640921 | \n", + "1.400418 | \n", + "1.0 | \n", + "1.239600 | \n", + "4.174589 | \n", + "100.0 | \n", + "0.021474 | \n", + "1.811519 | \n", + "... | \n", + "0.000626 | \n", + "0.000764 | \n", + "0.000349 | \n", + "0.000248 | \n", + "0.000546 | \n", + "0.000598 | \n", + "0.000678 | \n", + "0.995838 | \n", + "-1.195712 | \n", + "0.493420 | \n", + "
| 50% | \n", + "77.855050 | \n", + "0.495405 | \n", + "2.797660 | \n", + "1.713712 | \n", + "1.0 | \n", + "1.872876 | \n", + "4.902235 | \n", + "100.0 | \n", + "0.031892 | \n", + "2.390296 | \n", + "... | \n", + "0.000822 | \n", + "0.001077 | \n", + "0.000437 | \n", + "0.000305 | \n", + "0.000725 | \n", + "0.000758 | \n", + "0.000851 | \n", + "1.283603 | \n", + "-0.958729 | \n", + "0.695247 | \n", + "
| 75% | \n", + "108.452179 | \n", + "0.700400 | \n", + "2.968752 | \n", + "1.969656 | \n", + "1.0 | \n", + "2.543439 | \n", + "5.504270 | \n", + "100.0 | \n", + "0.056310 | \n", + "2.975050 | \n", + "... | \n", + "0.001052 | \n", + "0.001480 | \n", + "0.000543 | \n", + "0.000384 | \n", + "0.000906 | \n", + "0.000939 | \n", + "0.001102 | \n", + "1.710658 | \n", + "-0.754349 | \n", + "0.847084 | \n", + "
| max | \n", + "198.119995 | \n", + "0.999339 | \n", + "3.767904 | \n", + "2.872905 | \n", + "1.0 | \n", + "5.352800 | \n", + "5.997097 | \n", + "100.0 | \n", + "0.980254 | \n", + "5.810054 | \n", + "... | \n", + "0.005163 | \n", + "0.005151 | \n", + "0.002085 | \n", + "0.001093 | \n", + "0.002337 | \n", + "0.003297 | \n", + "0.003742 | \n", + "5.907584 | \n", + "-0.501113 | \n", + "0.999243 | \n", + "
8 rows × 164 columns
\n", + "| \n", + " | alpha | \n", + "beta | \n", + "mmin | \n", + "mmax | \n", + "lam | \n", + "mpp | \n", + "sigpp | \n", + "alpha_chi | \n", + "beta_chi | \n", + "xi_spin | \n", + "sigma_spin | \n", + "H0 | \n", + "Om0 | \n", + "w0 | \n", + "lamb | \n", + "log_likelihood | \n", + "variance | \n", + "
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| alpha | \n", + "1.000000 | \n", + "0.102840 | \n", + "0.150330 | \n", + "0.044006 | \n", + "-0.399092 | \n", + "0.029724 | \n", + "-0.028080 | \n", + "-0.042796 | \n", + "-0.045329 | \n", + "-0.061009 | \n", + "0.003468 | \n", + "-0.023973 | \n", + "0.038009 | \n", + "0.137577 | \n", + "0.510388 | \n", + "0.057235 | \n", + "0.095880 | \n", + "
| beta | \n", + "0.102840 | \n", + "1.000000 | \n", + "-0.052229 | \n", + "-0.102618 | \n", + "-0.117226 | \n", + "0.081481 | \n", + "-0.137837 | \n", + "-0.076178 | \n", + "-0.014593 | \n", + "-0.061626 | \n", + "-0.028563 | \n", + "-0.072214 | \n", + "0.107780 | \n", + "0.002857 | \n", + "-0.103693 | \n", + "0.093911 | \n", + "0.179609 | \n", + "
| mmin | \n", + "0.150330 | \n", + "-0.052229 | \n", + "1.000000 | \n", + "0.098832 | \n", + "-0.083318 | \n", + "0.076833 | \n", + "-0.033327 | \n", + "-0.067174 | \n", + "-0.046322 | \n", + "-0.048248 | \n", + "0.056417 | \n", + "-0.093888 | \n", + "0.056433 | \n", + "0.101209 | \n", + "0.014229 | \n", + "0.037321 | \n", + "-0.001154 | \n", + "
| mmax | \n", + "0.044006 | \n", + "-0.102618 | \n", + "0.098832 | \n", + "1.000000 | \n", + "0.171459 | \n", + "-0.191954 | \n", + "0.137365 | \n", + "-0.022086 | \n", + "-0.078806 | \n", + "0.093225 | \n", + "0.077121 | \n", + "0.162287 | \n", + "-0.027380 | \n", + "0.010086 | \n", + "0.093685 | \n", + "-0.259266 | \n", + "-0.130577 | \n", + "
| lam | \n", + "-0.399092 | \n", + "-0.117226 | \n", + "-0.083318 | \n", + "0.171459 | \n", + "1.000000 | \n", + "-0.768351 | \n", + "0.548668 | \n", + "0.045615 | \n", + "-0.020203 | \n", + "0.029992 | \n", + "0.091211 | \n", + "0.539933 | \n", + "-0.076569 | \n", + "-0.058193 | \n", + "-0.048806 | \n", + "-0.473377 | \n", + "-0.166637 | \n", + "
| mpp | \n", + "0.029724 | \n", + "0.081481 | \n", + "0.076833 | \n", + "-0.191954 | \n", + "-0.768351 | \n", + "1.000000 | \n", + "-0.680014 | \n", + "-0.073610 | \n", + "0.049858 | \n", + "0.019664 | \n", + "-0.072891 | \n", + "-0.757970 | \n", + "0.126399 | \n", + "-0.044086 | \n", + "-0.095138 | \n", + "0.586912 | \n", + "0.265031 | \n", + "
| sigpp | \n", + "-0.028080 | \n", + "-0.137837 | \n", + "-0.033327 | \n", + "0.137365 | \n", + "0.548668 | \n", + "-0.680014 | \n", + "1.000000 | \n", + "0.000565 | \n", + "-0.098283 | \n", + "-0.059140 | \n", + "0.058377 | \n", + "0.191383 | \n", + "-0.088490 | \n", + "0.011920 | \n", + "0.126626 | \n", + "-0.622818 | \n", + "-0.526754 | \n", + "
| alpha_chi | \n", + "-0.042796 | \n", + "-0.076178 | \n", + "-0.067174 | \n", + "-0.022086 | \n", + "0.045615 | \n", + "-0.073610 | \n", + "0.000565 | \n", + "1.000000 | \n", + "0.639341 | \n", + "-0.081536 | \n", + "0.166228 | \n", + "-0.003404 | \n", + "-0.051070 | \n", + "0.015877 | \n", + "0.036550 | \n", + "-0.091780 | \n", + "-0.096983 | \n", + "
| beta_chi | \n", + "-0.045329 | \n", + "-0.014593 | \n", + "-0.046322 | \n", + "-0.078806 | \n", + "-0.020203 | \n", + "0.049858 | \n", + "-0.098283 | \n", + "0.639341 | \n", + "1.000000 | \n", + "-0.006663 | \n", + "-0.083432 | \n", + "-0.093991 | \n", + "-0.015513 | \n", + "0.046653 | \n", + "0.011073 | \n", + "0.147665 | \n", + "0.381339 | \n", + "
| xi_spin | \n", + "-0.061009 | \n", + "-0.061626 | \n", + "-0.048248 | \n", + "0.093225 | \n", + "0.029992 | \n", + "0.019664 | \n", + "-0.059140 | \n", + "-0.081536 | \n", + "-0.006663 | \n", + "1.000000 | \n", + "0.118800 | \n", + "0.025346 | \n", + "0.105407 | \n", + "0.110062 | \n", + "-0.086955 | \n", + "0.144875 | \n", + "0.078301 | \n", + "
| sigma_spin | \n", + "0.003468 | \n", + "-0.028563 | \n", + "0.056417 | \n", + "0.077121 | \n", + "0.091211 | \n", + "-0.072891 | \n", + "0.058377 | \n", + "0.166228 | \n", + "-0.083432 | \n", + "0.118800 | \n", + "1.000000 | \n", + "0.062329 | \n", + "-0.041574 | \n", + "0.070895 | \n", + "0.077787 | \n", + "-0.328869 | \n", + "-0.285673 | \n", + "
| H0 | \n", + "-0.023973 | \n", + "-0.072214 | \n", + "-0.093888 | \n", + "0.162287 | \n", + "0.539933 | \n", + "-0.757970 | \n", + "0.191383 | \n", + "-0.003404 | \n", + "-0.093991 | \n", + "0.025346 | \n", + "0.062329 | \n", + "1.000000 | \n", + "-0.296857 | \n", + "0.012281 | \n", + "-0.066959 | \n", + "-0.410175 | \n", + "-0.066856 | \n", + "
| Om0 | \n", + "0.038009 | \n", + "0.107780 | \n", + "0.056433 | \n", + "-0.027380 | \n", + "-0.076569 | \n", + "0.126399 | \n", + "-0.088490 | \n", + "-0.051070 | \n", + "-0.015513 | \n", + "0.105407 | \n", + "-0.041574 | \n", + "-0.296857 | \n", + "1.000000 | \n", + "-0.077209 | \n", + "0.010703 | \n", + "0.083483 | \n", + "0.099786 | \n", + "
| w0 | \n", + "0.137577 | \n", + "0.002857 | \n", + "0.101209 | \n", + "0.010086 | \n", + "-0.058193 | \n", + "-0.044086 | \n", + "0.011920 | \n", + "0.015877 | \n", + "0.046653 | \n", + "0.110062 | \n", + "0.070895 | \n", + "0.012281 | \n", + "-0.077209 | \n", + "1.000000 | \n", + "0.069446 | \n", + "-0.017337 | \n", + "-0.052709 | \n", + "
| lamb | \n", + "0.510388 | \n", + "-0.103693 | \n", + "0.014229 | \n", + "0.093685 | \n", + "-0.048806 | \n", + "-0.095138 | \n", + "0.126626 | \n", + "0.036550 | \n", + "0.011073 | \n", + "-0.086955 | \n", + "0.077787 | \n", + "-0.066959 | \n", + "0.010703 | \n", + "0.069446 | \n", + "1.000000 | \n", + "-0.011215 | \n", + "-0.014705 | \n", + "
| log_likelihood | \n", + "0.057235 | \n", + "0.093911 | \n", + "0.037321 | \n", + "-0.259266 | \n", + "-0.473377 | \n", + "0.586912 | \n", + "-0.622818 | \n", + "-0.091780 | \n", + "0.147665 | \n", + "0.144875 | \n", + "-0.328869 | \n", + "-0.410175 | \n", + "0.083483 | \n", + "-0.017337 | \n", + "-0.011215 | \n", + "1.000000 | \n", + "0.414488 | \n", + "
| variance | \n", + "0.095880 | \n", + "0.179609 | \n", + "-0.001154 | \n", + "-0.130577 | \n", + "-0.166637 | \n", + "0.265031 | \n", + "-0.526754 | \n", + "-0.096983 | \n", + "0.381339 | \n", + "0.078301 | \n", + "-0.285673 | \n", + "-0.066856 | \n", + "0.099786 | \n", + "-0.052709 | \n", + "-0.014705 | \n", + "0.414488 | \n", + "1.000000 | \n", + "
"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "source": [
+ "## Compare `wcosmo` and `astropy` timing.\n",
+ "\n",
+ "The two primary pieces of functionality we use are converting from luminosity distance to redshift, calculating the distance to redshift jacobian, and calculating the differential comoving volume.\n",
+ "\n",
+ "Timing the `wcosmo` implementation is non trivial as we rely on JIT compilation and also need to make sure we wait until the evaluation is complete.\n",
+ "The steps are:\n",
+ "\n",
+ "- jit compile a wrapper function to call.\n",
+ "- burn an evaluation for the compilation.\n",
+ "- run the function and use `block_until_ready` to ensure we get the full time.\n",
+ "\n",
+ "We also time `wcosmo` with the `numpy` and `cupy` backends.\n",
+ "Note that `cupy` also requires burning a call to compile the underlying `CUDA` code.\n",
+ "\n",
+ "We manually switch backends, although this can be done automatically using `GWPopulation`."
+ ],
+ "metadata": {
+ "id": "nHKQkswiTuBQ"
+ }
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {
+ "id": "kcb71ObqYwC8"
+ },
+ "outputs": [],
+ "source": [
+ "!pip install wcosmo --quiet --progress-bar off"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "import numpy as np\n",
+ "import wcosmo\n",
+ "\n",
+ "\n",
+ "def set_backend(backend):\n",
+ " from importlib import import_module\n",
+ " np_modules = dict(\n",
+ " numpy=\"numpy\",\n",
+ " jax=\"jax.numpy\",\n",
+ " cupy=\"cupy\",\n",
+ " )\n",
+ " linalg_modules = dict(\n",
+ " numpy=\"scipy.linalg\",\n",
+ " jax=\"jax.scipy.linalg\",\n",
+ " cupy=\"cupyx.scipy.linalg\",\n",
+ " )\n",
+ " setattr(wcosmo.wcosmo, \"xp\", import_module(np_modules[backend]))\n",
+ " setattr(wcosmo.utils, \"xp\", import_module(np_modules[backend]))\n",
+ " toeplitz = getattr(import_module(linalg_modules[backend]), \"toeplitz\")\n",
+ " setattr(wcosmo.utils, \"toeplitz\", toeplitz)\n",
+ "\n",
+ "\n",
+ "ndata = np.random.uniform(1, 10, 1000000)"
+ ],
+ "metadata": {
+ "id": "z-o1RrCPdXRz"
+ },
+ "execution_count": 2,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "source": [
+ "### wcosmo + jax + GPU"
+ ],
+ "metadata": {
+ "id": "dseFs4ZVXXNu"
+ }
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "import jax.numpy as jnp\n",
+ "import numpy as np\n",
+ "from jax import jit\n",
+ "\n",
+ "\n",
+ "set_backend(\"jax\")\n",
+ "\n",
+ "jdata = jnp.array(ndata)\n",
+ "\n",
+ "\n",
+ "@jit\n",
+ "def time_jax_redshift(jdata):\n",
+ " return wcosmo.z_at_value(wcosmo.FlatwCDM(67, 0.3, -1).luminosity_distance, jdata)\n",
+ "\n",
+ "\n",
+ "@jit\n",
+ "def time_jax_dvcdz(jdata):\n",
+ " return wcosmo.FlatwCDM(67, 0.3, -1).differential_comoving_volume(jdata)\n",
+ "\n",
+ "\n",
+ "burn_vals = time_jax_redshift(jdata)\n",
+ "burn_vals = time_jax_dvcdz(jdata)"
+ ],
+ "metadata": {
+ "id": "IY_u2ZI3PRua"
+ },
+ "execution_count": 3,
+ "outputs": []
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "%%time\n",
+ "\n",
+ "_ = time_jax_redshift(jdata).block_until_ready()"
+ ],
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "xjppj7LCPU_B",
+ "outputId": "fb0c6ff9-96e6-4771-cb6b-67f1d5e40b85"
+ },
+ "execution_count": 4,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "CPU times: user 1.56 ms, sys: 0 ns, total: 1.56 ms\n",
+ "Wall time: 4.83 ms\n"
+ ]
+ }
+ ]
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "%%time\n",
+ "\n",
+ "_ = time_jax_dvcdz(jdata).block_until_ready()"
+ ],
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "JgS3_9RVU1nH",
+ "outputId": "ce2f0c85-0161-4f45-9cd7-3494077e7a83"
+ },
+ "execution_count": 5,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "CPU times: user 632 µs, sys: 173 µs, total: 805 µs\n",
+ "Wall time: 3.69 ms\n"
+ ]
+ }
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "source": [
+ "### astropy + cpu\n",
+ "\n",
+ "Note that this is very slow in this case so we only use one percent of the full data.\n",
+ "Since this is `numpy`-based the time scales linearly with the amount of data.\n",
+ "\n",
+ "In practice, most people when using `astropy` use intepolation to evaluate `z_at_value` at many points as is done in `wcosmo`."
+ ],
+ "metadata": {
+ "id": "ozTDHo9iXgWi"
+ }
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "from astropy import cosmology, units"
+ ],
+ "metadata": {
+ "id": "_kOnSJMbXn9Q"
+ },
+ "execution_count": 6,
+ "outputs": []
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "%%time\n",
+ "\n",
+ "_ = cosmology.z_at_value(\n",
+ " cosmology.FlatwCDM(67, 0.3, -1).luminosity_distance,\n",
+ " ndata[:10000] * units.Mpc,\n",
+ ").value"
+ ],
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "V-mMVWAhPgfp",
+ "outputId": "2272fbab-0bdb-4cad-be8d-65bb6ecc249e"
+ },
+ "execution_count": 7,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "CPU times: user 40.9 s, sys: 418 ms, total: 41.3 s\n",
+ "Wall time: 52 s\n"
+ ]
+ }
+ ]
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "%%time\n",
+ "\n",
+ "_ = cosmology.FlatwCDM(67, 0.3, -1).differential_comoving_volume(\n",
+ " ndata[:10000],\n",
+ ").value"
+ ],
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "vF6Ab5jMVBdX",
+ "outputId": "3f176351-bd05-4a09-f665-183cecb2e105"
+ },
+ "execution_count": 8,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "CPU times: user 106 ms, sys: 27 µs, total: 106 ms\n",
+ "Wall time: 107 ms\n"
+ ]
+ }
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "source": [
+ "### wcosmo + numpy + cpu"
+ ],
+ "metadata": {
+ "id": "3KTFoNDkXzTp"
+ }
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "set_backend(\"numpy\")"
+ ],
+ "metadata": {
+ "id": "qGl50ks_PpGQ"
+ },
+ "execution_count": 9,
+ "outputs": []
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "%%time\n",
+ "\n",
+ "_ = wcosmo.z_at_value(\n",
+ " wcosmo.FlatwCDM(67, 0.3, -1).luminosity_distance, ndata\n",
+ ")"
+ ],
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "bQl0yn4uVSXp",
+ "outputId": "6901b6ba-de14-43b0-dca7-6909127dfed9"
+ },
+ "execution_count": 10,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "CPU times: user 89.5 ms, sys: 75.2 ms, total: 165 ms\n",
+ "Wall time: 86.7 ms\n"
+ ]
+ }
+ ]
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "%%time\n",
+ "\n",
+ "_ = wcosmo.FlatwCDM(67, 0.3, -1).differential_comoving_volume(ndata)"
+ ],
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "XjsGLM2sVXC-",
+ "outputId": "1fc8817c-a0a1-44fa-dcf0-05b8f724b76f"
+ },
+ "execution_count": 11,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "CPU times: user 65.8 ms, sys: 89.1 ms, total: 155 ms\n",
+ "Wall time: 79.4 ms\n"
+ ]
+ }
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "source": [
+ "### wcosmo + cupy + gpu\n",
+ "\n",
+ "The final test is using the `cupy` backend on the GPU.\n",
+ "Typically this is much faster than `numpy` but slower than the `JAX` GPU code.\n",
+ "Especially, not tested here is transfer between CPU/GPU which can be quite slow for `cupy`."
+ ],
+ "metadata": {
+ "id": "Vn27QdGSYC3Q"
+ }
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "import cupy\n",
+ "\n",
+ "set_backend(\"cupy\")\n",
+ "\n",
+ "cdata = cupy.asarray(ndata)\n",
+ "\n",
+ "_ = wcosmo.z_at_value(\n",
+ " wcosmo.FlatwCDM(67, 0.3, -1).luminosity_distance, cdata\n",
+ ")\n",
+ "_ = wcosmo.FlatwCDM(67, 0.3, -1).differential_comoving_volume(cdata)"
+ ],
+ "metadata": {
+ "id": "wU3XF5PpVbwn"
+ },
+ "execution_count": 12,
+ "outputs": []
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "%%time\n",
+ "\n",
+ "_ = wcosmo.z_at_value(\n",
+ " wcosmo.FlatwCDM(67, 0.3, -1).luminosity_distance, cdata\n",
+ ")\n",
+ "cupy.cuda.stream.get_current_stream().synchronize()"
+ ],
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "Q4LUkMDuWDsk",
+ "outputId": "a874883f-54e6-41e0-c8fc-50ce7a05039e"
+ },
+ "execution_count": 13,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "CPU times: user 5.42 ms, sys: 17 µs, total: 5.44 ms\n",
+ "Wall time: 5.45 ms\n"
+ ]
+ }
+ ]
+ },
+ {
+ "cell_type": "code",
+ "source": [
+ "%%time\n",
+ "\n",
+ "_ = wcosmo.FlatwCDM(67, 0.3, -1).differential_comoving_volume(cdata)\n",
+ "cupy.cuda.stream.get_current_stream().synchronize()"
+ ],
+ "metadata": {
+ "colab": {
+ "base_uri": "https://localhost:8080/"
+ },
+ "id": "BCr3UfbaWEgJ",
+ "outputId": "fefccd5c-5683-4e66-ecb0-d6635d1e5f80"
+ },
+ "execution_count": 14,
+ "outputs": [
+ {
+ "output_type": "stream",
+ "name": "stdout",
+ "text": [
+ "CPU times: user 73.6 ms, sys: 0 ns, total: 73.6 ms\n",
+ "Wall time: 83.6 ms\n"
+ ]
+ }
+ ]
+ },
+ {
+ "cell_type": "code",
+ "source": [],
+ "metadata": {
+ "id": "TjD_9bTFboru"
+ },
+ "execution_count": null,
+ "outputs": []
+ }
+ ]
+}
\ No newline at end of file