Multivariate Z-prime/factor (Kummel et al, 2010) in python

multi_z_factor.py

import numpy as np
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis

def z_factor(p, n):
    numerator = 3 * (np.std(p) + np.std(n))
    denominator = np.abs(np.mean(p) - np.mean(n))
    return 1 - (numerator / denominator)


def multivariate_z_factor(data, labels):
    assert len(set(labels)) == 2, "only allowed 2 types of label"
    labels = np.array(labels)
    label_1, label_2 = set(labels)
    lda = LinearDiscriminantAnalysis(n_components=1)
    lda_out = lda.fit_transform(data, labels)
    x = lda_out[labels == label_1]
    y = lda_out[labels == label_2]
    return z_factor(x, y)