visualTools/pyplotters.py at main · gardner-lab/visualTools · 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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

# Matplotlib based plotters


import numpy as np
import matplotlib.pyplot as plt


def colorScaler(inputEmbedding,metric = 'time', colormap = 'hsv',customSeries = None):

	N = inputEmbedding.shape[0]

	if N < 4:
		print('Warning, small embedding dimension to plot')

	if metric == 'time':
		array = np.linspace(0, 1, N)
		cmap = plt.get_cmap(colormap)  # Get the desired colormap
		rgbArray = cmap(array)[:, :3]

	if metric == 'RGB':
		print('fail')


	return rgbArray


'''

Code to create a pixelated UMAP image
-Add advanced colorization methods, or write wrapper/function
-Generalize Label Text
'''
def pixelatedUMAP(inputEmbedding,imageSize = 500,mode = 'averaging',boundary = .5):

	# Obtain X and Y coordinates
	x = inputEmbedding[:,0]
	y = inputEmbedding[:,1]

	# Set X and Y ranges of output image
	range_min_x, range_max_x = np.min(x) - boundary, np.max(x) + boundary
	range_min_y, range_max_y = np.min(y) - boundary, np.max(y) + boundary

	# Reference function to colorize
	colors = colorScaler(inputEmbedding,'time')

	# Create an empty image grid
	image = np.zeros((imageSize, imageSize, 3))
	imageCount = np.zeros((imageSize, imageSize, 3))

	# Scale and shift the coordinates to match the image grid
	scaled_x = (x - range_min_x) / (range_max_x - range_min_x) * imageSize
	scaled_y = (y - range_min_y) / (range_max_y - range_min_y) * imageSize

	# Iterate over points and add colors to the image
	for xi, yi, color in zip(scaled_x, scaled_y, colors):
		xi, yi = int(xi), int(yi)
		image[xi - 1, yi - 1, :] += color  # Verify indexing
		imageCount[xi - 1, yi - 1, :] += 1

	if mode == 'averaging':
		image = image/imageCount

	plt.imshow(image, origin='lower', extent=[range_min_x, range_max_x, range_min_y, range_max_y])
	plt.xlabel('UMAP Embedding Dim1')
	plt.ylabel('UMAP Embedding Dim1')
	plt.title('UMAP Embedding -- Pixelated')

def threePlotter(inputEmbedding):

	minIndex = 0
	maxIndex = 5000

	manualTimeCorrection = 5/1000 #150 --> 200HZ and a step of one --> each slice moves 1/200 = 5ms
	time = np.arange(maxIndex - minIndex) * manualTimeCorrection # ADD TIME... better

	# spkToColors = rescale_array(spk_emb)


	# Generate some random data
	x = inputEmbedding[minIndex:maxIndex,0]
	y = inputEmbedding[minIndex:maxIndex,1]
	z = time
	#color = spkToColors[minIndex:maxIndex,:]

	# Create a 3D scatter plot
	fig = plt.figure(facecolor='white')
	ax = fig.add_subplot(111, projection='3d')
	scatter = ax.scatter(x, y, z, cmap='viridis') #, c=color

	# Add a colorbar
	#cbar = fig.colorbar(scatter)
	#cbar.set_label('Spike Embedding')

	# Set labels for each axis
	ax.set_xlabel('UMAP Embedding Dim1')
	ax.set_ylabel('UMAP Embedding Dim1')
	ax.set_zlabel('Time [S]')


if __name__ == '__main__':

	embeddingPath = '/Users/ethanmuchnik/Desktop/SuperData/results/B119/splitsA-standardAttemptA-UMAP.npz'
	embedding = np.load(embeddingPath)['spk_emb']
	pixelatedUMAP(embedding,imageSize = 500,mode = 'additive')
	plt.show()

	# embeddingPath = '/Users/ethanmuchnik/Desktop/SuperData/results/B119/splitsA-standardAttemptA-UMAP.npz'
	# embedding = np.load(embeddingPath)['spk_emb']
	# threePlotter(embedding)
	# plt.show()