Repository: spikeinterface/spikeinterface
Issue: not fitted error from running export_to_phy with compute_pc_features=True

Hi,

I tried to use `export_to_phy` with `compute_pc_features=True`, but I got `nofittederror`.
Please find the complete error message below.
Have anyone had this issue and how could we solve it? Thanks!

```
---------------------------------------------------------------------------                                                                                                                                                                   
_RemoteTraceback                          Traceback (most recent call last)                                                                                                                                                                   
_RemoteTraceback:                                                                                                                                                                                                                             
"""                                                                                                                                                                                                                                           
Traceback (most recent call last):                                                                                                                                                                                                            
  File "/home/amz/.conda/envs/si_tuto/lib/python3.9/concurrent/futures/process.py", line 246, in _process_worker                                                                                                                              
    r = call_item.fn(*call_item.args, **call_item.kwargs)                                                                                                                                                                                     
  File "/home/amz/.conda/envs/si_tuto/lib/python3.9/concurrent/futures/process.py", line 205, in _process_chunk                                                                                                                               
    return [fn(*args) for args in chunk]                                                                                                                                                                                                      
  File "/home/amz/.conda/envs/si_tuto/lib/python3.9/concurrent/futures/process.py", line 205, in <listcomp>                                                                                                                                   
    return [fn(*args) for args in chunk]                                                                                                                                                                                                      
  File "/home/amz/.conda/envs/si_tuto/lib/python3.9/site-packages/spikeinterface/core/job_tools.py", line 337, in function_wrapper                                                                                                            
    return _func(segment_index, start_frame, end_frame, _worker_ctx)                                                                                                                                                                          
  File "/home/amz/.conda/envs/si_tuto/lib/python3.9/site-packages/spikeinterface/postprocessing/principal_component.py", line 591, in _all_pc_extractor_chunk                                                                                 
    all_pcs[i, :, c] = pca_model[chan_ind].transform(w)                                                                                                                                                                                       
  File "/home/amz/.conda/envs/si_tuto/lib/python3.9/site-packages/sklearn/decomposition/_incremental_pca.py", line 392, in transform                                                                                                          
    return super().transform(X)                                                                                                                                                                                                               
  File "/home/amz/.conda/envs/si_tuto/lib/python3.9/site-packages/sklearn/decomposition/_base.py", line 118, in transform                                                                                                                     
    check_is_fitted(self)                                                                                                                                                                                                                     
  File "/home/amz/.conda/envs/si_tuto/lib/python3.9/site-packages/sklearn/utils/validation.py", line 1345, in check_is_fitted                                                                                                                 
    raise NotFittedError(msg % {"name": type(estimator).__name__})                                                                                                                                                                            
sklearn.exceptions.NotFittedError: This IncrementalPCA instance is not fitted yet. Call 'fit' with appropriate arguments before using this estimator.                                                                                         
"""                                                                                                                                                                                                                                           
                                                                                                                                                                                                                                              
The above exception was the direct cause of the following exception:                                                                                                                                                                          
                                                                                                                                                                                                                                              
NotFittedError                            Traceback (most recent call last)       
File ~/Documents/git/spiketutorials/Official_Tutorial_SI_0.95_Oct2022/SpikeInterface_Tutorial.py:106                                                                                                                                          
    102 print(we_curated)                                                                                                                                                                                                                     
    103 # In[5]:                                                                                                                                                                                                                              
    104                                                                                                                                                                                                                                       
    105 #we.fit()                                                                                                                                                                                                                             
--> 106 sexp.export_to_phy(we, output_folder=base_folder / 'phy_KS3',                                                                                                                                                                         
    107                    compute_amplitudes=False, copy_binary=True, remove_if_exists = True,                                                                                                                                               
    108                    **job_kwargs)                                                                                                                                                                                                      
    110 assert 0                                                                                                                                                                                                                              
    111 """                                                                                                                                                                                                                                   
    112 recording_nwb = se.read_nwb_recording(file_path=nwb_file_path)                                                                                                                                                                        
    113 sorting_nwb = se.read_nwb_sorting(file_path=nwb_file_path)                                                                                                                                                                            
   (...)                                                   
    369 # In[36]:                                          
    370 """                                                

File ~/.conda/envs/si_tuto/lib/python3.9/site-packages/spikeinterface/exporters/to_phy.py:197, in export_to_phy(waveform_extractor, output_folder, compute_pc_features, compute_amplitudes, sparsity_dict, copy_binary, max_channels_per_templ
ate, remove_if_exists, peak_sign, template_mode, dtype, verbose, **job_kwargs)                                         
    194     pc = compute_principal_components(waveform_extractor, n_components=5, mode='by_channel_local')             
    196 max_channels_per_template = min(max_channels_per_template, len(channel_ids))                                   
--> 197 pc.run_for_all_spikes(output_folder / 'pc_features.npy',                                                       
    198                       max_channels_per_template=max_channels_per_template, peak_sign=peak_sign,                
    199                       **job_kwargs)
    201 pc_feature_ind = np.zeros((len(unit_ids), max_channels_per_template), dtype='int64')                           
    202 best_channels_index = get_template_channel_sparsity(waveform_extractor, method='best_channels',                
    203                                                     peak_sign=peak_sign, num_channels=max_channels_per_template,
    204                                                     outputs='index')                                           

File ~/.conda/envs/si_tuto/lib/python3.9/site-packages/spikeinterface/postprocessing/principal_component.py:320, in WaveformPrincipalComponent.run_for_all_spikes(self, file_path, max_channels_per_template, peak_sign, **job_kwargs)
    317 init_args = init_args + (all_pcs_args, spike_times, spike_labels, we.nbefore, we.nafter,                       
    318                          unit_channels, pca_model)
    319 processor = ChunkRecordingExecutor(recording, func, init_func, init_args, job_name='extract PCs', **job_kwargs) 
--> 320 processor.run()                                    

File ~/.conda/envs/si_tuto/lib/python3.9/site-packages/spikeinterface/core/job_tools.py:312, in ChunkRecordingExecutor.run(self)
    310                 returns.append(res)
    311         else:                                      
--> 312             for res in results:                                                                                                                                                                                                       
    313                 pass                                                                                                                                                                                                                  
    315 return returns                                     

File ~/.conda/envs/si_tuto/lib/python3.9/site-packages/tqdm/std.py:1195, in tqdm.__iter__(self)                        
   1192 time = self._time                                  
   1194 try:                                               
-> 1195     for obj in iterable:
   1196         yield obj                                  
   1197         # Update and possibly print the progressbar.                                                           
   1198         # Note: does not call self.update(1) for speed optimisation.                                           

File ~/.conda/envs/si_tuto/lib/python3.9/concurrent/futures/process.py:562, in _chain_from_iterable_of_lists(iterable) 
    556 def _chain_from_iterable_of_lists(iterable):
    557     """                                            
    558     Specialized implementation of itertools.chain.from_iterable.                                               
    559     Each item in *iterable* should be a list.  This function is                                                
    560     careful not to keep references to yielded objects.                                                         
    561     """                                            
--> 562     for element in iterable:
    563         element.reverse()
    564         while element:

File ~/.conda/envs/si_tuto/lib/python3.9/concurrent/futures/_base.py:609, in Executor.map.<locals>.result_iterator()   
    606 while fs:                                          
    607     # Careful not to keep a reference to the popped future                                                     
    608     if timeout is None:
--> 609         yield fs.pop().result()
    610     else:                                          
    611         yield fs.pop().result(end_time - time.monotonic())                                                     

File ~/.conda/envs/si_tuto/lib/python3.9/concurrent/futures/_base.py:446, in Future.result(self, timeout)              
    444     raise CancelledError()
    445 elif self._state == FINISHED:
--> 446     return self.__get_result()
    447 else:                                              
    448     raise TimeoutError()

File ~/.conda/envs/si_tuto/lib/python3.9/concurrent/futures/_base.py:391, in Future.__get_result(self)                 
    389 if self._exception:                                
    390     try:                                           
--> 391         raise self._exception
    392     finally:                                       
    393         # Break a reference cycle with the exception in self._exception                                        
    394         self = None                                

NotFittedError: This IncrementalPCA instance is not fitted yet. Call 'fit' with appropriate arguments before using this estimator.
> /home/amz/.conda/envs/si_tuto/lib/python3.9/concurrent/futures/_base.py(391)__get_result()                           
    389         if self._exception:
    390             try:          
--> 391                 raise self._exception
    392             finally:
    393                 # Break a reference cycle with the exception in self._exception

```



Changed files:
spikeinterface/exporters/to_phy.py, spikeinterface/postprocessing/principal_component.py, spikeinterface/exporters/tests/test_export_to_phy.py, spikeinterface/postprocessing/tests/test_principal_component.py

diff --git a/spikeinterface/exporters/tests/test_export_to_phy.py b/spikeinterface/exporters/tests/test_export_to_phy.py
index 45a68d27..9a97ecfd 100644
--- a/spikeinterface/exporters/tests/test_export_to_phy.py
+++ b/spikeinterface/exporters/tests/test_export_to_phy.py
@@ -7,6 +7,7 @@ import numpy as np
 from spikeinterface import extract_waveforms, download_dataset
 import spikeinterface.extractors as se
 from spikeinterface.exporters import export_to_phy
+from spikeinterface.postprocessing import get_template_channel_sparsity, compute_principal_components
 
 if hasattr(pytest, "global_test_folder"):
     cache_folder = pytest.global_test_folder / "exporters"
@@ -33,8 +34,8 @@ def test_export_to_phy():
     export_to_phy(waveform_extractor, output_folder,
                   compute_pc_features=True,
                   compute_amplitudes=True,
-                  max_channels_per_template=8,
-                  n_jobs=1, chunk_size=10000, progress_bar=True)
+                  n_jobs=1, chunk_size=10000,
+                  progress_bar=True)
 
 
 def test_export_to_phy_by_property():
@@ -58,12 +59,11 @@ def test_export_to_phy_by_property():
     sorting = sorting.save(folder=sort_folder)
 
     waveform_extractor = extract_waveforms(recording, sorting, waveform_folder)
-
+    sparsity_group = get_template_channel_sparsity(waveform_extractor, method="by_property", by_property="group")
     export_to_phy(waveform_extractor, output_folder,
                   compute_pc_features=True,
                   compute_amplitudes=True,
-                  max_channels_per_template=8,
-                  sparsity_dict=dict(method="by_property", by_property="group"),
+                  sparsity=sparsity_group,
                   n_jobs=1, chunk_size=10000, progress_bar=True)
 
     template_inds = np.load(output_folder / "template_ind.npy")
@@ -72,13 +72,14 @@ def test_export_to_phy_by_property():
     # Remove one channel
     recording_rm = recording.channel_slice([0, 2, 3, 4, 5, 6, 7])
     waveform_extractor_rm = extract_waveforms(recording_rm, sorting, waveform_folder_rm)
+    sparsity_group = get_template_channel_sparsity(waveform_extractor_rm, method="by_property", by_property="group")
 
     export_to_phy(waveform_extractor_rm, output_folder_rm,
                   compute_pc_features=True,
                   compute_amplitudes=True,
-                  max_channels_per_template=8,
-                  sparsity_dict=dict(method="by_property", by_property="group"),
-                  n_jobs=1, chunk_size=10000, progress_bar=True)
+                  sparsity=sparsity_group,
+                  n_jobs=1, chunk_size=10000,
+                  progress_bar=True)
 
     template_inds = np.load(output_folder_rm / "template_ind.npy")
     assert template_inds.shape == (num_units, 4)
@@ -94,35 +95,43 @@ def test_export_to_phy_by_sparsity():
 
     waveform_folder = cache_folder / 'waveforms'
     output_folder_radius = cache_folder / 'phy_output_radius'
-    output_folder_thr = cache_folder / 'phy_output_thr'
+    output_folder_multi_sparse = cache_folder / 'phy_output_multi_sparse'
 
-    for f in (waveform_folder, output_folder_radius, output_folder_thr):
+    for f in (waveform_folder, output_folder_radius, output_folder_multi_sparse):
         if f.is_dir():
             shutil.rmtree(f)
 
     waveform_extractor = extract_waveforms(recording, sorting, waveform_folder)
-
+    sparsity_radius = get_template_channel_sparsity(waveform_extractor, method="radius", radius_um=50)
     export_to_phy(waveform_extractor, output_folder_radius,
                   compute_pc_features=True,
                   compute_amplitudes=True,
-                  max_channels_per_template=None,
-                  sparsity_dict=dict(method="radius", radius_um=50),
+                  sparsity=sparsity_radius,
                   n_jobs=1, chunk_size=10000, progress_bar=True)
 
     template_ind = np.load(output_folder_radius / "template_ind.npy")
+    pc_ind = np.load(output_folder_radius / "pc_feature_ind.npy")
     # templates have different shapes!
     assert -1 in template_ind
+    assert -1 in pc_ind
 
-    export_to_phy(waveform_extractor, output_folder_thr,
+    # pre-compute PC with sparsity
+    sparsity_radius_small = get_template_channel_sparsity(waveform_extractor, method="radius", radius_um=30)
+    pc = compute_principal_components(waveform_extractor, sparsity=sparsity_radius_small)
+    export_to_phy(waveform_extractor, output_folder_multi_sparse,
                   compute_pc_features=True,
                   compute_amplitudes=True,
-                  max_channels_per_template=None,
-                  sparsity_dict=dict(method="threshold", threshold=2),
-                  n_jobs=1, chunk_size=10000, progress_bar=True)
+                  sparsity=sparsity_radius,
+                  n_jobs=1, chunk_size=10000,
+                  progress_bar=True)
 
-    template_ind = np.load(output_folder_thr / "template_ind.npy")
+    template_ind = np.load(output_folder_multi_sparse / "template_ind.npy")
+    pc_ind = np.load(output_folder_multi_sparse / "pc_feature_ind.npy")
     # templates have different shapes!
     assert -1 in template_ind
+    assert -1 in pc_ind
+    # PC sparsity is more stringent than teplate sparsity
+    assert pc_ind.shape[1] < template_ind.shape[1]
 
 
 if __name__ == '__main__':
diff --git a/spikeinterface/exporters/to_phy.py b/spikeinterface/exporters/to_phy.py
index 8b5c62bf..8873d2f2 100644
--- a/spikeinterface/exporters/to_phy.py
+++ b/spikeinterface/exporters/to_phy.py
@@ -14,9 +14,8 @@ from spikeinterface.postprocessing import (get_template_channel_sparsity,
 
 
 def export_to_phy(waveform_extractor, output_folder, compute_pc_features=True,
-                  compute_amplitudes=True, sparsity_dict=None, copy_binary=True,
-                  max_channels_per_template=16, remove_if_exists=False,
-                  peak_sign='neg', template_mode='median',
+                  compute_amplitudes=True, sparsity=None, copy_binary=True,
+                  remove_if_exists=False, peak_sign='neg', template_mode='median',
                   dtype=None, verbose=True, **job_kwargs):
     """
     Exports a waveform extractor to the phy template-gui format.
@@ -31,17 +30,8 @@ def export_to_phy(waveform_extractor, output_folder, compute_pc_features=True,
         If True (default), pc features are computed
     compute_amplitudes: bool
         If True (default), waveforms amplitudes are computed
-    sparsity_dict: dict or None
-        If given, the dictionary should contain a sparsity method (e.g. "best_channels") and optionally
-        arguments associated with the method (e.g. "num_channels" for "best_channels" method).
-        Other examples are:
-           * by radius: sparsity_dict=dict(method="radius", radius_um=100)
-           * by SNR threshold: sparsity_dict=dict(method="threshold", threshold=2)
-           * by property: sparsity_dict=dict(method="by_property", by_property="group")
-        Default is sparsity_dict=dict(method="best_channels", num_channels=16)
-        For more info, see the postprocessing.get_template_channel_sparsity() function.
-    max_channels_per_template: int or None
-        Maximum channels per unit to return. If None, all channels are returned
+    sparsity: dict or None
+        The sparsity dictionary (key: unit_id, value: list of channel_ids)
     copy_binary: bool
         If True, the recording is copied and saved in the phy 'output_folder'
     remove_if_exists: bool
@@ -59,23 +49,24 @@ def export_to_phy(waveform_extractor, output_folder, compute_pc_features=True,
     """
     assert isinstance(waveform_extractor, spikeinterface.core.waveform_extractor.WaveformExtractor), \\
         'waveform_extractor must be a WaveformExtractor object'
-    recording = waveform_extractor.recording
     sorting = waveform_extractor.sorting
+    assert waveform_extractor.has_recording(), "Export to phy is not supported for 'recordingless' waveform extractors"
+    recording = waveform_extractor.recording
 
-    assert recording.get_num_segments() == sorting.get_num_segments(), \\
-        "The recording and sorting objects must have the same number of segments!"
-
-    assert recording.get_num_segments() == 1, "Export to phy work only with one segment"
-
-    if sparsity_dict is None:
-        sparsity_dict = dict(method="best_channels", num_channels=16)
-
+    assert waveform_extractor.get_num_segments() == 1, "Export to phy only works with one segment"
     channel_ids = recording.channel_ids
     num_chans = recording.get_num_channels()
-    fs = recording.get_sampling_frequency()
-
-    if max_channels_per_template is None:
-        max_channels_per_template = num_chans
+    fs = recording.sampling_frequency
+
+    # handle sparsity
+    if sparsity is None:
+        if num_chans > 64:
+            print(
+                "WARNING: export to Phy with many channels and without sparsity might result in a heavy and less "
+                "informative visualization. You can use the 'sparsity' argument to enforce sparsity "
+                "(see spikeinterface.postprocessing.get_template_channel_sparsity() function for sparsity details.)"
+            )
+        sparsity = {unit_id: channel_ids for unit_id in sorting.unit_ids}
 
     empty_flag = False
     non_empty_units = []
@@ -137,22 +128,20 @@ def export_to_phy(waveform_extractor, output_folder, compute_pc_features=True,
     np.save(str(output_folder / 'spike_clusters.npy'), spike_labels[:, np.newaxis])
 
     # export templates/templates_ind/similar_templates
-    # shape (num_units, num_samples, num_channels)
+    # shape (num_units, num_samples, max_num_channels)
     templates = []
     templates_ind = []
 
-    template_sparsity = get_template_channel_sparsity(waveform_extractor,
-                                                      outputs="id", **sparsity_dict)
-    num_sparse_chans = np.max([len(channels) for channels in template_sparsity.values()])
-    num_channels = np.min([max_channels_per_template, num_sparse_chans])
+    # here we pad template inds with -1 if len of sparse channels is unequal
+    max_num_channels_templates = np.max([len(channels) for channels in sparsity.values()])
     for unit_id in unit_ids:
-        template = waveform_extractor.get_template(unit_id, mode=template_mode, sparsity=template_sparsity)
-        inds = waveform_extractor.recording.ids_to_indices(template_sparsity[unit_id])
-        if template.shape[-1] < num_channels:
+        template = waveform_extractor.get_template(unit_id, mode=template_mode, sparsity=sparsity)
+        inds = waveform_extractor.recording.ids_to_indices(sparsity[unit_id])
+        if template.shape[-1] < max_num_channels_templates:
             # fill missing channels
-            template_full = np.zeros((template.shape[0], num_channels))
+            template_full = np.zeros((template.shape[0], max_num_channels_templates))
             template_full[:, :template.shape[-1]] = template
-            inds_full = np.concatenate((inds, np.array([-1] * (num_channels - template.shape[-1]))))
+            inds_full = np.concatenate((inds, np.array([-1] * (max_num_channels_templates - template.shape[-1]))))
         else:
             template_full = template
             inds_full = inds
@@ -192,18 +181,24 @@ def export_to_phy(waveform_extractor, output_folder, compute_pc_features=True,
             pc = waveform_extractor.load_extension('principal_components')
         else:
             pc = compute_principal_components(waveform_extractor, n_components=5, mode='by_channel_local')
+
+        if pc.get_sparsity() is None:
+            pc_sparsity = sparsity
+            pc_feature_ind = templates_ind
+        else:
+            pc_sparsity = pc.get_sparsity()
+            pc_feature_ind = None
         
-        max_channels_per_template = min(max_channels_per_template, len(channel_ids))
+        max_num_channels_pc = np.max([len(channels) for channels in pc_sparsity.values()])
         pc.run_for_all_spikes(output_folder / 'pc_features.npy',
-                              max_channels_per_template=max_channels_per_template, peak_sign=peak_sign,
+                              sparsity=pc_sparsity, peak_sign=peak_sign,
                               **job_kwargs)
 
-        pc_feature_ind = np.zeros((len(unit_ids), max_channels_per_template), dtype='int64')
-        best_channels_index = get_template_channel_sparsity(waveform_extractor, method='best_channels',
-                                                            peak_sign=peak_sign, num_channels=max_channels_per_template,
-                                                            outputs='index')
-        for u, unit_id in enumerate(unit_ids):
-            pc_feature_ind[u, :] = best_channels_index[unit_id]
+        if pc_feature_ind is None:
+            pc_feature_ind = -np.ones((len(unit_ids), max_num_channels_pc), dtype='int64')
+            for u, unit_id in enumerate(unit_ids):
+                sparse_channel_indices = recording.ids_to_indices(pc_sparsity[unit_id])
+                pc_feature_ind[u, :len(sparse_channel_indices)] = sparse_channel_indices
         np.save(str(output_folder / 'pc_feature_ind.npy'), pc_feature_ind)
 
     # Save .tsv metadata
diff --git a/spikeinterface/postprocessing/principal_component.py b/spikeinterface/postprocessing/principal_component.py
index 2943634b..bc549c8d 100644
--- a/spikeinterface/postprocessing/principal_component.py
+++ b/spikeinterface/postprocessing/principal_component.py
@@ -198,7 +198,17 @@ class WaveformPrincipalComponent(BaseWaveformExtractorExtension):
 
         return projections
 
-    #TODO
+    def get_sparsity(self):
+        """
+        Returns sparsity used in PC computation
+
+        Returns
+        -------
+        dict
+            The sparsity dictionary (key: unit_id, value: list of channel_ids)
+        """
+        return self._params["sparsity"]
+
     def _run(self, n_jobs=1, progress_bar=False):
         """
         Compute the PCs on waveforms extacted within the WaveformExtarctor.
@@ -255,8 +265,8 @@ class WaveformPrincipalComponent(BaseWaveformExtractorExtension):
     def get_extension_function():
         return compute_principal_components
 
-    def run_for_all_spikes(self, file_path, max_channels_per_template=16, peak_sign='neg',
-                           **job_kwargs):
+    def run_for_all_spikes(self, file_path=None, peak_sign='neg',
+                           sparsity=None, **job_kwargs):
         """
         Project all spikes from the sorting on the PCA model.
         This is a long computation because waveform need to be extracted from each spikes.
@@ -267,12 +277,19 @@ class WaveformPrincipalComponent(BaseWaveformExtractorExtension):
 
         Parameters
         ----------
-        file_path : str or Path
-            Path to npy file that will store the PCA projections
-        max_channels_per_template : int, optionl
-            Maximum number of best channels to compute PCA projections on
+        file_path : str or Path or None
+            Path to npy file that will store the PCA projections.
+            If None, output is saved in principal_components/all_pcs.npy
         peak_sign : str, optional
             Peak sign to get best channels ('neg', 'pos', 'both'), by default 'neg'
+        sparsity : dict, optional
+            The sparsity to use for PC projection computation (by default None):
+            * If None and the principal_component extension has been computed with sparsity, the same
+            sparsity is used. 
+            * If given and sparsity has not been used to compute principal components,
+            the provided sparsity is used (dictionary with unit_id as key and list of channel_ids as values).
+            * If None and principal components are not sparse, no sparsity is applied and PC projections are 
+            computed on all channels
         {}
         """
         p = self._params
@@ -286,24 +303,39 @@ class WaveformPrincipalComponent(BaseWaveformExtractorExtension):
         assert sorting.get_num_segments() == 1
         assert p['mode'] in ('by_channel_local', 'by_channel_global')
 
+        if file_path is None:
+            file_path = self.extension_folder / "all_pcs.npy"
         file_path = Path(file_path)
 
         all_spikes = sorting.get_all_spike_trains(outputs='unit_index')
         spike_times, spike_labels = all_spikes[0]
 
-        max_channels_per_template = min(max_channels_per_template, we.get_num_channels())
+        if sparsity is not None:
+            if p["sparsity"] is not None:
+                for unit_id, sparse_channels in sparsity.items():
+                    assert np.all(sparse_channels == p["sparsity"][unit_id]), \\
+                        ("The provided sparsity does not match the sparsity used to estimate PC models. "
+                         "Use sparsity=None to use correct sparsity")
+        else:
+            sparsity = p["sparsity"]
 
-        best_channels_index = get_template_channel_sparsity(we, outputs="index", peak_sign=peak_sign,
-                                                            num_channels=max_channels_per_template)
+        if sparsity is None:
+            sparse_channels_indices = {unit_id: np.arange(we.get_num_channels()) for unit_id in we.unit_ids}
+            max_channels_per_template = we.get_num_channels()
+        else:
+            sparse_channels_indices = {unit_id: self.waveform_extractor.channel_ids_to_indices(sparsity[unit_id])
+                                       for unit_id in sorting.unit_ids}
+            max_sparse_channels = np.max([len(best_channels) for best_channels in sparse_channels_indices.values()])
+            max_channels_per_template = max_sparse_channels
 
-        unit_channels = [best_channels_index[unit_id] for unit_id in sorting.unit_ids]
+        unit_channels = [sparse_channels_indices[unit_id] for unit_id in sorting.unit_ids]
 
         pca_model = self.get_pca_model()
         if p['mode'] in ['by_channel_global', 'concatenated']:
             pca_model = [pca_model] * recording.get_num_channels()
 
         # nSpikes, nFeaturesPerChannel, nPCFeatures
-        # this come from  phy template-gui
+        # this comes from  phy template-gui
         # https://github.com/kwikteam/phy-contrib/blob/master/docs/template-gui.md#datasets
         shape = (spike_times.size, p['n_components'], max_channels_per_template)
         all_pcs = np.lib.format.open_memmap(filename=file_path, mode='w+', dtype='float32', shape=shape)
diff --git a/spikeinterface/postprocessing/tests/test_principal_component.py b/spikeinterface/postprocessing/tests/test_principal_component.py
index b6f866ec..5943e805 100644
--- a/spikeinterface/postprocessing/tests/test_principal_component.py
+++ b/spikeinterface/postprocessing/tests/test_principal_component.py
@@ -1,5 +1,5 @@
 import unittest
-import shutil
+import pytest
 from pathlib import Path
 
 import numpy as np
@@ -11,6 +11,11 @@ from spikeinterface.postprocessing import (WaveformPrincipalComponent, compute_p
                                            get_template_channel_sparsity)
 from spikeinterface.postprocessing.tests.common_extension_tests import WaveformExtensionCommonTestSuite
 
+if hasattr(pytest, "global_test_folder"):
+    cache_folder = pytest.global_test_folder / "postprocessing"
+else:
+    cache_folder = Path("cache_folder") / "postprocessing"
+
 
 DEBUG = False
 
@@ -48,16 +53,31 @@ class PrincipalComponentsExtensionTest(WaveformExtensionCommonTestSuite, unittes
 
         pc_file1 = pc.extension_folder / 'all_pc1.npy'
         pc.run_for_all_spikes(
-            pc_file1, max_channels_per_template=7, chunk_size=10000, n_jobs=1)
+            pc_file1, sparsity=None, chunk_size=10000, n_jobs=1)
         all_pc1 = np.load(pc_file1)
 
         pc_file2 = pc.extension_folder / 'all_pc2.npy'
         pc.run_for_all_spikes(
-            pc_file2, max_channels_per_template=7, chunk_size=10000, n_jobs=2)
+            pc_file2, sparsity=None, chunk_size=10000, n_jobs=2)
         all_pc2 = np.load(pc_file2)
 
         assert np.array_equal(all_pc1, all_pc2)
 
+        # test with sparsity
+        sparsity = get_template_channel_sparsity(we, method="radius",
+                                                 radius_um=50)
+        we_copy = we.save(folder=cache_folder / "we_copy")
+        pc_sparse = self.extension_class.get_extension_function()(we_copy, sparsity=sparsity, load_if_exists=False)
+        pc_file_sparse = pc.extension_folder / 'all_pc_sparse.npy'
+        pc_sparse.run_for_all_spikes(pc_file_sparse, sparsity=sparsity, chunk_size=10000, n_jobs=1)
+        all_pc_sparse = np.load(pc_file_sparse)
+        all_spikes = we_copy.sorting.get_all_spike_trains(outputs='unit_id')
+        _, spike_labels = all_spikes[0]
+        for unit_id, sparse_channel_ids in sparsity.items():
+            # check dimensions
+            pc_unit = all_pc_sparse[spike_labels == unit_id]
+            assert np.allclose(pc_unit[:, :, len(sparse_channel_ids):], 0)
+
     def test_sparse(self):
         we = self.we2
         unit_ids = we.unit_ids