Coverage for parallel_bilby/analysis/plotting.py: 89%

1import dynesty.plotting as dyplot 

2import matplotlib 

3import matplotlib.pyplot as plt 

4import numpy as np 

5from bilby.core.sampler.dynesty import dynesty_stats_plot 

6from bilby.core.utils import logger 

7 

8from ..utils import stopwatch 

9 

10matplotlib.use("Agg") 

11 

12 

13@stopwatch 

14def plot_current_state(sampler, search_parameter_keys, outdir, label): 

15 labels = [label.replace("_", " ") for label in search_parameter_keys] 

16 try: 

17 filename = f"{outdir}/{label}_checkpoint_trace.png" 

18 fig = dyplot.traceplot(sampler.results, labels=labels)[0] 

19 fig.tight_layout() 

20 fig.savefig(filename) 

21 except ( 

22 AssertionError, 

23 RuntimeError, 

24 np.linalg.linalg.LinAlgError, 

25 ValueError, 

26 ) as e: 

27 logger.warning(e) 

28 logger.warning("Failed to create dynesty state plot at checkpoint") 

29 finally: 

30 plt.close("all") 

31 try: 

32 filename = f"{outdir}/{label}_checkpoint_run.png" 

33 fig, axs = dyplot.runplot(sampler.results, mark_final_live=False) 

34 fig.tight_layout() 

35 plt.savefig(filename) 

36 except (RuntimeError, np.linalg.linalg.LinAlgError, ValueError) as e: 

37 logger.warning(e) 

38 logger.warning("Failed to create dynesty run plot at checkpoint") 

39 finally: 

40 plt.close("all") 

41 try: 

42 filename = f"{outdir}/{label}_checkpoint_stats.png" 

43 fig, _ = dynesty_stats_plot(sampler) 

44 fig.tight_layout() 

45 plt.savefig(filename) 

46 except (RuntimeError, ValueError) as e: 

47 logger.warning(e) 

48 logger.warning("Failed to create dynesty stats plot at checkpoint") 

49 finally: 

50 plt.close("all")