Search code examples
pythonplotellipseeigenvector

Plotting confidence ellipse with eigen vectors in 2D


I am trying to replicate the plot demonstrated in this guide, where a confidence ellipse is shown on top of a scatter plot, along with the eigen vectors. After scavenging the internet for a Python example without success.


Solution

  • I decided to port the Matlab code provided in the article to Python.

    I am sharing my work as I could not find a similar example. Here is the resulting plot:

    enter image description here

    Here is my Python code:

    import traceback
    import numpy as np
    import matplotlib.pyplot as plt
    from matplotlib.patches import Ellipse
    from numpy import linalg as LA
    import matplotlib.transforms as transforms
    
    def plot_ellipse():
    
        # Create some random data
        s = np.array([2,2])
        x = np.random.randn(334);
        y1 = np.random.normal(s[0] * x, 1)
        y2 = np.random.normal(s[1] * x, 1)
        data = np.array([y1,y2])
    
        # Calculate the eigenvectors and eigenvalues
        covariance = np.cov(data)
        eigenval, eigenvec = LA.eig(covariance)
    
        # Get the largest eigenvalue
        largest_eigenval = max(eigenval)
    
        # Get the index of the largest eigenvector
        largest_eigenvec_ind_c = np.argwhere(eigenval == max(eigenval))[0][0]
        largest_eigenvec = eigenvec[:,largest_eigenvec_ind_c]
    
    
        # Get the smallest eigenvector and eigenvalue
        smallest_eigenval = min(eigenval)
        if largest_eigenvec_ind_c == 0:
            smallest_eigenvec = eigenvec[:,1]
        else:
            smallest_eigenvec = eigenvec[:,0]
    
        # Calculate the angle between the x-axis and the largest eigenvector
        angle = np.arctan2(largest_eigenvec[1], largest_eigenvec[0]);
    
        # This angle is between -pi and pi.
        # Let's shift it such that the angle is between 0 and 2pi
        if (angle < 0):
            angle = angle + 2*np.pi;
    
        # Get the coordinates of the data mean
        avg_0 = np.mean(data[0]);
        avg_1 = np.mean(data[1]);
    
        #% Get the 95% confidence interval error ellipse
        chisquare_val = 2.4477;
        X0=avg_0;
        Y0=avg_1;
    
        ax=plt.gca()
    
        pearson = covariance[0, 1]/np.sqrt(covariance[0, 0] * covariance[1, 1])
    
        ell_radius_x = np.sqrt(1 + pearson)
        ell_radius_y = np.sqrt(1 - pearson)
        ellipse = Ellipse((0, 0), width=ell_radius_x * 2, height=ell_radius_y * 2, facecolor='none', edgecolor=(1, 0, 0.2))
    
    
        scale_x = np.sqrt(covariance[0, 0] * chisquare_val)
        mean_x = X0
        scale_y = np.sqrt(covariance[1, 1] * chisquare_val)
        mean_y = Y0
    
        transf = transforms.Affine2D().rotate_deg(45).scale(scale_x, scale_y).translate(mean_x, mean_y)
        ellipse.set_transform(transf + ax.transData)
        ax.add_patch(ellipse)
    
    
        # Plot the original data
        plt.plot(data[0],data[1], '.');
        plt.xlim([min(data[0])-3, max(data[0])+3]);
        plt.ylim([min(data[1])-3, max(data[1])+3]);
    
        # Plot the eigenvectors
        quiveropts = dict(headaxislength=0, color='red', headlength=0, units='xy',angles='xy',scale=1)
    
        plt.quiver(X0, Y0, largest_eigenvec[0]*np.sqrt(largest_eigenval), largest_eigenvec[1]*np.sqrt(largest_eigenval), **quiveropts)
        plt.quiver(X0, Y0, smallest_eigenvec[0]*np.sqrt(smallest_eigenval), smallest_eigenvec[1]*np.sqrt(smallest_eigenval), **quiveropts)
    
        plt.show()
    
        return
    
    if __name__ == '__main__':
        try:
            plot_ellipse()
        except Exception as e:
            print("Caught exception:" + str(e))
            print(traceback.format_exc())