Merge pull request #1263 from PyAutoLabs/feature/priors-jax-native

feat: JAX-native priors — xp dispatch on value_for / log_prior_from_value / vector_from_unit_vector

Author: Jammy2211

autofit/mapper/prior/abstract.py

@@ -4,6 +4,8 @@
 from copy import copy
 from typing import Union, Tuple, Optional, Dict
 
+import numpy as np
+
 from autoconf import conf
 
 from autofit.mapper.prior.arithmetic import ArithmeticMixin
@@ -137,7 +139,7 @@ def random(
             )
         )
 
-    def value_for(self, unit: float) -> float:
+    def value_for(self, unit, xp=np):
         """
         Return a physical value for a value between 0 and 1 with the transformation
         described by this prior.
@@ -146,6 +148,11 @@ def value_for(self, unit: float) -> float:
         ----------
         unit
             A unit value between 0 and 1.
+        xp
+            Array-module to dispatch on (``numpy`` or ``jax.numpy``). Default ``numpy``.
+            Concrete subclasses override this method to provide a JAX-traceable
+            closed-form when ``xp`` is ``jax.numpy``; the base path delegates to
+            the message stack which is scipy-backed and NumPy-only.
 
         Returns
         -------

autofit/mapper/prior/gaussian.py

@@ -113,3 +113,21 @@ def parameter_string(self) -> str:
         Return a human-readable string summarizing the GaussianPrior parameters.
         """
         return f"mean = {self.mean}, sigma = {self.sigma}"
+
+    def value_for(self, unit, xp=np):
+        """
+        Map a unit value in [0, 1] to a physical value drawn from this Gaussian prior.
+
+        Parameters
+        ----------
+        unit
+            A unit value between 0 and 1.
+        xp
+            Array-module to dispatch on (``numpy`` or ``jax.numpy``). Default ``numpy``.
+            The NumPy path delegates to the message stack (``erfinv`` via scipy); the
+            JAX path uses the same closed-form via ``jax.scipy.special.erfinv``.
+        """
+        if xp is np:
+            return self.message.value_for(unit)
+        from jax.scipy.special import erfinv
+        return self.mean + self.sigma * xp.sqrt(2.0) * erfinv(2.0 * unit - 1.0)

autofit/mapper/prior/log_gaussian.py

@@ -113,7 +113,7 @@ def _new_for_base_message(self, message):
             id_=self.instance().id,
         )
 
-    def value_for(self, unit: float) -> float:
+    def value_for(self, unit, xp=np):
         """
         Return a physical value for a value between 0 and 1 with the transformation
         described by this prior.
@@ -122,21 +122,33 @@ def value_for(self, unit: float) -> float:
         ----------
         unit
             A unit value between 0 and 1.
+        xp
+            Array-module to dispatch on (``numpy`` or ``jax.numpy``). Default ``numpy``.
+            The NumPy path delegates to the message stack; the JAX path uses the
+            closed-form ``exp(mean + sigma * sqrt(2) * erfinv(2*unit - 1))``.
 
         Returns
         -------
         A physical value, mapped from the unit value accoridng to the prior.
         """
-        return super().value_for(unit)
+        if xp is np:
+            return super().value_for(unit)
+        from jax.scipy.special import erfinv
+        log_value = self.mean + self.sigma * xp.sqrt(2.0) * erfinv(2.0 * unit - 1.0)
+        return xp.exp(log_value)
 
     @property
     def parameter_string(self) -> str:
         return f"mean = {self.mean}, sigma = {self.sigma}"
 
     def log_prior_from_value(self, value, xp=np):
-        if value <= 0:
-            return float("-inf")
-
-        return self.message.base_message.log_prior_from_value(np.log(value)) - np.log(
-            value
-        )
+        if xp is np:
+            if value <= 0:
+                return float("-inf")
+            return self.message.base_message.log_prior_from_value(
+                np.log(value)
+            ) - np.log(value)
+
+        log_value = xp.log(value)
+        base_log_prior = (log_value - self.mean) ** 2 / (2 * self.sigma ** 2)
+        return xp.where(value > 0, base_log_prior - log_value, xp.inf)

autofit/mapper/prior/log_uniform.py

@@ -121,17 +121,23 @@ def log_prior_from_value(self, value, xp=np) -> float:
         ----------
         value : float
             The physical value of this prior's corresponding parameter in a `NonLinearSearch` sample.
+        xp
+            Array-module to dispatch on (``numpy`` or ``jax.numpy``). Default ``numpy``.
         """
         return 1.0 / value
 
-    def value_for(self, unit: float) -> float:
+    def value_for(self, unit, xp=np):
         """
         Returns a physical value from an input unit value according to the limits of the log10 uniform prior.
 
         Parameters
         ----------
         unit
             A unit value between 0 and 1.
+        xp
+            Array-module to dispatch on (``numpy`` or ``jax.numpy``). Default ``numpy``.
+            The NumPy path delegates to the message stack (scipy-backed); the JAX
+            path uses the closed-form ``lower * (upper / lower) ** unit``.
 
         Returns
         -------
@@ -145,7 +151,9 @@ def value_for(self, unit: float) -> float:
 
         physical_value = prior.value_for(unit=0.2)
         """
-        return super().value_for(unit)
+        if xp is np:
+            return super().value_for(unit)
+        return self.lower_limit * (self.upper_limit / self.lower_limit) ** unit
 
     def dict(self) -> dict:
         """

autofit/mapper/prior/truncated_gaussian.py

@@ -1,5 +1,7 @@
 from typing import Optional, Tuple
 
+import numpy as np
+
 from autofit.messages.truncated_normal import TruncatedNormalMessage
 from .abstract import Prior
 
@@ -130,3 +132,32 @@ def parameter_string(self) -> str:
                 f"lower_limit = {self.lower_limit}, "
                 f"upper_limit = {self.upper_limit}"
                 )
+
+    def value_for(self, unit, xp=np):
+        """
+        Map a unit value in [0, 1] to a physical value drawn from this truncated Gaussian prior.
+
+        Parameters
+        ----------
+        unit
+            A unit value between 0 and 1.
+        xp
+            Array-module to dispatch on (``numpy`` or ``jax.numpy``). Default ``numpy``.
+            Both paths share the standard truncated-normal inverse-CDF construction
+            via ``norm.cdf`` / ``norm.ppf`` from the matching ``scipy.stats`` /
+            ``jax.scipy.stats`` namespace.
+        """
+        if xp is np:
+            from scipy.stats import norm
+        else:
+            from jax.scipy.stats import norm
+
+        a = (self.lower_limit - self.mean) / self.sigma
+        b = (self.upper_limit - self.mean) / self.sigma
+
+        lower_cdf = norm.cdf(a)
+        upper_cdf = norm.cdf(b)
+        truncated_cdf = lower_cdf + unit * (upper_cdf - lower_cdf)
+
+        x_standard = norm.ppf(truncated_cdf)
+        return self.mean + self.sigma * x_standard

autofit/mapper/prior/uniform.py

@@ -132,14 +132,19 @@ def parameter_string(self) -> str:
         """A human-readable string summarizing the prior's lower and upper limits."""
         return f"lower_limit = {self.lower_limit}, upper_limit = {self.upper_limit}"
 
-    def value_for(self, unit: float) -> float:
+    def value_for(self, unit, xp=np):
         """
         Returns a physical value from an input unit value according to the limits of the uniform prior.
 
         Parameters
         ----------
         unit
             A unit value between 0 and 1.
+        xp
+            Array-module to dispatch on (``numpy`` or ``jax.numpy``). Default ``numpy``.
+            The NumPy path preserves the historical ``float(round(..., 14))`` snap
+            (used as a hash key in ``model.priors``); the JAX path uses the
+            closed-form ``lower + (upper - lower) * unit`` so the trace stays symbolic.
 
         Returns
         -------
@@ -153,9 +158,11 @@ def value_for(self, unit: float) -> float:
 
         physical_value = prior.value_for(unit=0.2)
         """
-        return float(
-            round(super().value_for(unit), 14)
-        )
+        if xp is np:
+            return float(
+                round(super().value_for(unit), 14)
+            )
+        return self.lower_limit + (self.upper_limit - self.lower_limit) * unit
 
     def log_prior_from_value(self, value, xp=np):
         """
@@ -166,7 +173,10 @@ def log_prior_from_value(self, value, xp=np):
 
         For a UniformPrior this is always zero, provided the value is between the lower and upper limit.
         """
-        return 0.0
+        if xp is np:
+            return 0.0
+        in_bounds = (value >= self.lower_limit) & (value <= self.upper_limit)
+        return xp.where(in_bounds, xp.zeros_like(value), -xp.inf)
 
     @property
     def limits(self) -> Tuple[float, float]:

autofit/mapper/prior_model/abstract.py

@@ -645,26 +645,41 @@ def priors_ordered_by_id(self):
         """Unique priors sorted by their id, defining the canonical parameter ordering."""
         return [prior for _, prior in self.prior_tuples_ordered_by_id]
 
-    def vector_from_unit_vector(self, unit_vector):
+    def vector_from_unit_vector(self, unit_vector, xp=np):
         """
         Parameters
         ----------
         unit_vector: [float]
             A unit hypercube vector
+        xp
+            Array-module to dispatch on (``numpy`` or ``jax.numpy``). Default ``numpy``.
+            When ``xp is numpy`` the return type stays a plain ``list`` to preserve
+            the existing contract used by Nautilus / Dynesty / Emcee / Zeus. When
+            ``xp is jax.numpy`` an ``xp.stack`` 1-D array is returned so the call
+            is JIT-traceable end-to-end.
 
         Returns
         -------
         values: [float]
             A vector with values output by priors
         """
-        return list(
-            map(
-                lambda prior_tuple, unit: prior_tuple.prior.value_for(
-                    unit,
-                ),
-                self.prior_tuples_ordered_by_id,
-                unit_vector,
+        if xp is np:
+            return list(
+                map(
+                    lambda prior_tuple, unit: prior_tuple.prior.value_for(
+                        unit,
+                    ),
+                    self.prior_tuples_ordered_by_id,
+                    unit_vector,
+                )
             )
+        return xp.stack(
+            [
+                prior_tuple.prior.value_for(unit, xp=xp)
+                for prior_tuple, unit in zip(
+                    self.prior_tuples_ordered_by_id, unit_vector
+                )
+            ]
         )
 
     def random_unit_vector_within_limits(

autofit/messages/normal.py

@@ -399,14 +399,16 @@ def logpdf_gradient_hessian(self, x: np.ndarray) -> Tuple[np.ndarray, np.ndarray
 
     __default_fields__ = ("log_norm", "id_")
 
-    def value_for(self, unit: float) -> float:
+    def value_for(self, unit, xp=np):
         """
         Map a unit value in [0, 1] to a physical value drawn from this Gaussian prior.
 
         Parameters
         ----------
         unit
             A unit value between 0 and 1 representing a uniform draw.
+        xp
+            Array-module to dispatch on (``numpy`` or ``jax.numpy``). Default ``numpy``.
 
         Returns
         -------
@@ -417,15 +419,13 @@ def value_for(self, unit: float) -> float:
         >>> prior = af.GaussianPrior(mean=1.0, sigma=2.0)
         >>> physical_value = prior.value_for(unit=0.5)
         """
-        if isinstance(unit, (np.ndarray, np.float64, float, int, list)):
-            from scipy.special import erfinv as scipy_erfinv
-            inv = scipy_erfinv(1 - 2.0 * (1.0 - unit))
+        if xp is np:
+            from scipy.special import erfinv
         else:
-            import jax.numpy as jnp
-            from jax._src.scipy.special import erfinv
-            inv = erfinv(1 - 2.0 * (1.0 - unit))
+            from jax.scipy.special import erfinv
 
-        return self.mean + (self.sigma * np.sqrt(2) * inv)
+        inv = erfinv(2.0 * unit - 1.0)
+        return self.mean + self.sigma * xp.sqrt(2.0) * inv
 
     def log_prior_from_value(self, value: float, xp=np) -> float:
         """

autofit/messages/truncated_normal.py

@@ -440,7 +440,7 @@ def logpdf_gradient_hessian(self, x: np.ndarray) -> Tuple[np.ndarray, np.ndarray
 
     __default_fields__ = ("log_norm", "id_")
 
-    def value_for(self, unit: float) -> float:
+    def value_for(self, unit, xp=np):
         """
         Map a unit value in [0, 1] to a physical value drawn from this truncated Gaussian prior.
 
@@ -451,6 +451,8 @@ def value_for(self, unit: float) -> float:
         ----------
         unit
             A unit value between 0 and 1 representing a uniform draw.
+        xp
+            Array-module to dispatch on (``numpy`` or ``jax.numpy``). Default ``numpy``.
 
         Returns
         -------
@@ -461,18 +463,18 @@ def value_for(self, unit: float) -> float:
         >>> prior = af.TruncatedNormalMessage(mean=1.0, sigma=2.0, lower_limit=0.0, upper_limit=2.0)
         >>> physical_value = prior.value_for(unit=0.5)
         """
-        from scipy.stats import norm
+        if xp is np:
+            from scipy.stats import norm
+        else:
+            from jax.scipy.stats import norm
 
-        # Standardized truncation bounds
         a = (self.lower_limit - self.mean) / self.sigma
         b = (self.upper_limit - self.mean) / self.sigma
 
-        # Interpolate unit into [Phi(a), Phi(b)]
         lower_cdf = norm.cdf(a)
         upper_cdf = norm.cdf(b)
         truncated_cdf = lower_cdf + unit * (upper_cdf - lower_cdf)
 
-        # Map back to x using inverse CDF, then rescale
         x_standard = norm.ppf(truncated_cdf)
         return self.mean + self.sigma * x_standard