def check_mst_vs_scipy(stages_full: Dict[str, Stage],
mst: Dict[str, Stage]) -> None:
"""
compare our MST solution to scipy baseline
:param stages_full: full network
:param mst: mst solution
"""
assign_ids = np.array(
[stage_id for stage_id in sorted(stages_full.keys())])
N = assign_ids.size
# create cost array
cost = np.zeros((len(stages_full), len(stages_full)))
for stage_id, stage in stages_full.items():
for down_stage_id in stage.down_stages:
down_stage = stages_full[down_stage_id]
i = np.flatnonzero(assign_ids == stage_id)
j = np.flatnonzero(assign_ids == down_stage_id)
cost[j, i] = cost[i, j] = GuaranteedServiceModelDAG._get_edge_cost(
stage, down_stage)
C = minimum_spanning_tree(cost).toarray()
assert C.sum() == calc_edge_cost(
mst), 'different costs for MST and scipy solvers'
for i, c in enumerate(C):
for d in range(N):
if assign_ids[d] in mst[assign_ids[i]].down_stages.keys():
assert C[i, d] != 0, 'scipy said there is no edge'
else:
assert C[i, d] == 0, 'scipy said there must be an edge'
def calc_edge_cost(stages: Dict[str, Stage]):
cost = 0
for stage_id, stage in stages.items():
for down_stage_id in stage.down_stages:
down_stage = stages[down_stage_id]
cost += GuaranteedServiceModelDAG._get_edge_cost(stage, down_stage)
return cost