find_MAP#

pymc_extras.inference.find_MAP(method: Literal['nelder-mead', 'powell', 'CG', 'BFGS', 'Newton-CG', 'L-BFGS-B', 'TNC', 'COBYLA', 'SLSQP', 'trust-constr', 'dogleg', 'trust-ncg', 'trust-exact', 'trust-krylov', 'basinhopping'] = 'L-BFGS-B', *, model: Model | None = None, use_grad: bool | None = None, use_hessp: bool | None = None, use_hess: bool | None = None, initvals: dict | None = None, random_seed: int | Generator | None = None, jitter_rvs: list[TensorVariable] | None = None, progressbar: bool = True, include_transformed: bool = True, freeze_model: bool = True, gradient_backend: Literal['pytensor', 'jax'] = 'pytensor', compile_kwargs: dict | None = None, compute_hessian: bool = False, **optimizer_kwargs) → dict[str, ndarray] | tuple[dict[str, ndarray], ndarray] | tuple[dict[str, ndarray], OptimizeResult] | tuple[dict[str, ndarray], OptimizeResult, ndarray][source]#

Fit a PyMC model via maximum a posteriori (MAP) estimation using JAX and scipy.optimize.

Parameters:

  • model (pm.Model) – The PyMC model to be fit. If None, the current model context is used.

  • 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.

    See scipy.optimize.minimize documentation for details.

  • use_grad (bool | None, optional) – Whether to use gradients in the optimization. Defaults to None, which determines this automatically based on the method.

  • use_hessp (bool | None, optional) – Whether to use Hessian-vector products in the optimization. Defaults to None, which determines this automatically based on the method.

  • use_hess (bool | None, optional) – Whether to use the Hessian matrix in the optimization. Defaults to None, which determines this automatically based on the method.

  • initvals (None | dict, optional) –

    Initial values for the model parameters, as str:ndarray key-value pairs. Partial initialization is permitted.

    If None, the model’s default initial values are used.

  • random_seed (None | int | np.random.Generator, optional) – Seed for the random number generator or a numpy Generator for reproducibility

  • jitter_rvs (list of TensorVariables, optional) – Variables whose initial values should be jittered. If None, all variables are jittered.

  • progressbar (bool, optional) – Whether to display a progress bar during optimization. Defaults to True.

  • include_transformed (bool, optional) – Whether to include transformed variable values in the returned dictionary. Defaults to True.

  • freeze_model (bool, optional) – If True, freeze_dims_and_data will be called on the model before compiling the loss functions. This is sometimes necessary for JAX, and can sometimes improve performance by allowing constant folding. Defaults to True.

  • gradient_backend (str, default "pytensor") – Which backend to use to compute gradients. Must be one of “pytensor” or “jax”.

  • compute_hessian (bool) – If True, the inverse Hessian matrix at the optimum will be computed and included in the returned DataTree object. This is needed for the Laplace approximation, but can be computationally expensive for high-dimensional problems. Defaults to False.

  • compile_kwargs (dict, optional) – Additional options to pass to the pytensor.function function when compiling loss functions.

  • **optimizer_kwargs – Additional keyword arguments to pass to the scipy.optimize function being used. Unless method = "basinhopping", scipy.optimize.minimize will be used. For basinhopping, scipy.optimize.basinhopping will be used. See the documentation of these functions for details.

Returns:

map_result – Results of Maximum A Posteriori (MAP) estimation, including the optimized point, inverse Hessian, transformed latent variables, and optimizer results.

Return type:

DataTree