Mext try

Author: FlorianPfaff

pyrecest/distributions/abstract_manifold_specific_distribution.py

@@ -156,7 +156,8 @@ def sample_metropolis_hastings_jax(
     n:          number of samples to return (after burn-in and thinning)
     """
     import jax.numpy as _jnp
-    from jax import _random, _lax
+    from jax import lax as _lax
+    from jax import random as _random
 
 
     start_point = _jnp.asarray(start_point)

pyrecest/distributions/hypersphere_subset/abstract_hyperhemispherical_distribution.py

@@ -69,28 +69,39 @@ def sample_metropolis_hastings(
                 def proposal(_):
                     return HyperhemisphericalUniformDistribution(self.dim).sample(1)
             else:
+                # JAX backend: proposal(key, x) -> x_prop
                 import jax as _jax
                 import jax.numpy as _jnp
-                def proposal(key: _jax.Array, _) -> tuple[_jax.Array, _jnp.ndarray]:
-                    # Sample on full sphere
-                    if self.dim == 2:
-                        key, key_phi = _jnp.random.split(key)
-                        key, key_sz = _jnp.random.split(key)
-
-                        phi = 2.0 * pi * _jnp.random.uniform(key_phi, shape=(1,))
-                        sz = 2.0 * _jnp.random.uniform(key_sz, shape=(1,)) - 1.0
-                        r = sqrt(1.0 - sz**2)
 
-                        s = stack([r * cos(phi), r * sin(phi), sz], axis=1)
+                def proposal(key, _):
+                    """JAX independence proposal: uniform on upper hemisphere."""
+                    if self.dim == 2:
+                        # Explicit S² sampling
+                        key, key_phi = _jax.random.split(key)
+                        key, key_sz = _jax.random.split(key)
+
+                        phi = 2.0 * _jnp.pi * _jax.random.uniform(key_phi, shape=(1,))
+                        sz = 2.0 * _jax.random.uniform(key_sz, shape=(1,)) - 1.0
+                        r = _jnp.sqrt(1.0 - sz**2)
+
+                        # Shape (1, 3)
+                        s = _jnp.stack(
+                            [r * _jnp.cos(phi), r * _jnp.sin(phi), sz],
+                            axis=1,
+                        )
                     else:
-                        key, subkey = random.split(key)
-                        samples_unnorm = random.normal(subkey, shape=(1, self.dim + 1))
-                        norms = linalg.norm(samples_unnorm, axis=1, keepdims=True)
+                        # General S^d: sample N(0, I) in R^{d+1} and normalize
+                        key, subkey = _jax.random.split(key)
+                        samples_unnorm = _jax.random.normal(subkey, shape=(1, self.dim + 1))
+                        norms = _jnp.linalg.norm(samples_unnorm, axis=1, keepdims=True)
                         s = samples_unnorm / norms
 
-                    # To upper hemisphere
-                    s = (1 - 2 * (s[-1, :] < 0)) * s
-                    return key, s
+                    # Project to upper hemisphere: last coordinate >= 0
+                    # s shape: (1, dim+1); last coord is s[..., -1:]
+                    sign = _jnp.where(s[..., -1:] < 0.0, -1.0, 1.0)
+                    s = sign * s
+
+                    return s
 
 
         if start_point is None:

pyrecest/distributions/hypersphere_subset/abstract_hyperspherical_distribution.py

@@ -21,6 +21,9 @@
     sin,
     vstack,
     zeros,
+    sqrt,
+    stack,
+    linalg,
 )
 from scipy.optimize import minimize
 
@@ -74,9 +77,37 @@ def sample_metropolis_hastings(
                 HypersphericalUniformDistribution,
             )
 
-            def proposal(_):
-                return HypersphericalUniformDistribution(self.dim).sample(1)
-
+            if pyrecest.backend.__backend_name__ in ("numpy", "pytorch"):
+                def proposal(_):
+                    return HypersphericalUniformDistribution(self.dim).sample(1)
+            else:
+                import jax as _jax
+                import jax.numpy as _jnp
+                def proposal(key, _):
+                    """JAX independence proposal: uniform on hypersphere."""
+                    if self.dim == 2:
+                        # Explicit S² sampling
+                        key, key_phi = _jax.random.split(key)
+                        key, key_sz = _jax.random.split(key)
+
+                        phi = 2.0 * _jnp.pi * _jax.random.uniform(key_phi, shape=(1,))
+                        sz = 2.0 * _jax.random.uniform(key_sz, shape=(1,)) - 1.0
+                        r = _jnp.sqrt(1.0 - sz**2)
+
+                        # Shape (1, 3)
+                        s = _jnp.stack(
+                            [r * _jnp.cos(phi), r * _jnp.sin(phi), sz],
+                            axis=1,
+                        )
+                    else:
+                        # General S^d: sample N(0, I) in R^{d+1} and normalize
+                        key, subkey = _jax.random.split(key)
+                        samples_unnorm = _jax.random.normal(subkey, shape=(1, self.dim + 1))
+                        norms = _jnp.linalg.norm(samples_unnorm, axis=1, keepdims=True)
+                        s = samples_unnorm / norms
+
+                    return s
+                
         if start_point is None:
             start_point = HypersphericalUniformDistribution(self.dim).sample(1)
         # Call the sample_metropolis_hastings method of AbstractDistribution