MCMCSamplingAlgorithm -> SamplingAlgorithm (#542)

* MCMCSamplingAlgorithm -> SamplingAlgorithm * include stochastic gradient algorithms
blackjax-devs · Jun 13, 2023 · fd2cae8 · fd2cae8
1 parent a90ca8e
commit fd2cae8
Show file tree

Hide file tree

Showing 16 changed files with 67 additions and 67 deletions.
diff --git a/blackjax/base.py b/blackjax/base.py
@@ -85,7 +85,7 @@ def __call__(self, rng_key: PRNGKey, state: State) -> Tuple[State, Info]:
         """
 
 
-class MCMCSamplingAlgorithm(NamedTuple):
+class SamplingAlgorithm(NamedTuple):
     """A pair of functions that represents a MCMC sampling algorithm.
 
     Blackjax sampling algorithms are implemented as a pair of pure functions: a

diff --git a/blackjax/mcmc/elliptical_slice.py b/blackjax/mcmc/elliptical_slice.py
@@ -17,7 +17,7 @@
 import jax
 import jax.numpy as jnp
 
-from blackjax.base import MCMCSamplingAlgorithm
+from blackjax.base import SamplingAlgorithm
 from blackjax.types import Array, ArrayLikeTree, ArrayTree, PRNGKey
 from blackjax.util import generate_gaussian_noise
 
@@ -149,7 +149,7 @@ class elliptical_slice:
 
     Returns
     -------
-    A ``MCMCSamplingAlgorithm``.
+    A ``SamplingAlgorithm``.
     """
 
     init = staticmethod(init)
@@ -161,7 +161,7 @@ def __new__(  # type: ignore[misc]
         *,
         mean: Array,
         cov: Array,
-    ) -> MCMCSamplingAlgorithm:
+    ) -> SamplingAlgorithm:
         kernel = cls.build_kernel(cov, mean)
 
         def init_fn(position: ArrayLikeTree):
@@ -174,7 +174,7 @@ def step_fn(rng_key: PRNGKey, state):
                 loglikelihood_fn,
             )
 
-        return MCMCSamplingAlgorithm(init_fn, step_fn)
+        return SamplingAlgorithm(init_fn, step_fn)
 
 
 def elliptical_proposal(

diff --git a/blackjax/mcmc/ghmc.py b/blackjax/mcmc/ghmc.py
@@ -21,7 +21,7 @@
 import blackjax.mcmc.integrators as integrators
 import blackjax.mcmc.metrics as metrics
 import blackjax.mcmc.proposal as proposal
-from blackjax.base import MCMCSamplingAlgorithm
+from blackjax.base import SamplingAlgorithm
 from blackjax.types import ArrayLikeTree, ArrayTree, PRNGKey
 from blackjax.util import generate_gaussian_noise, pytree_size
 
@@ -253,7 +253,7 @@ class ghmc:
 
     Returns
     -------
-    A ``MCMCSamplingAlgorithm``.
+    A ``SamplingAlgorithm``.
     """
 
     init = staticmethod(init)
@@ -269,7 +269,7 @@ def __new__(  # type: ignore[misc]
         *,
         divergence_threshold: int = 1000,
         noise_gn: Callable = lambda _: 0.0,
-    ) -> MCMCSamplingAlgorithm:
+    ) -> SamplingAlgorithm:
         kernel = cls.build_kernel(noise_gn, divergence_threshold)
 
         def init_fn(position: ArrayLikeTree, rng_key: PRNGKey):
@@ -286,4 +286,4 @@ def step_fn(rng_key: PRNGKey, state):
                 delta,
             )
 
-        return MCMCSamplingAlgorithm(init_fn, step_fn)  # type: ignore[arg-type]
+        return SamplingAlgorithm(init_fn, step_fn)  # type: ignore[arg-type]
diff --git a/blackjax/mcmc/hmc.py b/blackjax/mcmc/hmc.py
@@ -20,7 +20,7 @@
 import blackjax.mcmc.metrics as metrics
 import blackjax.mcmc.proposal as proposal
 import blackjax.mcmc.trajectory as trajectory
-from blackjax.base import MCMCSamplingAlgorithm
+from blackjax.base import SamplingAlgorithm
 from blackjax.mcmc.trajectory import hmc_energy
 from blackjax.types import Array, ArrayLikeTree, ArrayTree, PRNGKey
 
@@ -205,7 +205,7 @@ class hmc:
 
     Returns
     -------
-    A ``MCMCSamplingAlgorithm``.
+    A ``SamplingAlgorithm``.
     """
 
     init = staticmethod(init)
@@ -220,7 +220,7 @@ def __new__(  # type: ignore[misc]
         *,
         divergence_threshold: int = 1000,
         integrator: Callable = integrators.velocity_verlet,
-    ) -> MCMCSamplingAlgorithm:
+    ) -> SamplingAlgorithm:
         kernel = cls.build_kernel(integrator, divergence_threshold)
 
         def init_fn(position: ArrayLikeTree):
@@ -236,7 +236,7 @@ def step_fn(rng_key: PRNGKey, state):
                 num_integration_steps,
             )
 
-        return MCMCSamplingAlgorithm(init_fn, step_fn)
+        return SamplingAlgorithm(init_fn, step_fn)
 
 
 def hmc_proposal(

diff --git a/blackjax/mcmc/mala.py b/blackjax/mcmc/mala.py
@@ -20,7 +20,7 @@
 
 import blackjax.mcmc.diffusions as diffusions
 import blackjax.mcmc.proposal as proposal
-from blackjax.base import MCMCSamplingAlgorithm
+from blackjax.base import SamplingAlgorithm
 from blackjax.types import ArrayLikeTree, ArrayTree, PRNGKey
 
 __all__ = ["MALAState", "MALAInfo", "init", "build_kernel", "mala"]
@@ -165,7 +165,7 @@ class mala:
 
     Returns
     -------
-    A ``MCMCSamplingAlgorithm``.
+    A ``SamplingAlgorithm``.
 
     """
 
@@ -176,7 +176,7 @@ def __new__(  # type: ignore[misc]
         cls,
         logdensity_fn: Callable,
         step_size: float,
-    ) -> MCMCSamplingAlgorithm:
+    ) -> SamplingAlgorithm:
         kernel = cls.build_kernel()
 
         def init_fn(position: ArrayLikeTree):
@@ -185,4 +185,4 @@ def init_fn(position: ArrayLikeTree):
         def step_fn(rng_key: PRNGKey, state):
             return kernel(rng_key, state, logdensity_fn, step_size)
 
-        return MCMCSamplingAlgorithm(init_fn, step_fn)
+        return SamplingAlgorithm(init_fn, step_fn)
diff --git a/blackjax/mcmc/marginal_latent_gaussian.py b/blackjax/mcmc/marginal_latent_gaussian.py
@@ -18,7 +18,7 @@
 import jax.numpy as jnp
 import jax.scipy.linalg as linalg
 
-from blackjax.base import MCMCSamplingAlgorithm
+from blackjax.base import SamplingAlgorithm
 from blackjax.types import Array, PRNGKey
 
 __all__ = ["MarginalState", "MarginalInfo", "init_and_kernel", "mgrad_gaussian"]
@@ -173,7 +173,7 @@ class mgrad_gaussian:
 
     Returns
     -------
-    A ``MCMCSamplingAlgorithm``.
+    A ``SamplingAlgorithm``.
 
     """
 
@@ -182,7 +182,7 @@ def __new__(  # type: ignore[misc]
         logdensity_fn: Callable,
         covariance: Array,
         mean: Optional[Array] = None,
-    ) -> MCMCSamplingAlgorithm:
+    ) -> SamplingAlgorithm:
         init, kernel = init_and_kernel(logdensity_fn, covariance, mean)
 
         def init_fn(position: Array):
@@ -195,4 +195,4 @@ def step_fn(rng_key: PRNGKey, state, delta: float):
                 delta,
             )
 
-        return MCMCSamplingAlgorithm(init_fn, step_fn)  # type: ignore[arg-type]
+        return SamplingAlgorithm(init_fn, step_fn)  # type: ignore[arg-type]
diff --git a/blackjax/mcmc/nuts.py b/blackjax/mcmc/nuts.py
@@ -24,7 +24,7 @@
 import blackjax.mcmc.proposal as proposal
 import blackjax.mcmc.termination as termination
 import blackjax.mcmc.trajectory as trajectory
-from blackjax.base import MCMCSamplingAlgorithm
+from blackjax.base import SamplingAlgorithm
 from blackjax.types import Array, ArrayLikeTree, ArrayTree, PRNGKey
 
 __all__ = ["NUTSInfo", "init", "build_kernel", "nuts"]
@@ -207,7 +207,7 @@ class nuts:
 
     Returns
     -------
-    A ``MCMCSamplingAlgorithm``.
+    A ``SamplingAlgorithm``.
 
     """
 
@@ -223,7 +223,7 @@ def __new__(  # type: ignore[misc]
         max_num_doublings: int = 10,
         divergence_threshold: int = 1000,
         integrator: Callable = integrators.velocity_verlet,
-    ) -> MCMCSamplingAlgorithm:
+    ) -> SamplingAlgorithm:
         kernel = cls.build_kernel(integrator, divergence_threshold, max_num_doublings)
 
         def init_fn(position: ArrayLikeTree):
@@ -238,7 +238,7 @@ def step_fn(rng_key: PRNGKey, state):
                 inverse_mass_matrix,
             )
 
-        return MCMCSamplingAlgorithm(init_fn, step_fn)
+        return SamplingAlgorithm(init_fn, step_fn)
 
 
 def iterative_nuts_proposal(

diff --git a/blackjax/mcmc/periodic_orbital.py b/blackjax/mcmc/periodic_orbital.py
@@ -19,7 +19,7 @@
 
 import blackjax.mcmc.integrators as integrators
 import blackjax.mcmc.metrics as metrics
-from blackjax.base import MCMCSamplingAlgorithm
+from blackjax.base import SamplingAlgorithm
 from blackjax.types import Array, ArrayLikeTree, ArrayTree, PRNGKey
 
 __all__ = ["PeriodicOrbitalState", "init", "build_kernel", "orbital_hmc"]
@@ -259,7 +259,7 @@ class orbital_hmc:
 
     Returns
     -------
-    A ``MCMCSamplingAlgorithm``.
+    A ``SamplingAlgorithm``.
     """
 
     init = staticmethod(init)
@@ -273,7 +273,7 @@ def __new__(  # type: ignore[misc]
         period: int,
         *,
         bijection: Callable = integrators.velocity_verlet,
-    ) -> MCMCSamplingAlgorithm:
+    ) -> SamplingAlgorithm:
         kernel = cls.build_kernel(bijection)
 
         def init_fn(position: ArrayLikeTree):
@@ -289,7 +289,7 @@ def step_fn(rng_key: PRNGKey, state):
                 period,
             )
 
-        return MCMCSamplingAlgorithm(init_fn, step_fn)
+        return SamplingAlgorithm(init_fn, step_fn)
 
 
 def periodic_orbital_proposal(

diff --git a/blackjax/mcmc/random_walk.py b/blackjax/mcmc/random_walk.py
@@ -59,7 +59,7 @@
 import numpy as np
 from jax import numpy as jnp
 
-from blackjax.base import MCMCSamplingAlgorithm
+from blackjax.base import SamplingAlgorithm
 from blackjax.mcmc import proposal
 from blackjax.types import Array, ArrayLikeTree, ArrayTree, PRNGKey
 from blackjax.util import generate_gaussian_noise
@@ -210,7 +210,7 @@ class additive_step_random_walk:
 
     Returns
     -------
-    A ``MCMCSamplingAlgorithm``.
+    A ``SamplingAlgorithm``.
     """
 
     init = staticmethod(init)
@@ -227,13 +227,13 @@ def normal_random_walk(cls, logdensity_fn: Callable, sigma):
             The value of the covariance matrix of the gaussian proposal distribution.
         Returns
         -------
-             A ``MCMCSamplingAlgorithm``.
+             A ``SamplingAlgorithm``.
         """
         return cls(logdensity_fn, normal(sigma))
 
     def __new__(  # type: ignore[misc]
         cls, logdensity_fn: Callable, random_step: Callable
-    ) -> MCMCSamplingAlgorithm:
+    ) -> SamplingAlgorithm:
         kernel = cls.build_kernel()
 
         def init_fn(position: ArrayLikeTree):
@@ -242,7 +242,7 @@ def init_fn(position: ArrayLikeTree):
         def step_fn(rng_key: PRNGKey, state):
             return kernel(rng_key, state, logdensity_fn, random_step)
 
-        return MCMCSamplingAlgorithm(init_fn, step_fn)
+        return SamplingAlgorithm(init_fn, step_fn)
 
 
 def build_irmh() -> Callable:
@@ -313,7 +313,7 @@ class irmh:
 
     Returns
     -------
-    A ``MCMCSamplingAlgorithm``.
+    A ``SamplingAlgorithm``.
 
     """
 
@@ -324,7 +324,7 @@ def __new__(  # type: ignore[misc]
         cls,
         logdensity_fn: Callable,
         proposal_distribution: Callable,
-    ) -> MCMCSamplingAlgorithm:
+    ) -> SamplingAlgorithm:
         kernel = cls.build_kernel()
 
         def init_fn(position: ArrayLikeTree):
@@ -333,7 +333,7 @@ def init_fn(position: ArrayLikeTree):
         def step_fn(rng_key: PRNGKey, state):
             return kernel(rng_key, state, logdensity_fn, proposal_distribution)
 
-        return MCMCSamplingAlgorithm(init_fn, step_fn)
+        return SamplingAlgorithm(init_fn, step_fn)
 
 
 def build_rmh():
@@ -428,7 +428,7 @@ class rmh:
 
     Returns
     -------
-    A ``MCMCSamplingAlgorithm``.
+    A ``SamplingAlgorithm``.
     """
 
     init = staticmethod(init)
@@ -439,7 +439,7 @@ def __new__(  # type: ignore[misc]
         logdensity_fn: Callable,
         proposal_generator: Callable[[PRNGKey, ArrayLikeTree], ArrayTree],
         proposal_logdensity_fn: Optional[Callable[[ArrayLikeTree], ArrayTree]] = None,
-    ) -> MCMCSamplingAlgorithm:
+    ) -> SamplingAlgorithm:
         kernel = cls.build_kernel()
 
         def init_fn(position: ArrayLikeTree):
@@ -454,7 +454,7 @@ def step_fn(rng_key: PRNGKey, state):
                 proposal_logdensity_fn,
             )
 
-        return MCMCSamplingAlgorithm(init_fn, step_fn)
+        return SamplingAlgorithm(init_fn, step_fn)
 
 
 def build_rmh_transition_energy(proposal_logdensity_fn: Optional[Callable]) -> Callable:

diff --git a/blackjax/sgmcmc/csgld.py b/blackjax/sgmcmc/csgld.py
@@ -19,7 +19,7 @@
 import jax
 import jax.numpy as jnp
 
-from blackjax.base import MCMCSamplingAlgorithm
+from blackjax.base import SamplingAlgorithm
 from blackjax.sgmcmc.diffusions import overdamped_langevin
 from blackjax.types import Array, ArrayLikeTree, ArrayTree, PRNGKey
 
@@ -206,7 +206,7 @@ class csgld:
 
     Returns
     -------
-    A ``MCMCSamplingAlgorithm``.
+    A ``SamplingAlgorithm``.
 
     """
     init = staticmethod(init)
@@ -220,7 +220,7 @@ def __new__(  # type: ignore[misc]
         num_partitions: int = 512,
         energy_gap: float = 100,
         min_energy: float = 0,
-    ) -> MCMCSamplingAlgorithm:
+    ) -> SamplingAlgorithm:
         kernel = cls.build_kernel(num_partitions, energy_gap, min_energy)
 
         def init_fn(position: ArrayLikeTree):
@@ -246,4 +246,4 @@ def step_fn(
                 temperature,
             )
 
-        return MCMCSamplingAlgorithm(init_fn, step_fn)  # type: ignore[arg-type]
+        return SamplingAlgorithm(init_fn, step_fn)  # type: ignore[arg-type]