Support non-numpy array backends

.github/workflows/basic-install.yml

@@ -20,8 +20,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        # disable windows build test as bilby_cython is currently broken there
-        os: [ubuntu-latest, macos-latest]
+        os: [ubuntu-latest, macos-latest, windows-latest]
         python-version: ["3.10", "3.11", "3.12"]
     steps:
       - uses: actions/checkout@v4
@@ -52,8 +51,8 @@ jobs:
           python -c "import bilby.hyper"
           python -c "import cli_bilby"
           python test/import_test.py
-      # - if: ${{ matrix.os != "windows-latest" }}
-      #   run: |
-      #     for script in $(pip show -f bilby | grep "bin\/" | xargs -I {} basename {}); do
-      #         ${script} --help;
-      #     done
+      - if: runner.os != 'Windows'
+        run: |
+          for script in $(pip show -f bilby | grep "bin\/" | xargs -I {} basename {}); do
+              ${script} --help;
+          done

.github/workflows/unit-tests.yml

@@ -62,6 +62,15 @@ jobs:
     - name: Run unit tests
       run: |
         python -m pytest --cov=bilby --cov-branch --durations 10 -ra --color yes --cov-report=xml --junitxml=pytest.xml
+    - name: Run jax-backend unit tests
+      run: |
+        python -m pip install .[jax]
+        SCIPY_ARRAY_API=1 pytest --array-backend jax --durations 10
+    - name: Run torch-backend unit tests
+      # there are scipy version issues with python 3.10 and torch
+      if: matrix.python.version > 3.10
+      run: |
+        SCIPY_ARRAY_API=1 pytest --array-backend torch --durations 10
     - name: Run sampler tests
       run: |
         pytest test/integration/sampler_run_test.py --durations 10 -v        

bilby/compat/jax.py

@@ -0,0 +1,40 @@
+import jax
+import jax.numpy as jnp
+from ..core.likelihood import Likelihood
+
+
+class JittedLikelihood(Likelihood):
+    """
+    A wrapper to just-in-time compile a :code:`Bilby` likelihood for use with :code:`jax`.
+
+    Parameters
+    ==========
+    likelihood: bilby.core.likelihood.Likelihood
+        The likelihood to wrap.
+    cast_to_float: bool
+        Whether to return a float instead of a :code:`jax.Array`.
+    """
+
+    def __init__(self, likelihood, cast_to_float=True):
+        self._likelihood = likelihood
+        self._ll = jax.jit(likelihood.log_likelihood)
+        self._llr = jax.jit(likelihood.log_likelihood_ratio)
+        self.cast_to_float = cast_to_float
+        super().__init__()
+
+    def __getattr__(self, name):
+        return getattr(self._likelihood, name)
+
+    def log_likelihood(self, parameters):
+        parameters = {k: jnp.array(v) for k, v in parameters.items()}
+        ln_l = self._ll(parameters)
+        if self.cast_to_float:
+            ln_l = float(ln_l)
+        return ln_l
+
+    def log_likelihood_ratio(self, parameters):
+        parameters = {k: jnp.array(v) for k, v in parameters.items()}
+        ln_l = self._llr(parameters)
+        if self.cast_to_float:
+            ln_l = float(ln_l)
+        return ln_l

bilby/compat/patches.py

@@ -0,0 +1,50 @@
+import array_api_compat as aac
+
+from .utils import BackendNotImplementedError
+
+
+def erfinv_import(xp):
+    if aac.is_numpy_namespace(xp):
+        from scipy.special import erfinv
+    elif aac.is_jax_namespace(xp):
+        from jax.scipy.special import erfinv
+    elif aac.is_torch_namespace(xp):
+        from torch.special import erfinv
+    elif aac.is_cupy_namespace(xp):
+        from cupyx.scipy.special import erfinv
+    else:
+        raise BackendNotImplementedError
+    return erfinv
+
+
+def multivariate_logpdf(xp, mean, cov):
+    if aac.is_numpy_namespace(xp):
+        from scipy.stats import multivariate_normal
+
+        logpdf = multivariate_normal(mean=mean, cov=cov).logpdf
+    elif aac.is_jax_namespace(xp):
+        from functools import partial
+        from jax.scipy.stats.multivariate_normal import logpdf
+
+        logpdf = partial(logpdf, mean=mean, cov=cov)
+    elif aac.is_torch_namespace(xp):
+        from torch.distributions.multivariate_normal import MultivariateNormal
+
+        mvn = MultivariateNormal(loc=mean, covariance_matrix=xp.asarray(cov))
+        logpdf = mvn.log_prob
+    else:
+        raise BackendNotImplementedError
+    return logpdf
+
+
+def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False, *, xp=None):
+    if xp is None:
+        xp = aac.get_namespace(a)
+
+    # the scipy version of logsumexp cannot be vmapped
+    if aac.is_jax_namespace(xp):
+        from jax.scipy.special import logsumexp as lse
+    else:
+        from scipy.special import logsumexp as lse
+
+    return lse(a=a, axis=axis, b=b, keepdims=keepdims, return_sign=return_sign)

bilby/compat/types.py

@@ -0,0 +1,4 @@
+import numpy as np
+
+Real = float | int
+ArrayLike = np.ndarray | list | tuple

bilby/compat/utils.py

@@ -0,0 +1,105 @@
+import inspect
+from collections.abc import Iterable
+
+import numpy as np
+from array_api_compat import array_namespace
+
+from ..core.utils.log import logger
+
+__all__ = ["array_module", "promote_to_array"]
+
+
+def array_module(arr):
+    if isinstance(arr, tuple) and len(arr) == 1:
+        arr = arr[0]
+    try:
+        return array_namespace(arr)
+    except TypeError:
+        if isinstance(arr, dict):
+            try:
+                return array_namespace(*[val for val in arr.values() if not isinstance(val, str)])
+            except TypeError:
+                return np
+        elif arr.__class__.__module__ == "builtins" and isinstance(arr, Iterable):
+            try:
+                return array_namespace(*arr)
+            except TypeError:
+                return np
+        elif arr.__class__.__module__ == "builtins":
+            return np
+        elif arr.__module__.startswith("pandas"):
+            return np
+        else:
+            logger.warning(
+                f"Unknown array module for type: {type(arr)} Defaulting to numpy."
+            )
+            return np
+
+
+def promote_to_array(args, backend, skip=None):
+    if skip is None:
+        skip = len(args)
+    else:
+        skip = len(args) - skip
+    if backend.__name__ != "numpy":
+        args = tuple(backend.array(arg) for arg in args[:skip]) + args[skip:]
+    return args
+
+
+def xp_wrap(func, no_xp=False):
+    """
+    A decorator that will figure out the array module from the input
+    arguments and pass it to the function as the 'xp' keyword argument.
+
+    Parameters
+    ==========
+    func: function
+        The function to be decorated.
+    no_xp: bool
+        If True, the decorator will not attempt to add the 'xp' keyword
+        argument and so the wrapper is a no-op.
+
+    Returns
+    =======
+    function
+        The decorated function.
+    """
+    def parse_args_kwargs_for_xp(*args, xp=None, **kwargs):
+        if not no_xp and xp is None:
+            try:
+                # if the user specified the target arrays in kwargs
+                # we need to be able to support this, if there is
+                # only one kwargs, pass it through alone, this is
+                # sometimes a dictionary of arrays so this is needed
+                # to remove a level of nesting
+                if len(args) > 0:
+                    xp = array_module(args)
+                elif len(kwargs) == 1:
+                    xp = array_module(next(iter(kwargs.values())))
+                elif len(kwargs) > 1:
+                    xp = array_module(kwargs)
+                else:
+                    xp = np
+                kwargs["xp"] = xp
+            except TypeError as e:
+                print("type failed", e)
+                kwargs["xp"] = np
+        elif not no_xp:
+            kwargs["xp"] = xp
+        return args, kwargs
+
+    sig = inspect.signature(func)
+    if any(name in sig.parameters for name in ("self", "cls")):
+        def wrapped(self, *args, xp=None, **kwargs):
+            args, kwargs = parse_args_kwargs_for_xp(*args, xp=xp, **kwargs)
+            return func(self, *args, **kwargs)
+    else:
+        def wrapped(*args, xp=None, **kwargs):
+            args, kwargs = parse_args_kwargs_for_xp(*args, xp=xp, **kwargs)
+            return func(*args, **kwargs)
+
+    return wrapped
+
+
+class BackendNotImplementedError(NotImplementedError):
+    pass