performance of PPS in non-linear situation from blog about ppscore HOT 7 CLOSED

from scipy.stats import pearsonr, spearmanr, kendalltau
import ppscore
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
sns.set()


def plot_data(xy, corrs, ax=None, xlim=None, ylim=None):
    """
    
    :param xy: np.array, shape: n * 2
    :param corrs: [(corr_name, corr_func), ...], corr_func(x, y) -> corr_xy
    :param ax:
    :return: 
    """
    if ax is None:
        _, ax = plt.subplots()
    if xlim is not None:
        ax.set_xlim(xlim)
    if ylim is not None:
        ax.set_ylim(ylim)
    ax.set_frame_on(False)
    ax.axes.get_xaxis().set_visible(False)
    ax.axes.get_yaxis().set_visible(False)
    # 
    ax.plot(*xy.T, ',')
    title = []
    for name, func in corrs:
        corr_xy = func(xy[:, 0], xy[:, 1])
        corr_yx = func(xy[:, 1], xy[:, 0]) 
        title.append(f"{name}: {np.round(corr_xy, 3)}, {np.round(corr_yx, 3)}")
    ax.set_title('\n'.join(title))


def pearson(x, y):
    return pearsonr(x, y)[0]

def spearman(x, y):
    return spearmanr(x, y)[0]

def kendall(x, y):
    return kendalltau(x, y)[0]

def pps(x, y):
    df = pd.DataFrame({"x": x, "y": y})
    return ppscore.score(df, 'x', 'y')['ppscore']

def rotation(xy, t):
    return np.dot(xy, [[np.cos(t), -np.sin(t)],
                       [np.sin(t), np.cos(t)]])


corr_config = [('pearson', pearson), ('spearman', spearman), ('kendall', kendall), ('pps', pps)]

# image1
corrs = [1.0, 0.8, 0.4, 0.0, -0.4, -0.8, -1.0]
n = 800
_, axes = plt.subplots(1, len(corrs), figsize=(2 * len(corrs), 2.5))
for i, corr in enumerate(corrs):
    cov = [[1, corr], [corr, 1]]  # 协方差矩阵
    xy = np.random.multivariate_normal([0, 0], cov, n)  # 多维正态分布
    plot_data(xy, corr_config, axes[i])
plt.tight_layout();


# image2 rotation normal
ts = np.array([0, 1/12, 1/6, 1/4, 1/3, 5/12, 1/2]) * np.pi
n = 1000
xy = np.random.multivariate_normal([0, 0], [[1, 1], [1, 1]], n)
_, axes = plt.subplots(1, len(ts), figsize=(2 * len(ts), 2.5))
for i, t in enumerate(ts):
    xy_rot = rotation(xy, t)
    plot_data(xy_rot, corr_config,
              axes[i], xlim=(-4, 4), ylim=(-4, 4))
plt.tight_layout();


# image3 non-linear situations
_, axes = plt.subplots(1, len(ts), figsize=(2 * len(ts), 2.5))
n = 1000

# fig1
x = np.random.uniform(-1, 1, n)
y = 4 * (x**2 - 0.5)**2 + np.random.uniform(-1, 1, n) / 3
plot_data(np.array([x, y]).T, corr_config, ax=axes[0])

# fig2
y = np.random.uniform(-1, 1, n)
xy = np.array([x, y]).T
xy = rotation(xy, -np.pi / 8)
plot_data(xy, corr_config, ax=axes[1])

# fig3
xy = rotation(xy, -np.pi / 8)
plot_data(xy, corr_config, ax=axes[2])

# fig4
y = 2 * x**2 + np.random.uniform(-1, 1, n)
plot_data(np.array([x, y]).T, corr_config, ax=axes[3])

# fig5
y = (x**2 + np.random.uniform(0, 0.5, n)) * \
        np.array([-1, 1])[np.random.randint(0, 2, size=n)]
plot_data(np.array([x, y]).T, corr_config, ax=axes[4])

# fig6
y = np.cos(x * np.pi) + np.random.uniform(0, 1/8, n)
x = np.sin(x * np.pi) + np.random.uniform(0, 1/8, n)
plot_data(np.array([x, y]).T, corr_config, ax=axes[5])

# fig7
xy1 = np.random.multivariate_normal([3, 3], [[1, 0], [0, 1]], int(n/4))
xy2 = np.random.multivariate_normal([-3, 3], [[1, 0], [0, 1]], int(n/4))
xy3 = np.random.multivariate_normal([-3, -3], [[1, 0], [0, 1]], int(n/4))
xy4 = np.random.multivariate_normal([3, -3], [[1, 0], [0, 1]], int(n/4))
xy = np.concatenate((xy1, xy2, xy3, xy4), axis=0)
plot_data(xy, corr_config, ax=axes[6])

plt.tight_layout();

corrs = [1.0, 0.8, 0.4, 0.0, -0.4, -0.8, -1.0]
n = 800
_, axes = plt.subplots(1, len(corrs), figsize=(2 * len(corrs), 2.5))
for i, corr in enumerate(corrs):
    cov = [[1, corr], [corr, 1]]  # 协方差矩阵
    xy = np.random.multivariate_normal([0, 0], cov, n)  # 多维正态分布
    plot_data(xy, corr_config, axes[i])
plt.tight_layout();