Examples

[ColmTalbot]

Add examples demonstrating

• Comparison with approximation in Adachi & Kasai 2011
• Comparison of pre-specified models with astropy
• Comparison across parameter space with astropy
• Demonstration of using jax/numpy backends
• other things?

A decent amount of this is already being done in #1

[afarah18]

It might also be good to do a bare-bones example of cosmological inference with this on a GPU, since that's the main reason we developed this.

[ColmTalbot]

Good idea, I can probably modify my notebook from last week.

[ColmTalbot]

See #7

[Qazalbash]

Hi, I am struggling to select the backend. I used the set_backend function mentioned in the example, but it uses numpy while I specified jax.

import numpy as np
import wcosmo


def set_backend(backend):
    from importlib import import_module
    np_modules = dict(
        numpy="numpy",
        jax="jax.numpy",
        cupy="cupy",
    )
    linalg_modules = dict(
        numpy="scipy.linalg",
        jax="jax.scipy.linalg",
        cupy="cupyx.scipy.linalg",
    )
    setattr(wcosmo.wcosmo, "xp", import_module(np_modules[backend]))
    setattr(wcosmo.utils, "xp", import_module(np_modules[backend]))
    toeplitz = getattr(import_module(linalg_modules[backend]), "toeplitz")
    setattr(wcosmo.utils, "toeplitz", toeplitz)

[ColmTalbot]

I suspect this example is out of date, there are now additional submodules that need to have these updated. The easiest solution if you're happy with some additional dependencies is to install gwpopulation and use gwpopulation.set_backend. As a temporary measure, the full set that need to be updated is specified in the pyproject.toml. The gwpopulation code to do this is here, we could quite easily adapt something like this, but I'm concerned about code duplication.

[project.entry-points."gwpopulation.xp"]
wcosmo = "wcosmo.wcosmo"
wcosmo-analytic = "wcosmo.analytic"
wcosmo-astropy = "wcosmo.astropy"
wcosmo-integrate = "wcosmo.integrate"
wcosmo-taylor = "wcosmo.taylor"
wcosmo-utils = "wcosmo.utils"

[project.entry-points."gwpopulation.scs"]
wcosmo-analytic = "wcosmo.analytic"

[project.entry-points."gwpopulation.other"]
wcosmo-taylor = "wcosmo.taylor:scipy.linalg.toeplitz"

[Qazalbash]

Can you also guide me on how to use units inside jit/grad/vmap transforms?