Skip to content

Instantly share code, notes, and snippets.

@mbrenig
Created May 19, 2017 17:02
Show Gist options
  • Save mbrenig/66eccd7e36c1107e27c7a4929895d2bc to your computer and use it in GitHub Desktop.
Save mbrenig/66eccd7e36c1107e27c7a4929895d2bc to your computer and use it in GitHub Desktop.
Assign week to quarters.
"""
Map the next 1000 weeks to quarters in a sensible way so that
most quarters are 13 weeks long... and occasionally we introduce
a 14 week quater, so as to maximise the amount of working days that
fall in their correct quarter
"""
from datetime import date, timedelta
from collections import defaultdict
import networkx as nx
JAN, FEB, MAR = 1, 2, 3
APR, MAY, JUN = 4, 5, 6
JUL, AUG, SEP = 7, 8, 9
OCT, NOV, DEC = 10, 11, 12
Q1 = 'Q1'
Q2 = 'Q2'
Q3 = 'Q3'
Q4 = 'Q4'
NEXT_YEAR = +1
THIS_YEAR = 0
LAST_YEAR = -1
MO2QTR = {SEP: (Q1, NEXT_YEAR),
OCT: (Q1, NEXT_YEAR),
NOV: (Q1, NEXT_YEAR),
DEC: (Q2, NEXT_YEAR),
JAN: (Q2, THIS_YEAR),
FEB: (Q2, THIS_YEAR),
MAR: (Q3, THIS_YEAR),
APR: (Q3, THIS_YEAR),
MAY: (Q3, THIS_YEAR),
JUN: (Q4, THIS_YEAR),
JUL: (Q4, THIS_YEAR),
AUG: (Q4, THIS_YEAR)}
class Week:
""" Represents a single week """
def __init__(self, monday):
""" Initialize a week object """
if monday.weekday() != 0:
raise ValueError('%s is not a Monday' % monday)
self.monday = monday
self.weekdays = [self.monday+timedelta(i) for i in range(5)]
# Establish the candidate quarters for this week.
self.quarters = defaultdict(int)
for weekday in self.weekdays:
year, month = weekday.year, weekday.month
qtr, ydelta = MO2QTR[month]
quarter = '%s-%s' % (year+ydelta, qtr)
self.quarters[quarter] += 1
ordered_quarters = sorted(self.quarters.items(),
key=lambda x: x[1],
reverse=True)
self.pref_qtr = ordered_quarters[0][0]
self.alt_qtr = None
if len(ordered_quarters) > 1:
self.alt_qtr = ordered_quarters[1][0]
self.selected_qtr = None
def min_qtr(self):
""" Return the earliest quarter """
return min(self.pref_qtr, self.alt_qtr)
def max_qtr(self):
""" Return the latest quarter """
return max(self.pref_qtr, self.alt_qtr)
def __lt__(self, other):
return self.monday < other.monday
def __repr__(self):
if self.alt_qtr:
return '<Week %s *%s* %s>' % (self.monday, self.pref_qtr, self.alt_qtr)
return '<Week %s %s>' % (self.monday, self.pref_qtr)
def next_weekday(dtime, weekday=0):
""" Get the next weekday after date d
0 = Monday, 1 = Tuesday, ... 6=Sunday """
days_ahead = weekday - dtime.weekday()
if days_ahead <= 0: # Target day already happened this week
days_ahead += 7
nextwd = dtime + timedelta(days_ahead)
return date(year=nextwd.year, month=nextwd.month, day=nextwd.day)
def gen_weeks(num=1000, start_date=date.today()):
""" Create n weeks from the monday before start_date, onwards """
last_monday = next_weekday(start_date) - timedelta(days=7)
weeks = []
for _ in range(num):
week = Week(last_monday)
weeks.append(week)
last_monday += timedelta(7)
return weeks
def make_graph(undecided):
""" Make a graph of the weeks listed in undecided. Edges are matching qtrs """
graph = nx.Graph()
common_qtrs = defaultdict(set)
for week in undecided:
graph.add_node(week)
common_qtrs[week.pref_qtr].add(week)
common_qtrs[week.alt_qtr].add(week)
for nodes in common_qtrs.values():
if len(nodes) == 1:
continue
left, right = list(nodes)
graph.add_edge(left, right)
return graph
def select_qtrs(weeks):
""" given a sequence of weeks, select the quarters they'll go into """
qtow = defaultdict(set)
undecided = []
for week in weeks:
if len(week.quarters) == 1:
# Must pick preferred quarter
qtow[week.pref_qtr].add(week)
else:
undecided.append(week)
trivial_progress = True
while trivial_progress:
trivial_progress = False
decided = []
for week in undecided:
if len(qtow[week.pref_qtr]) == 12 and len(qtow[week.alt_qtr]) == 13:
# Trivial
qtow[week.pref_qtr].add(week)
decided.append(week)
trivial_progress = True
for week in decided:
undecided.remove(week)
print("Need to decide where %s of %s weeks fall" % (len(undecided), len(weeks)))
# Lets make a graph of the undecided weeks to see which are connected runs.
graph = make_graph(undecided)
runs = sorted(nx.connected_component_subgraphs(graph), key=lambda x: -len(x.nodes()))
print("Need to decide where %s runs of weeks fall" % len(runs))
decided_cnt = 0
for run in runs:
# Each connected run of qtr spanning weeks needs to be scheduled early or late.
# Pick whichever leads to the largest number of correclty aligned days.
early, late = 0, 0
for week in run.nodes():
# NB this breaks if you name the quartes something that doesn't sort alphabetically.
early += week.quarters[week.min_qtr()]
late += week.quarters[week.max_qtr()]
for week in run.nodes():
if early > late:
qtow[week.min_qtr()].add(week)
elif late > early:
qtow[week.max_qtr()].add(week)
decided_cnt += 1
return qtow
def print_assignments(qtow):
""" Pretty print the quarter to week mapping. """
for qtr in sorted(qtow.keys()):
weeks = qtow[qtr]
if not weeks:
continue
first_week = min(weeks)
last_week = max(weeks)
num = len(weeks)
print('%s\tWeek 1 starting %s\tWeek %s starting %s' % (qtr,
first_week.monday,
num,
last_week.monday))
def calculate():
""" Create loads of week objects and select quarters for them. """
weeks = gen_weeks(num=3000, start_date=date(2010, 9, 10))
assignments = select_qtrs(weeks)
print_assignments(assignments)
if __name__ == '__main__':
calculate()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment