grapevine-mcmc 0.1.0__tar.gz

grapevine_mcmc-0.1.0/.github/workflows/release.yml

@@ -0,0 +1,46 @@
+name: Release
+
+on:
+  push:
+    # Sequence of patterns matched against refs/tags
+    tags:
+      - 'v*' # Push events to matching v*, i.e. v1.0, v20.15.10
+
+jobs:
+  release-build:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: build release distributions
+        run: |
+          python -m pip install pdm
+          pdm build
+
+      - name: upload windows dists
+        uses: actions/upload-artifact@v4
+        with:
+          name: release-dists
+          path: dist/
+
+  pypi-publish:
+    runs-on: ubuntu-latest
+    needs:
+      - release-build
+    permissions:
+      id-token: write
+
+    steps:
+      - name: Retrieve release distributions
+        uses: actions/download-artifact@v4
+        with:
+          name: release-dists
+          path: dist/
+
+      - name: Publish release distributions to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

grapevine_mcmc-0.1.0/.github/workflows/run_tests.yml

@@ -0,0 +1,42 @@
+name: Run tests
+
+on:
+  push:
+    branches:
+      - '**'
+    tags-ignore:
+      - '**'
+
+jobs:
+  build:
+
+    runs-on: ${{ matrix.os }}
+
+    strategy:
+        matrix:
+            os: [ubuntu-latest, windows-latest]
+            python-version: [3.12]
+
+    steps:
+
+    - name: checkout code
+      uses: actions/checkout@v2
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v2
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Install pdm
+      run: pip install pdm
+
+    - name: pre-commit checks
+      uses: pre-commit/action@v2.0.3
+
+    - name: Run tests
+      run: |
+        pdm install --dev
+        pdm run pytest tests --cov=src
+
+    - name: Upload coverage reports to Codecov
+      uses: codecov/codecov-action@v3

grapevine_mcmc-0.1.0/.gitignore

@@ -0,0 +1,164 @@
+.DS_Store
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm-project.org/#use-with-ide
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

grapevine_mcmc-0.1.0/.pre-commit-config.yaml

@@ -0,0 +1,15 @@
+repos:
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.6.0
+    hooks:
+    -   id: check-yaml
+    -   id: end-of-file-fixer
+    -   id: trailing-whitespace
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  # Ruff version.
+  rev: v0.6.0
+  hooks:
+    # Run the linter.
+    - id: ruff
+    # Run the formatter.
+    - id: ruff-format

grapevine_mcmc-0.1.0/PKG-INFO

@@ -0,0 +1,99 @@
+Metadata-Version: 2.3
+Name: grapevine-mcmc
+Version: 0.1.0
+Summary: Faster Hamiltonian trajectories for problems with guesses, using Blackjax
+Author-email: Teddy Groves <tedgro@dtu.dk>
+License: MIT
+Requires-Python: >=3.12
+Requires-Dist: blackjax>=1.2.4
+Requires-Dist: jax==0.4.33
+Provides-Extra: demos
+Requires-Dist: optimistix>=0.0.8; extra == 'demos'
+Description-Content-Type: text/markdown
+
+# grapevine
+[![Tests](https://github.com/dtu-qmcm/grapevine/actions/workflows/run_tests.yml/badge.svg)](https://github.com/dtu-qmcm/grapevine/actions/workflows/run_tests.yml)
+[![Project Status: WIP – Initial development is in progress, but there has not yet been a stable, usable release suitable for the public.](https://www.repostatus.org/badges/latest/wip.svg)](https://www.repostatus.org/#wip)
+[![Supported Python versions: 3.12 and newer](https://img.shields.io/badge/python->=3.12-blue.svg)](https://www.python.org/)
+
+JAX/Blackjax implementation of the grapevine method for reusing the solutions of guessing problems embedded in Hamiltonian trajectories.
+
+The grapevine method can dramatically speed up MCMC for statistical with embedded equation solving problems.
+
+## Installation
+
+```sh
+pip install grapevine-mcmc
+```
+
+## Usage
+
+First make a suitable log density function.
+
+This function should have two arguments: a set of parameters (a [Pytree](https://jax.readthedocs.io/en/latest/pytrees.html)) and a guess (also a Pytree). It should return the log density of these parameters (a number) and a new guess. It should also be generally compatible with JAX, and will probalbly involve some differentiable numerical solving, for example using [optimistix](https://docs.kidger.site/optimistix/).
+
+Here is a simple example of such a function:
+
+```python
+from functools import partial
+
+import jax
+
+from jax.scipy.stats import norm
+from jax.scipy.special import expit
+from jax import numpy as jnp
+
+import optimistix as optx
+
+# equation solving problems often need 64 bit floats
+jax.config.update("jax_enable_x64", True)
+
+solver = optx.Newton(rtol=1e-8, atol=1e-8)
+obs = jnp.array(0.7)
+
+
+def fn(y, args):
+    """Equation defining a root-finding problem."""
+    a = args
+    return y - jnp.tanh(y * expit(a) + 1)
+
+
+def joint_logdensity(a, obs, guess):
+    """An example log density."""
+    sol = optx.root_find(fn, solver, guess, args=a)
+    log_prior = norm.logpdf(a, loc=0.0, scale=1.0)
+    log_likelihood = norm.logpdf(obs, loc=sol.value, scale=0.5)
+    return log_prior + log_likelihood, sol.value
+
+
+posterior_logdensity = partial(joint_logdensity, obs=obs)
+posterior_logdensity(a=0.0, guess=0.01)
+# (Array(-1.22095095, dtype=float64), Array(0.8952192, dtype=float64))
+```
+
+Now you can run MCMC on your model using GrapeNUTS, the grapevine version of the [NUTS](http://www.stat.columbia.edu/~gelman/research/published/nuts.pdf) sampler!
+
+```python
+from grapevine import run_grapenuts
+
+INITIAL_POSITION = jnp.array(0.0)
+DEFAULT_GUESS = jnp.array(0.01)
+SEED = 1234
+
+key = jax.random.key(SEED)
+samples, info = run_grapenuts(
+    logdensity_fn=posterior_logdensity,
+    rng_key=key,
+    init_parameters=INITIAL_POSITION,
+    num_warmup=10,
+    num_samples=10,
+    default_guess=DEFAULT_GUESS,
+    progress_bar=False,
+    initial_step_size=0.01,
+    max_num_doublings=4,
+    is_mass_matrix_diagonal=True,
+    target_acceptance_rate=0.8,
+)
+jnp.quantile(samples.position, jnp.array([0.01, 0.5, 0.99]))
+# Array([-1.26712677,  0.12950684,  0.93903677], dtype=float64)
+```

grapevine_mcmc-0.1.0/README.md

@@ -0,0 +1,86 @@
+# grapevine
+[![Tests](https://github.com/dtu-qmcm/grapevine/actions/workflows/run_tests.yml/badge.svg)](https://github.com/dtu-qmcm/grapevine/actions/workflows/run_tests.yml)
+[![Project Status: WIP – Initial development is in progress, but there has not yet been a stable, usable release suitable for the public.](https://www.repostatus.org/badges/latest/wip.svg)](https://www.repostatus.org/#wip)
+[![Supported Python versions: 3.12 and newer](https://img.shields.io/badge/python->=3.12-blue.svg)](https://www.python.org/)
+
+JAX/Blackjax implementation of the grapevine method for reusing the solutions of guessing problems embedded in Hamiltonian trajectories.
+
+The grapevine method can dramatically speed up MCMC for statistical with embedded equation solving problems.
+
+## Installation
+
+```sh
+pip install grapevine-mcmc
+```
+
+## Usage
+
+First make a suitable log density function.
+
+This function should have two arguments: a set of parameters (a [Pytree](https://jax.readthedocs.io/en/latest/pytrees.html)) and a guess (also a Pytree). It should return the log density of these parameters (a number) and a new guess. It should also be generally compatible with JAX, and will probalbly involve some differentiable numerical solving, for example using [optimistix](https://docs.kidger.site/optimistix/).
+
+Here is a simple example of such a function:
+
+```python
+from functools import partial
+
+import jax
+
+from jax.scipy.stats import norm
+from jax.scipy.special import expit
+from jax import numpy as jnp
+
+import optimistix as optx
+
+# equation solving problems often need 64 bit floats
+jax.config.update("jax_enable_x64", True)
+
+solver = optx.Newton(rtol=1e-8, atol=1e-8)
+obs = jnp.array(0.7)
+
+
+def fn(y, args):
+    """Equation defining a root-finding problem."""
+    a = args
+    return y - jnp.tanh(y * expit(a) + 1)
+
+
+def joint_logdensity(a, obs, guess):
+    """An example log density."""
+    sol = optx.root_find(fn, solver, guess, args=a)
+    log_prior = norm.logpdf(a, loc=0.0, scale=1.0)
+    log_likelihood = norm.logpdf(obs, loc=sol.value, scale=0.5)
+    return log_prior + log_likelihood, sol.value
+
+
+posterior_logdensity = partial(joint_logdensity, obs=obs)
+posterior_logdensity(a=0.0, guess=0.01)
+# (Array(-1.22095095, dtype=float64), Array(0.8952192, dtype=float64))
+```
+
+Now you can run MCMC on your model using GrapeNUTS, the grapevine version of the [NUTS](http://www.stat.columbia.edu/~gelman/research/published/nuts.pdf) sampler!
+
+```python
+from grapevine import run_grapenuts
+
+INITIAL_POSITION = jnp.array(0.0)
+DEFAULT_GUESS = jnp.array(0.01)
+SEED = 1234
+
+key = jax.random.key(SEED)
+samples, info = run_grapenuts(
+    logdensity_fn=posterior_logdensity,
+    rng_key=key,
+    init_parameters=INITIAL_POSITION,
+    num_warmup=10,
+    num_samples=10,
+    default_guess=DEFAULT_GUESS,
+    progress_bar=False,
+    initial_step_size=0.01,
+    max_num_doublings=4,
+    is_mass_matrix_diagonal=True,
+    target_acceptance_rate=0.8,
+)
+jnp.quantile(samples.position, jnp.array([0.01, 0.5, 0.99]))
+# Array([-1.26712677,  0.12950684,  0.93903677], dtype=float64)
+```

grapevine_mcmc-0.1.0/benchmarks/optimistix_example.py

@@ -0,0 +1,191 @@
+"""An example comparing GrapeNUTS and NUTS on a representative problem.
+
+This is supposed to be a complete example, mirroring how the grapevine method is used in practice.
+
+
+"""
+
+from functools import partial
+import timeit
+
+from blackjax.util import run_inference_algorithm
+from grapevine.grapenuts import GrapeNUTSState
+import jax
+import jax.numpy as jnp
+import optimistix as optx
+
+from blackjax import nuts
+from blackjax import window_adaptation as nuts_window_adaptation
+from jax.scipy.stats import norm
+
+from grapevine import run_grapenuts
+
+# Use 64 bit floats
+jax.config.update("jax_enable_x64", True)
+
+SEED = 1234
+SD = 0.05
+TRUE_PARAMS = {
+    "log_km_s": 2.0,
+    "log_km_p": 3.0,
+    "log_vmax": -1.0,
+    "log_k_eq": 5.0,
+    "log_s1": 2.0,
+    "log_s2": 2.9,
+    "log_s3": 0.9,
+    "log_s4": 0.1,
+}
+TRUE_PARAMS_ARR = jnp.array(list(TRUE_PARAMS.values()))
+DEFAULT_GUESS = jnp.array([0.01, 0.01, 0.01, 0.01])
+
+# hack the timeit module to not destroy the timed function's return value
+timeit.template = """
+def inner(_it, _timer{init}):
+    {setup}
+    _t0 = _timer()
+    for _i in _it:
+        retval = {stmt}
+    _t1 = _timer()
+    return _t1 - _t0, retval
+"""
+
+
+def rmm(p, km_s, km_p, vmax, k_eq, s):
+    num = vmax * (s - p / k_eq) / km_s
+    denom = 1 + s / km_s + p / km_p
+    return num / denom
+
+
+def fn(y, args):
+    p1, p2, p3, p4 = y
+    km_s, km_p, vmax, k_eq, s1, s2, s3, s4 = args
+    v1 = rmm(p1, km_s, km_p, vmax, k_eq, s1)
+    v2 = rmm(p2, km_s, km_p, vmax, k_eq, s2)
+    v3 = rmm(p3, km_s, km_p, vmax, k_eq, s3)
+    v4 = rmm(p4, km_s, km_p, vmax, k_eq, s4)
+    return jnp.array([v1, v2, v3, v4])
+
+
+solver = optx.Newton(rtol=1e-8, atol=1e-8)
+
+
+def grapenuts_state_from_nuts_state(nuts_state, guess):
+    position, logdensity, logdensity_grad = nuts_state
+    return GrapeNUTSState(position, logdensity, logdensity_grad, guess)
+
+
+def joint_logdensity_grapenuts(params, obs, guess):
+    sol = optx.root_find(fn, solver, guess, args=jnp.exp(params))
+    log_prior = norm.logpdf(
+        params,
+        loc=TRUE_PARAMS_ARR,
+        scale=jnp.full(params.shape, 1),
+    ).sum()
+    log_likelihood = norm.logpdf(
+        jnp.log(obs), loc=jnp.log(sol.value), scale=jnp.full(obs.shape, SD)
+    ).sum()
+    return log_prior + log_likelihood, sol.value
+
+
+def joint_logdensity_nuts(params, obs):
+    sol = optx.root_find(fn, solver, DEFAULT_GUESS, args=jnp.exp(params))
+    log_prior = norm.logpdf(
+        params,
+        loc=TRUE_PARAMS_ARR,
+        scale=jnp.full(params.shape, 1),
+    ).sum()
+    log_likelihood = norm.logpdf(
+        jnp.log(obs), loc=jnp.log(sol.value), scale=jnp.full(obs.shape, SD)
+    ).sum()
+    return log_prior + log_likelihood
+
+
+def simulate(key, params, guess):
+    sol = optx.root_find(fn, solver, guess, args=jnp.exp(params))
+    return sol.value, jnp.exp(
+        jnp.log(sol.value) + jax.random.normal(key, shape=sol.value.shape) * SD
+    )
+
+
+def main():
+    key = jax.random.key(SEED)
+    key, sim_key = jax.random.split(key)
+    true_p, sim = simulate(sim_key, TRUE_PARAMS_ARR, DEFAULT_GUESS)
+    print("True substrate concs: " + str(true_p))
+    print("Simulated measurements: " + str(sim))
+    posterior_logdensity_gn = partial(joint_logdensity_grapenuts, obs=sim)
+    posterior_logdensity_nuts = partial(joint_logdensity_nuts, obs=sim)
+    key, grapenuts_key = jax.random.split(key)
+    key, nuts_key_warmup = jax.random.split(key)
+    key, nuts_key_sampling = jax.random.split(key)
+
+    def run_grapenuts_example():
+        return run_grapenuts(
+            posterior_logdensity_gn,
+            key,
+            init_parameters=(TRUE_PARAMS_ARR),
+            default_guess=DEFAULT_GUESS,
+            num_warmup=1000,
+            num_samples=1000,
+            initial_step_size=0.0001,
+            max_num_doublings=10,
+            is_mass_matrix_diagonal=False,
+            target_acceptance_rate=0.95,
+            progress_bar=False,
+        )
+
+    def run_nuts_example():
+        warmup = nuts_window_adaptation(
+            nuts,
+            posterior_logdensity_nuts,
+            progress_bar=False,
+            initial_step_size=0.0001,
+            max_num_doublings=10,
+            is_mass_matrix_diagonal=False,
+            target_acceptance_rate=0.95,
+        )
+        (initial_state, tuned_parameters), _ = warmup.run(
+            nuts_key_warmup,
+            TRUE_PARAMS_ARR,
+            num_steps=1000,  #  type: ignore
+        )
+        kernel = nuts(posterior_logdensity_nuts, **tuned_parameters)
+        return run_inference_algorithm(
+            nuts_key_sampling,
+            kernel,
+            1000,
+            initial_state,
+        )
+
+    time_grapenuts, (state_grapenuts, _) = timeit.timeit(
+        run_grapenuts_example,
+        number=1,
+    )  #  type: ignore
+    time_nuts, (_, (state_nuts, _)) = timeit.timeit(
+        run_nuts_example,
+        number=1,
+    )  #  type: ignore
+    __import__("pdb").set_trace()
+    print("True param vals: " + str(TRUE_PARAMS_ARR))
+    print("GrapeNUTS quantiles:")
+    print(
+        jnp.quantile(
+            state_grapenuts.position,
+            jnp.array([0.01, 0.5, 0.99]),
+            axis=0,
+        ).round(4)
+    )
+    print("NUTS quantiles:")
+    print(
+        jnp.quantile(
+            state_nuts.position,
+            jnp.array([0.01, 0.5, 0.99]),
+            axis=0,
+        ).round(4)
+    )
+    print(f"Runtime for grapenuts: {round(time_grapenuts, 4)}")
+    print(f"Runtime for nuts: {round(time_nuts, 4)}")
+
+
+if __name__ == "__main__":
+    main()