Linear Regression on multidimensional arrays

linear_regression.py

def linear_regression(y, comp_data):
    '''
    Performs linear regression on single pixels as well 
    as multidimensional arrays

    Parameters
    ----------
    y : array_like, shape: (signal_axis) or (nav_shape, signal_axis)
        The data to be fit to
    comp_data : array_like, shape: (number_of_comp, signal_axis)
        The components to fit to the data

    Returns:
    ----------
    fit_coefficients : array_like, shape: (number_of_comp) or (nav_shape, number_of_comp)

    '''
    square = np.matmul(comp_data, comp_data.T)
    square_inv = np.linalg.inv(square)
    comp_data2 = np.matmul(square_inv,comp_data)
    return np.dot(y,comp_data2.T)