Replication code for the paper "Nonlinear Granger Causality using Kernel Ridge Regression" that introduces and highlights the usage of mlcausality.
import numpy as np
from sklearn import metrics
from copy import deepcopy


def brier_loss_func(proba, label):
    proba = proba.copy()
    label = label.copy()
    N = len(proba)
    brier_all = np.zeros((N), dtype=np.float32)
    for i in range(N):
        label_matrix = deepcopy(label[i]).astype(np.float32)
        np.fill_diagonal(label_matrix, np.nan)
        label_matrix = label_matrix.flatten()
        label_matrix = label_matrix[~np.isnan(label_matrix)]
        label_matrix = label_matrix.astype(int)
        proba_matrix = deepcopy(proba[i])
        proba_matrix = np.round(proba_matrix, 16)
        proba_matrix = proba_matrix.astype(np.float32)
        np.fill_diagonal(proba_matrix, np.nan)
        proba_matrix = proba_matrix.flatten()
        proba_matrix = proba_matrix[~np.isnan(proba_matrix)]
        proba_matrix = np.clip(proba_matrix, a_min=0, a_max=1)
        brier_all[i] = metrics.brier_score_loss(label_matrix, proba_matrix, pos_label=0)
    return brier_all

brier_loss.py
WojtekFulmyk/mlcausality-krr-paper-replication
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Input: RSS_R,RSS_U,n,pu,pr
Output: f_GC as values of f-statistic

- "RSS_R" and unrestricted "RSS_U" models.
- "pu" and "pr" are the number of parameters to be estimated for the unrestricted and restricted model, respectively. 
- "n" is the number of time-delayed vectors. 

@Reference: 
Wismüller, A., Dsouza, A.M., Vosoughi, M.A., and Abidin, Anas.
Large-scale nonlinear Granger causality for inferring directed dependence from short multivariate time-series data. Sci Rep 11, 7817 (2021).

"""



def calc_f_stat(RSS_R,RSS_U,n,pu,pr):
    f_GC = ((RSS_R-RSS_U)/(RSS_U))*((n-pu-1)/(pu-pr));
    return f_GC

calc_f_stat.py
WojtekFulmyk/mlcausality-krr-paper-replication
import numpy as np
import numba


# Liangyue Cao's method for finding the minimum embedding dimension
# https://doi.org/10.1016/S0167-2789(97)00118-8
# Python implementation by Wojciech Fulmyk, 2023
# Returns d (a list of embedding dimensions), E1, and E2
@numba.jit(nopython=True)
def cao_min_embedding_dimension(
    x, maxd=10, tau=1, threshold=0.95, max_relative_change=0.1
):
    if maxd < 3:
        raise ValueError("maxd must be greater than or equal to 3")
    v = np.zeros(maxd)
    a = np.zeros(maxd)
    ya = np.zeros(maxd)
    E1 = np.zeros(maxd - 1)
    E2 = np.zeros(maxd - 1)
    ndp = x.flatten().shape[0]
    for d in range(1, maxd + 1):
        a[d - 1] = 0.0
        ya[d - 1] = 0.0
        ng = ndp - d * tau
        for i in range(ng):
            v0 = 1.0e30
            for j in range(ng):
                if j == i:
                    continue
                for k in range(d):
                    v[k] = abs(x[i + k * tau] - x[j + k * tau])
                if d != 1:
                    for k in range(d - 1):
                        v[k + 1] = max(v[k], v[k + 1])
                vij = v[d - 1]
                if vij < v0 and vij != 0.0:
                    n0 = j
                    v0 = vij
            a[d - 1] += max(v0, abs(x[i + d * tau] - x[n0 + d * tau])) / v0
            ya[d - 1] += abs(x[i + d * tau] - x[n0 + d * tau])
        a[d - 1] /= float(ng)
        ya[d - 1] /= float(ng)
        if d >= 2:
            E1[d - 2] = a[d - 1] / a[d - 2]
            E2[d - 2] = ya[d - 1] / ya[d - 2]
    d = np.array(list(range(1, maxd)))
    embedding_dim = np.nan
    for dimension in range(3, maxd + 1):
        relative_error = abs(E1[dimension - 2] - E1[dimension - 3]) / E1[dimension - 3]
        if (
            np.isnan(embedding_dim)
            and not np.isnan(E1[dimension - 2])
            and E1[dimension - 2] >= threshold
            and relative_error < max_relative_change
        ):
            embedding_dim = dimension - 2
            break
    return d, E1, E2, embedding_dim


# Confirm implementation works by reconstructing Fig 1
# from Cao's paper, which plots E1 and E2 for the
# Hénon attractor
# Note that some small discrepancies in Fig 1 of
# Cao's paper and the implementation here are normal
# because the points used here may not be the
# exact same ones used by Cao. Nonetheless,
# the conclusion is the same: the minimum
# embedding dimension is 2 and the figures
# are very similar, which confirms that the
# python implementation is valid.
def cao_plot_henon():
    import matplotlib.pyplot as plt

    def henon_map(xn, yn, a, b):
        return yn + 1 - (a * xn**2), b * xn

    xt = []
    # Initial x and y for henon_map
    x = 0
    y = 0
    for i in range(10000):
        xn, yn = henon_map(x, y, 1.4, 0.3)
        xt.append(x)
        x, y = xn, yn
    xt_10t = np.array(xt)
    xt_1t = xt_10t[-1000:]
    result_1t = cao_min_embedding_dimension(xt_1t, maxd=10, tau=1)
    result_10t = cao_min_embedding_dimension(xt_10t, maxd=10, tau=1)
    plt.plot(result_1t[0], result_1t[1], "D-", color="green", label="E1-1T")
    plt.plot(result_10t[0], result_10t[1], "+:", color="red", label="E1-10T")
    plt.plot(result_1t[0], result_1t[2], "s--", color="blue", label="E2-1T")
    plt.plot(result_10t[0], result_10t[2], "x-.", color="black", label="E2-10T")
    plt.legend(loc="lower left")
    print("result_1t embedding dimension: " + str(result_1t[3]))
    print("result_10t embedding dimension: " + str(result_10t[3]))
    plt.show(block=False)
    plt.pause(0.001)

cao_min_embedding_dimension.py
WojtekFulmyk/mlcausality-krr-paper-replication
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.spatial import distance
from scipy.stats import pearsonr

import itertools

class ccm:
    """
    We're checking causality X -> Y       
    Args
        X: timeseries for variable X that could cause Y
        Y: timeseries for variable Y that could be caused by X
        tau: time lag. default = 1
        E: shadow manifold embedding dimension. default = 2
        L: time period/duration to consider (longer = more data). default = length of X
    """
    def __init__(self, X, Y, tau=1, E=2, L=None):
        '''
        X: timeseries for variable X that could cause Y
        Y: timeseries for variable Y that could be caused by X
        tau: time lag
        E: shadow manifold embedding dimension
        L: time period/duration to consider (longer = more data)
        We're checking for X -> Y
        '''
        self.X = X
        self.Y = Y
        self.tau = tau
        self.E = E
        if L == None:
            self.L = len(X)
        else:
            self.L = L
        self.My = self.shadow_manifold(Y) # shadow manifold for Y (we want to know if info from X is in Y)
        self.t_steps, self.dists = self.get_distances(self.My) # for distances between points in manifold

    def shadow_manifold(self, V):
        """
        Given
            V: some time series vector
            tau: lag step
            E: shadow manifold embedding dimension
            L: max time step to consider - 1 (starts from 0)
        Returns
            {t:[t, t-tau, t-2*tau ... t-(E-1)*tau]} = Shadow attractor manifold, dictionary of vectors
        """
        V = V[:self.L] # make sure we cut at L
        M = {t:[] for t in range((self.E-1) * self.tau, self.L)} # shadow manifold
        for t in range((self.E-1) * self.tau, self.L):
            v_lag = [] # lagged values
            for t2 in range(0, self.E-1 + 1): # get lags, we add 1 to E-1 because we want to include E
                v_lag.append(V[t-t2*self.tau])
            M[t] = v_lag
        return M

    # get pairwise distances between vectors in the time series
    def get_distances(self, M):
        """
        Args
            M: The shadow manifold from the time series
        Returns
            t_steps: timesteps
            dists: n x n matrix showing distances of each vector at t_step (rows) from other vectors (columns)
        """

        # we extract the time indices and vectors from the manifold M
        # we just want to be safe and convert the dictionary to a tuple (time, vector)
        # to preserve the time inds when we separate them
        t_vec = [(k, v) for k,v in M.items()]
        t_steps = np.array([i[0] for i in t_vec])
        vecs = np.array([i[1] for i in t_vec])
        dists = distance.cdist(vecs, vecs)
        return t_steps, dists

    def get_nearest_distances(self, t, t_steps, dists):
        """
        Args:
            t: timestep of vector whose nearest neighbors we want to compute
            t_teps: time steps of all vectors in the manifold M, output of get_distances()
            dists: distance matrix showing distance of each vector (row) from other vectors (columns). output of get_distances()
            E: embedding dimension of shadow manifold M
        Returns:
            nearest_timesteps: array of timesteps of E+1 vectors that are nearest to vector at time t
            nearest_distances: array of distances corresponding to vectors closest to vector at time t
        """
        t_ind = np.where(t_steps == t) # get the index of time t
        dist_t = dists[t_ind].squeeze() # distances from vector at time t (this is one row)

        # get top closest vectors
        nearest_inds = np.argsort(dist_t)[1:self.E+1 + 1] # get indices sorted, we exclude 0 which is distance from itself
        nearest_timesteps = t_steps[nearest_inds] # index column-wise, t_steps are same column and row-wise
        nearest_distances = dist_t[nearest_inds]

        return nearest_timesteps, nearest_distances

    def predict(self, t):
        """
        Args
            t: timestep at manifold of y, My, to predict X at same time step
        Returns
            X_true: the true value of X at time t
            X_hat: the predicted value of X at time t using the manifold My
        """
        eps = 0.000001 # epsilon minimum distance possible
        t_ind = np.where(self.t_steps == t) # get the index of time t
        dist_t = self.dists[t_ind].squeeze() # distances from vector at time t (this is one row)
        nearest_timesteps, nearest_distances = self.get_nearest_distances(t, self.t_steps, self.dists)

        # get weights
        #u = np.exp(-nearest_distances/np.max([eps, nearest_distances[0]])) # we divide by the closest distance to scale
        # Fix FloatingPointError: underflow encountered in exp
        u = np.exp(np.clip( -nearest_distances/np.max([eps, nearest_distances[0]]), -500, 500 )) # we divide by the closest distance to scale
        w = u / np.sum(u)

        # get prediction of X
        X_true = self.X[t] # get corresponding true X
        X_cor = np.array(self.X)[nearest_timesteps] # get corresponding Y to cluster in Mx
        X_hat = (w * X_cor).sum() # get X_hat

    #     DEBUGGING
    #     will need to check why nearest_distances become nan
    #     if np.isnan(X_hat):
    #         print(nearest_timesteps)
    #         print(nearest_distances)

        return X_true, X_hat


    def causality(self):
        '''
        Args:
            None
        Returns:
            (r, p): how much X causes Y. as a correlation between predicted X and true X and the p-value (significance)
        '''

        # run over all timesteps in M
        # X causes Y, we can predict X using My
        # X puts some info into Y that we can use to reverse engineer X from Y via My
        X_true_list = []
        X_hat_list = []

        for t in list(self.My.keys()): # for each time step in My
            X_true, X_hat = self.predict(t) # predict X from My
            X_true_list.append(X_true)
            X_hat_list.append(X_hat)

        x, y = X_true_list, X_hat_list
        r, p = pearsonr(x, y)

        return r, p

    def visualize_cross_mapping(self):
        """
        Visualize the shadow manifolds and some cross mappings
        """
        # we want to check cross mapping from Mx to My and My to Mx

        f, axs = plt.subplots(1, 2, figsize=(12, 6))

        for i, ax in zip((0, 1), axs): # i will be used in switching Mx and My in Cross Mapping visualization
            #===============================================
            # Shadow Manifolds Visualization

            X_lag, Y_lag = [], []
            for t in range(1, len(self.X)):
                X_lag.append(self.X[t-self.tau])
                Y_lag.append(self.Y[t-self.tau])
            X_t, Y_t = self.X[1:], self.Y[1:] # remove first value

            ax.scatter(X_t, X_lag, s=5, label='$M_x$')
            ax.scatter(Y_t, Y_lag, s=5, label='$M_y$', c='y')

            #===============================================
            # Cross Mapping Visualization

            A, B = [(self.Y, self.X), (self.X, self.Y)][i]
            cm_direction = ['Mx to My', 'My to Mx'][i]

            Ma = self.shadow_manifold(A)
            Mb = self.shadow_manifold(B)

            t_steps_A, dists_A = self.get_distances(Ma) # for distances between points in manifold
            t_steps_B, dists_B = self.get_distances(Mb) # for distances between points in manifold

            # Plot cross mapping for different time steps
            timesteps = list(Ma.keys())
            for t in np.random.choice(timesteps, size=3, replace=False):
                Ma_t = Ma[t]
                near_t_A, near_d_A = self.get_nearest_distances(t, t_steps_A, dists_A)

                for i in range(self.E+1):
                    # points on Ma
                    A_t = Ma[near_t_A[i]][0]
                    A_lag = Ma[near_t_A[i]][1]
                    ax.scatter(A_t, A_lag, c='b', marker='s')

                    # corresponding points on Mb
                    B_t = Mb[near_t_A[i]][0]
                    B_lag = Mb[near_t_A[i]][1]
                    ax.scatter(B_t, B_lag, c='r', marker='*', s=50)

                    # connections
                    ax.plot([A_t, B_t], [A_lag, B_lag], c='r', linestyle=':')

            ax.set_title(f'{cm_direction} cross mapping. time lag, tau = {self.tau}, E = 2')
            ax.legend(prop={'size': 14})

            ax.set_xlabel('$X_t$, $Y_t$', size=15)
            ax.set_ylabel('$X_{t-1}$, $Y_{t-1}$', size=15)
        plt.show()

    def plot_ccm_correls(self):
        """
        Args
            X: X time series
            Y: Y time series
            tau: time lag
            E: shadow manifold embedding dimension
            L: time duration
        Returns
            None. Just correlation plots between predicted X|M_y and true X
        """
        X_My_true, X_My_pred = [], []
        for t in range((self.E-1) * self.tau, self.L):
            true, pred = self.predict(t)
            X_My_true.append(true)
            X_My_pred.append(pred)

        # predicting X from My
        r, p = np.round(pearsonr(X_My_true, X_My_pred), 4)

        plt.scatter(X_My_true, X_My_pred, s=10)
        plt.xlabel('$X(t)$ (observed)', size=15)
        plt.ylabel('$\hat{X}(t)|M_y$ (estimated)', size=15)
        plt.title(f'tau={self.tau}, E={self.E}, L={self.L}, Correlation coeff = {r}')

        plt.show()


def loop_causal_cmm(data, tau=1, E=2, L=None):
    results_list = []
    permute_list = list(itertools.permutations(range(data.shape[1]),2))
    for X_idx, y_idx in permute_list:
        ccm1 = ccm(X=data[:,X_idx], Y=data[:,y_idx], tau=tau, E=E, L=L)
        score, pval = ccm1.causality()
        results_list.append([X_idx,y_idx,score,pval])
    out_df = pd.DataFrame(results_list, columns=['X','y','score','pvalue'])
    return out_df

causal_ccm.py
WojtekFulmyk/mlcausality-krr-paper-replication
# Get plot type to make
import sys

plot_type = sys.argv[1]
if plot_type.lower() == "brier":
    plot_type_cap = "Brier Score"
elif plot_type.lower() == "auc":
    plot_type_cap = "AUC"
elif plot_type.lower() == "accuracy":
    plot_type_cap = "Accuracy (thresh = 0.05)"
elif plot_type.lower() == "balanced_accuracy":
    plot_type_cap = "Bal.Accuracy (thresh = 0.05)"
elif plot_type.lower() == "f1wgt":
    plot_type_cap = "Weighted F1 (thresh = 0.05)"
elif plot_type.lower() == "sensitivity":
    plot_type_cap = "Sensitivity (thresh = 0.05)"
elif plot_type.lower() == "specificity":
    plot_type_cap = "Specificity (thresh = 0.05)"
elif plot_type.lower() == "gmean":
    plot_type_cap = "G-mean (thresh = 0.05)"
elif plot_type.lower() == "accuracy2":
    plot_type_cap = "Accuracy (max G-mean thresh)"
elif plot_type.lower() == "balanced_accuracy2":
    plot_type_cap = "Bal. Accuracy (max G-mean thresh)"
elif plot_type.lower() == "f1wgt2":
    plot_type_cap = "Weighted F1 (max G-mean thresh)"
elif plot_type.lower() == "sensitivity2":
    plot_type_cap = "Sensitivity (max G-mean thresh)"
elif plot_type.lower() == "specificity2":
    plot_type_cap = "Specificity (max G-mean thresh)"
elif plot_type.lower() == "gmean2":
    plot_type_cap = "G-mean (max G-mean thresh)"

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import seaborn as sns
import pickle

from copy import deepcopy
from decimal import Decimal, localcontext, ROUND_DOWN


# This truncate function is based on:
# https://stackoverflow.com/a/28323804
def truncate(number, places):
    if not isinstance(places, int):
        raise ValueError("Decimal places must be an integer.")
    if places < 1:
        raise ValueError("Decimal places must be at least 1.")
    # If you want to truncate to 0 decimal places, just do int(number).
    with localcontext() as context:
        context.rounding = ROUND_DOWN
        exponent = Decimal(str(10**-places))
        return Decimal(str(number)).quantize(exponent)


with open(r"pickled_data/5_linear_" + plot_type + "_500.pickle", "rb") as input_file:
    df_init1 = pickle.load(input_file)

df_init1["Length of time-series"] = 500


with open(r"pickled_data/5_linear_" + plot_type + "_1000.pickle", "rb") as input_file:
    df_init2 = pickle.load(input_file)

df_init2["Length of time-series"] = 1000

with open(r"pickled_data/5_linear_" + plot_type + "_1500.pickle", "rb") as input_file:
    df_init3 = pickle.load(input_file)

df_init3["Length of time-series"] = 1500

with open(r"pickled_data/5_linear_" + plot_type + "_2000.pickle", "rb") as input_file:
    df_init4 = pickle.load(input_file)

df_init4["Length of time-series"] = 2000
df1 = pd.concat([df_init1, df_init2, df_init3, df_init4]).reset_index(drop=True)
df1 = df1.melt("Length of time-series", var_name="Model", value_name="Val")

with open(r"pickled_data/5_nonlinear_" + plot_type + "_500.pickle", "rb") as input_file:
    df_init1 = pickle.load(input_file)

df_init1["Length of time-series"] = 500


with open(
    r"pickled_data/5_nonlinear_" + plot_type + "_1000.pickle", "rb"
) as input_file:
    df_init2 = pickle.load(input_file)

df_init2["Length of time-series"] = 1000

with open(
    r"pickled_data/5_nonlinear_" + plot_type + "_1500.pickle", "rb"
) as input_file:
    df_init3 = pickle.load(input_file)

df_init3["Length of time-series"] = 1500

with open(
    r"pickled_data/5_nonlinear_" + plot_type + "_2000.pickle", "rb"
) as input_file:
    df_init4 = pickle.load(input_file)

df_init4["Length of time-series"] = 2000
df2 = pd.concat([df_init1, df_init2, df_init3, df_init4]).reset_index(drop=True)
df2 = df2.melt("Length of time-series", var_name="Model", value_name="Val")

with open(r"pickled_data/7_nonlinear_" + plot_type + "_500.pickle", "rb") as input_file:
    df_init1 = pickle.load(input_file)

df_init1["Length of time-series"] = 500


with open(
    r"pickled_data/7_nonlinear_" + plot_type + "_1000.pickle", "rb"
) as input_file:
    df_init2 = pickle.load(input_file)

df_init2["Length of time-series"] = 1000

with open(
    r"pickled_data/7_nonlinear_" + plot_type + "_1500.pickle", "rb"
) as input_file:
    df_init3 = pickle.load(input_file)

df_init3["Length of time-series"] = 1500

with open(
    r"pickled_data/7_nonlinear_" + plot_type + "_2000.pickle", "rb"
) as input_file:
    df_init4 = pickle.load(input_file)

df_init4["Length of time-series"] = 2000
df3 = pd.concat([df_init1, df_init2, df_init3, df_init4]).reset_index(drop=True)
df3 = df3.melt("Length of time-series", var_name="Model", value_name="Val")

with open(r"pickled_data/9_nonlinear_" + plot_type + "_500.pickle", "rb") as input_file:
    df_init1 = pickle.load(input_file)

df_init1["Length of time-series"] = 500


with open(
    r"pickled_data/9_nonlinear_" + plot_type + "_1000.pickle", "rb"
) as input_file:
    df_init2 = pickle.load(input_file)

df_init2["Length of time-series"] = 1000

with open(
    r"pickled_data/9_nonlinear_" + plot_type + "_1500.pickle", "rb"
) as input_file:
    df_init3 = pickle.load(input_file)

df_init3["Length of time-series"] = 1500

with open(
    r"pickled_data/9_nonlinear_" + plot_type + "_2000.pickle", "rb"
) as input_file:
    df_init4 = pickle.load(input_file)

df_init4["Length of time-series"] = 2000
df4 = pd.concat([df_init1, df_init2, df_init3, df_init4]).reset_index(drop=True)
df4 = df4.melt("Length of time-series", var_name="Model", value_name="Val")

with open(
    r"pickled_data/11_nonlinear_" + plot_type + "_500.pickle", "rb"
) as input_file:
    df_init1 = pickle.load(input_file)

df_init1["Length of time-series"] = 500


with open(
    r"pickled_data/11_nonlinear_" + plot_type + "_1000.pickle", "rb"
) as input_file:
    df_init2 = pickle.load(input_file)

df_init2["Length of time-series"] = 1000

with open(
    r"pickled_data/11_nonlinear_" + plot_type + "_1500.pickle", "rb"
) as input_file:
    df_init3 = pickle.load(input_file)

df_init3["Length of time-series"] = 1500

with open(
    r"pickled_data/11_nonlinear_" + plot_type + "_2000.pickle", "rb"
) as input_file:
    df_init4 = pickle.load(input_file)

df_init4["Length of time-series"] = 2000
df5 = pd.concat([df_init1, df_init2, df_init3, df_init4]).reset_index(drop=True)
df5 = df5.melt("Length of time-series", var_name="Model", value_name="Val")

with open(r"pickled_data/34_zachary1_" + plot_type + "_500.pickle", "rb") as input_file:
    df_init1 = pickle.load(input_file)

df_init1["Length of time-series"] = 500


with open(
    r"pickled_data/34_zachary1_" + plot_type + "_1000.pickle", "rb"
) as input_file:
    df_init2 = pickle.load(input_file)

df_init2["Length of time-series"] = 1000

with open(
    r"pickled_data/34_zachary1_" + plot_type + "_1500.pickle", "rb"
) as input_file:
    df_init3 = pickle.load(input_file)

df_init3["Length of time-series"] = 1500

with open(
    r"pickled_data/34_zachary1_" + plot_type + "_2000.pickle", "rb"
) as input_file:
    df_init4 = pickle.load(input_file)

df_init4["Length of time-series"] = 2000
df6 = pd.concat([df_init1, df_init2, df_init3, df_init4]).reset_index(drop=True)
df6 = df6.melt("Length of time-series", var_name="Model", value_name="Val")

with open(r"pickled_data/34_zachary2_" + plot_type + "_500.pickle", "rb") as input_file:
    df_init1 = pickle.load(input_file)

df_init1["Length of time-series"] = 500


with open(
    r"pickled_data/34_zachary2_" + plot_type + "_1000.pickle", "rb"
) as input_file:
    df_init2 = pickle.load(input_file)

df_init2["Length of time-series"] = 1000

with open(
    r"pickled_data/34_zachary2_" + plot_type + "_1500.pickle", "rb"
) as input_file:
    df_init3 = pickle.load(input_file)

df_init3["Length of time-series"] = 1500

with open(
    r"pickled_data/34_zachary2_" + plot_type + "_2000.pickle", "rb"
) as input_file:
    df_init4 = pickle.load(input_file)

df_init4["Length of time-series"] = 2000
df7 = pd.concat([df_init1, df_init2, df_init3, df_init4]).reset_index(drop=True)
df7 = df7.melt("Length of time-series", var_name="Model", value_name="Val")

title_font = {
    "size": "32",
    "color": "black",
    "weight": "bold",
    "verticalalignment": "bottom",
}
title_font2 = {
    "size": "52",
    "color": "black",
    "weight": "bold",
    "verticalalignment": "bottom",
}
f, axes = plt.subplots(1, 7, sharex=False, sharey=True, figsize=(30, 14))
f.text(0.5, 1.025, plot_type_cap, ha="center", va="center", **title_font2)
f.text(-0.01, 0.5, plot_type_cap, ha="center", va="center", rotation=90, **title_font)

sns.boxplot(
    x="Model",
    y="Val",
    hue="Length of time-series",
    data=df1,
    palette="Set3",
    ax=axes[0],
    linewidth=2,
    saturation=1,
    width=0.7,
    color="black",
    notch=False,
    medianprops=dict(color="black", linewidth=10, alpha=1, solid_capstyle="butt"),
    whis=(0, 100),
    showfliers=False,
    bootstrap=10000,
    whiskerprops=dict(color="black", alpha=1),
    boxprops=dict(edgecolor="black", alpha=1),
    capprops=dict(color="black", alpha=1),
)
sns.boxplot(
    x="Model",
    y="Val",
    hue="Length of time-series",
    data=df2,
    palette="Set3",
    ax=axes[1],
    linewidth=2,
    saturation=1,
    width=0.7,
    color="black",
    notch=False,
    medianprops=dict(color="black", linewidth=10, alpha=1, solid_capstyle="butt"),
    whis=(0, 100),
    showfliers=False,
    bootstrap=10000,
    whiskerprops=dict(color="black", alpha=1),
    boxprops=dict(edgecolor="black", alpha=1),
    capprops=dict(color="black", alpha=1),
)
sns.boxplot(
    x="Model",
    y="Val",
    hue="Length of time-series",
    data=df3,
    palette="Set3",
    ax=axes[2],
    linewidth=2,
    saturation=1,
    width=0.7,
    color="black",
    notch=False,
    medianprops=dict(color="black", linewidth=10, alpha=1, solid_capstyle="butt"),
    whis=(0, 100),
    showfliers=False,
    bootstrap=10000,
    whiskerprops=dict(color="black", alpha=1),
    boxprops=dict(edgecolor="black", alpha=1),
    capprops=dict(color="black", alpha=1),
)
sns.boxplot(
    x="Model",
    y="Val",
    hue="Length of time-series",
    data=df4,
    palette="Set3",
    ax=axes[3],
    linewidth=2,
    saturation=1,
    width=0.7,
    color="black",
    notch=False,
    medianprops=dict(color="black", linewidth=10, alpha=1, solid_capstyle="butt"),
    whis=(0, 100),
    showfliers=False,
    bootstrap=10000,
    whiskerprops=dict(color="black", alpha=1),
    boxprops=dict(edgecolor="black", alpha=1),
    capprops=dict(color="black", alpha=1),
)
sns.boxplot(
    x="Model",
    y="Val",
    hue="Length of time-series",
    data=df5,
    palette="Set3",
    ax=axes[4],
    linewidth=2,
    saturation=1,
    width=0.7,
    color="black",
    notch=False,
    medianprops=dict(color="black", linewidth=10, alpha=1, solid_capstyle="butt"),
    whis=(0, 100),
    showfliers=False,
    bootstrap=10000,
    whiskerprops=dict(color="black", alpha=1),
    boxprops=dict(edgecolor="black", alpha=1),
    capprops=dict(color="black", alpha=1),
)
sns.boxplot(
    x="Model",
    y="Val",
    hue="Length of time-series",
    data=df6,
    palette="Set3",
    ax=axes[5],
    linewidth=2,
    saturation=1,
    width=0.7,
    color="black",
    notch=False,
    medianprops=dict(color="black", linewidth=10, alpha=1, solid_capstyle="butt"),
    whis=(0, 100),
    showfliers=False,
    bootstrap=10000,
    whiskerprops=dict(color="black", alpha=1),
    boxprops=dict(edgecolor="black", alpha=1),
    capprops=dict(color="black", alpha=1),
)
sns.boxplot(
    x="Model",
    y="Val",
    hue="Length of time-series",
    data=df7,
    palette="Set3",
    ax=axes[6],
    linewidth=2,
    saturation=1,
    width=0.7,
    color="black",
    notch=False,
    medianprops=dict(color="black", linewidth=10, alpha=1, solid_capstyle="butt"),
    whis=(0, 100),
    showfliers=False,
    bootstrap=10000,
    whiskerprops=dict(color="black", alpha=1),
    boxprops=dict(edgecolor="black", alpha=1),
    capprops=dict(color="black", alpha=1),
)

axes[0].set_title("5-linear", fontweight="bold", fontsize=48)
axes[1].set_title("5-nonlin", fontweight="bold", fontsize=48)
axes[2].set_title("7-nonlin", fontweight="bold", fontsize=48)
axes[3].set_title("9-nonlin", fontweight="bold", fontsize=48)
axes[4].set_title("11-nonlin", fontweight="bold", fontsize=48)
axes[5].set_title("Zachary1", fontweight="bold", fontsize=48)
axes[6].set_title("Zachary2", fontweight="bold", fontsize=48)

for i in range(7):
    if i < 6:
        axes[i].get_legend().remove()
    else:
        axes[i].legend(
            prop=dict(size=26, weight="bold"),
            title="$\\bf{Num. obs.}$",
            title_fontsize=26,
        )
    x_axis = axes[i].axes.get_xaxis()
    x_label = x_axis.get_label()
    x_label.set_visible(False)
    y_axis = axes[i].axes.get_yaxis()
    y_label = y_axis.get_label()
    y_label.set_visible(False)
    if i == 0:
        start, end = axes[i].get_ylim()
        new_start = float(truncate(deepcopy(start), 1))
        if plot_type.lower() == "accuracy" or plot_type.lower() == "accuracy2":
            new_start = abs(round(max(new_start - 0.1, 0), 1))
        else:
            new_start = abs(round(max(new_start, 0), 1))
    axes[i].yaxis.set_ticks(np.arange(new_start, end, 0.1))
    axes[i].yaxis.set_major_formatter(ticker.FormatStrFormatter("%0.1f"))
    if i == 0:
        axes[i].tick_params(width=5, size=10)
    else:
        axes[i].tick_params(
            top=False, bottom=True, left=False, right=False, width=5, size=10
        )
    axes[i].tick_params(axis="x", rotation=90, labelsize=48)
    axes[i].tick_params(axis="y", labelsize=48)
    [
        axes[i].axvline(x, color="0.3", linestyle="--", linewidth=3)
        for x in [0.5, 1.5, 2.5]
    ]
    axes[i].grid(visible=True, which="major", axis="y", linewidth=3)
    axes[i].set(axisbelow=True)
    for axis in ["top", "bottom", "left", "right"]:
        axes[i].spines[axis].set_linewidth(5)  # change width
    if i == 0:
        start2, end2 = axes[i].get_ylim()
    for axis in ["left", "right"]:
        axes[i].spines[axis].set_bounds(
            low=start2 - (end2 - start2) * 0.21153846153846154,
            high=end2 + (end2 - start2) * 0.06937799043062202,
        )  # Extend spine bounds

f.tight_layout()
f.subplots_adjust(wspace=0, hspace=0)
plt.rcParams["svg.fonttype"] = "none"
plt.savefig(
    "plots/" + plot_type.replace("_", "") + "plotmulti.pdf", bbox_inches="tight"
)
plt.savefig(
    "plots/" + plot_type.replace("_", "") + "plotmulti.eps", bbox_inches="tight"
)
# plt.show()

make_all_multi_plots.py
WojtekFulmyk/mlcausality-krr-paper-replication
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
The function separates input time-series data into pieces of data, based on Taken's theorem to represent system dynamics within the state-space.

@Reference: 
Wismüller, A., Dsouza, A.M., Vosoughi, M.A. et al. 
Large-scale nonlinear Granger causality for inferring directed dependence from short multivariate time-series data. Sci Rep 11, 7817 (2021).
"""
import torch
import numpy as np


def multivariate_split(X,ar_order, valid_percent=0):
    X=X.copy()
    TS=np.shape(X)[1]
    n_vars=np.shape(X)[0]
    val_num=int(valid_percent*TS)
    my_data_train=torch.zeros((TS-ar_order-val_num,ar_order,n_vars))
    my_data_y_train=torch.zeros((TS-ar_order-val_num,1,n_vars))
    my_data_val=torch.zeros((val_num,ar_order,n_vars))
    my_data_y_val=torch.zeros((val_num,1,n_vars))
    for i in range(TS-ar_order-val_num):
        my_data_train[i]=torch.from_numpy(X.transpose()[i:i+ar_order,:])
        my_data_y_train[i]=torch.from_numpy(X.transpose()[i+ar_order,:])

    for i in range(TS-ar_order-val_num, TS-ar_order,1):
        my_data_val[i-(TS-ar_order-val_num)]=torch.from_numpy(X.transpose()[i:i+ar_order,:])
        my_data_y_val[i-(TS-ar_order-val_num)]=torch.from_numpy(X.transpose()[i+ar_order,:])
    return my_data_train, my_data_y_train, my_data_val, my_data_y_val

multivariate_split.py
WojtekFulmyk/mlcausality-krr-paper-replication
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""


@Reference: 
Wismüller, A., Dsouza, A.M., Vosoughi, M.A. et al. 
Large-scale nonlinear Granger causality for inferring directed dependence from short multivariate time-series data. Sci Rep 11, 7817 (2021).
"""
import numpy as np

def normalize_0_mean_1_std(inp_series):
    inp_series=inp_series.copy()
    mean_ts=np.array([inp_series.mean(axis=1)]).transpose()
    mean_ts_mtrx = mean_ts*np.ones((1,inp_series.shape[1]));
    unb_data_mtrx = inp_series - mean_ts_mtrx
    p = np.power(unb_data_mtrx,2)
    s=np.array([p.sum(axis=1)]).transpose()
    sc=np.sqrt(s/p.shape[1])
    sc2=sc*(np.ones((1,p.shape[1])))
    nrm= np.divide(unb_data_mtrx,sc2)
    return nrm

normalize_0_mean_1_std.py
WojtekFulmyk/mlcausality-krr-paper-replication
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
The code returns performance measures (in terms of AUROC) given a list of 
"Adjs": data matrices and 
"label": label matrices. 

@Reference: 

Originally created by:
Wismüller, A., Dsouza, A.M., Vosoughi, M.A. et al. 
Large-scale nonlinear Granger causality for inferring directed dependence from short multivariate time-series data. Sci Rep 11, 7817 (2021).

Modified by Wojciech (Victor) Fulmyk to account for the diagonal which is always labelled correctly.
"""
import numpy as np
from sklearn import metrics
from copy import deepcopy


def recovery_performance(Adjs,label):
    Adjs=Adjs.copy()
    label=label.copy()
    N=len(Adjs)
    auc_all = np.zeros((N), dtype=np.float32)   
    for i in range(N):
        label_matrix = deepcopy(label[i]).astype(float)
        np.fill_diagonal(label_matrix, np.nan)
        label_matrix = label_matrix.flatten()
        label_matrix = label_matrix[~np.isnan(label_matrix)]
        Adj_matrix = deepcopy(Adjs[i]).astype(float)
        np.fill_diagonal(Adj_matrix, np.nan)
        Adj_matrix = Adj_matrix.flatten()
        Adj_matrix = Adj_matrix[~np.isnan(Adj_matrix)]
        auc_all[i] = metrics.roc_auc_score(label_matrix, Adj_matrix)
    return auc_all


recovery_performance_mod.py
WojtekFulmyk/mlcausality-krr-paper-replication
# Get time samples and seed from arguments passed when this script is called
import sys

network_to_check = sys.argv[1]
n_time_samples_str = sys.argv[2]
n_time_samples_int = int(n_time_samples_str)
seed = int(sys.argv[3])

import os

# Set environment variables useful for multiprocessing.
# This will only work if your BLAS is MKL or openblas.
# If using other BLAS implementations, these have no effect
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"

os.makedirs("plots", exist_ok=True)
os.makedirs("pickled_data", exist_ok=True)

import warnings

warnings.simplefilter(action="ignore", category=FutureWarning)
warnings.simplefilter(action="ignore", category=UserWarning)

import numpy as np

np.seterr(all="raise")

import pandas as pd

import cao_min_embedding_dimension
import mlcausality

from utils_mod import lsNGC as granger
from recovery_performance_mod import recovery_performance

from brier_loss import brier_loss_func

import causal_ccm

from tigramite import data_processing as pcmci_data_processing
from tigramite.pcmci import PCMCI
from tigramite.independence_tests.parcorr import ParCorr

import tqdm

import matplotlib.pyplot as plt

from copy import deepcopy
import itertools
from pprint import pprint
import pickle

import networkx as nx

from sklearn.metrics import (
    confusion_matrix,
    accuracy_score,
    balanced_accuracy_score,
    f1_score,
)

# Get number of cpu cores
import psutil

num_cores = psutil.cpu_count(logical=True)
# Import multriprocessing plugins from tqdm
# These rely on concurrent.futures
from tqdm.contrib.concurrent import ensure_lock, process_map

# Set seed.
np.random.seed(seed)

if network_to_check == "5_linear":
    # from spectral_connectivity.simulate import simulate_MVAR
    def simulate_MVAR(
        coefficients,
        noise_covariance=None,
        n_time_samples=100,
        n_trials=1,
        n_burnin_samples=100,
    ):
        """
        Simulate multivariate autoregressive (MVAR) process.
        Parameters
        ----------
        coefficients : array, shape (n_time_samples, n_lags, n_signals, n_signals)
        noise_covariance : array, shape (n_signals, n_signals)
        Returns
        -------
        time_series : array, shape (n_time_samples - n_burnin_samples,
                                    n_trials, n_signals)
        """
        n_lags, n_signals, _ = coefficients.shape
        if noise_covariance is None:
            noise_covariance = np.eye(n_signals)
        time_series = np.random.multivariate_normal(
            np.zeros((n_signals,)),
            noise_covariance,
            size=(n_time_samples + n_burnin_samples, n_trials),
        )

        for time_ind in np.arange(n_lags, n_time_samples + n_burnin_samples):
            for lag_ind in np.arange(n_lags):
                time_series[time_ind, ...] += np.matmul(
                    coefficients[np.newaxis, np.newaxis, lag_ind, ...],
                    time_series[time_ind - (lag_ind + 1), ..., np.newaxis],
                ).squeeze()
        return time_series[n_burnin_samples:, ...]

    # Generate data using the spectral_connectivity package

    def baccala_example3():
        """Baccalá, L.A., and Sameshima, K. (2001). Partial directed coherence:
        a new concept in neural structure determination. Biological
        Cybernetics 84, 463–474.
        """
        np.random.seed(seed)
        n_time_samples, n_lags, n_signals = n_time_samples_int, 3, 5
        coefficients = np.zeros((n_lags, n_signals, n_signals))

        coefficients[0, 0, 0] = 0.95 * np.sqrt(2)
        coefficients[1, 0, 0] = -0.9025

        coefficients[1, 1, 0] = 0.50
        coefficients[2, 2, 0] = -0.40

        coefficients[1, 3, 0] = -0.5
        coefficients[0, 3, 3] = 0.5 * np.sqrt(2)
        coefficients[0, 3, 4] = 0.25 * np.sqrt(2)

        coefficients[0, 4, 3] = -0.5 * np.sqrt(2)
        coefficients[0, 4, 4] = 0.5 * np.sqrt(2)

        noise_covariance = None

        return (
            simulate_MVAR(
                coefficients,
                noise_covariance=noise_covariance,
                n_time_samples=n_time_samples,
                n_trials=50,
                n_burnin_samples=500,
            ),
            coefficients,
        )

    data, coefficients = baccala_example3()

    data_list = [data[:, i, :] for i in range(data.shape[1])]

    # Find the adjacency matrix (which are essentially the true labels)
    adjacency_matrix = np.zeros((coefficients.shape[1], coefficients.shape[2]))
    for i in range(coefficients.shape[0]):
        adjacency_matrix += np.ceil(np.abs(coefficients[i]))

    adjacency_matrix = np.clip(adjacency_matrix, 0, 1)
    np.fill_diagonal(adjacency_matrix, 0)
    adjacency_matrix = adjacency_matrix.T.astype(int)
    true_labels = adjacency_matrix
elif network_to_check == "5_nonlinear":
    # 5-node nonlinear network
    true_labels = np.array(
        [
            [0, 1, 1, 1, 0],
            [0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0],
            [0, 0, 0, 0, 1],
            [0, 0, 0, 1, 0],
        ]
    )
    n_burnin_samples = 500
    n_trials = 50
    n_lags = 3
    data_list = []
    while True:
        np.random.seed(seed)
        seed += 1
        data = np.random.normal(0, 1, [n_lags, 5])
        try:
            for j in range(n_lags, n_time_samples_int + n_burnin_samples):
                x1 = (
                    0.95 * np.sqrt(2) * data[-1, 0]
                    - 0.9025 * data[-2, 0]
                    + np.random.normal(0, 1)
                )
                x2 = 0.5 * data[-2, 0] ** 2 + np.random.normal(0, 1)
                x3 = -0.4 * data[-3, 0] + np.random.normal(0, 1)
                x4 = (
                    -0.5 * data[-2, 0] ** 2
                    + 0.5 * np.sqrt(2) * data[-1, 3]
                    + 0.25 * np.sqrt(2) * data[-1, 4]
                    + np.random.normal(0, 1)
                )
                x5 = (
                    -0.5 * np.sqrt(2) * data[-1, 3]
                    + 0.5 * np.sqrt(2) * data[-1, 4]
                    + np.random.normal(0, 1)
                )
                data = np.vstack([data, np.array([x1, x2, x3, x4, x5])])
        except Exception:
            continue
        data = np.array(data[n_burnin_samples:])
        data_list.append(data)
        if len(data_list) == n_trials:
            break
elif network_to_check == "7_nonlinear":
    # 7-node nonlinear network
    true_labels = np.array(
        [
            [0, 1, 0, 0, 0, 1, 1],
            [0, 0, 1, 0, 0, 0, 0],
            [0, 0, 0, 1, 0, 1, 0],
            [0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 1, 0, 1],
            [0, 0, 0, 1, 0, 0, 0],
        ]
    )
    n_burnin_samples = 500
    n_trials = 50
    n_lags = 3
    data_list = []
    while True:
        np.random.seed(seed)
        seed += 1
        data = np.random.normal(0, 1, [n_lags, 7])
        try:
            for j in range(n_lags, n_time_samples_int + n_burnin_samples):
                x1 = (
                    0.95 * np.sqrt(2) * data[-1, 0]
                    - 0.9025 * data[-2, 0]
                    + np.random.normal(0, 1)
                )
                x2 = (
                    -0.04 * data[-3, 0] ** 3
                    + 0.04 * data[-1, 0] ** 3
                    + np.random.normal(0, 1)
                )
                x3 = (
                    -0.04 * np.sqrt(2) * data[-1, 1] ** 3
                    + 0.04 * np.sqrt(2) * data[-2, 1] ** 3
                    + np.random.normal(0, 1)
                )
                x4 = (
                    np.log1p(np.abs(data[-1, 2])) * np.sign(data[-1, 2])
                    + 0.001 * data[-2, 6] ** 3
                    - 0.001 * data[-3, 6] ** 3
                    + np.random.normal(0, 1)
                )
                x5 = np.clip(np.random.normal(0, 1), -1, 1) * 0.04 * data[
                    -2, 5
                ] ** 5 + np.random.normal(0, 1)
                x6 = (
                    0.04 * data[-2, 0] ** 3
                    + 0.04 * data[-1, 2] ** 3
                    + np.random.normal(0, 1)
                )
                x7 = np.clip(np.random.normal(0, 1), -0.5, 0.5) * (
                    0.04 * data[-2, 0] ** 3
                    + 0.1 * data[-1, 5] ** 2
                    - 0.1 * data[-2, 5] ** 2
                ) + np.random.normal(0, 1)
                data = np.vstack([data, np.array([x1, x2, x3, x4, x5, x6, x7])])
        except Exception:
            continue
        data = np.array(data[n_burnin_samples:])
        data_list.append(data)
        if len(data_list) == n_trials:
            break
elif network_to_check == "9_nonlinear":
    # 9-node nonlinear network
    true_labels = np.array(
        [
            [0, 1, 1, 1, 0, 0, 0, 1, 1],
            [0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 1, 0],
            [0, 0, 0, 0, 1, 1, 0, 0, 0],
            [0, 0, 0, 1, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 1, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 1],
            [0, 0, 0, 0, 0, 0, 0, 0, 0],
        ]
    )
    n_burnin_samples = 500
    n_trials = 50
    n_lags = 3
    data_list = []
    while True:
        np.random.seed(seed)
        seed += 1
        data = np.random.normal(0, 1, [n_lags, 9])
        try:
            for j in range(n_lags, n_time_samples_int + n_burnin_samples):
                x1 = (
                    0.95 * np.sqrt(2) * data[-1, 0]
                    - 0.9025 * data[-2, 0]
                    + np.random.normal(0, 1)
                )
                x2 = (
                    0.5 * data[-2, 0] ** 2
                    + 0.5 * data[-1, 1]
                    - 0.4 * data[-2, 1]
                    + np.random.normal(0, 1)
                )
                x3 = (
                    -0.4 * data[-3, 0]
                    + 0.5 * data[-1, 2]
                    - 0.4 * data[-2, 2]
                    + np.random.normal(0, 1)
                )
                x4 = (
                    -0.5 * data[-2, 0] ** 2
                    + 0.5 * data[-1, 3]
                    - 0.4 * data[-2, 3]
                    + 0.5 * np.sqrt(2) * data[-1, 3]
                    + 0.25 * np.sqrt(2) * data[-1, 4]
                    + np.random.normal(0, 1)
                )
                x5 = (
                    -0.5 * np.sqrt(2) * data[-1, 3]
                    + 0.5 * np.sqrt(2) * data[-1, 4]
                    + np.random.normal(0, 1)
                )
                x6 = (
                    np.log1p(np.abs(data[-1, 3])) * np.sign(data[-1, 3])
                    + 0.5 * data[-1, 5]
                    - 0.4 * data[-2, 5]
                    + np.random.normal(0, 1)
                )
                x7 = (
                    np.clip(np.random.normal(0, 1), -1, 1) * 0.04 * data[-2, 5] ** 5
                    + 0.5 * data[-1, 6]
                    - 0.4 * data[-2, 6]
                    + np.random.normal(0, 1)
                )
                x8 = (
                    0.4 * data[-2, 0]
                    + 0.25 * data[-1, 2] ** 3
                    + 0.5 * data[-1, 7]
                    - 0.4 * data[-2, 7]
                    + np.random.normal(0, 1)
                )
                x9 = (
                    np.clip(np.random.normal(0, 1), -0.5, 0.5)
                    * (
                        0.2 * data[-2, 0]
                        + 0.1 * data[-1, 7] ** 2
                        - 0.1 * data[-2, 7] ** 2
                    )
                    + 0.5 * data[-1, 8]
                    - 0.4 * data[-2, 8]
                    + np.random.normal(0, 1)
                )
                data = np.vstack([data, np.array([x1, x2, x3, x4, x5, x6, x7, x8, x9])])
        except Exception:
            continue
        data = np.array(data[n_burnin_samples:])
        data_list.append(data)
        if len(data_list) == n_trials:
            break
elif network_to_check == "11_nonlinear":
    # 11-node nonlinear network
    true_labels = np.array(
        [
            [0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0],
            [0, 0, 1, 1, 0, 0, 0, 0, 0, 1, 1],
            [0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0],
            [0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0],
            [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        ]
    )
    n_burnin_samples = 500
    n_trials = 50
    n_lags = 3
    data_list = []
    while True:
        np.random.seed(seed)
        seed += 1
        data = np.random.normal(0, 1, [n_lags, 11])
        try:
            for j in range(n_lags, n_time_samples_int + n_burnin_samples):
                x1 = (
                    0.25 * data[-1, 0] ** 2
                    - 0.25 * data[-2, 0] ** 2
                    + np.random.normal(0, 1)
                )
                x2 = np.log1p(np.abs(data[-2, 0])) * np.sign(
                    data[-2, 0]
                ) + np.random.normal(0, 1)
                x3 = -0.1 * data[-3, 1] ** 3 + np.random.normal(0, 1)
                x4 = (
                    -0.5 * data[-2, 1] ** 2
                    + 0.5 * np.sqrt(2) * data[-1, 3]
                    + 0.25 * np.sqrt(2) * data[-1, 4]
                    + np.random.normal(0, 1)
                )
                x5 = (
                    -0.5 * np.sqrt(2) * data[-1, 3]
                    + 0.5 * np.sqrt(2) * data[-1, 4]
                    + np.random.normal(0, 1)
                )
                x6 = np.log1p(np.abs(data[-1, 3])) * np.sign(
                    data[-1, 3]
                ) + np.random.normal(0, 1)
                x7 = np.clip(np.random.normal(0, 1), -1, 1) * 0.04 * data[
                    -2, 5
                ] ** 5 + np.random.normal(0, 1)
                x8 = (
                    0.4 * data[-2, 0] + 0.25 * data[-1, 2] ** 3 + np.random.normal(0, 1)
                )
                x9 = np.clip(np.random.normal(0, 1), -0.5, 0.5) * (
                    0.2 * data[-2, 0] + 0.1 * data[-1, 7] ** 2 - 0.1 * data[-2, 7] ** 2
                ) + np.random.normal(0, 1)
                x10 = (
                    0.25 * data[-3, 0] ** 2
                    - 0.01 * data[-3, 1] ** 2
                    + 0.15 * data[-3, 2] ** 3
                    + np.random.normal(0, 1)
                )
                x11 = (
                    0.1 * data[-1, 1] ** 4
                    - 0.1 * data[-2, 1] ** 4
                    + 0.1 * data[-3, 5] ** 3
                    + np.random.normal(0, 1)
                )
                data = np.vstack(
                    [data, np.array([x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11])]
                )
        except Exception:
            continue
        data = np.array(data[n_burnin_samples:])
        data_list.append(data)
        if len(data_list) == n_trials:
            break
elif network_to_check == "34_zachary1":
    # Generate data
    n_burnin_samples = 500
    n_trials = 50
    n_lags = 1
    true_labels_list = []
    data_list = []
    while True:
        np.random.seed(seed)
        seed += 1
        # Load the Zachary karate club dataset
        G_orig = nx.karate_club_graph()
        adjacency_matrix = nx.adjacency_matrix(G_orig).todense()

        true_labels = np.clip(adjacency_matrix, 0, 1).astype(int)
        true_labels_list.append(true_labels)

        data = 0.01 * np.random.normal(0, 1, [n_lags, true_labels.shape[0]])
        try:
            for t in range(n_lags, n_time_samples_int + n_burnin_samples):
                data_new = np.zeros_like(data[[-1]])
                for i in range(data_new.shape[1]):
                    data_new[0, i] = (
                        (1 - 0.025 * true_labels[:, i].sum())
                        * (1 - 1.8 * data[-1, i] ** 2)
                        + np.matmul(
                            0.025 * true_labels[:, i], (1 - 1.8 * np.square(data[-1]))
                        )
                        + 0.01 * np.random.normal(0, 1)
                    )
                data = np.vstack([data, data_new])
        except Exception:
            continue
        data = np.array(data[n_burnin_samples:])
        data_list.append(data)
        if len(data_list) == n_trials:
            break
elif network_to_check == "34_zachary2":
    # Generate data
    n_burnin_samples = 500
    n_trials = 50
    n_lags = 1
    true_labels_list = []
    data_list = []
    while True:
        np.random.seed(seed)
        seed += 1
        # Load the Zachary karate club dataset
        G_orig = nx.karate_club_graph()
        adjacency_matrix = nx.adjacency_matrix(G_orig).todense()
        comb_list = deepcopy(
            list(itertools.combinations(range(adjacency_matrix.shape[0]), 2))
        )
        list_of_edges = [(i, j) for (i, j) in comb_list if adjacency_matrix[i, j] != 0]
        # Randomly select 5 edges to keep as bidirectional
        idx_bidirectional = np.random.choice(range(len(list_of_edges)), 5)
        edges_nonbidirectional = [
            list_of_edges[i]
            for i in range(len(list_of_edges))
            if i not in idx_bidirectional
        ]
        # Randomly assign a direction
        for a, b in edges_nonbidirectional:
            if np.random.random() < 0.5:
                adjacency_matrix[a, b] = 0
            else:
                adjacency_matrix[b, a] = 0

        true_labels = np.clip(adjacency_matrix, 0, 1).astype(int)
        true_labels_list.append(true_labels)

        data = 0.01 * np.random.normal(0, 1, [n_lags, true_labels.shape[0]])
        try:
            for t in range(n_lags, n_time_samples_int + n_burnin_samples):
                data_new = np.zeros_like(data[[-1]])
                for i in range(data_new.shape[1]):
                    data_new[0, i] = (
                        (1 - 0.05 * true_labels[:, i].sum())
                        * (1 - 1.8 * data[-1, i] ** 2)
                        + np.matmul(
                            0.05 * true_labels[:, i], (1 - 1.8 * np.square(data[-1]))
                        )
                        + 0.01 * np.random.normal(0, 1)
                    )
                data = np.vstack([data, data_new])
        except Exception as e:
            print(e)
            continue
        data = np.array(data[n_burnin_samples:])
        data_list.append(data)
        if len(data_list) == n_trials:
            break


if network_to_check != "34_zachary1" and network_to_check != "34_zachary2":
    true_labels_list = [true_labels for i in range(len(data_list))]


G = nx.from_numpy_array(true_labels_list[0], create_using=nx.DiGraph)
if network_to_check != "34_zachary1" and network_to_check != "34_zachary2":
    G_pos = nx.kamada_kawai_layout(G)
else:
    G_pos = nx.circular_layout(G)

G_d = dict(G.degree)
G_labels = {node: str(node + 1) for node in G_d.keys()}
nx.draw(
    G,
    G_pos,
    labels=G_labels,
    with_labels=True,
    node_color="white",
    node_size=750,
    edgecolors="black",
    linewidths=2,
    arrows=True,
    arrowstyle="->",
    connectionstyle="arc3, rad = 0.2",
)
plt.rcParams["svg.fonttype"] = "none"
if n_time_samples_str == "500":
    if network_to_check == "5_linear":
        plt.savefig("plots/5_linear_nonlinear_network.pdf", bbox_inches="tight")
        plt.savefig("plots/5_linear_nonlinear_network.eps", bbox_inches="tight")
    elif network_to_check == "5_nonlinear":
        pass
    else:
        plt.savefig("plots/" + network_to_check + "_network.pdf", bbox_inches="tight")
        plt.savefig("plots/" + network_to_check + "_network.eps", bbox_inches="tight")
# plt.show(block=False)
# plt.pause(0.001)


# Cao's minimum embedding dimension
pprint("### Find Cao's minimum embedding dimension ###")


def caoloop(d):
    d_list = [[] for i in range(d.shape[1])]
    E1_list = [[] for i in range(d.shape[1])]
    E2_list = [[] for i in range(d.shape[1])]
    embedding_dim_list = [[] for i in range(d.shape[1])]
    for c in range(d.shape[1]):
        (
            cur_d,
            cur_E1,
            cur_E2,
            embedding_dim,
        ) = cao_min_embedding_dimension.cao_min_embedding_dimension(
            d[:, c], maxd=20, tau=1, threshold=0.8, max_relative_change=0.3
        )
        d_list[c].append(cur_d)
        E1_list[c].append(cur_E1)
        E2_list[c].append(cur_E2)
        embedding_dim_list[c].append(embedding_dim)
    return [d_list, E1_list, E2_list, embedding_dim_list]


cao_loop_output = process_map(caoloop, data_list, max_workers=num_cores)
# Ensure process_map finished before continuing:
# the following will not pass until the lock can be acquired
with ensure_lock(tqdm.auto.tqdm, lock_name="mp_lock") as lk:
    pass

# Extract output into lists of lists for easier processing downstream
d_list = [[] for i in range(data_list[0].shape[1])]
E1_list = [[] for i in range(data_list[0].shape[1])]
E2_list = [[] for i in range(data_list[0].shape[1])]
embedding_dim_list = [[] for i in range(data_list[0].shape[1])]
for o in range(len(cao_loop_output)):
    for v in range(len(cao_loop_output[0][0])):
        d_list[v].append(cao_loop_output[o][0][v][0])
        E1_list[v].append(cao_loop_output[o][1][v][0])
        E2_list[v].append(cao_loop_output[o][2][v][0])
        embedding_dim_list[v].append(cao_loop_output[o][3][v][0])


min_embed_dim = round(np.nanmax(np.array(embedding_dim_list).flatten()))

pprint("Minimum embedding dimension: " + str(min_embed_dim))


# Function to calculate metrics
def calc_metrics(pvalue_matricies, thresh):
    # Check sensitivity and specificity
    sensitivities = []
    specificities = []
    accuracies = []
    balanced_accuracies = []
    f1wgt = []
    listlen = deepcopy(len(pvalue_matricies))
    for i in range(listlen):
        preds = deepcopy(np.array(pvalue_matricies[i])).astype(float)
        np.fill_diagonal(preds, np.nan)
        preds = preds.flatten()
        preds = preds[~np.isnan(preds)]
        preds = (preds < thresh).astype(int)
        true_labels_curr = deepcopy(true_labels_list[i]).astype(float)
        np.fill_diagonal(true_labels_curr, np.nan)
        true_labels_curr = true_labels_curr.flatten()
        true_labels_curr = true_labels_curr[~np.isnan(true_labels_curr)]
        true_labels_curr = true_labels_curr.astype(int)
        tn, fp, fn, tp = confusion_matrix(
            true_labels_curr.flatten(), preds.flatten()
        ).ravel()
        sensitivity = tp / (tp + fn)
        sensitivities.append(sensitivity)
        specificity = tn / (tn + fp)
        specificities.append(specificity)
        accuracy = accuracy_score(true_labels_curr.flatten(), preds.flatten())
        accuracies.append(accuracy)
        balanced_accuracy = balanced_accuracy_score(
            true_labels_curr.flatten(), preds.flatten()
        )
        balanced_accuracies.append(balanced_accuracy)
        f1 = f1_score(true_labels_curr.flatten(), preds.flatten(), average="weighted")
        f1wgt.append(f1)
    return sensitivities, specificities, accuracies, balanced_accuracies, f1wgt


# Run mlcausality:
pprint("### Run mlcausality ###")
# Note: use minimum embedding dimension - 1 for lag if no logdiff is used and
# dimension - 2 for lag if logdiff is used
lag = max(min_embed_dim - 1, 1)


def locoloop(d):
    mlcausality_output = mlcausality.multiloco_mlcausality(
        d, lags=[lag], scaler_init_1="quantile", return_pvalue_matrix_only=True
    )
    return mlcausality_output


all_mlcausality_output = process_map(locoloop, data_list, max_workers=num_cores)
# Ensure process_map finished before continuing:
# the following will not pass until the lock can be acquired
with ensure_lock(tqdm.auto.tqdm, lock_name="mp_lock") as lk:
    pass

mlcausality_pvalue_matricies = all_mlcausality_output
mlcausality_adjacency_matricies = [1 - i for i in mlcausality_pvalue_matricies]

mlcausality_recovery_performance = recovery_performance(
    mlcausality_adjacency_matricies, true_labels_list
)
pprint("mlcausality recovery performance:")
pprint("Mean: " + str(mlcausality_recovery_performance.mean()))
pprint(
    pd.DataFrame(mlcausality_recovery_performance).quantile([0, 0.25, 0.5, 0.75, 1]).T
)

mlcausality_brier_score_loss = brier_loss_func(
    mlcausality_pvalue_matricies, true_labels_list
)
pprint("mlcausality brier score:")
pprint("Mean: " + str(mlcausality_brier_score_loss.mean()))
pprint(pd.DataFrame(mlcausality_brier_score_loss).quantile([0, 0.25, 0.5, 0.75, 1]).T)

# Check sensitivity and specificity for the mlcausality model
mlc_sens, mlc_spec, mlc_acc, mlc_bal_acc, mlc_f1wgt = calc_metrics(
    mlcausality_pvalue_matricies, 0.05
)


pprint("mlcausality sensitivities:")
pprint(pd.DataFrame(mlc_sens).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("mlcausality specificities:")
pprint(pd.DataFrame(mlc_spec).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("mlcausality (sensitivities*specificities)**0.5:")
pprint(
    pd.DataFrame((np.array(mlc_sens) * np.array(mlc_spec)) ** 0.5)
    .quantile([0, 0.25, 0.5, 0.75, 1])
    .T
)
pprint("mlcausality accuracies:")
pprint(pd.DataFrame(mlc_acc).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("mlcausality balanced accuracies:")
pprint(pd.DataFrame(mlc_bal_acc).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("mlcausality f1 scores:")
pprint(pd.DataFrame(mlc_f1wgt).quantile([0, 0.25, 0.5, 0.75, 1]).T)

optimal_j = 0
best_score = 0
for j in [round(r * 0.01, 2) for r in range(1, 100)]:
    mlc_sens2, mlc_spec2, mlc_acc2, mlc_bal_acc2, mlc_f1wgt2 = calc_metrics(
        mlcausality_pvalue_matricies, j
    )
    best_score_new = np.median((np.array(mlc_sens2) * np.array(mlc_spec2)) ** 0.5)
    if best_score_new > best_score:
        best_score = deepcopy(best_score_new)
        optimal_j = deepcopy(j)


mlc_sens2, mlc_spec2, mlc_acc2, mlc_bal_acc2, mlc_f1wgt2 = calc_metrics(
    mlcausality_pvalue_matricies, optimal_j
)


pprint("Optimal cutoff: " + str(optimal_j))
pprint("Optimal mlcausality sensitivities:")
pprint(pd.DataFrame(mlc_sens2).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("Optimal mlcausality specificities:")
pprint(pd.DataFrame(mlc_spec2).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("Optimal mlcausality (sensitivities*specificities)**0.5:")
pprint(
    pd.DataFrame((np.array(mlc_sens2) * np.array(mlc_spec2)) ** 0.5)
    .quantile([0, 0.25, 0.5, 0.75, 1])
    .T
)
pprint("Optimal mlcausality accuracies:")
pprint(pd.DataFrame(mlc_acc2).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("Optimal mlcausality balanced accuracies:")
pprint(pd.DataFrame(mlc_bal_acc2).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("Optimal mlcausality f1 scores:")
pprint(pd.DataFrame(mlc_f1wgt2).quantile([0, 0.25, 0.5, 0.75, 1]).T)


# Run lsNGC
pprint("### Run lsNGC ###")

k_f = 25
k_g = 5
ar_order = max(min_embed_dim - 1, 1)
while k_f >= 2:
    try:
        _ = granger(data_list[0].T, k_f=k_f, k_g=k_g, ar_order=ar_order, normalize=1)
        break
    except Exception:
        k_f -= 1

pprint(
    "lsNGC parameters: k_f="
    + str(k_f)
    + ", k_g="
    + str(k_g)
    + ", ar_order="
    + str(ar_order)
)


def lsNGCloop(d):
    lsNGC_output = granger(d.T, k_f=k_f, k_g=k_g, ar_order=ar_order, normalize=1)
    return lsNGC_output


all_lsNGC_output = process_map(lsNGCloop, data_list, max_workers=num_cores)
# Ensure process_map finished before continuing:
# the following will not pass until the lock can be acquired
with ensure_lock(tqdm.auto.tqdm, lock_name="mp_lock") as lk:
    pass

lsgc_adjacency_matricies = []
lsgc_fstat_matricies = []
lsgc_pvalue_matricies = []
for out in all_lsNGC_output:
    lsgc_adjacency_matricies.append(out[0])
    lsgc_fstat_matricies.append(out[1])
    pvalues_lsgc = out[2]
    np.fill_diagonal(pvalues_lsgc, 1)
    # lsgc_adjacency_matricies.append(1 - pvalues_lsgc)
    lsgc_pvalue_matricies.append(pvalues_lsgc)

lsgc_recovery_performance = recovery_performance(lsgc_fstat_matricies, true_labels_list)
# lsgc_recovery_performance = recovery_performance(lsgc_adjacency_matricies,true_labels_list)
pprint("lsNGC recovery performance:")
pprint("Mean: " + str(lsgc_recovery_performance.mean()))
pprint(pd.DataFrame(lsgc_recovery_performance).quantile([0, 0.25, 0.5, 0.75, 1]).T)

lsgc_brier_score_loss = brier_loss_func(lsgc_pvalue_matricies, true_labels_list)
pprint("lsgc brier score:")
pprint("Mean: " + str(lsgc_brier_score_loss.mean()))
pprint(pd.DataFrame(lsgc_brier_score_loss).quantile([0, 0.25, 0.5, 0.75, 1]).T)

# Check sensitivity and specificity for the lsNGC model
lsgc_sens, lsgc_spec, lsgc_acc, lsgc_bal_acc, lsgc_f1wgt = calc_metrics(
    lsgc_pvalue_matricies, 0.05
)


pprint("lsgc sensitivities:")
pprint(pd.DataFrame(lsgc_sens).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("lsgc specificities:")
pprint(pd.DataFrame(lsgc_spec).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("lsgc (sensitivities*specificities)**0.5:")
pprint(
    pd.DataFrame((np.array(lsgc_sens) * np.array(lsgc_spec)) ** 0.5)
    .quantile([0, 0.25, 0.5, 0.75, 1])
    .T
)
pprint("lsgc accuracies:")
pprint(pd.DataFrame(lsgc_acc).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("lsgc balanced accuracies:")
pprint(pd.DataFrame(lsgc_bal_acc).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("lsgc f1 scores:")
pprint(pd.DataFrame(lsgc_f1wgt).quantile([0, 0.25, 0.5, 0.75, 1]).T)

optimal_j = 0
best_score = 0
for j in [round(r * 0.01, 2) for r in range(1, 100)]:
    lsgc_sens2, lsgc_spec2, lsgc_acc2, lsgc_bal_acc2, lsgc_f1wgt2 = calc_metrics(
        lsgc_pvalue_matricies, j
    )
    best_score_new = np.median((np.array(lsgc_sens2) * np.array(lsgc_spec2)) ** 0.5)
    if best_score_new > best_score:
        best_score = deepcopy(best_score_new)
        optimal_j = deepcopy(j)


lsgc_sens2, lsgc_spec2, lsgc_acc2, lsgc_bal_acc2, lsgc_f1wgt2 = calc_metrics(
    lsgc_pvalue_matricies, optimal_j
)


pprint("Optimal cutoff: " + str(optimal_j))
pprint("Optimal lsgc sensitivities:")
pprint(pd.DataFrame(lsgc_sens2).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("Optimal lsgc specificities:")
pprint(pd.DataFrame(lsgc_spec2).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("Optimal lsgc (sensitivities*specificities)**0.5:")
pprint(
    pd.DataFrame((np.array(lsgc_sens2) * np.array(lsgc_spec2)) ** 0.5)
    .quantile([0, 0.25, 0.5, 0.75, 1])
    .T
)
pprint("Optimal lsgc accuracies:")
pprint(pd.DataFrame(lsgc_acc2).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("Optimal lsgc balanced accuracies:")
pprint(pd.DataFrame(lsgc_bal_acc2).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("Optimal lsgc f1 scores:")
pprint(pd.DataFrame(lsgc_f1wgt2).quantile([0, 0.25, 0.5, 0.75, 1]).T)


# Run causal_ccm
pprint("### Run causal_ccm ###")


def causal_ccm_loop(d):
    causal_ccm_output = causal_ccm.loop_causal_cmm(d.astype(np.float64), tau=1, E=2)
    return causal_ccm_output


all_causal_ccm_output = process_map(causal_ccm_loop, data_list, max_workers=num_cores)
# Ensure process_map finished before continuing:
# the following will not pass until the lock can be acquired
with ensure_lock(tqdm.auto.tqdm, lock_name="mp_lock") as lk:
    pass


# Generate pvalue adjacency matricies
permute_lists = list(itertools.permutations(range(data_list[0].shape[1]), 2))
causal_ccm_pvalue_matricies = []
for mlout in all_causal_ccm_output:
    cur_pvalue_matrix = np.ones([data_list[0].shape[1], data_list[0].shape[1]])
    for X_idx, y_idx in permute_lists:
        cur_pvalue_matrix[X_idx, y_idx] = mlout.loc[
            ((mlout.X == X_idx) & (mlout.y == y_idx)), "pvalue"
        ].iloc[0]
    np.fill_diagonal(cur_pvalue_matrix, 1)
    causal_ccm_pvalue_matricies.append(cur_pvalue_matrix)


# Generate adjacency matricies and check recovery performance
permute_lists = list(itertools.permutations(range(data_list[0].shape[1]), 2))
causal_ccm_adjacency_matricies = []
for mlout in all_causal_ccm_output:
    cur_pvalue_adjacency_matrix = np.zeros(
        [data_list[0].shape[1], data_list[0].shape[1]]
    )
    for X_idx, y_idx in permute_lists:
        cur_pvalue_adjacency_matrix[X_idx, y_idx] = mlout.loc[
            ((mlout.X == X_idx) & (mlout.y == y_idx)), "score"
        ].iloc[0]
    causal_ccm_adjacency_matricies.append(cur_pvalue_adjacency_matrix)

causal_ccm_recovery_performance = recovery_performance(
    causal_ccm_adjacency_matricies, true_labels_list
)
pprint("causal_ccm recovery performance:")
pprint("Mean: " + str(causal_ccm_recovery_performance.mean()))
pprint(
    pd.DataFrame(causal_ccm_recovery_performance).quantile([0, 0.25, 0.5, 0.75, 1]).T
)

causal_ccm_brier_score_loss = brier_loss_func(
    causal_ccm_pvalue_matricies, true_labels_list
)
pprint("causal_ccm brier score:")
pprint("Mean: " + str(causal_ccm_brier_score_loss.mean()))
pprint(pd.DataFrame(causal_ccm_brier_score_loss).quantile([0, 0.25, 0.5, 0.75, 1]).T)

# Check sensitivity and specificity for the causal_ccm model
(
    causal_ccm_sens,
    causal_ccm_spec,
    causal_ccm_acc,
    causal_ccm_bal_acc,
    causal_ccm_f1wgt,
) = calc_metrics(causal_ccm_pvalue_matricies, 0.05)


pprint("causal_ccm sensitivities:")
pprint(pd.DataFrame(causal_ccm_sens).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("causal_ccm specificities:")
pprint(pd.DataFrame(causal_ccm_spec).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("causal_ccm (sensitivities*specificities)**0.5:")
pprint(
    pd.DataFrame((np.array(causal_ccm_sens) * np.array(causal_ccm_spec)) ** 0.5)
    .quantile([0, 0.25, 0.5, 0.75, 1])
    .T
)
pprint("causal_ccm accuracies:")
pprint(pd.DataFrame(causal_ccm_acc).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("causal_ccm balanced accuracies:")
pprint(pd.DataFrame(causal_ccm_bal_acc).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("causal_ccm f1 scores:")
pprint(pd.DataFrame(causal_ccm_f1wgt).quantile([0, 0.25, 0.5, 0.75, 1]).T)

optimal_j = 0
best_score = 0
for j in [round(r * 0.01, 2) for r in range(1, 100)]:
    (
        causal_ccm_sens2,
        causal_ccm_spec2,
        causal_ccm_acc2,
        causal_ccm_bal_acc2,
        causal_ccm_f1wgt2,
    ) = calc_metrics(causal_ccm_pvalue_matricies, j)
    best_score_new = np.median(
        (np.array(causal_ccm_sens2) * np.array(causal_ccm_spec2)) ** 0.5
    )
    if best_score_new > best_score:
        best_score = deepcopy(best_score_new)
        optimal_j = deepcopy(j)


(
    causal_ccm_sens2,
    causal_ccm_spec2,
    causal_ccm_acc2,
    causal_ccm_bal_acc2,
    causal_ccm_f1wgt2,
) = calc_metrics(causal_ccm_pvalue_matricies, optimal_j)


pprint("Optimal cutoff: " + str(optimal_j))
pprint("Optimal causal_ccm sensitivities:")
pprint(pd.DataFrame(causal_ccm_sens2).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("Optimal causal_ccm specificities:")
pprint(pd.DataFrame(causal_ccm_spec2).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("Optimal causal_ccm (sensitivities*specificities)**0.5:")
pprint(
    pd.DataFrame((np.array(causal_ccm_sens2) * np.array(causal_ccm_spec2)) ** 0.5)
    .quantile([0, 0.25, 0.5, 0.75, 1])
    .T
)
pprint("Optimal causal_ccm accuracies:")
pprint(pd.DataFrame(causal_ccm_acc2).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("Optimal causal_ccm balanced accuracies:")
pprint(pd.DataFrame(causal_ccm_bal_acc2).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("Optimal causal_ccm f1 scores:")
pprint(pd.DataFrame(causal_ccm_f1wgt2).quantile([0, 0.25, 0.5, 0.75, 1]).T)


# Run pcmci
pprint("### Run pcmci ###")


def pmciloop(d):
    dataframe = pcmci_data_processing.DataFrame(d)
    parcorr = ParCorr()
    pcmci = PCMCI(dataframe=dataframe, cond_ind_test=parcorr, verbosity=0)
    results = pcmci.run_pcmci(
        tau_min=1, tau_max=min_embed_dim - 1, pc_alpha=None, alpha_level=0.05
    )
    pcmci_pvalues_matrix = results["p_matrix"].min(axis=2)
    np.fill_diagonal(pcmci_pvalues_matrix, 1)
    return pcmci_pvalues_matrix


pcmci_pvalue_matricies = process_map(pmciloop, data_list, max_workers=num_cores)
# Ensure process_map finished before continuing:
# the following will not pass until the lock can be acquired
with ensure_lock(tqdm.auto.tqdm, lock_name="mp_lock") as lk:
    pass

pcmci_adjacency_matricies = []
for i in range(len(pcmci_pvalue_matricies)):
    pcmci_pvalues_matrix_oneminus = 1 - pcmci_pvalue_matricies[i]
    np.fill_diagonal(pcmci_pvalues_matrix_oneminus, 0)
    pcmci_adjacency_matricies.append(pcmci_pvalues_matrix_oneminus)

pcmci_recovery_performance = recovery_performance(
    pcmci_adjacency_matricies, true_labels_list
)
pprint("pcmci recovery performance:")
pprint("Mean: " + str(pcmci_recovery_performance.mean()))
pprint(pd.DataFrame(pcmci_recovery_performance).quantile([0, 0.25, 0.5, 0.75, 1]).T)

pcmci_brier_score_loss = brier_loss_func(pcmci_pvalue_matricies, true_labels_list)
pprint("pcmci brier score:")
pprint("Mean: " + str(pcmci_brier_score_loss.mean()))
pprint(pd.DataFrame(pcmci_brier_score_loss).quantile([0, 0.25, 0.5, 0.75, 1]).T)

# Check sensitivity and specificity for the pcmci model
pcmci_sens, pcmci_spec, pcmci_acc, pcmci_bal_acc, pcmci_f1wgt = calc_metrics(
    pcmci_pvalue_matricies, 0.05
)


pprint("pcmci sensitivities:")
pprint(pd.DataFrame(pcmci_sens).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("pcmci specificities:")
pprint(pd.DataFrame(pcmci_spec).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("pcmci (sensitivities*specificities)**0.5:")
pprint(
    pd.DataFrame((np.array(pcmci_sens) * np.array(pcmci_spec)) ** 0.5)
    .quantile([0, 0.25, 0.5, 0.75, 1])
    .T
)
pprint("pcmci accuracies:")
pprint(pd.DataFrame(pcmci_acc).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("pcmci balanced accuracies:")
pprint(pd.DataFrame(pcmci_bal_acc).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("pcmci f1 scores:")
pprint(pd.DataFrame(pcmci_f1wgt).quantile([0, 0.25, 0.5, 0.75, 1]).T)

optimal_j = 0
best_score = 0
for j in [round(r * 0.01, 2) for r in range(1, 100)]:
    pcmci_sens2, pcmci_spec2, pcmci_acc2, pcmci_bal_acc2, pcmci_f1wgt2 = calc_metrics(
        pcmci_pvalue_matricies, j
    )
    best_score_new = np.median((np.array(pcmci_sens2) * np.array(pcmci_spec2)) ** 0.5)
    if best_score_new > best_score:
        best_score = deepcopy(best_score_new)
        optimal_j = deepcopy(j)


pcmci_sens2, pcmci_spec2, pcmci_acc2, pcmci_bal_acc2, pcmci_f1wgt2 = calc_metrics(
    pcmci_pvalue_matricies, optimal_j
)


pprint("Optimal cutoff: " + str(optimal_j))
pprint("Optimal pcmci sensitivities:")
pprint(pd.DataFrame(pcmci_sens2).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("Optimal pcmci specificities:")
pprint(pd.DataFrame(pcmci_spec2).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("Optimal pcmci (sensitivities*specificities)**0.5:")
pprint(
    pd.DataFrame((np.array(pcmci_sens2) * np.array(pcmci_spec2)) ** 0.5)
    .quantile([0, 0.25, 0.5, 0.75, 1])
    .T
)
pprint("Optimal pcmci accuracies:")
pprint(pd.DataFrame(pcmci_acc2).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("Optimal pcmci balanced accuracies:")
pprint(pd.DataFrame(pcmci_bal_acc2).quantile([0, 0.25, 0.5, 0.75, 1]).T)
pprint("Optimal pcmci f1 scores:")
pprint(pd.DataFrame(pcmci_f1wgt2).quantile([0, 0.25, 0.5, 0.75, 1]).T)


df_all_recovery_performance = pd.DataFrame(
    np.array(
        [
            mlcausality_recovery_performance,
            lsgc_recovery_performance,
            causal_ccm_recovery_performance,
            pcmci_recovery_performance,
        ]
    ).T,
    columns=["MLC", "lsNGC", "LM", "PCMCI"],
)
with open(
    "pickled_data/" + network_to_check + "_auc_" + n_time_samples_str + ".pickle", "wb"
) as f:  # should be 'wb' rather than 'w'
    pickle.dump(df_all_recovery_performance, f)


df_all_brier_score_loss = pd.DataFrame(
    np.array(
        [
            mlcausality_brier_score_loss,
            lsgc_brier_score_loss,
            causal_ccm_brier_score_loss,
            pcmci_brier_score_loss,
        ]
    ).T,
    columns=["MLC", "lsNGC", "LM", "PCMCI"],
)
with open(
    "pickled_data/" + network_to_check + "_brier_" + n_time_samples_str + ".pickle",
    "wb",
) as f:  # should be 'wb' rather than 'w'
    pickle.dump(df_all_brier_score_loss, f)


df_all_sensitivities = pd.DataFrame(
    np.array(
        [
            mlc_sens,
            lsgc_sens,
            causal_ccm_sens,
            pcmci_sens,
        ]
    ).T,
    columns=["MLC", "lsNGC", "LM", "PCMCI"],
)
with open(
    "pickled_data/"
    + network_to_check
    + "_sensitivity_"
    + n_time_samples_str
    + ".pickle",
    "wb",
) as f:  # should be 'wb' rather than 'w'
    pickle.dump(df_all_sensitivities, f)


df_all_specificities = pd.DataFrame(
    np.array(
        [
            mlc_spec,
            lsgc_spec,
            causal_ccm_spec,
            pcmci_spec,
        ]
    ).T,
    columns=["MLC", "lsNGC", "LM", "PCMCI"],
)
with open(
    "pickled_data/"
    + network_to_check
    + "_specificity_"
    + n_time_samples_str
    + ".pickle",
    "wb",
) as f:  # should be 'wb' rather than 'w'
    pickle.dump(df_all_specificities, f)


df_all_gmean = pd.DataFrame(
    np.array(
        [
            list((np.array(mlc_sens) * np.array(mlc_spec)) ** (1 / 2)),
            list((np.array(lsgc_sens) * np.array(lsgc_spec)) ** (1 / 2)),
            list((np.array(causal_ccm_sens) * np.array(causal_ccm_spec)) ** (1 / 2)),
            list((np.array(pcmci_sens) * np.array(pcmci_spec)) ** (1 / 2)),
        ]
    ).T,
    columns=["MLC", "lsNGC", "LM", "PCMCI"],
)
with open(
    "pickled_data/" + network_to_check + "_gmean_" + n_time_samples_str + ".pickle",
    "wb",
) as f:  # should be 'wb' rather than 'w'
    pickle.dump(df_all_gmean, f)


df_all_accuracies = pd.DataFrame(
    np.array([mlc_acc, lsgc_acc, causal_ccm_acc, pcmci_acc]).T,
    columns=["MLC", "lsNGC", "LM", "PCMCI"],
)
with open(
    "pickled_data/" + network_to_check + "_accuracy_" + n_time_samples_str + ".pickle",
    "wb",
) as f:  # should be 'wb' rather than 'w'
    pickle.dump(df_all_accuracies, f)


df_all_balanced_accuracies = pd.DataFrame(
    np.array(
        [
            mlc_bal_acc,
            lsgc_bal_acc,
            causal_ccm_bal_acc,
            pcmci_bal_acc,
        ]
    ).T,
    columns=["MLC", "lsNGC", "LM", "PCMCI"],
)
with open(
    "pickled_data/"
    + network_to_check
    + "_balanced_accuracy_"
    + n_time_samples_str
    + ".pickle",
    "wb",
) as f:  # should be 'wb' rather than 'w'
    pickle.dump(df_all_balanced_accuracies, f)


df_all_f1wgt = pd.DataFrame(
    np.array([mlc_f1wgt, lsgc_f1wgt, causal_ccm_f1wgt, pcmci_f1wgt]).T,
    columns=["MLC", "lsNGC", "LM", "PCMCI"],
)
with open(
    "pickled_data/" + network_to_check + "_f1wgt_" + n_time_samples_str + ".pickle",
    "wb",
) as f:  # should be 'wb' rather than 'w'
    pickle.dump(df_all_f1wgt, f)


df_all_sensitivities2 = pd.DataFrame(
    np.array(
        [
            mlc_sens2,
            lsgc_sens2,
            causal_ccm_sens2,
            pcmci_sens2,
        ]
    ).T,
    columns=["MLC", "lsNGC", "LM", "PCMCI"],
)
with open(
    "pickled_data/"
    + network_to_check
    + "_sensitivity2_"
    + n_time_samples_str
    + ".pickle",
    "wb",
) as f:  # should be 'wb' rather than 'w'
    pickle.dump(df_all_sensitivities2, f)


df_all_specificities2 = pd.DataFrame(
    np.array(
        [
            mlc_spec2,
            lsgc_spec2,
            causal_ccm_spec2,
            pcmci_spec2,
        ]
    ).T,
    columns=["MLC", "lsNGC", "LM", "PCMCI"],
)
with open(
    "pickled_data/"
    + network_to_check
    + "_specificity2_"
    + n_time_samples_str
    + ".pickle",
    "wb",
) as f:  # should be 'wb' rather than 'w'
    pickle.dump(df_all_specificities2, f)


df_all_gmean2 = pd.DataFrame(
    np.array(
        [
            list((np.array(mlc_sens2) * np.array(mlc_spec2)) ** (1 / 2)),
            list((np.array(lsgc_sens2) * np.array(lsgc_spec2)) ** (1 / 2)),
            list((np.array(causal_ccm_sens2) * np.array(causal_ccm_spec2)) ** (1 / 2)),
            list((np.array(pcmci_sens2) * np.array(pcmci_spec2)) ** (1 / 2)),
        ]
    ).T,
    columns=["MLC", "lsNGC", "LM", "PCMCI"],
)
with open(
    "pickled_data/" + network_to_check + "_gmean2_" + n_time_samples_str + ".pickle",
    "wb",
) as f:  # should be 'wb' rather than 'w'
    pickle.dump(df_all_gmean2, f)


df_all_accuracies2 = pd.DataFrame(
    np.array([mlc_acc2, lsgc_acc2, causal_ccm_acc2, pcmci_acc2]).T,
    columns=["MLC", "lsNGC", "LM", "PCMCI"],
)
with open(
    "pickled_data/" + network_to_check + "_accuracy2_" + n_time_samples_str + ".pickle",
    "wb",
) as f:  # should be 'wb' rather than 'w'
    pickle.dump(df_all_accuracies2, f)


df_all_balanced_accuracies2 = pd.DataFrame(
    np.array(
        [
            mlc_bal_acc2,
            lsgc_bal_acc2,
            causal_ccm_bal_acc2,
            pcmci_bal_acc2,
        ]
    ).T,
    columns=["MLC", "lsNGC", "LM", "PCMCI"],
)
with open(
    "pickled_data/"
    + network_to_check
    + "_balanced_accuracy2_"
    + n_time_samples_str
    + ".pickle",
    "wb",
) as f:  # should be 'wb' rather than 'w'
    pickle.dump(df_all_balanced_accuracies2, f)


df_all_f1wgt2 = pd.DataFrame(
    np.array([mlc_f1wgt2, lsgc_f1wgt2, causal_ccm_f1wgt2, pcmci_f1wgt2]).T,
    columns=["MLC", "lsNGC", "LM", "PCMCI"],
)
with open(
    "pickled_data/" + network_to_check + "_f1wgt2_" + n_time_samples_str + ".pickle",
    "wb",
) as f:  # should be 'wb' rather than 'w'
    pickle.dump(df_all_f1wgt2, f)


pprint("Program finished running!")

run_all_models.py
WojtekFulmyk/mlcausality-krr-paper-replication
import numpy as np
from scipy import stats, io, signal, linalg
import matplotlib.pyplot as plt
from scipy.spatial.distance import squareform, pdist
from sklearn.model_selection import StratifiedShuffleSplit
from sklearn.svm import SVC
from sklearn import metrics
from scipy.stats import f as scipyf
import warnings
from os.path import join
import os
import math
from sklearn.cluster import KMeans
import torch
from normalize_0_mean_1_std import normalize_0_mean_1_std
from calc_f_stat import calc_f_stat
from multivariate_split import multivariate_split

"""
This code includes the main functionality of the proposed method. 
"""

"""
The file includes the primary function which calculates Granger causality using lsNGC. 
"""

def lsNGC(inp_series, ar_order=1, k_f=3, k_g=2, normalize=1):
    if normalize:
        X_normalized=normalize_0_mean_1_std(inp_series)
    else:
        X_normalized=inp_series.copy()
    
    X_train, Y_train , X_test, Y_test=multivariate_split(X=X_normalized,ar_order=ar_order)
    
    X_train=torch.flatten(X_train, start_dim=1)

    km= KMeans(n_clusters= k_f, max_iter= 100, random_state=123)
    km.fit(X_train)
    cent= km.cluster_centers_


    max=0 

    for i in range(k_f):
        for j in range(k_f):
            d= np.linalg.norm(cent[i]-cent[j])
            if(d> max):
                max= d
    d= max

    sigma= d/math.sqrt(2*k_f)

    sig_d=np.zeros((np.shape(X_normalized)[0],np.shape(X_normalized)[0]));
    sig=np.zeros((np.shape(X_normalized)[0],np.shape(X_normalized)[0]));

#    Z_train_label=Y_train
    for i in range(X_normalized.shape[0]):
        Z_temp=X_normalized.copy()
        Z_train, Z_train_label , _ , _=multivariate_split(X=Z_temp,ar_order=ar_order)
        Z_train=torch.flatten(Z_train, start_dim=1)
        Z_train_label=torch.flatten(Z_train_label, start_dim=1)

        # Obtain phase space Z_s by exclusing time series of of x_s
        Z_s_train, Z_s_train_label , _ , _=multivariate_split(X=np.delete(Z_temp,[i],axis=0),ar_order=ar_order)
        # Obtain phase space reconstruction of x_s
        W_s_train, W_s_train_label , _ , _=multivariate_split(X=np.array([Z_temp[i]]),ar_order=ar_order)

        # Flatten data
        Z_s_train=torch.flatten(Z_s_train, start_dim=1)
        Z_s_train_label=torch.flatten(Z_s_train_label, start_dim=1)

        W_s_train=torch.flatten(W_s_train, start_dim=1)
        W_s_train_label=torch.flatten(W_s_train_label, start_dim=1)
        # Obtain k_g number of cluster centers in the phase space W_s with k-means clustering, will have dim=(k_g * d)
        kmg= KMeans(n_clusters= k_g, max_iter= 100, random_state=123)
        kmg.fit(W_s_train)
        cent_W_s= kmg.cluster_centers_
        # Calculate activations for each of the k_g neurons
        shape= W_s_train.shape
        row= shape[0]
        column= k_g
        G= np.empty((row,column), dtype= float)
        maxg=0 

        for ii in range(k_g):
            for jj in range(k_g):
                dg= np.linalg.norm(cent_W_s[ii]-cent_W_s[jj])
                if(dg> maxg):
                    maxg= dg
        dg= maxg

        sigmag= dg/math.sqrt(2*k_g)
        if sigmag==0:
            sigmag=1
        for ii in range(row):
            for jj in range(column):
                dist= np.linalg.norm(W_s_train[ii]-cent_W_s[jj])
                G[ii][jj]= math.exp(-math.pow(dist,2)/math.pow(2*sigmag,2))
        # Generalized radial basis function
        g_ws=np.array([G[ii]/sum(G[ii]) for ii in range(len(G))])
        # Calculate activations for each of the k_f neurons 
        shape= Z_s_train.shape
        row= shape[0]
        column= k_f
        F= np.empty((row,column), dtype= float)
        for ii in range(row):
            for jj in range(column):
                cent_temp=cent.copy()
                cent_temp=np.delete(cent_temp,np.arange(jj,jj+ar_order),axis=1)
                dist= np.linalg.norm(Z_s_train[ii]-cent_temp)
                F[ii][jj]= math.exp(-math.pow(dist,2)/math.pow(2*sigma,2))
        # Generalized radial basis function
        f_zs=np.array([F[ii]/sum(F[ii]) for ii in range(len(F))])

        # Prediction in the presence of x_s
        num_samples=f_zs.shape[0]

        f_new=np.concatenate((0.5*f_zs,0.5*g_ws),axis=1)
        GTG= np.dot(f_new.T,f_new)
        GTG_inv= np.linalg.pinv(GTG)
        fac= np.dot(GTG_inv,f_new.T)
        W_presence= np.dot(fac,Z_train_label)
        
        prediction_presence= np.dot(f_new,W_presence)
        error_presence=prediction_presence-np.array(Z_train_label)
        sig[i,:]=np.diag(np.cov(error_presence.T))

        # Prediction without x_s
        GTG= np.dot(f_zs.T,f_zs)
        GTG_inv= np.linalg.pinv(GTG)
        fac= np.dot(GTG_inv,f_zs.T)
        W_absence= np.dot(fac,Z_train_label)

        prediction_absence= np.dot(f_zs,W_absence)
        error_absence=prediction_absence-np.array(Z_train_label)
        sig_d[i,:]=np.diag(np.cov(error_absence.T))
    # Comupte the Granger causality index

    Aff=np.log(np.divide(sig_d,sig))
    Aff=(Aff>0)*Aff
    np.fill_diagonal(Aff,0)
    f_stat=calc_f_stat(sig_d, sig, n=num_samples+1, pu=k_f+k_g, pr=k_f)
    np.fill_diagonal(f_stat,0)
    n=num_samples+1
    pu=k_f+k_g
    pr=k_f
    f_dfn = pu-pr
    f_dfd = n-pu-1
    pvalue = np.empty_like(f_stat)
    for i in range(f_stat.shape[0]):
        for j in range(f_stat.shape[1]):
            pvalue[i,j] = scipyf.sf(f_stat[i,j], f_dfn, f_dfd)
    
    return Aff, f_stat, pvalue

utils_mod.py
WojtekFulmyk/mlcausality-krr-paper-replication