Improve pathfinder error handling and type hints

- Add proper type hints throughout pathfinder module
- Improve error handling in concurrent execution paths
- Better handling of when all paths are fail by displaying results before Assertion
- Changed Australian English spelling to US
- Update compile_pymc usage to handle deprecation warning
- Add tests for concurrent execution and seed reproducibility
- Clean up imports and remove redundant code
- Improve docstrings and error messages

Author: aphc14

pymc_extras/inference/pathfinder/importance_sampling.py

@@ -1,5 +1,5 @@
 import logging
-import warnings
+import warnings as _warnings
 
 from dataclasses import dataclass, field
 from typing import Literal
@@ -30,7 +30,7 @@ def importance_sampling(
     num_draws: int,
     method: Literal["psis", "psir", "identity", "none"] | None,
     random_seed: int | None = None,
-):
+) -> ImportanceSamplingResult:
     """Pareto Smoothed Importance Resampling (PSIR)
     This implements the Pareto Smooth Importance Resampling (PSIR) method, as described in Algorithm 5 of Zhang et al. (2022). The PSIR follows a similar approach to Algorithm 1 PSIS diagnostic from Yao et al., (2018). However, before computing the the importance ratio r_s, the logP and logQ are adjusted to account for the number multiple estimators (or paths). The process involves resampling from the original sample with replacement, with probabilities proportional to the computed importance weights from PSIS.
 
@@ -50,8 +50,8 @@ def importance_sampling(
 
     Returns
     -------
-    NDArray
-        importance sampled draws
+    ImportanceSamplingResult
+        importance sampled draws and other info based on the specified method
 
     Future work!
     ----------
@@ -68,14 +68,14 @@ def importance_sampling(
     Zhang, L., Carpenter, B., Gelman, A., & Vehtari, A. (2022). Pathfinder: Parallel quasi-Newton variational inference. Journal of Machine Learning Research, 23(306), 1-49.
     """
 
-    warning_msgs = []
+    warnings = []
     num_paths, _, N = samples.shape
 
     if method == "none":
-        warning_msgs.append(
+        warnings.append(
             "Importance sampling is disabled. The samples are returned as is which may include samples from failed paths with non-finite logP or logQ values. It is recommended to use importance_sampling='psis' for better stability."
         )
-        return ImportanceSamplingResult(samples=samples, warnings=warning_msgs)
+        return ImportanceSamplingResult(samples=samples, warnings=warnings)
     else:
         samples = samples.reshape(-1, N)
         logP = logP.ravel()
@@ -87,8 +87,8 @@ def importance_sampling(
         logQ -= log_I
         logiw = logP - logQ
 
-        with warnings.catch_warnings():
-            warnings.filterwarnings(
+        with _warnings.catch_warnings():
+            _warnings.filterwarnings(
                 "ignore", category=RuntimeWarning, message="overflow encountered in exp"
             )
             if method == "psis":
@@ -113,27 +113,27 @@ def importance_sampling(
     try:
         resampled = rng.choice(samples, size=num_draws, replace=replace, p=p, shuffle=False, axis=0)
         return ImportanceSamplingResult(
-            samples=resampled, pareto_k=pareto_k, warnings=warning_msgs, method=method
+            samples=resampled, pareto_k=pareto_k, warnings=warnings, method=method
         )
     except ValueError as e1:
         if "Fewer non-zero entries in p than size" in str(e1):
             num_nonzero = np.where(np.nonzero(p)[0], 1, 0).sum()
-            warning_msgs.append(
+            warnings.append(
                 f"Not enough valid samples: {num_nonzero} available out of {num_draws} requested. Switching to psir importance sampling."
             )
             try:
                 resampled = rng.choice(
                     samples, size=num_draws, replace=True, p=p, shuffle=False, axis=0
                 )
                 return ImportanceSamplingResult(
-                    samples=resampled, pareto_k=pareto_k, warnings=warning_msgs, method=method
+                    samples=resampled, pareto_k=pareto_k, warnings=warnings, method=method
                 )
             except ValueError as e2:
                 logger.error(
                     "Importance sampling failed even with psir importance sampling. "
                     "This might indicate invalid probability weights or insufficient valid samples."
                 )
                 raise ValueError(
-                    "Importance sampling failed with both with and without replacement"
+                    "Importance sampling failed for both with and without replacement"
                 ) from e2
         raise

pymc_extras/inference/pathfinder/lbfgs.py

@@ -106,7 +106,7 @@ def __init__(self, message=None):
 
 
 class LBFGS:
-    """L-BFGS optimiser wrapper around scipy's implementation.
+    """L-BFGS optimizer wrapper around scipy's implementation.
 
     Parameters
     ----------
@@ -124,16 +124,18 @@ class LBFGS:
         maximum number of line search steps, defaults to 1000
     """
 
-    def __init__(self, value_grad_fn, maxcor, maxiter=1000, ftol=1e-5, gtol=1e-8, maxls=1000):
+    def __init__(
+        self, value_grad_fn, maxcor, maxiter=1000, ftol=1e-5, gtol=1e-8, maxls=1000
+    ) -> None:
         self.value_grad_fn = value_grad_fn
         self.maxcor = maxcor
         self.maxiter = maxiter
         self.ftol = ftol
         self.gtol = gtol
         self.maxls = maxls
 
-    def minimise(self, x0):
-        """minimises objective function starting from initial position.
+    def minimize(self, x0) -> tuple[NDArray, NDArray, int, LBFGSStatus]:
+        """minimizes objective function starting from initial position.
 
         Parameters
         ----------