Coverage for adrs/search/random.py: 79%

1import math 

2import numpy as np 

3from typing import Sequence, override, cast 

4 

5from adrs.search import Search, ParameterGrid, Permutation 

6 

7from enum import StrEnum 

8 

9 

10class Distribution(StrEnum): 

11 UNIFORM = "uniform" 

12 NORMAL = "normal" 

13 LOGNORMAL = "lognormal" 

14 EXPONENTIAL = "exponential" 

15 BETA = "beta" 

16 GAMMA = "gamma" 

17 

18 

19class RandomSearch(Search): 

20 def __init__( 

21 self, 

22 samples: int | float | np.float64, 

23 seed: int, 

24 dist: Distribution, 

25 **dist_kwargs, 

26 ): 

27 self.samples = samples 

28 self.rng = np.random.default_rng(seed) 

29 self.dist = dist 

30 self.dist_kwargs = dist_kwargs 

31 

32 @staticmethod 

33 def id() -> str: 

34 return "randomsearch" 

35 

36 def _dist(self, n: int) -> np.ndarray: 

37 match self.dist: 

38 case Distribution.UNIFORM: 

39 return self.rng.uniform(0, 1, n) 

40 case Distribution.NORMAL: 

41 return self.rng.normal(0, 1, n) 

42 case Distribution.LOGNORMAL: 

43 return self.rng.lognormal(0, 1, n) 

44 case Distribution.EXPONENTIAL: 

45 return self.rng.exponential(1, n) 

46 case Distribution.BETA: 

47 return self.rng.beta( 

48 self.dist_kwargs.get("a", 2), self.dist_kwargs.get("b", 5), n 

49 ) 

50 case Distribution.GAMMA: 

51 return self.rng.gamma( 

52 self.dist_kwargs.get("shape", 2), 

53 self.dist_kwargs.get("scale", 1), 

54 n, 

55 ) 

56 case dist: 

57 raise ValueError(f"Invalid distribution: {dist}") 

58 

59 @override 

60 def search(self, grid: ParameterGrid) -> Sequence[Permutation]: 

61 if len(grid) == 0: 

62 return [] 

63 

64 # determine number of samples 

65 total_perms = math.prod(map(lambda x: len(x), grid.values())) 

66 n = ( 

67 round(total_perms * self.samples) 

68 if isinstance(self.samples, float | np.float64) 

69 else cast(int, self.samples) 

70 ) 

71 

72 permutations = [] 

73 

74 for _ in range(n): 

75 perm: Permutation = {} 

76 for k, v in grid.items(): 

77 # sample from distribution and avoid negatives 

78 weights = np.abs(self._dist(len(v))) # type: ignore 

79 probs = weights / weights.sum() # normalize to (0, 1) 

80 perm[k] = self.rng.choice(a=v, p=probs) 

81 permutations.append(perm) 

82 

83 return permutations 

84 

85 @override 

86 def filter[T](self, permutations: Sequence[T]) -> Sequence[T]: 

87 # determine number of samples 

88 n = ( 

89 round(len(permutations) * self.samples) 

90 if isinstance(self.samples, float | np.float64) 

91 else cast(int, self.samples) 

92 ) 

93 

94 weights = np.abs(self._dist(len(permutations))) 

95 probs = weights / weights.sum() 

96 return self.rng.choice(a=permutations, p=probs, size=n, replace=False) # type: ignore