Added environment.yml for smoother development installation

Author: jacobpennington

README.md

@@ -102,18 +102,12 @@ phy template-gui params.py
 
 ## Developer instructions
 
-To get the most up-to-date changes to the code, clone the repository and install in editable mode in your Kilosort environment, along with the other installation steps mentioned above.
+To get the most up-to-date changes to the code, clone the repository and install in editable mode in your Kilosort environment, along with the other installation steps mentioned above. Using the included `environment.yml` to create your environment is recommended, which will also install pytest and the cuda version of Pytorch.
 ~~~
 git clone [email protected]:MouseLand/Kilosort.git
+conda env create -f environment.yml
 conda activate kilosort
 pip install -e Kilosort[gui]
-pip uninstall torch
-conda install pytorch pytorch-cuda=11.8 -c pytorch -c nvidia
-~~~
-
-For unit testing, you will need to install pytest
-~~~
-pip install pytest
 ~~~
 
 Then run all tests with:

environment.yml

@@ -0,0 +1,37 @@
+name: kilosort
+channels:
+  - conda-forge
+  - pytorch
+  - nvidia
+dependencies:
+  - python=3.9.*
+  - pip
+  - psutil
+  - numpy>=1.20.0,<2.0.0
+  - scipy
+  - scikit-learn
+  - tqdm
+  - numba
+  - faiss-cpu
+  - libblas=*=*mkl
+  - pytest
+  # CUDA
+  - pytorch-cuda=11.8
+  - pytorch
+  - pynvml
+  # GUI
+  - matplotlib
+  - pip:
+    - pyqtgraph>=0.13.0
+    - qtpy
+    - pyqt6
+    - pyqt6.sip
+
+
+# TODO first thing tomorrow, create fresh env from this now that I've made all
+# the adjustments. Goal is should only need to do the following:
+# ```
+# conda create -f environment.yml
+# conda activate kilosort
+# pip install -e .
+# ```

setup.py

@@ -9,6 +9,7 @@
     'torch>=1.6',
     'numba',
     'faiss-cpu',
+    'psutil'
 ]
 
 gui_deps = [