Modernize find_MAP

.github/workflows/tests.yml

@@ -92,6 +92,7 @@ jobs:
           - |
             tests/tuning/test_scaling.py
             tests/tuning/test_starting.py
+            tests/tuning/test_scipy_interface.py
             tests/distributions/test_dist_math.py
             tests/distributions/test_transform.py
             tests/sampling/test_mcmc.py

ARCHITECTURE.md

@@ -37,6 +37,10 @@ of the topics below refer to that specific library
 * Model comparison, particularly efficient leave-one-out cross-validation approximation
 * Data structures for Bayesian inference data storage and manipulation
 
+### SciPy and better-optimize
+* Numerical optimization for {func}`pymc.find_MAP`, through `scipy.optimize` wrapped by
+  `better-optimize` (fused objectives, progress bars, early stopping)
+
 
 # Modules
 The codebase of PyMC is split among single Python file modules at the root

conda-envs/environment-alternative-backends.yml

@@ -6,6 +6,7 @@ channels:
 dependencies:
 # Base dependencies
 - arviz>=1.1.0,<2.0
+- better-optimize>=0.4.2,<0.5
 - blas
 - cachetools>=4.2.1,<7
 - cloudpickle

conda-envs/environment-dev.yml

@@ -6,6 +6,7 @@ channels:
 dependencies:
 # Base dependencies
 - arviz>=1.1.0,<2.0
+- better-optimize>=0.4.2,<0.5
 - blas
 - cachetools>=4.2.1,<7
 - cloudpickle

conda-envs/environment-docs.yml

@@ -6,6 +6,7 @@ channels:
 dependencies:
 # Base dependencies
 - arviz>=1.1.0,<2.0
+- better-optimize>=0.4.2,<0.5
 - cachetools>=4.2.1,<7
 - cloudpickle
 - numpy>=1.25.0

conda-envs/environment-test.yml

@@ -6,6 +6,7 @@ channels:
 dependencies:
 # Base dependencies
 - arviz>=1.1.0,<2.0
+- better-optimize>=0.4.2,<0.5
 - blas
 - cachetools>=4.2.1,<7
 - cloudpickle

conda-envs/windows-environment-dev.yml

@@ -6,6 +6,7 @@ channels:
 dependencies:
 # Base dependencies (see install guide for Windows)
 - arviz>=1.1.0,<2.0
+- better-optimize>=0.4.2,<0.5
 - blas
 - cachetools>=4.2.1,<7
 - cloudpickle

conda-envs/windows-environment-test.yml

@@ -6,6 +6,7 @@ channels:
 dependencies:
 # Base dependencies (see install guide for Windows)
 - arviz>=1.1.0,<2.0
+- better-optimize>=0.4.2,<0.5
 - blas
 - cachetools>=4.2.1,<7
 - cloudpickle

pymc/gp/gp.py

@@ -610,7 +610,7 @@ def predict(
         R"""
         Return mean and covariance of the conditional distribution given a `point`.
 
-        The `point` might be the MAP estimate or a sample from a trace.
+        The `point` is a ``{name: value}`` dict, e.g. one draw of ``pm.find_MAP()`` or ``pm.sample()``.
 
         Parameters
         ----------
@@ -1270,7 +1270,7 @@ def predict(self, Xnew, point=None, diag=False, pred_noise=False, model=None):
         R"""
         Return mean and covariance of the conditional distribution given a `point`.
 
-        The `point` might be the MAP estimate or a sample from a trace.
+        The `point` is a ``{name: value}`` dict, e.g. one draw of ``pm.find_MAP()`` or ``pm.sample()``.
 
         Parameters
         ----------

pymc/model/transform/conditioning.py

@@ -262,13 +262,13 @@ def change_value_transforms(
             w = pm.Binomial("w", n=9, p=p, observed=6)
 
         with change_value_transforms(base_m, {"p": logodds}) as transformed_p:
-            mean_q = pm.find_MAP()
+            mean_q = pm.find_MAP(return_inferencedata=True, jitter=True).posterior["p"].item()
 
         with change_value_transforms(transformed_p, {"p": None}) as untransformed_p:
             new_p = untransformed_p["p"]
-            std_q = ((1 / pm.find_hessian(mean_q, vars=[new_p])) ** 0.5)[0]
+            std_q = ((1 / pm.find_hessian({"p": mean_q}, vars=[new_p])) ** 0.5)[0]
 
-        print(f"  Mean, Standard deviation\\np {mean_q['p']:.2}, {std_q[0]:.2}")
+        print(f"  Mean, Standard deviation\\np {mean_q:.2}, {std_q[0]:.2}")
         #   Mean, Standard deviation
         # p 0.67, 0.16
 
@@ -343,12 +343,12 @@ def remove_value_transforms(
         with pm.Model() as transformed_m:
             p = pm.Uniform("p", 0, 1)
             w = pm.Binomial("w", n=9, p=p, observed=6)
-            mean_q = pm.find_MAP()
+            mean_q = pm.find_MAP(return_inferencedata=True, jitter=True).posterior["p"].item()
 
         with remove_value_transforms(transformed_m) as untransformed_m:
             new_p = untransformed_m["p"]
-            std_q = ((1 / pm.find_hessian(mean_q, vars=[new_p])) ** 0.5)[0]
-            print(f"  Mean, Standard deviation\\np {mean_q['p']:.2}, {std_q[0]:.2}")
+            std_q = ((1 / pm.find_hessian({"p": mean_q}, vars=[new_p])) ** 0.5)[0]
+            print(f"  Mean, Standard deviation\\np {mean_q:.2}, {std_q[0]:.2}")
 
         #   Mean, Standard deviation
         # p 0.67, 0.16

pymc/sampling/forward.py

@@ -672,7 +672,7 @@ def sample_posterior_predictive(
     Parameters
     ----------
     trace : backend, list, Dataset, DataTree, or MultiTrace
-        Trace generated from MCMC sampling, or a list of dicts (eg. points or from :func:`~pymc.find_MAP`),
+        Trace generated from MCMC sampling or :func:`~pymc.find_MAP`, or a list of dicts (eg. points),
         or :class:`xarray.Dataset` (eg. DataTree.posterior or DataTree.prior)
     model : BaseModel (optional if in ``with`` context)
         Model to be used to generate the posterior predictive samples. It will

pymc/sampling/mcmc.py

@@ -78,6 +78,7 @@
 from pymc.step_methods.arraystep import BlockedStep, PopulationArrayStepShared
 from pymc.step_methods.compound import flatten_steps
 from pymc.step_methods.hmc import quadpotential
+from pymc.tuning.starting import _find_MAP_point
 from pymc.util import (
     RandomSeed,
     RandomState,
@@ -722,8 +723,8 @@ def sample(
         Only applicable to the pymc nuts sampler.
     jitter_max_retries : int
         Maximum number of repeated attempts (per chain) at creating an initial matrix with uniform
-        jitter that yields a finite probability. This applies to ``jitter+adapt_diag`` and
-        ``jitter+adapt_full`` init methods.
+        jitter that yields a finite probability. This applies to ``jitter+adapt_diag``,
+        ``jitter+adapt_full`` and ``jitter+map`` init methods.
     n_init : int
         Number of iterations of initializer. Only works for 'ADVI' init methods.
     trace : backend, optional
@@ -1700,7 +1701,7 @@ def _init_jitter(
     jitter_max_retries: int,
     logp_fn: Callable[[PointType], np.ndarray] | None = None,
 ) -> list[PointType]:
-    """Apply a uniform jitter in [-1, 1] to the test value as starting point in each chain.
+    """Apply a uniform jitter in [-1, 1] to the initial point as starting point in each chain.
 
     ``model.check_start_vals`` is used to test whether the jittered starting
     values produce a finite log probability. Invalid values are resampled
@@ -1786,22 +1787,25 @@ def init_nuts(
           Currently, this is ``jitter+adapt_diag``, but this can change in the future. If you
           depend on the exact behaviour, choose an initialization method explicitly.
         * adapt_diag: Start with a identity mass matrix and then adapt a diagonal based on the
-          variance of the tuning samples. All chains use the test value (usually the prior mean)
+          variance of the tuning samples. All chains use the initial point (usually the prior mean)
           as starting point.
-        * jitter+adapt_diag: Same as ``adapt_diag``, but use test value plus a uniform jitter in
-          [-1, 1] as starting point in each chain.
+        * jitter+adapt_diag: Same as ``adapt_diag``, but use the initial point plus a uniform
+          jitter in [-1, 1] as starting point in each chain.
         * jitter+adapt_diag_grad:
           An experimental initialization method that uses information from gradients and samples
           during tuning.
         * advi+adapt_diag: Run ADVI and then adapt the resulting diagonal mass matrix based on the
           sample variance of the tuning samples.
         * advi: Run ADVI to estimate posterior mean and diagonal mass matrix.
         * advi_map: Initialize ADVI with MAP and use MAP as starting point.
-        * map: Use the MAP as starting point. This is discouraged.
+        * map: Use the MAP, searched for from the model's initial point, as starting point.
+          This is discouraged.
+        * jitter+map: Same as ``map``, but search from the initial point plus a uniform jitter
+          in [-1, 1].
         * adapt_full: Adapt a dense mass matrix using the sample covariances. All chains use the
-          test value (usually the prior mean) as starting point.
-        * jitter+adapt_full: Same as ``adapt_full``, but use test value plus a uniform jitter in
-          [-1, 1] as starting point in each chain.
+          initial point (usually the prior mean) as starting point.
+        * jitter+adapt_full: Same as ``adapt_full``, but use the initial point plus a uniform
+          jitter in [-1, 1] as starting point in each chain.
 
     chains : int
         Number of jobs to start.
@@ -1817,8 +1821,8 @@ def init_nuts(
         Whether or not to display a progressbar for advi sampling.
     jitter_max_retries : int
         Maximum number of repeated attempts (per chain) at creating an initial matrix with uniform jitter
-        that yields a finite probability. This applies to ``jitter+adapt_diag`` and ``jitter+adapt_full``
-        init methods.
+        that yields a finite probability. This applies to ``jitter+adapt_diag``, ``jitter+adapt_full``
+        and ``jitter+map`` init methods.
     **kwargs : keyword arguments
         Extra keyword arguments are forwarded to pymc.NUTS.
 
@@ -1879,6 +1883,8 @@ def model_logp_fn(ip: PointType) -> np.ndarray:
     )
 
     apoints = [DictToArrayBijection.map(point) for point in initial_points]
+    # MAP-based inits run a single search, from the first chain's initvals if given per chain
+    map_initvals = initvals if initvals is None or isinstance(initvals, dict) else initvals[0]
     apoints_data = [apoint.data for apoint in apoints]
     potential: quadpotential.QuadPotential
 
@@ -1958,7 +1964,15 @@ def model_logp_fn(ip: PointType) -> np.ndarray:
         cov = approx.std.eval() ** 2
         potential = quadpotential.QuadPotentialDiag(cov, rng=random_seed_list[0])
     elif init == "advi_map":
-        start = pm.find_MAP(include_transformed=True, seed=random_seed_list[0])
+        start = _find_MAP_point(
+            model=model,
+            initvals=map_initvals,
+            jitter=False,
+            jitter_max_retries=jitter_max_retries,
+            random_seed=random_seed_list[0],
+            progressbar=progressbar and not quiet,
+            compile_kwargs=compile_kwargs,
+        )
         approx = pm.MeanField(model=model, start=start)
         pm.fit(
             random_seed=random_seed_list[0],
@@ -1978,8 +1992,16 @@ def model_logp_fn(ip: PointType) -> np.ndarray:
         ]
         cov = approx.std.eval() ** 2
         potential = quadpotential.QuadPotentialDiag(cov, rng=random_seed_list[0])
-    elif init == "map":
-        start = pm.find_MAP(include_transformed=True, seed=random_seed_list[0])
+    elif init in ("map", "jitter+map"):
+        start = _find_MAP_point(
+            model=model,
+            initvals=map_initvals,
+            jitter=init == "jitter+map",
+            jitter_max_retries=jitter_max_retries,
+            random_seed=random_seed_list[0],
+            progressbar=progressbar and not quiet,
+            compile_kwargs=compile_kwargs,
+        )
         cov = -pm.find_hessian(point=start, negate_output=False)
         initial_points = [start] * chains
         potential = quadpotential.QuadPotentialFull(cov, rng=random_seed_list[0])