# Source code of HuiduRep

```tree
.
│   README.md
│   train.ipynb
│   validate.ipynb
│
├───model
│       CMAES.py
│       FeatureDecoder.py
│       FeatureEncoder.py
│       Prediction.py
│       ProjectionHead.py
│
├───resources
│   └───checkpoint
│           HuiduRep.pt
│
├───train
│       CMAES_trainer.py
│       spike_dataset.py
│
└───utils
        CMAES_utils.py
        monitor_utils.py
        scheduler_utils.py
        spike_augment.py
        test_utils.py
        spikeforest_utils.py
```

Due to storage limitations, only the trained model weights are provided. The training data and other dataset details used in our experiments can be accessed from the publicly available **International Brain Laboratory (IBL)** dataset, which have mentioned in our paper. Please refer to the documentation of the IBL dataset for information on data access and preprocessing.

- **train.ipynb**: overall training code of HuiduRep.
- **validate.ipynb**: overall validation code of HuiduRep.
- **model**
  - **CMAES.py**: Model definition of Huiduirep.
- **resources**
  - **HuiduRep.pt**: Weights of HuiduRep with $\alpha = 0.2$ and representation dim  $=32$.
- **train**
  - **CMAES_trainer.py**: Define the trainer function for HuiduRep, including the optimiser, learning rate scheduler and other hyperparameters.
  - **spike_dataset.py**: Define the dataset loader used during training, including the use of view augmentation strategies.
- **utils**
  - **CMAES_utils.py**: Contains the `load_model` function to load HuiduRep from its HuiduRep.pt.
  - **monitor_utils.py**: Contains a series of monitor functions with different cluster methods like GMM, Kmeans, HDBSCAN to monitor the effect of model along with evaluate the performance of HuiduRep.
  - **scheduler_utils.py**: Define some learning rate schedulers like `WarmupCosineLR`.
  - **spike_augment.py**: Define whole view augmentation strategies announced in our paper, include jitter, noise, crop, and collision.
  - **test_utils.py**: Define the utils for testing the performance of HuiduRep.
  - **spikeforest_utils.py**: Define the noticed preprocessing steps in the supplementary material as well as the threshold based spike detection method used in the HuiduRep Pipeline.