Make fit_pathfinder more similar to fit_laplace and pm.sample

docs/api_reference.rst

@@ -23,7 +23,10 @@ Inference
 .. autosummary::
    :toctree: generated/
 
+   find_MAP
    fit
+   fit_laplace
+   fit_pathfinder
 
 
 Distributions

pymc_extras/__init__.py

@@ -15,9 +15,7 @@
 
 from pymc_extras import gp, statespace, utils
 from pymc_extras.distributions import *
-from pymc_extras.inference.find_map import find_MAP
-from pymc_extras.inference.fit import fit
-from pymc_extras.inference.laplace import fit_laplace
+from pymc_extras.inference import find_MAP, fit, fit_laplace, fit_pathfinder
 from pymc_extras.model.marginal.marginal_model import (
     MarginalModel,
     marginalize,

pymc_extras/inference/__init__.py

@@ -12,7 +12,9 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
-
+from pymc_extras.inference.find_map import find_MAP
 from pymc_extras.inference.fit import fit
+from pymc_extras.inference.laplace import fit_laplace
+from pymc_extras.inference.pathfinder.pathfinder import fit_pathfinder
 
-__all__ = ["fit"]
+__all__ = ["fit", "fit_pathfinder", "fit_laplace", "find_MAP"]

pymc_extras/inference/fit.py

@@ -15,19 +15,20 @@
 
 def fit(method, **kwargs):
     """
-    Fit a model with an inference algorithm
+    Fit a model with an inference algorithm.
+    See :func:`fit_pathfinder` and :func:`fit_laplace` for more details.
 
     Parameters
     ----------
     method : str
         Which inference method to run.
         Supported: pathfinder or laplace
 
-    kwargs are passed on.
+    kwargs: keyword arguments are passed on to the inference method.
 
     Returns
     -------
-    arviz.InferenceData
+    :class:`~arviz.InferenceData`
     """
     if method == "pathfinder":
         from pymc_extras.inference.pathfinder import fit_pathfinder

pymc_extras/inference/laplace.py

@@ -509,7 +509,7 @@ def fit_laplace(
 
     Returns
     -------
-    idata: az.InferenceData
+    :class:`~arviz.InferenceData`
         An InferenceData object containing the approximated posterior samples.
 
     Examples

pymc_extras/inference/pathfinder/pathfinder.py

@@ -12,6 +12,7 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
+
 import collections
 import logging
 import time
@@ -67,6 +68,7 @@
 # TODO: change to typing.Self after Python versions greater than 3.10
 from typing_extensions import Self
 
+from pymc_extras.inference.laplace import add_data_to_inferencedata
 from pymc_extras.inference.pathfinder.importance_sampling import (
     importance_sampling as _importance_sampling,
 )
@@ -1630,6 +1632,7 @@ def fit_pathfinder(
     inference_backend: Literal["pymc", "blackjax"] = "pymc",
     pathfinder_kwargs: dict = {},
     compile_kwargs: dict = {},
+    initvals: dict | None = None,
 ) -> az.InferenceData:
     """
     Fit the Pathfinder Variational Inference algorithm.
@@ -1665,12 +1668,12 @@ def fit_pathfinder(
     importance_sampling : str, None, optional
         Method to apply sampling based on log importance weights (logP - logQ).
         Options are:
-        "psis" : Pareto Smoothed Importance Sampling (default)
-                Recommended for more stable results.
-        "psir" : Pareto Smoothed Importance Resampling
-                Less stable than PSIS.
-        "identity" : Applies log importance weights directly without resampling.
-        None : No importance sampling weights. Returns raw samples of size (num_paths, num_draws_per_path, N) where N is number of model parameters. Other methods return samples of size (num_draws, N).
+
+        - **"psis"** : Pareto Smoothed Importance Sampling (default). Usually most stable.
+        - **"psir"** : Pareto Smoothed Importance Resampling. Less stable than PSIS.
+        - **"identity"** : Applies log importance weights directly without resampling.
+        - **None** : No importance sampling weights. Returns raw samples of size (num_paths, num_draws_per_path, N) where N is number of model parameters. Other methods return samples of size (num_draws, N).
+
     progressbar : bool, optional
         Whether to display a progress bar (default is True). Setting this to False will likely reduce the computation time.
     random_seed : RandomSeed, optional
@@ -1685,17 +1688,24 @@ def fit_pathfinder(
         Additional keyword arguments for the Pathfinder algorithm.
     compile_kwargs
         Additional keyword arguments for the PyTensor compiler. If not provided, the default linker is "cvm_nogc".
+    initvals: dict | None = None
+        Initial values for the model parameters, as str:ndarray key-value pairs. Paritial initialization is permitted.
+        If None, the model's default initial values are used.
 
     Returns
     -------
-    arviz.InferenceData
+    :class:`~arviz.InferenceData`
         The inference data containing the results of the Pathfinder algorithm.
 
     References
     ----------
     Zhang, L., Carpenter, B., Gelman, A., & Vehtari, A. (2022). Pathfinder: Parallel quasi-Newton variational inference. Journal of Machine Learning Research, 23(306), 1-49.
     """
 
+    if initvals is not None:
+        for rv_name, ivals in initvals.items():
+            model.set_initval(model.named_vars[rv_name], ivals)
+
     model = modelcontext(model)
 
     valid_importance_sampling = {"psis", "psir", "identity", None}
@@ -1775,4 +1785,7 @@ def fit_pathfinder(
         model=model,
         importance_sampling=importance_sampling,
     )
+
+    idata = add_data_to_inferencedata(idata, progressbar, model, compile_kwargs)
+
     return idata