Polynomial fit with iterative rejection of points.

    This function makes use of function polfit_residuals for display
    purposes.

    Parameters
    ----------
    x : 1d numpy array, float
        X coordinates of the data being fitted.
    y : 1d numpy array, float
        Y coordinates of the data being fitted.
    deg : int
        Degree of the fitting polynomial.
    times_sigma_cook : float or None
        Number of times the standard deviation of Cook's distances
        above the median value to reject points iteratively.
    color : single character or 1d numpy array of characters
        Color for all the symbols (single character) or for each
        individual symbol (array of color names with the same length as
        'x' or 'y'). If 'color' is a single character, the rejected
        points are displayed in red color, whereas when 'color' is an
        array of color names, rejected points are displayed with the
        color provided in this array.
    size : int
        Marker size for all the symbols (single character) or for each
        individual symbol (array of integers with the same length as
        'x' or 'y').
    xlim : tuple (floats)
        Plot limits in the X axis.
    ylim : tuple (floats)
        Plot limits in the Y axis.
    xlabel : string
        Character string for label in X axis.
    ylabel : string
        Character string for label in y axis.
    title : string
        Character string for graph title.
    use_r : bool
        If True, the function computes several fits, using R, to
        polynomials of degree deg, deg+1 and deg+2 (when possible).
    geometry : tuple (4 integers) or None
        x, y, dx, dy values employed to set the window geometry.
    debugplot : int
        Determines whether intermediate computations and/or plots
        are displayed. The valid codes are defined in
        numina.array.display.pause_debugplot.

    Return
    ------
    poly : instance of Polynomial (numpy)
        Result from the polynomial fit using numpy Polynomial. Only
        points not flagged as rejected are employed in the fit.
    yres : 1d numpy array, float
        Residuals from polynomial fit. Note that the residuals are
        computed for all the points, including the rejected ones. In
        this way the dimension of this array is the same as the
        dimensions of the input 'x' and 'y' arrays.
    reject : 1d numpy array, bool
        Boolean array indicating rejected points.

def polfit_residuals_with_cook_rejection(
        x, y, deg, times_sigma_cook,
        color='b', size=75,
        xlim=None, ylim=None,
        xlabel=None, ylabel=None, title=None,
        use_r=None,
        geometry=(0,0,640,480),
        debugplot=0):
    """Polynomial fit with iterative rejection of points.

    This function makes use of function polfit_residuals for display
    purposes.

    Parameters
    ----------
    x : 1d numpy array, float
        X coordinates of the data being fitted.
    y : 1d numpy array, float
        Y coordinates of the data being fitted.
    deg : int
        Degree of the fitting polynomial.
    times_sigma_cook : float or None
        Number of times the standard deviation of Cook's distances
        above the median value to reject points iteratively.
    color : single character or 1d numpy array of characters
        Color for all the symbols (single character) or for each
        individual symbol (array of color names with the same length as
        'x' or 'y'). If 'color' is a single character, the rejected
        points are displayed in red color, whereas when 'color' is an
        array of color names, rejected points are displayed with the
        color provided in this array.
    size : int
        Marker size for all the symbols (single character) or for each
        individual symbol (array of integers with the same length as
        'x' or 'y').
    xlim : tuple (floats)
        Plot limits in the X axis.
    ylim : tuple (floats)
        Plot limits in the Y axis.
    xlabel : string
        Character string for label in X axis.
    ylabel : string
        Character string for label in y axis.
    title : string
        Character string for graph title.
    use_r : bool
        If True, the function computes several fits, using R, to
        polynomials of degree deg, deg+1 and deg+2 (when possible).
    geometry : tuple (4 integers) or None
        x, y, dx, dy values employed to set the window geometry.
    debugplot : int
        Determines whether intermediate computations and/or plots
        are displayed. The valid codes are defined in
        numina.array.display.pause_debugplot.

    Return
    ------
    poly : instance of Polynomial (numpy)
        Result from the polynomial fit using numpy Polynomial. Only
        points not flagged as rejected are employed in the fit.
    yres : 1d numpy array, float
        Residuals from polynomial fit. Note that the residuals are
        computed for all the points, including the rejected ones. In
        this way the dimension of this array is the same as the
        dimensions of the input 'x' and 'y' arrays.
    reject : 1d numpy array, bool
        Boolean array indicating rejected points.

    """

    # protections
    if type(x) is not np.ndarray:
        raise ValueError("x=" + str(x) + " must be a numpy.ndarray")
    elif x.ndim != 1:
        raise ValueError("x.ndim=" + str(x.ndim) + " must be 1")
    if type(y) is not np.ndarray:
        raise ValueError("y=" + str(y) + " must be a numpy.ndarray")
    elif y.ndim != 1:
        raise ValueError("y.ndim=" + str(y.ndim) + " must be 1")
    npoints = x.size
    if npoints != y.size:
        raise ValueError("x.size != y.size")
    if type(deg) not in [np.int, np.int64]:
        raise ValueError("deg=" + str(deg) +
                         " is not a valid integer")
    if deg >= npoints:
        raise ValueError("Polynomial degree=" + str(deg) +
                         " can't be fitted with npoints=" + str(npoints))

    # initialize boolean rejection array
    reject = np.zeros(npoints, dtype=np.bool)

    # if there is no room to remove two points, compute a fit without
    # rejection
    if deg == npoints - 1 or deg == npoints - 2:
        poly, yres = polfit_residuals(x=x, y=y, deg=deg, reject=None,
                                      color=color, size=size,
                                      xlim=xlim, ylim=ylim,
                                      xlabel=xlabel, ylabel=ylabel,
                                      title=title,
                                      use_r=use_r,
                                      geometry=geometry,
                                      debugplot=debugplot)
        return poly, yres, reject

    # main loop to reject points iteratively
    loop_to_reject_points = True
    poly = None
    yres = None
    any_point_removed = False
    while loop_to_reject_points:
        # fit to compute residual variance (neglecting already
        # rejected points)
        poly, yres = polfit_residuals(x=x, y=y, deg=deg, reject=reject,
                                      color=color, size=size,
                                      xlim=xlim, ylim=ylim,
                                      xlabel=xlabel, ylabel=ylabel,
                                      title=title,
                                      use_r=use_r,
                                      geometry=geometry,
                                      debugplot=debugplot)
        npoints_effective = npoints - np.sum(reject)
        residual_variance = np.sum(yres*yres)/float(npoints_effective-deg-1)
        # check that there is room to remove two points with the
        # current polynomial degree
        if deg <= npoints_effective - 2:
            cook_distance = np.zeros(npoints)
            for i in range(npoints):
                if not reject[i]:
                    reject_cook = np.copy(reject)
                    reject_cook[i] = True
                    poly_cook, yres_cook = polfit_residuals(
                        x=x, y=y, deg=deg, reject=reject_cook,
                        color=color, size=size,
                        xlim=xlim, ylim=ylim,
                        xlabel=xlabel, ylabel=ylabel,
                        title="Computing Cook's distance for point " +
                              str(i+1),
                        use_r=False,
                        debugplot=0)
                    yres_cook_fitted = yres_cook[np.logical_not(reject)]
                    cook_distance[i] = \
                        np.sum(yres_cook_fitted*yres_cook_fitted) / \
                        (2*residual_variance)
                else:
                    cook_distance[i] = np.inf
                if abs(debugplot) >= 10:
                    print('i, cook_distance[i]:', i, cook_distance[i])
            # determine median absolute cook distance, excluding points
            # already rejected
            dist_cook_fitted = np.abs(cook_distance[np.logical_not(reject)])
            q50 = np.median(dist_cook_fitted)
            # rms computed from the previous data after removing the
            # point with the largest deviation
            rms = np.std(np.sort(dist_cook_fitted)[:-2])
            if abs(debugplot) >= 10:
                print("--> median.......:", q50)
                print("--> rms -2 points:", rms)
            # reject fitted point exceeding the threshold with the
            # largest Cook distance (note: with this method only one
            # point is removed in each iteration of the loop). If the
            # polynomial degree is larger than 1, only intermediate
            # points can be discarded (i.e., the first and last point
            # are never rejected because the curvature of the
            # extrapolated polynomials leads to false outliers)
            index_to_remove = []
            if deg > 1:
                n1 = 1
                n2 = npoints - 1
            else:
                n1 = 0
                n2 = npoints
            for i in range(n1, n2):
                if not reject[i]:
                    if np.abs(cook_distance[i]-q50) > times_sigma_cook * rms:
                        index_to_remove.append(i)
                        if abs(debugplot) >= 10:
                            print('--> suspicious point #', i + 1)
            if len(index_to_remove) == 0:
                if abs(debugplot) >= 10:
                    if any_point_removed:
                        print('==> no need to remove any additional point')
                    else:
                        print('==> no need to remove any point')
                loop_to_reject_points = False
            else:
                imax = np.argmax(np.abs(cook_distance[index_to_remove]))
                reject[index_to_remove[imax]] = True
                any_point_removed = True
                if abs(debugplot) >= 10:
                    print('==> removing point #', index_to_remove[imax] + 1)
        else:
            loop_to_reject_points = False

    # return result
    return poly, yres, reject