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))