Coverage for parallel_bilby/utils.py: 93%

1import functools 

2import logging 

3import os 

4import sys 

5import time 

6 

7import bilby 

8from dynesty.utils import get_print_fn_args 

9 

10logger = bilby.core.utils.logger 

11 

12 

13def get_cli_args(): 

14 """Tool to get CLI args (also makes testing easier)""" 

15 return sys.argv[1:] 

16 

17 

18def get_version_information(): 

19 version_file = os.path.join( 

20 os.path.dirname(os.path.dirname(__file__)), "parallel_bilby/.version" 

21 ) 

22 try: 

23 with open(version_file, "r") as f: 

24 return f.readline().rstrip() 

25 except EnvironmentError: 

26 print("No version information file '.version' found") 

27 

28 

29def safe_file_dump(data, filename, module): 

30 """Safely dump data to a .pickle file 

31 

32 Parameters 

33 ---------- 

34 data: 

35 data to dump 

36 filename: str 

37 The file to dump to 

38 module: pickle, dill 

39 The python module to use 

40 """ 

41 

42 temp_filename = filename + ".temp" 

43 with open(temp_filename, "wb") as file: 

44 module.dump(data, file) 

45 os.rename(temp_filename, filename) 

46 

47 

48def stopwatch(_func=None, *, log_level=logging.DEBUG): 

49 def decorator_stopwatch(func): 

50 @functools.wraps(func) 

51 def wrapper(*args, **kwargs): 

52 t0 = time.time() 

53 result = func(*args, **kwargs) 

54 dur = time.time() - t0 

55 logger.log(log_level, f"{func.__name__} took {dur}s") 

56 return result 

57 

58 return wrapper 

59 

60 if _func is None: 

61 return decorator_stopwatch 

62 else: 

63 return decorator_stopwatch(_func) 

64 

65 

66def stdout_sampling_log(**kwargs): 

67 """Logs will look like: 

68 #:282|eff(%):26.406|logl*:-inf<-160.2<inf|logz:-165.5+/-0.1|dlogz:1038.1>0.1 

69 

70 Adapted from dynesty 

71 https://github.com/joshspeagle/dynesty/blob/bb1c5d5f9504c9c3bbeffeeba28ce28806b42273/py/dynesty/utils.py#L349 

72 """ 

73 niter, short_str, mid_str, long_str = get_print_fn_args(**kwargs) 

74 custom_str = [f"#: {niter:d}"] + mid_str 

75 custom_str = "|".join(custom_str).replace(" ", "") 

76 sys.stdout.write("\033[K" + custom_str + "\r") 

77 sys.stdout.flush()