Update and Refactor find_MAP and fit_laplace (pymc-devs#531)

* Move laplace and find_map to submodule

* Split idata utilities into `idata.py`

* Refactor find_MAP

* Refactor fit_laplace

* Update better-optimize version pin

* Handle labeling of non-scalar RVs without dims

* Add unconstrained posterior draws/points to unconstrained_posterior

Author: jessegrabowski

pymc_extras/inference/find_map.py

@@ -326,7 +326,7 @@ def find_MAP(
     )
 
     raveled_optimized = RaveledVars(optimizer_result.x, initial_params.point_map_info)
-    unobserved_vars = get_default_varnames(model.unobserved_value_vars, include_transformed=True)
+    unobserved_vars = get_default_varnames(model.unobserved_value_vars, include_transformed)
     unobserved_vars_values = model.compile_fn(unobserved_vars, mode="FAST_COMPILE")(
         DictToArrayBijection.rmap(raveled_optimized)
     )
@@ -335,20 +335,13 @@ def find_MAP(
         var.name: value for var, value in zip(unobserved_vars, unobserved_vars_values)
     }
 
-    idata = map_results_to_inference_data(
-        map_point=optimized_point, model=frozen_model, include_transformed=include_transformed
-    )
-
-    idata = add_fit_to_inference_data(
-        idata=idata, mu=raveled_optimized, H_inv=H_inv, model=frozen_model
-    )
-
+    idata = map_results_to_inference_data(optimized_point, frozen_model)
+    idata = add_fit_to_inference_data(idata, raveled_optimized, H_inv)
     idata = add_optimizer_result_to_inference_data(
-        idata=idata, result=optimizer_result, method=method, mu=raveled_optimized, model=model
+        idata, optimizer_result, method, raveled_optimized, model
     )
-
     idata = add_data_to_inference_data(
-        idata=idata, progressbar=False, model=model, compile_kwargs=compile_kwargs
+        idata, progressbar=False, model=model, compile_kwargs=compile_kwargs
     )
 
     return idata

pymc_extras/inference/laplace.py

@@ -59,7 +59,6 @@ def make_unpacked_variable_names(names: list[str], model: pm.Model) -> list[str]
 def map_results_to_inference_data(
     map_point: dict[str, float | int | np.ndarray],
     model: pm.Model | None = None,
-    include_transformed: bool = True,
 ):
     """
     Add the MAP point to an InferenceData object in the posterior group.
@@ -69,13 +68,13 @@ def map_results_to_inference_data(
 
     Parameters
     ----------
+    idata: az.InferenceData
+        An InferenceData object to which the MAP point will be added.
     map_point: dict
         A dictionary containing the MAP point estimates for each variable. The keys should be the variable names, and
         the values should be the corresponding MAP estimates.
     model: Model, optional
         A PyMC model. If None, the model is taken from the current model context.
-    include_transformed: bool
-        Whether to return transformed (unconstrained) variables in the constrained_posterior group. Default is True.
 
     Returns
     -------
@@ -119,7 +118,7 @@ def map_results_to_inference_data(
         dims=dims,
     )
 
-    if unconstrained_names and include_transformed:
+    if unconstrained_names:
         unconstrained_posterior = az.from_dict(
             posterior={
                 k: np.expand_dims(v, (0, 1))

pymc_extras/inference/laplace_approx/find_map.py

@@ -302,7 +302,7 @@ def fit_laplace(
     ----------
     model : pm.Model
         The PyMC model to be fit. If None, the current model context is used.
-    optimize_method : str
+    method : str
         The optimization method to use. Valid choices are: Nelder-Mead, Powell, CG, BFGS, L-BFGS-B, TNC, SLSQP,
         trust-constr, dogleg, trust-ncg, trust-exact, trust-krylov, and basinhopping.
 
@@ -441,11 +441,9 @@ def fit_laplace(
             .rename({"temp_chain": "chain", "temp_draw": "draw"})
         )
 
-        if include_transformed:
-            idata.unconstrained_posterior = unstack_laplace_draws(
-                new_posterior.laplace_approximation.values, model, chains=chains, draws=draws
-            )
-
+        idata.unconstrained_posterior = unstack_laplace_draws(
+            new_posterior.laplace_approximation.values, model, chains=chains, draws=draws
+        )
         idata.posterior = new_posterior.drop_vars(
             ["laplace_approximation", "unpacked_variable_names"]
         )

pymc_extras/inference/laplace_approx/idata.py

@@ -36,7 +36,7 @@ dynamic = ["version"]  # specify the version in the __init__.py file
 dependencies = [
   "pymc>=5.21.1",
   "scikit-learn",
-  "better-optimize>=0.1.2",
+  "better-optimize>=0.1.4",
   "pydantic>=2.0.0",
 ]
 

pymc_extras/inference/laplace_approx/laplace.py

@@ -133,19 +133,12 @@ def compute_z(x):
     ],
 )
 @pytest.mark.parametrize(
-    "backend, gradient_backend, include_transformed",
-    [("jax", "jax", True), ("jax", "pytensor", False)],
+    "backend, gradient_backend",
+    [("jax", "jax"), ("jax", "pytensor")],
     ids=str,
 )
 def test_find_MAP(
-    method,
-    use_grad,
-    use_hess,
-    use_hessp,
-    backend,
-    gradient_backend: GradientBackend,
-    include_transformed,
-    rng,
+    method, use_grad, use_hess, use_hessp, backend, gradient_backend: GradientBackend, rng
 ):
     pytest.importorskip("jax")
 
@@ -161,12 +154,12 @@ def test_find_MAP(
             use_hessp=use_hessp,
             progressbar=False,
             gradient_backend=gradient_backend,
-            include_transformed=include_transformed,
             compile_kwargs={"mode": backend.upper()},
             maxiter=5,
         )
 
     assert hasattr(idata, "posterior")
+    assert hasattr(idata, "unconstrained_posterior")
     assert hasattr(idata, "fit")
     assert hasattr(idata, "optimizer_result")
     assert hasattr(idata, "observed_data")
@@ -176,29 +169,9 @@ def test_find_MAP(
     assert posterior["mu"].shape == ()
     assert posterior["sigma"].shape == ()
 
-    if include_transformed:
-        assert hasattr(idata, "unconstrained_posterior")
-        unconstrained_posterior = idata.unconstrained_posterior.squeeze(["chain", "draw"])
-        assert "sigma_log__" in unconstrained_posterior
-        assert unconstrained_posterior["sigma_log__"].shape == ()
-    else:
-        assert not hasattr(idata, "unconstrained_posterior")
-
-
-def test_find_map_outside_model_context():
-    """
-    Test that find_MAP can be called outside of a model context.
-    """
-    with pm.Model() as m:
-        mu = pm.Normal("mu", 0, 1)
-        sigma = pm.Exponential("sigma", 1)
-        y_hat = pm.Normal("y_hat", mu=mu, sigma=sigma, observed=np.random.normal(size=10))
-
-    idata = find_MAP(model=m, method="L-BFGS-B", use_grad=True, progressbar=False)
-
-    assert hasattr(idata, "posterior")
-    assert hasattr(idata, "fit")
-    assert hasattr(idata, "optimizer_result")
+    unconstrained_posterior = idata.unconstrained_posterior.squeeze(["chain", "draw"])
+    assert "sigma_log__" in unconstrained_posterior
+    assert unconstrained_posterior["sigma_log__"].shape == ()
 
 
 @pytest.mark.parametrize(

pyproject.toml

@@ -83,32 +83,7 @@ def test_fit_laplace_basic(mode, gradient_backend: GradientBackend):
     np.testing.assert_allclose(idata.fit["covariance_matrix"].values, bda_cov, rtol=1e-3, atol=1e-3)
 
 
-def test_fit_laplace_outside_model_context():
-    with pm.Model() as m:
-        mu = pm.Normal("mu", 0, 1)
-        sigma = pm.Exponential("sigma", 1)
-        y_hat = pm.Normal("y_hat", mu=mu, sigma=sigma, observed=np.random.normal(size=10))
-
-    idata = fit_laplace(
-        model=m,
-        optimize_method="L-BFGS-B",
-        use_grad=True,
-        progressbar=False,
-        chains=1,
-        draws=100,
-    )
-
-    assert hasattr(idata, "posterior")
-    assert hasattr(idata, "fit")
-    assert hasattr(idata, "optimizer_result")
-
-    assert idata.posterior["mu"].shape == (1, 100)
-
-
-@pytest.mark.parametrize(
-    "include_transformed", [True, False], ids=["include_transformed", "no_transformed"]
-)
-def test_fit_laplace_coords(include_transformed, rng):
+def test_fit_laplace_coords(rng):
     coords = {"city": ["A", "B", "C"], "obs_idx": np.arange(100)}
     with pm.Model(coords=coords) as model:
         mu = pm.Normal("mu", mu=3, sigma=0.5, dims=["city"])
@@ -127,7 +102,6 @@ def test_fit_laplace_coords(include_transformed, rng):
             chains=1,
             draws=1000,
             optimizer_kwargs=dict(tol=1e-20),
-            include_transformed=include_transformed,
         )
 
     np.testing.assert_allclose(
@@ -146,11 +120,6 @@ def test_fit_laplace_coords(include_transformed, rng):
         "sigma_log__[C]",
     ]
 
-    assert hasattr(idata, "unconstrained_posterior") == include_transformed
-    if include_transformed:
-        assert "sigma_log__" in idata.unconstrained_posterior
-        assert "city" in idata.unconstrained_posterior.coords
-
 
 def test_fit_laplace_ragged_coords(rng):
     coords = {"city": ["A", "B", "C"], "feature": [0, 1], "obs_idx": np.arange(100)}
@@ -230,50 +199,6 @@ def test_model_with_nonstandard_dimensionality(rng):
     assert "class" in list(idata.unconstrained_posterior.sigma_log__.coords.keys())
 
 
-def test_laplace_nonstandard_dims_2d():
-    true_P = np.array([[0.5, 0.3, 0.2], [0.1, 0.6, 0.3], [0.2, 0.4, 0.4]])
-    y_obs = pm.draw(
-        pmx.DiscreteMarkovChain.dist(
-            P=true_P,
-            init_dist=pm.Categorical.dist(
-                logit_p=np.ones(
-                    3,
-                )
-            ),
-            shape=(100, 5),
-        )
-    )
-
-    with pm.Model(
-        coords={
-            "time": range(y_obs.shape[0]),
-            "state": list("ABC"),
-            "next_state": list("ABC"),
-            "unit": [1, 2, 3, 4, 5],
-        }
-    ) as model:
-        y = pm.Data("y", y_obs, dims=["time", "unit"])
-        init_dist = pm.Categorical.dist(
-            logit_p=np.ones(
-                3,
-            )
-        )
-        P = pm.Dirichlet("P", a=np.eye(3) * 2 + 1, dims=["state", "next_state"])
-        y_hat = pmx.DiscreteMarkovChain(
-            "y_hat", P=P, init_dist=init_dist, dims=["time", "unit"], observed=y_obs
-        )
-
-        idata = pmx.fit_laplace(progressbar=True)
-
-        # The simplex transform should drop from the right-most dimension, so the left dimension should be unmodified
-        assert "state" in list(idata.unconstrained_posterior.P_simplex__.coords.keys())
-
-        # The mutated dimension should be unknown coords
-        assert "P_simplex___dim_1" in list(idata.unconstrained_posterior.P_simplex__.coords.keys())
-
-        assert idata.unconstrained_posterior.P_simplex__.shape[-2:] == (3, 2)
-
-
 def test_laplace_nonscalar_rv_without_dims():
     with pm.Model(coords={"test": ["A", "B", "C"]}) as model:
         x_loc = pm.Normal("x_loc", mu=0, sigma=1, dims=["test"])