def test_build_patterns(trips):
time_windows = get_time_windows(48, 3)
patterns = tdc.build_patterns(trips, time_windows)
patterns = patterns.sort_values(['tour_id', 'outbound', 'trip_num'])
assert patterns.shape[0] == 34
assert patterns.shape[1] == 6
assert patterns.index.name == tdc.TOUR_LEG_ID
output_columns = [
tdc.TOUR_ID, tdc.PATTERN_ID, tdc.TRIP_NUM, tdc.STOP_TIME_DURATION,
tdc.TOUR_ID, tdc.OUTBOUND
]
assert set(output_columns).issubset(patterns.columns)
def test_get_time_windows(duration, levels, expected):
time_windows = get_time_windows(duration, levels)
if levels == 1:
assert time_windows.ndim == 1
assert len(time_windows) == expected
assert np.sum(time_windows <= duration) == expected
else:
assert len(time_windows) == levels
assert len(time_windows[0]) == expected
total_duration = np.sum(time_windows, axis=0)
assert np.sum(total_duration <= duration) == expected
df = pd.DataFrame(np.transpose(time_windows))
assert len(df) == len(df.drop_duplicates())
def get_pattern_index_and_arrays(tour_indexes, durations, one_way=True):
"""
A helper method to quickly calculate all of the potential time windows
for a given set of tour indexes and durations.
:param tour_indexes: List of tour indexes
:param durations: List of tour durations
:param one_way: If True, calculate windows for only one tour leg. If False,
calculate tour windows for both legs
:return: np.array: Tour indexes repeated for valid pattern
np.array: array with a column for main tour leg, outbound leg, and inbound leg
np.array: array with the number of patterns for each tour
"""
max_columns = 2 if one_way else 3
max_duration = np.max(durations)
time_windows = get_time_windows(max_duration, max_columns)
patterns = []
pattern_sizes = []
for duration in durations:
possible_windows = time_windows[:max_columns,
np.where(
time_windows.sum(
axis=0) == duration)[0]]
possible_windows = np.unique(possible_windows, axis=1).transpose()
patterns.append(possible_windows)
pattern_sizes.append(possible_windows.shape[0])
indexes = np.repeat(tour_indexes, pattern_sizes)
patterns = np.concatenate(patterns)
# If we've done everything right, the indexes
# calculated above should be the same length as
# the pattern options
assert patterns.shape[0] == len(indexes)
return indexes, patterns, pattern_sizes
def apply_stage_two_model(omnibus_spec, trips, chunk_size, trace_label):
if not trips.index.is_monotonic:
trips = trips.sort_index()
# Assign the duration of the appropriate leg to the trip
trips[TRIP_DURATION] = np.where(trips[OUTBOUND], trips[OB_DURATION], trips[IB_DURATION])
trips['depart'] = -1
# If this is the first outbound trip, the choice is easy, assign the depart time
# to equal the tour start time.
trips.loc[(trips['trip_num'] == 1) & (trips[OUTBOUND]), 'depart'] = trips['start']
# If its the first return leg, it is easy too. Just assign the trip start time to the
# end time minus the IB duration
trips.loc[(trips['trip_num'] == 1) & (~trips[OUTBOUND]), 'depart'] = trips['end'] - trips[IB_DURATION]
# The last leg of the outbound tour needs to begin at the start plus OB duration
trips.loc[(trips['trip_count'] == trips['trip_num']) & (trips[OUTBOUND]), 'depart'] = \
trips['start'] + trips[OB_DURATION]
# The last leg of the inbound tour needs to begin at the end time of the tour
trips.loc[(trips['trip_count'] == trips['trip_num']) & (~trips[OUTBOUND]), 'depart'] = \
trips['end']
# Slice off the remaining trips with an intermediate stops to deal with.
# Hopefully, with the tricks above we've sliced off a lot of choices.
# This slice should only include trip numbers greater than 2 since the
side_trips = trips[(trips['trip_num'] != 1) & (trips['trip_count'] != trips['trip_num'])]
# No processing needs to be done because we have simple trips / tours
if side_trips.empty:
assert trips['depart'].notnull().all
return trips['depart'].astype(int)
# Get the potential time windows
time_windows = get_time_windows(side_trips[TRIP_DURATION].max(), side_trips[TRIP_COUNT].max() - 1)
trip_list = []
for i, chooser_chunk, chunk_trace_label in \
chunk.adaptive_chunked_choosers_by_chunk_id(side_trips, chunk_size, trace_label):
for is_outbound, trip_segment in chooser_chunk.groupby(OUTBOUND):
direction = OUTBOUND if is_outbound else 'inbound'
spec = get_spec_for_segment(omnibus_spec, direction)
segment_trace_label = '{}_{}'.format(direction, chunk_trace_label)
patterns = build_patterns(trip_segment, time_windows)
choices = choose_tour_leg_pattern(trip_segment,
patterns, spec, trace_label=segment_trace_label)
choices = pd.merge(choices.reset_index(), patterns.reset_index(),
on=[TOUR_LEG_ID, PATTERN_ID], how='left')
choices = choices[['trip_id', 'stop_time_duration']].copy()
trip_list.append(choices)
trip_list = pd.concat(trip_list, sort=True).set_index('trip_id')
trips['stop_time_duration'] = 0
trips.update(trip_list)
trips.loc[trips['trip_num'] == 1, 'stop_time_duration'] = trips['depart']
trips.sort_values(['tour_id', 'outbound', 'trip_num'])
trips['stop_time_duration'] = trips.groupby(['tour_id', 'outbound'])['stop_time_duration'].cumsum()
trips.loc[trips['trip_num'] != trips['trip_count'], 'depart'] = trips['stop_time_duration']
return trips['depart'].astype(int)