Parameter draws is not optional with default value 1000

Author: carsten-j

pymc_experimental/inference/laplace.py

@@ -32,7 +32,7 @@
 
 def laplace(
     vars: Sequence[Variable],
-    draws=1_000,
+    draws: Optional[int] = 1000,
     model=None,
     random_seed: Optional[RandomSeed] = None,
     progressbar=True,
@@ -49,9 +49,9 @@ def laplace(
     vars : Sequence[Variable]
         A sequence of variables for which the Laplace approximation of the posterior distribution
         is to be created.
-    draws : int, optional, default=1_000
+    draws : Optional[int] with default=1_000
         The number of draws to sample from the posterior distribution for creating the approximation.
-        For draws=0 only the fit of the Laplace approximation is returned
+        For draws=None only the fit of the Laplace approximation is returned
     model : object, optional, default=None
         The model object that defines the posterior distribution. If None, the default model will be used.
     random_seed : Optional[RandomSeed], optional, default=None
@@ -103,7 +103,8 @@ def laplace(
     # See https://www.pymc.io/projects/docs/en/stable/api/model/generated/pymc.model.transform.conditioning.remove_value_transforms.html
     untransformed_m = remove_value_transforms(transformed_m)
     untransformed_vars = [untransformed_m[v.name] for v in vars]
-    hessian = pm.find_hessian(point=map, vars=untransformed_vars, model=untransformed_m)
+    hessian = pm.find_hessian(
+        point=map, vars=untransformed_vars, model=untransformed_m)
 
     if np.linalg.det(hessian) == 0:
         raise np.linalg.LinAlgError("Hessian is singular.")
@@ -113,12 +114,13 @@ def laplace(
 
     chains = 1
 
-    if draws != 0:
+    if draws is not None:
         samples = rng.multivariate_normal(mean, cov, size=(chains, draws))
 
         data_vars = {}
         for i, var in enumerate(vars):
-            data_vars[str(var)] = xr.DataArray(samples[:, :, i], dims=("chain", "draw"))
+            data_vars[str(var)] = xr.DataArray(
+                samples[:, :, i], dims=("chain", "draw"))
 
         coords = {"chain": np.arange(chains), "draw": np.arange(draws)}
         ds = xr.Dataset(data_vars, coords=coords)
@@ -136,13 +138,15 @@ def laplace(
 def addFitToInferenceData(vars, idata, mean, covariance):
     coord_names = [v.name for v in vars]
     # Convert to xarray DataArray
-    mean_dataarray = xr.DataArray(mean, dims=["rows"], coords={"rows": coord_names})
+    mean_dataarray = xr.DataArray(
+        mean, dims=["rows"], coords={"rows": coord_names})
     cov_dataarray = xr.DataArray(
         covariance, dims=["rows", "columns"], coords={"rows": coord_names, "columns": coord_names}
     )
 
     # Create xarray dataset
-    dataset = xr.Dataset({"mean_vector": mean_dataarray, "covariance_matrix": cov_dataarray})
+    dataset = xr.Dataset({"mean_vector": mean_dataarray,
+                         "covariance_matrix": cov_dataarray})
 
     idata.add_groups(fit=dataset)
 

pymc_experimental/tests/test_laplace.py

@@ -60,4 +60,44 @@ def test_laplace():
     bda_cov = np.array([[y.var() / n, 0], [0, 1 / (2 * n)]])
 
     assert np.allclose(idata.fit["mean_vector"].values, bda_map)
-    assert np.allclose(idata.fit["covariance_matrix"].values, bda_cov, atol=1e-4)
+    assert np.allclose(
+        idata.fit["covariance_matrix"].values, bda_cov, atol=1e-4)
+
+
+@pytest.mark.filterwarnings(
+    "ignore:Model.model property is deprecated. Just use Model.:FutureWarning",
+    "ignore:hessian will stop negating the output in a future version of PyMC.\n"
+    + "To suppress this warning set `negate_output=False`:FutureWarning",
+)
+def test_laplace_only_fit():
+
+    # Example originates from Bayesian Data Analyses, 3rd Edition
+    # By Andrew Gelman, John Carlin, Hal Stern, David Dunson,
+    # Aki Vehtari, and Donald Rubin.
+    # See section. 4.1
+
+    y = np.array([2642, 3503, 4358], dtype=np.float64)
+    n = y.size
+
+    with pm.Model() as m:
+        logsigma = pm.Uniform("logsigma", 1, 100)
+        mu = pm.Uniform("mu", -10000, 10000)
+        yobs = pm.Normal("y", mu=mu, sigma=pm.math.exp(logsigma), observed=y)
+        vars = [mu, logsigma]
+
+    idata = pmx.fit(
+        method="laplace",
+        vars=vars,
+        model=m,
+        random_seed=173300,
+    )
+
+    assert idata.fit["mean_vector"].shape == (len(vars),)
+    assert idata.fit["covariance_matrix"].shape == (len(vars), len(vars))
+
+    bda_map = [y.mean(), np.log(y.std())]
+    bda_cov = np.array([[y.var() / n, 0], [0, 1 / (2 * n)]])
+
+    assert np.allclose(idata.fit["mean_vector"].values, bda_map)
+    assert np.allclose(
+        idata.fit["covariance_matrix"].values, bda_cov, atol=1e-4)