def vcycle(
    A,
    shape,
    num_pre_iter=1,
    num_post_iter=1,
    omega=0.8,
    direct_solve_size=64,
    cache=None,
):
    """
    Implements the V-Cycle preconditioner.
    The V-Cycle solver was recommended by :cite:`lee2014scalable` to solve the alpha matting problem.

    Parameters
    ----------
    A: numpy.ndarray
        Input matrix
    shape: tuple of ints
        Describing the height and width of the image
    num_pre_iter: int
        Number of Jacobi iterations before each V-Cycle, defaults to 1
    num_post_iter: int
        Number of Jacobi iterations after each V-Cycle, defaults to 1
    omega: float
        Weight parameter for the Jacobi method. If method fails to converge, try different values.

    Returns
    -------
    precondition: function
        Function which applies the V-Cycle preconditioner to a vector

    Example
    -------
    >>> from pymatting import *
    >>> import numpy as np
    >>> from scipy.sparse import csc_matrix
    >>> A = np.array([[2, 3], [3, 5]])
    >>> preconditioner = vcycle(A, (2, 2))
    >>> preconditioner(np.array([1, 2]))
    array([-1.,  1.])
    """

    if cache is None:
        cache = {}

    def precondition(r):
        return _vcycle_step(
            A, r, shape, cache, num_pre_iter, num_post_iter, omega, direct_solve_size
        )

    return precondition

    Implements the V-Cycle preconditioner.
    The V-Cycle solver was recommended by :cite:`lee2014scalable` to solve the alpha matting problem.

    Parameters
    ----------
    A: numpy.ndarray
        Input matrix
    shape: tuple of ints
        Describing the height and width of the image
    num_pre_iter: int
        Number of Jacobi iterations before each V-Cycle, defaults to 1
    num_post_iter: int
        Number of Jacobi iterations after each V-Cycle, defaults to 1
    omega: float
        Weight parameter for the Jacobi method. If method fails to converge, try different values.

    Returns
    -------
    precondition: function
        Function which applies the V-Cycle preconditioner to a vector

    Example
    -------
    >>> from pymatting import *
    >>> import numpy as np
    >>> from scipy.sparse import csc_matrix
    >>> A = np.array([[2, 3], [3, 5]])
    >>> preconditioner = vcycle(A, (2, 2))
    >>> preconditioner(np.array([1, 2]))
    array([-1.,  1.])
    
pymatting/preconditioner/vcycle.py