-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathArena.py
232 lines (198 loc) · 9.16 KB
/
Arena.py
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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
###########################################
## Arena class source code - everything ##
## related to the state of the arena ##
## (same as field, same as architecture) ##
## lives in this class ##
## (c) Artem Pashchinskiy, UCLA, 2019 ##
###########################################
import numpy as np
#import nibabel as nib # NIFTI images processing
import scipy.misc
import imageio
import scipy.stats as stats
from math import log, sin, cos
# generate field .npy mask from a .gif black and white image
# (used to work with a .nii.gz segmentations before)
def generate(filename, targetfolder = None):
if not targetfolder:
targetfolder = 'Simulation resources/fields_processed'
ext = filename.split('.')[-1]
if ext == 'gif':
array2d = np.asarray(imageio.imread(filename)).transpose()
#avoid grey zones
array2d[array2d > 0] = 255
# LEGACY FROM .nii.gz arenas processing
# elif ext == 'gz':
# img = nib.load(filename)
# data = img.get_data()
# # get a 2d slice of the pseudo 3d NIFTI image
# array2d = np.squeeze(data)
# # make all selected pixels white
# array2d[array2d == 1] = 255;
# END OF LEGACY PIECE
# LEGACY 2: uncomment below to create a completely empty field with no walls
#array2d[array2d == 0] = 255;
#report results
print("Nest map sucesfully generated")
print("Area avaliable for agents: ", np.sum(array2d)/255)
#save result as a NumPy array file and as a .gif image
filename = filename.split('/')[-1]
newfilename = targetfolder + '/' + filename[:max(filename.find('_raw'), filename.find('.'))]
np.save(newfilename+'_Field.npy', array2d)
scipy.misc.imsave(newfilename+'_Field.gif', array2d.transpose())
# a translator function that converts abbreviated name of the field to the
# filename of the .npy mask generated by the "generate()" function as well
# as the coordinates of the heated (starting) nest center
# new entries shuld be added to the body of the function directly
def getTypeFile(type, fieldfolder = './../Simulation resources/fields_processed/'):
translator = {
#old_set, ORIGIN = (1100, 550)
'4o_2_4o_old' : ("4ch_2_4ch_Field.npy", (1100, 550)),
'3o_2_3o_old' : ("3ch_2_3ch_Field.npy", (1100, 550)),
'4c_2_4c_old' : ("4chcl_2_4chcl_Field.npy", (1100, 550)),
'3o_2_4o_old' : ("3ch_2_4ch_Field.npy", (1100, 550)),
'1sq_2_1sq' : ('1sq_2_1sq_Field.npy', (1100, 550)),
#new set, ORIGIN = (1260, 570)
'4o_2_4o' : ('4ch_2_4ch_fixed_Field.npy', (1260, 570)),
'3o_2_3o' : ('3ch_2_3ch_fixed_Field.npy', (1260, 570)),
'3o_2_4o' : ("3ch_2_4ch_fixed_Field.npy", (1260, 570)),
'4o_2_3o' : ("4ch_2_3ch_fixed_Field.npy", (1260, 570)),
'4o_2_3cl' : ("4ch_2_3ch_cl_Field.npy", (1260, 570)),
'3o_2_4cl' : ("3ch_2_4ch_cl_Field.npy", (1260, 570)),
#virtual set, ORIGIN = (1470, 540), area = 690,000
'c0d0' : ("c0d0_Field.npy",(1470, 540)),
'c1d0' : ("c1d0_Field.npy",(1470, 540)),
'c0d1' : ("c0d1_Field.npy",(1470, 540)),
'c1d1' : ("c1d1_Field.npy",(1470, 540)),
'2tun' : ("2tun_Field.npy",(1475, 540)),
'c0d1-m' : ("c0d1-m_Field.npy",(1470, 540)),
}
try:
return fieldfolder+translator[type][0], translator[type][1]
except KeyError:
return "FIELD NOT FOUND"
class Arena:
# constructor works with specific types of fields that are specified
# in the "translator function" getTypeFile()
def __init__(self, field):
self.type = field
self.arFile, self.origin = getTypeFile(self.type)
self.nestDim = self.loadField(self.arFile)[2:]
# essentially copies the consructor
# call to return the object to its default state
def reset(self, field = "c0d0"):
self.type = field
self.arFile, self.origin = getTypeFile(field)
self.nestDim = self.loadField(self.arFile)[2:]
def getNestDim(self):
return self.nestDim
# load the .npy files into the memory and compute its spatial properties
def loadField(self, filename, verbose = False):
self.field = np.load(filename)
self.dim = self.field.shape
self.dimX = self.dim[0] # x-dimension of the whole area
self.dimY = self.dim[1] # y-dimension of the whole area
self.minY = np.argwhere(sum(self.field) > 0)[0][0] # min y-coordinate of the avaliable nest
self.maxY = np.argwhere(sum(self.field) > 0)[-1][0] # max y-coordinate of the avaliable nest
self.minX = np.argwhere(sum(self.field.transpose()) > 0)[0][0] # min x-coordinate of the avaliable nest
self.maxX = np.argwhere(sum(self.field.transpose()) > 0)[-1][0] # max x-coordinate of the avaliable nest
if verbose:
print('x:', self.minX, self.maxX, 'y:', self.minY, self.maxY)
return self.dim, self.field, self.maxX-self.minX, self.maxY-self.minY
# look up the state of a certain location on the field
def getFieldVal(self, x_or_tuple, y = None, default = 0):
try:
if y == None:
return self.field[int(round(x_or_tuple[0])), int(round(x_or_tuple[1]))]
else:
return self.field[int(round(x_or_tuple)), int(round(y))]
except IndexError:
return default
# assign a state to a certain location on the field
def setFieldVal(self, val, x_or_tuple, y = None):
try:
if y == None:
self.field[int(round(x_or_tuple[0])), int(round(x_or_tuple[1]))] = val
else:
self.field[int(round(x_or_tuple)), int(round(y))] = val
except IndexError:
pass
def Xmin(self):
return self.minX
def Xmid(self):
return (self.maxX + self.minX) // 2
# generate a list of starting positions inside the cold nest based on the number of agents,
# properties of the loaded field, and a requested distribution type along x-axis
# (e for exponential, ls for left-skewed, rs for right skewed, u for uniform)
# returns a list of the positions (as tuples) and a PDF of the requested distriburtion
def getStartPos(self, pars, dist = 'e'):
try:
dist = pars["INIT_DISTR"]
except KeyError:
pass
shift = pars["SHIFT"]
mu, sigma = log(100), 0.9 # parameters for lognormal distr
positions = np.empty((2, pars['NUM']), dtype = int) # array with NUM columns; j-th in the end represents starting (x, y) of j-th ant
if dist == 'e':
lower = 0 # left bound of exponential x-distribution
upper = self.maxX - self.origin[0] - 5 # length of non-truncated region ofexponential x-distr
scale = pars["ESCALE"]
X = stats.truncexpon(b = (upper-lower)/scale, loc=lower, scale=scale)
positions[0] = X.rvs(pars['NUM']) + self.origin[0]
def pdf(t):
return 0;
elif dist == 'ls':
i = 0
while i < pars['NUM']:
s = self.origin[0] - shift + np.random.lognormal(mu, sigma)
while s > self.maxX - 5 or s < self.minX + 5:
s = self.origin[0] - shift + np.random.lognormal(mu, sigma)
positions[0][i] = int(s)
i += 1
def pdf(t):
# np.random.lognormal pdf taken from documenatation
arr = [0 if (x < self.origin[0] - shift) or (x > self.maxX - 5) else np.exp(-(np.log(x) - mu)**2 / (2 * sigma**2)) / (x * sigma * np.sqrt(2 * np.pi)) for x in t]
return arr
elif dist == 'rs':
i = 0
while i < pars['NUM']:
s = self.maxX - np.random.lognormal(mu, sigma)
while s > self.maxX - 5 or s < self.minX + 5:
s = self.maxX - np.random.lognormal(mu, sigma)
positions[0][i] = int(s)
i += 1
def pdf(t):
x = self.maxX - t
return np.exp(-(np.log(x) - mu)**2 / (2 * sigma**2)) / (x * sigma * np.sqrt(2 * np.pi))
elif dist == 'u':
positions[0] = self.origin[0] - shift + np.random.randint(self.maxX - self.origin[0] + shift - 5, size = pars['NUM'])
def pdf(t):
arr = [0 if (x < self.origin[0] - shift) or (x > self.maxX - 5) else 1/(self.maxX - self.origin[0] + shift - 5) for x in t]
return arr
positions[1] = [np.random.choice((np.nonzero(self.field[positions[0][i]])[0])) for i in range(pars['NUM'])]
return positions.transpose(), pdf #return as list of coordinate pairs and return the pdf function
# find all agents within a square (with side 2*interrad) around the given location
def getInter(self, x, y, interrad):
neighbours = np.unique(self.field[x - interrad : x + interrad + 1, y - interrad : y + interrad + 1])
setIntNeigh = {int(i) for i in neighbours}
return setIntNeigh
# fill the arena locations covered by the body area of an agent with the corresponsing antID
def fillbodyspace(self, antID, headX, headY, direction, length, width):
if length == 0 and width == 0:
return -1
yy, xx = np.mgrid[-width//2 : width//2+1, -length : 1]
rot = np.array(((np.cos(direction), -np.sin(direction)), (np.sin(direction), np.cos(direction))))
xxr, yyr = np.zeros_like(xx), np.zeros_like(yy)
for i, j, ir, jr in np.nditer([xx, yy, xxr, yyr], op_flags=[['readonly'], ['readonly'], ['readwrite'], ['readwrite']]):
ir[...], jr[...] = np.matmul(np.array([i, j]), rot)[0]+headX, np.matmul(np.array([i, j]), rot)[1]+headY
if self.getFieldVal(int(ir), int(jr)) != 0:
self.setFieldVal(antID, int(ir), int(jr))
return (xxr, yyr)
### Usage example for the field generation procedure
### note: before masks created by this procedure can be further used
### corresponding lines should be manually added to the getTypeFile() func
if __name__ == "__main__":
generate('Simulation resources/fields_raw/4ch_2_4ch_fixed.gif')
generate('Simulation resources/fields_raw/3ch_2_4ch_fixed.gif')
generate('Simulation resources/fields_raw/2tun.gif')
generate('Simulation resources/fields_raw/c0d1-m.gif')