Coverage for ibllib/io/extractors/base.py: 89%

1"""Base Extractor classes. 

2 

3A module for the base Extractor classes. The Extractor, given a session path, will extract the 

4processed data from raw hardware files and optionally save them. 

5""" 

6 

7import abc (empty)

8from collections import OrderedDict (empty)

9import json (empty)

10from pathlib import Path (empty)

11 

12import numpy as np (empty)

13import pandas as pd (empty)

14from ibllib.io import raw_data_loaders as raw (empty)

15from ibllib.io.raw_data_loaders import load_settings, _logger (empty)

16 

17 

18class BaseExtractor(abc.ABC): (empty)

19 """ 

20 Base extractor class. 

21 

22 Writing an extractor checklist: 

23 

24 - on the child class, overload the _extract method 

25 - this method should output one or several numpy.arrays or dataframe with a consistent shape 

26 - save_names is a list or a string of filenames, there should be one per dataset 

27 - set save_names to None for a dataset that doesn't need saving (could be set dynamically in 

28 the _extract method) 

29 

30 :param session_path: Absolute path of session folder 

31 :type session_path: str/Path 

32 """ 

33 

34 session_path = None (empty)

35 """pathlib.Path: Absolute path of session folder.""" (empty)

36 

37 save_names = None (empty)

38 """tuple of str: The filenames of each extracted dataset, or None if array should not be saved.""" (empty)

39 

40 var_names = None (empty)

41 """tuple of str: A list of names for the extracted variables. These become the returned output keys.""" (empty)

42 

43 default_path = Path('alf') # relative to session (empty)

44 """pathlib.Path: The default output folder relative to `session_path`.""" (empty)

45 

46 def __init__(self, session_path=None): (empty)

47 # If session_path is None Path(session_path) will fail 

48 self.session_path = Path(session_path) 77 ctx1a'()*IJKLMNOPQRSTGUVWXbYE+,3clwxjdefH-8BC./gm:tu9;=?!Z0#$12y%vhkpnoq@[]^_4irs

49 

50 def extract(self, save=False, path_out=None, **kwargs): (empty)

51 """ 

52 :return: dict of numpy.array, list of filenames 

53 """ 

54 out = self._extract(**kwargs) 60 ctx1azIJKLMNOPQRSTGUVWXbY3clwxjdef8BCgmtu9!Z0#$1256y%vhkpnoq4irs

55 files = self._save(out, path_out=path_out) if save else None 60 ctx1azIJKLMNOPQRSTGUVWXbY3clwxjdef8BCgmtu9!Z0#$1256y%vhkpnoq4irs

56 return out, files 60 ctx1azIJKLMNOPQRSTGUVWXbY3clwxjdef8BCgmtu9!Z0#$1256y%vhkpnoq4irs

57 

58 def _save(self, data, path_out=None): (empty)

59 # Check if self.save_names is of the same length of out 

60 if not path_out: 21 ctx1zbclwxjdefBCgmtuyvhki

61 path_out = self.session_path.joinpath(self.default_path) 13 ctx1zbcldefBCtuvi

62 

63 def _write_to_disk(file_path, data): 21 ctx1zbclwxjdefBCgmtuyvhki

64 """Implements different save calls depending on file extension. 

65 

66 Parameters 

67 ---------- 

68 file_path : pathlib.Path 

69 The location to save the data. 

70 data : pandas.DataFrame, numpy.ndarray 

71 The data to save 

72 

73 """ 

74 csv_separators = { 21 ctx1zbclwxjdefBCgmtuyvhki

75 ".csv": ",", 

76 ".ssv": " ", 

77 ".tsv": "\t" 

78 } 

79 # Ensure empty files are not created; we expect all datasets to have a non-zero size 

80 if getattr(data, 'size', len(data)) == 0: 21 ctx1zbclwxjdefBCgmtuyvhki

81 filename = file_path.relative_to(self.session_path).as_posix() 

82 raise ValueError(f'Data for {filename} appears to be empty') 

83 file_path = Path(file_path) 21 ctx1zbclwxjdefBCgmtuyvhki

84 file_path.parent.mkdir(exist_ok=True, parents=True) 21 ctx1zbclwxjdefBCgmtuyvhki

85 if file_path.suffix == ".npy": 21 ctx1zbclwxjdefBCgmtuyvhki

86 np.save(file_path, data) 21 ctx1zbclwxjdefBCgmtuyvhki

87 elif file_path.suffix in [".parquet", ".pqt"]: 15 ctx1zbcwxjdefgyvhki

88 if not isinstance(data, pd.DataFrame): 10 ctx1bcjdefghki

89 _logger.error("Data is not a panda's DataFrame object") 

90 raise TypeError("Data is not a panda's DataFrame object") 

91 data.to_parquet(file_path) 10 ctx1bcjdefghki

92 elif file_path.suffix in csv_separators: 5 ctx1zwxyv

93 sep = csv_separators[file_path.suffix] 5 ctx1zwxyv

94 data.to_csv(file_path, sep=sep) 5 ctx1zwxyv

95 # np.savetxt(file_path, data, delimiter=sep) 

96 else: 

97 _logger.error(f"Don't know how to save {file_path.suffix} files yet") 

98 

99 if self.save_names is None: 21 ctx1zbclwxjdefBCgmtuyvhki

100 file_paths = [] 

101 elif isinstance(self.save_names, str): 21 ctx1zbclwxjdefBCgmtuyvhki

102 file_paths = path_out.joinpath(self.save_names) 2 ctx1tu

103 _write_to_disk(file_paths, data) 2 ctx1tu

104 elif isinstance(data, dict): 19 ctx1zbclwxjdefBCgmyvhki

105 file_paths = [] 12 ctx1bcljdefgmhki

106 for var, value in data.items(): 12 ctx1bcljdefgmhki

107 if fn := self.save_names[self.var_names.index(var)]: 12 ctx1bcljdefgmhki

108 fpath = path_out.joinpath(fn) 12 ctx1bcljdefgmhki

109 _write_to_disk(fpath, value) 12 ctx1bcljdefgmhki

110 file_paths.append(fpath) 12 ctx1bcljdefgmhki

111 else: # Should be list or tuple... 

112 assert len(data) == len(self.save_names) 7 ctx1zwxBCyv

113 file_paths = [] 7 ctx1zwxBCyv

114 for data, fn in zip(data, self.save_names): 7 ctx1zwxBCyv

115 if fn: 7 ctx1zwxBCyv

116 fpath = path_out.joinpath(fn) 7 ctx1zwxBCyv

117 _write_to_disk(fpath, data) 7 ctx1zwxBCyv

118 file_paths.append(fpath) 7 ctx1zwxBCyv

119 return file_paths 21 ctx1zbclwxjdefBCgmtuyvhki

120 

121 @abc.abstractmethod (empty)

122 def _extract(self): (empty)

123 pass 

124 

125 

126class BaseBpodTrialsExtractor(BaseExtractor): (empty)

127 """ 

128 Base (abstract) extractor class for bpod jsonable data set. 

129 

130 Wraps the _extract private method. 

131 

132 :param session_path: Absolute path of session folder. 

133 :type session_path: str 

134 :param bpod_trials 

135 :param settings 

136 """ 

137 

138 bpod_trials = None (empty)

139 settings = None (empty)

140 task_collection = None (empty)

141 frame2ttl = None (empty)

142 audio = None (empty)

143 

144 def extract(self, bpod_trials=None, settings=None, **kwargs): (empty)

145 """ 

146 :param: bpod_trials (optional) bpod trials from jsonable in a dictionary 

147 :param: settings (optional) bpod iblrig settings json file in a dictionary 

148 :param: save (bool) write output ALF files, defaults to False 

149 :param: path_out (pathlib.Path) output path (defaults to `{session_path}/alf`) 

150 :return: numpy.ndarray or list of ndarrays, list of filenames 

151 :rtype: dtype('float64') 

152 """ 

153 self.bpod_trials = bpod_trials 45 ctx1aIJKLMNOPQRSTGUVWXbY3cljdefgmZ01256hkpnoq4irs

154 self.settings = settings 45 ctx1aIJKLMNOPQRSTGUVWXbY3cljdefgmZ01256hkpnoq4irs

155 self.task_collection = kwargs.pop('task_collection', 'raw_behavior_data') 45 ctx1aIJKLMNOPQRSTGUVWXbY3cljdefgmZ01256hkpnoq4irs

156 if self.bpod_trials is None: 45 ctx1aIJKLMNOPQRSTGUVWXbY3cljdefgmZ01256hkpnoq4irs

157 self.bpod_trials = raw.load_data(self.session_path, task_collection=self.task_collection) 43 ctx1aIJKLMNOPQRSTGUVWXbYcljdefgmZ01256hkpnoqirs

158 if not self.settings: 45 ctx1aIJKLMNOPQRSTGUVWXbY3cljdefgmZ01256hkpnoq4irs

159 self.settings = raw.load_settings(self.session_path, task_collection=self.task_collection) 45 ctx1aIJKLMNOPQRSTGUVWXbY3cljdefgmZ01256hkpnoq4irs

160 if self.settings is None: 45 ctx1aIJKLMNOPQRSTGUVWXbY3cljdefgmZ01256hkpnoq4irs

161 self.settings = {"IBLRIG_VERSION": "100.0.0"} 

162 elif self.settings.get("IBLRIG_VERSION", "") == "": 45 ctx1aIJKLMNOPQRSTGUVWXbY3cljdefgmZ01256hkpnoq4irs

163 self.settings["IBLRIG_VERSION"] = "100.0.0" 3 ctx1Ggh

164 # Get all detected TTLs. These are stored for QC purposes 

165 self.frame2ttl, self.audio = raw.load_bpod_fronts(self.session_path, data=self.bpod_trials) 45 ctx1aIJKLMNOPQRSTGUVWXbY3cljdefgmZ01256hkpnoq4irs

166 

167 return super(BaseBpodTrialsExtractor, self).extract(**kwargs) 45 ctx1aIJKLMNOPQRSTGUVWXbY3cljdefgmZ01256hkpnoq4irs

168 

169 @property (empty)

170 def alf_path(self): (empty)

171 """pathlib.Path: The full task collection filepath.""" 

172 if self.session_path: 

173 return self.session_path.joinpath(self.task_collection or '').absolute() 

174 

175 

176def run_extractor_classes(classes, session_path=None, **kwargs): (empty)

177 """ 

178 Run a set of extractors with the same inputs. 

179 

180 :param classes: list of Extractor class 

181 :param save: True/False 

182 :param path_out: (defaults to alf path) 

183 :param kwargs: extractor arguments (session_path...) 

184 :return: dictionary of arrays, list of files 

185 """ 

186 files = [] 21 ctx1abcljdefgmtuhkpnoqirs

187 outputs = OrderedDict({}) 21 ctx1abcljdefgmtuhkpnoqirs

188 assert session_path 21 ctx1abcljdefgmtuhkpnoqirs

189 # if a single class is passed, convert as a list 

190 try: 21 ctx1abcljdefgmtuhkpnoqirs

191 iter(classes) 21 ctx1abcljdefgmtuhkpnoqirs

192 except TypeError: 2 ctx1no

193 classes = [classes] 2 ctx1no

194 for classe in classes: 21 ctx1abcljdefgmtuhkpnoqirs

195 cls = classe(session_path=session_path) 21 ctx1abcljdefgmtuhkpnoqirs

196 out, fil = cls.extract(**kwargs) 21 ctx1abcljdefgmtuhkpnoqirs

197 if isinstance(fil, list): 21 ctx1abcljdefgmtuhkpnoqirs

198 files.extend(fil) 

199 elif fil is not None: 21 ctx1abcljdefgmtuhkpnoqirs

200 files.append(fil) 2 ctx1tu

201 if isinstance(out, dict): 21 ctx1abcljdefgmtuhkpnoqirs

202 outputs.update(out) 

203 elif isinstance(cls.var_names, str): 21 ctx1abcljdefgmtuhkpnoqirs

204 outputs[cls.var_names] = out 21 ctx1abcljdefgmtuhkpnoqirs

205 else: 

206 for i, k in enumerate(cls.var_names): 19 ctx1abcljdefgmhkpnoqirs

207 outputs[k] = out[i] 19 ctx1abcljdefgmhkpnoqirs

208 return outputs, files 21 ctx1abcljdefgmtuhkpnoqirs

209 

210 

211def get_task_protocol(session_path, task_collection='raw_behavior_data'): (empty)

212 """ 

213 Return the task protocol name from task settings. 

214 

215 If the session path and/or task collection do not exist, the settings file is missing or 

216 otherwise can not be parsed, or if the 'PYBPOD_PROTOCOL' key is absent, None is returned. 

217 A warning is logged if the session path or settings file doesn't exist. An error is logged if 

218 the settings file can not be parsed. 

219 

220 Parameters 

221 ---------- 

222 session_path : str, pathlib.Path 

223 The absolute session path. 

224 task_collection : str 

225 The session path directory containing the task settings file. 

226 

227 Returns 

228 ------- 

229 str or None 

230 The Pybpod task protocol name or None if not found. 

231 """ 

232 try: 21 ctx1abAcljdefgmhkpnoqDirs

233 settings = load_settings(session_path, task_collection=task_collection) 21 ctx1abAcljdefgmhkpnoqDirs

234 except json.decoder.JSONDecodeError: 

235 _logger.error(f'Can\'t read settings for {session_path}') 

236 return 

237 if settings: 21 ctx1abAcljdefgmhkpnoqDirs

238 return settings.get('PYBPOD_PROTOCOL', None) 21 ctx1abAcljdefgmhkpnoqDirs

239 else: 

240 return 1 ctx1A

241 

242 

243def _get_task_extractor_map(): (empty)

244 """ 

245 Load the task protocol extractor map. 

246 

247 Returns 

248 ------- 

249 Dict[str, str] 

250 A map of task protocol to Bpod trials extractor class. 

251 """ 

252 FILENAME = 'task_extractor_map.json' 25 ctx1abEA7FcljdefHgmhkpnoqDirs

253 with open(Path(__file__).parent.joinpath(FILENAME)) as fp: 25 ctx1abEA7FcljdefHgmhkpnoqDirs

254 task_extractors = json.load(fp) 25 ctx1abEA7FcljdefHgmhkpnoqDirs

255 try: 25 ctx1abEA7FcljdefHgmhkpnoqDirs

256 # look if there are custom extractor types in the personal projects repo 

257 import projects 25 ctx1abEA7FcljdefHgmhkpnoqDirs

258 custom_extractors = Path(projects.__file__).parent.joinpath(FILENAME) 25 ctx1abEA7FcljdefHgmhkpnoqDirs

259 with open(custom_extractors, 'r') as fp: 25 ctx1abEA7FcljdefHgmhkpnoqDirs

260 custom_task_types = json.load(fp) 25 ctx1abEA7FcljdefHgmhkpnoqDirs

261 task_extractors.update(custom_task_types) 25 ctx1abEA7FcljdefHgmhkpnoqDirs

262 except (ModuleNotFoundError, FileNotFoundError): 

263 pass 

264 return task_extractors 25 ctx1abEA7FcljdefHgmhkpnoqDirs

265 

266 

267def get_bpod_extractor_class(session_path, task_collection='raw_behavior_data'): (empty)

268 """ 

269 Get the Bpod trials extractor class associated with a given Bpod session. 

270 

271 Note that unlike :func:`get_session_extractor_type`, this function maps directly to the Bpod 

272 trials extractor class name. This is hardware invariant and is purly to determine the Bpod only 

273 trials extractor. 

274 

275 Parameters 

276 ---------- 

277 session_path : str, pathlib.Path 

278 The session path containing Bpod behaviour data. 

279 task_collection : str 

280 The session_path sub-folder containing the Bpod settings file. 

281 

282 Returns 

283 ------- 

284 str 

285 The extractor class name. 

286 """ 

287 # Attempt to get protocol name from settings file 

288 protocol = get_task_protocol(session_path, task_collection=task_collection) 21 ctx1abAcljdefgmhkpnoqDirs

289 if not protocol: 21 ctx1abAcljdefgmhkpnoqDirs

290 raise ValueError(f'No task protocol found in {Path(session_path) / task_collection}') 1 ctx1A

291 return protocol2extractor(protocol) 21 ctx1abAcljdefgmhkpnoqDirs

292 

293 

294def protocol2extractor(protocol): (empty)

295 """ 

296 Get the Bpod trials extractor class associated with a given Bpod task protocol. 

297 

298 The Bpod task protocol can be found in the 'PYBPOD_PROTOCOL' field of the 

299 _iblrig_taskSettings.raw.json file. 

300 

301 Parameters 

302 ---------- 

303 protocol : str 

304 A Bpod task protocol name. 

305 

306 Returns 

307 ------- 

308 str 

309 The extractor class name. 

310 """ 

311 # Attempt to get extractor class from protocol 

312 extractor_map = _get_task_extractor_map() 24 ctx1abEAFcljdefHgmhkpnoqDirs

313 extractor = extractor_map.get(protocol, None) 24 ctx1abEAFcljdefHgmhkpnoqDirs

314 if extractor is None: # Try lazy matching of name 24 ctx1abEAFcljdefHgmhkpnoqDirs

315 extractor = next((extractor_map[tt] for tt in extractor_map if tt in protocol), None) 23 ctx1abEAFcljdefgmhkpnoqDirs

316 if extractor is None: 24 ctx1abEAFcljdefHgmhkpnoqDirs

317 raise ValueError(f'No extractor associated with "{protocol}"') 2 ctx1EF

318 return extractor 24 ctx1abEAFcljdefHgmhkpnoqDirs