PyMC/PyTensor Implementation of Pathfinder VI (#387)

* renamed samples argument name and pathfinder variables to avoid confusion

* Minor changes made to the `fit_pathfinder` function and added test

`fit_pathfinder`
- Edited `fit_pathfinder` to produce `pathfinder_state`, `pathfinder_info`, `pathfinder_samples` and `pathfinder_idata` for closer examination of the outputs.
- Changed the `num_samples` argument name to `num_draws` to avoid `TypeError` got multiple values for keyword argument 'num_samples'.
- Initial points are automatically set to jitter as jitter is required for pathfinder.

Extras
- New function 'get_jaxified_logp_ravel_inputs' to simplify previous code structure in fit_pathfinder.

Tests
- Added extra test for pathfinder to test pathfinder_info variables and pathfinder_idata  are consistent for a given random seed.

* extract additional pathfinder objects from high level API for debugging

* changed pathfinder samples argument to  num_draws

* feat(pathfinder): add PyMC-based Pathfinder VI implementation

Add a new PyMC-based implementation of Pathfinder VI that uses PyTensor operations which provides support for both PyMC and BlackJAX backends in fit_pathfinder.

* Multipath Pathfinder VI implementation in pymc-experimental

- Implemented  in  to support running multiple Pathfinder instances in parallel.
- Implemented  function in  for Pareto Smoothed Importance Resampling (PSIR).
- Moved relevant pathfinder files into the  directory.
- Updated tests to reflect changes in the Pathfinder implementation and added tests for new functionalities.

* Added type hints and epsilon parameter to fit_pathfinder

* Removed initial point values (l=0) to reduce iterations. Simplified  and .

* Added placeholder/reminder to remove jax dependency when converting trace data to InferenceData

* Sync updates with draft PR #386. \n- Added pytensor.function for bfgs_sample

* Reduced size of compute graph with pathfinder_body_fn

Summaryh of changes:
- Remove multiprocessing code in favour of reusing compiled  for each path
-  takes only random_seed as argument for each path
- Compute graph significantly smaller by using pure pytensor op and symoblic variables
- Added LBFGSOp to compile with pytensor.function
- Cleaned up codes using pytensor variables

* - Added TODO comments for implementing Taylor approximation methods:  and .
- Corrected the dimensions in comments for matrices Q and R in the  function.
- Uumerical stability in the  calculation by changing from  to .

* fix: correct posterior approximations in Pathfinder VI

Fixed incorrect and inconsistent posterior approximations in the Pathfinder VI
algorithm by:

1. Adding missing parentheses in the phi calculation to ensure proper order
   of operations in matrix multiplications
2. Changing the sign in mu calculation from 'x +' to 'x -' to match Stan's
   implementation (which differs from the original paper)

The resulting changes now make the posterior approximations more reliable.

* feat: Add dense BFGS sampling for Pathfinder VI

Implements both sparse and dense BFGS sampling approaches for Pathfinder VI:
- Adds bfgs_sample_dense for cases where 2*maxcor >= num_params.
- Moved existing  and  computations to bfgs_sample_sparse, making the sparse use cases more explicit.

Other changes:
- Sets default maxcor=5 instead of dynamic sizing based on parameters

Dense approximations are recommended when the target distribution has higher dependencies among the parameters.

* feat: improve Pathfinder performance and compatibility

Bigger changes:
- Made pmx.fit compatible with method='pathfinder'
- Remove JAX dependency when inference_backend='pymc' to support Windows users
- Improve runtime performance by setting trust_input=True for compiled functions

Minor changes:
- Change default num_paths from 1 to 4 for stable and reliable approximations
- Change LBFGS code using dataclasses
- Update tests to handle both PyMC and BlackJAX backends

* minor: improve error handling in Pathfinder VI

- Add LBFGSInitFailed exception for failed LBFGS initialisation
- Skip failed paths in multipath_pathfinder and track number of failures
- Handle NaN values from Cholesky decompsition in bfgs_sample
- Add checks for numericl stabilty in matrix operations

Slight performance improvements:
- Set allow_gc=False in scan ops
- Use FAST_RUN mode consistently

* Progress bar and other minor changes

Major:
  - Added progress bar support.

Minor
  - Added  exception for non-finite log prob values
  - Removed .
  - Allowed maxcor argument to be None, and dynamically set based on the number of model parameters.
  - Improved logging to inform users about failed paths and lbfgs initialisation.

* set maxcor to max(5, floor(N / 1.9)). max=1 will cause error

* Refactor Pathfinder VI: Default to PSIS, Add Concurrency, and Improved Computational Performance

- Significantly computational efficiency by combining 3 computational graphs into 1 larger compile. Removed non-shared inputs and used  with  for significant performance gains.
- Set default importance sampling method to 'psis' for more stable posterior results, avoiding local peaks seen with 'psir'.
- Introduce concurrency options ('thread' and 'process') for multithreading and multiprocessing. Defaults to No concurrency as there haven't been any/or much reduction to the compute time.
- Adjusted default  from 8 to 4 and  from 1.0 to 2.0 and maxcor to max(3*log(N), 5). This default setting lessens computational time and and the degree by which the posterior variance is being underestimated.

* Improvements to Importance Sampling and InferenceData shape

- Handle different importance sampling methods for reshaping and adjusting log densities.
- Modified  to return InferenceData with chain dim of size num_paths when

* Display summary of results, Improve error handling, General improvements

Changes:
- Add rich table summary display for results
- Added PathStatus and LBFGSStatus for error handling, status tracking and displaying results
- Changed importance_sampling return type to ImportanceSamplingResult
- Changed multipath_pathfinder return type to MultiPathfinderResult
- Added dataclass containers for results (ImportanceSamplingResult, PathfinderResult, MultiPathfinderResult)
- Refactored LBFGS by removing PyTensor Op classes in favor of pure functions
- Added timing and configuration tracking
- Improve concurrency with better error handling
- Improved docstrings and type hints
- Simplified logp and gradient computation by combining into single function
- Added compile_kwargs parameter for pytensor compilation options

* Move pathfinder module to pymc_extras

- Move pathfinder module from pymc_experimental to pymc_extras
- Update directory structure to match upstream repository

* 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

* fix: Use typing_extensions.Self for Python 3.10 compatibility

Author: aphc14

pymc_extras/inference/fit.py

@@ -11,7 +11,6 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
-from importlib.util import find_spec
 
 
 def fit(method, **kwargs):
@@ -31,9 +30,6 @@ def fit(method, **kwargs):
     arviz.InferenceData
     """
     if method == "pathfinder":
-        if find_spec("blackjax") is None:
-            raise RuntimeError("Need BlackJAX to use `pathfinder`")
-
         from pymc_extras.inference.pathfinder import fit_pathfinder
 
         return fit_pathfinder(**kwargs)

pymc_extras/inference/pathfinder.py

@@ -0,0 +1,3 @@
+from pymc_extras.inference.pathfinder.pathfinder import fit_pathfinder
+
+__all__ = ["fit_pathfinder"]

pymc_extras/inference/pathfinder/__init__.py

@@ -0,0 +1,139 @@
+import logging
+import warnings as _warnings
+
+from dataclasses import dataclass, field
+from typing import Literal
+
+import arviz as az
+import numpy as np
+
+from numpy.typing import NDArray
+from scipy.special import logsumexp
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class ImportanceSamplingResult:
+    """container for importance sampling results"""
+
+    samples: NDArray
+    pareto_k: float | None = None
+    warnings: list[str] = field(default_factory=list)
+    method: str = "none"
+
+
+def importance_sampling(
+    samples: NDArray,
+    logP: NDArray,
+    logQ: NDArray,
+    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.
+
+    Parameters
+    ----------
+    samples : NDArray
+        samples from proposal distribution, shape (L, M, N)
+    logP : NDArray
+        log probability values of target distribution, shape (L, M)
+    logQ : NDArray
+        log probability values of proposal distribution, shape (L, M)
+    num_draws : int
+        number of draws to return where num_draws <= samples.shape[0]
+    method : str, optional
+        importance sampling method to use. Options are "psis" (default), "psir", "identity", "none. Pareto Smoothed Importance Sampling (psis) is recommended in many cases for more stable results than Pareto Smoothed Importance Resampling (psir). identity applies the log importance weights directly without resampling. none applies no importance sampling weights and returns the samples as is of size num_draws_per_path * num_paths.
+    random_seed : int | None
+
+    Returns
+    -------
+    ImportanceSamplingResult
+        importance sampled draws and other info based on the specified method
+
+    Future work!
+    ----------
+    - Implement the 3 sampling approaches and 5 weighting functions from Elvira et al. (2019)
+    - Implement Algorithm 2 VSBC marginal diagnostics from Yao et al. (2018)
+    - Incorporate these various diagnostics, sampling approaches and weighting functions into VI algorithms.
+
+    References
+    ----------
+    Elvira, V., Martino, L., Luengo, D., & Bugallo, M. F. (2019). Generalized Multiple Importance Sampling. Statistical Science, 34(1), 129-155. https://doi.org/10.1214/18-STS668
+
+    Yao, Y., Vehtari, A., Simpson, D., & Gelman, A. (2018). Yes, but Did It Work?: Evaluating Variational Inference. arXiv:1802.02538 [Stat]. http://arxiv.org/abs/1802.02538
+
+    Zhang, L., Carpenter, B., Gelman, A., & Vehtari, A. (2022). Pathfinder: Parallel quasi-Newton variational inference. Journal of Machine Learning Research, 23(306), 1-49.
+    """
+
+    warnings = []
+    num_paths, _, N = samples.shape
+
+    if method == "none":
+        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=warnings)
+    else:
+        samples = samples.reshape(-1, N)
+        logP = logP.ravel()
+        logQ = logQ.ravel()
+
+        # adjust log densities
+        log_I = np.log(num_paths)
+        logP -= log_I
+        logQ -= log_I
+        logiw = logP - logQ
+
+        with _warnings.catch_warnings():
+            _warnings.filterwarnings(
+                "ignore", category=RuntimeWarning, message="overflow encountered in exp"
+            )
+            if method == "psis":
+                replace = False
+                logiw, pareto_k = az.psislw(logiw)
+            elif method == "psir":
+                replace = True
+                logiw, pareto_k = az.psislw(logiw)
+            elif method == "identity":
+                replace = False
+                pareto_k = None
+            else:
+                raise ValueError(f"Invalid importance sampling method: {method}")
+
+    # NOTE: Pareto k is normally bad for Pathfinder even when the posterior is close to the NUTS posterior or closer to NUTS than ADVI.
+    # Pareto k may not be a good diagnostic for Pathfinder.
+    # TODO: Find replacement diagnostics for Pathfinder.
+
+    p = np.exp(logiw - logsumexp(logiw))
+    rng = np.random.default_rng(random_seed)
+
+    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=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()
+            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=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 for both with and without replacement"
+                ) from e2
+        raise