sdes typing, dataset test

Author: homerjed

sbgm/sde/_sde.py

@@ -1,11 +1,11 @@
-from typing import Sequence, Tuple, Self, Callable, Optional, Union
+from typing import Sequence, Tuple, Self, Callable, Optional
 import jax
 import jax.numpy as jnp
 import equinox as eqx
-from jaxtyping import Key, Array, Float, jaxtyped
+from jaxtyping import Key, Array, Float, Scalar
 
-TimeFn = Callable[[float | Float[Array, ""]], Float[Array, ""]]
-Time = Float[Array, ""] | float
+Time = Scalar | float
+TimeFn = Callable[[Time], Scalar]
 
 
 def default_weight_fn(t, *, beta_integral=None, sigma_fn=None): 
@@ -15,35 +15,71 @@ def default_weight_fn(t, *, beta_integral=None, sigma_fn=None):
 
 class SDE(eqx.Module):
     """
-        SDE abstract class.
+        Abstract base class for Stochastic Differential Equations (SDEs) used in 
+        score-based generative modeling and related diffusion models.
+
+        This class defines the required interface and provides base functionality for 
+        forward and reverse-time SDEs, prior sampling, and log-probability computation. 
+        The user should subclass `SDE` and implement the following methods:
+            - `sde()`: returns drift and diffusion coefficients
+            - `marginal_prob()`: returns the parameters of the marginal distribution at time t
+            - `prior_sample()`: returns samples from the terminal distribution p_T(x)
+            - `prior_log_prob()`: returns log-probabilities under p_T(x)
+            - `weight()`: returns weighting for loss functions
+
+        Attributes:
+            dt (float): Time discretization step size.
+            t0 (float): Start time of the diffusion process.
+            t1 (float): End time of the diffusion process.
     """
     dt: float
     t0: float
     t1: float
 
     def __init__(self, dt: float = 0.01, t0: float = 0., t1: float = 1.):
         """
-            Construct an SDE.
+            Initialize the base SDE with time parameters.
+
+            Args:
+                dt (float): Time step for the SDE solver.
+                t0 (float): Initial time of the process.
+                t1 (float): Terminal time of the process.
         """
         super().__init__()
         self.t0 = t0
         self.t1 = t1
         self.dt = dt
 
-    def sde(self, x: Array, t: Union[float, Array]) -> Tuple[Array, Array]:
-        pass
+    def sde(self, x: Array, t: Time) -> Tuple[Array, Array]:
+        """
+            Return the drift and diffusion coefficients f(x, t), g(t) of the SDE.
+
+            Must be implemented by subclass.
+        """
+        ...
 
-    def marginal_prob(self, x: Array, t: Union[float, Array]) -> Tuple[Array, Array]:
+    def marginal_prob(self, x: Array, t: Time) -> Tuple[Array, Array]:
         """ Parameters to determine the marginal distribution of the SDE, $p_t(x)$. """
-        pass
+        ...
 
     def prior_sample(self, key: Key, shape: Sequence[int]) -> Array:
-        """ Generate one sample from the prior distribution, $p_T(x)$. """
-        pass
+        """ 
+            Generate one sample from the prior distribution, $p_T(x)$. 
+        """
+        ...
+
+    def weight(self, t: Time, likelihood_weight: bool = False) -> Array:
+        """
+            Return the training loss weight at time t.
 
-    def weight(self, t: Union[float, Array], likelihood_weight: bool = False) -> Array:
-        """ Weighting for loss """
-        pass
+            Args:
+                t (float or Array): Time value(s).
+                likelihood_weight (bool): Whether to use likelihood weighting (optional).
+
+            Returns:
+                Array: Scalar or array of weights.
+        """
+        ...
 
     def prior_log_prob(self, z: Array) -> Array:
         """
@@ -52,20 +88,25 @@ def prior_log_prob(self, z: Array) -> Array:
             Useful for computing the log-likelihood via probability flow ODE.
 
             Args:
-            z: latent code
+                z: latent code
+
             Returns:
-            log probability density
+                log probability density
         """
-        pass
+        ...
 
     def reverse(self, score_fn: eqx.Module, probability_flow: bool = False) -> Self:
         """
             Create the reverse-time SDE/ODE.
 
             Args:
-            score_fn: A time-dependent score-based model that takes x and t and returns the score.
-            probability_flow: If `True`, create the reverse-time ODE used for probability flow sampling.
+                score_fn: A time-dependent score-based model that takes x and t and returns the score.
+                probability_flow: If `True`, create the reverse-time ODE used for probability flow sampling.
+
+            Returns:
+                SDE: A subclass implementing the reverse-time SDE.
         """
+
         sde_fn = self.sde
 
         if hasattr(self, "beta_integral_fn"):
@@ -77,7 +118,7 @@ def reverse(self, score_fn: eqx.Module, probability_flow: bool = False) -> Self:
         _t0 = self.t0
         _t1 = self.t1
 
-        # Build the class for reverse-time SDE.
+        # Build the class for the reverse-time SDE.
         class RSDE(self.__class__, SDE):
             probability_flow: bool
 
@@ -91,7 +132,7 @@ def sde(
                 t: Time, 
                 q: Optional[Float[Array, "..."]] = None,
                 a: Optional[Float[Array, "..."]] = None
-            ) -> Tuple[Float[Array, "..."], Float[Array, ""]]:
+            ) -> Tuple[Float[Array, "..."], Scalar]:
                 """ 
                     Create the drift and diffusion functions for the reverse SDE/ODE. 
                     - forward time SDE:
@@ -102,11 +143,11 @@ def sde(
                         dx = [f(x, t) - 0.5 * g^2(t) * score(x, t)] * dt (ODE => No dw)
                 """
                 t = jnp.asarray(t)
-                coeff = 0.5 if self.probability_flow else 1.
+                c = 0.5 if self.probability_flow else 1.
                 drift, diffusion = sde_fn(x, t)
                 score = score_fn(t, x, q, a)
                 # Drift coefficient of reverse SDE and probability flow only different by a factor
-                drift = drift - jnp.square(diffusion) * score * coeff
+                drift = drift - jnp.square(diffusion) * score * c
                 # Set the diffusion function to zero for ODEs (dw=0)
                 diffusion = 0. if self.probability_flow else diffusion
                 return drift, diffusion
@@ -117,49 +158,4 @@ def sde(
 def _get_log_prob_fn(scale: float = 1.) -> Callable:
     def _log_prob_fn(z: Array) -> Array:
         return jax.scipy.stats.norm.logpdf(z, loc=0., scale=scale).sum()
-    return _log_prob_fn
-
-
-# if __name__ == "__main__":
-#     import os 
-#     import matplotlib.pyplot as plt 
-#     import numpy as np
-
-#     figs_dir = "/project/ls-gruen/users/jed.homer/1pt_pdf/little_studies/sgm_lib/sgm/figs/"
-
-#     # Plot SDEs with time
-#     beta_integral_fn = lambda t: t
-#     beta_fn = get_beta_fn(beta_integral_fn)
-#     sigma_fn = lambda t: jnp.exp(t)
-
-#     times = dict(t0=0., t1=4., dt=0.1)
-
-#     vp_sde = VPSDE(beta_integral_fn, **times)
-#     ve_sde = VESDE(sigma_fn=sigma_fn)
-#     subvp_sde = SubVPSDE(beta_integral_fn, **times)
-
-#     x = jnp.ones((1,))
-#     T = jnp.linspace(1e-5, times["t1"], 1000)
-
-#     def get_sde_drift_and_diffusion_fn(sde):
-#         return jax.vmap(sde.sde, in_axes=(None, 0))
-
-#     def get_sde_mean_and_std(sde):
-#         return jax.vmap(sde.marginal_prob, in_axes=(None, 0))
-
-#     fig, axs = plt.subplots(1, 4, figsize=(21., 4.), dpi=200)
-#     ax = axs[0]
-#     ax.plot(T, jax.vmap(beta_fn)(T), linestyle=":", label=r"$\beta(t)$")
-#     ax_ = ax.twinx()
-#     ax.legend(frameon=False, loc="upper left")
-#     ax_.plot(T, jax.vmap(beta_integral_fn)(T), label=r"$\int_0^t\beta(s)ds$")
-#     ax_.legend(frameon=False, loc="lower right")
-#     plt.title("SDEs")
-#     for ax, _sde in zip(axs[1:], [ve_sde, vp_sde, subvp_sde]):
-#         mu, std = get_sde_mean_and_std(_sde)(x, T)
-#         ax.set_title(str(_sde.__class__.__name__))
-#         ax.plot(T, mu, label=r"$\mu(t)$")
-#         ax.plot(T, std, label=r"$\sigma(t)$")
-#         ax.legend(frameon=False)
-#     plt.savefig(os.path.join(figs_dir, "sdes.png"), bbox_inches="tight")
-#     plt.close()
+    return _log_prob_fn

sbgm/sde/_subvp.py

@@ -3,15 +3,15 @@
 import jax.numpy as jnp
 import jax.random as jr
 import equinox as eqx
-from jaxtyping import Key, Array, Float, jaxtyped
+from jaxtyping import PRNGKeyArray, Array, Float, Scalar, jaxtyped
 from beartype import beartype as typechecker
 
 from ._sde import SDE, _get_log_prob_fn, Time, TimeFn
 
 
 def get_beta_fn(beta_integral_fn: TimeFn | eqx.Module) -> TimeFn:
     """ Obtain beta function from a beta integral. """
-    def _beta_fn(t: Time) -> Float[Array, ""]:
+    def _beta_fn(t: Time) -> Scalar:
         _, beta = jax.jvp(
             beta_integral_fn, 
             primals=(t,), 
@@ -47,7 +47,7 @@ def __init__(
         self.weight_fn = weight_fn
 
     @jaxtyped(typechecker=typechecker)
-    def sde(self, x: Float[Array, "..."], t: Time) -> Tuple[Float[Array, "..."], Float[Array, ""]]:
+    def sde(self, x: Float[Array, "..."], t: Time) -> Tuple[Float[Array, "..."], Scalar]:
         """ 
             dx = f(x, t) * dt + g(t) * dw 
             dx = -0.5 * beta(t) * x * dt + sqrt(beta(t) * (1 - exp(-2 * int[beta(s)]))) * dw
@@ -59,7 +59,7 @@ def sde(self, x: Float[Array, "..."], t: Time) -> Tuple[Float[Array, "..."], Flo
         return drift, diffusion
 
     @jaxtyped(typechecker=typechecker)
-    def marginal_prob(self, x: Float[Array, "..."], t: Time) -> Tuple[Float[Array, "..."], Float[Array, ""]]:
+    def marginal_prob(self, x: Float[Array, "..."], t: Time) -> Tuple[Float[Array, "..."], Scalar]:
         """ 
             Sub-VP SDE p_t(x(t)|x(0)) is 
                 x(t) ~ G[x(t)|mu(x(0), t), sigma^2(t)]
@@ -73,7 +73,7 @@ def marginal_prob(self, x: Float[Array, "..."], t: Time) -> Tuple[Float[Array, "
         return mean, std
 
     @jaxtyped(typechecker=typechecker)
-    def weight(self, t: Time, likelihood_weight: bool = False) -> Float[Array, ""]:
+    def weight(self, t: Time, likelihood_weight: bool = False) -> Scalar:
         # Likelihood weighting: above Eq 8 https://arxiv.org/pdf/2101.09258.pdf
         if self.weight_fn is not None and not likelihood_weight:
             weight = self.weight_fn(t)
@@ -84,8 +84,8 @@ def weight(self, t: Time, likelihood_weight: bool = False) -> Float[Array, ""]:
                 weight = jnp.square(1. - jnp.exp(-self.beta_integral_fn(t)))
         return weight 
 
-    def prior_sample(self, key: Key[jnp.ndarray, "..."], shape: Sequence[int]) -> Float[Array, "..."]:
+    def prior_sample(self, key: PRNGKeyArray, shape: Sequence[int]) -> Float[Array, "..."]:
         return jr.normal(key, shape)
 
-    def prior_log_prob(self, z: Float[Array, "..."]) -> Float[Array, ""]:
+    def prior_log_prob(self, z: Float[Array, "..."]) -> Scalar:
         return _get_log_prob_fn(scale=1.)(z)

sbgm/sde/_ve.py

@@ -1,17 +1,17 @@
-from typing import Callable, Optional, Sequence, Tuple, Union
+from typing import Optional, Sequence, Tuple, Union
 import jax
 import jax.numpy as jnp
 import jax.random as jr
 import equinox as eqx
-from jaxtyping import Key, Array, Float, jaxtyped
+from jaxtyping import PRNGKeyArray, Array, Float, Scalar, jaxtyped
 from beartype import beartype as typechecker
 
 from ._sde import SDE, _get_log_prob_fn, Time, TimeFn
 
 
 def get_diffusion_fn(sigma_fn: Union[TimeFn, eqx.Module]) -> TimeFn:
     """ Get diffusion coefficient function for VE SDE: dx = sqrt(d[sigma^2(t)]/dt)dw """
-    def _diffusion_fn(t: Time) -> Float[Array, ""]:
+    def _diffusion_fn(t: Time) -> Scalar:
         _, dsigmadt = jax.jvp(
             lambda t: jnp.square(sigma_fn(t)), 
             primals=(t,), 
@@ -41,15 +41,15 @@ def __init__(
             dx = sqrt(d[sigma_fn(t) ** 2]/dt)
 
             Args:
-            sigma: default variance value
-            dt: timestep width
+                sigma: default variance value
+                dt: timestep width
         """
         super().__init__(dt=dt, t0=t0, t1=t1)
         self.sigma_fn = sigma_fn
         self.weight_fn = weight_fn
 
     @jaxtyped(typechecker=typechecker)
-    def sde(self, x: Float[Array, "..."], t: Time) -> Tuple[Float[Array, "..."], Float[Array, ""]]:
+    def sde(self, x: Float[Array, "..."], t: Time) -> Tuple[Float[Array, "..."], Scalar]:
         drift = jnp.zeros_like(x)
         _, dsigma2dt = jax.jvp(
             lambda t: jnp.square(self.sigma_fn(t)), 
@@ -61,7 +61,7 @@ def sde(self, x: Float[Array, "..."], t: Time) -> Tuple[Float[Array, "..."], Flo
         return drift, diffusion
 
     @jaxtyped(typechecker=typechecker)
-    def marginal_prob(self, x: Float[Array, "..."], t: Time) -> Tuple[Float[Array, "..."], Float[Array, ""]]:
+    def marginal_prob(self, x: Float[Array, "..."], t: Time) -> Tuple[Float[Array, "..."], Scalar]:
         """ 
             SDE:
                 dx = sqrt(d[sigma^2(t)]/dt) * dw
@@ -74,7 +74,7 @@ def marginal_prob(self, x: Float[Array, "..."], t: Time) -> Tuple[Float[Array, "
         return x, std 
 
     @jaxtyped(typechecker=typechecker)
-    def weight(self, t: Time, likelihood_weight: bool = False) -> Float[Array, ""]:
+    def weight(self, t: Time, likelihood_weight: bool = False) -> Scalar:
         if self.weight_fn is not None and not likelihood_weight:
             weight = self.weight_fn(t)
         else:
@@ -84,8 +84,8 @@ def weight(self, t: Time, likelihood_weight: bool = False) -> Float[Array, ""]:
                 weight = jnp.square(self.sigma_fn(t)) # Same for likelihood weighting
         return weight
 
-    def prior_sample(self, key: Key[jnp.ndarray, "..."], shape: Sequence[int]) -> Float[Array, "..."]:
+    def prior_sample(self, key: PRNGKeyArray, shape: Sequence[int]) -> Float[Array, "..."]:
         return jr.normal(key, shape) * self.sigma_fn(self.t1) 
 
-    def prior_log_prob(self, z: Float[Array, "..."]) -> Float[Array, ""]:
+    def prior_log_prob(self, z: Float[Array, "..."]) -> Scalar:
         return _get_log_prob_fn(scale=self.sigma_fn(self.t1))(z)