Added metropolis hastings sampling for jax

Author: FlorianPfaff

pyrecest/distributions/abstract_manifold_specific_distribution.py

@@ -5,7 +5,7 @@
 import pyrecest.backend
 
 # pylint: disable=no-name-in-module,no-member
-from pyrecest.backend import empty, int32, int64, log, random, squeeze
+from pyrecest.backend import empty, int32, int64, log, random, squeeze, where
 
 
 class AbstractManifoldSpecificDistribution(ABC):
@@ -81,9 +81,15 @@ def sample_metropolis_hastings(
     ):
         # jscpd:ignore-end
         """Metropolis Hastings sampling algorithm."""
-        assert (
-            pyrecest.backend.__backend_name__ != "jax"
-        ), "Not supported on this backend"
+        if pyrecest.backend.__backend_name__ == "jax":
+            return sample_metropolis_hastings_jax(
+                log_pdf=self.ln_pdf,
+                proposal=proposal,
+                start_point=start_point,
+                n=n,
+                burn_in=burn_in,
+                skipping=skipping,
+            )
         if proposal is None or start_point is None:
             raise NotImplementedError(
                 "Default proposals and starting points should be set in inheriting classes."
@@ -115,3 +121,54 @@ def sample_metropolis_hastings(
 
         relevant_samples = s[burn_in::skipping, :]
         return squeeze(relevant_samples)
+
+
+def sample_metropolis_hastings_jax(
+    log_pdf,               # function: R^d -> scalar log density
+    proposal,              # function: (key, x) -> x_proposed
+    start_point,
+    n: int,
+    burn_in: int = 10,
+    skipping: int = 5,
+):
+    """
+    Metropolis-Hastings sampler in JAX, jittable and vectorizable.
+
+    Args:
+        key: jax.random.PRNGKey
+        log_pdf: callable taking an array x and returning log p(x)
+        proposal: callable taking (key, x) and returning a proposed x'
+        start_point: initial state x_0 (array)
+        n: number of samples to return (after burn-in and thinning)
+        burn_in: number of initial states to discard
+        skipping: thinning factor (keep every `skipping`-th state)
+    """
+    from jax import lax
+
+    total_steps = burn_in + n * skipping
+
+    def one_step(carry, _):
+        x, log_px = carry
+
+        # Propose new state
+        x_prop = proposal(x)
+        log_px_prop = log_pdf(x_prop)
+
+        # Acceptance probability in log-space
+        log_alpha = log_px_prop - log_px          # assumes symmetric proposal
+        log_alpha = min(0.0, log_alpha)   # cap at log(1) = 0
+        u = random.uniform()                   # in (0, 1)
+        accept = log(u) < log_alpha
+
+        # If accepted, move to proposal; else stay
+        x_new = where(accept, x_prop, x)
+        log_px_new = where(accept, log_px_prop, log_px)
+
+        return (x_new, log_px_new), x_new
+
+    init_carry = (start_point, log_pdf(start_point))
+    (_, _, _), chain = lax.scan(one_step, init_carry, None, length=total_steps)
+
+    # Burn-in + thinning
+    samples = chain[burn_in::skipping]
+    return samples  # shape: (n, *start_point.shape)