Merge pull request #375 from lnccbrown/update-pymc

Improve documentation for Go-Nogo, bug fixes and test improvement
lnccbrown · Mar 22, 2024 · 77e3aad · 77e3aad
2 parents f2bfa62 + 47cc43f
commit 77e3aad
Show file tree

Hide file tree

Showing 19 changed files with 912 additions and 451 deletions.
diff --git a/.github/workflows/build_and_publish.yml b/.github/workflows/build_and_publish.yml
@@ -8,6 +8,9 @@ on:
 jobs:
   run_tests:
     runs-on: ubuntu-latest
+    defaults:
+      run:
+        shell: bash -el {0}
 
     strategy:
       fail-fast: true
@@ -17,41 +20,37 @@ jobs:
     steps:
       - uses: actions/checkout@v4
 
-      - name: Setup Python
-        uses: actions/setup-python@v5
+      - name: Setup miniconda
+        uses: conda-incubator/setup-miniconda@v3
         with:
           python-version: ${{ matrix.python-version }}
+          miniforge-variant: Mambaforge
+          mamba-version: "*"
+          channels: conda-forge,defaults
+          channel-priority: true
 
-      - name: Install poetry
-        uses: snok/install-poetry@v1
-        with:
-          version: latest
-          virtualenvs-create: true
-          virtualenvs-in-project: true
-          installer-parallel: true
-
-      - name: Load cache
-        uses: actions/cache@v4
-        id: cache
-        with:
-          path: .venv
-          key: venv-${{ runner.os }}-${{ matrix.python-version }}-${{ hashFiles('**/pyproject.toml') }}
+      - name: install pymc and poetry
+        run: |
+          mamba info
+          mamba install -c conda-forge pymc=5.10 poetry
 
-      - name: Install dependencies
-        if: steps.cache.outputs.cache-hit != 'true'
-        run: poetry install --no-interaction
+      - name: install hssm
+        run: |
+          poetry export --with dev --output=requirements.txt
+          pip install -r requirements.txt
+          pip install -e .
 
       - name: Run mypy
-        run: poetry run mypy src
+        run: mypy src
 
       - name: Check styling
-        run: poetry run black . --check
+        run: black . --check
 
       - name: Linting
-        run: poetry run ruff check .
+        run: ruff check .
 
       - name: Run tests
-        run: poetry run pytest -n auto -s
+        run: pytest -n auto -s
 
   publish:
     name: Build wheel and publish to test-PyPI, and then PyPI, and publish docs

diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml
@@ -6,6 +6,9 @@ on:
 jobs:
   run_tests:
     runs-on: ubuntu-latest
+    defaults:
+      run:
+        shell: bash -el {0}
 
     strategy:
       fail-fast: true
@@ -15,41 +18,37 @@ jobs:
     steps:
       - uses: actions/checkout@v4
 
-      - name: Setup Python
-        uses: actions/setup-python@v5
+      - name: Setup miniconda
+        uses: conda-incubator/setup-miniconda@v3
         with:
           python-version: ${{ matrix.python-version }}
-
-      - name: Install poetry
-        uses: snok/install-poetry@v1
-        with:
-          version: latest
-          virtualenvs-create: true
-          virtualenvs-in-project: true
-          installer-parallel: true
-
-      - name: Load cache
-        uses: actions/cache@v4
-        id: cache
-        with:
-          path: .venv
-          key: venv-${{ runner.os }}-${{ matrix.python-version }}-${{ hashFiles('**/pyproject.toml') }}
-
-      - name: Install dependencies
-        if: steps.cache.outputs.cache-hit != 'true'
-        run: poetry install --no-interaction
+          miniforge-variant: Mambaforge
+          mamba-version: "*"
+          channels: conda-forge,defaults
+          channel-priority: true
+
+      - name: install pymc and poetry
+        run: |
+          mamba info
+          mamba install -c conda-forge pymc=5.10 poetry
+
+      - name: install hssm
+        run: |
+          poetry export --with dev --output=requirements.txt
+          pip install -r requirements.txt
+          pip install -e .
 
       - name: Run mypy
-        run: poetry run mypy src
+        run: mypy src
 
       - name: Check styling
-        run: poetry run black . --check
+        run: black . --check
 
       - name: Linting
-        run: poetry run ruff check .
+        run: ruff check .
 
       - name: Run tests
-        run: poetry run pytest -n auto -s
+        run: pytest -n auto -s
 
       - name: build docs
-        run: poetry run mkdocs build
+        run: mkdocs build
diff --git a/README.md b/README.md
@@ -31,21 +31,58 @@ HSSM is a Python toolbox that provides a seamless combination of state-of-the-ar
 
 ## Installation
 
-`hssm` is available through PyPI. You can install it with pip via:
+**Update**: From HSSM 0.2.1 on, we recommend `conda` as the virtual environment manager for HSSM. We will also gradually make HSSM available directly through `conda-forge` in the near future. For now, please follow the instruction below to install HSSM:
 
+### Step 1: Create a conda environment
+
+If you haven't already, please follow the [Anaconda official website](https://www.anaconda.com/download) to install anaconda. We assume that you already have one of [Anaconda](https://www.anaconda.com/download), [Miniconda](https://docs.anaconda.com/free/miniconda/index.html), [miniforge](https://github.com/conda-forge/miniforge/releases), or [mambaforge](https://github.com/conda-forge/miniforge/releases) installed on your system and have access to either `conda` or `mamba` available on your command line.
+
+To create a conda environment, use the following command. Substitute `mamba` for `conda` if `mamba` is available:
+
+```bash
+conda create -n <your-env-name> python=3.11
+conda activate <your-env-name>
+```
+
+Substitute `<your-env-name>` with the name of the virtual environment that you choose. HSSM 0.2.0 and above supports Python versions 3.10 and 3.11.
+
+### Step 2: Install PyMC through conda-forge
+
+Installation through `conda-forge` is the official way of installing PyMC. This will also install other libraries such as `libblas` that PyMC requires to run properly.
+
+```bash
+conda install -c conda-forge pymc
+```
+
+As of HSSM 0.2.1, HSSM supports PyMC 5.10.4. If a future newer version of PyMC causes compatibility issues, please specify the version of PyMC:
+
+```bash
+conda install -c conda-forge pymc=5.10
 ```
+
+### Step 3: Install `hssm` through `pip`
+
+In the same environment, install `hssm` through `pip`.
+
+```bash
 pip install hssm
 ```
 
-You can also install the bleeding-edge version of `hssm` directly from this repo:
+## Advanced Installation
 
+`hssm` is also available through PyPI. You can directly install it with pip into any virtual environment via:
+
+```bash
+pip install hssm
 ```
+
+You can also install the bleeding-edge version of `hssm` directly from this repo:
+
+```bash
 pip install git+https://github.com/lnccbrown/HSSM.git
 ```
 
-You will need optional dependencies to use JAX-based samplers and graph the models.
-Please refer to our [installation guide](https://lnccbrown.github.io/HSSM/getting_started/installation/)
-for more detailed instructions.
+While this installation is much simpler, you might need optional dependencies to use JAX-based samplers and to produce model graphs. You might also encounter this warning message `WARNING (pytensor.tensor.blas): Using NumPy C-API based implementation for BLAS functions.` Please refer to our [advanced installation guide](https://lnccbrown.github.io/HSSM/getting_started/installation/) for more details.
 
 **Note**: Possible solutions to any issues with installations with hssm can be located
 [here](https://github.com/lnccbrown/HSSM/discussions). We recommend leveraging an

diff --git a/docs/changelog.md b/docs/changelog.md
@@ -2,6 +2,19 @@
 
 ## 0.2.x
 
+### 0.2.1
+
+We added a few new features in 0.2.1:
+
+1. We have finished updating the HSSM code base to support go-nogo data and deadline. We will provide documentation once the networks are added to our huggingface repo.
+2. We updated `hssm.distribution_utils` to streamline the creation of `pm.Distribution`s.
+3. We now support response variables other than `rt` and `response`. They can be specified through `model_config` via the new `response` field.
+
+Other minor updates
+
+- Fixed an incompatible shape error during posterior predictive sampling when `p_outlier` is estimated as a parameter.
+- Updated documentation for using `make_distribution` with PyMC.
+
 ### 0.2.0
 
 This is a major version update! Many changes have taken place in this version:

diff --git a/docs/getting_started/installation.md b/docs/getting_started/installation.md
@@ -1,33 +1,69 @@
 # Installation
 
-## Install the HSSM package
+**Update**: From HSSM 0.2.1 on, we recommend `conda` as the virtual environment manager for HSSM. We will also gradually make HSSM available directly through `conda-forge` in the near future. For now, please follow the instruction below to install HSSM:
 
-### Install via `pip` from PyPI
+## Recommended: install HSSM in a conda environment
 
-HSSM is available through [PyPI](https://pypi.org/project/hssm/). The easiest way to
-install HSSM is through pip.
+### Step 1: Create a conda environment
+
+If you haven't already, please follow the [Anaconda official website](https://www.anaconda.com/download) to install anaconda. We assume that you already have one of [Anaconda](https://www.anaconda.com/download), [Miniconda](https://docs.anaconda.com/free/miniconda/index.html), [miniforge](https://github.com/conda-forge/miniforge/releases), or [mambaforge](https://github.com/conda-forge/miniforge/releases) installed on your system and have access to either `conda` or `mamba` available on your command line.
+
+To create a conda environment, use the following command. Substitute `mamba` for `conda` if `mamba` is available:
+
+```bash
+conda create -n <your-env-name> python=3.11
+conda activate <your-env-name>
+```
+
+Substitute `<your-env-name>` with the name of the virtual environment that you choose. HSSM 0.2.0 and above support Python versions 3.10 and 3.11.
+
+### Step 2: Install PyMC through conda-forge
+
+Installation through `conda-forge` is the official way of installing PyMC. This will also install other libraries such as `libblas` that PyMC requires to run properly.
+
+```bash
+conda install -c conda-forge pymc
+```
+
+As of HSSM 0.2.1, HSSM supports PyMC 5.10.4. If a future newer version of PyMC causes compatibility issues, please specify the version of PyMC:
+
+```bash
+conda install -c conda-forge pymc=5.10
+```
+
+### Step 3: Install `hssm` through `pip`
+
+In the same environment, install `hssm` through `pip`.
 
 ```bash
 pip install hssm
 ```
 
-Because HSSM depends on very specific versions of PyMC, JAX and Bambi, we recommend that
-you install HSSM into a dedicated virtual environment to avoid dependency conflicts.
+## Advanced: install via `pip` from PyPI or GitHub
 
-### Install from GitHub
+`hssm` is also available through PyPI. You can directly install it with pip into any virtual environment via:
 
-You can also install the bleeding edge version of `hssm` directly from
-[our repo](https://github.com/lnccbrown/HSSM):
+```bash
+pip install hssm
+```
+
+You can also install the bleeding-edge version of `hssm` directly from this repo:
 
 ```
 pip install git+https://github.com/lnccbrown/HSSM.git
 ```
 
-## Install optional dependencies
+Because HSSM depends on very specific versions of PyMC, JAX and Bambi, we recommend that
+you install HSSM into a dedicated virtual environment to avoid dependency conflicts.
+
+## Advanced: install optional dependencies
 
-Some functionalities in HSSM are available through optional dependencies.
+Whether you have installed HSSM via `conda`, `pip`, or GitHub, you might still need additional
+packages installed for additional features such as sampling with `blackjax` or GPU support
+for `JAX`. Please follow the instructions below if you need any of these additional
+features:
 
-### Sampling with JAX through `numpyro` or `blackjax`
+### 1. Sampling with JAX through `numpyro` or `blackjax`
 
 JAX-based sampling is done through `numpyro` and `blackjax`. `numpyro` is installed as
 a dependency by default. You need to have `blackjax` installed if you want to use the
@@ -37,21 +73,14 @@ a dependency by default. You need to have `blackjax` installed if you want to us
 pip install blackjax
 ```
 
-### Sampling with JAX support for GPU
+### 2. Sampling with JAX support for GPU
 
 The `nuts_numpyro` sampler uses JAX as the backend and thus can support sampling on nvidia
-GPU. The only thing you need to do to take advantage of this is to install JAX with CUDA
-support before installing HSSM. Here's one example:
+GPUs. The only thing you need to do to take advantage of this is to install JAX with CUDA
+support before installing HSSM. This works whether `conda` is used or not. Here's one example:
 
 ```bash
-python -m venv .venv  # Create a virtual environment
-source .venv/bin/activate  # Activate the virtual environment
-
-pip install --upgrade pip
-
-# We need to limit the version of JAX for now due to some breaking
-# changes introduced in JAX 0.4.16.
-pip install --upgrade "jax[cuda11_pip]<0.4.16" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 pip install hssm
 ```
 
@@ -62,10 +91,16 @@ details on installing JAX on different platforms with GPU or TPU support.
 Note that on Google Colab, JAX support for GPU is enabled by default if the Colab backend
 has GPU enabled. You simply need only install HSSM.
 
-### Visualizing the model
+### Visualizing the model with `graphviz`
 
-Model graphs are created with `model.graph()` through `graphviz`. In order to use it,
-you need to have `graphviz` installed system-wide and then install its Python binding:
+Model graphs are created with `model.graph()` through `graphviz`. If you have installed
+hssm in a conda environment, you can simply install `graphviz` in conda:
+
+```bash
+conda install -c conda-forge graphviz
+```
+
+If you have installed hssm in a non-conda environment, you need to have `graphviz` installed system-wide and then install its Python binding:
 
 #### Install graphviz system-wide
 
@@ -83,7 +118,19 @@ pip install graphviz
 
 ## Common issues
 
-1. `pip` installation fails with missing dependencies:
+1. I run into warnings such as
+
+```
+WARNING (pytensor.tensor.blas): Using NumPy C-API based implementation for BLAS functions.
+```
+
+   This is because `pytensor`, the compute backend of PyMC, cannot find a BLAS library on your
+   system to optimize its computation. This can be resolved by following the recommended
+   steps to install HSSM into a conda environment. If conda cannot be used, you can follow
+   [this discussion](https://github.com/pymc-devs/pytensor/issues/524) to link a BLAS library
+   with `pytensor`.
+
+2. `pip` installation fails with missing dependencies:
 
    Here's an example: