SNF.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/8/8 14:01
# @Author  : Li Xiao
# @File    : SNF.py
import snf
import pandas as pd
import numpy as np
import argparse
import seaborn as sns

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--path', '-p', type=str, nargs=3, required=True,
                        help='Location of input files, must be 3 files')
    parser.add_argument('--metric', '-m', type=str, choices=['braycurtis', 'canberra', 'chebyshev', 'cityblock',
                        'correlation', 'cosine', 'dice', 'euclidean', 'hamming', 'jaccard', 'kulsinski',
                        'mahalanobis', 'matching', 'minkowski', 'rogerstanimoto', 'russellrao', 'seuclidean',
                        'sokalmichener', 'sokalsneath', 'sqeuclidean', 'wminkowski', 'yule'], default='sqeuclidean',
                        help='Distance metric to compute. Must be one of available metrics in :py:func scipy.spatial.distance.pdist.')
    parser.add_argument('--K', '-k', type=int, default=20,
                        help='(0, N) int, number of neighbors to consider when creating affinity matrix. See Notes of :py:func snf.compute.affinity_matrix for more details. Default: 20.')
    parser.add_argument('--mu', '-mu', type=int, default=0.5,
                        help='(0, 1) float, Normalization factor to scale similarity kernel when constructing affinity matrix. See Notes of :py:func snf.compute.affinity_matrix for more details. Default: 0.5.')
    args = parser.parse_args()

    print('Load data files...')
    omics_data_1 = pd.read_csv(args.path[0], header=0, index_col=None)
    omics_data_2 = pd.read_csv(args.path[1], header=0, index_col=None)
    omics_data_3 = pd.read_csv(args.path[2], header=0, index_col=None)
    print(omics_data_1.shape, omics_data_2.shape, omics_data_3.shape)

    if omics_data_1.shape[0] != omics_data_2.shape[0] or omics_data_1.shape[0] != omics_data_3.shape[0]:
        print('Input files must have same samples.')
        exit(1)

    omics_data_1.rename(columns={omics_data_1.columns.tolist()[0]: 'Sample'}, inplace=True)
    omics_data_2.rename(columns={omics_data_2.columns.tolist()[0]: 'Sample'}, inplace=True)
    omics_data_3.rename(columns={omics_data_3.columns.tolist()[0]: 'Sample'}, inplace=True)

    # align samples of different data
    omics_data_1.sort_values(by='Sample', ascending=True, inplace=True)
    omics_data_2.sort_values(by='Sample', ascending=True, inplace=True)
    omics_data_3.sort_values(by='Sample', ascending=True, inplace=True)

    print('Start similarity network fusion...')
    affinity_nets = snf.make_affinity([omics_data_1.iloc[:, 1:].values.astype(np.float), omics_data_2.iloc[:, 1:].values.astype(np.float), omics_data_3.iloc[:, 1:].values.astype(np.float)],
                                      metric=args.metric, K=args.K, mu=args.mu)

    fused_net =snf.snf(affinity_nets, K=args.K)

    print('Save fused adjacency matrix...')
    fused_df = pd.DataFrame(fused_net)
    fused_df.columns = omics_data_1['Sample'].tolist()
    fused_df.index = omics_data_1['Sample'].tolist()
    fused_df.to_csv('result/SNF_fused_matrix.csv', header=True, index=True)

    np.fill_diagonal(fused_df.values, 0)
    fig = sns.clustermap(fused_df.iloc[:, :], cmap='vlag', figsize=(8,8),)
    fig.savefig('result/SNF_fused_clustermap.png', dpi=300)
    print('Success! Results can be seen in result file')