Coverage for src/shephex/study/study.py: 98%

64 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2025-03-29 18:45 +0100

1from pathlib import Path 

2from typing import List, Optional, Tuple, Union 

3 

4from shephex.experiment.experiment import Experiment 

5from shephex.experiment.options import Options 

6from shephex.study.table import LittleTable as Table 

7 

8 

9class Study: 

10 """ 

11 A study is a series of related experiments. The role of the study object is 

12 to manage experiments and provide a common interface for interacting with 

13 them. 

14 """ 

15 

16 def __init__(self, path: Union[Path, str], refresh: bool = True, avoid_duplicates: bool = True) -> None: 

17 self.path = Path(path) 

18 self.avoid_duplicates = avoid_duplicates 

19 if refresh: 

20 self.refresh(clear_table=True) 

21 

22 def add_experiment( 

23 self, experiment: Experiment, verbose: bool = True, check_contain: bool = True 

24 ) -> bool: 

25 """ 

26 Add an experiment to the study. 

27 """ 

28 if check_contain: 

29 contained = self.contains_experiment(experiment) 

30 elif not check_contain: 

31 contained = False 

32 

33 if not contained: 

34 self.table.add_row(experiment.to_dict(), add_columns=True) 

35 if not experiment.directory.exists(): 

36 experiment.dump() 

37 

38 def update_experiment(self, experiment: Experiment) -> None: 

39 """ 

40 Update the experiment in the study. 

41 """ 

42 self.table.update_row(experiment.to_dict()) 

43 

44 def contains_experiment(self, experiment: Experiment) -> bool: 

45 """ 

46 Check if the experiment is already in the study. 

47 

48 Returns 

49 -------- 

50 contains: bool 

51 True if the experiment is in the study, False otherwise. 

52 """ 

53 if not self.avoid_duplicates: 

54 return False 

55 

56 contains = self.table.contains_row(experiment.to_dict()) 

57 return contains 

58 

59 def discover_experiments(self) -> List[Path]: 

60 return self.path.glob(f'*-{Experiment.extension}') 

61 

62 def refresh(self, clear_table: bool = False) -> None: 

63 # Check if the study directory exists 

64 if not self.path.exists(): 

65 self.path.mkdir(parents=True) # pragma: no cover 

66 

67 if clear_table: 

68 self.table = Table() 

69 

70 # Get the list of experiments 

71 experiments_paths = self.discover_experiments() 

72 

73 # Add the experiments to the table 

74 for experiment_path in experiments_paths: 

75 experiment = Experiment.load(experiment_path, load_procedure=False) 

76 if not self.contains_experiment(experiment): 

77 self.add_experiment(experiment, verbose=True, check_contain=True) 

78 else: 

79 self.update_experiment(experiment) 

80 

81 def report(self) -> None: 

82 from shephex.study.renderer import StudyRenderer 

83 StudyRenderer().render_study(self) 

84 

85 def get_experiments( 

86 self, 

87 status: str = None, 

88 load_procedure: bool = True, 

89 loaded_experiments: Optional[List[Experiment]] = None, 

90 ) -> List[Experiment]: 

91 """ 

92 Get the experiments in the study. 

93 

94 Todo: This is probably quite inefficient, so should be improved. 

95 """ 

96 experiments = [] 

97 

98 if status != 'all' or status is None: 

99 identifiers = self.table.where(status=status) 

100 else: 

101 identifiers = [row.identifier for row in self.table.table] 

102 

103 if loaded_experiments is not None: 

104 loaded_ids = [experiment.identifier for experiment in loaded_experiments] 

105 else: 

106 loaded_ids = [] 

107 

108 for identifier in identifiers: 

109 if identifier in loaded_ids: 

110 experiment = loaded_experiments[loaded_ids.index(identifier)] 

111 else: 

112 path = self.path / f'{identifier}-{Experiment.extension}' 

113 experiment = Experiment.load(path, load_procedure=load_procedure) 

114 

115 experiments.append(experiment) 

116 return experiments 

117 

118 def where(self, load_shephex_options: bool = True, *args, **kwargs) -> Union[List[str], Tuple[List[str], List[Options]]]: 

119 identifiers = self.table.where(*args, **kwargs) 

120 if load_shephex_options: 

121 paths = [Path(self.path) / f"{id_}-{Experiment.extension}" for id_ in identifiers] 

122 options = [Options.load(path / "shephex") for path in paths] 

123 return identifiers, options 

124 

125 return identifiers 

126