Coverage for bmm_multitask_learning/coalescent/plot_coalescent.py: 0%

1import matplotlib.pyplot as plt 

2 

3 

4def plot_coalesce(y_history, pairs): 

5 fig, ax = plt.subplots(1, 1) 

6 old_t = 0 

7 for i, layer in enumerate(y_history): 

8 new_item = min(layer, key=lambda x: x.t) 

9 max_t = new_item.t 

10 for item in layer: 

11 ax.plot([max_t], item.mean[0], ".") 

12 if i > 0: 

13 

14 # рисуем предков 

15 predecessors = pairs[i-1] 

16 for predecessor in predecessors: 

17 src = y_history[i-1][predecessor] 

18 x, y, dx, dy = old_t, src.mean[0], \ 

19 max_t - old_t, \ 

20 -(src.mean - new_item.mean)[0] 

21 

22 ax.arrow(x, y, dx, dy, width=-dx/1000) 

23 for j in range(len(layer) + 1): 

24 if j in predecessors: 

25 continue 

26 src = y_history[i-1][j] 

27 x, y, dx, dy = max_t, src.mean[0], -(max_t - old_t), 0 

28 ax.arrow(x, y, dx, dy, width=-dx/1000) 

29 

30 old_t = max_t 

31 return fig 

32