Fisher¤

def covariance_matrix(
    pytree: PyTree,
    parameters: Params,
    loglike_fn: callable,
    *loglike_args: Any,
    shape_dict: dict = {},
    save_memory: bool = False,
    **loglike_kwargs: Any,
) -> Array:
    """
    Calculates the covariance matrix of the pytree parameters. The
    `shaped_dict` parameter is used to specify the shape of the differentiated
    vector for specific parameters. For example, if the parameter `param` is
    a 1D array of shape (5,) and we wanted to calculate the covariance matrix
    of the mean, we can pass in `shape_dict={'param': (1,)}`. This will
    differentiate the log likelihood with respect to the mean of the parameters.

    Parameters
    ----------
    pytree : PyTree
        Pytree with a .model() function.
    parameters : Params
        A path or list of paths or list of nested paths.
    loglike_fn : callable
        The log likelihood function to differentiate.
    shape_dict : dict = {}
        A dictionary specifying the shape of the differentiated vector for
        specific parameters.
    save_memory : bool = False
        If True, the Hessian is calculated column by column to save memory.
        This is slightly slower.
    *loglike_args : Any
        The args to pass to the log likelihood function.
    **loglike_kwargs : Any
        The kwargs to pass to the log likelihood function.

    Returns
    -------
    covariance_matrix : Array
        The covariance matrix of the pytree parameters.
    """
    warnings.warn(
        "covariance_matrix is deprecated as of v0.5.0 and will be removed in v0.6.0",
        DeprecationWarning,
    )

    return np.linalg.inv(
        fisher_matrix(
            pytree,
            parameters,
            loglike_fn,
            shape_dict=shape_dict,
            save_memory=save_memory,
            *loglike_args,
            **loglike_kwargs,
        )
    )

def fisher_matrix(
    pytree: PyTree,
    parameters: Params,
    loglike_fn: callable,
    *loglike_args,
    shape_dict: dict = {},
    save_memory: bool = False,
    **loglike_kwargs,
) -> Array:
    """
    Calculates the Fisher information matrix of the log likelihood function with
    respect to the parameters of the pytree. It is evaluated at the current values
    of the parameters as listed in the pytree. Simply returns the negative Hessian.
    """
    warnings.warn(
        "fisher_matrix is deprecated as of v0.5.0 and will be removed in v0.6.0",
        DeprecationWarning,
    )

    return -hessian(
        pytree,
        parameters,
        loglike_fn,
        shape_dict=shape_dict,
        save_memory=save_memory,
        *loglike_args,
        **loglike_kwargs,
    )

def hessian(
    pytree: PyTree,
    parameters: Params,
    fn: callable,
    *fn_args,
    shape_dict: dict = {},
    save_memory: bool = False,
    **fn_kwargs,
) -> Array:
    """
    Calculates the Hessian of the function with respect to the parameters of the pytree.
    It is evaluated at the current values of the parameters as listed in the pytree.

    Parameters
    ----------
    pytree : PyTree
        Pytree holding the parameters values.
    parameters : Params
        Names of parameters to be used in Hessian calculation.
    fn : callable
        The function to differentiate.
    *fn_args : Any
        The args to pass to the log likelihood function.
    shape_dict : dict = {}
        A dictionary specifying the shape of the differentiated vector for
        specific parameters.
    save_memory : bool = False
        If True, the Hessian is calculated column by column to save memory.
    **fn__kwargs : Any
        The kwargs to pass to the log likelihood function.

    Returns
    -------
    hessian : Array
        The Hessian of the log likelihood function with respect to the
        parameters of the pytree.
    """
    warnings.warn(
        "fisher.hessian is deprecated as of v0.5.0 and will be removed in v0.6.0",
        DeprecationWarning,
    )

    # If only one parameter is passed, make it a list
    if len(parameters) == 1:
        parameters = [parameters]

    # Build empty vector to perturb
    pytree = set_array(pytree)
    shapes, lengths = _shapes_and_lengths(pytree, parameters, shape_dict)
    X = np.zeros(_lengths_to_N(lengths))

    # Build perturbation function to differentiate
    def loglike_fn_vec(X):
        parametric_pytree = _perturb(X, pytree, parameters, shapes, lengths)
        return fn(parametric_pytree, *fn_args, **fn_kwargs)

    # optional column by column hessian calculation for RAM mercy
    if save_memory:
        return _hessian_col_by_col(loglike_fn_vec, X)

    elif not save_memory:
        return jax.hessian(loglike_fn_vec)(X)