def test_stitcher_wrong_inputs(fake_tracklet):
with pytest.raises(IOError):
_ = TrackletStitcher([], n_tracks=2)
with pytest.raises(ValueError):
_ = TrackletStitcher([fake_tracklet], n_tracks=1)
with pytest.raises(ValueError):
_ = TrackletStitcher([fake_tracklet], n_tracks=2, min_length=2)
def fake_stitcher():
inds = np.arange(TRACKLET_LEN)
data = np.random.rand(inds.size, N_DETS, 3)
track = Tracklet(data, inds)
idx = np.linspace(0, inds.size, N_TRACKLETS + 1, dtype=int)
tracklets = TrackletStitcher.split_tracklet(track, idx[1:-1])
return TrackletStitcher(tracklets, n_tracks=2)
def test_stitcher_with_identity(real_tracklets):
# Add fake IDs
for i in range(3):
tracklet = real_tracklets[i]
for v in tracklet.values():
v[:, -1] = i
stitcher = TrackletStitcher.from_dict_of_dict(real_tracklets, n_tracks=3)
tracklets = sorted(stitcher, key=lambda t: t.identity)
assert all(tracklet.identity == i for i, tracklet in enumerate(tracklets))
# Split all tracklets in half
tracklets = [
t for track in stitcher for t in stitcher.split_tracklet(track, [25])
]
stitcher = TrackletStitcher(tracklets, n_tracks=3)
assert len(stitcher) == 6
stitcher.build_graph()
weight = stitcher.G.edges[('0out', '3in')]['weight']
def weight_func(t1, t2):
w = 0.01 if t1.identity == t2.identity else 1
return w * t1.distance_to(t2)
stitcher.build_graph(weight_func=weight_func)
assert stitcher.G.number_of_edges() == 27
new_weight = stitcher.G.edges[('0out', '3in')]['weight']
assert new_weight == weight // 100
stitcher.stitch()
assert len(stitcher.tracks) == 3
assert all(len(track) == 50 for track in stitcher.tracks)
assert all(0.998 <= track.likelihood <= 1 for track in stitcher.tracks)
tracks = sorted(stitcher.tracks, key=lambda t: t.identity)
assert all(track.identity == i for i, track in enumerate(tracks))