def vet_cost_function(
    sector_displacement_1d,
    input_images,
    blocks_shape,
    mask,
    smooth_gain,
    debug=False,
    gradient=False,
):
    """
    .. _`scipy minimization`: \
    https://docs.scipy.org/doc/scipy-0.18.1/reference/generated/scipy.optimize.minimize.html
    
    Variational Echo Tracking Cost Function.

    This function is designed to be used with the `scipy minimization`_.
    The function first argument is the variable to be used in the
    minimization procedure.

    The sector displacement must be a flat array compatible with the
    dimensions of the input image and sectors shape (see parameters section
    below for more details).

    .. _ndarray:\
    https://docs.scipy.org/doc/numpy/reference/generated/numpy.ndarray.html

    Parameters
    ----------
    sector_displacement_1d: ndarray_
        Array of displacements to apply to each sector. The dimensions are:
        sector_displacement_2d
        [ x (0) or y (1) displacement, i index of sector, j index of sector ].
        The shape of the sector displacements must be compatible with the
        input image and the block shape.
        The shape should be (2, mx, my) where mx and my are the numbers of
        sectors in the x and the y dimension.
    input_images: ndarray_
        Input images, sequence of 2D arrays, or 3D arrays.
        The first dimension represents the images time dimension.

        The template_image (first element in first dimensions) denotes the
        reference image used to obtain the displacement (2D array).
        The second is the target image.

        The expected dimensions are (2,nx,ny).
        Be aware the the 2D images dimensions correspond to (lon,lat) or (x,y).
    blocks_shape: ndarray_ (ndim=2)
        Number of sectors in each dimension (x and y).
        blocks_shape.shape = (mx,my)
    mask: ndarray_ (ndim=2)
        Data mask. If is True, the data is marked as not valid and is not
        used in the computations.
    smooth_gain: float
        Smoothness constrain gain
    debug: bool, optional
        If True, print debugging information.
    gradient: bool, optional
        If True, the gradient of the morphing function is returned.

    Returns
    -------
    penalty or gradient values.

    penalty: float
        Value of the cost function
    gradient_values: ndarray (float64 ,ndim = 3), optional
        If gradient keyword is True, the gradient of the function is also
        returned.
    """

    sector_displacement_2d = sector_displacement_1d.reshape(
        *((2,) + tuple(blocks_shape))
    )

    if input_images.shape[0] == 3:
        three_times = True
        previous_image = input_images[0]
        center_image = input_images[1]
        next_image = input_images[2]

    else:
        previous_image = None
        center_image = input_images[0]
        next_image = input_images[1]
        three_times = False

    if gradient:
        gradient_values = _cost_function(
            sector_displacement_2d,
            center_image,
            next_image,
            mask,
            smooth_gain,
            gradient=True,
        )
        if three_times:
            gradient_values += _cost_function(
                sector_displacement_2d,
                previous_image,
                center_image,
                mask,
                smooth_gain,
                gradient=True,
            )

        return gradient_values.ravel()

    else:
        residuals, smoothness_penalty = _cost_function(
            sector_displacement_2d,
            center_image,
            next_image,
            mask,
            smooth_gain,
            gradient=False,
        )

        if three_times:
            _residuals, _smoothness = _cost_function(
                sector_displacement_2d,
                previous_image,
                center_image,
                mask,
                smooth_gain,
                gradient=False,
            )

            residuals += _residuals
            smoothness_penalty += _smoothness

        if debug:
            print("\nresiduals", residuals)
            print("smoothness_penalty", smoothness_penalty)

        return residuals + smoothness_penalty
.. _`scipy minimization`: \
https://docs.scipy.org/doc/scipy-0.18.1/reference/generated/scipy.optimize.minimize.html

Variational Echo Tracking Cost Function.

This function is designed to be used with the `scipy minimization`_.
The function first argument is the variable to be used in the
minimization procedure.

The sector displacement must be a flat array compatible with the
dimensions of the input image and sectors shape (see parameters section
below for more details).

.. _ndarray:\
https://docs.scipy.org/doc/numpy/reference/generated/numpy.ndarray.html

Parameters
----------
sector_displacement_1d: ndarray_
Array of displacements to apply to each sector. The dimensions are:
sector_displacement_2d
[ x (0) or y (1) displacement, i index of sector, j index of sector ].
The shape of the sector displacements must be compatible with the
input image and the block shape.
The shape should be (2, mx, my) where mx and my are the numbers of
sectors in the x and the y dimension.
input_images: ndarray_
Input images, sequence of 2D arrays, or 3D arrays.
The first dimension represents the images time dimension.

The template_image (first element in first dimensions) denotes the
reference image used to obtain the displacement (2D array).
The second is the target image.

The expected dimensions are (2,nx,ny).
Be aware the the 2D images dimensions correspond to (lon,lat) or (x,y).
blocks_shape: ndarray_ (ndim=2)
Number of sectors in each dimension (x and y).
blocks_shape.shape = (mx,my)
mask: ndarray_ (ndim=2)
Data mask. If is True, the data is marked as not valid and is not
used in the computations.
smooth_gain: float
Smoothness constrain gain
debug: bool, optional
If True, print debugging information.
gradient: bool, optional
If True, the gradient of the morphing function is returned.

Returns
-------
penalty or gradient values.

penalty: float
Value of the cost function
gradient_values: ndarray (float64 ,ndim = 3), optional
If gradient keyword is True, the gradient of the function is also
returned.
pysteps/motion/vet.py