mst.py

import numpy as np
import tqdm

from scipy.sparse.csgraph import minimum_spanning_tree

def prerank_minimum_spanning_tree(ensemble):
    n_samples = ensemble.shape[0]
    n_members = ensemble.shape[1]
    preranks = np.zeros((n_samples, n_members))

    for i in tqdm.tqdm(list(range(n_samples))):
        for j in range(n_members):
            ensemble_omit_j = np.delete(ensemble[i], j, axis=0)
            ensemble_omit_j_t1 = ensemble_omit_j[:, np.newaxis, :]
            ensemble_omit_j_t2 = ensemble_omit_j[np.newaxis, :, :]

            distances = np.sqrt(
                np.square(ensemble_omit_j_t1 - ensemble_omit_j_t2).sum(axis=-1)
            )

            mst = minimum_spanning_tree(distances).toarray()
            preranks[i, j] = mst.sum()

    return preranks