host_predict/old_amr_split.py at master · superphy/host_predict · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
#!/usr/bin/env python

import numpy as np
import pandas as pd
from sklearn.externals import joblib
from sklearn.cross_validation import train_test_split
from sklearn.model_selection import StratifiedKFold
from math import floor

def decode_categories(data, class_dict):
	#print(class_dict)
	arry = np.array([])
	for item in data:
		arry = np.append(arry,class_dict[item])
	#print(arry)
	return arry


def encode_categories(data, class_dict):
	arry = np.array([], dtype = 'i4')
	for item in data:
		temp = int(float((item.decode('utf-8')).split("=")[-1]))
		# 0=1, 1=2, 2=4, 3=8, 4=16, 5=32
		# CHANGE ALL OF THIS CODE TO AN ENCODER
		for index in range(len(class_dict)):
			check = (int(float(class_dict[index].split("=")[-1])))
			#print(temp, check)
			if temp == check:
				temp = index
			#else:
				#print("AHHHHHHHHHHHHHHHHHHHHHHHHH")
		arry = np.append(arry,temp)
	print(arry)
	return arry


if __name__ == "__main__":
	# Matrix of experimental MIC values
	df = joblib.load("amr_data/mic_class_dataframe.pkl")

	# Matrix of classes for each drug
	mic_class_dict = joblib.load("amr_data/mic_class_order_dict.pkl")

	df_cols = df.columns # Col names are drugs
	for drug in df_cols:

		print("start: making train/test data for ", drug)

		num_classes = len(mic_class_dict[drug])

		matrix   = np.load('amr_data/'+drug+'/kmer_matrix.npy')
		rows_gen = np.load('amr_data/'+drug+'/kmer_rows_genomes.npy')
		rows_mic = np.load('amr_data/'+drug+'/kmer_rows_mic.npy')
		cols     = np.load('amr_data/'+drug+'/kmer_cols.npy')

		num_rows = len(rows_gen)

		## Create the training and testing data sets
		# Determine the size of the sets
		chunk = floor(num_rows/5)
		remainder = num_rows%5
		if remainder == 0 :
			train_size = chunk*4
			test_size  = chunk
		else:
			train_size = (chunk*4)+remainder
			test_size  = chunk

		# Create masks
		train_mask   = [1]*train_size
		train_mask_b = [0]*test_size
		train_mask   = train_mask + train_mask_b

		test_mask   = [0]*train_size
		test_mask_b = [1]*test_size
		test_mask   = test_mask + test_mask_b

		# Make data sets
		train_list = [bool(x) for x in train_mask]
		test_list  = [bool(x) for x in test_mask]

		train_data  = matrix[train_list, :]
		train_names = rows_mic[train_list]
		test_data   = matrix[test_list, :]
		test_names  = rows_mic[test_list]

		print(matrix.shape)
		print(train_data.shape)
		print(test_data.shape)

		np.save('amr_data/'+drug+'/train_data.npy', train_data)
		np.save('amr_data/'+drug+'/train_names.npy', train_names)
		np.save('amr_data/'+drug+'/test_data.npy', test_data)
		np.save('amr_data/'+drug+'/test_names.npy', test_names)

		print("end: making train/test data for ",drug)