Merge pull request #109 from talmolab/aidan/uv

Added uv support

Author: asirbu77

README.md

@@ -4,6 +4,70 @@
 This is part of the Virtual Neuroscience Lab (VNL) project.
 
 ## Installation
+
+### Option 1: `uv`
+
+#### Prerequisites
+
+- Python 3.11 or 3.12
+- [uv](https://docs.astral.sh/uv/) package manager (recommended) or pip
+- CUDA 12.x or 13.x (for GPU support, optional)
+
+#### Installing `uv`
+
+If you don't have uv installed:
+
+```bash
+# Linux/macOS
+curl -LsSf https://astral.sh/uv/install.sh | sh
+
+# Or using pip
+pip install uv
+```
+
+#### Installation Steps
+
+1. Clone the repository:
+```bash
+git clone https://github.com/talmolab/stac-mjx.git
+```
+2. Create and activate a virutal environment:
+```bash
+uv venv
+source .venv/bin/activate  # On Windows: .venv\Scripts\activate
+```
+3. Install the package with optional dependencies based on your hardware. CUDA 12, CUDA 13, and CPU-only configurations are supported:
+
+For CUDA 12.x:
+```bash
+uv pip install -e ".[cuda12]"
+```
+
+For CUDA 13.x:
+```bash
+uv pip install -e ".[cuda13]"
+```
+
+For CPU-only:
+```bash
+uv pip install -e .
+```
+
+For development, include the `[dev]` extras in addition to the hardware optional dependencies:
+```bash
+uv pip install -e ".[cuda13,dev]"
+```
+4. Verify the installation:
+```bash
+python -c "import jax; print(f'JAX version: {jax.__version__}'); print(f'Available devices: {jax.devices()}')"
+```
+5. Register the environment as a Jupyter kernel:
+```bash
+python -m ipykernel install --user --name=stac-mjx --display-name="Python (stac-mjx)"
+```
+
+### Option 2: `conda`
+
 stac-mjx relies on many prerequisites, therefore we suggest installing in a new conda environment, using the provided `environment.yaml`:
 [Local installation before package is officially published]
 1. Clone the repository `git clone https://github.com/talmolab/stac-mjx.git` and `cd` into it

demos/rodent_demo.ipynb

@@ -42,6 +42,8 @@ version = {attr = "stac_mjx.version.__version__"}
 readme = {file = ["README.md"], content-type="text/markdown"}
 
 [project.optional-dependencies]
+cuda12 = ["jax[cuda12]==0.7.2 ; sys_platform == 'linux'"]
+cuda13 = ["jax[cuda13]==0.7.2 ; sys_platform == 'linux'"]
 dev = [
     "pytest",
     "pytest-cov",