Support non-numpy array backends

.github/workflows/basic-install.yml

@@ -20,9 +20,8 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        # disable windows build test as bilby_cython is currently broken there
-        os: [ubuntu-latest, macos-latest]
-        python-version: ["3.10", "3.11", "3.12", "3.13"]
+        os: [ubuntu-latest, macos-latest, windows-latest]
+        python-version: ["3.11", "3.12", "3.13"]
     steps:
       - uses: actions/checkout@v4
         with:
@@ -39,6 +38,7 @@ jobs:
       - name: Test imports
         run: bash test/ci_test_imports.sh
       - name: Test entry points
+        if: matrix.os != 'windows-latest'
         run: |
           for script in $(pip show -f bilby | grep "bin\/" | xargs -I {} basename {}); do
               ${script} --help;

.github/workflows/unit-tests.yml

@@ -32,9 +32,6 @@ jobs:
       fail-fast: false
       matrix:
         python:
-          - name: Python 3.10
-            version: 3.10
-            short-version: 310
           - name: Python 3.11
             version: 3.11
             short-version: 311
@@ -84,3 +81,59 @@ jobs:
       with:
         name: pytest-${{ matrix.python.short-version }}
         path: pytest.xml
+
+  array-backend:
+
+    name: ${{ matrix.python.name }} array backend (${{ matrix.backend.name }})
+    runs-on: ubuntu-latest
+    container: ghcr.io/bilby-dev/bilby-python${{ matrix.python.short-version }}:latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python:
+          - name: Python 3.13
+            version: 3.13
+            short-version: 313
+        backend:
+          - name: numpy
+            install-args: .
+          - name: jax
+            install-args: .[jax]
+          - name: torch
+            install-args: .
+    env:
+      BILBY_ARRAY_API: 1
+      SCIPY_ARRAY_API: 1
+
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        fetch-depth: 0
+        fetch-tags: true
+    - name: Install package
+      run: |
+        # activate env so that conda list shows the correct environment
+        apt-get update
+        apt-get install -y gnupg curl
+        source $CONDA_PATH/bin/activate python${{ matrix.python.short-version }}
+        conda install -c conda-forge -y liblal!=7.7.0 python-lal!=7.7.0
+        python -m pip install ${{ matrix.backend.install-args }}
+        python -m pip install orng
+        conda list --show-channel-urls
+      shell: bash
+    - name: Run array-backend unit tests
+      run: |
+        pytest --array-backend ${{ matrix.backend.name }} --durations 10 --junitxml=pytest-array.xml
+    - name: Publish coverage to Codecov
+      uses: codecov/codecov-action@v5
+      with:
+        token: ${{ secrets.CODECOV_TOKEN }}
+        files: coverage.xml
+        flags: python${{ matrix.backend.name }}
+        slug: bilby-dev/bilby
+    - name: Upload array-backend test results
+      if: always()
+      uses: actions/upload-artifact@v4
+      with:
+        name: pytest-array-${{ matrix.python.short-version }}-${{ matrix.backend.name }}
+        path: pytest-array.xml

.gitignore

@@ -16,3 +16,4 @@ MANIFEST
 **/outdir
 .idea/*
 bilby/_version.py
+uv.lock

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,104 @@
+import array_api_compat as aac
+
+from .utils import xp_wrap, BackendNotImplementedError, BILBY_ARRAY_API
+
+
+def multivariate_logpdf(xp, mean, cov):
+    """
+    Return a function to evaluate the log probability density of a multivariate
+    Gaussian with given mean vector and covariance matrix for the provided
+    array backend.
+
+    Parameters
+    ==========
+    xp: numpy, torch, jax.numpy
+        A module that will resolve to :code:`numpy`, :code:`torch`, or
+        :code:`jax.numpy` in :code:`array_api_compat.is_..._namespace`.
+    mean: array-like
+        A one-dimensional array providing the mean of the distribution.
+    cov: array-like
+        A two-dimensional array providing the covariance matrix of the
+        distribution.
+
+    Returns
+    =======
+    logpdf: callable
+        A callable that provides the log probaility density provided an array
+        of points to evaluate at.
+    """
+    from scipy.stats import multivariate_normal
+
+    if not BILBY_ARRAY_API or aac.is_numpy_namespace(xp):
+        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(
+            f"Unable to import multivariate_logpdf for {xp}"
+        )
+    return logpdf
+
+
+@xp_wrap
+def interp(x, xs, fs, /, left=None, right=None, period=None, *, xp=None):
+    """
+    A simple implementation of numpy-style linear interpolation
+
+    The logic is copied from
+    https://github.com/pytorch/pytorch/issues/50334#issuecomment-1000917964
+
+    Parameters
+    ==========
+    x: array-like
+        The values to evaluate the interpolant at.
+    xs: array-like
+        The x-values for setting up the interpolant.
+    ys: array-like
+        The values of the function for setting up the interpolant.
+    left: float
+        The value to use for x < xs[0]. Default is fs[0]
+    right: float
+        The value to use for x > xs[-1]. Default is fs[-1].
+    period: float
+        The period of the interpolant.
+        Parameters left and right are ignored if period is specified.
+
+    Notes
+    =====
+    To avoid overlap with the ``xp`` variable, the second and third variable
+    names from differ from numpy.
+    These arguments are enforced to be positional only.
+    """
+    if not BILBY_ARRAY_API or hasattr(xp, "interp"):
+        return xp.interp(x, xs, fs, left=left, right=right, period=period)
+
+    if period is not None:
+        x = x % period
+    if left is None:
+        left = fs[0]
+    if right is None:
+        right = fs[-1]
+
+    x = xp.atleast_1d(x)
+
+    m = (fs[1:] - fs[:-1]) / (xs[1:] - xs[:-1])
+    b = fs[:-1] - (m * xs[:-1])
+
+    indices = xp.sum(xp.ge(x[:, None], xs[None, :]), axis=1) - 1
+    indices = xp.clip(indices, 0, len(m) - 1)
+
+    ret = m[indices] * x + b[indices]
+
+    if period is None:
+        ret = xp.where(x < xs[0], xp.asarray(left), ret)
+        ret = xp.where(x > xs[-1], xp.asarray(right), ret)
+
+    return ret.squeeze()

bilby/compat/types.py

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