DRAFT: Support non-numpy array backends

.github/workflows/basic-install.yml

@@ -19,8 +19,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
@@ -48,8 +47,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: ${{ matrix.os != "windows-latest" }}
+        run: |
+          for script in $(pip show -f bilby | grep "bin\/" | xargs -I {} basename {}); do
+              ${script} --help;
+          done

bilby/compat/jax.py

@@ -0,0 +1,73 @@
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+from ..core.likelihood import Likelihood
+
+
+def generic_bilby_likelihood_function(likelihood, parameters, use_ratio=True):
+    """
+    A wrapper to allow a :code:`Bilby` likelihood to be used with :code:`jax`.
+
+    Parameters
+    ==========
+    likelihood: bilby.core.likelihood.Likelihood
+        The likelihood to evaluate.
+    parameters: dict
+        The parameters to evaluate the likelihood at.
+    use_ratio: bool, optional
+        Whether to evaluate the likelihood ratio or the full likelihood.
+        Default is :code:`True`.
+    """
+    parameters = {k: jnp.array(v) for k, v in parameters.items()}
+    likelihood.parameters.update(parameters)
+    if use_ratio:
+        return likelihood.log_likelihood_ratio()
+    else:
+        return likelihood.log_likelihood()
+
+
+class JittedLikelihood(Likelihood):
+    """
+    A wrapper to just-in-time compile a :code:`Bilby` likelihood for use with :code:`jax`.
+
+    .. note::
+
+        This is currently hardcoded to return the log likelihood ratio, regardless of
+        the input.
+
+    Parameters
+    ==========
+    likelihood: bilby.core.likelihood.Likelihood
+        The likelihood to wrap.
+    likelihood_func: callable, optional
+        The function to use to evaluate the likelihood. Default is
+        :code:`generic_bilby_likelihood_function`. This function should take the
+        likelihood and parameters as arguments along with additional keyword arguments.
+    kwargs: dict, optional
+        Additional keyword arguments to pass to the likelihood function.
+    """
+
+    def __init__(
+        self,
+        likelihood,
+        likelihood_func=generic_bilby_likelihood_function,
+        kwargs=None,
+        cast_to_float=True,
+    ):
+        if kwargs is None:
+            kwargs = dict()
+        self.kwargs = kwargs
+        self._likelihood = likelihood
+        self.likelihood_func = jax.jit(partial(likelihood_func, likelihood))
+        self.cast_to_float = cast_to_float
+        super().__init__(dict())
+
+    def __getattr__(self, name):
+        return getattr(self._likelihood, name)
+
+    def log_likelihood_ratio(self):
+        ln_l = jnp.nan_to_num(self.likelihood_func(self.parameters, **self.kwargs))
+        if self.cast_to_float:
+            ln_l = float(ln_l)
+        return ln_l

bilby/compat/types.py

@@ -0,0 +1,6 @@
+from typing import Union
+import numpy as np
+
+Real = Union[float, int]
+ArrayLike = Union[np.ndarray, list, tuple]
+

bilby/compat/utils.py

@@ -0,0 +1,34 @@
+import numpy as np
+from scipy._lib._array_api import array_namespace
+
+__all__ = ["array_module", "promote_to_array"]
+
+
+def array_module(arr):
+    if arr.__class__.__module__ == "builtins":
+        return np
+    else:
+        return array_namespace(arr)
+
+
+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):
+
+    def wrapped(self, *args, **kwargs):
+        if "xp" not in kwargs:
+            try:
+                kwargs["xp"] = array_module(*args)
+            except TypeError:
+                pass
+        return func(self, *args, **kwargs)
+
+    return wrapped

bilby/core/likelihood.py

@@ -5,6 +5,7 @@
 from scipy.stats import multivariate_normal
 
 from .utils import infer_parameters_from_function, infer_args_from_function_except_n_args
+from ..compat.utils import array_module
 
 
 class Likelihood(object):
@@ -465,10 +466,16 @@ class AnalyticalMultidimensionalCovariantGaussian(Likelihood):
         """
 
     def __init__(self, mean, cov):
-        self.cov = np.atleast_2d(cov)
-        self.mean = np.atleast_1d(mean)
-        self.sigma = np.sqrt(np.diag(self.cov))
-        self.pdf = multivariate_normal(mean=self.mean, cov=self.cov)
+        xp = array_module(cov)
+        self.cov = xp.atleast_2d(cov)
+        self.mean = xp.atleast_1d(mean)
+        self.sigma = xp.sqrt(np.diag(self.cov))
+        if xp == np:
+            self.logpdf = multivariate_normal(mean=self.mean, cov=self.cov).logpdf
+        else:
+            from functools import partial
+            from jax.scipy.stats.multivariate_normal import logpdf
+            self.logpdf = partial(logpdf, mean=self.mean, cov=self.cov)
         parameters = {"x{0}".format(i): 0 for i in range(self.dim)}
         super(AnalyticalMultidimensionalCovariantGaussian, self).__init__(parameters=parameters)
 
@@ -477,8 +484,9 @@ def dim(self):
         return len(self.cov[0])
 
     def log_likelihood(self):
-        x = np.array([self.parameters["x{0}".format(i)] for i in range(self.dim)])
-        return self.pdf.logpdf(x)
+        xp = array_module(self.cov)
+        x = xp.array([self.parameters["x{0}".format(i)] for i in range(self.dim)])
+        return self.logpdf(x)
 
 
 class AnalyticalMultidimensionalBimodalCovariantGaussian(Likelihood):
@@ -496,12 +504,19 @@ class AnalyticalMultidimensionalBimodalCovariantGaussian(Likelihood):
         """
 
     def __init__(self, mean_1, mean_2, cov):
-        self.cov = np.atleast_2d(cov)
-        self.sigma = np.sqrt(np.diag(self.cov))
-        self.mean_1 = np.atleast_1d(mean_1)
-        self.mean_2 = np.atleast_1d(mean_2)
-        self.pdf_1 = multivariate_normal(mean=self.mean_1, cov=self.cov)
-        self.pdf_2 = multivariate_normal(mean=self.mean_2, cov=self.cov)
+        xp = array_module(cov)
+        self.cov = xp.atleast_2d(cov)
+        self.sigma = xp.sqrt(np.diag(self.cov))
+        self.mean_1 = xp.atleast_1d(mean_1)
+        self.mean_2 = xp.atleast_1d(mean_2)
+        if xp == np:
+            self.logpdf_1 = multivariate_normal(mean=self.mean_1, cov=self.cov).logpdf
+            self.logpdf_2 = multivariate_normal(mean=self.mean_2, cov=self.cov).logpdf
+        else:
+            from functools import partial
+            from jax.scipy.stats.multivariate_normal import logpdf
+            self.logpdf_1 = partial(logpdf, mean=self.mean_1, cov=self.cov)
+            self.logpdf_2 = partial(logpdf, mean=self.mean_2, cov=self.cov)
         parameters = {"x{0}".format(i): 0 for i in range(self.dim)}
         super(AnalyticalMultidimensionalBimodalCovariantGaussian, self).__init__(parameters=parameters)
 
@@ -510,8 +525,9 @@ def dim(self):
         return len(self.cov[0])
 
     def log_likelihood(self):
-        x = np.array([self.parameters["x{0}".format(i)] for i in range(self.dim)])
-        return -np.log(2) + np.logaddexp(self.pdf_1.logpdf(x), self.pdf_2.logpdf(x))
+        xp = array_module(self.cov)
+        x = xp.array([self.parameters["x{0}".format(i)] for i in range(self.dim)])
+        return -xp.log(2) + xp.logaddexp(self.logpdf_1(x), self.logpdf_2(x))
 
 
 class JointLikelihood(Likelihood):