addresses batch size bug and also resolves open file warnings
@@ -115,18 +115,21 @@ class WaveformPrincipalComponent(BaseWaveformExtractorExtension):
pca_file = self.extension_folder / f"pca_model_{mode}_{chan_id}.pkl"
if not pca_file.is_file() and chan_ind == 0:
_ = self._fit_by_channel_local()
- pca = pickle.load(pca_file.open("rb"))
+ with open(pca_file, 'rb') as fid:
+ pca = pickle.load(fid)
pca_model.append(pca)
elif mode == "by_channel_global":
pca_file = self.extension_folder / f"pca_model_{mode}.pkl"
if not pca_file.is_file():
_ = self._fit_by_channel_global()
- pca_model = pickle.load(pca_file.open("rb"))
+ with open(pca_file, 'rb') as fid:
+ pca_model = pickle.load(fid)
elif mode == "concatenated":
pca_file = self.extension_folder / f"pca_model_{mode}.pkl"
if not pca_file.is_file():
_ = self._fit_concatenated()
- pca_model = pickle.load(pca_file.open("rb"))
+ with open(pca_file, 'rb') as fid:
+ pca_model = pickle.load(fid)
self._pca_model = pca_model
def get_pca_model(self):
@@ -354,7 +357,7 @@ class WaveformPrincipalComponent(BaseWaveformExtractorExtension):
# fit
for unit_id in unit_ids:
wfs = we.get_waveforms(unit_id)
- if wfs.size == 0:
+ if len(wfs) < p['n_components']:
continue
for chan_ind, chan_id in enumerate(channel_ids):
pca = pca_model[chan_ind]
