<filename>code/baselines/utils_phy12.py

import numpy as np
import matplotlib.pyplot as plt
import torch
import random
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import time

from sklearn.metrics import roc_auc_score
from sklearn.metrics import average_precision_score

def get_data_split(base_path, split_path, split_type='random', reverse=False, baseline=True, dataset='P12', predictive_label='mortality'):
    # load data
    if dataset == 'P12':
        Pdict_list = np.load(base_path + '/processed_data/PTdict_list.npy', allow_pickle=True)
        arr_outcomes = np.load(base_path + '/processed_data/arr_outcomes.npy', allow_pickle=True)
        dataset_prefix = ''
    elif dataset == 'P19':
        Pdict_list = np.load(base_path + '/processed_data/PT_dict_list_6.npy', allow_pickle=True)
        arr_outcomes = np.load(base_path + '/processed_data/arr_outcomes_6.npy', allow_pickle=True)
        dataset_prefix = 'P19_'
    elif dataset == 'eICU':
        Pdict_list = np.load(base_path + '/processed_data/PTdict_list.npy', allow_pickle=True)
        arr_outcomes = np.load(base_path + '/processed_data/arr_outcomes.npy', allow_pickle=True)
        dataset_prefix = 'eICU_'
    elif dataset == 'PAM':
        Pdict_list = np.load(base_path + '/processed_data/PTdict_list.npy', allow_pickle=True)
        arr_outcomes = np.load(base_path + '/processed_data/arr_outcomes.npy', allow_pickle=True)
        dataset_prefix = ''  # not applicable

    show_statistics = False
    if show_statistics:
        idx_under_65 = []
        idx_over_65 = []

        idx_male = []
        idx_female = []

        # variables for statistics
        all_ages = []
        female_count = 0
        male_count = 0
        all_BMI = []

        X_static = np.zeros((len(Pdict_list), len(Pdict_list[0]['extended_static'])))
        for i in range(len(Pdict_list)):
            X_static[i] = Pdict_list[i]['extended_static']
            age, gender_0, gender_1, height, _, _, _, _, weight = X_static[i]
            if age > 0:
                all_ages.append(age)
                if age < 65:
                    idx_under_65.append(i)
                else:
                    idx_over_65.append(i)
            if gender_0 == 1:
                female_count += 1
                idx_female.append(i)
            if gender_1 == 1:
                male_count += 1
                idx_male.append(i)
            if height > 0 and weight > 0:
                all_BMI.append(weight / ((height / 100) ** 2))

        # plot statistics
        plt.hist(all_ages, bins=[i * 10 for i in range(12)])
        plt.xlabel('Years')
        plt.ylabel('# people')
        plt.title('Histogram of patients ages, age known in %d samples.\nMean: %.1f, Std: %.1f, Median: %.1f' %
                  (len(all_ages), np.mean(np.array(all_ages)), np.std(np.array(all_ages)), np.median(np.array(all_ages))))
        plt.show()

        plt.hist(all_BMI, bins=[5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60])
        all_BMI = np.array(all_BMI)
        all_BMI = all_BMI[(all_BMI > 10) & (all_BMI < 65)]
        plt.xlabel('BMI')
        plt.ylabel('# people')
        plt.title('Histogram of patients BMI, height and weight known in %d samples.\nMean: %.1f, Std: %.1f, Median: %.1f' %
                  (len(all_BMI), np.mean(all_BMI), np.std(all_BMI), np.median(all_BMI)))
        plt.show()
        print('\nGender known: %d,  Male count: %d,  Female count: %d\n' % (male_count + female_count, male_count, female_count))

    # np.save('saved/idx_under_65.npy', np.array(idx_under_65), allow_pickle=True)
    # np.save('saved/idx_over_65.npy', np.array(idx_over_65), allow_pickle=True)
    # np.save('saved/idx_male.npy', np.array(idx_male), allow_pickle=True)
    # np.save('saved/idx_female.npy', np.array(idx_female), allow_pickle=True)

    # transformer_path = True
    if baseline==True:
        BL_path = ''
    else:
        BL_path = 'baselines/'

    if split_type == 'random':
        # load random indices from a split
        idx_train, idx_val, idx_test = np.load(base_path + split_path, allow_pickle=True)
    elif split_type == 'age':
        if reverse == False:
            idx_train = np.load(BL_path+'saved/' + dataset_prefix + 'idx_under_65.npy', allow_pickle=True)
            idx_vt = np.load(BL_path+'saved/' + dataset_prefix + 'idx_over_65.npy', allow_pickle=True)
        elif reverse == True:
            idx_train = np.load(BL_path+'saved/' + dataset_prefix + 'idx_over_65.npy', allow_pickle=True)
            idx_vt = np.load(BL_path+'saved/' + dataset_prefix + 'idx_under_65.npy', allow_pickle=True)

        np.random.shuffle(idx_vt)
        idx_val = idx_vt[:round(len(idx_vt) / 2)]
        idx_test = idx_vt[round(len(idx_vt) / 2):]
    elif split_type == 'gender':
        if reverse == False:
            idx_train = np.load(BL_path+'saved/' + dataset_prefix + 'idx_male.npy', allow_pickle=True)
            idx_vt = np.load(BL_path+'saved/' + dataset_prefix + 'idx_female.npy', allow_pickle=True)
        elif reverse == True:
            idx_train = np.load(BL_path+'saved/' + dataset_prefix + 'idx_female.npy', allow_pickle=True)
            idx_vt = np.load(BL_path+'saved/' + dataset_prefix + 'idx_male.npy', allow_pickle=True)

        np.random.shuffle(idx_vt)
        idx_val = idx_vt[:round(len(idx_vt) / 2)]
        idx_test = idx_vt[round(len(idx_vt) / 2):]

    # extract train/val/test examples
    Ptrain = Pdict_list[idx_train]
    Pval = Pdict_list[idx_val]
    Ptest = Pdict_list[idx_test]

    # extract mortality labels
    if dataset == 'P12' or dataset == 'P19' or dataset == 'PAM':
        if predictive_label == 'mortality':
            y = arr_outcomes[:, -1].reshape((-1, 1))
        elif predictive_label == 'LoS':  # for P12 only
            y = arr_outcomes[:, 3].reshape((-1, 1))
            y = np.array(list(map(lambda los: 0 if los <= 3 else 1, y)))[..., np.newaxis]
    elif dataset == 'eICU':
        y = arr_outcomes[..., np.newaxis]
    ytrain = y[idx_train]
    yval = y[idx_val]
    ytest = y[idx_test]

    return Ptrain, Pval, Ptest, ytrain, yval, ytest

def getStats(P_tensor):
    N, T, F = P_tensor.shape
    Pf = P_tensor.transpose((2, 0, 1)).reshape(F, -1)
    mf = np.zeros((F, 1))
    stdf = np.ones((F, 1))
    eps = 1e-7
    for f in range(F):
        vals_f = Pf[f, :]
        vals_f = vals_f[vals_f > 0]
        mf[f] = np.mean(vals_f)
        stdf[f] = np.std(vals_f)
        stdf[f] = np.max([stdf[f], eps])
    return mf, stdf

def get_features_mean(X_features):
    """
    Calculate means of all time series features (36 features in P12 dataset).

    :param X_features: time series features for all samples in training set
    :return: list of means for all features
    """
    samples, timesteps, features = X_features.shape
    X = np.reshape(X_features, newshape=(samples*timesteps, features)).T
    means = []
    for row in X:
        row = row[row > 0]
        means.append(np.mean(row))
    return means

def mean_imputation(X_features, X_time, mean_features, missing_value_num):
    """
    Fill X_features missing values with mean values of all train samples.

    :param X_features: time series features for all samples
    :param X_time: times, when observations were measured
    :param mean_features: mean values of features from the training set
    :return: X_features, filled with mean values instead of zeros (missing observations)
    """
    time_length = []
    for times in X_time:
        if np.where(times == missing_value_num)[0].size == 0:
            time_length.append(times.shape[0])
        elif np.where(times == missing_value_num)[0][0] == 0:
            time_length.append(np.where(times == missing_value_num)[0][1])
        else:
            time_length.append(np.where(times == missing_value_num)[0][0])

    # check for inconsistency
    for i in range(len(X_features)):
        if np.any(X_features[i, time_length[i]:, :]):
            print('Inconsistency between X_features and X_time: features are measured without time stamp.')

    # impute times series features
    for i, sample in enumerate(X_features):
        X_features_relevant = sample[:time_length[i], :]
        missing_values_idx = np.where(X_features_relevant == missing_value_num)
        for row, col in zip(*missing_values_idx):
            X_features[i, row, col] = mean_features[col]

    return X_features

def forward_imputation(X_features, X_time, missing_value_num):
    """
    Fill X_features missing values with values, which are the same as its last measurement.

    :param X_features: time series features for all samples
    :param X_time: times, when observations were measured
    :return: X_features, filled with last measurements instead of zeros (missing observations)
    """
    time_length = []
    for times in X_time:
        if np.where(times == missing_value_num)[0].size == 0:
            time_length.append(times.shape[0])
        elif np.where(times == missing_value_num)[0][0] == 0:
            time_length.append(np.where(times == missing_value_num)[0][1])
        else:
            time_length.append(np.where(times == missing_value_num)[0][0])

    # impute times series features
    for i, sample in enumerate(X_features):
        for j, ts in enumerate(sample.T):   # note the transposed matrix
            first_observation = True
            current_value = -1
            for k, observation in enumerate(ts[:time_length[i]]):
                if X_features[i, k, j] == missing_value_num and first_observation:
                    continue
                elif X_features[i, k, j] != missing_value_num:
                    current_value = X_features[i, k, j]
                    first_observation = False
                elif X_features[i, k, j] == missing_value_num and not first_observation:
                    X_features[i, k, j] = current_value

    return X_features

def cubic_spline_imputation(X_features, X_time, missing_value_num):
    """
    Fill X_features missing values with cubic spline interpolation.

    :param X_features: time series features for all samples
    :param X_time: times, when observations were measured
    :return: X_features, filled with interpolated values
    """
    from scipy.interpolate import CubicSpline

    time_length = []
    for times in X_time:
        if np.where(times == missing_value_num)[0].size == 0:
            time_length.append(times.shape[0])
        elif np.where(times == missing_value_num)[0][0] == 0:
            time_length.append(np.where(times == missing_value_num)[0][1])
        else:
            time_length.append(np.where(times == missing_value_num)[0][0])

    # impute times series features
    for i, sample in enumerate(X_features):
        for j, ts in enumerate(sample.T):   # note the transposed matrix
            valid_ts = ts[:time_length[i]]
            zero_idx = np.where(valid_ts == missing_value_num)[0]
            non_zero_idx = np.nonzero(valid_ts)[0]
            y = valid_ts[non_zero_idx]

            if len(y) > 1:   # we need at least 2 observations to fit cubic spline
                x = X_time[i, :time_length[i], 0][non_zero_idx]
                x2interpolate = X_time[i, :time_length[i], 0][zero_idx]

                cs = CubicSpline(x, y)
                interpolated_ts = cs(x2interpolate)
                valid_ts[zero_idx] = interpolated_ts

                # set values before first measurement to the value of first measurement
                first_obs_index = non_zero_idx[0]
                valid_ts[:first_obs_index] = np.full(shape=first_obs_index, fill_value=valid_ts[first_obs_index])

                # set values after last measurement to the value of last measurement
                last_obs_index = non_zero_idx[-1]
                valid_ts[last_obs_index:] = np.full(shape=time_length[i] - last_obs_index, fill_value=valid_ts[last_obs_index])

                X_features[i, :time_length[i], j] = valid_ts

    return X_features

def mask_normalize(P_tensor, mf, stdf):
    """ Normalize time series variables. Missing ones are set to zero after normalization. """
    N, T, F = P_tensor.shape
    Pf = P_tensor.transpose((2,0,1)).reshape(F,-1)
    M = 1*(P_tensor>0) + 0*(P_tensor<=0)
    M_3D = M.transpose((2, 0, 1)).reshape(F, -1)
    for f in range(F):
        Pf[f] = (Pf[f]-mf[f])/(stdf[f]+1e-18)
    Pf = Pf * M_3D
    Pnorm_tensor = Pf.reshape((F,N,T)).transpose((1,2,0))
    Pfinal_tensor = np.concatenate([Pnorm_tensor, M], axis=2)
    return Pfinal_tensor

def getStats_static(P_tensor, dataset='P12'):
    N, S = P_tensor.shape
    Ps = P_tensor.transpose((1, 0))
    ms = np.zeros((S, 1))
    ss = np.ones((S, 1))

    if dataset == 'P12':
        # ['Age' 'Gender=0' 'Gender=1' 'Height' 'ICUType=1' 'ICUType=2' 'ICUType=3' 'ICUType=4' 'Weight']
        bool_categorical = [0, 1, 1, 0, 1, 1, 1, 1, 0]
    elif dataset == 'P19':
        # ['Age' 'Gender' 'Unit1' 'Unit2' 'HospAdmTime' 'ICULOS']
        bool_categorical = [0, 1, 0, 0, 0, 0]
    elif dataset == 'eICU':
        # ['apacheadmissiondx' 'ethnicity' 'gender' 'admissionheight' 'admissionweight'] -> 399 dimensions
        bool_categorical = [1] * 397 + [0] * 2

    for s in range(S):
        if bool_categorical[s] == 0:  # if not categorical
            vals_s = Ps[s, :]
            vals_s = vals_s[vals_s > 0]
            ms[s] = np.mean(vals_s)
            ss[s] = np.std(vals_s)
    return ms, ss

def mask_normalize_static(P_tensor, ms, ss):
    N, S = P_tensor.shape
    Ps = P_tensor.transpose((1, 0))

    # input normalization
    for s in range(S):
        Ps[s] = (Ps[s] - ms[s]) / (ss[s] + 1e-18)

    # set missing values to zero after normalization
    for s in range(S):
        idx_missing = np.where(Ps[s, :] <= 0)
        Ps[s, idx_missing] = 0

    # reshape back
    Pnorm_tensor = Ps.reshape((S, N)).transpose((1, 0))
    return Pnorm_tensor

def tensorize_normalize(P, y, mf, stdf, ms, ss):
    T, F = P[0]['arr'].shape
    D = len(P[0]['extended_static'])

    P_tensor = np.zeros((len(P), T, F))
    P_time = np.zeros((len(P), T, 1))
    P_static_tensor = np.zeros((len(P), D))
    for i in range(len(P)):
        P_tensor[i] = P[i]['arr']
        P_time[i] = P[i]['time']
        P_static_tensor[i] = P[i]['extended_static']
    P_tensor = mask_normalize(P_tensor, mf, stdf)
    P_tensor = torch.Tensor(P_tensor)

    P_time = torch.Tensor(P_time) / 60.0  # convert mins to hours
    P_static_tensor = mask_normalize_static(P_static_tensor, ms, ss)
    P_static_tensor = torch.Tensor(P_static_tensor)

    y_tensor = y
    y_tensor = torch.Tensor(y_tensor[:, 0]).type(torch.LongTensor)
    return P_tensor, P_static_tensor, P_time, y_tensor

def tensorize_normalize_other(P, y, mf, stdf):
    T, F = P[0].shape

    P_time = np.zeros((len(P), T, 1))
    for i in range(len(P)):
        tim = torch.linspace(0, T, T).reshape(-1, 1)
        P_time[i] = tim
    P_tensor = mask_normalize(P, mf, stdf)
    P_tensor = torch.Tensor(P_tensor)

    P_time = torch.Tensor(P_time) / 60.0

    y_tensor = y
    y_tensor = torch.Tensor(y_tensor[:, 0]).type(torch.LongTensor)
    return P_tensor, None, P_time, y_tensor

def masked_softmax(A, epsilon=0.000000001):
    A_max = torch.max(A, dim=1, keepdim=True)[0]
    A_exp = torch.exp(A - A_max)
    A_exp = A_exp * (A != 0).float()
    A_softmax = A_exp / (torch.sum(A_exp, dim=0, keepdim=True) + epsilon)
    return A_softmax

def random_sample(idx_0, idx_1, B, replace=False):
    """ Returns a balanced sample of tensors by randomly sampling without replacement. """
    idx0_batch = np.random.choice(idx_0, size=int(B / 2), replace=replace)
    idx1_batch = np.random.choice(idx_1, size=int(B / 2), replace=replace)
    idx = np.concatenate([idx0_batch, idx1_batch], axis=0)
    return idx

def random_sample_8(ytrain, B, replace=False):
    """ Returns a balanced sample of tensors by randomly sampling without replacement. """
    idx0_batch = np.random.choice(np.where(ytrain == 0)[0], size=int(B / 8), replace=replace)
    idx1_batch = np.random.choice(np.where(ytrain == 1)[0], size=int(B / 8), replace=replace)
    idx2_batch = np.random.choice(np.where(ytrain == 2)[0], size=int(B / 8), replace=replace)
    idx3_batch = np.random.choice(np.where(ytrain == 3)[0], size=int(B / 8), replace=replace)
    idx4_batch = np.random.choice(np.where(ytrain == 4)[0], size=int(B / 8), replace=replace)
    idx5_batch = np.random.choice(np.where(ytrain == 5)[0], size=int(B / 8), replace=replace)
    idx6_batch = np.random.choice(np.where(ytrain == 6)[0], size=int(B / 8), replace=replace)
    idx7_batch = np.random.choice(np.where(ytrain == 7)[0], size=int(B / 8), replace=replace)
    idx = np.concatenate([idx0_batch, idx1_batch, idx2_batch, idx3_batch, idx4_batch, idx5_batch, idx6_batch, idx7_batch], axis=0)
    return idx

def evaluate(model, P_tensor, P_time_tensor, P_static_tensor, batch_size=100, n_classes=2, static=1):
    model.eval()
    P_tensor = P_tensor.cuda()
    P_time_tensor = P_time_tensor.cuda()
    if static is None:
        Pstatic = None
    else:
        P_static_tensor = P_static_tensor.cuda()
        N, Fs = P_static_tensor.shape

    T, N, Ff = P_tensor.shape

    n_batches, rem = N // batch_size, N % batch_size

    out = torch.zeros(N, n_classes)
    start = 0
    for i in range(n_batches):
        P = P_tensor[:, start:start + batch_size, :]
        Ptime = P_time_tensor[:, start:start + batch_size]
        if P_static_tensor is not None:
            Pstatic = P_static_tensor[start:start + batch_size]
        lengths = torch.sum(Ptime > 0, dim=0)
        out[start:start + batch_size] = model.forward(P, Pstatic, Ptime, lengths).detach().cpu()
        start += batch_size
    if rem > 0:
        P = P_tensor[:, start:start + rem, :]
        Ptime = P_time_tensor[:, start:start + rem]
        if P_static_tensor is not None:
            Pstatic = P_static_tensor[start:start + batch_size]
        lengths = torch.sum(Ptime > 0, dim=0)
        out[start:start + rem] = model.forward(P, Pstatic, Ptime, lengths).detach().cpu()
    return out

def evaluate_standard(model, P_tensor, P_time_tensor, P_static_tensor, batch_size=100, n_classes=2, static=1):
    P_tensor = P_tensor.cuda()
    P_time_tensor = P_time_tensor.cuda()
    if static is None:
        P_static_tensor = None
    else:
        P_static_tensor = P_static_tensor.cuda()

    lengths = torch.sum(P_time_tensor > 0, dim=0)
    out = model.forward(P_tensor, P_static_tensor, P_time_tensor, lengths)
    return out

def evaluate_MTGNN(model, P_tensor, P_static_tensor, static=1):
    P_tensor = P_tensor.cuda()

    P_tensor = torch.permute(P_tensor, (1, 0, 2))
    P_tensor = torch.unsqueeze(P_tensor, dim=1)
    P_tensor = P_tensor.transpose(2, 3)

    if static is None:
        P_static_tensor = None
    else:
        P_static_tensor = P_static_tensor.cuda()

    out = model.forward(P_tensor, P_static_tensor)
    return out

def evaluate_DGM2(model, P_tensor, P_static_tensor, static=1):
    # suppose P_time is equal in all patients
    P_time = torch.arange(P_tensor.size()[0])

    P_tensor = P_tensor.cuda()
    P_tensor = torch.permute(P_tensor, (1, 0, 2))

    if static is None:
        P_static_tensor = None
    else:
        P_static_tensor = P_static_tensor.cuda()

    out = model.forward(P_tensor, P_time, P_static_tensor)
    return out

def linspace_vector(start, end, n_points):
    # start is either one value or a vector
    size = np.prod(start.size())

    assert(start.size() == end.size())
    if size == 1:
        # start and end are 1d-tensors
        res = torch.linspace(start, end, n_points)
    else:
        # start and end are vectors
        res = torch.Tensor()
        for i in range(0, start.size(0)):
            res = torch.cat((res,
                torch.linspace(start[i], end[i], n_points)),0)
        res = torch.t(res.reshape(start.size(0), n_points))
    return res

# Adam using warmup
class NoamOpt:
    "Optim wrapper that implements rate."

    def __init__(self, model_size, factor, warmup, optimizer):
        self.optimizer = optimizer
        self._step = 0
        self.warmup = warmup
        self.factor = factor
        self.model_size = model_size
        self._rate = 0

    def step(self):
        "Update parameters and rate"
        self._step += 1
        rate = self.rate()
        for p in self.optimizer.param_groups:
            p['lr'] = rate
        self._rate = rate
        self.optimizer.step()

    def rate(self, step=None):
        "Implement `lrate` above"
        if step is None:
            step = self._step
        return self.factor * \
               (self.model_size ** (-0.5) *
                min(step ** (-0.5), step * self.warmup ** (-1.5)))

    def zero_grad(self):
        self.optimizer.zero_grad()


from tqdm import tqdm
from itertools import chain

from torch.nn import BCEWithLogitsLoss
from torch_geometric.data import DataLoader
from sklearn.metrics import roc_auc_score, average_precision_score
from tensorboardX import SummaryWriter
import os
from models import DGCNN
import argparse
import os.path as osp

import numpy as np
import pandas as pd
import scipy.sparse as sp
from scipy.sparse.csgraph import shortest_path
import torch
import torch.nn.functional as F
# from gae_huawei import train, test
from models import GCN, SGC, GAT

from torch_geometric.nn import GAE
from torch_geometric.data import Data
from torch_geometric.datasets import Planetoid
from torch_geometric.utils import (negative_sampling, add_self_loops,
                                   train_test_split_edges, k_hop_subgraph,
                                   to_scipy_sparse_matrix)

import warnings
warnings.filterwarnings("ignore")

max_z = 0

def set_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    torch.backends.cudnn.deterministic = True

def load_data(path):
    alarm_graph = np.load(path, allow_pickle=True)
    node_list = list(alarm_graph.nodes)
    edge_tmp = []
    for edge in list(alarm_graph.edges):
        if edge[0] != edge[1]:
            edge_tmp.append(edge)  # 这一步是为了保证没有self-loops
    edge_list = []
    for i in range(len(edge_tmp)):
        a = node_list.index(edge_tmp[i][0])
        b = node_list.index(edge_tmp[i][1])
        edge_list.append([a, b])

    alarm_names = []
    for ne_name in list(alarm_graph.nodes):
        for alarm in alarm_graph.nodes[ne_name].keys():
            if alarm != 'NE_TYPE' and alarm not in alarm_names:
                alarm_names.append(alarm)

    labels = np.zeros([len(node_list), 3])
    for i in range(len(alarm_graph.nodes)):
        if alarm_graph.nodes[list(alarm_graph.nodes)[i]]['NE_TYPE'] == 'NODEB':
            labels[i][0] = 1
        elif alarm_graph.nodes[list(alarm_graph.nodes)[i]]['NE_TYPE'] == 'MICROWAVE':
            labels[i][1] = 1
        elif alarm_graph.nodes[list(alarm_graph.nodes)[i]]['NE_TYPE'] == 'ROUTER':
            labels[i][2] = 1

    label_list = []
    for i in range(len(alarm_graph.nodes)):
        if alarm_graph.nodes[list(alarm_graph.nodes)[i]]['NE_TYPE'] == 'NODEB':
            label_list.append(1)
        elif alarm_graph.nodes[list(alarm_graph.nodes)[i]]['NE_TYPE'] == 'MICROWAVE':
            label_list.append(2)
        elif alarm_graph.nodes[list(alarm_graph.nodes)[i]]['NE_TYPE'] == 'ROUTER':
            label_list.append(3)

    attribute_length = len(alarm_names)
    num_of_nodes = len(alarm_graph.nodes)
    attribute_one_hot = np.zeros([num_of_nodes, attribute_length])

    # one-hot
    for i in range(len(alarm_graph.nodes)):
        for alarm in alarm_graph.nodes[list(alarm_graph.nodes)[i]].keys():
            if alarm != 'NE_TYPE':
                attribute_one_hot[i][alarm_names.index(alarm)] = 1
    return node_list, edge_list, attribute_one_hot, labels, label_list

def drnl_node_labeling(edge_index, src, dst, num_nodes=None):
    global max_z
    # Double-radius node labeling (DRNL).
    src, dst = (dst, src) if src > dst else (src, dst)
    adj = to_scipy_sparse_matrix(edge_index, num_nodes=num_nodes).tocsr()

    idx = list(range(src)) + list(range(src + 1, adj.shape[0]))
    adj_wo_src = adj[idx, :][:, idx]

    idx = list(range(dst)) + list(range(dst + 1, adj.shape[0]))
    adj_wo_dst = adj[idx, :][:, idx]

    dist2src = shortest_path(adj_wo_dst, directed=False, unweighted=True,
                             indices=src)
    dist2src = np.insert(dist2src, dst, 0, axis=0)
    dist2src = torch.from_numpy(dist2src)

    dist2dst = shortest_path(adj_wo_src, directed=False, unweighted=True,
                             indices=dst - 1)
    dist2dst = np.insert(dist2dst, src, 0, axis=0)
    dist2dst = torch.from_numpy(dist2dst)

    dist = dist2src + dist2dst
    dist_over_2, dist_mod_2 = dist // 2, dist % 2

    z = 1 + torch.min(dist2src, dist2dst)
    z += dist_over_2 * (dist_over_2 + dist_mod_2 - 1)
    z[src] = 1.
    z[dst] = 1.
    z[torch.isnan(z)] = 0.

    max_z = max(int(z.max()), max_z)

    return z.to(torch.long)

def extract_enclosing_subgraphs(data, link_index, edge_index, y):
    data_list = []
    for src, dst in tqdm(link_index.t().tolist(), desc='Extracting...'):
        # src: source   dst: destination
        sub_nodes, sub_edge_index, mapping, _ = k_hop_subgraph(
            [src, dst], num_hops=2, edge_index=edge_index, relabel_nodes=True
        )
        src, dst = mapping.tolist()

        # remove target link from the subgraph
        mask1 = (sub_edge_index[0] != src) | (sub_edge_index[1] != dst)
        mask2 = (sub_edge_index[0] != dst) | (sub_edge_index[1] != src)
        sub_edge_index = sub_edge_index[:, mask1 & mask2]

        # calculate node labeling
        z = drnl_node_labeling(sub_edge_index, src, dst, num_nodes=sub_nodes.size(0))

        sub_data = Data(x=data.x[sub_nodes], z=z, edge_index=sub_edge_index, y=y)
        if 'pretrained_features' in data.keys:
            sub_data.pretrained_features = data.pretrained_features[sub_nodes]
        if 'alarm_features' in data.keys:
            sub_data.alarm_features = data.alarm_features[sub_nodes]

        data_list.append(sub_data)

    return data_list

def load_huawei_dataset():
    path = '../data/alarm_project_hitsz/preprocessed/G'
    nodes, edge_list, attribute, node_labels, labels = load_data(path)
    dataset = Data(x=torch.tensor(attribute, dtype=torch.float),
                   edge_index=torch.tensor(edge_list, dtype=torch.long).t().contiguous(),
                   y=torch.tensor(node_labels, dtype=torch.float), labels=labels)
    return dataset

def load_disease_dataset():
    path = '../data/disease_lp/'
    edges = pd.read_csv(path + 'disease_lp.edges.csv')
    labels = np.load(path + 'disease_lp.labels.npy')
    features = sp.load_npz(path + 'disease_lp.feats.npz').todense()
    dataset = Data(
        x=torch.tensor(features, dtype=torch.float),
        edge_index=torch.tensor(edges.values).t().contiguous(),
        y=F.one_hot(torch.tensor(labels))
    )
    return dataset

def load_cora_dataset():
    path = osp.join(osp.dirname(osp.relpath('seal_dataset.py')), '..', 'data', 'Planetoid')
    dataset = Planetoid(path, 'Cora')[0]
    dataset.y = F.one_hot(dataset.y).to(torch.float)
    # 这里做了一个处理，将标签转化维one-hot向量
    return dataset

def pre_train(model, data, optimizer):
    model.train()
    optimizer.zero_grad()
    z = model.encode(data.x, data.train_pos_edge_index)
    loss = model.recon_loss(z, data.train_pos_edge_index)
    loss.backward()
    optimizer.step()
    return float(loss), z

def pre_test(model, x, train_pos_edge_index, pos_edge_index, neg_edge_index):
    model.eval()
    with torch.no_grad():
        z = model.encode(x, train_pos_edge_index)
    return model.test(z, pos_edge_index, neg_edge_index)

def process(args):

    seed = 327
    set_seed(seed)

    print('Loading dataset~~~')
    if args.dataset == 'huawei':
        dataset = load_huawei_dataset()
        if args.use_alarm:
            alarm_feature_path = '../data/alarm_construct_graph/embedding_10.pt'
            dataset.alarm_features = torch.load(alarm_feature_path)
    elif args.dataset == 'disease':
        dataset = load_disease_dataset()
    elif args.dataset == 'cora':
        dataset = load_cora_dataset()
    else:
        raise ValueError("Invalid dataset type")

    data = train_test_split_edges(dataset, val_ratio=args.val_ratio, test_ratio=args.test_ratio)

    # =============================================================================================
    # the flowing step is just for huawei dataset
    if args.dataset == 'huawei' and args.hierarchical:
        # data.val_pos_edge_index
        edge_index = data.val_pos_edge_index.t().tolist()
        mask = [False] * len(edge_index)
        for i in range(len(edge_index)):
            if (data.labels[edge_index[i][0]] == 1 and data.labels[edge_index[i][1]] != 1) \
                    or (data.labels[edge_index[i][0]] != 1 and data.labels[edge_index[i][1]] == 1):
                mask[i] = True
        data.val_pos_edge_index = torch.tensor(edge_index, dtype=torch.long)[mask].t().contiguous()

        # data.val_neg_edge_index
        edge_index = data.val_neg_edge_index.t().tolist()
        mask = [False] * len(edge_index)
        for i in range(len(edge_index)):
            if (data.labels[edge_index[i][0]] == 1 and data.labels[edge_index[i][1]] != 1) \
                    or (data.labels[edge_index[i][0]] != 1 and data.labels[edge_index[i][1]] == 1):
                mask[i] = True
        data.val_neg_edge_index = torch.tensor(edge_index, dtype=torch.long)[mask].t().contiguous()

        # data.test_neg_edge_index
        edge_index = data.test_pos_edge_index.t().tolist()
        mask = [False] * len(edge_index)
        for i in range(len(edge_index)):
            if (data.labels[edge_index[i][0]] == 1 and data.labels[edge_index[i][1]] != 1) \
                    or (data.labels[edge_index[i][0]] != 1 and data.labels[edge_index[i][1]] == 1):
                mask[i] = True
        data.test_pos_edge_index = torch.tensor(edge_index, dtype=torch.long)[mask].t().contiguous()

        # data.test_neg_edge_index
        edge_index = data.test_neg_edge_index.t().tolist()
        mask = [False] * len(edge_index)
        for i in range(len(edge_index)):
            if (data.labels[edge_index[i][0]] == 1 and data.labels[edge_index[i][1]] != 1) \
                    or (data.labels[edge_index[i][0]] != 1 and data.labels[edge_index[i][1]] == 1):
                mask[i] = True
        data.test_neg_edge_index = torch.tensor(edge_index, dtype=torch.long)[mask].t().contiguous()
    # =================================================================================================

    edge_index, _ = add_self_loops(data.train_pos_edge_index)
    data.train_neg_edge_index = negative_sampling(
        edge_index=edge_index,
        num_nodes=data.num_nodes,
        num_neg_samples=data.train_pos_edge_index.size(1)
    )
    print('The dataset and the split edges are done!!!')

    if args.pretrain:
        pretrained_data = data.clone()
        pretrained_data.train_pos_edge_index = torch.cat(
            (pretrained_data.train_pos_edge_index, pretrained_data.train_neg_edge_index), dim=1
        )
        pretrained_data.train_neg_edge_index = None
        # 这个地方如果选择pretrain，应该采取negative injection的方式，重新训练得到特征
        print('-' * 60)
        print('Pretraining')
        if args.pre_encoder == 'GCN':
            pre_model = GAE(GCN(dataset.num_features, 32))
        else:
            raise ValueError('Invalid model type!')

        optimizer = torch.optim.Adam(pre_model.parameters(), lr=0.001)

        best_auc = 0
        patience = 0
        for pretrained_epoch in range(1, args.pretrained_epochs):
            train_loss, node_embedding = pre_train(pre_model, pretrained_data, optimizer)
            val_auc, val_ap = pre_test(pre_model, data.x, data.train_pos_edge_index,
                                       data.val_pos_edge_index, data.val_neg_edge_index)
            print(f"Epoch: {pretrained_epoch:03d}, Loss: {train_loss:.4f}, Val_AUC: {val_auc:.4f}, Val_AP: {val_ap:.4f}")
            if val_auc > best_auc:
                best_auc = val_auc
                patience = 0
            else:
                patience += 1
            if patience > args.patience:
                break

        print('-' * 60)
        print('Finished pretraining')
        data.pretrained_features = node_embedding.detach()

    if args.embedding == 'DRNL':
        pass
    else:
        data.x = data.y

    print('Starting extracting subgraphs~~~')
    # collect a list of subgraphs of training, validation and test
    train_pos_list = extract_enclosing_subgraphs(
        data, data.train_pos_edge_index, data.train_pos_edge_index, 1
    )
    train_neg_list = extract_enclosing_subgraphs(
        data, data.train_neg_edge_index, data.train_pos_edge_index, 0
    )

    val_pos_list = extract_enclosing_subgraphs(
        data, data.val_pos_edge_index, data.train_pos_edge_index, 1
    )
    val_neg_list = extract_enclosing_subgraphs(
        data, data.val_neg_edge_index, data.train_pos_edge_index, 0
    )

    test_pos_list = extract_enclosing_subgraphs(
        data, data.test_pos_edge_index, data.train_pos_edge_index, 1
    )
    test_neg_list = extract_enclosing_subgraphs(
        data, data.test_neg_edge_index, data.train_pos_edge_index, 0
    )
    print('Finished extracting subgraphs.')

    if args.embedding == 'DRNL':
        # convert labels to one-hot features
        for data in chain(train_pos_list, train_neg_list,
                          val_pos_list, val_neg_list,
                          test_pos_list, test_neg_list):
            data.x = F.one_hot(data.z, max_z + 1).to(torch.float)
    elif args.embedding == 'DRNL_SelfFeat':
        for data in chain(train_pos_list, train_neg_list,
                          val_pos_list, val_neg_list,
                          test_pos_list, test_neg_list):
            data.x = torch.cat((F.one_hot(data.z, max_z + 1).to(torch.float), data.x), dim=1)
    elif args.embedding == 'SelfFeat':
        pass
    else:
        raise ValueError("Unsupported embedding type.")

    if args.pretrain:
        for data in chain(train_pos_list, train_neg_list,
                          val_pos_list, val_neg_list,
                          test_pos_list, test_neg_list):
            data.x = torch.cat((data.x, data.pretrained_features), dim=1)
            data.pretrained_features = None

    if args.use_alarm:
        for data in chain(train_pos_list, train_neg_list,
                          val_pos_list, val_neg_list,
                          test_pos_list, test_neg_list):
            data.x = torch.cat((data.x, data.alarm_features), dim=1)
            data.alarm_features = None

    return train_pos_list + train_neg_list, val_pos_list + val_neg_list, test_pos_list + test_neg_list

def train(model, train_loader, device, optimizer, train_dataset):
    model.train()

    total_loss = 0
    for data in train_loader:
        data = data.to(device)
        optimizer.zero_grad()
        logits = model(data.x, data.edge_index, data.batch)
        loss = BCEWithLogitsLoss()(logits.view(-1), data.y.to(torch.float))
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * data.num_graphs

    return total_loss / len(train_dataset)

@torch.no_grad()
def test(loader, model, device):
    model.eval()

    y_pred, y_true = [], []
    for data in loader:
        data = data.to(device)
        logits = model(data.x, data.edge_index, data.batch)
        y_pred.append(logits.view(-1).cpu())
        y_true.append(data.y.view(-1).cpu().to(torch.float))

    return roc_auc_score(torch.cat(y_true), torch.cat(y_pred)), \
           average_precision_score(torch.cat(y_true), torch.cat(y_pred))

def run():
    parser = argparse.ArgumentParser("Configurations for seal")
    parser.add_argument('--dataset', default='huawei', type=str, help='dataset')
    parser.add_argument('--embedding', default='DRNL', type=str,
                        help='node encoding(["DRNL", "DRNL_SelfFeat", "SelfFeat"])')
    parser.add_argument('--epochs', default=101, type=int, help='training epochs')
    parser.add_argument('--cuda', default=torch.cuda.is_available(), type=bool, help='cuda')
    parser.add_argument('--lr', default=0.0001, type=float, help='learning rate')
    parser.add_argument('--weight_decay', default=5e-4, type=float, help='weight decay')
    parser.add_argument('--val_ratio', default=0.05, type=float)
    parser.add_argument('--test_ratio', default=0.10, type=float)
    parser.add_argument('--batch_size', default=32, type=int, help='batch size')
    parser.add_argument('--pretrain', action='store_true')
    parser.add_argument('--pretrained_epochs', default=401, type=int)
    parser.add_argument('--pre_encoder', default='GCN', type=str, choices=['GCN'])
    parser.add_argument('--patience', default=50, type=int, help='early stop steps')
    parser.add_argument('--use_alarm', action='store_true')
    parser.add_argument('--hierarchical', action='store_true')
    args = parser.parse_args()
    print(args)

    args.split_ratio = str(int((1-args.val_ratio-args.test_ratio)*100)) \
                       + str(int(args.val_ratio*100)) + str(int(args.test_ratio*100))

    train_dataset, val_dataset, test_dataset = process(args)

    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False)
    test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False)

    device = torch.device('cuda:1' if args.cuda else 'cpu')
    model = DGCNN(train_dataset, hidden_channels=32, num_layers=3).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)

    trainWriter = SummaryWriter('../{}/{}/{}/{}/{}'.format(
        'runs', 'SEAL', args.dataset + '_' + args.split_ratio + '_pretrained_' + str(args.pretrain),
        args.embedding, 'Train'
    ))
    valWriter = SummaryWriter('../{}/{}/{}/{}/{}'.format(
        'runs', 'SEAL', args.dataset + '_' + args.split_ratio + '_pretrained_' + str(args.pretrain),
        args.embedding, 'Val'
    ))

    best_val_auc = test_auc = test_ap = 0

    for epoch in range(1, args.epochs):
        loss = train(model, train_loader, device, optimizer, train_dataset)
        trainWriter.add_scalar(tag='Train Loss', scalar_value=loss, global_step=epoch)
        val_auc, val_ap = test(val_loader, model, device)
        valWriter.add_scalar(tag='Val AUC', scalar_value=val_auc, global_step=epoch)
        valWriter.add_scalar(tag='Val AP', scalar_value=val_ap, global_step=epoch)
        if val_auc > best_val_auc:
            best_val_auc = val_auc
            test_auc, test_ap = test(test_loader, model, device)
            # saving model parameters
            state = {'model': model.state_dict(), 'auc': test_auc, 'ap': test_ap, 'epoch': epoch}
            save_path = '../checkpoint/SEAL/'
            if not osp.exists(save_path):
                os.mkdir(save_path)
            torch.save(state, osp.join(save_path, args.dataset+'-'+args.split_ratio+'-'+args.embedding+'-'+'ckpt.pth'))

        print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}, Val_AUC: {val_auc:.4f}, Val_AP: {val_ap:.4f},'
              f'Test_AUC: {test_auc:.4f}, Test_AP: {test_ap:.4f}')

if __name__ == "__main__":
    run()


<filename>python/dataset/common_voice_loader.py
"""Load the Mozilla Common Voice dataset."""

import sys
import os
import csv
import subprocess

from multiprocessing import Pool, Lock, cpu_count
from tqdm import tqdm
from scipy.io import wavfile

from python.params import MIN_EXAMPLE_LENGTH, MAX_EXAMPLE_LENGTH
from python.dataset.config import CACHE_DIR, CORPUS_DIR
from python.util.storage import delete_file_if_exists
from python.dataset import download
from python.dataset.txt_files import generate_txt

# Path to the Mozilla Common Voice dataset.
__URL = 'https://common-voice-data-download.s3.amazonaws.com/cv_corpus_v1.tar.gz'
__MD5 = 'f1007e78cf91ab76b7cd3f1e8f554110'
__NAME = 'commonvoice'
__FOLDER_NAME = 'cv_corpus_v1'
__SOURCE_PATH = os.path.join(CACHE_DIR, __FOLDER_NAME)
__TARGET_PATH = os.path.realpath(os.path.join(CORPUS_DIR, __FOLDER_NAME))

# Define valid accents.
__VALID_ACCENTS = ['us',
                   'england',
                   'canada',
                   'australia',
                   'wales',
                   'newzealand',
                   'ireland',
                   'scotland',
                   'wales',
                   '']

def common_voice_loader(keep_archive):
    """Download and extract the common voice archive. Then build the output strings that can be
    written to the desired TXT files.

    Uses only the valid datasets, additional constraints are:
    * Downvotes must be at maximum 1/4 of upvotes.
    * Valid accents are: 'us', 'england', 'canada', 'australia'.
    * Accepting samples with only 1 upvote at the moment.

    Args:
        keep_archive (bool): Keep or delete the downloaded archive afterwards.

    Returns:
        Tuple[str]: Tuple containing the output strings that can be written to TXT files.
    """

    # Download and extract the dataset if necessary.
    download.maybe_download(__URL, md5=__MD5, cache_archive=keep_archive)
    if not os.path.isdir(__SOURCE_PATH):
        raise ValueError('"{}" is not a directory.'.format(__SOURCE_PATH))

    # Folders for each target.
    targets = [
        {
            'name': 'train',
            'folders': ['cv-valid-train']
        }, {
            'name': 'test',
            'folders': ['cv-valid-test']
        }, {
            'name': 'dev',
            'folders': ['cv-valid-dev']
        }
    ]

    txt_paths = []
    for target in targets:
        # Generate the WAV and a string for the `<target>.txt` file.
        output = __common_voice_loader(target['folders'])
        # Generate the `<target>.txt` file.
        txt_paths.append(generate_txt(__NAME, target['name'], output))

    # Cleanup extracted folder.
    download.cleanup_cache(__FOLDER_NAME)

    return tuple(txt_paths)

def __common_voice_loader(folders):
    """Build the output string that can be written to the desired TXT file.

    Uses only the valid datasets, additional constraints are:
    * Downvotes must be at maximum 1/4 of upvotes.
    * Valid accents are: 'us', 'england', 'canada', 'australia'.
    * Accepting samples with only 1 upvote at the moment.

    Args:
        folders (List(str)): A list containing folder names, e.g. `['train-valid', 'train-other']`.

    Returns:
        List[str]: List containing the output string that can be written to *.txt file.
    """

    output = []
    for folder in tqdm(folders, desc='Converting Common Voice data', total=len(folders),
                       file=sys.stdout, unit='CSVs', dynamic_ncols=True):
        # Open .csv file.
        with open('{}.csv'.format(os.path.join(__SOURCE_PATH, folder)), 'r') as csv_file:
            csv_reader = csv.reader(csv_file, delimiter=',')
            csv_lines = list(csv_reader)
            # print('csv_header:', csv_lines[0])
            # filename,text,up_votes,down_votes,age,gender,accent,duration

            lock = Lock()
            with Pool(processes=cpu_count()) as pool:
                # Create target folder if necessary.
                target_directory = os.path.join(__TARGET_PATH, folder)
                if not os.path.exists(target_directory):
                    os.makedirs(target_directory)

                # First line contains header.
                for result in pool.imap_unordered(__common_voice_loader_helper,
                                                  csv_lines[1:], chunksize=1):
                    if result is not None:
                        lock.acquire()
                        output.append(result)
                        lock.release()

    return output

def __common_voice_loader_helper(line):
    # Helper method for thread pool.

    # Cleanup label text.
    text = line[1].strip().replace('  ', ' ')
    # Enforce min label length.
    if len(text) > 1:
        # Check upvotes vs downvotes.
        if int(line[2]) >= 1 and int(line[3]) / int(line[2]) <= 1 / 4:
            # Check if speaker accent is valid.
            if line[6] in __VALID_ACCENTS:
                mp3_path = os.path.join(__SOURCE_PATH, line[0])
                assert os.path.isfile(mp3_path)
                wav_path = os.path.relpath('{}.wav'.format(mp3_path[:-4]), __SOURCE_PATH)
                wav_path = os.path.join(__TARGET_PATH, wav_path)

                delete_file_if_exists(wav_path)
                # Convert MP3 to WAV, reduce volume to 0.95, downsample to 16kHz and mono sound.
                subprocess.call(['sox', '-v', '0.95', mp3_path, '-r', '16k', wav_path,
                                 'remix', '1'])
                assert os.path.isfile(wav_path)

                # Validate that the example length is within boundaries.
                (sr, y) = wavfile.read(wav_path)
                length_sec = len(y) / sr
                if not MIN_EXAMPLE_LENGTH <= length_sec <= MAX_EXAMPLE_LENGTH:
                    return None

                # Add dataset relative to dataset path, label to TXT file buffer.
                wav_path = os.path.relpath(wav_path, CORPUS_DIR)
                return '{} {}\n'.format(wav_path, text)

    return None

# Test download script.
if __name__ == '__main__':
    print('Common Voice txt_paths: ', common_voice_loader(True))
    print('\nDone.')


import scipy.sparse as sp
import numpy as np
from typing import List

def create_adj_matrix(a_item2b_item: dict,
                      ordered_list_a_item: list,
                      ordered_list_b_item: list) -> sp.csr_matrix:
    """

    Parameters
    ----------
    a_item2b_item: dict
        dictionary from a_item to set of corresponding b_items
    ordered_list_a_item: list
        list with required a_items (unique)
    ordered_list_b_item: list
        list with required b_items (unique)

    Returns
    -------
    sp.csc_matrix
        adjacency matrix for a_items and b_items.
        adjacency matrix[i, j] = 1 <=> b_item in a_item2b_item[a_item]
        (a_item = ordered_list_a_item[i], b_item = ordered_list_b_item[i])

    """
    num_a_item = len(ordered_list_a_item)
    num_b_item = len(ordered_list_b_item)
    b_item2idx = dict(zip(ordered_list_b_item, range(num_b_item)))
    rows = []  # Non zero row indexes in adj_matrix
    cols = []  # Non zero column indexes in adj_matrix
    for i, a_item in enumerate(ordered_list_a_item):
        correspond_b_items = a_item2b_item[a_item].intersection(
            set(ordered_list_b_item))
        for b_item in correspond_b_items:
            rows.append(i)
            cols.append(b_item2idx[b_item])

    data = np.ones_like(rows)
    adj_matrix = sp.csr_matrix((data, (rows, cols)),
                               shape=(num_a_item, num_b_item))
    return adj_matrix

def create_combo_adj(combo_a_item2b_item: dict,
                     combo_a_item2a_item: dict,
                     ordered_list_a_item: list,
                     ordered_list_b_item: list) -> List[sp.csr_matrix]:
    """

    Parameters
    ----------
    combo_a_item2b_item: dict
         dict from combo a_items name to set of corresponding b_items
         (e.g. drug-drug pair name -> names of corresponding side effects)
    combo_a_item2a_item: dict
        dict from combo a_items name to sequence of two individual a_item names
        (e.g. combo drug-drug pair name -> np.array of two individual names)
    ordered_list_a_item: list
        list with required a_items (unique)
        (e.g. ordered list of drugs)
    ordered_list_b_item: list
        list with required b_items (unique)
        (e.g. ordered list of combo se)

    Returns
    -------
    List[sp.csr_matrix]
        list of adjacency matrices for a_items.
        For each b_item this matrix creates separately.
        (e.g. drug-drug adjacency matrix for each side effect.
        se adjacency matrix element is nonzero for two drugs <=>
        <=> corresponding two drugs have this se)

    """
    num_a_item = len(ordered_list_a_item)
    num_b_item = len(ordered_list_b_item)
    a_item2idx = dict(zip(ordered_list_a_item, range(num_a_item)))
    b_item2idx = dict(zip(ordered_list_b_item, range(num_b_item)))

    # rows[i], cols[i] -- lists of non zero indexes in adj_matrix of ith b_item
    rows = [[] for _ in range(num_b_item)]
    cols = [[] for _ in range(num_b_item)]
    for combo_a_item, set_b_item in combo_a_item2b_item.items():
        a_item1, a_item2 = combo_a_item2a_item[combo_a_item]
        if a_item1 not in ordered_list_a_item:
            continue
        if a_item2 not in ordered_list_a_item:
            continue
        idx_a_item1 = a_item2idx[a_item1]
        idx_a_item2 = a_item2idx[a_item2]
        correspond_b_items = combo_a_item2b_item[combo_a_item].intersection(
            set(ordered_list_b_item))
        for b_item in correspond_b_items:
            idx_b_item = b_item2idx[b_item]

            rows[idx_b_item].append(idx_a_item1)
            cols[idx_b_item].append(idx_a_item2)

            rows[idx_b_item].append(idx_a_item2)
            cols[idx_b_item].append(idx_a_item1)
    drug_drug_adj_list = [sp.csr_matrix((np.ones_like(rs), (rs, cs)),
                                        shape=(num_a_item, num_a_item))
                          for rs, cs in zip(rows, cols)]
    return drug_drug_adj_list


<filename>tayloralgebra.py
import operator
import numpy as np

from scipy import special
from scipy.special import factorial

class TaylorExpansion(object):

    def __init__(self, coefficients, center=0, cut_off=10):
        assert isinstance(coefficients, list) and all([isinstance(a, (float, int)) for a in coefficients])
        self.coefficients = [float(a) for a in coefficients[:cut_off]]
        self.center = center
        self.cut_off = cut_off

    @property
    def degree(self):
        return len(self) - 1

    def __str__(self):
        return self.coefficients.__str__()

    def __call__(self, x):
        if isinstance(x, TaylorExpansion):
            return self.compose(x)
        else:
            return sum([a * (x - self.center)**k for k, a in enumerate(self.coefficients)])

    def __len__(self):
        return len(self.coefficients)

    def __iter__(self):
        return iter(self.coefficients)

    def __getitem__(self, index):
        return self.coefficients[index]

    def __add__(self, other):
        if isinstance(other, TaylorExpansion):
            #assert self.center == other.center, "You cannot sum Taylor series defined at different points (centers)"
            cut_off = self._get_cut_off(other)
            longer, shorter = self._sort(other)
            return TaylorExpansion([a + shorter[k] if k < len(shorter) else a
                                    for k, a in enumerate(longer)], center=self.center, cut_off=cut_off)
        elif isinstance(other, (int, float)):
            return TaylorExpansion([a + other if k == 0 else a
                                    for k, a in enumerate(self)], center=self.center, cut_off=self.cut_off)
        else:
            raise ValueError

    def __neg__(self):
        return -1 * self

    def __sub__(self, other):
        return self + -other

    def __radd__(self, other):
        return self + other

    def __mul__(self, other):
        if isinstance(other, (float, int)):
            cut_off = self.cut_off
            new_coefficients = [other * a for a in self.coefficients]
        elif isinstance(other, (TaylorExpansion)):
            #assert self.center == other.center, "You cannot multiply Taylor series defined at different points (centers)"
            new_coefficients = []
            longer, shorter = self._sort(other)
            cut_off = self._get_cut_off(other)
            for n in range(min(len(self) + len(other) - 1, cut_off)):
                nth_coeff = sum([longer[k] * shorter[n - k] for k in range(len(longer))
                                 if len(shorter) > (n - k) >= 0])
                new_coefficients.append(nth_coeff)
        else:
            raise ValueError
        return TaylorExpansion(new_coefficients, center=self.center, cut_off=cut_off)

    def __rmul__(self, other):
        return self * other

    def __truediv__(self, other):
        cut_off = self._get_cut_off(other)
        inv = MultiplicativeInverse(other[0], cut_off)
        return self * inv(other - other[0])

    def compose(self, other):
        assert self.center == other.coefficients[0], "Error during composition: The constant coefficient of the internal function should be equal to the center of the external function"
        cut_off = self._get_cut_off(other)
        power = (other - self.center)
        result = self[0]
        for k in range(1, min(len(self), cut_off)):
            result += self[k] * power
            power = power * (other - self.center)
        return result if isinstance(result, TaylorExpansion) else TaylorExpansion(coefficients=[result],
                                                                                center=self.center,
                                                                                cut_off=self.cut_off)

    def _sort(self, other):
        longer = self if len(self) > len(other) else other
        shorter = other if self is longer else self
        return longer, shorter

    def _get_cut_off(self, other):
        return min(self.cut_off, other.cut_off)

#class MultiplicativeInverse(TaylorExpansion):
    #
    # def __init__(self, a, cut_off):
    #     if a == 0:
    #         raise ZeroDivisionError
    #     coefficients = [(-a) ** (-n) / a for n in range(cut_off)]
    #     super().__init__(coefficients, cut_off)

def compose_operator(left, right):
    return left(right)

class SymbolicFunction(object):

    def get_expansion(self, center, cut_off):
        pass #Abstract

    def evaluate(self, x):
        pass #abstract

    def _apply_operator(self, other, op):
        operation_list = [op, self, other]
        return CompositeFunction(operation_list)

    def __call__(self, x):
        if isinstance(x, SymbolicFunction):
            return self.compose(x)
        else:
            return self.evaluate(x)

    def compose(self, other):
        return self._apply_operator(other, compose_operator)

    def __neg__(self):
        return -1*self

    def __add__(self, other):
        return self._apply_operator(other, operator.add)

    def __radd__(self, other):
        return self.__add__(other)

    def __sub__(self, other):
        return self._apply_operator(other, operator.sub)

    def __rsub__(self, other):
        return -1*self.__sub__(other)

    def __mul__(self, other):
        return self._apply_operator(other, operator.mul)

    def __rmul__(self, other):
        return self.__mul__(other)

    def __truediv__(self, other):
        if isinstance(other, SymbolicFunction):
            return self*(MultiplicativeInverse()(other))
        else:
            return self*(1/other)

    def __rtruediv__(self, other):
        return self.__truediv__(other)**(-1)

    def __pow__(self, other):
        if isinstance(other, int):
            return PowerFunction(other)(self)
        else:
            return ExponentialFunction()(other*LogarithmFunction()(self))

    def __rpow__(self, other):
        raise NotImplementedError

class PrimitiveFunction(SymbolicFunction):

    def __init__(self, coeff_generator, fn, inv_fn=None):
        self.coeff_generator = coeff_generator
        self.fn = fn
        self.inv_fn = inv_fn
        self.taylor_expansion = {}

    def evaluate(self, x):
        return self.fn(x)

    # def invert(self):
    #     assert self.inv_fn is not None, "The function cannot be inverted"
    #     inv_derivative_function = MultiplicativeInverse()(self.derivative())
    #
    #     def inv_coeff(c, n):
    #         if n == 0:
    #             return self.inv_fn(c)
    #         elif n == 1:
    #             return 1 / self.coeff_generator(self.inv_fn(c), n)
    #         else:
    #             return inv_derivative_function.get_expansion(self.inv_fn(c), n)[n-1]
    #
    #     return PrimitiveFunction(inv_coeff, self.inv_fn, self.fn)

    def get_expansion(self, center, cut_off):
        if (center, cut_off) not in self.taylor_expansion:
            coefficients = [self.coeff_generator(center, n)
                            for n in range(cut_off)]
            self.taylor_expansion.update({(center, cut_off): TaylorExpansion(coefficients, center, cut_off)})
        return self.taylor_expansion[(center, cut_off)]

    def derivative(self):
        pass #abstract

class ExponentialFunction(PrimitiveFunction):

    def __init__(self):
        coeff_generator = lambda c, n: np.exp(c)/factorial(n)
        fn = lambda x: np.exp(x)
        inv_fn = lambda x: np.log(x)
        super().__init__(coeff_generator, fn, inv_fn)

    def derivative(self):
        return ExponentialFunction()

class LogarithmFunction(PrimitiveFunction):

    def __init__(self):
        coeff_generator = lambda c, n: (-1)**(n+1)*c**(-n)/n if n>0 else np.log(c)
        fn = lambda x: np.log(x)
        inv_fn = lambda x: np.exp(x)
        super().__init__(coeff_generator, fn, inv_fn)

    def derivative(self):
        return MultiplicativeInverse()

class SineFunction(PrimitiveFunction):

    def __init__(self):
        coeff_generator = lambda c, n: np.sin(n*np.pi/2. + c)/factorial(n)
        fn = lambda x: np.sin(x)
        inv_fn = lambda x: np.arcsin(x)
        super().__init__(coeff_generator, fn, inv_fn)

    def derivative(self):
        return CosineFunction()

class CosineFunction(PrimitiveFunction):

    def __init__(self):
        coeff_generator = lambda c, n: np.cos(n*np.pi/2. + c)/factorial(n)
        fn = lambda x: np.cos(x)
        inv_fn = lambda x: np.arccos(x)
        super().__init__(coeff_generator, fn, inv_fn)

    def derivative(self):
        return SineFunction()

class LogGammaFunction(PrimitiveFunction):

    def __init__(self):
        fn = lambda x: np.log(special.gamma(x))
        inv_fn = None

        def coeff_generator(c, n):
            if n == 0:
                return fn(c)
            else:
                return special.polygamma(n-1, c)/factorial(n)

        super().__init__(coeff_generator, fn, inv_fn)

    def derivative(self):
        return SineFunction()

class MultiplicativeInverse(PrimitiveFunction):

    def __init__(self):
        coeff_generator = lambda c, n: (-1)**n*c**(-(n+1))
        fn = lambda x: 1/x
        inv_fn = lambda x: 1/x
        super().__init__(coeff_generator, fn, inv_fn)

    def derivative(self):
        return -PowerFunction(2)(MultiplicativeInverse())

class PowerFunction(PrimitiveFunction):

    def __init__(self, k):
        def decreasing_factorial(k, n):
            if n < 0:
                return 0
            elif n == 0:
                return 1
            elif n == 1:
                return k
            else:
                return k*decreasing_factorial(k-1, n-1)
        coeff_generator = lambda c, n: (decreasing_factorial(k, n)/factorial(n))*c**(k-n) if n <= k else 0.
        fn = lambda x: x**k
        inv_fn = lambda x: x**(1/k)
        self.k = k
        super().__init__(coeff_generator, fn, inv_fn)

    def derivative(self):
        return self.k*PowerFunction(self.k - 1)

class CompositeFunction(SymbolicFunction):

    def __init__(self, operation_list):
        self.operation_list = operation_list

    def get_expansion(self, center, cut_off):
        op, left, right = self.operation_list
        if op is compose_operator:
            left_center = right.get_expansion(center, cut_off)[0]
        else:
            left_center = center
        return op(left.get_expansion(left_center, cut_off) if isinstance(left, SymbolicFunction) else left,
                  right.get_expansion(center, cut_off) if isinstance(right, SymbolicFunction) else right)

    def evaluate(self, x):
        op, left, right = self.operation_list
        if op is compose_operator:
            return left.evaluate(right.evaluate(x))
        else:
            return op(left.evaluate(x) if isinstance(left, SymbolicFunction) else left,
                      right.evaluate(x) if isinstance(right, SymbolicFunction) else right)


<filename>suncasa/utils/plot_mapX.py
# import os
# import matplotlib as mpl
import matplotlib.pyplot as plt
# import matplotlib.colorbar as colorbar
# import sunpy.cm.cm as cm  ## to bootstrap sdoaia color map
# import matplotlib.cm as cm
# import matplotlib.colors as colors
import astropy.units as u
# from astropy.io import fits
# import matplotlib.dates as mdates
# from astropy.time import Time
from sunpy import map as smap
# import matplotlib.gridspec as gridspec
# import numpy.ma as ma
# import matplotlib.patches as patches
# from suncasa.utils import stackplot as stp
# from IPython import embed
# from astropy.coordinates import SkyCoord
import numpy as np
from suncasa.utils import DButil
import warnings
import matplotlib.patches as patches
import numpy.ma as ma

class Sunmap():

    def __init__(self, sunmap, aia=False):
        if aia:
            try:
                sunmap = DButil.normalize_aiamap(sunmap)
            except:
                pass
            data = sunmap.data
            data[data < 1.0] = 1.0
            self.sunmap = smap.Map(data, sunmap.meta)
        else:
            self.sunmap = sunmap

        try:
            top_right_coord = self.sunmap.top_right_coord
            bottom_left_coord = self.sunmap.bottom_left_coord
            self.xrange = np.array([bottom_left_coord.Tx.to(u.arcsec).value, top_right_coord.Tx.to(u.arcsec).value])*u.arcsec
            self.yrange = np.array([bottom_left_coord.Ty.to(u.arcsec).value, top_right_coord.Ty.to(u.arcsec).value])*u.arcsec
        except:
            self.xrange = self.sunmap.xrange
            self.yrange = self.sunmap.yrange

    def map2wcsgrids(self, sunpymap=None, cell=False):
        '''

        :param sunpymap:
        :param cell: if True, return the coordinates of the pixel centers. if False, return the coordinates of the pixel boundaries
        :return:
        '''
        # embed()
        import astropy.units as u
        if sunpymap is None:
            sunpymap = self.sunmap
        ny, nx = sunpymap.data.shape
        x0, x1 = self.xrange.to(u.arcsec).value
        y0, y1 = self.yrange.to(u.arcsec).value
        dx = sunpymap.scale.axis1.to(u.arcsec / u.pix).value
        dy = sunpymap.scale.axis2.to(u.arcsec / u.pix).value

        if cell:
            mapx, mapy = np.linspace(x0, x1, nx) - dx / 2.0, np.linspace(y0, y1, ny) - dy / 2.0
            mapx = np.tile(mapx, ny).reshape(ny, nx)
            mapy = np.tile(mapy, nx).reshape(nx, ny).transpose()
        else:
            nx += 1
            ny += 1
            mapx, mapy = np.linspace(x0 - dx, x1 + dx, nx), np.linspace(y0 - dy, y1 + dy, ny)
            mapx = np.tile(mapx, ny).reshape(ny, nx)
            mapy = np.tile(mapy, nx).reshape(nx, ny).transpose()
        return mapx, mapy

    def get_map_extent(self, sunpymap=None, rot=0, rangereverse=False):
        if sunpymap is None:
            sunpymap = self.sunmap
        rot = rot % 360
        if rot == 90:
            extent = np.array(
                self.yrange.to(u.arcsec).value[::-1].tolist() + self.xrange.to(u.arcsec).value.tolist())
            extent = extent - np.array([sunpymap.scale.axis2.value] * 2 + [sunpymap.scale.axis1.value] * 2) / 2.0
        elif rot == 180:
            extent = np.array(
                self.xrange.to(u.arcsec).value[::-1].tolist() + self.yrange.to(u.arcsec).value[::-1].tolist())
            extent = extent - np.array([sunpymap.scale.axis1.value] * 2 + [sunpymap.scale.axis2.value] * 2) / 2.0
        elif rot == 270:
            extent = np.array(
                self.yrange.to(u.arcsec).value.tolist() + self.xrange.to(u.arcsec).value[::-1].tolist())
            extent = extent - np.array([sunpymap.scale.axis1.value] * 2 + [sunpymap.scale.axis2.value] * 2) / 2.0
        else:
            extent = np.array(self.xrange.to(u.arcsec).value.tolist() + self.yrange.to(u.arcsec).value.tolist())
            extent = extent - np.array([sunpymap.scale.axis1.value] * 2 + [sunpymap.scale.axis2.value] * 2) / 2.0
        if rangereverse:
            x0, x1, y0, y1 = extent
            extent = -np.array([x1, x0, y1, y0])
        return extent

    def imshow(self, axes=None, rot=0, rangereverse=False, maskon=False, image_enhance=False, **kwargs):
        '''
        :param sunpymap:
        :param axes:
        :param rot: rotation angle in degrees counter-clockwise. Must be an integer multiple of 90.
        :param kwargs:
        :return:
        '''
        sunpymap = self.sunmap
        if axes is None:
            axes = plt.subplot()
        rot = rot % 360
        if rot == 0:
            imdata = sunpymap.data
        elif rot == 90:
            imdata = sunpymap.data.transpose()[:, ::-1]
        elif rot == 180:
            imdata = sunpymap.data[::-1, ::-1]
        elif rot == 270:
            imdata = sunpymap.data.transpose()[::-1, :]
        else:
            warnings.warn('rot must be an integer multiple of 90. rot not implemented!')
            imdata = sunpymap.data
            rot = 0
        extent = self.get_map_extent(rot=rot, rangereverse=rangereverse)

        if maskon:
            if isinstance(maskon, bool):
                imdataplt = ma.masked_invalid(imdata)
                immask = imdataplt.mask
            elif isinstance(maskon, dict):
                if 'masked_equal' in maskon.keys():
                    imdataplt = ma.masked_equal(imdata, maskon['masked_equal'])
                elif 'masked_greater' in maskon.keys():
                    imdataplt = ma.masked_greater(imdata, maskon['masked_greater'])
                elif 'masked_less' in maskon.keys():
                    imdataplt = ma.masked_less(imdata, maskon['masked_less'])
                elif 'masked_greater_equal' in maskon.keys():
                    imdataplt = ma.masked_greater_equal(imdata, maskon['masked_greater_equal'])
                elif 'masked_less_equal' in maskon.keys():
                    imdataplt = ma.masked_less_equal(imdata, maskon['masked_less_equal'])
                elif 'masked_outside' in maskon.keys():
                    v1, v2 = maskon['masked_outside']
                    imdataplt = ma.masked_outside(imdata, v1, v2)
                elif 'masked_inside' in maskon.keys():
                    v1, v2 = maskon['masked_inside']
                    imdataplt = ma.masked_inside(imdata, v1, v2)
                elif 'masked_invalid' in maskon.keys():
                    imdataplt = ma.masked_invalid(imdata)
                else:
                    raise ValueError('maskon key wrong.')
                immask = imdataplt.mask
            else:
                raise TypeError('maskon must be bool or dict type.')
        else:
            imdataplt = imdata.copy()

        if image_enhance:
            dmax = np.nanmax(imdataplt)
            dmin = np.nanmin(imdataplt)
            from skimage.exposure import equalize_adapthist
            if isinstance(image_enhance, dict):
                imdataplt = equalize_adapthist(imdataplt, **image_enhance) * (dmax - dmin) + dmin
            else:
                imdataplt = equalize_adapthist(imdataplt) * (dmax - dmin) + dmin

        if maskon:
            imdataplt = ma.masked_array(imdataplt, immask)

        if isinstance(axes, list):
            ims = []
            for ax in axes:
                im = ax.imshow(imdataplt, extent=extent, origin='lower', **kwargs)
                ims.append(im)
                if rot == 0:
                    ax.set_xlabel('Solar X [arcsec]')
                    ax.set_ylabel('Solar Y [arcsec]')
                elif rot == 90:
                    ax.set_xlabel('-Solar Y [arcsec]')
                    ax.set_ylabel('Solar X [arcsec]')
                elif rot == 180:
                    ax.set_xlabel('-Solar X [arcsec]')
                    ax.set_ylabel('-Solar Y [arcsec]')
                elif rot == 270:
                    ax.set_xlabel('Solar Y [arcsec]')
                    ax.set_ylabel('-Solar X [arcsec]')
            return ims
        else:
            im = axes.imshow(imdataplt, extent=extent, origin='lower', **kwargs)

            if rot == 0:
                axes.set_xlabel('Solar X [arcsec]')
                axes.set_ylabel('Solar Y [arcsec]')
            elif rot == 90:
                axes.set_xlabel('-Solar Y [arcsec]')
                axes.set_ylabel('Solar X [arcsec]')
            elif rot == 180:
                axes.set_xlabel('-Solar X [arcsec]')
                axes.set_ylabel('-Solar Y [arcsec]')
            elif rot == 270:
                axes.set_xlabel('Solar Y [arcsec]')
                axes.set_ylabel('-Solar X [arcsec]')
            return im

    def contour(self, axes=None, rot=0, mapx=None, mapy=None, rangereverse=False, **kwargs):
        sunpymap = self.sunmap
        if axes is None:
            axes = plt.subplot()
        rot = rot % 360
        if (mapx is None) or (mapy is None):
            if rot == 0:
                mapx, mapy = self.map2wcsgrids(cell=True)
            elif rot == 90:
                mapy, mapx = self.map2wcsgrids(cell=True)
            elif rot == 180:
                mapx, mapy = self.map2wcsgrids(cell=True)
            elif rot == 270:
                mapy, mapx = self.map2wcsgrids(cell=True)

        if isinstance(axes, list):
            ims = []
            for ax in axes:
                im = ax.contour(mapx, mapy, sunpymap.data, **kwargs)
                ims.append(im)
                extent = self.get_map_extent(rot=rot, rangereverse=rangereverse)
                ax.set_xlim(extent[:2])
                ax.set_ylim(extent[2:])
            return ims
        else:
            im = axes.contour(mapx, mapy, sunpymap.data, **kwargs)
            extent = self.get_map_extent(rot=rot, rangereverse=rangereverse)
            axes.set_xlim(extent[:2])
            axes.set_ylim(extent[2:])
            return im

    def contourf(self, axes=None, rot=0, mapx=None, mapy=None, rangereverse=False, **kwargs):
        sunpymap = self.sunmap
        if axes is None:
            axes = plt.subplot()
        rot = rot % 360
        if (mapx is None) or (mapy is None):
            if rot == 0:
                mapx, mapy = self.map2wcsgrids(cell=True)
            elif rot == 90:
                mapy, mapx = self.map2wcsgrids(cell=True)
            elif rot == 180:
                mapx, mapy = self.map2wcsgrids(cell=True)
            elif rot == 270:
                mapy, mapx = self.map2wcsgrids(cell=True)

        if isinstance(axes, list):
            ims = []
            for ax in axes:
                im = ax.contourf(mapx, mapy, sunpymap.data, **kwargs)
                ims.append(im)
                extent = self.get_map_extent(rot=rot, rangereverse=rangereverse)
                ax.set_xlim(extent[:2])
                ax.set_ylim(extent[2:])
            return ims
        else:
            im = axes.contourf(mapx, mapy, sunpymap.data, **kwargs)
            extent = self.get_map_extent(rot=rot, rangereverse=rangereverse)
            axes.set_xlim(extent[:2])
            axes.set_ylim(extent[2:])
            return im

    def draw_limb(self, axes=None, rangereverse=False, **kwargs):
        if 'c' not in kwargs and 'color' not in kwargs:
            kwargs['c'] = 'w'
        if 'ls' not in kwargs and 'linestyle' not in kwargs:
            kwargs['ls'] = 'solid'
        sunpymap = self.sunmap
        if axes is None:
            axes = plt.gca()

        rsun = sunpymap.rsun_obs
        phi = np.linspace(-180, 180, num=181) * u.deg
        x = np.cos(phi) * rsun
        y = np.sin(phi) * rsun
        if isinstance(axes, list):
            ims = []
            for ax in axes:
                ax.set_autoscale_on(False)
                im = ax.plot(x, y, **kwargs)
                ims.append(im)
            return ims
        else:
            axes.set_autoscale_on(False)
            im = axes.plot(x, y, **kwargs)
            return im

    def draw_grid(self, axes=None, rot=0, grid_spacing=None, **kwargs):
        sunpymap = self.sunmap
        if grid_spacing is None:
            grid_spacing = 15 * u.deg

        def hgs2hcc(rsun, lon, lat, B0, L0):
            lon_L0 = lon - L0
            x = rsun * np.cos(lat) * np.sin(lon)
            y = rsun * (np.sin(lat) * np.cos(B0) - np.cos(lat) * np.cos(lon_L0) * np.sin(B0))
            z = rsun * (np.sin(lat) * np.sin(B0) + np.cos(lat) * np.cos(lon_L0) * np.cos(B0))
            return x, y, z

        def hcc2hpc(x, y, z, dsun):
            d = np.sqrt(x ** 2 + y ** 2 + (dsun - z) ** 2)
            Tx = np.arctan2(x, dsun - z)
            Ty = np.arcsin(y / d)
            return Tx, Ty

        if 'c' not in kwargs and 'color' not in kwargs:
            kwargs['c'] = 'w'
        if 'ls' not in kwargs and 'linestyle' not in kwargs:
            kwargs['ls'] = 'dotted'
        dsun = sunpymap.dsun
        rsun = sunpymap.rsun_meters
        if axes is None:
            axes = plt.gca()
        im = []
        b0 = sunpymap.heliographic_latitude.to(u.deg)
        l0 = sunpymap.heliographic_longitude.to(u.deg)
        hg_longitude_deg = np.linspace(-90, 90, num=91) * u.deg
        hg_latitude_deg = np.arange(0, 90, grid_spacing.to(u.deg).value)
        hg_latitude_deg = np.hstack([-hg_latitude_deg[1:][::-1], hg_latitude_deg]) * u.deg
        for lat in hg_latitude_deg:
            c = hgs2hcc(rsun, hg_longitude_deg, lat * np.ones(91), b0, l0)
            coords = hcc2hpc(c[0], c[1], c[2], dsun)
            if rot in [90, 270]:
                coords_ = [coords[1], coords[0]]
            else:
                coords_ = coords
            if isinstance(axes, list):
                for ax in axes:
                    im += ax.plot(coords_[0].to(u.arcsec), coords_[1].to(u.arcsec), **kwargs)
            else:
                im += axes.plot(coords_[0].to(u.arcsec), coords_[1].to(u.arcsec), **kwargs)

        hg_longitude_deg = np.arange(0, 90, grid_spacing.to(u.deg).value)
        hg_longitude_deg = np.hstack([-hg_longitude_deg[1:][::-1], hg_longitude_deg]) * u.deg
        hg_latitude_deg = np.linspace(-90, 90, num=91) * u.deg

        for lon in hg_longitude_deg:
            c = hgs2hcc(rsun, lon * np.ones(91), hg_latitude_deg, b0, l0)
            coords = hcc2hpc(c[0], c[1], c[2], dsun)
            if rot in [90, 270]:
                coords_ = [coords[1], coords[0]]
            else:
                coords_ = coords
            if isinstance(axes, list):
                for ax in axes:
                    im += ax.plot(coords_[0].to(u.arcsec), coords_[1].to(u.arcsec), **kwargs)
            else:
                im += axes.plot(coords_[0].to(u.arcsec), coords_[1].to(u.arcsec), **kwargs)
        return im

    def draw_rectangle(self, bottom_left, width, height, axes=None, **kwargs):
        if 'ec' not in kwargs and 'edgecolor' not in kwargs:
            kwargs['ec'] = 'w'
        if 'ls' not in kwargs and 'linestyle' not in kwargs:
            kwargs['ls'] = 'solid'
        if 'fill' not in kwargs:
            kwargs['fill'] = False
        if axes is None:
            axes = plt.gca()

        if isinstance(axes, list):
            ims = []
            for ax in axes:
                ax.set_autoscale_on(False)
                im = ax.add_patch(patches.Rectangle(bottom_left, width, height, **kwargs))
                ims.append(im)
            return ims
        else:
            axes.set_autoscale_on(False)
            im = axes.add_patch(patches.Rectangle(bottom_left, width, height, **kwargs))
            return im

    def imshow_RGB(self, maps, axes=None, returndataonly=False, rangereverse=False):
        from scipy import ndimage
        from astropy.coordinates import SkyCoord
        mapR = maps[0]
        znewR = mapR.data
        aiamapx, aiamapy = self.map2wcsgrids(sunpymap=mapR, cell=False)
        mapG = maps[1]
        XX, YY = mapG.data_to_pixel(SkyCoord(aiamapx * u.arcsec, aiamapy * u.arcsec, frame=mapG.coordinate_frame))
        znewG = ndimage.map_coordinates(mapG.data, [YY, XX], order=1)
        mapB = maps[2]
        XX, YY = mapB.data_to_pixel(SkyCoord(aiamapx * u.arcsec, aiamapy * u.arcsec, frame=mapB.coordinate_frame))
        znewB = ndimage.map_coordinates(mapB.data, [YY, XX], order=1)

        znewR = np.sqrt(znewR)
        znewG = np.sqrt(znewG)
        znewB = np.sqrt(znewB)

        vmax, vmin = np.sqrt(5000), np.sqrt(10)
        # clrange=DButil.sdo_aia_scale_dict(304)
        znewR[znewR > vmax] = vmax
        znewR[znewR < vmin] = vmin
        # clrange=DButil.sdo_aia_scale_dict(94)
        vmax, vmin = np.sqrt(20000), np.sqrt(200)
        znewG[znewG > vmax] = vmax
        znewG[znewG < vmin] = vmin
        # clrange=DButil.sdo_aia_scale_dict(211)
        vmax, vmin = np.sqrt(5000), np.sqrt(100)
        znewB[znewB > vmax] = vmax
        znewB[znewB < vmin] = vmin
        znewR = (znewR - np.nanmin(znewR)) / (np.nanmax(znewR) - np.nanmin(znewR))
        znewG = (znewG - np.nanmin(znewG)) / (np.nanmax(znewG) - np.nanmin(znewG))
        znewB = (znewB - np.nanmin(znewB)) / (np.nanmax(znewB) - np.nanmin(znewB))
        # znew1 = np.sqrt(znew1)
        # znew2 = np.sqrt(znew2)
        # imshow(np.sqrt(np.stack([znew0, znew1, znew2], axis=-1)), extent=list(aiamap.xrange.value) + list(aiamap.yrange.value),origin='lower')
        if returndataonly:
            return np.stack([znewR, znewG, znewB], axis=-1)
        else:
            if axes:
                pass
            else:
                axes = plt.subplot()
            extent = self.get_map_extent(sunpymap=mapR, rangereverse=rangereverse)
            if isinstance(axes, list):
                ims = []
                for ax in axes:
                    im = ax.imshow(np.stack([znewR, znewG, znewB], axis=-1),
                                   extent=extent, origin='lower')
                    ims.append(im)
                return ims
            else:
                return axes.imshow(np.stack([znewR, znewG, znewB], axis=-1),
                                   extent=extent, origin='lower')


<filename>tests.py
from scipy.io import wavfile

fps, bowl_sound = wavfile.read('voices/enathu_peyar.wav')
tones = range(-25,25)
transposed = [picthshift(bowl_sound,n) for n in tones]


<reponame>Comp-Sound-Music/interactive_piano
import unittest
import numpy as np
from scipy import fftpack
from piano import create_harmonies
from conversion_table import name_to_key as table

class CreateHarmonies(unittest.TestCase):
    def test_create_harmonies(self):
        # reference (with permission): https://cs510sound-spring2021.zulip.cs.pdx.edu/#narrow/stream/7-general/topic/HW.203.3A.20Chord.3F/near/2239
        samp_rt = 48000
        eps = 2.
        act_freqs = [
            523.25,
            554.37,
            587.33,
            622.25,
            659.26,
            698.46,
            739.99,
            783.99,
            830.61,
            880.0,
            932.33,
            987.77,
            1046.5
        ]
        for i,k in enumerate(table.keys()):
            v = table[k]
            ret = create_harmonies([k])
            x = fftpack.fft(ret)
            freqs = fftpack.fftfreq(len(x)) * samp_rt
            peaks = np.argsort(np.abs(x))[-6:]
            peak = round(np.sort(freqs[peaks])[-1],2)
            diff = round(np.abs(peak-act_freqs[i]),2)
            print(f"Test {k}  -- |{peak} - {act_freqs[i]}| = {diff}")
            self.assertTrue(diff <= eps)
        return

if __name__=="__main__":
    unittest.main()


<reponame>yosagaf/medical-biometrics<filename>assignement2/viewing3DBrainMRI.py
# import necessary packages

import imageio
import scipy.ndimage as ndi
import numpy as np
import SimpleITK as sitk
import matplotlib.pyplot as plt

# the path of a T1-weighted brain .nii image
path = "data/BRAIN.nii"

# read the .nii image containing the volume with the SimpleITK
sitk_f = sitk.ReadImage(path)

# access to the numpy array
slices = sitk.GetArrayFromImage(sitk_f)
print("[INFOS] 2D Array slice data type :", type(slices)) #-> numpy array
print("[INFOS] 3D sitk object type      :", type(sitk_f)) #-> numpy array
print("[INFOS] Shape of 3D image array  :", slices.shape)
print("[INFOS] Shape of 2D slice array  :", slices[0].shape)
print("[INFOS] Number of slices         :", slices.shape[0])

def removeKeymapConflicts(new_keys_set):
    for prop in plt.rcParams:
        if prop.startswith('keymap.'):
            keys = plt.rcParams[prop]
            remove_list = set(keys) & new_keys_set
            for key in remove_list:
                keys.remove(key)

def viewer3D(volume):
    removeKeymapConflicts({'n', 'l'})
    fig, ax = plt.subplots()
    ax.volume = volume
    ax.index = volume.shape[0] // 2
    ax.imshow(volume[ax.index], cmap='gray')
    fig.canvas.mpl_connect('key_press_event', processKey)

def processKey(event):
    fig = event.canvas.figure
    ax = fig.axes[0]
    if event.key == 'n':
        lSlice(ax)
    elif event.key == 'l':
        nSlice(ax)
    fig.canvas.draw()

def lSlice(ax):
    volume = ax.volume
    ax.index = (ax.index - 1) % volume.shape[0]  # wrap around using %
    ax.images[0].set_array(volume[ax.index])

def nSlice(ax):
    volume = ax.volume
    ax.index = (ax.index + 1) % volume.shape[0]
    ax.images[0].set_array(volume[ax.index])

mask3d_array = slices > 900
mask3d_array = ndi.binary_dilation(mask3d_array, iterations=8)
mask3d_array = ndi.binary_closing(mask3d_array, iterations=8)

weights_edge = [[[1, 1, 1],
                 [0, 0, 0],
                 [1, -1, -1]],

                [[1, 1, 1],
                 [0, 0, 0],
                 [-1, -1, -1]],

                [[1, 1, 1],
                 [0, 0, 0],
                 [-1, -1, -1]]]

im3d_edge = ndi.convolve(slices, weights_edge)

#viewer3D(mask_array)
#viewer3D(slices)
viewer3D(im3d_edge)
plt.show()


"""
Functions for testing brainlocate.
"""

import pytest
import numpy as np
import os
from numpy.testing import assert_array_equal, assert_allclose
import brainload as bl
import brainload.brainlocate as loc

THIS_DIR = os.path.dirname(os.path.abspath(__file__))
TEST_DATA_DIR = os.path.join(THIS_DIR, os.pardir, 'test_data')

# Respect the environment variable BRAINLOAD_TEST_DATA_DIR if it is set. If not, fall back to default.
TEST_DATA_DIR = os.getenv('BRAINLOAD_TEST_DATA_DIR', TEST_DATA_DIR)

def test_closest_vertex_to_very_close_point_known_dist():
    try:
        from scipy.spatial.distance import cdist
    except ImportError:
        pytest.skip("Optional dependency scipy not installed, skipping tests which require scipy.")
    vert_coords = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0]])
    faces = np.array([0, 1, 2])
    locator = loc.BrainLocate(vert_coords, faces)
    query_coords = np.array([[1.0, 0.0, 0.0], [1.0, 1.0, 0.0], [1.0, 1.1, 0.0], [0.1, 0.1, 0.1]])
    res = locator.get_closest_vertex(query_coords)
    assert res.shape == (4, )
    assert res[0] == 1     # the closest vertex index in the mesh for query coordinate at index 0
    assert res[1] == 2     # the closest vertex index in the mesh for query coordinate at index 1
    assert res[2] == 2     # the closest vertex index in the mesh for query coordinate at index 2
    assert res[3] == 0     # the closest vertex index in the mesh for query coordinate at index 3
    dist_matrix = cdist(vert_coords, query_coords)
    assert dist_matrix.shape == (3, 4)
    assert dist_matrix[0][0] == pytest.approx(1.0, 0.00001)
    assert dist_matrix[0][1] == pytest.approx(1.4142135623730951, 0.00001)   # This is sqrt(2)
    assert dist_matrix[1][0] == pytest.approx(0.0, 0.00001)
    min_index = np.argmin(dist_matrix, axis=0)
    assert min_index.shape == (4, )    # we queried for 4 coordinates.
    assert min_index[0] == res[0]
    assert min_index[1] == res[1]
    assert min_index[2] == res[2]
    assert min_index[3] == res[3]

def test_closest_vertex_to_very_close_point():
    try:
        from scipy.spatial.distance import cdist
    except ImportError:
        pytest.skip("Optional dependency scipy not installed, skipping tests which require scipy.")
    vert_coords, faces, _ = bl.subject_mesh('subject1', TEST_DATA_DIR, surf='white', hemi='both')
    locator = loc.BrainLocate(vert_coords, faces)
    query_coords = np.array([[58.0 , -45.0, 75.0]])
    res = locator.get_closest_vertex(query_coords)
    assert res.shape == (1, )
    assert res[0] == 210683     # the vertex index in the mesh
    expected_vert_210683_coords = (58.005173, -44.736935,  74.418076)
    assert_allclose(vert_coords[210683], np.array(expected_vert_210683_coords))
    dist = cdist(np.array([expected_vert_210683_coords]), query_coords)
    assert dist[0][0] == pytest.approx(0.6386434810831467, 0.001)

def test_closest_vertex_to_far_away_point():
    try:
        from scipy.spatial.distance import cdist
    except ImportError:
        pytest.skip("Optional dependency scipy not installed, skipping tests which require scipy.")
    vert_coords, faces, _ = bl.subject_mesh('subject1', TEST_DATA_DIR, surf='white', hemi='both')
    locator = loc.BrainLocate(vert_coords, faces)
    query_coords = np.array([[134.37332 , -57.259495, 149.267631], [134.37332 , -57.259495, 149.267631], [58.0 , -45.0, 75.0]])
    res = locator.get_closest_vertex(query_coords)
    assert res.shape == (3, )
    assert res[0] == 209519     # the vertex index in the mesh
    assert res[1] == 209519
    assert res[2] == 210683
    assert_allclose(vert_coords[209519], np.array((58.258751, -45.213722,  74.348068)))
    dist = cdist(np.array([[58.258751, -45.213722,  74.348068]]), np.array([[134.37332 , -57.259495, 149.267631]]))
    assert dist[0][0] == pytest.approx(107.47776133, 0.001)

def test_get_closest_vertex_and_distance_to_far_away_point():
    try:
        from scipy.spatial.distance import cdist
    except ImportError:
        pytest.skip("Optional dependency scipy not installed, skipping tests which require scipy.")
    vert_coords, faces, _ = bl.subject_mesh('subject1', TEST_DATA_DIR, surf='white', hemi='both')
    locator = loc.BrainLocate(vert_coords, faces)
    query_coords = np.array([[134.37332 , -57.259495, 149.267631], [134.37332 , -57.259495, 149.267631], [134.37332 , -57.259495, 149.267631]])  # just query 3 times for the same coord to see whether results and consistent
    res = locator.get_closest_vertex_and_distance(query_coords)
    assert res.shape == (3, 2)
    assert res[0,0] == 209519              # the vertex index in the mesh
    assert res[1,0] == 209519
    assert res[2,0] == 209519
    assert res[0,1] == 107.47776120258028  # the distance
    assert res[1,1] == 107.47776120258028
    assert res[2,1] == 107.47776120258028

def test_get_closest_vertex_to_vertex_0_coordinate():
    try:
        from scipy.spatial.distance import cdist
    except ImportError:
        pytest.skip("Optional dependency scipy not installed, skipping tests which require scipy.")
    vert_coords, faces, _ = bl.subject_mesh('subject1', TEST_DATA_DIR, surf='white', hemi='both')
    locator = loc.BrainLocate(vert_coords, faces)
    known_vertex_0_coord = (-1.85223234, -107.98274994, 22.76972961)    # coordinate for vertex 0 in the test data.
    assert_allclose(vert_coords[0], np.array(known_vertex_0_coord))
    query_coords = np.array([known_vertex_0_coord])
    res = locator.get_closest_vertex_and_distance(query_coords)
    assert res.shape == (1, 2)
    assert res[0, 0] == 0    # vertex index of closest vertex. The query coordinate is the known coordinate of vertex 0, so this must be 0.
    assert abs(res[0, 1]) <= 0.001   # distance to closest vertex. The query coordinate is the known coordinate of vertex 0, so this must be very close to 0.0