From fad60f6dbb6760a211e8dada3124f8eebaaa1f7d Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Wed, 15 May 2024 10:09:14 -0600 Subject: [PATCH] add method to load into memory --- src/spikeinterface/core/template.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/src/spikeinterface/core/template.py b/src/spikeinterface/core/template.py index 4eb82be2d6..644f4b29a8 100644 --- a/src/spikeinterface/core/template.py +++ b/src/spikeinterface/core/template.py @@ -188,6 +188,19 @@ def _are_passed_templates_sparse(self) -> bool: return are_templates_sparse + def get_channel_locations(self) -> np.ndarray: + assert self.probe is not None, "Templates.get_channel_locations() needs a set probe to be defined." + channel_locations = self.probe.contact_positions + return channel_locations + + def load_all_arrays_into_memory(self): + "This is useful to avoid making requests to the disk when the templates are stored in disk" + self.templates_array = np.asarray(self.templates_array) + if self.sparsity_mask is not None: + self.sparsity_mask = np.asarray(self.sparsity_mask) + self.channel_ids = np.asarray(self.channel_ids) + self.unit_ids = np.asarray(self.unit_ids.copy()) + def to_dict(self): return { "templates_array": self.templates_array, @@ -385,8 +398,3 @@ def __eq__(self, other): return False return True - - def get_channel_locations(self): - assert self.probe is not None, "Templates.get_channel_locations() needs a probe to be set" - channel_locations = self.probe.contact_positions - return channel_locations