Merge branch 'master' of github.com:rossant/pykilosort

* 'master' of github.com:rossant/pykilosort:
  Typed Configuration (#28)
  Comparison between MATLAB & Python Versions (#25)
  Add dockerfile for running tests with GPU (#22)

Author: alexmorley

pykilosort/default_params.py

@@ -17,7 +17,7 @@
 from phy.gui.qt import QSlider, Qt, QLabel, QScrollArea, QVBoxLayout, QWidget
 from phy.gui import Actions
 
-from .default_params import default_params
+from .params import KilosortParams
 from .main import run
 
 logger = logging.getLogger(__name__)
@@ -168,7 +168,7 @@ def create_probe_view(self, gui):
     def create_params_widget(self, gui):
         """Create the widget that allows to enter parameters for KS2."""
         widget = KeyValueWidget(gui)
-        for name, default in default_params.items():
+        for name, default in KilosortParams().dict().items():
             # HACK: None default params in KS2 are floats
             vtype = 'float' if default is None else None
             widget.add_pair(name, default, vtype=vtype)

pykilosort/gui.py

@@ -3,15 +3,16 @@
 from pathlib import Path
 from phylib.io.traces import get_ephys_reader
 
-from pprint import pprint
 import numpy as np
+from pprint import pprint
+from pydantic import BaseModel
 
 from .preprocess import preprocess, get_good_channels, get_whitening_matrix, get_Nbatch
 from .cluster import clusterSingleBatches
 from .learn import learnAndSolve8b
 from .postprocess import find_merges, splitAllClusters, set_cutoff, rezToPhy
 from .utils import Bunch, Context, memmap_large_array, load_probe
-from .default_params import default_params, set_dependent_params
+from .params import KilosortParams
 
 logger = logging.getLogger(__name__)
 
@@ -61,10 +62,8 @@ def run(
     assert probe
 
     # Get params.
-    user_params = params or {}
-    params = default_params.copy()
-    set_dependent_params(params)
-    params.update(user_params)
+    if not isinstance(params, BaseModel):
+        params = KilosortParams(**params or {})
     assert params
 
     # dir path

pykilosort/main.py

@@ -0,0 +1,88 @@
+import typing as t
+from math import ceil
+
+from pydantic import BaseModel, Field
+
+from .utils import Bunch
+# TODO: design - Let's move all of this to a yaml file with sections so that its easier to read.
+#              - We can then just parse the yaml file to generate this.
+
+
+class KilosortParams(BaseModel):
+    fs: float = Field(30000., description="sample rate")
+
+    fshigh: float = Field(150., description="high pass filter frequency")
+    fslow: t.Optional[float] = Field(None, description="low pass filter frequency")
+    minfr_goodchannels: float = Field(0.1, description="minimum firing rate on a 'good' channel (0 to skip)")
+
+    Th: t.List[float] = Field([10, 4], description="""
+        threshold on projections (like in Kilosort1, can be different for last pass like [10 4])
+    """)
+    ThPre: float = Field(8, description="threshold crossings for pre-clustering (in PCA projection space)")
+
+    lam: float = Field(10, description="""
+        how important is the amplitude penalty (like in Kilosort1, 0 means not used,
+        10 is average, 50 is a lot)
+    """)
+
+    AUCsplit: float = Field(0.9, description="""
+        splitting a cluster at the end requires at least this much isolation for each sub-cluster (max=1)
+    """)
+
+    minFR: float = Field(1. / 50, description="""
+        minimum spike rate (Hz), if a cluster falls below this for too long it gets removed
+    """)
+
+    momentum: t.List[float] = Field([20, 400], description="""
+        number of samples to average over (annealed from first to second value)
+    """)
+
+    sigmaMask: float = Field(30, description="""
+        spatial constant in um for computing residual variance of spike
+    """)
+
+    # danger, changing these settings can lead to fatal errors
+    # options for determining PCs
+    spkTh: float = Field(-6, description="spike threshold in standard deviations")
+    reorder: int = Field(1, description="whether to reorder batches for drift correction.")
+    nskip: int = Field(5, description="how many batches to skip for determining spike PCs")
+    nSkipCov: int = Field(25, description="compute whitening matrix from every nth batch")
+
+    # GPU = 1  # has to be 1, no CPU version yet, sorry
+    # Nfilt = 1024 # max number of clusters
+    nfilt_factor: int = Field(4, description="max number of clusters per good channel (even temporary ones)")
+    ntbuff = Field(64, description="""
+    samples of symmetrical buffer for whitening and spike detection
+    
+    Must be multiple of 32 + ntbuff. This is the batch size (try decreasing if out of memory).
+    """)
+
+    whiteningRange: int = Field(32, description="number of channels to use for whitening each channel")
+    nSkipCov: int = Field(25, description="compute whitening matrix from every N-th batch")
+    scaleproc: int = Field(200, description="int16 scaling of whitened data")
+    nPCs: int = Field(3, description="how many PCs to project the spikes into")
+
+    nt0: int = 61
+    nup: int = 10
+    sig: int = 1
+    gain: int = 1
+
+    templateScaling: float = 20.0
+
+    loc_range: t.List[int] = [5, 4]
+    long_range: t.List[int] = [30, 6]
+
+    Nfilt: t.Optional[int] = None # This should be a computed property once we add the probe to the config 
+
+    # Computed properties
+    @property
+    def NT(self) -> int:
+        return 64 * 1024 + self.ntbuff
+
+    @property
+    def NTbuff(self) -> int:
+        return self.NT + 4 * self.ntbuff
+
+    @property
+    def nt0min(self) -> int:
+        return int(ceil(20 * self.nt0 / 61))

pykilosort/params.py

@@ -252,7 +252,6 @@ def get_good_channels(raw_data=None, probe=None, params=None):
     minfr_goodchannels = params.minfr_goodchannels
 
     chanMap = probe.chanMap
-    # Nchan = probe.Nchan
     NchanTOT = len(chanMap)
 
     ich = []