-
Notifications
You must be signed in to change notification settings - Fork 0
/
function_only.py
119 lines (79 loc) · 3.3 KB
/
function_only.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
import pickle
import pandas as pd
from sklearn.preprocessing import LabelEncoder
import xgboost as xgb
import datetime
import json
import optparse
def calculate_sum_time(config_file, source_location, target_location, first_ata, reverse=False):
config = json.load(open(config_file))
xgb_model_prime = pickle.load(open(config["prime_model"], "rb"))
xgb_model_second = pickle.load(open(config["second_model"], "rb"))
dd = pd.read_csv(config["loc_dataset"])
df = pd.read_csv(config["loc_predictions"])
locations_be = dd.reset_index().loc[dd[dd["location"] == source_location].index.min():dd[dd["location"] == target_location].index.min(), ["location"]]
locations_ne = dd.reset_index().loc[dd[dd["location"] == target_location].index.min():dd[dd["location"] == source_location].index.min(), ["location"]]
if locations_ne.shape[0] > 1:
locations = locations_ne[::-1]
else:
locations = locations_be
le = LabelEncoder()
le.fit(dd.loc[:, ["location"]].values.ravel())
locations_label = le.transform(locations.values.ravel())
atas = [first_ata]
atds = []
predicted_a = []
predicted_b = []
for loc in locations_label:
X_prime = pd.DataFrame.from_dict({"location": [loc], "actual_ta": [atas[-1]], "gbtt_ptd": [int(df[df["loc"] == le.inverse_transform([loc])[0]]["ptd"])]})
X_prime_pred = int(xgb_model_prime.predict(X_prime)[0])
ata_ls = list(str(atas[-1]))
if ata_ls[0:2] == ["2", "4"]:
ata_ls[0:2] = ["0", "0"]
ata_n = "".join(ata_ls)
ata_dt = datetime.datetime.strptime(ata_n, "%H%M")
dt_1 = ata_dt + datetime.timedelta(minutes=X_prime_pred)
atds.append(int(dt_1.strftime("%H%M")))
X_second = pd.DataFrame.from_dict({"location": [loc], "actual_td": [atds[-1]], "gbtt_pta": [int(df[df["loc"] == le.inverse_transform([loc])[0]]["pta"])]})
X_second_pred = int(xgb_model_second.predict(X_second)[0])
dt_2 = ata_dt + datetime.timedelta(minutes=X_second_pred)
if dt_2.strftime("%H%M")[:2] == "00":
ata_ls = list(dt_2.strftime("%H%M"))
ata_ls[:2] = ["2", "4"]
ata = int("".join(ata_ls))
else:
ata = int(dt_2.strftime("%H%M"))
atas.append(ata)
predicted_a.append(X_prime_pred), predicted_b.append(X_second_pred)
atas = atas[:-1]
locat_out = [l[0] for l in list(locations.values)]
if reverse:
atas.reverse()
atds.reverse()
predicted_a.reverse()
predicted_b.reverse()
locat_out.reverse()
output = {}
for locat, atd in zip(locat_out, atds):
output[locat] = atd
return output
if __name__ == "__main__":
parser = optparse.OptionParser()
parser.add_option('-c', '--config',
action="store", dest="config",
help="Config file", default="config.json")
parser.add_option('-s', '--source',
action="store", dest="source",
help="Source")
parser.add_option('-d', '--dest',
action="store", dest="dest",
help="Destination")
parser.add_option('-a', '--arrival',
action="store", dest="arrival",
help="Arrival time", type="int")
parser.add_option('-i', '--inverse',
const=True, dest="inverse", action="store_const",
help="Inverse", default=False)
options, args = parser.parse_args()
args = parser.parse_args()
print(calculate_sum_time(options.config, options.source, options.dest, options.arrival, reverse=options.inverse))