add generate data set func

This commit is contained in:
matixezor 2021-05-16 19:41:09 +02:00
parent dc3086307b
commit c6e2f5b796
2 changed files with 60 additions and 0 deletions

View File

View File

@ -0,0 +1,60 @@
from itertools import product
visibility = ('bad', 'medium', 'good')
stability = ('unstable', 'stable')
ground = ('mud', 'grass', 'dry ground', 'concrete')
mine_type = ('AT', 'AP', 'ADM')
armed = (True, False)
pressure_gt_two = (True, False)
attributes = ['visibility', 'stability', 'ground', 'mine_type', 'armed', 'pressure_gt_two']
def generate_data_set():
data_list = list(product(visibility, stability, ground, mine_type, armed, pressure_gt_two))
data = [
{
'visibility': entry[0],
'stability': entry[1],
'ground': entry[2],
'mine_type': entry[3],
'armed': entry[4],
'pressure_gt_two': entry[5]
} for entry in data_list]
for index, item in enumerate(data):
if not item['armed']:
data[index] = {**item, 'action': 'defuse'}
else:
if item['stability'] == 'unstable':
data[index] = {**item, 'action': 'detonation'}
else:
if item['visibility'] == 'bad':
data[index] = {**item, 'action': 'detonation'}
elif item['visibility'] == 'medium':
if item['ground'] in ['mud', 'grass']:
data[index] = {**item, 'action': 'detonation'}
else:
if item['mine_type'] in ['AT', 'ADM']:
data[index] = {**item, 'action': 'defuse'}
else:
if item['pressure_gt_two']:
data[index] = {**item, 'action': 'defuse'}
else:
data[index] = {**item, 'action': 'detonation'}
else:
if item['ground'] == 'mud':
data[index] = {**item, 'action': 'detonation'}
else:
if item['mine_type'] in ['AT', 'ADM']:
data[index] = {**item, 'action': 'defuse'}
else:
if item['pressure_gt_two']:
data[index] = {**item, 'action': 'defuse'}
else:
data[index] = {**item, 'action': 'detonation'}
return data
data_set = generate_data_set()
training_set = data_set[:50] + data_set[55:90] + data_set[95:140] + data_set[145:200] + data_set[205:]
test_set = data_set[50:55] + data_set[90:95] + data_set[140:145]