Coverage for ibllib/qc/alignment_qc.py: 67%

1import logging (empty)

2 

3import numpy as np (empty)

4from pathlib import Path (empty)

5 

6from neuropixel import trace_header (empty)

7import spikeglx (empty)

8 

9from iblatlas.atlas import AllenAtlas (empty)

10from ibllib.pipes import histology (empty)

11from ibllib.pipes.ephys_alignment import EphysAlignment (empty)

12from ibllib.qc import base (empty)

13from ibllib.oneibl.patcher import FTPPatcher (empty)

14from ibllib.qc.base import CRITERIA as CRITERIA_BASE (empty)

15 

16_log = logging.getLogger(__name__) (empty)

17CRITERIA = {"PASS": 0.8} (empty)

18 

19 

20class AlignmentQC(base.QC): (empty)

21 """ 

22 Class that is used to update the extended_qc of the probe insertion fields with the results 

23 from the ephys alignment procedure 

24 """ 

25 def __init__(self, probe_id, one=None, brain_atlas=None, channels=True, collection=None): (empty)

26 super().__init__(probe_id, one=one, log=_log, endpoint='insertions') 2 ctx1ab

27 

28 # Data 

29 self.alignments = None 2 ctx1ab

30 self.xyz_picks = None 2 ctx1ab

31 self.depths = None 2 ctx1ab

32 self.cluster_chns = None 2 ctx1ab

33 self.chn_coords = None 2 ctx1ab

34 self.align_keys_sorted = None 2 ctx1ab

35 

36 # Metrics and passed trials 

37 self.sim_matrix = None 2 ctx1ab

38 self.criteria = CRITERIA 2 ctx1ab

39 

40 # Get the brain atlas 

41 self.brain_atlas = brain_atlas or AllenAtlas(25) 2 ctx1ab

42 # Flag for uploading channels to alyx. For testing purposes 

43 self.channels_flag = channels 2 ctx1ab

44 

45 self.insertion = self.one.alyx.get(f'/insertions/{self.eid}', clobber=True) 2 ctx1ab

46 self.resolved = (self.insertion.get('json', {'temp': 0}).get('extended_qc'). 2 ctx1ab

47 get('alignment_resolved', False)) 

48 

49 self.probe_collection = collection 2 ctx1ab

50 

51 def load_data(self, prev_alignments=None, xyz_picks=None, depths=None, cluster_chns=None, (empty)

52 chn_coords=None): 

53 """" 

54 Load data required to assess alignment qc and compute similarity matrix. If no arguments 

55 are given load_data will fetch all the relevant data required 

56 """ 

57 if not np.any(prev_alignments): 2 ctx1ab

58 aligned_traj = self.one.alyx.get(f'/trajectories?&probe_insertion={self.eid}' 1 ctx1a

59 '&provenance=Ephys aligned histology track', 

60 clobber=True) 

61 if len(aligned_traj) > 0: 1 ctx1a

62 self.alignments = aligned_traj[0].get('json', {}) 1 ctx1a

63 else: 

64 self.alignments = {} 1 ctx1a

65 return 1 ctx1a

66 else: 

67 self.alignments = prev_alignments 2 ctx1ab

68 

69 align_keys = [*self.alignments.keys()] 2 ctx1ab

70 self.align_keys_sorted = sorted(align_keys, reverse=True) 2 ctx1ab

71 

72 if not np.any(xyz_picks): 2 ctx1ab

73 self.xyz_picks = np.array(self.insertion['json']['xyz_picks']) / 1e6 1 ctx1a

74 else: 

75 self.xyz_picks = xyz_picks 2 ctx1ab

76 

77 if len(self.alignments) < 2: 2 ctx1ab

78 return 1 ctx1a

79 

80 if not self.probe_collection: 2 ctx1ab

81 all_collections = self.one.list_collections(self.insertion['session']) 2 ctx1ab

82 if f'alf/{self.insertion["name"]}/pykilosort' in all_collections: 2 ctx1ab

83 self.probe_collection = f'alf/{self.insertion["name"]}/pykilosort' 

84 else: 

85 self.probe_collection = f'alf/{self.insertion["name"]}' 2 ctx1ab

86 

87 if not np.any(chn_coords): 2 ctx1ab

88 self.chn_coords = self.one.load_dataset(self.insertion['session'], 

89 'channels.localCoordinates.npy', 

90 collection=self.probe_collection) 

91 else: 

92 self.chn_coords = chn_coords 2 ctx1ab

93 

94 if not np.any(depths): 2 ctx1ab

95 self.depths = self.chn_coords[:, 1] 

96 else: 

97 self.depths = depths 2 ctx1ab

98 

99 if not np.any(cluster_chns): 2 ctx1ab

100 self.cluster_chns = self.one.load_dataset(self.insertion['session'], 

101 'clusters.channels.npy', 

102 collection=self.probe_collection) 

103 else: 

104 self.cluster_chns = cluster_chns 2 ctx1ab

105 

106 def compute(self): (empty)

107 """ 

108 Computes the similarity matrix if > 2 alignments. If no data loaded, wraps around load_data 

109 to get all relevant data needed 

110 """ 

111 

112 if self.alignments is None: 2 ctx1ab

113 self.load_data() 1 ctx1a

114 

115 if len(self.alignments) < 2: 2 ctx1ab

116 self.log.info(f"Insertion {self.eid}: One or less alignment found...") 1 ctx1a

117 self.sim_matrix = np.array([len(self.alignments)]) 1 ctx1a

118 else: 

119 self.log.info(f"Insertion {self.eid}: Running QC on alignment data...") 2 ctx1ab

120 self.sim_matrix = self.compute_similarity_matrix() 2 ctx1ab

121 

122 return self.sim_matrix 2 ctx1ab

123 

124 def run(self, update=True, upload_alyx=True, upload_flatiron=True): (empty)

125 """ 

126 Compute alignment_qc for a specified probe insertion and updates extended qc field in alyx. 

127 If alignment is resolved and upload flags set to True channels from resolved 

128 alignment will be updated to alyx and datasets sent to ibl-ftp-patcher to be uploaded to 

129 flatiron 

130 """ 

131 if self.sim_matrix is None: 2 ctx1ab

132 self.compute() 2 ctx1ab

133 

134 # Case where the alignment has already been resolved 

135 if self.resolved: 2 ctx1ab

136 self.log.info(f"Alignment for insertion {self.eid} already resolved, channels won't be" 

137 f" updated. To force update of channels use " 

138 f"resolve_manual method with force=True") 

139 results = {'alignment_count': len(self.alignments)} 

140 if update: 

141 self.update_extended_qc(results) 

142 results.update({'alignment_resolved': True}) 

143 

144 # Case where no alignments have been made 

145 elif np.all(self.sim_matrix == 0) and self.sim_matrix.shape[0] == 1: 2 ctx1ab

146 # We don't update database 

147 results = {'alignment_resolved': False} 1 ctx1a

148 

149 # Case where only one alignment 

150 elif np.all(self.sim_matrix == 1) and self.sim_matrix.shape[0] == 1: 2 ctx1ab

151 results = {'alignment_count': len(self.alignments), 1 ctx1a

152 'alignment_stored': self.align_keys_sorted[0], 

153 'alignment_resolved': False} 

154 if update: 1 ctx1a

155 self.update_extended_qc(results) 1 ctx1a

156 

157 # Case where 2 or more alignments and alignments haven't been resolved 

158 else: 

159 results = self.compute_alignment_status() 2 ctx1ab

160 

161 if update: 2 ctx1ab

162 self.update_extended_qc(results) 2 ctx1ab

163 

164 if results['alignment_resolved'] and np.bitwise_or(upload_alyx, upload_flatiron): 2 ctx1ab

165 self.upload_channels(results['alignment_stored'], upload_alyx, upload_flatiron) 1 ctx1a

166 

167 return results 2 ctx1ab

168 

169 def resolve_manual(self, align_key, update=True, upload_alyx=True, upload_flatiron=True, (empty)

170 force=False): 

171 """ 

172 Method to manually resolve the alignment of a probe insertion with a given alignment 

173 regardless of the number of alignments or the alignment qc value. Channels from specified 

174 alignment will be uploaded to alyx and datasets sent to ibl-ftp-patcher to be uploaded to 

175 flatiron. If alignment already resolved will only upload if force flag set to True 

176 """ 

177 

178 if self.sim_matrix is None: 1 ctx1b

179 self.compute() 1 ctx1b

180 assert align_key in self.align_keys_sorted, 'align key not recognised' 1 ctx1b

181 

182 if self.resolved == 1 and not force: 1 ctx1b

183 self.log.info(f"Alignment for insertion {self.eid} already resolved, channels won't be" 

184 f"updated. To overwrite stored channels with alignment {align_key} " 

185 f"set 'force=True'") 

186 file_paths = [] 

187 else: 

188 if self.sim_matrix.shape[0] > 1: 1 ctx1b

189 results = self.compute_alignment_status() 1 ctx1b

190 else: 

191 results = {'alignment_count': self.sim_matrix.shape[0]} 

192 

193 results['alignment_resolved'] = True 1 ctx1b

194 results['alignment_stored'] = align_key 1 ctx1b

195 results['alignment_resolved_by'] = 'experimenter' 1 ctx1b

196 

197 if update: 1 ctx1b

198 self.update_extended_qc(results) 1 ctx1b

199 file_paths = [] 1 ctx1b

200 

201 if upload_alyx or upload_flatiron: 1 ctx1b

202 file_paths = self.upload_channels(align_key, upload_alyx, upload_flatiron) 1 ctx1b

203 

204 return file_paths 1 ctx1b

205 

206 def compute_similarity_matrix(self): (empty)

207 """ 

208 Computes the similarity matrix between each alignment stored in the ephys aligned 

209 trajectory. Similarity matrix based on number of clusters that share brain region and 

210 parent brain region 

211 """ 

212 

213 r = self.brain_atlas.regions 2 ctx1ab

214 

215 clusters = dict() 2 ctx1ab

216 for iK, key in enumerate(self.align_keys_sorted): 2 ctx1ab

217 # Location of reference lines used for alignment 

218 feature = np.array(self.alignments[key][0]) 2 ctx1ab

219 track = np.array(self.alignments[key][1]) 2 ctx1ab

220 

221 # Instantiate EphysAlignment object 

222 ephysalign = EphysAlignment(self.xyz_picks, self.depths, track_prev=track, 2 ctx1ab

223 feature_prev=feature, 

224 brain_atlas=self.brain_atlas) 

225 

226 # Find xyz location of all channels 

227 xyz_channels = ephysalign.get_channel_locations(feature, track) 2 ctx1ab

228 brain_regions = ephysalign.get_brain_locations(xyz_channels) 2 ctx1ab

229 

230 # Find the location of clusters along the alignment 

231 cluster_info = dict() 2 ctx1ab

232 cluster_info['brain_id'] = brain_regions['id'][self.cluster_chns] 2 ctx1ab

233 cluster_info['parent_id'] = r.get(ids=cluster_info['brain_id']).parent.astype(int) 2 ctx1ab

234 clusters.update({key: cluster_info}) 2 ctx1ab

235 

236 sim_matrix = np.zeros((len(self.align_keys_sorted), len(self.align_keys_sorted))) 2 ctx1ab

237 

238 for ik, key in enumerate(self.align_keys_sorted): 2 ctx1ab

239 for ikk, key2 in enumerate(self.align_keys_sorted): 2 ctx1ab

240 same_id = np.where(clusters[key]['brain_id'] == clusters[key2]['brain_id'])[0] 2 ctx1ab

241 not_same_id = \ 2 ctx1ab

242 np.where(clusters[key]['brain_id'] != clusters[key2]['brain_id'])[0] 

243 same_parent = np.where(clusters[key]['parent_id'][not_same_id] == 2 ctx1ab

244 clusters[key2]['parent_id'][not_same_id])[0] 

245 sim_matrix[ik, ikk] = len(same_id) + (len(same_parent) * 0.5) 2 ctx1ab

246 # Normalise 

247 sim_matrix_norm = sim_matrix / np.max(sim_matrix) 2 ctx1ab

248 

249 return sim_matrix_norm 2 ctx1ab

250 

251 def compute_alignment_status(self): (empty)

252 """ 

253 Determine whether alignments agree based on value in similarity matrix. If any alignments 

254 have similarity of 0.8 set the alignment to be resolved 

255 """ 

256 # Set diagonals to zero so we don't use those to find max 

257 np.fill_diagonal(self.sim_matrix, 0) 2 ctx1ab

258 

259 max_sim = np.max(self.sim_matrix) 2 ctx1ab

260 

261 results = {'alignment_qc': max_sim, 2 ctx1ab

262 'alignment_count': self.sim_matrix.shape[0]} 

263 

264 if max_sim > CRITERIA['PASS']: 2 ctx1ab

265 location = np.where(self.sim_matrix == max_sim) 1 ctx1a

266 results.update({'alignment_stored': self.align_keys_sorted[np.min(location)]}) 1 ctx1a

267 results.update({'alignment_resolved': True}) 1 ctx1a

268 results.update({'alignment_resolved_by': 'qc'}) 1 ctx1a

269 

270 else: 

271 results.update({'alignment_stored': self.align_keys_sorted[0]}) 2 ctx1ab

272 results.update({'alignment_resolved': False}) 2 ctx1ab

273 

274 return results 2 ctx1ab

275 

276 def upload_channels(self, alignment_key, upload_alyx, upload_flatiron): (empty)

277 """ 

278 Upload channels to alyx and flatiron based on the alignment specified by the alignment key 

279 """ 

280 

281 feature = np.array(self.alignments[alignment_key][0]) 2 ctx1ab

282 track = np.array(self.alignments[alignment_key][1]) 2 ctx1ab

283 

284 try: 2 ctx1ab

285 meta_dset = self.one.list_datasets(self.insertion['session'], '*ap.meta', 2 ctx1ab

286 collection=f'raw_ephys_data/{self.insertion["name"]}') 

287 

288 meta_file = self.one.load_dataset(self.insertion['session'], meta_dset[0].split('/')[-1], 2 ctx1ab

289 collection=f'raw_ephys_data/{self.insertion["name"]}', 

290 download_only=True) 

291 geometry = spikeglx.read_geometry(meta_file) 

292 chns = np.c_[geometry['x'], geometry['y']] 

293 except Exception as err: 2 ctx1ab

294 self.log.warning(f"Could not compute channel locations from meta file, errored with message: {err}. " 2 ctx1ab

295 f"Will use default Neuropixel 1 channels") 

296 geometry = trace_header(version=1) 2 ctx1ab

297 chns = np.c_[geometry['x'], geometry['y']] 2 ctx1ab

298 

299 ephysalign = EphysAlignment(self.xyz_picks, chns[:, 1], 2 ctx1ab

300 track_prev=track, 

301 feature_prev=feature, 

302 brain_atlas=self.brain_atlas) 

303 channels_mlapdv = np.int32(ephysalign.get_channel_locations(feature, track) * 1e6) 2 ctx1ab

304 channels_atlas_id = ephysalign.get_brain_locations(channels_mlapdv / 1e6)['id'] 2 ctx1ab

305 

306 # Need to change channels stored on alyx as well as the stored key is not the same as the latest key 

307 if upload_alyx: 2 ctx1ab

308 if alignment_key != self.align_keys_sorted[0]: 2 ctx1ab

309 histology.register_aligned_track(self.eid, channels_mlapdv / 1e6, 2 ctx1ab

310 chn_coords=chns, one=self.one, 

311 overwrite=True, channels=self.channels_flag, 

312 brain_atlas=self.brain_atlas) 

313 

314 ephys_traj = self.one.alyx.get(f'/trajectories?&probe_insertion={self.eid}' 2 ctx1ab

315 '&provenance=Ephys aligned histology track', 

316 clobber=True) 

317 patch_dict = {'json': self.alignments} 2 ctx1ab

318 self.one.alyx.rest('trajectories', 'partial_update', id=ephys_traj[0]['id'], 2 ctx1ab

319 data=patch_dict) 

320 

321 files_to_register = [] 2 ctx1ab

322 if upload_flatiron: 2 ctx1ab

323 ftp_patcher = FTPPatcher(one=self.one) 

324 

325 alf_path = self.one.eid2path(self.insertion['session']).joinpath('alf', self.insertion["name"]) 

326 alf_path.mkdir(exist_ok=True, parents=True) 

327 

328 f_name = alf_path.joinpath('electrodeSites.mlapdv.npy') 

329 np.save(f_name, channels_mlapdv) 

330 files_to_register.append(f_name) 

331 

332 f_name = alf_path.joinpath('electrodeSites.brainLocationIds_ccf_2017.npy') 

333 np.save(f_name, channels_atlas_id) 

334 files_to_register.append(f_name) 

335 

336 f_name = alf_path.joinpath('electrodeSites.localCoordinates.npy') 

337 np.save(f_name, chns) 

338 files_to_register.append(f_name) 

339 

340 probe_collections = self.one.list_collections(self.insertion['session'], filename='channels*', 

341 collection=f'alf/{self.insertion["name"]}*') 

342 

343 for collection in probe_collections: 

344 chns = self.one.load_dataset(self.insertion['session'], 'channels.localCoordinates', collection=collection) 

345 ephysalign = EphysAlignment(self.xyz_picks, chns[:, 1], 

346 track_prev=track, 

347 feature_prev=feature, 

348 brain_atlas=self.brain_atlas) 

349 channels_mlapdv = np.int32(ephysalign.get_channel_locations(feature, track) * 1e6) 

350 channels_atlas_id = ephysalign.get_brain_locations(channels_mlapdv / 1e6)['id'] 

351 

352 alf_path = self.one.eid2path(self.insertion['session']).joinpath(collection) 

353 alf_path.mkdir(exist_ok=True, parents=True) 

354 

355 f_name = alf_path.joinpath('channels.mlapdv.npy') 

356 np.save(f_name, channels_mlapdv) 

357 files_to_register.append(f_name) 

358 

359 f_name = alf_path.joinpath('channels.brainLocationIds_ccf_2017.npy') 

360 np.save(f_name, channels_atlas_id) 

361 files_to_register.append(f_name) 

362 

363 self.log.info("Writing datasets to FlatIron") 

364 ftp_patcher.create_dataset(path=files_to_register, 

365 created_by=self.one.alyx.user) 

366 

367 return files_to_register 2 ctx1ab

368 

369 def update_experimenter_evaluation(self, prev_alignments=None, override=False): (empty)

370 

371 if not np.any(prev_alignments) and not np.any(self.alignments): 

372 aligned_traj = self.one.alyx.get(f'/trajectories?&probe_insertion={self.eid}' 

373 '&provenance=Ephys aligned histology track', 

374 clobber=True) 

375 if len(aligned_traj) > 0: 

376 self.alignments = aligned_traj[0].get('json', {}) 

377 else: 

378 self.alignments = {} 

379 return 

380 else: 

381 self.alignments = prev_alignments 

382 

383 outcomes = [align[2].split(':')[0] for key, align in self.alignments.items() 

384 if len(align) == 3] 

385 if len(outcomes) > 0: 

386 vals = [CRITERIA_BASE[out] for out in outcomes] 

387 max_qc = np.argmax(vals) 

388 outcome = outcomes[max_qc] 

389 self.update(outcome, namespace='experimenter', override=override) 

390 else: 

391 self.log.warning(f'No experimenter qc found, qc field of probe insertion {self.eid} ' 

392 f'will not be updated') 

393 

394 

395def get_aligned_channels(ins, chn_coords, one, ba=None, save_dir=None): (empty)

396 

397 ba = ba or AllenAtlas(25) 

398 depths = chn_coords[:, 1] 

399 xyz = np.array(ins['json']['xyz_picks']) / 1e6 

400 traj = one.alyx.rest('trajectories', 'list', probe_insertion=ins['id'], 

401 provenance='Ephys aligned histology track')[0] 

402 align_key = ins['json']['extended_qc']['alignment_stored'] 

403 feature = traj['json'][align_key][0] 

404 track = traj['json'][align_key][1] 

405 ephysalign = EphysAlignment(xyz, depths, track_prev=track, 

406 feature_prev=feature, 

407 brain_atlas=ba, speedy=True) 

408 channels_mlapdv = np.int32(ephysalign.get_channel_locations(feature, track) * 1e6) 

409 channels_atlas_ids = ephysalign.get_brain_locations(channels_mlapdv / 1e6)['id'] 

410 

411 out_files = [] 

412 if save_dir is not None: 

413 f_name = Path(save_dir).joinpath('channels.mlapdv.npy') 

414 np.save(f_name, channels_mlapdv) 

415 out_files.append(f_name) 

416 

417 f_name = Path(save_dir).joinpath('channels.brainLocationIds_ccf_2017.npy') 

418 np.save(f_name, channels_atlas_ids) 

419 out_files.append(f_name) 

420 

421 return out_files