Coverage for src / cosmic_toolbox / utils.py: 72%

1# Copyright (C) 2023 ETH Zurich 

2# Institute for Particle Physics and Astrophysics 

3# Author: Silvan Fischbacher 

4 

5""" 

6General utilities for argument parsing, multiprocessing, and convenience functions. 

7 

8Provides functions for: 

9- Converting argument strings to dictionaries 

10- Parsing command-line sequences (lists, tuples) 

11- Running functions in parallel with multiprocessing 

12- Miscellaneous helper functions 

13 

14""" 

15 

16import argparse 

17import ast 

18import re 

19import time 

20from multiprocessing import Pool 

21 

22import numpy as np 

23 

24from cosmic_toolbox.logger import get_logger 

25 

26LOGGER = get_logger(__file__) 

27 

28 

29def arg_str_to_dict(arg_str): 

30 """ 

31 Converts a string in the format "{arg:value}" or "{arg1:value1,arg2:value2,...}" 

32 to a dictionary with keys and values. 

33 Note: strings should not contain ' or ". 

34 

35 :param arg_str: A string in the format "{arg:value}" or 

36 "{arg1:value1, arg2:value2,...}". 

37 :return arg_dict: dictionary with keys and values corresponding to the input string. 

38 """ 

39 arg_dict = {} 

40 # Use regular expression to extract argument pairs 

41 arg_pairs = re.findall(r"(\w+?):((\d+\.\d+)|(\d+)|([-\w.\[\],()]+))", arg_str) 

42 # Loop over argument pairs and add to dictionary 

43 for pair in arg_pairs: 

44 arg_name, arg_value = pair[0], pair[1] 

45 # If the value is a tuple, convert it to a tuple 

46 if ("(" in arg_value and ")" in arg_value) or ( 

47 "[" in arg_value and "]" in arg_value 

48 ): 

49 arg_value = ast.literal_eval(arg_value) 

50 else: 

51 # If the value is an integer, convert it to an integer 

52 try: 

53 arg_value = int(arg_value) 

54 # If the value is a float, convert it to a float 

55 except ValueError: 

56 try: 

57 arg_value = float(arg_value) 

58 except ValueError: 

59 # If the value is a string, strip the quotes 

60 arg_value = arg_value.strip("'\"") 

61 arg_dict[arg_name] = arg_value 

62 return arg_dict 

63 

64 

65def parse_sequence(s): 

66 """ 

67 Parses a string to a list/tuple for argparse. Can be used as type for argparse. 

68 

69 :param s: String to parse. 

70 :return: tuple or list. 

71 :raises argparse.ArgumentTypeError: If the string cannot be parsed to a tuple. 

72 """ 

73 try: 

74 # Using ast.literal_eval to safely evaluate the string as a Python literal 

75 return ast.literal_eval(s) 

76 except (ValueError, SyntaxError) as err: 

77 raise argparse.ArgumentTypeError(f"Invalid list value: {s}") from err 

78 

79 

80def parse_list(s): 

81 """ 

82 Parses a string to a list for argparse. Can be used as type for argparse. 

83 

84 :param s: String to parse. 

85 :return: list. 

86 """ 

87 return list(parse_sequence(s)) 

88 

89 

90def is_between(x, min, max): 

91 """ 

92 Checks if x is between min and max. 

93 

94 :param x: Value to check. 

95 :param min: Minimum value. 

96 :param max: Maximum value. 

97 :return: True if x is between min and max, False otherwise. 

98 """ 

99 

100 return (x > min) & (x < max) 

101 

102 

103def run_imap_multiprocessing(func, argument_list, num_processes, verb=True): 

104 """ 

105 Runs a function with a list of arguments in parallel using multiprocessing. 

106 

107 :param func: Function to run. 

108 :param argument_list: List of arguments to run the function with. 

109 :param num_processes: Number of processes to use. 

110 :param verb: If True, show progress bar. 

111 :return: List of results from the function. 

112 """ 

113 

114 pool = Pool(processes=num_processes) 

115 

116 result_list = [] 

117 if verb: 

118 for result in LOGGER.progressbar( 

119 pool.imap(func=func, iterable=argument_list), 

120 total=len(argument_list), 

121 at_level="info", 

122 ): 

123 result_list.append(result) 

124 else: 

125 for result in pool.imap(func=func, iterable=argument_list): 

126 result_list.append(result) 

127 

128 return result_list 

129 

130 

131def random_sleep(max_seconds=0, min_seconds=0): 

132 """ 

133 Sleeps for a random amount of time between min_seconds and max_seconds. 

134 

135 :param max_seconds: Maximum number of seconds to sleep. 

136 :param min_seconds: Minimum number of seconds to sleep. 

137 """ 

138 sec = np.random.uniform(min_seconds, max_seconds) 

139 

140 # Format time in a more readable way 

141 if sec < 60: 

142 time_str = f"{sec:.2f} seconds" 

143 elif sec < 3600: 

144 minutes = int(sec // 60) 

145 seconds = sec % 60 

146 time_str = f"{minutes} min {seconds:.2f} sec" 

147 else: 

148 hours = int(sec // 3600) 

149 minutes = int((sec % 3600) // 60) 

150 seconds = sec % 60 

151 time_str = f"{hours} h {minutes} min {seconds:.2f} sec" 

152 

153 if sec > 0: 

154 LOGGER.critical(f"Sleeping for {time_str}") 

155 time.sleep(sec)