[ { "path": ".codecov.yml", "content": "comment: false\n" }, { "path": ".github/ISSUE_TEMPLATE/generic-issue-template.md", "content": "---\nname: generic issue template\nabout: question, feature request, or bug report\n\n---\n\n*Please make sure you are using **the latest tagged versions** and **the last stable release of Julia** before proceeding.* Support for other versions is very limited.\n\nIf you need *help* using this package for Bayesian inference, please provide a self-contained description of your inference problem (simplifying if possible) and preferably the code you have written so far to code the log likelihood. If you found a *bug*, please provide a self-contained working example, complete with (simulated) data. Please make sure you set a random seed (`Random.seed!`) at the beginning to make your example reproducible. If you are requesting a new *feature*, please provide a description and a rationale.\n\n## Self-contained example that demonstrates the problem\n\n```julia\nusing DynamicHMC\n```\n\n## Output, expected outcome, comparison to other samplers\n\nDid the sampler fail to run, produce incorrect results, …?\n\n## Contributing code for tests\n\nYou can contribute to the development of this package by allowing that your example is used as a test. Please indicate whether your code can be incorporated into this package under the MIT \"Expat\" license, found in the root directory of this package.\n" }, { "path": ".github/dependabot.yml", "content": "# see https://docs.github.com/en/code-security/dependabot/working-with-dependabot/keeping-your-actions-up-to-date-with-dependabot\nversion: 2\nupdates:\n - package-ecosystem: \"github-actions\"\n directory: \"/\"\n schedule:\n # Check for updates to GitHub Actions every month\n interval: \"monthly\"\n" }, { "path": ".github/workflows/CI.yml", "content": "# from https://discourse.julialang.org/t/easy-workflow-file-for-setting-up-github-actions-ci-for-your-julia-package/49765\n\nname: CI\non:\n pull_request:\n branches:\n - master\n push:\n branches:\n - master\n tags: '*'\njobs:\n test:\n name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}\n runs-on: ${{ matrix.os }}\n strategy:\n fail-fast: false\n matrix:\n version:\n - '1.10' # Replace this with the minimum Julia version that your package supports. E.g. if your package requires Julia 1.9 or higher, change this to '1.9'.\n - '1' # Leave this line unchanged. '1' will automatically expand to the latest stable 1.x release of Julia.\n # NOTE: CI on nightly is not enabled by default, uncomment if you need it. Cf the discussion at\n # https://discourse.julialang.org/t/why-do-packages-run-continuous-integration-tests-on-julia-nightly/101208/3\n # - 'nightly'\n os:\n - ubuntu-latest\n arch:\n - x64\n steps:\n - uses: actions/checkout@v6\n - uses: julia-actions/setup-julia@v2\n with:\n version: ${{ matrix.version }}\n arch: ${{ matrix.arch }}\n - uses: actions/cache@v5\n env:\n cache-name: cache-artifacts\n with:\n path: ~/.julia/artifacts\n key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}\n restore-keys: |\n ${{ runner.os }}-test-${{ env.cache-name }}-\n ${{ runner.os }}-test-\n ${{ runner.os }}-\n - uses: julia-actions/julia-buildpkg@v1\n - uses: julia-actions/julia-runtest@v1\n - uses: julia-actions/julia-processcoverage@v1\n - uses: codecov/codecov-action@v5\n with:\n file: lcov.info\n # NOTE: you need to add the token to Github secrets, see\n # https://docs.codecov.com/docs/adding-the-codecov-token\n token: ${{ secrets.CODECOV_TOKEN }}\n" }, { "path": ".github/workflows/CompatHelper.yml", "content": "# see the docs at https://github.com/JuliaRegistries/CompatHelper.jl\n\nname: CompatHelper\non:\n schedule:\n - cron: 0 0 * * *\n workflow_dispatch:\npermissions:\n contents: write\n pull-requests: write\njobs:\n CompatHelper:\n runs-on: ubuntu-latest\n steps:\n - name: Check if Julia is already available in the PATH\n id: julia_in_path\n run: which julia\n continue-on-error: true\n - name: Install Julia, but only if it is not already available in the PATH\n uses: julia-actions/setup-julia@v2\n with:\n version: '1'\n arch: ${{ runner.arch }}\n if: steps.julia_in_path.outcome != 'success'\n - name: \"Add the General registry via Git\"\n run: |\n import Pkg\n ENV[\"JULIA_PKG_SERVER\"] = \"\"\n Pkg.Registry.add(\"General\")\n shell: julia --color=yes {0}\n - name: \"Install CompatHelper\"\n run: |\n import Pkg\n name = \"CompatHelper\"\n uuid = \"aa819f21-2bde-4658-8897-bab36330d9b7\"\n version = \"3\"\n Pkg.add(; name, uuid, version)\n shell: julia --color=yes {0}\n - name: \"Run CompatHelper\"\n run: |\n import CompatHelper\n CompatHelper.main()\n shell: julia --color=yes {0}\n env:\n GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}\n COMPATHELPER_PRIV: ${{ secrets.DOCUMENTER_KEY }}\n # COMPATHELPER_PRIV: ${{ secrets.COMPATHELPER_PRIV }}\n" }, { "path": ".github/workflows/TagBot.yml", "content": "# see the docs at https://github.com/JuliaRegistries/TagBot\n\nname: TagBot\non:\n issue_comment:\n types:\n - created\n workflow_dispatch:\n inputs:\n lookback:\n default: 3\npermissions:\n actions: read\n checks: read\n contents: write\n deployments: read\n issues: read\n discussions: read\n packages: read\n pages: read\n pull-requests: read\n repository-projects: read\n security-events: read\n statuses: read\njobs:\n TagBot:\n if: github.event_name == 'workflow_dispatch' || github.actor == 'JuliaTagBot'\n runs-on: ubuntu-latest\n steps:\n - uses: JuliaRegistries/TagBot@v1\n with:\n token: ${{ secrets.GITHUB_TOKEN }}\n # Edit the following line to reflect the actual name of the GitHub Secret containing your private key\n ssh: ${{ secrets.DOCUMENTER_KEY }}\n # ssh: ${{ secrets.NAME_OF_MY_SSH_PRIVATE_KEY_SECRET }}\n" }, { "path": ".github/workflows/documentation.yml", "content": "# see https://juliadocs.github.io/Documenter.jl/dev/man/hosting/#GitHub-Actions\n# add this file to the repository once you set up authentication as described in the Documenter manual\n\nname: Documentation\n\non:\n push:\n branches:\n - master\n tags: '*'\n pull_request:\n\njobs:\n build:\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v6\n - uses: julia-actions/setup-julia@latest\n with:\n version: '1.10'\n - name: Install dependencies\n run: julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'\n - name: Build and deploy\n env:\n GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # If authenticating with GitHub Actions token\n DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }} # If authenticating with SSH deploy key\n run: julia --project=docs/ docs/make.jl\n" }, { "path": ".gitignore", "content": "*.jl.cov\n*.jl.*.cov\n*.jl.mem\n/notes.org\n/Manifest.toml\n/deps/deps.jl\n/docs/build\n/docs/Manifest.toml\n/test/coverage/Manifest.toml\n" }, { "path": "CHANGELOG.md", "content": "# Changelog\n\nAll notable changes to this project will be documented in this file.\n\nThe format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),\nand this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).\n\n## Unreleased\n\n### Added\n### Changed\n### Deprecated\n### Removed\n### Fixed\n### Security\n\n## v3.6.0\n\n- add `logdensities` to `mcmc` results (#194)\n- minor cleanup of `Project.toml`\n\n## v3.5.1\n\n- Cleanup tests, tighten tolerances. (#190)\n- Use OhMyThreads for threaded example. (#191)\n\n## v3.5.0\n\n- remove `@unpack`, use `(;)`\n- require Julia 1.10\n- minor test fixes\n\n## v3.4.8\n\n- refresh Github actions\n- loosen test tolerances a bit to test cleanly on Julia 1.11\n- remove vestigial code from tests\n\n## v3.4.7\n\n- fix compat entries\n\n## v3.4.6\n\n- test with Aqua, minor fixes\n\n## v3.4.5\n\n- Provide more details when initial stepsize search fails. (#180)\n- Simplify stepsize search. (#181)\n- replace `Parameters: @unpack` with `SimpleUnPack` (#182)\n- replace `GLOBAL_RNG` with `default_rng()` (#183)\n\n## v3.4.4\n\n- compat bumps\n\n## 3.4.3\n\n- remove workaround for JET.jl\n\n## 3.4.2\n\n- change stacked ordering\n\n## 3.4.1\n\n- Support Julia LTS and update CompatHelper\n\n## 3.4.0\n\n- transition to MCMCDiagnosticTools for diagnostics\n- minor fixes\n- JET tests up front so that method extensions do not confuse it\n\n## v3.3.0\n\n- remove `position_matrix` [https://github.com/tpapp/DynamicHMC.jl/pull/165](#165), add utility functions for posterior matrices\n\n## v3.2.1\n\n- compat version bumps\n\n## v3.2.0\n\n- compat version bumps\n\n## v3.1.2\n\n- compat version bumps\n\n## v3.1.1\n\n- minor doc and export list fixes (follow-up to ([#145](https://github.com/tpapp/DynamicHMC.jl/pull/145)))\n\n## v3.1.0\n\n- more robust U-turn checking ([#145](https://github.com/tpapp/DynamicHMC.jl/pull/145))\n\n## v3.0.0\n\n- get rid of `local_optimization` in warmup ([#146](https://github.com/tpapp/DynamicHMC.jl/pull/146))\n\n## v2.2.0\n\n- add a progress bar ([#136](https://github.com/tpapp/DynamicHMC.jl/pull/136))\n- compat bounds, minor changes\n\n## v2.1.4\n\n- compat bumps extension\n\n## v2.1.3\n\n- relax test bounds a bit ([#116](https://github.com/tpapp/DynamicHMC.jl/pull/116))\n\n## v2.1.2\n\nTechnical release (compat version bounds extended).\n\n## v2.1.1\n\n- re-enable support for Julia 1.0 ([#107](https://github.com/tpapp/DynamicHMC.jl/pull/107))\n\n- fix penalty sign in initial optimization ([#97](https://github.com/tpapp/DynamicHMC.jl/pull/97))\n\n- add example for skipping stepsize search ([#104](https://github.com/tpapp/DynamicHMC.jl/pull/104))\n\n## v2.1.0\n\n- add experimental “iterator” interface ([#94](https://github.com/tpapp/DynamicHMC.jl/pull/94))\n\n- use `randexp` for Metropolis acceptance draws\n\n- remove dependence on StatsFuns.jl\n\n## v2.0.2\n\nDefault keyword arguments for LogProgressReport.\n\n## v2.0.1\n\nDon't print `chain_id` when it is `nothing`.\n\n## v2.0.0\n\nNote: the interface was redesigned. You probably want to review the docs, especially the worked example.\n\n### API changes\n\n- major API change: entry point is now `mcmc_with_warmup`\n\n- refactor warmup code, add initial optimizer\n\n- use the LogDensityProblems v0.9.0 API\n\n- use Julia's Logging module for progress messages\n\n- diagnostics moved to `DynamicHMC.Diagnostics`\n\n - report turning and divergence positions\n\n - add `leapfrog_trajectory` for exploration\n\n### Implementation changes\n\n- factor out the tree traversal code\n\n - abstract trajectory interface\n\n - separate random and non-random parts\n\n - stricter and more exact unit tests\n\n- refactor Hamiltonian code slightly\n\n - caching is now in EvaluatedLogDensity\n\n - functions renamed\n\n- misc\n\n - remove dependency on DataStructures, Suppressor\n\n - cosmetic changes to dual averaging code\n\n - large test cleanup\n\n## v1.0.6\n\n- fix LogDensityProblems version bounds\n\n## v1.0.5\n\n- fix tuning with singular covariance matrices\n\n## v1.0.4\n\n- minor fixes in tests and coverage\n\n## v1.0.3 and prior\n\nNo CHANGELOG available.\n" }, { "path": "LICENSE.md", "content": "The DynamicHMC.jl package is licensed under the MIT \"Expat\" License:\n\n> Copyright (c) 2020: Tamas K. Papp.\n>\n> Permission is hereby granted, free of charge, to any person obtaining a copy\n> of this software and associated documentation files (the \"Software\"), to deal\n> in the Software without restriction, including without limitation the rights\n> to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n> copies of the Software, and to permit persons to whom the Software is\n> furnished to do so, subject to the following conditions:\n>\n> The above copyright notice and this permission notice shall be included in all\n> copies or substantial portions of the Software.\n>\n> THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n> IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n> FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n> AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n> LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n> OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n> SOFTWARE.\n>\n" }, { "path": "Project.toml", "content": "name = \"DynamicHMC\"\nuuid = \"bbc10e6e-7c05-544b-b16e-64fede858acb\"\nversion = \"3.6.0\"\nauthors = [\"Tamas K. Papp \"]\n\n[deps]\nArgCheck = \"dce04be8-c92d-5529-be00-80e4d2c0e197\"\nDocStringExtensions = \"ffbed154-4ef7-542d-bbb7-c09d3a79fcae\"\nFillArrays = \"1a297f60-69ca-5386-bcde-b61e274b549b\"\nLinearAlgebra = \"37e2e46d-f89d-539d-b4ee-838fcccc9c8e\"\nLogDensityProblems = \"6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c\"\nLogExpFunctions = \"2ab3a3ac-af41-5b50-aa03-7779005ae688\"\nProgressMeter = \"92933f4c-e287-5a05-a399-4b506db050ca\"\nRandom = \"9a3f8284-a2c9-5f02-9a11-845980a1fd5c\"\nStatistics = \"10745b16-79ce-11e8-11f9-7d13ad32a3b2\"\nTensorCast = \"02d47bb6-7ce6-556a-be16-bb1710789e2b\"\n\n[compat]\nArgCheck = \"1, 2\"\nDocStringExtensions = \"0.8, 0.9\"\nFillArrays = \"0.13, 1\"\nLinearAlgebra = \"1.6\"\nLogDensityProblems = \"1, 2\"\nLogExpFunctions = \"0.3\"\nProgressMeter = \"1\"\nRandom = \"1.6\"\nStatistics = \"1\"\nTensorCast = \"0.4\"\njulia = \"1.10\"\n\n[workspace]\nprojects = [\"test\", \"docs\"]\n" }, { "path": "README.md", "content": "# DynamicHMC\n\nImplementation of robust dynamic Hamiltonian Monte Carlo methods in Julia.\n\n[![Project Status: Active – The project has reached a stable, usable state and is being actively developed.](https://www.repostatus.org/badges/latest/active.svg)](https://www.repostatus.org/#active)\n[![build](https://github.com/tpapp/DynamicHMC.jl/workflows/CI/badge.svg)](https://github.com/tpapp/DynamicHMC.jl/actions?query=workflow%3ACI)\n[![codecov](https://codecov.io/github/tpapp/DynamicHMC.jl/graph/badge.svg?token=idfCiPzxjL)](https://codecov.io/github/tpapp/DynamicHMC.jl)\n[![Documentation](https://img.shields.io/badge/docs-stable-blue.svg)](https://tpapp.github.io/DynamicHMC.jl/stable)\n[![Documentation](https://img.shields.io/badge/docs-master-blue.svg)](https://tpapp.github.io/DynamicHMC.jl/dev)\n[![DOI](https://zenodo.org/badge/93741413.svg)](https://zenodo.org/badge/latestdoi/93741413)\n[![Aqua QA](https://raw.githubusercontent.com/JuliaTesting/Aqua.jl/master/badge.svg)](https://github.com/JuliaTesting/Aqua.jl)\n\n## Overview\n\nThis package implements a modern version of the “No-U-turn sampler” in the Julia language, mostly as described in [Betancourt (2017)](https://arxiv.org/abs/1701.02434), with some tweaks.\n\nIn contrast to frameworks which utilize a directed acyclic graph to build a posterior for a Bayesian model from small components, this package requires that you code a *log-density function* of the posterior in Julia. Derivatives can be provided manually, or using [automatic differentiation](http://www.juliadiff.org/).\n\nConsequently, this package requires that the user is comfortable with the basics of the theory of Bayesian inference, to the extent of coding a (log) posterior density in Julia. This approach allows the use of standard tools like [profiling](https://docs.julialang.org/en/v1/manual/profile/) and [benchmarking](https://github.com/JuliaCI/BenchmarkTools.jl) to optimize its [performance](https://docs.julialang.org/en/v1/manual/performance-tips/).\n\nThe building blocks of the algorithm are implemented using a *functional* (non-modifying) approach whenever possible, allowing extensive unit testing of components, while at the same time also intended to serve as a transparent, pedagogical introduction to the low-level mechanics of current Hamiltonian Monte Carlo samplers, and as a platform for research into MCMC methods.\n\nPlease start with the [documentation](https://tamaspapp.eu/DynamicHMC.jl/dev/).\n\n## Examples\n\n- Some basic examples are available in [DynamicHMCExamples.jl](https://github.com/tpapp/DynamicHMCExamples.jl).\n\n- [DynamicHMCModels.jl](https://github.com/StatisticalRethinkingJulia/DynamicHMCModels.jl) contains worked examples from the [Statistical Rethinking](https://xcelab.net/rm/statistical-rethinking/) book.\n\n## Support and participation\n\nFor general questions, open an issue or ask on [the Discourse forum](https://discourse.julialang.org/). I am happy to help with models.\n\nUsers who rely on this package and want to participate in discussions are recommended to subscribe to the Github notifications (“watching” the package). Also, I will do my best to accommodate feature requests, just open issues.\n\n## Bibliography\n\nBetancourt, M. J., Byrne, S., & Girolami, M. (2014). Optimizing the integrator step size for Hamiltonian Monte Carlo. [arXiv preprint arXiv:1411.6669](https://arxiv.org/pdf/1411.6669).\n\nBetancourt, M. (2016). Diagnosing suboptimal cotangent disintegrations in Hamiltonian Monte Carlo. [arXiv preprint arXiv:1604.00695](https://arxiv.org/abs/1604.00695).\n\nBetancourt, M. (2017). A Conceptual Introduction to Hamiltonian Monte Carlo. [arXiv preprint arXiv:1701.02434](https://arxiv.org/abs/1701.02434).\n\nGelman, A., Carlin, J. B., Stern, H. S., Dunson, D. B., Vehtari, A., & Rubin, D. B. (2013). Bayesian data analysis. : CRC Press.\n\nGelman, A., & Hill, J. (2007). Data analysis using regression and multilevel/hierarchical models.\n\nHoffman, M. D., & Gelman, A. (2014). The No-U-turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo. Journal of Machine Learning Research, 15(1), 1593-1623.\n\nMcElreath, R. (2018). Statistical rethinking: A Bayesian course with examples in R and Stan. Chapman and Hall/CRC.\n" }, { "path": "appveyor.yml", "content": "environment:\n matrix:\n - julia_version: 1\n - julia_version: nightly\n\nplatform:\n - x86 # 32-bit\n - x64 # 64-bit\n\n# # Uncomment the following lines to allow failures on nightly julia\n# # (tests will run but not make your overall status red)\n# matrix:\n# allow_failures:\n# - julia_version: nightly\n\nbranches:\n only:\n - master\n - /release-.*/\n\nnotifications:\n - provider: Email\n on_build_success: false\n on_build_failure: false\n on_build_status_changed: false\n\ninstall:\n - ps: iex ((new-object net.webclient).DownloadString(\"https://raw.githubusercontent.com/JuliaCI/Appveyor.jl/version-1/bin/install.ps1\"))\n\nbuild_script:\n - echo \"%JL_BUILD_SCRIPT%\"\n - C:\\julia\\bin\\julia -e \"%JL_BUILD_SCRIPT%\"\n\ntest_script:\n - echo \"%JL_TEST_SCRIPT%\"\n - C:\\julia\\bin\\julia -e \"%JL_TEST_SCRIPT%\"\n\n# # Uncomment to support code coverage upload. Should only be enabled for packages\n# # which would have coverage gaps without running on Windows\n# on_success:\n# - echo \"%JL_CODECOV_SCRIPT%\"\n# - C:\\julia\\bin\\julia -e \"%JL_CODECOV_SCRIPT%\"\n" }, { "path": "docs/Project.toml", "content": "[deps]\nDocumenter = \"e30172f5-a6a5-5a46-863b-614d45cd2de4\"\nLogDensityProblems = \"6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c\"\nLogDensityProblemsAD = \"996a588d-648d-4e1f-a8f0-a84b347e47b1\"\nMCMCDiagnosticTools = \"be115224-59cd-429b-ad48-344e309966f0\"\nOhMyThreads = \"67456a42-1dca-4109-a031-0a68de7e3ad5\"\nRandom = \"9a3f8284-a2c9-5f02-9a11-845980a1fd5c\"\nStatistics = \"10745b16-79ce-11e8-11f9-7d13ad32a3b2\"\nTransformVariables = \"84d833dd-6860-57f9-a1a7-6da5db126cff\"\nTransformedLogDensities = \"f9bc47f6-f3f8-4f3b-ab21-f8bc73906f26\"\n\n[compat]\nDocumenter = \"1\"\nLogDensityProblems = \"2\"\n" }, { "path": "docs/make.jl", "content": "using Documenter, DynamicHMC\n\nmakedocs(modules = [DynamicHMC],\n format = Documenter.HTML(; prettyurls = get(ENV, \"CI\", nothing) == \"true\"),\n clean = true,\n sitename = \"DynamicHMC.jl\",\n authors = \"Tamás K. Papp\",\n checkdocs = :exports,\n linkcheck = true,\n linkcheck_ignore = [r\"^.*xcelab\\.net.*$\", r\"^.*stat\\.columbia\\.edu.*$\"],\n pages = [\"Introduction\" => \"index.md\",\n \"A worked example\" => \"worked_example.md\",\n \"Documentation\" => \"interface.md\"])\n\ndeploydocs(repo = \"github.com/tpapp/DynamicHMC.jl.git\",\n push_preview = true)\n" }, { "path": "docs/src/index.md", "content": "# Introduction\n\n[DynamicHMC.jl](https://github.com/tpapp/DynamicHMC.jl/) implements a variant of the “No-U-Turn Sampler” of [Hoffmann and Gelman (2014)](https://arxiv.org/abs/1111.4246), as described in [Betancourt (2017)](https://arxiv.org/abs/1701.02434).[^1] This package is mainly useful for Bayesian inference.\n\n[^1]: In order to make the best use of this package, you should read at least the latter paper thoroughly.\n\nIn order to use it, you should be familiar with the conceptual building blocks of Bayesian inference, most importantly, you should be able to code a (log) posterior as a function in Julia.[^2] The package aims to “[do one thing and do it well](https://en.wikipedia.org/wiki/Unix_philosophy#Do_One_Thing_and_Do_It_Well)”: given a log density function\n\n```math\n\\ell: \\mathbb{R}^k \\to \\mathbb{R}\n```\n\nfor which you have values ``\\ell(x)`` and the gradient ``\\nabla \\ell(x)``, it samples values from a density\n\n```math\np(x) \\propto \\exp(\\ell(x))\n```\n\nusing the algorithm above.\n\n[^2]: For various techniques and a discussion of MCMC methods (eg domain transformations, or integrating out discrete parameters), you may find the [Stan Modeling Language manual](https://mc-stan.org/users/documentation/index.html) helpful. If you are unfamiliar with Bayesian methods, I would recommend [Bayesian Data Analysis](http://www.stat.columbia.edu/~gelman/book/) and [Statistical Rethinking](https://xcelab.net/rm/statistical-rethinking/).\n\nThe interface of DynamicHMC.jl expects that you code ``\\ell(x), \\nabla\\ell(x)`` using the interface of the [LogDensityProblems.jl](https://github.com/tpapp/LogDensityProblems.jl) package. The latter package also allows you to just code ``\\ell(x)`` and obtain ``\\nabla\\ell(x)`` via [automatic differentiation](https://en.wikipedia.org/wiki/Automatic_differentiation).\n\nWhile the NUTS algorithm operates on an *unrestricted* domain ``\\mathbb{R}^k``, some parameters have natural restrictions: for example, standard deviations are positive, valid correlation matrices are a subset of all matrices, and structural econometric models can have parameter restrictions for stability. In order to sample for posteriors with parameters like these, *domain transformations* are required.[^3] Also, it is advantageous to decompose a flat vector `x` to a collection of parameters in a disciplined manner. It is recommended that you use [TransformVariables.jl](https://github.com/tpapp/TransformVariables.jl) in combination with [TransformedLogDensities.jl](https://github.com/tpapp/TransformedLogDensities.jl) for this purpose: it has built-in transformations for common cases, and also allows decomposing vectors into tuples, named tuples, and arrays of parameters, combined with these transformations.\n\n[^3]: For nonlinear transformations, correcting with the logarithm of the determinant of the Jacobian is required.\n\n## Use cases\n\nThis package has the following intended use cases:\n\n1. A robust and simple engine for MCMC. The intended audience is users who like to code their (log)posteriors directly, optimize and benchmark them as Julia code, and at the same time want to have access detailed diagnostic information from the NUTS sampler.\n\n2. A *backend* for another interface that needs a NUTS implementation.\n\n3. A *research platform* for advances in MCMC methods. The code of this package is extensively documented, and should allow extensions and experiments easily using multiple dispatch. Contributions are always welcome.\n\n## Support\n\nIf you have questions, feature requests, or bug reports, please [open an issue](https://github.com/tpapp/DynamicHMC.jl/issues/new). I would like to emphasize that it is perfectly fine to open issues just to ask questions. You can also address questions to [`@Tamas_Papp`](https://discourse.julialang.org/u/Tamas_Papp) on the Julia discourse forum.\n\n## Versioning and interface changes\n\nPackage versioning follows [Semantic Versioning 2.0](https://semver.org/). Only major version increments change the API in a breaking manner, but there is no deprecation cycle. You are strongly advised to add a [compatibility section](https://pkgdocs.julialang.org/dev/compatibility/) to your `Project.toml`, eg\n\n```toml\n[compat]\nDynamicHMC = \"3\"\n```\n\nOnly symbols (functions and types) exported directly or indirectly from the `DynamicHMC` module are considered part of the interface. Importantly, the [`DynamicHMC.Diagnostics`](@ref Diagnostics) submodule is not considered part of the interface with respect to semantic versioning, and may be changed with just a minor version increment. The rationale for this is that a good generic diagnostics interface is much harder to get right, so some experimental improvements, occasionally reverted or redesigned, will be normal for this package in the medium run. If you depend on this explicitly in non-interactive code, use\n\n```toml\n[compat]\nDynamicHMC = \"~3.3.0\"\n```\n\nThere is an actively maintained [CHANGELOG](https://github.com/tpapp/DynamicHMC.jl/blob/master/CHANGELOG.md) which is worth reading after every release, especially major ones.\n" }, { "path": "docs/src/interface.md", "content": "# User interface\n\n```@docs\nDynamicHMC\n```\n\n## Sampling\n\nThe main entry point for sampling is\n\n```@docs\nmcmc_with_warmup\n```\n\n## Utilities for working with the posterior\n\n```@docs\nstack_posterior_matrices\npool_posterior_matrices\n```\n\n## Warmup\n\nWarmup can be customized.\n\n### Default warmup sequences\n\nA warmup sequence is just a tuple of [warmup building blocks](@ref wbb). Two commonly used sequences are predefined.\n\n```@docs\ndefault_warmup_stages\nfixed_stepsize_warmup_stages\n```\n\n### [Warmup building blocks](@id wbb)\n\n```@docs\nInitialStepsizeSearch\nDualAveraging\nTuningNUTS\nGaussianKineticEnergy\n```\n\n## Progress report\n\nProgress reports can be explicit or silent.\n\n```@docs\nNoProgressReport\nLogProgressReport\nProgressMeterReport\n```\n\n## Algorithm and parameters\n\nYou probably won't need to change these options with normal usage, except possibly increasing the maximum tree depth.\n\n```@docs\nDynamicHMC.NUTS\n```\n\n## Inspecting warmup\n\n!!! note\n The warmup interface below is not considered part of the exposed API, and may change with just minor version bump. It is intended for interactive use; the docstrings and the field names of results should be informative.\n\n```@docs\nDynamicHMC.mcmc_keep_warmup\n```\n\n## Stepwise sampling\n\n!!! note\n The stepwise sampling interface below is not considered part of the exposed API, and may change with just minor version bump.\n\nAn experimental interface is available to users who wish to do MCMC one step at a time, eg until some desired criterion about effective sample size or mixing is satisfied. See the docstrings below for an example.\n\n```@docs\nDynamicHMC.mcmc_steps\nDynamicHMC.mcmc_next_step\n```\n\n# Diagnostics\n\n!!! note\n Strictly speaking the `Diagnostics` submodule API is not considered part of the exposed interface, and may change with just minor version bump. It is intended for interactive use.\n\n```@docs\nDynamicHMC.Diagnostics.explore_log_acceptance_ratios\nDynamicHMC.Diagnostics.summarize_tree_statistics\nDynamicHMC.Diagnostics.leapfrog_trajectory\nDynamicHMC.Diagnostics.EBFMI\nDynamicHMC.PhasePoint\n```\n" }, { "path": "docs/src/worked_example.md", "content": "# A worked example\n\n!!! note\n An extended version of this example can be found [in the DynamicHMCExamples.jl package](https://github.com/tpapp/DynamicHMCExamples.jl/blob/master/src/example_independent_bernoulli.jl).\n\n## Problem statement\n\nConsider estimating the parameter ``0 \\le \\alpha \\le 1`` from ``n`` IID observations[^4]\n\n[^4]: Note that NUTS is not especially suitable for low-dimensional parameter spaces, but this example works fine.\n\n```math\ny_i \\sim \\mathrm{Bernoulli}(\\alpha)\n```\nWe will code this with the help of TransformVariables.jl, and obtain the gradient with ForwardDiff.jl (in practice, at the moment I would recommend [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) for small models, and [Flux.jl](https://github.com/FluxML/Flux.jl) for larger ones — consider benchmarking a single evaluation of the log density with gradient).[^5]\n\n[^5]: An example of how you can benchmark a log density with gradient `∇P`, obtained as described below:\n ```julia\n using BenchmarkTools, LogDensityProblems\n x = randn(LogDensityProblems.dimension(∇P))\n @benchmark LogDensityProblems.logdensity_and_gradient($∇P, $x)\n ```\n\n## Coding up the log density\n\nFirst, we load the packages we use.\n\n```@example bernoulli\nusing TransformVariables, TransformedLogDensities, LogDensityProblems, LogDensityProblemsAD,\n DynamicHMC, DynamicHMC.Diagnostics, Statistics, Random\nnothing # hide\n```\n\nGenerally, I would recommend defining an immutable composite type (ie `struct`) to hold the data and all parameters relevant for the log density (eg the prior). This allows you to test your code in a modular way before sampling. For this model, the number of draws equal to `1` is a sufficient statistic.\n\n```@example bernoulli\nstruct BernoulliProblem\n n::Int # Total number of draws in the data\n s::Int # Number of draws `==1` in the data\nend\n```\n\nThen we make this problem *callable* with the parameters. Here, we have a single parameter `α`, but pass this in a `NamedTuple` to demonstrate a generally useful pattern. Then, we define an instance of this problem with the data, called `p`.[^6]\n\n[^6]: Note that here we used a *flat prior*. This is generally not a good idea for variables with non-finite support: one would usually make priors parameters of the `struct` above, and add the log prior to the log likelihood above.\n\n```@example bernoulli\nfunction (problem::BernoulliProblem)(θ)\n (; α) = θ # extract the parameters\n (; n, s) = problem # extract the data\n # log likelihood, with constant terms dropped\n s * log(α) + (n-s) * log(1-α)\nend\n```\n\nIt is generally a good idea to test that your code works by calling it with the parameters; it should return a likelihood. For more complex models, you should benchmark and [optimize](https://docs.julialang.org/en/v1/manual/performance-tips/) this callable directly.\n\n```@example bernoulli\np = BernoulliProblem(20, 10)\np((α = 0.5, )) # make sure that it works\n```\n\nWith [TransformVariables.jl](https://github.com/tpapp/TransformVariables.jl), we set up a *transformation* ``\\mathbb{R} \\to [0,1]`` for ``\\alpha``, and use the convenience function `TransformedLogDensity` to obtain a log density in ``\\mathbb{R}^1``. Finally, we obtain a log density that supports gradients using automatic differentiation, with [LogDensityProblemsAD.jl](https://github.com/tpapp/LogDensityProblemsAD.jl).\n\n```@example bernoulli\ntrans = as((α = as𝕀,))\nP = TransformedLogDensity(trans, p)\n∇P = ADgradient(:ForwardDiff, P)\n```\n\nFinally, we run MCMC with warmup. Note that you have to specify the *random number generator* explicitly — this is good practice for parallel code. The last parameter is the number of samples.\n\n```@example bernoulli\nresults = mcmc_with_warmup(Random.default_rng(), ∇P, 1000; reporter = NoProgressReport())\nnothing # hide\n```\n\nThe returned value is a `NamedTuple`. Most importantly, it contains the field `posterior_matrix`. You should use the transformation you defined above to retrieve the parameters (here, only `α`). We display the mean here to check that it was recovered correctly.\n\n```@example bernoulli\nposterior = transform.(trans, eachcol(results.posterior_matrix))\nposterior_α = first.(posterior)\nmean(posterior_α)\n```\n\nUsing the [`DynamicHMC.Diagnostics`](@ref Diagnostics) submodule, you can obtain various useful diagnostics. The *tree statistics* in particular contain a lot of useful information about turning, divergence, acceptance rates, and tree depths for each step of the chain. Here we just obtain a summary.\n\n```@example bernoulli\nusing DynamicHMC.Diagnostics\nsummarize_tree_statistics(results.tree_statistics)\n```\n\n## Parallel chains and diagnostics\n\nUsually one would run multiple chains and check convergence and mixing using generic MCMC diagnostics not specific to NUTS.\n\nThe specifics of running multiple chains is up to the user: various forms of [parallel computing](https://docs.julialang.org/en/v1/manual/parallel-computing/) can be utilized depending on the problem scale and the hardware available. In the example below we use [multi-threading](https://docs.julialang.org/en/v1/manual/multi-threading/), using [OhMyThreads.jl](https://github.com/JuliaFolds2/OhMyThreads.jl); other excellent packages are available for threading.\n\nIt is easy to obtain posterior results for use with [MCMCDiagnosticsTools.jl](https://github.com/TuringLang/MCMCDiagnosticTools.jl/) with [`stack_posterior_matrices`](@ref):\n\n```@example bernoulli\nusing OhMyThreads, MCMCDiagnosticTools\nresults5 = tcollect(mcmc_with_warmup(Random.default_rng(), ∇P, 1000; reporter = NoProgressReport()) for _ in 1:5)\ness_rhat(stack_posterior_matrices(results5))\n```\n\nUse [`pool_posterior_matrices`](@ref) for a pooled sample:\n\n```@example bernoulli\nposterior5 = transform.(trans, eachcol(pool_posterior_matrices(results5)))\nposterior5_α = first.(posterior5)\nmean(posterior5_α)\n```\n" }, { "path": "src/DynamicHMC.jl", "content": "\"\"\"\nImplementation of the No U-Turn Sampler for MCMC.\n\nPlease [read the documentation](https://tamaspapp.eu/DynamicHMC.jl/dev/). For the\nimpatient: you probably want to\n\n1. define a log density problem (eg for Bayesian inference) using the `LogDensityProblems`\npackage, then\n\n2. use it with [`mcmc_with_warmup`](@ref).\n\"\"\"\nmodule DynamicHMC\n\nusing ArgCheck: @argcheck\nusing DocStringExtensions: FIELDS, FUNCTIONNAME, SIGNATURES, TYPEDEF\nusing FillArrays: Fill\nusing LinearAlgebra: checksquare, cholesky, diag, dot, Diagonal, Symmetric, UniformScaling\nusing LogDensityProblems: capabilities, LogDensityOrder, dimension, logdensity_and_gradient\nusing LogExpFunctions: logaddexp\nusing Random: AbstractRNG, randn, Random, randexp\nusing Statistics: cov, mean, median, middle, quantile, var\nusing TensorCast: @cast, TensorCast\n\ninclude(\"utilities.jl\")\ninclude(\"trees.jl\")\ninclude(\"hamiltonian.jl\")\ninclude(\"stepsize.jl\")\ninclude(\"NUTS.jl\")\ninclude(\"reporting.jl\")\ninclude(\"mcmc.jl\")\ninclude(\"diagnostics.jl\")\n\nend # module\n" }, { "path": "src/NUTS.jl", "content": "#####\n##### NUTS tree sampler implementation.\n#####\n\nexport NUTS\n\n####\n#### Trajectory and implementation\n####\n\n\"\"\"\nRepresentation of a trajectory, ie a Hamiltonian with a discrete integrator that\nalso checks for divergence.\n\"\"\"\nstruct TrajectoryNUTS{TH,Tf,S}\n \"Hamiltonian.\"\n H::TH\n \"Log density of z (negative log energy) at initial point.\"\n π₀::Tf\n \"Stepsize for leapfrog.\"\n ϵ::Tf\n \"Smallest decrease allowed in the log density.\"\n min_Δ::Tf\n \"Turn statistic configuration.\"\n turn_statistic_configuration::S\nend\n\nfunction move(trajectory::TrajectoryNUTS, z, fwd)\n (; H, ϵ) = trajectory\n leapfrog(H, z, fwd ? ϵ : -ϵ)\nend\n\n###\n### proposals\n###\n\n\"\"\"\n$(SIGNATURES)\n\nRandom boolean which is `true` with the given probability `exp(logprob)`, which can be `≥ 1`\nin which case no random value is drawn.\n\"\"\"\nfunction rand_bool_logprob(rng::AbstractRNG, logprob)\n logprob ≥ 0 || (randexp(rng, Float64) > -logprob)\nend\n\nfunction calculate_logprob2(::TrajectoryNUTS, is_doubling, ω₁, ω₂, ω)\n biased_progressive_logprob2(is_doubling, ω₁, ω₂, ω)\nend\n\nfunction combine_proposals(rng, ::TrajectoryNUTS, z₁, z₂, logprob2::Real, is_forward)\n rand_bool_logprob(rng, logprob2) ? z₂ : z₁\nend\n\n###\n### statistics for visited nodes\n###\n\nstruct AcceptanceStatistic{T}\n \"\"\"\n Logarithm of the sum of metropolis acceptances probabilities over the whole trajectory\n (including invalid parts).\n \"\"\"\n log_sum_α::T\n \"Total number of leapfrog steps.\"\n steps::Int\nend\n\nfunction combine_acceptance_statistics(A::AcceptanceStatistic, B::AcceptanceStatistic)\n AcceptanceStatistic(logaddexp(A.log_sum_α, B.log_sum_α), A.steps + B.steps)\nend\n\n\"\"\"\n$(SIGNATURES)\n\nAcceptance statistic for a leaf. The initial leaf is considered not to be visited.\n\"\"\"\nfunction leaf_acceptance_statistic(Δ, is_initial)\n is_initial ? AcceptanceStatistic(oftype(Δ, -Inf), 0) : AcceptanceStatistic(min(Δ, 0), 1)\nend\n\n\"\"\"\n$(SIGNATURES)\n\nReturn the acceptance rate (a `Real` betwen `0` and `1`).\n\"\"\"\nacceptance_rate(A::AcceptanceStatistic) = min(exp(A.log_sum_α) / A.steps, 1)\n\ncombine_visited_statistics(::TrajectoryNUTS, v, w) = combine_acceptance_statistics(v, w)\n\n###\n### turn analysis\n###\n\n\"\"\"\nStatistics for the identification of turning points. See Betancourt (2017, appendix), and\nsubsequent discussion of improvements at\n.\n\nMomenta `p₋` and `p₊` are kept so that they can be added to `ρ` when combining turn\nstatistics.\n\nTurn detection is always done by [`combine_turn_statistics`](@ref), which returns `nothing`\nin case of turning. A `GeneralizedTurnStatistic` should always correspond to a trajectory\nthat is *not* turning (or a leaf node, where the concept does not apply).\n\"\"\"\nstruct GeneralizedTurnStatistic{T}\n \"momentum at the left edge of the trajectory\"\n p₋::T\n \"p♯ at the left edge of the trajectory\"\n p♯₋::T\n \"momentum at the right edge of the trajectory\"\n p₊::T\n \"p♯ at the right edge of the trajectory\"\n p♯₊::T\n \"sum of momenta along trajectory\"\n ρ::T\nend\n\nfunction leaf_turn_statistic(::Val{:generalized}, H, z)\n p♯ = calculate_p♯(H, z)\n GeneralizedTurnStatistic(z.p, p♯, z.p, p♯, z.p)\nend\n\n\"\"\"\n$(SIGNATURES)\n\nInternal test for turning. See Betancourt (2017, appendix).\n\"\"\"\n_is_turning(p♯₋, p♯₊, ρ) = dot(p♯₋, ρ) < 0 || dot(p♯₊, ρ) < 0\n\nfunction combine_turn_statistics(::TrajectoryNUTS,\n x::GeneralizedTurnStatistic, y::GeneralizedTurnStatistic)\n _is_turning(x.p♯₋, y.p♯₋, x.ρ + y.p₋) && return nothing\n _is_turning(x.p♯₊, y.p♯₊, x.p₊ + y.ρ) && return nothing\n ρ = x.ρ + y.ρ\n _is_turning(x.p♯₋, y.p♯₊, ρ) && return nothing\n GeneralizedTurnStatistic(x.p₋, x.p♯₋, y.p₊, y.p♯₊, ρ)\nend\n\nis_turning(::TrajectoryNUTS, ::GeneralizedTurnStatistic) = false\nis_turning(::TrajectoryNUTS, ::Nothing) = true\n\n###\n### leafs\n###\n\nfunction leaf(trajectory::TrajectoryNUTS, z, is_initial)\n (; H, π₀, min_Δ, turn_statistic_configuration) = trajectory\n Δ = is_initial ? zero(π₀) : logdensity(H, z) - π₀\n isdiv = Δ < min_Δ\n v = leaf_acceptance_statistic(Δ, is_initial)\n if isdiv\n nothing, v\n else\n τ = leaf_turn_statistic(turn_statistic_configuration, H, z)\n (z, Δ, τ), v\n end\nend\n\n####\n#### NUTS interface\n####\n\n\"Default maximum depth for trees.\"\nconst DEFAULT_MAX_TREE_DEPTH = 10\n\n\"\"\"\n$(TYPEDEF)\n\nImplementation of the “No-U-Turn Sampler” of Hoffman and Gelman (2014), including subsequent\ndevelopments, as described in Betancourt (2017).\n\n# Fields\n\n$(FIELDS)\n\"\"\"\nstruct NUTS{S}\n \"Maximum tree depth.\"\n max_depth::Int\n \"Threshold for negative energy relative to starting point that indicates divergence.\"\n min_Δ::Float64\n \"\"\"\n Turn statistic configuration. Currently only `Val(:generalized)` (the default) is\n supported.\n \"\"\"\n turn_statistic_configuration::S\n function NUTS(; max_depth = DEFAULT_MAX_TREE_DEPTH, min_Δ = -1000.0,\n turn_statistic_configuration = Val{:generalized}())\n @argcheck 0 < max_depth ≤ MAX_DIRECTIONS_DEPTH\n @argcheck min_Δ < 0\n S = typeof(turn_statistic_configuration)\n new{S}(Int(max_depth), Float64(min_Δ), turn_statistic_configuration)\n end\nend\n\n\"\"\"\n$(TYPEDEF)\n\nDiagnostic information for a single tree built with the No-U-turn sampler.\n\n# Fields\n\nAccessing fields directly is part of the API.\n\n$(FIELDS)\n\"\"\"\nstruct TreeStatisticsNUTS\n \"Log density of the Hamiltonian (negative energy).\"\n π::Float64\n \"Depth of the tree.\"\n depth::Int\n \"Reason for termination. See [`InvalidTree`](@ref) and [`REACHED_MAX_DEPTH`](@ref).\"\n termination::InvalidTree\n \"Acceptance rate statistic.\"\n acceptance_rate::Float64\n \"Number of leapfrog steps evaluated.\"\n steps::Int\n \"Directions for tree doubling (useful for debugging).\"\n directions::Directions\nend\n\n\"\"\"\n$(SIGNATURES)\n\nNo-U-turn Hamiltonian Monte Carlo transition, using Hamiltonian `H`, starting at evaluated\nlog density position `Q`, using stepsize `ϵ`. Parameters of `algorithm` govern the details\nof tree construction.\n\nReturn two values, the new evaluated log density position, and tree statistics.\n\"\"\"\nfunction sample_tree(rng, algorithm::NUTS, H::Hamiltonian, Q::EvaluatedLogDensity, ϵ;\n p = rand_p(rng, H.κ), directions = rand(rng, Directions))\n (; max_depth, min_Δ, turn_statistic_configuration) = algorithm\n z = PhasePoint(Q, p)\n trajectory = TrajectoryNUTS(H, logdensity(H, z), ϵ, min_Δ, turn_statistic_configuration)\n ζ, v, termination, depth = sample_trajectory(rng, trajectory, z, max_depth, directions)\n tree_statistics = TreeStatisticsNUTS(logdensity(H, ζ), depth, termination,\n acceptance_rate(v), v.steps, directions)\n ζ.Q, tree_statistics\nend\n" }, { "path": "src/diagnostics.jl", "content": "#####\n##### statistics and diagnostics for the NUTS-specific tree statistics.\n#####\n##### For general posterior diagnostics, see `stack_posterior_matrices`.\n\nmodule Diagnostics\n\nexport EBFMI, summarize_tree_statistics, explore_log_acceptance_ratios, leapfrog_trajectory,\n InvalidTree, REACHED_MAX_DEPTH, is_divergent\n\nusing DynamicHMC: GaussianKineticEnergy, Hamiltonian, evaluate_ℓ, InvalidTree,\n REACHED_MAX_DEPTH, is_divergent, local_log_acceptance_ratio, PhasePoint, rand_p, leapfrog,\n logdensity, MAX_DIRECTIONS_DEPTH\n\nusing ArgCheck: @argcheck\nusing DocStringExtensions: FIELDS, SIGNATURES, TYPEDEF\nusing LogDensityProblems: dimension\nimport Random\nusing Statistics: mean, quantile, var\n\n\"\"\"\n$(SIGNATURES)\n\nEnergy Bayesian fraction of missing information. Useful for diagnosing poorly\nchosen kinetic energies.\n\nLow values (`≤ 0.3`) are considered problematic. See Betancourt (2016).\n\"\"\"\nfunction EBFMI(tree_statistics)\n πs = map(x -> x.π, tree_statistics)\n mean(abs2, diff(πs)) / var(πs)\nend\n\n\"Acceptance quantiles for [`TreeStatisticsSummary`](@ref) diagnostic summary.\"\nconst ACCEPTANCE_QUANTILES = [0.05, 0.25, 0.5, 0.75, 0.95]\n\n\"\"\"\n$(TYPEDEF)\n\nStoring the output of [`NUTS_statistics`](@ref) in a structured way, for pretty\nprinting. Currently for internal use.\n\n# Fields\n\n$(FIELDS)\n\"\"\"\nstruct TreeStatisticsSummary{T <: Real, C <: NamedTuple}\n \"Sample length.\"\n N::Int\n \"average_acceptance\"\n a_mean::T\n \"acceptance quantiles\"\n a_quantiles::Vector{T}\n \"termination counts\"\n termination_counts::C\n \"depth counts (first element is for `0`)\"\n depth_counts::Vector{Int}\nend\n\n\"\"\"\n$(SIGNATURES)\n\nCount termination reasons in `tree_statistics`, return as a `NamedTuple`.\n\"\"\"\nfunction count_terminations(tree_statistics)\n max_depth = 0\n divergence = 0\n turning = 0\n for tree_statistic in tree_statistics\n it = tree_statistic.termination\n if it == REACHED_MAX_DEPTH\n max_depth += 1\n elseif is_divergent(it)\n divergence += 1\n else\n turning += 1\n end\n end\n (; max_depth, divergence, turning)\nend\n\n\"\"\"\n$(SIGNATURES)\n\nCount depths in tree statistics.\n\"\"\"\nfunction count_depths(tree_statistics)\n c = zeros(Int, MAX_DIRECTIONS_DEPTH + 1)\n for tree_statistic in tree_statistics\n c[tree_statistic.depth + 1] += 1\n end\n c[1:something(findlast(!iszero, c), 0)]\nend\n\n\"\"\"\n$(SIGNATURES)\n\nSummarize tree statistics. Mostly useful for NUTS diagnostics.\n\"\"\"\nfunction summarize_tree_statistics(tree_statistics)\n As = map(x -> x.acceptance_rate, tree_statistics)\n TreeStatisticsSummary(length(tree_statistics),\n mean(As), quantile(As, ACCEPTANCE_QUANTILES),\n count_terminations(tree_statistics),\n count_depths(tree_statistics))\nend\n\nfunction Base.show(io::IO, stats::TreeStatisticsSummary)\n (; N, a_mean, a_quantiles, termination_counts, depth_counts) = stats\n println(io, \"Hamiltonian Monte Carlo sample of length $(N)\")\n print(io, \" acceptance rate mean: $(round(a_mean; digits = 2)), 5/25/50/75/95%:\")\n for aq in a_quantiles\n print(io, \" \", round(aq; digits = 2))\n end\n println(io)\n function print_percentages(pairs)\n is_first = true\n for (key, value) in sort(collect(pairs), by = first)\n if is_first\n is_first = false\n else\n print(io, \",\")\n end\n print(io, \" $(key) => $(round(Int, 100*value/N))%\")\n end\n end\n print(io, \" termination:\")\n print_percentages(pairs(termination_counts))\n println(io)\n print(io, \" depth:\")\n print_percentages(zip(axes(depth_counts, 1) .- 1, depth_counts))\nend\n\n####\n#### Acceptance ratio diagnostics\n####\n\n\"\"\"\n$(SIGNATURES)\n\nFrom an initial position, calculate the uncapped log acceptance ratio for the given log2\nstepsizes and momentums `ps`, `N` of which are generated randomly by default.\n\"\"\"\nfunction explore_log_acceptance_ratios(ℓ, q, log2ϵs;\n rng = Random.default_rng(),\n κ = GaussianKineticEnergy(dimension(ℓ)),\n N = 20, ps = [rand_p(rng, κ) for _ in 1:N])\n H = Hamiltonian(κ, ℓ)\n Q = evaluate_ℓ(ℓ, q)\n As = [local_log_acceptance_ratio(H, PhasePoint(Q, p)) for p in ps]\n [A(2.0^log2ϵ) for log2ϵ in log2ϵs, A in As]\nend\n\n\"\"\"\n$(TYPEDEF)\n\nImplements an iterator on a leapfrog trajectory until the first non-finite log density.\n\n# Fields\n\n$(FIELDS)\n\"\"\"\nstruct LeapfrogTrajectory{TH,TZ,TF,Tϵ}\n \"Hamiltonian\"\n H::TH\n \"Initial position\"\n z₀::TZ\n \"Negative energy at initial position.\"\n π₀::TF\n \"Stepsize (negative: move backward).\"\n ϵ::Tϵ\nend\n\nBase.IteratorSize(::Type{<:LeapfrogTrajectory}) = Base.SizeUnknown()\n\nfunction Base.iterate(lft::LeapfrogTrajectory, zi = (lft.z₀, 0))\n (; H, ϵ, π₀) = lft\n z, i = zi\n if isfinite(z.Q.ℓq)\n z′ = leapfrog(H, z, ϵ)\n i′ = i + sign(ϵ)\n _position_information(lft, z′, i′), (z′, i′)\n else\n nothing\n end\nend\n\n\"\"\"\n$(SIGNATURES)\n\nPosition information returned by [`leapfrog_trajectory`](@ref), see documentation there.\nInternal function.\n\"\"\"\nfunction _position_information(lft::LeapfrogTrajectory, z, i)\n (; H, π₀) = lft\n (z = z, position = i, Δ = logdensity(H, z) - π₀)\nend\n\n\"\"\"\n$(SIGNATURES)\n\nCalculate a leapfrog trajectory visiting `positions` (specified as a `UnitRange`, eg `-5:5`)\nrelative to the starting point `q`, with stepsize `ϵ`. `positions` has to contain `0`, and\nthe trajectories are only tracked up to the first non-finite log density in each direction.\n\nReturns a vector of `NamedTuple`s, each containin\n\n- `z`, a [`PhasePoint`](@ref) object,\n\n- `position`, the corresponding position,\n\n- `Δ`, the log density + the kinetic energy relative to position `0`.\n\"\"\"\nfunction leapfrog_trajectory(ℓ, q, ϵ, positions::UnitRange{<:Integer};\n rng = Random.default_rng(),\n κ = GaussianKineticEnergy(dimension(ℓ)), p = rand_p(rng, κ))\n A, B = first(positions), last(positions)\n @argcheck A ≤ 0 ≤ B \"Positions has to contain `0`.\"\n Q = evaluate_ℓ(ℓ, q)\n H = Hamiltonian(κ, ℓ)\n z₀ = PhasePoint(Q, p)\n π₀ = logdensity(H, z₀)\n lft_fwd = LeapfrogTrajectory(H, z₀, π₀, ϵ)\n fwd_part = collect(Iterators.take(lft_fwd, B))\n bwd_part = collect(Iterators.take(LeapfrogTrajectory(H, z₀, π₀, -ϵ), -A))\n vcat(reverse!(bwd_part), _position_information(lft_fwd, z₀, 0), fwd_part)\nend\n\nend\n" }, { "path": "src/hamiltonian.jl", "content": "#####\n##### Building blocks for traversing a Hamiltonian deterministically, using the leapfrog\n##### integrator.\n#####\n\nexport GaussianKineticEnergy\n\n####\n#### kinetic energy\n####\n\n\"\"\"\n$(TYPEDEF)\n\nKinetic energy specifications. Implements the methods\n\n- `Base.size`\n\n- [`kinetic_energy`](@ref)\n\n- [`calculate_p♯`](@ref)\n\n- [`∇kinetic_energy`](@ref)\n\n- [`rand_p`](@ref)\n\nFor all subtypes, it is implicitly assumed that kinetic energy is symmetric in\nthe momentum `p`,\n\n```julia\nkinetic_energy(κ, p, q) == kinetic_energy(κ, .-p, q)\n```\n\nWhen the above is violated, the consequences are undefined.\n\"\"\"\nabstract type KineticEnergy end\n\n\"\"\"\n$(TYPEDEF)\n\nEuclidean kinetic energies (position independent).\n\"\"\"\nabstract type EuclideanKineticEnergy <: KineticEnergy end\n\n\"\"\"\n$(TYPEDEF)\n\nGaussian kinetic energy, with ``K(q,p) = p ∣ q ∼ 1/2 pᵀ⋅M⁻¹⋅p + log|M|`` (without constant),\nwhich is independent of ``q``.\n\nThe inverse covariance ``M⁻¹`` is stored.\n\n!!! note\n Making ``M⁻¹`` approximate the posterior variance is a reasonable starting point.\n\"\"\"\nstruct GaussianKineticEnergy{T <: AbstractMatrix,\n S <: AbstractMatrix} <: EuclideanKineticEnergy\n \"M⁻¹\"\n M⁻¹::T\n \"W such that W*W'=M. Used for generating random draws.\"\n W::S\n function GaussianKineticEnergy(M⁻¹::T, W::S) where {T, S}\n @argcheck checksquare(M⁻¹) == checksquare(W)\n new{T,S}(M⁻¹, W)\n end\nend\n\n\"\"\"\n$(SIGNATURES)\n\nGaussian kinetic energy with the given inverse covariance matrix `M⁻¹`.\n\"\"\"\nGaussianKineticEnergy(M⁻¹::AbstractMatrix) = GaussianKineticEnergy(M⁻¹, cholesky(inv(M⁻¹)).L)\n\n\"\"\"\n$(SIGNATURES)\n\nGaussian kinetic energy with the given inverse covariance matrix `M⁻¹`.\n\"\"\"\nGaussianKineticEnergy(M⁻¹::Diagonal) = GaussianKineticEnergy(M⁻¹, Diagonal(.√inv.(diag(M⁻¹))))\n\n\"\"\"\n$(SIGNATURES)\n\nGaussian kinetic energy with a diagonal inverse covariance matrix `M⁻¹=m⁻¹*I`.\n\"\"\"\nGaussianKineticEnergy(N::Integer, m⁻¹ = 1.0) = GaussianKineticEnergy(Diagonal(Fill(m⁻¹, N)))\n\nfunction Base.show(io::IO, κ::GaussianKineticEnergy{T}) where {T}\n print(io::IO, \"Gaussian kinetic energy ($(nameof(T))), √diag(M⁻¹): $(.√(diag(κ.M⁻¹)))\")\nend\n\n## NOTE about implementation: the 3 methods are callable without a third argument (`q`)\n## because they are defined for Gaussian (Euclidean) kinetic energies.\n\nBase.size(κ::GaussianKineticEnergy, args...) = size(κ.M⁻¹, args...)\n\n\"\"\"\n$(SIGNATURES)\n\nReturn kinetic energy `κ`, at momentum `p`.\n\"\"\"\nkinetic_energy(κ::GaussianKineticEnergy, p, q = nothing) = dot(p, κ.M⁻¹ * p) / 2\n\n\"\"\"\n$(SIGNATURES)\n\nReturn ``p♯ = M⁻¹⋅p``, used for turn diagnostics.\n\"\"\"\ncalculate_p♯(κ::GaussianKineticEnergy, p, q = nothing) = κ.M⁻¹ * p\n\n\"\"\"\n$(SIGNATURES)\n\nCalculate the gradient of the logarithm of kinetic energy in momentum `p`.\n\"\"\"\n∇kinetic_energy(κ::GaussianKineticEnergy, p, q = nothing) = calculate_p♯(κ, p)\n\n\"\"\"\n$(SIGNATURES)\n\nGenerate a random momentum from a kinetic energy at position `q`.\n\"\"\"\nrand_p(rng::AbstractRNG, κ::GaussianKineticEnergy, q = nothing) = κ.W * randn(rng, size(κ.W, 1))\n\n####\n#### Hamiltonian\n####\n\nstruct Hamiltonian{K,L}\n \"The kinetic energy specification.\"\n κ::K\n \"\"\"\n The (log) density we are sampling from. Supports the `LogDensityProblem` API.\n Technically, it is the negative of the potential energy.\n \"\"\"\n ℓ::L\n \"\"\"\n $(SIGNATURES)\n\n Construct a Hamiltonian from the log density `ℓ`, and the kinetic energy specification\n `κ`. `ℓ` with a vector are expected to support the `LogDensityProblems` API, with\n gradients.\n \"\"\"\n function Hamiltonian(κ::K, ℓ::L) where {K <: KineticEnergy,L}\n @argcheck capabilities(ℓ) ≥ LogDensityOrder(1)\n @argcheck dimension(ℓ) == size(κ, 1)\n new{K,L}(κ, ℓ)\n end\nend\n\nBase.show(io::IO, H::Hamiltonian) = print(io, \"Hamiltonian with $(H.κ)\")\n\n\"\"\"\n$(TYPEDEF)\n\nA log density evaluated at position `q`. The log densities and gradient are saved, so that\nthey are not calculated twice for every leapfrog step (both as start- and endpoints).\n\nBecause of caching, a `EvaluatedLogDensity` should only be used with a specific Hamiltonian,\npreferably constructed with the `evaluate_ℓ` constructor.\n\nIn composite types and arguments, `Q` is usually used for this type.\n\"\"\"\nstruct EvaluatedLogDensity{T,S}\n \"Position.\"\n q::T\n \"ℓ(q). Saved for reuse in sampling.\"\n ℓq::S\n \"∇ℓ(q). Cached for reuse in sampling.\"\n ∇ℓq::T\n function EvaluatedLogDensity(q::T, ℓq::S, ∇ℓq::T) where {T <: AbstractVector,S <: Real}\n @argcheck length(q) == length(∇ℓq)\n new{T,S}(q, ℓq, ∇ℓq)\n end\nend\n\n# general constructors below are necessary to sanitize input from eg Diagnostics, or an\n# initial position given as integers, etc\n\nfunction EvaluatedLogDensity(q::AbstractVector, ℓq::Real, ∇ℓq::AbstractVector)\n q, ∇ℓq = promote(q, ∇ℓq)\n EvaluatedLogDensity(q, ℓq, ∇ℓq)\nend\n\nEvaluatedLogDensity(q, ℓq::Real, ∇ℓq) = EvaluatedLogDensity(collect(q), ℓq, collect(∇ℓq))\n\n\"\"\"\n$(SIGNATURES)\n\nEvaluate log density and gradient and save with the position. Preferred interface for\ncreating `EvaluatedLogDensity` instances.\n\nNon-finite elements in `q` always throw an error.\n\nNon-finite and not `-Inf` elements in the log density throw an error if `strict`, otherwise\nreplace the log density with `-Inf`.\n\nNon-finite elements in the gradient throw an error if `strict`, otherwise replace\nthe log density with `-Inf`.\n\"\"\"\nfunction evaluate_ℓ(ℓ, q; strict::Bool = false)\n all(isfinite, q) || _error(\"Position vector has non-finite elements.\"; q)\n ℓq, ∇ℓq = logdensity_and_gradient(ℓ, q)\n if (isfinite(ℓq) && all(isfinite, ∇ℓq)) || ℓq == -Inf\n # everything is finite, or log density is -Inf, which will be rejected\n EvaluatedLogDensity(q, ℓq, ∇ℓq)\n elseif !strict\n # something went wrong, but proceed and replace log density with -Inf, so it is\n # rejected.\n EvaluatedLogDensity(q, oftype(ℓq, -Inf), ∇ℓq) # somew\n elseif isfinite(ℓq)\n _error(\"Gradient has non-finite elements.\"; q, ∇ℓq)\n else\n _error(\"Invalid log posterior.\"; q, ℓq)\n end\nend\n\n\"\"\"\n$(TYPEDEF)\n\nA point in phase space, consists of a position (in the form of an evaluated log density `ℓ`\nat `q`) and a momentum.\n\"\"\"\nstruct PhasePoint{T <: EvaluatedLogDensity,S}\n \"Evaluated log density.\"\n Q::T\n \"Momentum.\"\n p::S\n function PhasePoint(Q::T, p::S) where {T,S}\n @argcheck length(p) == length(Q.q)\n new{T,S}(Q, p)\n end\nend\n\n\"\"\"\n$(SIGNATURES)\n\nLog density for Hamiltonian `H` at point `z`.\n\nIf `ℓ(q) == -Inf` (rejected), skips the kinetic energy calculation.\n\nNon-finite values (incl `NaN`, `Inf`) are automatically converted to `-Inf`. This can happen\nif\n\n1. the log density is not a finite value,\n\n2. the kinetic energy is not a finite value (which usually happens when `NaN` or `Inf` got\nmixed in the leapfrog step, leading to an invalid position).\n\"\"\"\nfunction logdensity(H::Hamiltonian{<:EuclideanKineticEnergy}, z::PhasePoint)\n (; ℓq) = z.Q\n isfinite(ℓq) || return oftype(ℓq, -Inf)\n K = kinetic_energy(H.κ, z.p)\n ℓq - (isfinite(K) ? K : oftype(K, Inf))\nend\n\nfunction calculate_p♯(H::Hamiltonian{<:EuclideanKineticEnergy}, z::PhasePoint)\n calculate_p♯(H.κ, z.p)\nend\n\n\"\"\"\n leapfrog(H, z, ϵ)\n\nTake a leapfrog step of length `ϵ` from `z` along the Hamiltonian `H`.\n\nReturn the new phase point.\n\nThe leapfrog algorithm uses the gradient of the next position to evolve the momentum. If\nthis is not finite, the momentum won't be either, `logdensity` above will catch this and\nreturn an `-Inf`, making the point divergent.\n\"\"\"\nfunction leapfrog(H::Hamiltonian{<: EuclideanKineticEnergy}, z::PhasePoint, ϵ)\n (; ℓ, κ) = H\n (; p, Q) = z\n @argcheck isfinite(Q.ℓq) \"Internal error: leapfrog called from non-finite log density\"\n pₘ = p + ϵ/2 * Q.∇ℓq\n q′ = Q.q + ϵ * ∇kinetic_energy(κ, pₘ)\n Q′ = evaluate_ℓ(H.ℓ, q′)\n p′ = pₘ + ϵ/2 * Q′.∇ℓq\n PhasePoint(Q′, p′)\nend\n" }, { "path": "src/mcmc.jl", "content": "#####\n##### Sampling: high-level interface and building blocks\n#####\n\nexport InitialStepsizeSearch, DualAveraging, TuningNUTS, mcmc_with_warmup,\n default_warmup_stages, fixed_stepsize_warmup_stages, stack_posterior_matrices,\n pool_posterior_matrices\n\n\"Significant digits to display for reporting.\"\nconst REPORT_SIGDIGITS = 3\n\n####\n#### docstrings for reuse (not exported, internal)\n####\n\nconst _DOC_POSTERIOR_MATRIX =\n \"`posterior_matrix`, a matrix of position vectors, indexes by `[parameter_index, draw_index]`\"\n\nconst _DOC_TREE_STATISTICS =\n \"`tree_statistics`, a vector of tree statistics for each sample\"\n\nconst _DOC_EPSILONS = \"`ϵs`, a vector of step sizes for each sample\"\n\nconst _DOC_LOGDENSITIES =\n \"`logdensities`, a vector of log densities, each associated with a column of the posterior matrix\"\n\n####\n#### parts unaffected by warmup\n####\n\n\"\"\"\n$(TYPEDEF)\n\nA log density bundled with an RNG and options for sampling. Contains the parts of the\nproblem which are **not changed during warmup** (and thus the whole sampling).\n\n# Fields\n\n$(FIELDS)\n\"\"\"\nstruct SamplingLogDensity{R,L,O,S}\n \"Random number generator.\"\n rng::R\n \"Log density.\"\n ℓ::L\n \"\"\"\n Algorithm used for sampling, also contains the relevant parameters that are not affected\n by adaptation. See eg [`NUTS`](@ref).\n \"\"\"\n algorithm::O\n \"Reporting warmup information and chain progress.\"\n reporter::S\nend\n\n####\n#### warmup building blocks\n####\n\n###\n### warmup state\n###\n\n\"\"\"\n$(TYPEDEF)\n\nRepresentation of a warmup state. Not part of the API.\n\n# Fields\n\n$(FIELDS)\n\"\"\"\nstruct WarmupState{TQ <: EvaluatedLogDensity,Tκ <: KineticEnergy, Tϵ <: Union{Real,Nothing}}\n \"phasepoint\"\n Q::TQ\n \"kinetic energy\"\n κ::Tκ\n \"stepsize\"\n ϵ::Tϵ\nend\n\nfunction Base.show(io::IO, warmup_state::WarmupState)\n (; κ, ϵ) = warmup_state\n ϵ_display = ϵ ≡ nothing ? \"unspecified\" : \"≈ $(round(ϵ; sigdigits = REPORT_SIGDIGITS))\"\n print(io, \"adapted sampling parameters: stepsize (ϵ) $(ϵ_display), $(κ)\")\nend\n\n###\n### warmup interface and stages\n###\n\n\"\"\"\n$(SIGNATURES)\n\nReturn the *results* and the *next warmup state* after warming up/adapting according to\n`warmup_stage`, starting from `warmup_state`.\n\nUse `nothing` for a no-op.\n\"\"\"\nfunction warmup(sampling_logdensity::SamplingLogDensity, warmup_stage::Nothing, warmup_state)\n nothing, warmup_state\nend\n\n\"\"\"\n$(SIGNATURES)\n\nHelper function to create random starting positions in the `[-2,2]ⁿ` box.\n\"\"\"\nrandom_position(rng, N) = rand(rng, N) .* 4 .- 2\n\n\"Docstring for initial warmup arguments.\"\nconst DOC_INITIAL_WARMUP_ARGS =\n\"\"\"\n- `q`: initial position. *Default*: random (uniform [-2,2] for each coordinate).\n\n- `κ`: kinetic energy specification. *Default*: Gaussian with identity matrix.\n\n- `ϵ`: a scalar for initial stepsize, or `nothing` for heuristic finders.\n\"\"\"\n\n\"\"\"\n$(SIGNATURES)\n\nCreate an initial warmup state from a random position.\n\n# Keyword arguments\n\n$(DOC_INITIAL_WARMUP_ARGS)\n\"\"\"\nfunction initialize_warmup_state(rng, ℓ; q = random_position(rng, dimension(ℓ)),\n κ = GaussianKineticEnergy(dimension(ℓ)), ϵ = nothing)\n WarmupState(evaluate_ℓ(ℓ, q; strict = true), κ, ϵ)\nend\n\nfunction warmup(sampling_logdensity, stepsize_search::InitialStepsizeSearch, warmup_state)\n (; rng, ℓ, reporter) = sampling_logdensity\n (; Q, κ, ϵ) = warmup_state\n @argcheck ϵ ≡ nothing \"stepsize ϵ manually specified, won't perform initial search\"\n z = PhasePoint(Q, rand_p(rng, κ))\n try\n ϵ = find_initial_stepsize(stepsize_search, local_log_acceptance_ratio(Hamiltonian(κ, ℓ), z))\n report(reporter, \"found initial stepsize\",\n ϵ = round(ϵ; sigdigits = REPORT_SIGDIGITS))\n nothing, WarmupState(Q, κ, ϵ)\n catch e\n @info \"failed to find initial stepsize\" q = z.Q.q p = z.p κ\n rethrow(e)\n end\nend\n\n\"\"\"\n$(TYPEDEF)\n\nTune the step size `ϵ` during sampling, and the metric of the kinetic energy at the end of\nthe block. The method for the latter is determined by the type parameter `M`, which can be\n\n1. `Diagonal` for diagonal metric (the default),\n\n2. `Symmetric` for a dense metric,\n\n3. `Nothing` for an unchanged metric.\n\n# Results\n\nA `NamedTuple` with the following fields:\n\n- $(_DOC_POSTERIOR_MATRIX)\n\n- $(_DOC_TREE_STATISTICS)\n\n- $(_DOC_EPSILONS)\n\n- $(_DOC_LOGDENSITIES)\n\n# Fields\n\n$(FIELDS)\n\"\"\"\nstruct TuningNUTS{M,D}\n \"Number of samples.\"\n N::Int\n \"Dual averaging parameters.\"\n stepsize_adaptation::D\n \"\"\"\n Regularization factor for normalizing variance. An estimated covariance matrix `Σ` is\n rescaled by `λ` towards ``σ²I``, where ``σ²`` is the median of the diagonal. The\n constructor has a reasonable default.\n \"\"\"\n λ::Float64\n function TuningNUTS{M}(N::Integer, stepsize_adaptation::D,\n λ = 5.0/N) where {M <: Union{Nothing,Diagonal,Symmetric},D}\n @argcheck N ≥ 20 # variance estimator is kind of meaningless for few samples\n @argcheck λ ≥ 0\n new{M,D}(N, stepsize_adaptation, λ)\n end\nend\n\nfunction Base.show(io::IO, tuning::TuningNUTS{M}) where {M}\n (; N, stepsize_adaptation, λ) = tuning\n print(io, \"Stepsize and metric tuner, $(N) samples, $(M) metric, regularization $(λ)\")\nend\n\n\"\"\"\n$(SIGNATURES)\n\nEstimate the inverse metric from the chain.\n\nIn most cases, this should be regularized, see [`regularize_M⁻¹`](@ref).\n\"\"\"\nsample_M⁻¹(::Type{Diagonal}, posterior_matrix) = Diagonal(vec(var(posterior_matrix; dims = 2)))\n\nsample_M⁻¹(::Type{Symmetric}, posterior_matrix) = Symmetric(cov(posterior_matrix; dims = 2))\n\n\"\"\"\n$(SIGNATURES)\n\nAdjust the inverse metric estimated from the sample, using an *ad-hoc* shrinkage method.\n\"\"\"\nfunction regularize_M⁻¹(Σ::Symmetric, λ::Real)\n d = diag(Σ)\n (1 - λ) * Σ + λ * Diagonal(d) # ad-hoc “shrinkage”\nend\n\nregularize_M⁻¹(Σ::Union{Diagonal}, λ::Real) = Σ\n\n\"\"\"\n$(SIGNATURES)\n\nCreate an empty posterior matrix, based on `Q` (a logdensity evaluated at a position).\n\"\"\"\n_empty_posterior_matrix(Q, N) = Matrix{eltype(Q.q)}(undef, length(Q.q), N)\n\n\"\"\"\n$(SIGNATURES)\n\nCreate an empty vector for log densities, based on `Q` (a logdensity evaluated at a\nposition).\n\"\"\"\n_empty_logdensity_vector(Q, N) = Vector{typeof(Q.ℓq)}(undef, N)\n\n\"\"\"\n$(SIGNATURES)\n\nPerform a warmup on a given `sampling_logdensity`, using the specified `tuning`, starting\nfrom `warmup_state`.\n\nReturn two values. The first is either `nothing`, or a `NamedTuple` of\n\n- $(_DOC_POSTERIOR_MATRIX)\n\n- $(_DOC_TREE_STATISTICS)\n\n- $(_DOC_EPSILONS)\n\n- $(_DOC_LOGDENSITIES)\n\nThe second is the warmup state.\n\"\"\"\nfunction warmup(sampling_logdensity, tuning::TuningNUTS{M}, warmup_state) where {M}\n (; rng, ℓ, algorithm, reporter) = sampling_logdensity\n (; Q, κ, ϵ) = warmup_state\n (; N, stepsize_adaptation, λ) = tuning\n posterior_matrix = _empty_posterior_matrix(Q, N)\n logdensities = _empty_logdensity_vector(Q, N)\n tree_statistics = Vector{TreeStatisticsNUTS}(undef, N)\n H = Hamiltonian(κ, ℓ)\n ϵ_state = initial_adaptation_state(stepsize_adaptation, ϵ)\n ϵs = Vector{Float64}(undef, N)\n mcmc_reporter = make_mcmc_reporter(reporter, N;\n currently_warmup = true,\n tuning = M ≡ Nothing ? \"stepsize\" : \"stepsize and $(M) metric\")\n for i in 1:N\n ϵ = current_ϵ(ϵ_state)\n ϵs[i] = ϵ\n Q, stats = sample_tree(rng, algorithm, H, Q, ϵ)\n posterior_matrix[:, i] = Q.q\n logdensities[i] = Q.ℓq\n tree_statistics[i] = stats\n ϵ_state = adapt_stepsize(stepsize_adaptation, ϵ_state, stats.acceptance_rate)\n report(mcmc_reporter, i; ϵ = round(ϵ; sigdigits = REPORT_SIGDIGITS))\n end\n if M ≢ Nothing\n κ = GaussianKineticEnergy(regularize_M⁻¹(sample_M⁻¹(M, posterior_matrix), λ))\n report(mcmc_reporter, \"adaptation finished\", adapted_kinetic_energy = κ)\n end\n ((; posterior_matrix, tree_statistics, ϵs, logdensities), WarmupState(Q, κ, final_ϵ(ϵ_state)))\nend\n\n\"\"\"\n$(TYPEDEF)\n\nA composite type for performing MCMC stepwise after warmup.\n\nThe type is *not* part of the API, see [`mcmc_steps`](@ref) and [`mcmc_next_step`](@ref).\n\"\"\"\nstruct MCMCSteps{TR,TA,TH,TE}\n rng::TR\n algorithm::TA\n H::TH\n ϵ::TE\nend\n\n\"\"\"\n$(SIGNATURES)\n\nReturn a value which can be used to perform MCMC stepwise, eg until some criterion is\nsatisfied about the sample. See [`mcmc_next_step`](@ref).\n\nTwo constructors are available:\n\n1. Explicitly providing\n - `rng`, the random number generator,\n - `algorithm`, see [`mcmc_with_warmup`](@ref),\n - `κ`, the (adapted) metric,\n - `ℓ`, the log density callable (see [`mcmc_with_warmup`](@ref),\n - `ϵ`, the stepsize.\n\n2. Using the fields `sampling_logdensity` and `warmup_state`, eg from\n [`mcmc_keep_warmup`](@ref) (make sure you use eg `final_warmup_state`).\n\n# Example\n\n```julia\n# initialization\nresults = DynamicHMC.mcmc_keep_warmup(RNG, ℓ, 0; reporter = NoProgressReport())\nsteps = mcmc_steps(results.sampling_logdensity, results.final_warmup_state)\nQ = results.final_warmup_state.Q\n\n# a single update step\nQ, tree_stats = mcmc_next_step(steps, Q)\n\n# extract the position\nQ.q\n```\n\"\"\"\nmcmc_steps(rng, algorithm, κ, ℓ, ϵ) = MCMCSteps(rng, algorithm, Hamiltonian(κ, ℓ), ϵ)\n\nfunction mcmc_steps(sampling_logdensity::SamplingLogDensity, warmup_state)\n (; rng, ℓ, algorithm) = sampling_logdensity\n (; κ, ϵ) = warmup_state\n mcmc_steps(rng, algorithm, κ, ℓ, ϵ)\nend\n\n\"\"\"\n$(SIGNATURES)\n\nGiven `Q` (an evaluated log density at a position), return the next `Q` and tree statistics.\n\"\"\"\nfunction mcmc_next_step(mcmc_steps::MCMCSteps, Q::EvaluatedLogDensity)\n (; rng, algorithm, H, ϵ) = mcmc_steps\n sample_tree(rng, algorithm, H, Q, ϵ)\nend\n\n\"\"\"\n$(SIGNATURES)\n\nMarkov Chain Monte Carlo for `sampling_logdensity`, with the adapted `warmup_state`.\n\nReturn a `NamedTuple` of\n\n- $(_DOC_POSTERIOR_MATRIX)\n\n- $(_DOC_TREE_STATISTICS)\n\n- $(_DOC_LOGDENSITIES)\n\"\"\"\nfunction mcmc(sampling_logdensity, N, warmup_state)\n (; reporter) = sampling_logdensity\n (; Q) = warmup_state\n posterior_matrix = _empty_posterior_matrix(Q, N)\n logdensities = _empty_logdensity_vector(Q, N)\n tree_statistics = Vector{TreeStatisticsNUTS}(undef, N)\n mcmc_reporter = make_mcmc_reporter(reporter, N; currently_warmup = false)\n steps = mcmc_steps(sampling_logdensity, warmup_state)\n for i in 1:N\n Q, tree_statistics[i] = mcmc_next_step(steps, Q)\n posterior_matrix[:, i] = Q.q\n logdensities[i] = Q.ℓq\n report(mcmc_reporter, i)\n end\n (; posterior_matrix, tree_statistics, logdensities)\nend\n\n\"\"\"\n$(SIGNATURES)\n\nHelper function for constructing the “middle” doubling warmup stages in\n[`default_warmup_stages`](@ref).\n\"\"\"\nfunction _doubling_warmup_stages(M, stepsize_adaptation, middle_steps,\n doubling_stages::Val{D}) where {D}\n ntuple(i -> TuningNUTS{M}(middle_steps * 2^(i - 1), stepsize_adaptation), D)\nend\n\n\"\"\"\n$(SIGNATURES)\n\nA sequence of warmup stages:\n\n1. select an initial stepsize using `stepsize_search` (default: based on a heuristic),\n\n2. tuning stepsize with `init_steps` steps\n\n3. tuning stepsize and covariance: first with `middle_steps` steps, then repeat with twice\n the steps `doubling_stages` times\n\n4. tuning stepsize with `terminating_steps` steps.\n\n`M` (`Diagonal`, the default or `Symmetric`) determines the type of the metric adapted from\nthe sample.\n\nThis is the suggested tuner of most applications.\n\nUse `nothing` for `stepsize_adaptation` to skip the corresponding step.\n\"\"\"\nfunction default_warmup_stages(;\n stepsize_search = InitialStepsizeSearch(),\n M::Type{<:Union{Diagonal,Symmetric}} = Diagonal,\n stepsize_adaptation = DualAveraging(),\n init_steps = 75, middle_steps = 25, doubling_stages = 5,\n terminating_steps = 50)\n (stepsize_search,\n TuningNUTS{Nothing}(init_steps, stepsize_adaptation),\n _doubling_warmup_stages(M, stepsize_adaptation, middle_steps, Val(doubling_stages))...,\n TuningNUTS{Nothing}(terminating_steps, stepsize_adaptation))\nend\n\n\"\"\"\n$(SIGNATURES)\n\nA sequence of warmup stages for fixed stepsize, only tuning covariance: first with\n`middle_steps` steps, then repeat with twice the steps `doubling_stages` times\n\nVery similar to [`default_warmup_stages`](@ref), but omits the warmup stages with just\nstepsize tuning.\n\"\"\"\nfunction fixed_stepsize_warmup_stages(;\n M::Type{<:Union{Diagonal,Symmetric}} = Diagonal,\n middle_steps = 25, doubling_stages = 5)\n _doubling_warmup_stages(M, FixedStepsize(), middle_steps, Val(doubling_stages))\nend\n\n\"\"\"\n$(SIGNATURES)\n\nHelper function for implementing warmup.\n\n!!! note\n Changes may imply documentation updates in [`mcmc_keep_warmup`](@ref).\n\"\"\"\nfunction _warmup(sampling_logdensity, stages, initial_warmup_state)\n foldl(stages; init = ((), initial_warmup_state)) do acc, stage\n stages_and_results, warmup_state = acc\n results, warmup_state′ = warmup(sampling_logdensity, stage, warmup_state)\n stage_information = (stage, results, warmup_state = warmup_state′)\n (stages_and_results..., stage_information), warmup_state′\n end\nend\n\n\"Shared docstring part for the MCMC API.\"\nconst DOC_MCMC_ARGS =\n\"\"\"\n# Arguments\n\n- `rng`: the random number generator, eg `Random.default_rng()`.\n\n- `ℓ`: the log density, supporting the API of the `LogDensityProblems` package\n\n- `N`: the number of samples for inference, after the warmup.\n\n# Keyword arguments\n\n- `initialization`: see below.\n\n- `warmup_stages`: a sequence of warmup stages. See [`default_warmup_stages`](@ref) and\n [`fixed_stepsize_warmup_stages`](@ref); the latter requires an `ϵ` in initialization.\n\n- `algorithm`: see [`NUTS`](@ref). It is very unlikely you need to modify\n this, except perhaps for the maximum depth.\n\n- `reporter`: how progress is reported. By default, verbosely for interactive sessions using\n the log message mechanism (see [`LogProgressReport`](@ref), and no reporting for\n non-interactive sessions (see [`NoProgressReport`](@ref)).\n\n# Initialization\n\nThe `initialization` keyword argument should be a `NamedTuple` which can contain the\nfollowing fields (all of them optional and provided with reasonable defaults):\n\n$(DOC_INITIAL_WARMUP_ARGS)\n\"\"\"\n\n\"\"\"\n$(SIGNATURES)\n\nPerform MCMC with NUTS, keeping the warmup results. Returns a `NamedTuple` of\n\n- `initial_warmup_state`, which contains the initial warmup state\n\n- `warmup`, an iterable of `NamedTuple`s each containing fields\n\n - `stage`: the relevant warmup stage\n\n - `results`: results returned by that warmup stage (may be `nothing` if not applicable,\n or a chain, with tree statistics, etc; see the documentation of stages)\n\n - `warmup_state`: the warmup state *after* the corresponding stage.\n\n- `final_warmup_state`, which contains the final adaptation after all the warmup\n\n- `inference`, which has `posterior_matrix` and `tree_statistics`, see\n [`mcmc_with_warmup`](@ref).\n\n- `sampling_logdensity`, which contains information that is invariant to warmup\n\n!!! warning\n This function is not (yet) exported because the the warmup interface may change with\n minor versions without being considered breaking. Recommended for interactive use.\n\n$(DOC_MCMC_ARGS)\n\"\"\"\nfunction mcmc_keep_warmup(rng::AbstractRNG, ℓ, N::Integer;\n initialization = (),\n warmup_stages = default_warmup_stages(),\n algorithm = NUTS(),\n reporter = default_reporter())\n sampling_logdensity = SamplingLogDensity(rng, ℓ, algorithm, reporter)\n initial_warmup_state = initialize_warmup_state(rng, ℓ; initialization...)\n warmup, warmup_state = _warmup(sampling_logdensity, warmup_stages, initial_warmup_state)\n inference = mcmc(sampling_logdensity, N, warmup_state)\n (; initial_warmup_state, warmup, final_warmup_state = warmup_state, inference,\n sampling_logdensity)\nend\n\n\"\"\"\n$(SIGNATURES)\n\nPerform MCMC with NUTS, including warmup which is not returned. Return a `NamedTuple` of\n\n- $(_DOC_POSTERIOR_MATRIX)\n\n- $(_DOC_TREE_STATISTICS)\n\n- `κ` and `ϵ`, the adapted metric and stepsize.\n\n$(DOC_MCMC_ARGS)\n\n# Usage examples\n\nUsing a fixed stepsize:\n```julia\nmcmc_with_warmup(rng, ℓ, N;\n initialization = (ϵ = 0.1, ),\n warmup_stages = fixed_stepsize_warmup_stages())\n```\n\nStarting from a given position `q₀` and kinetic energy scaled down (will still be adapted):\n```julia\nmcmc_with_warmup(rng, ℓ, N;\n initialization = (q = q₀, κ = GaussianKineticEnergy(5, 0.1)))\n```\n\nUsing a dense metric:\n```julia\nmcmc_with_warmup(rng, ℓ, N;\n warmup_stages = default_warmup_stages(; M = Symmetric))\n```\n\nDisabling the initial stepsize search (provided explicitly, still adapted):\n```julia\nmcmc_with_warmup(rng, ℓ, N;\n initialization = (ϵ = 1.0, ),\n warmup_stages = default_warmup_stages(; stepsize_search = nothing))\n```\n\"\"\"\nfunction mcmc_with_warmup(rng, ℓ, N; initialization = (),\n warmup_stages = default_warmup_stages(),\n algorithm = NUTS(), reporter = default_reporter())\n (; final_warmup_state, inference) =\n mcmc_keep_warmup(rng, ℓ, N; initialization = initialization,\n warmup_stages = warmup_stages, algorithm = algorithm,\n reporter = reporter)\n (; κ, ϵ) = final_warmup_state\n (; inference..., κ, ϵ)\nend\n\n####\n#### utilities\n####\n\n\"\"\"\n$(SIGNATURES)\n\nGiven a vector of `results`, each containing a property `posterior_matrix` (eg obtained from\n[`mcmc_with_warmup`](@ref) with the same sample length), return a lazy view as an array\nindexed by `[draw_index, chain_index, parameter_index]`.\n\nThis is useful as an input for eg `MCMCDiagnosticTools.ess_rhat`.\n\n!!! note\n The ordering is not compatible with MCMCDiagnostictools version < 0.2.\n\"\"\"\nfunction stack_posterior_matrices(results)\n @cast _[i, k, j]:= results[k].posterior_matrix[j, i]\nend\n\n\"\"\"\n$(SIGNATURES)\n\nGiven a vector of `results`, each containing a property `posterior_matrix` (eg obtained from\n[`mcmc_with_warmup`](@ref) with the same sample length), return a lazy view as an array\nindexed by `[parameter_index, pooled_draw_index]`.\n\nThis is useful for posterior analysis after diagnostics (see eg `Base.eachcol`).\n\"\"\"\nfunction pool_posterior_matrices(results)\n @cast _[i, j ⊗ k] := results[k].posterior_matrix[i, j]\nend\n" }, { "path": "src/reporting.jl", "content": "#####\n##### Reporting progress.\n#####\n\nimport ProgressMeter\n\nexport NoProgressReport, LogProgressReport, ProgressMeterReport\n\n\"\"\"\n$(TYPEDEF)\n\nA placeholder type for not reporting any information.\n\"\"\"\nstruct NoProgressReport end\n\n\"\"\"\n$(SIGNATURES)\n\nReport to the given `reporter`.\n\nThe second argument can be\n\n1. a string, which is displayed as is (this is supported by all reporters).\n\n2. or a step in an MCMC chain with a known number of steps for progress reporters (see\n[`make_mcmc_reporter`](@ref)).\n\n`meta` arguments are key-value pairs.\n\nIn this context, a *step* is a NUTS transition, not a leapfrog step.\n\"\"\"\nreport(reporter::NoProgressReport, step::Union{AbstractString,Integer}; meta...) = nothing\n\n\"\"\"\n$(SIGNATURES)\n\nReturn a reporter which can be used for progress reports with a known number of\n`total_steps`. May return the same reporter, or a related object. Will display `meta` as\nkey-value pairs.\n\n## Arguments:\n- `reporter::NoProgressReport`: the original reporter\n- `total_steps`: total number of steps\n\n## Keyword arguments:\n- `currently_warmup::Bool`: `true` if we are currently doing warmup; `false` if we are currently doing MCMC\n- `meta`: key-value pairs that will be displayed by the reporter\n\"\"\"\nmake_mcmc_reporter(reporter::NoProgressReport, total_steps; currently_warmup::Bool = false, meta...) = reporter\n\n\"\"\"\n$(TYPEDEF)\n\nReport progress into the `Logging` framework, using `@info`.\n\nFor the information reported, a *step* is a NUTS transition, not a leapfrog step.\n\n# Fields\n\n$(FIELDS)\n\"\"\"\nBase.@kwdef struct LogProgressReport{T}\n \"ID of chain. Can be an arbitrary object, eg `nothing`.\"\n chain_id::T = nothing\n \"Always report progress past `step_interval` of the last report.\"\n step_interval::Int = 100\n \"Always report progress past this much time (in seconds) after the last report.\"\n time_interval_s::Float64 = 1000.0\nend\n\n\"\"\"\n$(SIGNATURES)\n\nAssemble log message metadata.\n\nCurrently, it adds `chain_id` *iff* it is not `nothing`.\n\"\"\"\n_log_meta(chain_id::Nothing, meta) = meta\n\n_log_meta(chain_id, meta) = (; chain_id, meta...)\n\nfunction report(reporter::LogProgressReport, message::AbstractString; meta...)\n @info message _log_meta(reporter.chain_id, meta)...\n nothing\nend\n\n\"\"\"\n$(TYPEDEF)\n\nA composite type for tracking the state for which the last log message was emitted, for MCMC\nreporting with a given total number of steps (see [`make_mcmc_reporter`](@ref).\n\n# Fields\n\n$(FIELDS)\n\"\"\"\nmutable struct LogMCMCReport{T}\n \"The progress report sink.\"\n log_progress_report::T\n \"Total steps for this stage.\"\n total_steps::Int\n \"Index of the last reported step.\"\n last_reported_step::Int\n \"The last time a report was logged (determined using `time_ns`).\"\n last_reported_time_ns::UInt64\nend\n\nfunction report(reporter::LogMCMCReport, message::AbstractString; meta...)\n @info message _log_meta(reporter.log_progress_report.chain_id, meta)...\n nothing\nend\n\nfunction make_mcmc_reporter(reporter::LogProgressReport, total_steps::Integer; currently_warmup::Bool = false, meta...)\n @info \"Starting MCMC\" total_steps = total_steps meta...\n LogMCMCReport(reporter, total_steps, -1, time_ns())\nend\n\nfunction report(reporter::LogMCMCReport, step::Integer; meta...)\n (; log_progress_report, total_steps, last_reported_step, last_reported_time_ns) = reporter\n (; chain_id, step_interval, time_interval_s) = log_progress_report\n @argcheck 1 ≤ step ≤ total_steps\n Δ_steps = step - last_reported_step\n t_ns = time_ns()\n Δ_time_s = (t_ns - last_reported_time_ns) / 1_000_000_000\n if last_reported_step < 0 || Δ_steps ≥ step_interval || Δ_time_s ≥ time_interval_s\n seconds_per_step = Δ_time_s / Δ_steps\n meta_progress = (step,\n seconds_per_step = round(seconds_per_step; sigdigits = 2),\n estimated_seconds_left = round((total_steps - step) *\n seconds_per_step; sigdigits = 2))\n @info \"MCMC progress\" merge(_log_meta(chain_id, meta_progress), meta)...\n reporter.last_reported_step = step\n reporter.last_reported_time_ns = t_ns\n end\n nothing\nend\n\n\"\"\"\n$(TYPEDEF)\n\nReport progress via a progress bar, using `ProgressMeter.jl`.\n\nExample usage:\n```julia\njulia> ProgressMeterReport()\n```\n\"\"\"\nstruct ProgressMeterReport\nend\n\nstruct ProgressMeterReportMCMC{T}\n currently_warmup::Bool\n progress_meter::T\nend\n\nfunction make_mcmc_reporter(reporter::ProgressMeterReport, total_steps::Integer; currently_warmup::Bool=false, meta...)\n description = currently_warmup ? \"Warmup: \" : \"MCMC: \"\n return ProgressMeterReportMCMC(currently_warmup, ProgressMeter.Progress(total_steps, 1, description))\nend\n\nfunction report(reporter::ProgressMeterReport, message::AbstractString; meta...)\n return nothing\nend\n\nfunction report(reporter::ProgressMeterReportMCMC, message::AbstractString; meta...)\n return nothing\nend\n\nfunction report(reporter::ProgressMeterReport, step::Integer; meta...)\n return nothing\nend\n\nfunction report(reporter::ProgressMeterReportMCMC, step::Integer; meta...)\n ProgressMeter.next!(reporter.progress_meter)\n return nothing\nend\n\n\"\"\"\n$(SIGNATURES)\n\nReturn a default reporter, taking the environment into account. Keyword arguments are passed\nto constructors when applicable.\n\"\"\"\nfunction default_reporter(; kwargs...)\n if isinteractive()\n LogProgressReport(; kwargs...)\n else\n NoProgressReport()\n end\nend\n" }, { "path": "src/stepsize.jl", "content": "#####\n##### stepsize heuristics and adaptation\n#####\n\n####\n#### initial stepsize\n####\n\n\"\"\"\n$(TYPEDEF)\n\nParameters for the search algorithm for the initial stepsize.\n\nThe algorithm finds an initial stepsize ``ϵ`` so that the local log acceptance ratio\n``A(ϵ)`` is near `params.log_threshold`.\n\n$FIELDS\n\n!!! NOTE\n\n The algorithm is from Hoffman and Gelman (2014), default threshold modified to `0.8` following later practice in Stan.\n\"\"\"\nstruct InitialStepsizeSearch\n \"The stepsize where the search is started.\"\n initial_ϵ::Float64\n \"Log of the threshold that needs to be crossed.\"\n log_threshold::Float64\n \"Maximum number of iterations for crossing the threshold.\"\n maxiter_crossing::Int\n function InitialStepsizeSearch(; log_threshold::Float64 = log(0.8), initial_ϵ = 0.1, maxiter_crossing = 400)\n @argcheck isfinite(log_threshold) && log_threshold < 0\n @argcheck isfinite(initial_ϵ) && 0 < initial_ϵ\n @argcheck maxiter_crossing ≥ 50\n new(initial_ϵ, log_threshold, maxiter_crossing)\n end\nend\n\n\"\"\"\n$(SIGNATURES)\n\nFind an initial stepsize that matches the conditions of `parameters` (see\n[`InitialStepsizeSearch`](@ref)).\n\n`A` is the local log acceptance ratio (uncapped). Cf [`local_log_acceptance_ratio`](@ref).\n\"\"\"\nfunction find_initial_stepsize(parameters::InitialStepsizeSearch, A)\n (; initial_ϵ, log_threshold, maxiter_crossing) = parameters\n ϵ = initial_ϵ\n Aϵ = A(ϵ)\n double = Aϵ > log_threshold # do we double?\n for _ in 1:maxiter_crossing\n ϵ′ = double ? 2 * ϵ : ϵ / 2\n Aϵ′ = A(ϵ′)\n (double ? Aϵ′ < log_threshold : Aϵ′ > log_threshold) && return ϵ′\n ϵ = ϵ′\n end\n dir = double ? \"below\" : \"above\"\n _error(\"Initial stepsize search reached maximum number of iterations from $(dir) without crossing.\";\n maxiter_crossing, initial_ϵ, ϵ)\nend\n\n\"\"\"\n$(SIGNATURES)\n\nReturn a function of the stepsize (``ϵ``) that calculates the local log acceptance\nratio for a single leapfrog step around `z` along the Hamiltonian `H`. Formally,\nlet\n\n```julia\nA(ϵ) = logdensity(H, leapfrog(H, z, ϵ)) - logdensity(H, z)\n```\n\nNote that the ratio is not capped by `0`, so it is not a valid (log) probability *per se*.\n\"\"\"\nfunction local_log_acceptance_ratio(H, z)\n ℓ0 = logdensity(H, z)\n isfinite(ℓ0) ||\n _error(\"Starting point has non-finite density.\";\n hamiltonian_logdensity = ℓ0, logdensity = z.Q.ℓq, position = z.Q.q)\n function(ϵ)\n z1 = leapfrog(H, z, ϵ)\n ℓ1 = logdensity(H, z1)\n ℓ1 - ℓ0\n end\nend\n\n\"\"\"\n$(TYPEDEF)\n\nParameters for the dual averaging algorithm of Gelman and Hoffman (2014, Algorithm 6).\n\nTo get reasonable defaults, initialize with `DualAveraging()`.\n\n# Fields\n\n$(FIELDS)\n\"\"\"\nstruct DualAveraging{T}\n \"target acceptance rate\"\n δ::T\n \"regularization scale\"\n γ::T\n \"relaxation exponent\"\n κ::T\n \"offset\"\n t₀::Int\n function DualAveraging(δ::T, γ::T, κ::T, t₀::Int) where {T <: Real}\n @argcheck 0 < δ < 1\n @argcheck γ > 0\n @argcheck 0.5 < κ ≤ 1\n @argcheck t₀ ≥ 0\n new{T}(δ, γ, κ, t₀)\n end\nend\n\nfunction DualAveraging(; δ = 0.8, γ = 0.05, κ = 0.75, t₀ = 10)\n DualAveraging(promote(δ, γ, κ)..., t₀)\nend\n\n\"Current state of adaptation for `ϵ`.\"\nBase.@kwdef struct DualAveragingState{T <: AbstractFloat}\n μ::T\n m::Int\n H̄::T\n logϵ::T\n logϵ̄::T\nend\n\n\"\"\"\n$(SIGNATURES)\n\nReturn an initial adaptation state for the adaptation method and a stepsize `ϵ`.\n\"\"\"\nfunction initial_adaptation_state(::DualAveraging, ϵ)\n @argcheck ϵ > 0\n logϵ = log(ϵ)\n DualAveragingState(; μ = log(10) + logϵ, m = 1, H̄ = zero(logϵ), logϵ, logϵ̄ = zero(logϵ))\nend\n\n\"\"\"\n$(SIGNATURES)\n\nUpdate the adaptation `A` of log stepsize `logϵ` with average Metropolis acceptance rate `a`\nover the whole visited trajectory, using the dual averaging algorithm of Gelman and Hoffman\n(2014, Algorithm 6). Return the new adaptation state.\n\"\"\"\nfunction adapt_stepsize(parameters::DualAveraging, A::DualAveragingState, a)\n @argcheck 0 ≤ a ≤ 1\n (; δ, γ, κ, t₀) = parameters\n (; μ, m, H̄, logϵ, logϵ̄) = A\n m += 1\n H̄ += (δ - a - H̄) / (m + t₀)\n logϵ = μ - √m/γ * H̄\n logϵ̄ += m^(-κ)*(logϵ - logϵ̄)\n DualAveragingState(; μ, m, H̄, logϵ, logϵ̄)\nend\n\n\"\"\"\n$(SIGNATURES)\n\nReturn the stepsize `ϵ` for the next HMC step while adapting.\n\"\"\"\ncurrent_ϵ(A::DualAveragingState, tuning = true) = exp(A.logϵ)\n\n\"\"\"\n$(SIGNATURES)\n\nReturn the final stepsize `ϵ` after adaptation.\n\"\"\"\nfinal_ϵ(A::DualAveragingState, tuning = true) = exp(A.logϵ̄)\n\n###\n### fixed stepsize adaptation placeholder\n###\n\n\"\"\"\n$(SIGNATURES)\n\nAdaptation with fixed stepsize. Leaves `ϵ` unchanged.\n\"\"\"\nstruct FixedStepsize end\n\ninitial_adaptation_state(::FixedStepsize, ϵ) = ϵ\n\nadapt_stepsize(::FixedStepsize, ϵ, a) = ϵ\n\ncurrent_ϵ(ϵ::Real) = ϵ\n\nfinal_ϵ(ϵ::Real) = ϵ\n" }, { "path": "src/trees.jl", "content": "#####\n##### Abstract tree/trajectory interface\n#####\n\n####\n#### Directions\n####\n\n\"Maximum number of iterations [`next_direction`](@ref) supports.\"\nconst MAX_DIRECTIONS_DEPTH = 32\n\n\"\"\"\nInternal type implementing random directions. Draw a new value with `rand`, see\n[`next_direction`](@ref).\n\nServes two purposes: a fixed value of `Directions` is useful for unit testing, and drawing a\nsingle bit flag collection economizes on the RNG cost.\n\"\"\"\nstruct Directions\n flags::UInt32\nend\n\nBase.rand(rng::AbstractRNG, ::Type{Directions}) = Directions(rand(rng, UInt32))\n\n\"\"\"\n$(SIGNATURES)\n\nReturn the next direction flag and the new state of directions. Results are undefined for\nmore than [`MAX_DIRECTIONS_DEPTH`](@ref) updates.\n\"\"\"\nfunction next_direction(directions::Directions)\n (; flags) = directions\n Bool(flags & 0x01), Directions(flags >>> 1)\nend\n\n####\n#### Trajectory interface\n####\n\n\"\"\"\n $(FUNCTIONNAME)(trajectory, z, is_forward)\n\nMove along the trajectory in the specified direction. Return the new position.\n\"\"\"\nfunction move end\n\n\"\"\"\n $(FUNCTIONNAME)(trajectory, τ)\n\nTest if the turn statistics `τ` indicate that the corresponding tree is turning.\n\nWill only be called on nontrivial trees (at least two nodes).\n\"\"\"\nfunction is_turning end\n\n\"\"\"\n $(FUNCTIONNAME)(trajectory, τ₁, τ₂)\n\nCombine turn statistics on trajectory. Implementation can assume that the trees that\ncorrespond to the turn statistics have the same ordering.\n\nWhen\n```julia\nτ = combine_turn_statistics(trajectory, τ₁, τ₂)\nis_turning(trajectory, τ)\n```\nthe combined turn statistic `τ` is guaranteed not to escape the caller, so it can eg change\ntype.\n\"\"\"\nfunction combine_turn_statistics end\n\n\"\"\"\n $(FUNCTIONNAME)(trajectory, v₁, v₂)\n\nCombine visited node statistics for adjacent trees trajectory. Implementation should be\ninvariant to the ordering of `v₁` and `v₂` (ie the operation is commutative).\n\"\"\"\nfunction combine_visited_statistics end\n\n\"\"\"\n $(FUNCTIONNAME)(trajectory, is_doubling::Bool, ω₁, ω₂, ω)\n\nCalculate the log probability if selecting the subtree corresponding to `ω₂`. Being the log\nof a probability, it is always `≤ 0`, but implementations are allowed to return and accept\nvalues `> 0` and treat them as `0`.\n\nWhen `is_doubling`, the tree corresponding to `ω₂` was obtained from a doubling step (this\ncan be relevant eg for biased progressive sampling).\n\nThe value `ω = logaddexp(ω₁, ω₂)` is provided for avoiding redundant calculations.\n\nSee [`biased_progressive_logprob2`](@ref) for an implementation.\n\"\"\"\nfunction calculate_logprob2 end\n\n\"\"\"\n $(FUNCTIONNAME)(rng, trajectory, ζ₁, ζ₂, logprob2::Real, is_forward::Bool)\n\nCombine two proposals `ζ₁, ζ₂` on `trajectory`, with log probability `logprob2` for\nselecting `ζ₂`.\n\n `ζ₁` is before `ζ₂` iff `is_forward`.\n\"\"\"\nfunction combine_proposals end\n\n\"\"\"\n ζωτ_or_nothing, v = $(FUNCTIONNAME)(trajectory, z, is_initial)\n\nInformation for a tree made of a single node. When `is_initial == true`, this is the first\nnode.\n\nThe first value is either\n\n1. `nothing` for a divergent node,\n\n2. a tuple containing the proposal `ζ`, the log weight (probability) of the node `ω`, the\nturn statistics `τ` (never tested with `is_turning` for leaf nodes).\n\nThe second value is the visited node information.\n\"\"\"\nfunction leaf end\n\n####\n#### utilities\n####\n\n\"\"\"\n$(SIGNATURES)\n\nCombine turn statistics with the given direction. When `is_forward`, `τ₁` is before `τ₂`,\notherwise after.\n\nInternal helper function.\n\"\"\"\n@inline function combine_turn_statistics_in_direction(trajectory, τ₁, τ₂, is_forward::Bool)\n if is_forward\n combine_turn_statistics(trajectory, τ₁, τ₂)\n else\n combine_turn_statistics(trajectory, τ₂, τ₁)\n end\nend\n\nfunction combine_proposals_and_logweights(rng, trajectory, ζ₁, ζ₂, ω₁::Real, ω₂::Real,\n is_forward::Bool, is_doubling::Bool)\n ω = logaddexp(ω₁, ω₂)\n logprob2 = calculate_logprob2(trajectory, is_doubling, ω₁, ω₂, ω)\n ζ = combine_proposals(rng, trajectory, ζ₁, ζ₂, logprob2, is_forward)\n ζ, ω\nend\n\n\"\"\"\n$(SIGNATURES)\n\nGiven (relative) log probabilities `ω₁` and `ω₂`, return the log probabiliy of\ndrawing a sample from the second (`logprob2`).\n\nWhen `bias`, biases towards the second argument, introducing anti-correlations.\n\"\"\"\nfunction biased_progressive_logprob2(bias::Bool, ω₁::Real, ω₂::Real, ω = logaddexp(ω₁, ω₂))\n ω₂ - (bias ? ω₁ : ω)\nend\n\n####\n#### abstract trajectory interface\n####\n\n\"\"\"\n$(SIGNATURES)\n\nInformation about an invalid (sub)tree, using positions relative to the starting node.\n\n1. When `left < right`, this tree was *turning*.\n\n2. When `left == right`, this is a *divergent* node.\n\n3. `left == 1 && right == 0` is used as a sentinel value for reaching maximum depth without\nencountering any invalid trees (see [`REACHED_MAX_DEPTH`](@ref). All other `left > right`\nvalues are disallowed.\n\"\"\"\nstruct InvalidTree\n left::Int\n right::Int\nend\n\nInvalidTree(i::Integer) = InvalidTree(i, i)\n\nis_divergent(invalid_tree::InvalidTree) = invalid_tree.left == invalid_tree.right\n\nfunction Base.show(io::IO, invalid_tree::InvalidTree)\n msg = if is_divergent(invalid_tree)\n \"divergence at position $(invalid_tree.left)\"\n elseif invalid_tree == REACHED_MAX_DEPTH\n \"reached maximum depth without divergence or turning\"\n else\n (; left, right) = invalid_tree\n \"turning at positions $(left):$(right)\"\n end\n print(io, msg)\nend\n\n\"Sentinel value for reaching maximum depth.\"\nconst REACHED_MAX_DEPTH = InvalidTree(1, 0)\n\n\"\"\"\n result, v = adjacent_tree(rng, trajectory, z, i, depth, is_forward)\n\nTraverse the tree of given `depth` adjacent to point `z` in `trajectory`.\n\n`is_forward` specifies the direction, `rng` is used for random numbers in\n[`combine_proposals`](@ref). `i` is an integer position relative to the initial node (`0`).\n\nThe *first value* is either\n\n1. an `InvalidTree`, indicating the first divergent node or turning subtree that was\nencounteted and invalidated this tree.\n\n2. a tuple of `(ζ, ω, τ, z′, i′), with\n\n - `ζ`: the proposal from the tree.\n\n - `ω`: the log weight of the subtree that corresponds to the proposal\n\n - `τ`: turn statistics\n\n - `z′`: the last node of the tree\n\n - `i′`: the position of the last node relative to the initial node.\n\nThe *second value* is always the visited node statistic.\n\"\"\"\nfunction adjacent_tree(rng, trajectory, z, i, depth, is_forward)\n i′ = i + (is_forward ? 1 : -1)\n if depth == 0\n z′ = move(trajectory, z, is_forward)\n ζωτ, v = leaf(trajectory, z′, false)\n if ζωτ ≡ nothing\n InvalidTree(i′), v\n else\n (ζωτ..., z′, i′), v\n end\n else\n # “left” tree\n t₋, v₋ = adjacent_tree(rng, trajectory, z, i, depth - 1, is_forward)\n t₋ isa InvalidTree && return t₋, v₋\n ζ₋, ω₋, τ₋, z₋, i₋ = t₋\n\n # “right” tree — visited information from left is kept even if invalid\n t₊, v₊ = adjacent_tree(rng, trajectory, z₋, i₋, depth - 1, is_forward)\n v = combine_visited_statistics(trajectory, v₋, v₊)\n t₊ isa InvalidTree && return t₊, v\n ζ₊, ω₊, τ₊, z₊, i₊ = t₊\n\n # turning invalidates\n τ = combine_turn_statistics_in_direction(trajectory, τ₋, τ₊, is_forward)\n is_turning(trajectory, τ) && return InvalidTree(i′, i₊), v\n\n # valid subtree, combine proposals\n ζ, ω = combine_proposals_and_logweights(rng, trajectory, ζ₋, ζ₊, ω₋, ω₊,\n is_forward, false)\n (ζ, ω, τ, z₊, i₊), v\n end\nend\n\n\"\"\"\n$(SIGNATURES)\n\nSample a `trajectory` starting at `z`, up to `max_depth`. `directions` determines the tree\nexpansion directions.\n\nReturn the following values\n\n- `ζ`: proposal from the tree\n\n- `v`: visited node statistics\n\n- `termination`: an `InvalidTree` (this includes the last doubling step turning, which is\n technically a valid tree) or `REACHED_MAX_DEPTH` when all subtrees were valid and no\n turning happens.\n\n- `depth`: the depth of the tree that was sampled from. Doubling steps that lead to an\n invalid adjacent tree do not contribute to `depth`.\n\"\"\"\nfunction sample_trajectory(rng, trajectory, z, max_depth::Integer, directions::Directions)\n @argcheck max_depth ≤ MAX_DIRECTIONS_DEPTH\n (ζ, ω, τ), v = leaf(trajectory, z, true)\n z₋ = z₊ = z\n depth = 0\n termination = REACHED_MAX_DEPTH\n i₋ = i₊ = 0\n while depth < max_depth\n is_forward, directions = next_direction(directions)\n t′, v′ = adjacent_tree(rng, trajectory, is_forward ? z₊ : z₋, is_forward ? i₊ : i₋,\n depth, is_forward)\n v = combine_visited_statistics(trajectory, v, v′)\n\n # invalid adjacent tree: stop\n t′ isa InvalidTree && (termination = t′; break)\n\n # extract information from adjacent tree\n ζ′, ω′, τ′, z′, i′ = t′\n\n # update edges and combine proposals\n if is_forward\n z₊, i₊ = z′, i′\n else\n z₋, i₋ = z′, i′\n end\n\n # tree has doubled successfully\n ζ, ω = combine_proposals_and_logweights(rng, trajectory, ζ, ζ′, ω, ω′,\n is_forward, true)\n depth += 1\n\n # when the combined tree is turning, stop\n τ = combine_turn_statistics_in_direction(trajectory, τ, τ′, is_forward)\n is_turning(trajectory, τ) && (termination = InvalidTree(i₋, i₊); break)\n end\n ζ, v, termination, depth\nend\n" }, { "path": "src/utilities.jl", "content": "#####\n##### utilities\n#####\n\n####\n#### error messages\n####\n\n\"\"\"\n$(TYPEDEF)\n\nThe error type used by this package. Debug information should be printed without truncation,\nwith full precision.\n\n$(FIELDS)\n\"\"\"\nstruct DynamicHMCError <: Exception\n message::String\n debug_information::NamedTuple\nend\n\n\"\"\"\n$(SIGNATURES)\n\nThrow a `DynamicHMCError` with given message, keyword arguments used for debug information.\n\"\"\"\n_error(message::AbstractString; kwargs...) = throw(DynamicHMCError(message, NamedTuple(kwargs)))\n\nfunction Base.showerror(io::IO, error::DynamicHMCError)\n (; message, debug_information) = error\n printstyled(io, \"DynamicHMC error: \", error; color = :red)\n for (key, value) in pairs(debug_information)\n print(io, \"\\n \")\n printstyled(io, string(key); color = :blue, bold = true)\n printstyled(io, \" = \", value)\n end\n nothing\nend\n" }, { "path": "test/Project.toml", "content": "[deps]\nAqua = \"4c88cf16-eb10-579e-8560-4a9242c79595\"\nArgCheck = \"dce04be8-c92d-5529-be00-80e4d2c0e197\"\nDocStringExtensions = \"ffbed154-4ef7-542d-bbb7-c09d3a79fcae\"\nDynamicHMC = \"bbc10e6e-7c05-544b-b16e-64fede858acb\"\nForwardDiff = \"f6369f11-7733-5829-9624-2563aa707210\"\nHypothesisTests = \"09f84164-cd44-5f33-b23f-e6b0d136a0d5\"\nJET = \"c3a54625-cd67-489e-a8e7-0a5a0ff4e31b\"\nLinearAlgebra = \"37e2e46d-f89d-539d-b4ee-838fcccc9c8e\"\nLogDensityProblems = \"6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c\"\nLogDensityTestSuite = \"feb245ec-c857-584e-a66a-22324acf10c6\"\nLogExpFunctions = \"2ab3a3ac-af41-5b50-aa03-7779005ae688\"\nLogging = \"56ddb016-857b-54e1-b83d-db4d58db5568\"\nMCMCDiagnosticTools = \"be115224-59cd-429b-ad48-344e309966f0\"\nOhMyThreads = \"67456a42-1dca-4109-a031-0a68de7e3ad5\"\nRandom = \"9a3f8284-a2c9-5f02-9a11-845980a1fd5c\"\nStatistics = \"10745b16-79ce-11e8-11f9-7d13ad32a3b2\"\nStatsBase = \"2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91\"\nTest = \"8dfed614-e22c-5e08-85e1-65c5234f0b40\"\nTransformVariables = \"84d833dd-6860-57f9-a1a7-6da5db126cff\"\n\n[sources]\nDynamicHMC = {path = \"..\"}\n\n[compat]\nLogDensityTestSuite = \"0.7\"\n" }, { "path": "test/runtests.jl", "content": "using DynamicHMC, Test, ArgCheck, DocStringExtensions, HypothesisTests, LinearAlgebra,\n Random, Statistics\nimport OhMyThreads\nusing LogExpFunctions: logaddexp, log1mexp\nusing StatsBase: mean_and_cov\nusing Logging: with_logger, NullLogger\nimport ForwardDiff, Random, TransformVariables\nusing DynamicHMC.Diagnostics\nusing DynamicHMC.Diagnostics: ACCEPTANCE_QUANTILES\nusing LogDensityProblems: logdensity_and_gradient, dimension, LogDensityProblems\nusing LogDensityTestSuite\n\n####\n#### static analysis and QA; before everything else as tests extend methods\n####\n\n@testset \"static analysis with JET.jl\" begin\n using JET\n @test isempty(JET.get_reports(report_package(DynamicHMC, target_modules=(DynamicHMC,))))\nend\n\n@testset \"Aqua\" begin\n import Aqua\n Aqua.test_all(DynamicHMC; ambiguities = false)\n # testing separately, cf https://github.com/JuliaTesting/Aqua.jl/issues/77\n Aqua.test_ambiguities(DynamicHMC)\nend\n\n###\n### general test environment\n###\n\nconst RNG = copy(Random.default_rng()) # shorthand\nRandom.seed!(RNG, UInt32[0x23ef614d, 0x8332e05c, 0x3c574111, 0x121aa2f4])\n\n\"Tolerant testing in a CI environment.\"\nconst RELAX = (k = \"CONTINUOUS_INTEGRATION\"; haskey(ENV, k) && ENV[k] == \"true\")\n\ninclude(\"utilities.jl\")\n\n####\n#### unit tests\n####\n\ninclude(\"test_trees.jl\")\ninclude(\"test_hamiltonian.jl\")\ninclude(\"test_NUTS.jl\")\ninclude(\"test_stepsize.jl\")\ninclude(\"test_mcmc.jl\")\ninclude(\"test_diagnostics.jl\")\ninclude(\"test_logging.jl\")\n\n####\n#### sample correctness tests\n####\n\ninclude(\"sample-correctness_tests.jl\")\n" }, { "path": "test/sample-correctness_tests.jl", "content": "include(\"sample-correctness_utilities.jl\")\n\n#####\n##### sample correctness tests\n#####\n##### Sample from well-characterized distributions using LogDensityTestSuite, check\n##### convergence and mixing, and compare.\n\n\"Adaptation with dense matrix\"\nconst MCMC_ARGS2 = (warmup_stages = default_warmup_stages(; M = Symmetric),)\n\n@testset \"NUTS tests with random normal\" begin\n for _ in 1:10\n K = rand(RNG, 3:10)\n μ = randn(RNG, K)\n d = abs.(randn(RNG, K))\n C = rand_C(K)\n ℓ = multivariate_normal(μ, Diagonal(d) * C)\n title = \"multivariate normal μ = $(μ) d = $(d) C = $(C)\"\n NUTS_tests(RNG, ℓ, title, 1000; mcmc_args = MCMC_ARGS2,\n R̂_alert = 1.02, τ_alert = 0.7)\n end\nend\n\n@testset \"ill-conditioned multivariate normal\" begin\n # this test case was isolated using random tests\n μ = [-1.729922440774685, -0.011762500688978205, 0.11423091067230899, 0.05085717388622323, 0.09102774773399233, -0.3769237300508154, -1.1645971596831883, -1.4196407006756644, 0.07406060991401947]\n d = [0.31285715405356296, 1.6321047397137334, 1.9304214045496948, 0.9408515651923572, 0.632832415315841, 0.3994529605030148, 0.9479547802750243, 0.000686699019868418, 0.14074551354895906]\n C = [1.0 -0.625893845478092 -0.8607538232958145 0.4906036948283603 -0.045129301268019346 -0.9798256449980116 -0.09448716779625055 0.1972478332046149 -0.38125524332165456; 0.0 0.7799082601131022 0.22963314745353192 -0.8390321758549951 -0.2940681265758735 0.05788305453491861 -0.30348581879657555 -0.3395815944065493 0.40817023926937634; 0.0 0.0 0.45428127109998945 0.07704183020878513 0.5013749270904165 0.09940288184055725 -0.4898077520422466 -0.04390387380845317 -0.39358273046921877; 0.0 0.0 0.0 0.22225566111771966 -0.5034002085122711 0.1540822287067389 -0.52831870161212 -0.20197326086456527 -0.4230725997740589; 0.0 0.0 0.0 0.0 0.6377293278924043 0.002108173376346147 -0.563819920556515 0.07024142256309863 0.20409522211102057; 0.0 0.0 0.0 0.0 0.0 0.05444765270890811 0.21770654511030652 0.4167989822452558 0.4096707796964533; 0.0 0.0 0.0 0.0 0.0 0.0 0.12102564140379203 0.6237333486866049 -0.1142510107612157; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.4851374500990013 -0.2027266958462243; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.30084429646746724]'\n ℓ = multivariate_normal(μ, Diagonal(d) * C)\n title = \"ill-conditioned multivariate normal (isolated test case 1)\"\n NUTS_tests(RNG, ℓ, title, 1000; mcmc_args = MCMC_ARGS2)\n\n d = [0.44940324099952655, 1.2470316880832284, 1.4254609657195896, 0.47414925026956667, 0.7208717869588667, 0.9012540329863461, 0.259210347514327, 0.48018821609980755, 0.036285320442367444]\n C = [1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.007468818792116497 0.999972107983943 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.9511843069109334 0.06094826193577815 0.30254540758929904 0.0 0.0 0.0 0.0 0.0 0.0; 0.5836451073483746 0.5224198876250752 -0.1567642318026896 0.6015486890596806 0.0 0.0 0.0 0.0 0.0; -0.04549583361258265 0.16604582867077644 -0.6573154635023393 0.5230837360874556 0.5144693366823966 0.0 0.0 0.0 0.0; 0.3090114014598978 0.21784144366429148 0.09455066936309542 0.7472520532986878 0.3661721405808872 0.39452447632098014 0.0 0.0 0.0; 0.27849576428755396 0.008203485989481384 -0.6289527864239539 0.5299626182310367 -0.18989119185086065 0.3458859908657774 0.30039148523055575 0.0 0.0; -0.7595504281026706 -0.6109486667620377 0.08322674440383553 -0.12441158714041263 -0.15879164203513468 -0.0032350588677425886 0.027740844099589795 0.03775094878848311 0.0; 0.8843786481850745 0.4137017432529274 0.19839646818921372 -0.07842556868606812 0.03458430271168502 0.0036393230648423818 0.0006870732712296159 -0.0015642900624311437 0.0011437266452138846]\n ℓ = multivariate_normal(μ, Diagonal(d) * C)\n title = \"ill-conditioned multivariate normal (isolated test case 2)\"\n NUTS_tests(RNG, ℓ, title, 1000; mcmc_args = MCMC_ARGS2)\n\n μ = [0.21062974278940136, -1.218937450424899, 0.06421875640449011, -0.8234583898758592, -2.31397504655407, -0.4751175796619936, -1.2623323961397874, 0.2150945580900463, 1.0797988499707567, 0.6923991470384713]\n d = [1.235510286986013, 0.25725289997297635, 0.39737933906879164, 1.2464348820193416, 0.3082850398698708, 0.9563709407505254, 1.6547932918031834, 1.9782388109071316, 0.38580150239677885, 0.45488559976648274]\n C = [1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.5858606519975413 0.8104118067013929 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; -0.3184163160259112 0.8041538301838452 0.501943888387077 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.3173460682399272 0.6771172525630316 -0.41159671670836784 0.520952821327462 0.0 0.0 0.0 0.0 0.0 0.0; -0.987376065017123 -0.0893955251935478 -0.1251983682331955 0.015871075518314355 0.03421145802664587 0.0 0.0 0.0 0.0 0.0; 0.37469357703269496 -0.8443427667670257 0.32370544135718116 -0.052396077029688945 -0.14292183643709977 0.13686782878290468 0.0 0.0 0.0 0.0; -0.6171193584146126 -0.6578898907477293 -0.39307408945037237 -0.1518878423897761 -0.04583110799414341 0.024372352823947997 0.0779290101096559 0.0 0.0 0.0; 0.5435692867326045 -0.6050903050824995 0.08910494475273394 -0.3209596162864902 0.39975938033524144 0.07516818530300905 -0.06448639900775556 0.24047260310743332 0.0 0.0; -0.06388905564192496 0.9843759627707926 -0.12367139895609519 -0.02886519073736079 0.08699952332803386 -0.020427021493780943 0.0227516163109634 0.010263085877575476 0.04674602752418515 0.0; -0.05914353971342278 0.5051281727293001 -0.0853459337837312 0.7320866937322082 0.42886052044809864 0.011574865047660135 0.10703394808902246 0.045502786672532804 -0.01539436089666275 0.017135804222740844]\n ℓ = multivariate_normal(μ, Diagonal(d) * C)\n title = \"ill-conditioned multivariate normal (isolated test case 2)\"\n NUTS_tests(RNG, ℓ, title, 1000; mcmc_args = MCMC_ARGS2)\nend\n\n@testset \"NUTS tests with specific normal distributions\" begin\n ℓ = multivariate_normal([0.0], 5e8)\n NUTS_tests(RNG, ℓ, \"univariate huge variance\", 1000)\n\n ℓ = multivariate_normal([1.0], 5e8)\n NUTS_tests(RNG, ℓ, \"univariate huge variance, offset\", 1000)\n\n ℓ = multivariate_normal([1.0], 5e-8)\n NUTS_tests(RNG, ℓ, \"univariate tiny variance, offset\", 1000)\n\n ℓ = multivariate_normal([1.0, 2.0, 3.0], Diagonal([1.0, 2.0, 3.0]))\n NUTS_tests(RNG, ℓ, \"mildly scaled diagonal\", 1000)\n\n # these tests are kept because they did produce errors for some code that turned out to\n # be buggy in the early development version; this does not meant that they are\n # particularly powerful or sensitive ones\n ℓ = multivariate_normal([-0.37833073009094703, -0.3973395239297558],\n cholesky([0.08108928067723374 -0.19742780267879112;\n -0.19742780267879112 1.2886298811010262]).L)\n NUTS_tests(RNG, ℓ, \"kept 2 dim\", 1000)\n\n ℓ = multivariate_normal(\n [-1.0960316317778482, -0.2779143641884689, -0.4566289703243874],\n cholesky([2.2367476976202463 1.4710084974801891 2.41285525745893;\n 1.4710084974801891 1.1684361535929932 0.9632367554302268;\n 2.41285525745893 0.9632367554302268 4.5595606374865785]).L)\n NUTS_tests(RNG, ℓ, \"kept 3 dim\", 1000)\n\n ℓ = multivariate_normal(\n [-1.42646, 0.94423, 0.852379, -1.12906, 0.0868619, 0.948781, -0.875067, 1.07243],\n cholesky([14.8357 2.42526 -2.97011 2.08363 -1.67358 4.02846 5.57947 7.28634;\n 2.42526 10.8874 -1.08992 1.99358 1.85011 -2.29754 -0.0540131 1.79718;\n -2.97011 -1.08992 3.05794 0.0321187 1.8052 -1.5309 1.78163 -0.0821483;\n 2.08363 1.99358 0.0321187 2.38112 -0.252784 0.666474 1.73862 2.55874;\n -1.67358 1.85011 1.8052 -0.252784 12.3109 -2.3913 -2.99741 -1.95031;\n 4.02846 -2.29754 -1.5309 0.666474 -2.3913 4.89957 3.6118 5.22626;\n 5.57947 -0.0540131 1.78163 1.73862 -2.99741 3.6118 10.215 9.60671;\n 7.28634 1.79718 -0.0821483 2.55874 -1.95031 5.22626 9.60671 11.5554]).L)\n NUTS_tests(RNG, ℓ, \"kept 8 dim\", 1000)\nend\n\n@testset \"NUTS tests with mixtures\" begin\n ℓ1 = multivariate_normal(zeros(3), 1.0)\n D2 = I * 0.4\n C2 = [1.0 -0.48058358598852935 0.39971148270854306;\n 0.0 0.876948924897229 -0.5361348433365906;\n 0.0 0.0 0.7434985947205197]\n ℓ2 = multivariate_normal(ones(3), D2 * C2)\n ℓ = mix(0.2, ℓ1, ℓ2)\n NUTS_tests(RNG, ℓ, \"mixture of two normals\", 1000; τ_alert = 0.15, p_alert = 0.005)\nend\n\n@testset \"NUTS tests with heavier tails and skewness\" begin\n K = 5\n 𝒩 = StandardMultivariateNormal(K)\n\n # somewhat nasty, relaxed requirements\n ℓ = elongate(1.1)(𝒩)\n NUTS_tests(RNG, ℓ, \"elongate(1.1, 𝑁)\",\n 10000; p_alert = 0.05, EBFMI_alert = 0.2, R̂_fail = 1.05, τ_fail = 0.3)\n\n # this has very nasty tails so we relax requirements a bit\n ℓ = (elongate(1.1) ∘ shift(ones(K)))(𝒩)\n NUTS_tests(RNG, ℓ, \"skew elongate(1.1, 𝑁)\",\n 10000; τ_alert = 0.1, EBFMI_alert = 0.2, R̂_fail = 1.05, p_fail = 0.001)\n\n # funnel, mixed with a normal\n ℓ = mix(0.8, funnel()(𝒩), 𝒩)\n NUTS_tests(RNG, ℓ, \"funnel\", 10000;\n EBFMI_alert = 0.2, τ_alert = 0.1, p_fail = 5e-3, R̂_fail = 1.05)\nend\n" }, { "path": "test/sample-correctness_utilities.jl", "content": "#####\n##### utilities for testing sample correctness\n#####\n\nusing MCMCDiagnosticTools: ess_rhat\n\n\"\"\"\n$(SIGNATURES)\n\nRun `K` chains of MCMC on `ℓ`, each for `N` samples, return a posterior matrices stacked\n(indexed by `[draw, parameter, chain]`) and concatenated (indexed by `[draw, parameter]`),\nand EBFMI statistics as fields of a `NamedTuple`.\n\nKeyword arguments are passed to `mcmc_with_warmup`.\n\"\"\"\nfunction run_chains(rng, ℓ, N, K; mcmc_args...)\n results = OhMyThreads.tcollect(mcmc_with_warmup(rng, ℓ, N; reporter = NoProgressReport(), mcmc_args...)\n for _ in 1:K)\n (stacked_posterior_matrices = stack_posterior_matrices(results),\n concat_posterior_matrices = pool_posterior_matrices(results),\n EBFMIs = map(r -> EBFMI(r.tree_statistics), results))\nend\n\n###\n### Multivariate normal ℓ for testing\n###\n\n\"Random Cholesky factor for correlation matrix.\"\nfunction rand_C(K)\n t = TransformVariables.CorrCholeskyFactor(K)\n TransformVariables.transform(t, randn(RNG, TransformVariables.dimension(t)) ./ 4)'\nend\n\n\"\"\"\n$(SIGNATURES)\n\n`R̂` (within/between variance) and `τ` (effective sample size coefficient) statistics for\nposterior matrices, eg the output of `run_chains`.\n\"\"\"\nfunction mcmc_statistics(stacked_posterior_matrices)\n (; ess, rhat) = ess_rhat(stacked_posterior_matrices)\n (R̂ = rhat, τ = ess ./ size(stacked_posterior_matrices, 1))\nend\n\n\"\"\"\n$(SIGNATURES)\n\nJitter in `(-ϵ, +ϵ)`. Useful for tiebreaking for the Kolmogorov-Smirnov tests.\n\"\"\"\njitter(rng, len, ϵ = 64*eps()) = (2 * ϵ) .* (rand(rng, len) .- 0.5)\n\n\"\"\"\n$(SIGNATURES)\n\nRun MCMC on `ℓ`, obtaining `N` samples from `K` independently adapted chains.\n\n`R̂`, `τ`, Kolmogorov-Smirnov and Anderson-Darling `p`, and EBFMIs are obtained and compared\nto thresholds that either *alert* or *fail*. The latter should be lax because of false\npositives, the tests can be rather hair-trigger.\n\nOutput is sent to `io`. Specifically, `title` is printed for the first alert.\n\n`mcmc_args` are passed down to `mcmc_with_warmup`.\n\"\"\"\nfunction NUTS_tests(rng, ℓ, title, N; K = 5, io = stdout, mcmc_args = NamedTuple(),\n R̂_alert = 1.01, R̂_fail = 2 * (R̂_alert - 1) + 1,\n τ_alert = 1.0, τ_fail = τ_alert * 0.5,\n p_alert = 0.1, p_fail = p_alert * 0.1,\n EBFMI_alert = 0.5, EBFMI_fail = EBFMI_alert / 2)\n @argcheck 1 < R̂_alert ≤ R̂_fail\n @argcheck 0 < τ_fail ≤ τ_alert\n @argcheck 0 < p_fail ≤ p_alert\n @argcheck 0 < EBFMI_fail < EBFMI_alert\n\n d = dimension(ℓ)\n _round(x) = round(x; sigdigits = 3) # for printing\n\n title_printed = false\n function _print_diagnostics(label, is_min, value, alert_threshold, error_threshold)\n if !title_printed\n printstyled(io, \"INFO while testing: $(title), dimension $(d)\\n\";\n color = :blue, bold = true)\n title_printed = true\n end\n if is_min\n vm = minimum(value)\n if vm ≥ alert_threshold\n mark, rel, vt, col = '✓', '≥', alert_threshold, :green\n elseif vm ≥ error_threshold\n mark, rel, vt, col = '!', '≱', alert_threshold, :yellow\n else\n mark, rel, vt, col = '✘', '≱', error_threshold, :red\n end\n else\n vm = maximum(value)\n if vm ≤ alert_threshold\n mark, rel, vt, col = '✓', '≤', alert_threshold, :green\n elseif vm ≤ error_threshold\n mark, rel, vt, col = '!', '≰', alert_threshold, :yellow\n else\n mark, rel, vt, col = '✘', '≰', error_threshold, :red\n end\n end\n printstyled(io, \"$(mark) $(label) = $(_round.(vm)) $(rel) $(_round(vt))\\n\";\n color = col)\n end\n (; stacked_posterior_matrices, concat_posterior_matrices, EBFMIs) =\n run_chains(RNG, ℓ, N, K; mcmc_args...)\n\n # mixing and autocorrelation diagnostics\n (; R̂, τ) = mcmc_statistics(stacked_posterior_matrices)\n _print_diagnostics(\"R̂\", false, R̂, R̂_alert, R̂_fail)\n @test all(maximum(R̂) ≤ R̂_fail)\n _print_diagnostics(\"τ\", true, τ, τ_alert, τ_fail)\n @test all(minimum(τ) ≥ τ_fail)\n _print_diagnostics(\"EBFMI\", true, EBFMIs, EBFMI_alert, EBFMI_fail)\n @test all(minimum(EBFMIs) ≥ EBFMI_fail)\n\n # distribution comparison tests\n Z = concat_posterior_matrices\n Z′ = samples(ℓ, 1000)\n pd_alert = p_alert / d # a simple Bonferroni correction\n pd_fail = p_fail / d\n ps = map((a, b) -> pvalue(KSampleADTest(a, b)), eachrow(Z), eachrow(Z′))\n _print_diagnostics(\"p\", true, ps, p_alert, p_fail)\n @test all(minimum(ps) ≥ p_fail)\nend\n" }, { "path": "test/test_NUTS.jl", "content": "using DynamicHMC: TrajectoryNUTS, rand_bool_logprob, GeneralizedTurnStatistic,\n AcceptanceStatistic, leaf_acceptance_statistic, acceptance_rate, TreeStatisticsNUTS,\n NUTS, sample_tree, combine_turn_statistics, combine_visited_statistics, evaluate_ℓ,\n Hamiltonian\n\n###\n### random booleans\n###\n\n@testset \"random booleans\" begin\n for prob in (1:9) ./ 10\n logprob = log(prob)\n @test abs(mean(rand_bool_logprob(RNG, logprob) for _ in 1:10000) - prob) ≤ 0.02\n end\n\n # these operations don't call the RNG, this is checked\n RNG′ = copy(RNG)\n @test all(rand_bool_logprob(RNG, 0) for _ in 1:10000)\n @test all(rand_bool_logprob(RNG, 10) for _ in 1:10000)\n @test rand(RNG′) == rand(RNG)\nend\n\n###\n### test turn statistics\n###\n\n@testset \"low-level turn statistics\" begin\n trajectory = TrajectoryNUTS(nothing, 0, 1, -1000, Val{:generalized})\n p = ones(3) # unit vector\n c = 0.1 # a constant, just for consistency checking of combination\n # turn statistics constructed so that τ₁ + τ₂ won't be turning, τ₁ + τ₃ will be\n τ₁ = GeneralizedTurnStatistic(p, p .- c, p, p .- c, p)\n τ₂ = GeneralizedTurnStatistic(3 .* p, 3 .* p .+ c, 3 .* p, 3 .* p .+ c, 3 .* p)\n τ₃ = GeneralizedTurnStatistic(2 .* p, 2 .* p .+ c, 2 .* p, 2 .* p .+ c, -2 .* p)\n τ = combine_turn_statistics(trajectory, τ₁, τ₂)\n # test mechanics of combination\n @test τ.ρ == τ₁.ρ .+ τ₂.ρ\n # test non-turning\n @test !is_turning(trajectory, τ)\n # test turning\n @test is_turning(trajectory, combine_turn_statistics(trajectory, τ₁, τ₃))\nend\n\n@testset \"low-level visited statistics\" begin\n trajectory = TrajectoryNUTS(nothing, 0, 1, -1000, Val{:generalized})\n vs(p, is_initial = false) = leaf_acceptance_statistic(log(p), is_initial)\n x = vs(0.3)\n @test acceptance_rate(x) ≈ 0.3\n y = vs(0.6)\n @test acceptance_rate(y) ≈ 0.6\n x0 = vs(10, true) # initial node, does not count\n z = reduce((x, y) -> combine_visited_statistics(trajectory, x, y),\n [x, x, y, x0])\n @test acceptance_rate(z) ≈ 0.4\nend\n\n# define a distribution which is divergent everywhere except at 0\nstruct AlwaysDivergentTest\n K::Int\nend\n\nfunction LogDensityProblems.capabilities(::Type{AlwaysDivergentTest})\n LogDensityProblems.LogDensityOrder{1}()\nend\nLogDensityProblems.dimension(d::AlwaysDivergentTest) = d.K\nfunction LogDensityProblems.logdensity_and_gradient(d::AlwaysDivergentTest, x)\n ∇ = ones(length(x))\n if all(iszero.(x))\n 0.0, ∇\n else\n -Inf, ∇\n end\nend\n\n@testset \"unconditional divergence\" begin\n # test NUTS sampler where all movements are divergent\n K = 3\n ℓ = AlwaysDivergentTest(K)\n Q, tree_statistics = sample_tree(RNG, NUTS(), Hamiltonian(GaussianKineticEnergy(K), ℓ),\n evaluate_ℓ(ℓ, zeros(K)), 1.0)\n @test is_divergent(tree_statistics.termination)\n @test iszero(tree_statistics.acceptance_rate)\n @test iszero(tree_statistics.depth)\n @test tree_statistics.steps == 1\nend\n\n@testset \"normal NUTS HMC transition mean and cov\" begin\n # A test for sample_tree with a fixed ϵ and κ, which is perfectly adapted and should\n # provide excellent mixing\n for _ in 1:10\n K = rand(RNG, 2:8)\n N = 10000\n μ = randn(RNG, K)\n Σ = rand_Σ(K)\n L = cholesky(Σ).L\n ℓ = multivariate_normal(μ, L)\n Q = evaluate_ℓ(ℓ, randn(RNG, K))\n H = Hamiltonian(GaussianKineticEnergy(Σ), ℓ)\n qs = Array{Float64}(undef, N, K)\n ϵ = 0.5\n algorithm = NUTS()\n for i in 1:N\n Q = first(sample_tree(RNG, algorithm, H, Q, ϵ))\n qs[i, :] = Q.q\n end\n m, C = mean_and_cov(qs, 1)\n tol = maximum(diag(C)) / 50\n @test vec(m) ≈ μ atol = tol rtol = tol norm = x -> norm(x,1)\n @test cov(qs, dims = 1) ≈ L*L' atol = 0.1 rtol = 0.1\n end\nend\n" }, { "path": "test/test_diagnostics.jl", "content": "#####\n##### test diagnostics\n#####\n\n@testset \"summarize tree statistics\" begin\n N = 1000\n directions = Directions(UInt32(0))\n function rand_invalidtree()\n if rand(RNG) < 0.1\n REACHED_MAX_DEPTH\n else\n left = rand(RNG, -5:5)\n right = left + rand(RNG, 0:5)\n InvalidTree(left, right)\n end\n end\n tree_statistics = [TreeStatisticsNUTS(randn(RNG), rand(RNG, 0:5), rand_invalidtree(),\n rand(RNG), rand(RNG, 1:30), directions) for _ in 1:N]\n stats = summarize_tree_statistics(tree_statistics)\n # acceptance rates\n @test stats.N == N\n @test stats.a_mean ≈ mean(x -> x.acceptance_rate, tree_statistics)\n @test stats.a_quantiles ==\n quantile((x -> x.acceptance_rate).(tree_statistics), ACCEPTANCE_QUANTILES)\n # termination counts\n @test stats.termination_counts.divergence ==\n count(x -> is_divergent(x.termination), tree_statistics)\n @test stats.termination_counts.max_depth ==\n count(x -> x.termination == REACHED_MAX_DEPTH, tree_statistics)\n @test stats.termination_counts.turning ==\n (N - stats.termination_counts.max_depth - stats.termination_counts.divergence)\n # depth counts\n for (i, c) in enumerate(stats.depth_counts)\n @test count(x -> x.depth == i - 1, tree_statistics) == c\n end\n @test sum(stats.depth_counts) == N\n # misc\n @test 1.8 ≤ EBFMI(tree_statistics) ≤ 2.2 # nonsensical value, just checking calculation\n @test repr(stats) isa AbstractString # just test that it prints w/o error\nend\n\n@testset \"log acceptance ratios\" begin\n ℓ = multivariate_normal(ones(5))\n log2ϵs = -5:5\n N = 13\n logA = explore_log_acceptance_ratios(ℓ, zeros(5), log2ϵs; N = N)\n @test all(isfinite.(logA))\n @test size(logA) == (length(log2ϵs), N)\nend\n\n@testset \"leapfrog trajectory\" begin\n # problem setup\n K = 2\n ℓ = multivariate_normal(ones(K))\n κ = GaussianKineticEnergy(K)\n q = zeros(K)\n Q = evaluate_ℓ(ℓ, q)\n p = ones(K) .* 0.98\n H = Hamiltonian(κ, ℓ)\n ϵ = 0.1\n ixs = 1:15\n ix0 = 5\n\n # calculate trajectory manually\n zs1 = let z = PhasePoint(Q, p)\n [(z1 = z; z = leapfrog(H, z, ϵ); z1) for _ in ixs]\n end\n πs1 = logdensity.(Ref(H), zs1)\n Δs1 = πs1 .- πs1[ix0]\n\n # calculate using function\n traj = leapfrog_trajectory(ℓ, zs1[ix0].Q.q, ϵ, ixs .- ix0; κ = κ, p = zs1[ix0].p)\n @test all(isapprox.(map(t -> t.Δ, traj), Δs1; atol = 1e-5))\n @test all(map((t, y) -> t.z.Q.q ≈ y.Q.q, traj, zs1))\n @test all(map((t, y) -> t.z.p ≈ y.p, traj, zs1))\nend\n" }, { "path": "test/test_hamiltonian.jl", "content": "using DynamicHMC: GaussianKineticEnergy, kinetic_energy, ∇kinetic_energy, rand_p,\n Hamiltonian, EvaluatedLogDensity, evaluate_ℓ, PhasePoint, logdensity, leapfrog,\n calculate_p♯, logdensity, find_initial_stepsize, DynamicHMCError, local_log_acceptance_ratio\n\n####\n#### utility functions\n####\n\n\"Test kinetic energy gradient by automatic differentiation.\"\nfunction test_KE_gradient(κ::DynamicHMC.EuclideanKineticEnergy, p)\n ∇ = ∇kinetic_energy(κ, p)\n ∇_AD = ForwardDiff.gradient(p -> kinetic_energy(κ, p), p)\n @test ∇ ≈ ∇_AD\nend\n\n####\n#### testsets\n####\n\n@testset \"Gaussian KE full\" begin\n for _ in 1:100\n K = rand(RNG, 2:10)\n Σ = rand_Σ(Symmetric, K)\n κ = GaussianKineticEnergy(inv(Σ))\n (; M⁻¹, W) = κ\n @test W isa LowerTriangular\n @test M⁻¹ * W * W' ≈ Diagonal(ones(K))\n m, C = simulated_meancov(()->rand_p(RNG, κ), 10000)\n @test Matrix(Σ) ≈ C rtol = 0.1\n test_KE_gradient(κ, randn(RNG, K))\n end\nend\n\n@testset \"Gaussian KE diagonal\" begin\n for _ in 1:100\n K = rand(RNG, 2:10)\n Σ = rand_Σ(Diagonal, K)\n κ = GaussianKineticEnergy(inv(Σ))\n (; M⁻¹, W) = κ\n @test W isa Diagonal\n # FIXME workaround for https://github.com/JuliaLang/julia/issues/28869\n @test M⁻¹ * (W * W') ≈ Diagonal(ones(K))\n m, C = simulated_meancov(()->rand_p(RNG, κ), 10000)\n @test Matrix(Σ) ≈ C rtol = 0.1\n test_KE_gradient(κ, randn(RNG, K))\n end\nend\n\n@testset \"phasepoint internal consistency\" begin\n # when this breaks, interface was modified, rewrite tests\n @test fieldnames(PhasePoint) == (:Q, :p)\n \"Test the consistency of cached values.\"\n function test_consistency(H, z)\n (; q, ℓq, ∇ℓq) = z.Q\n (; ℓ) = H\n ℓ2, ∇ℓ2 = logdensity_and_gradient(ℓ, q)\n @test ℓ2 == ℓq\n @test ∇ℓ2 == ∇ℓq\n end\n (; H, z, Σ ) = rand_Hz(rand(RNG, 3:10))\n test_consistency(H, z)\n ϵ = find_stable_ϵ(H.κ, Σ)\n for _ in 1:10\n z = leapfrog(H, z, ϵ)\n test_consistency(H, z)\n end\nend\n\n@testset \"leapfrog calculation\" begin\n # Simple leapfrog implementation. `q`: position, `p`: momentum, `ℓ`: neg_energy, `ϵ`:\n # stepsize. `m` is the diagonal of the kinetic energy ``K(p)=p'M⁻¹p``, defaults to `I`.\n function leapfrog_Gaussian(q, p, ℓ, ϵ, m = ones(length(p)))\n u = .√(1 ./ m)\n pₕ = p .+ ϵ/2 .* last(logdensity_and_gradient(ℓ, q))\n q′ = q .+ ϵ * u .* (u .* pₕ) # mimic numerical calculation leapfrog performs\n p′ = pₕ .+ ϵ/2 .* last(logdensity_and_gradient(ℓ, q′))\n q′, p′\n end\n\n n = 3\n M = rand_Σ(Diagonal, n)\n m = diag(M)\n κ = GaussianKineticEnergy(inv(M))\n q = randn(RNG, n)\n p = randn(RNG, n)\n Σ = rand_Σ(n)\n ℓ = multivariate_normal(randn(RNG, n), cholesky(Σ).L)\n H = Hamiltonian(κ, ℓ)\n ϵ = find_stable_ϵ(H.κ, Σ)\n z = PhasePoint(evaluate_ℓ(ℓ, q), p)\n\n @testset \"arguments not modified\" begin\n q₂, p₂ = copy(q), copy(p)\n q′, p′ = leapfrog_Gaussian(q, p, ℓ, ϵ, m)\n z′ = leapfrog(H, z, ϵ)\n @test p == p₂ # arguments not modified\n @test q == q₂\n @test z′.Q.q ≈ q′\n @test z′.p ≈ p′\n end\n\n @testset \"leapfrog steps\" begin\n for i in 1:100\n q, p = leapfrog_Gaussian(q, p, ℓ, ϵ, m)\n z = leapfrog(H, z, ϵ)\n @test z.Q.q ≈ q\n @test z.p ≈ p\n end\n end\n\n @testset \"invalid values\" begin\n n = 3\n ℓ = multivariate_normal(randn(RNG, n), I(n))\n @test_throws DynamicHMCError evaluate_ℓ(ℓ, fill(NaN, n))\n end\nend\n\n@testset \"leapfrog Hamiltonian invariance\" begin\n \"Test that the Hamiltonian is invariant using the leapfrog integrator.\"\n function test_hamiltonian_invariance(H, z, L, ϵ; atol)\n π₀ = logdensity(H, z)\n warned = false\n for i in 1:L\n z = leapfrog(H, z, ϵ)\n Δ = logdensity(H, z) - π₀\n if abs(Δ) ≥ atol && !warned\n @warn \"Hamiltonian invariance violated\" step = i L Δ\n show(H)\n show(z)\n warned = true\n end\n @test Δ ≈ 0 atol = atol\n end\n end\n\n for _ in 1:100\n (; H, z) = rand_Hz(rand(RNG, 2:5))\n ϵ = find_initial_stepsize(InitialStepsizeSearch(), local_log_acceptance_ratio(H, z))\n test_hamiltonian_invariance(H, z, 10, ϵ/100; atol = 0.5)\n end\nend\n\n@testset \"leapfrog back and forth\" begin\n for _ in 1:1000\n (; H, z) = rand_Hz(5)\n z1 = z\n N = 5\n ϵ = 0.1\n z1 = leapfrog(H, z1, ϵ)\n z1 = leapfrog(H, z1, -ϵ)\n @test z.p ≈ z1.p norm = x -> norm(x, Inf) atol = 1e-5\n @test z.Q.q ≈ z1.Q.q norm = x -> norm(x, Inf) atol = 1e-6\n end\n\n for _ in 1:100\n (; H, z, Σ) = rand_Hz(2)\n z1 = z\n N = 3\n\n # use something near the stable stepsize to avoid numerical issue, but perturb it a\n # bit for testing\n ϵ = find_stable_ϵ(H.κ, Σ) * (0.5 + rand(RNG))\n\n # forward\n for _ in 1:N\n z1 = leapfrog(H, z1, ϵ)\n end\n\n # backward\n for _ in 1:N\n z1 = leapfrog(H, z1, -ϵ)\n end\n\n @test z.p ≈ z1.p norm = x -> norm(x, Inf) rtol = 0.001\n @test z.Q.q ≈ z1.Q.q norm = x -> norm(x, Inf) rtol = 0.001\n end\nend\n\n@testset \"PhasePoint building blocks and infinite values\" begin\n # wrong gradient length\n @test_throws ArgumentError EvaluatedLogDensity([1.0, 2.0], 1.0, [1.0])\n\n # wrong p length\n Q = EvaluatedLogDensity([1.0, 2.0], 1.0, [1.0, 2.0])\n @test_throws ArgumentError PhasePoint(Q, [1.0])\n @test PhasePoint(Q, [1.0, 2.0]) isa PhasePoint\n\n # fallback constructors\n Q1 = EvaluatedLogDensity([1.0, 2.0], -2.0, [3.0, 3.0]) # standard\n Q2 = EvaluatedLogDensity([1, 2], -2.0, [3.0, 3.0]) # promote\n Q3 = EvaluatedLogDensity((i for i in 1:2), -2.0, [3.0, 3.0]) # generator\n @test Q1.q == Q2.q == Q3.q\n @test Q1.ℓq == Q2.ℓq == Q3.ℓq\n @test Q1.∇ℓq == Q2.∇ℓq == Q3.∇ℓq\n\n # infinity fallbacks\n h = Hamiltonian(GaussianKineticEnergy(1), multivariate_normal(zeros(1)))\n @test logdensity(h, PhasePoint(EvaluatedLogDensity([1.0], -Inf, [1.0]), [1.0])) == -Inf\n @test logdensity(h, PhasePoint(EvaluatedLogDensity([1.0], NaN, [1.0]), [1.0])) == -Inf\n @test logdensity(h, PhasePoint(EvaluatedLogDensity([1.0], 9.0, [1.0]), [NaN])) == -Inf\n\nend\n\n@testset \"Hamiltonian and KE printing\" begin\n κ = GaussianKineticEnergy(Diagonal([1.0, 4.0]))\n @test repr(κ) == \"Gaussian kinetic energy (Diagonal), √diag(M⁻¹): [1.0, 2.0]\"\n H = Hamiltonian(κ, multivariate_normal(zeros(2)))\n @test repr(H) ==\n \"Hamiltonian with Gaussian kinetic energy (Diagonal), √diag(M⁻¹): [1.0, 2.0]\"\n @test_throws ArgumentError Hamiltonian(κ, multivariate_normal(zeros(1)))\nend\n\n####\n#### test Hamiltonian/leapfrog using HMC\n####\n\n\"\"\"\n$(SIGNATURES)\n\nSimple Hamiltonian Monte Carlo transition, for testing.\n\"\"\"\nfunction HMC_transition(H, z::PhasePoint, ϵ, L)\n π₀ = logdensity(H, z)\n z′ = z\n for _ in 1:L\n z′ = leapfrog(H, z′, ϵ)\n end\n Δ = logdensity(H, z′) - π₀\n accept = Δ > 0 || (rand(RNG) < exp(Δ))\n accept ? z′ : z\nend\n\n\"\"\"\n$(SIGNATURES)\n\nSimple Hamiltonian Monte Carlo sample, for testing.\n\"\"\"\nfunction HMC_sample(H, q, N, ϵ; L = 10)\n qs = similar(q, N, length(q))\n for i in 1:N\n z = PhasePoint(evaluate_ℓ(H.ℓ, q), rand_p(RNG, H.κ))\n q = HMC_transition( H, z, ϵ, L).Q.q\n qs[i, :] = q\n end\n qs\nend\n\n@testset \"unit normal simple HMC\" begin\n # Tests the leapfrog and Hamiltonian code with HMC.\n K = 2\n ℓ = multivariate_normal(zeros(K), Diagonal(ones(K)))\n q = randn(RNG, K)\n H = Hamiltonian(GaussianKineticEnergy(Diagonal(ones(K))), ℓ)\n qs = HMC_sample(H, q, 10000, find_stable_ϵ(H.κ, Diagonal(ones(K))) / 5)\n m, C = mean_and_cov(qs, 1)\n @test vec(m) ≈ zeros(K) atol = 0.1\n @test C ≈ Matrix(Diagonal(ones(K))) atol = 0.1\nend\n" }, { "path": "test/test_logging.jl", "content": "# NOTE currently we just check that logging does not error, more explicit testing might make\n# sense\n\nℓ = multivariate_normal(ones(1))\nκ = GaussianKineticEnergy(1)\nQ = evaluate_ℓ(ℓ, [1.0])\n\nreporters_1 = [\n NoProgressReport(),\n ProgressMeterReport(),\n]\n\nreporters_2 = [\n LogProgressReport(),\n]\n\nwith_logger(NullLogger()) do # suppress logging in CI\n for reporter in deepcopy(vcat(reporters_1, reporters_2))\n results = mcmc_with_warmup(RNG, ℓ, 10_000; reporter = reporter)\n end\n\n for reporter in deepcopy(vcat(reporters_1, reporters_2))\n DynamicHMC.report(reporter, \"\")\n\n mcmc_reporter_1 = DynamicHMC.make_mcmc_reporter(reporter, 1_000; currently_warmup = true)\n DynamicHMC.report(mcmc_reporter_1, \"\")\n DynamicHMC.report(mcmc_reporter_1, 1)\n\n mcmc_reporter_2 = DynamicHMC.make_mcmc_reporter(reporter, 1_000; currently_warmup = false)\n DynamicHMC.report(mcmc_reporter_2, \"\")\n DynamicHMC.report(mcmc_reporter_2, 1)\n end\n\n for reporter in deepcopy(reporters_1)\n DynamicHMC.report(reporter, \"\")\n DynamicHMC.report(reporter, 1)\n\n mcmc_reporter_1 = DynamicHMC.make_mcmc_reporter(reporter, 1_000; currently_warmup = true)\n DynamicHMC.report(mcmc_reporter_1, \"\")\n DynamicHMC.report(mcmc_reporter_1, 1)\n\n mcmc_reporter_2 = DynamicHMC.make_mcmc_reporter(reporter, 1_000; currently_warmup = false)\n DynamicHMC.report(mcmc_reporter_2, \"\")\n DynamicHMC.report(mcmc_reporter_2, 1)\n end\nend\n" }, { "path": "test/test_mcmc.jl", "content": "using DynamicHMC: mcmc_steps, mcmc_next_step, mcmc_keep_warmup, WarmupState\n\n#####\n##### Test building blocks of MCMC\n#####\n\n@testset \"printing\" begin\n ℓ = multivariate_normal(ones(1))\n κ = GaussianKineticEnergy(1)\n Q = evaluate_ℓ(ℓ, [1.0])\n @test repr(WarmupState(Q, κ, 1.0)) isa String\n @test repr(WarmupState(Q, κ, nothing)) isa String\nend\n\n@testset \"mcmc\" begin\n ℓ = multivariate_normal(ones(5))\n\n @testset \"default warmup\" begin\n results = mcmc_with_warmup(RNG, ℓ, 10000; reporter = NoProgressReport())\n Z = results.posterior_matrix\n @test map(z -> LogDensityProblems.logdensity(ℓ, z), eachcol(Z)) ≈ results.logdensities\n @test norm(mean(Z; dims = 2) .- ones(5), Inf) < 0.04\n @test norm(std(Z; dims = 2) .- ones(5), Inf) < 0.04\n @test mean(x -> x.acceptance_rate, results.tree_statistics) ≥ 0.8\n @test 0.5 ≤ results.ϵ ≤ 2\n end\n\n @testset \"fixed stepsize\" begin\n results = mcmc_with_warmup(RNG, ℓ, 10000;\n initialization = (ϵ = 1.0, ),\n reporter = NoProgressReport(),\n warmup_stages = fixed_stepsize_warmup_stages())\n Z = results.posterior_matrix\n @test norm(mean(Z; dims = 2) .- ones(5), Inf) < 0.04\n @test norm(std(Z; dims = 2) .- ones(5), Inf) < 0.04\n @test mean(x -> x.acceptance_rate, results.tree_statistics) ≥ 0.7\n end\n\n @testset \"explicitly provided initial stepsize\" begin\n results = mcmc_with_warmup(RNG, ℓ, 10000;\n initialization = (ϵ = 1.0, ),\n reporter = NoProgressReport(),\n warmup_stages = default_warmup_stages(; stepsize_search = nothing))\n Z = results.posterior_matrix\n @test norm(mean(Z; dims = 2) .- ones(5), Inf) < 0.03\n @test norm(std(Z; dims = 2) .- ones(5), Inf) < 0.04\n @test mean(x -> x.acceptance_rate, results.tree_statistics) ≥ 0.7\n end\n\n @testset \"stepwise\" begin\n results = mcmc_keep_warmup(RNG, ℓ, 0; reporter = NoProgressReport())\n steps = mcmc_steps(results.sampling_logdensity, results.final_warmup_state)\n qs = let Q = results.final_warmup_state.Q\n [(Q = first(mcmc_next_step(steps, Q)); Q.q) for _ in 1:1000]\n end\n @test norm(mean(reduce(hcat, qs); dims = 2) .- ones(5), Inf) ≤ 0.1\n end\nend\n\n@testset \"robust U-turn tests\" begin\n # Cf https://github.com/tpapp/DynamicHMC.jl/issues/115\n function count_max_depth(rng, ℓ, max_depth; N = 1000)\n results = mcmc_with_warmup(rng, ℓ, N;\n algorithm = DynamicHMC.NUTS(max_depth = max_depth),\n reporter = NoProgressReport())\n sum(getfield.(results.tree_statistics, :depth) .≥ max_depth)\n end\n ℓ = multivariate_normal(zeros(200))\n max_depth = 12\n M = sum([count_max_depth(RNG, ℓ, max_depth) for _ in 1:20])\n @test M == 0\nend\n\n@testset \"posterior accessors sanity checks\" begin\n D, N, K = 5, 100, 7\n ℓ = multivariate_normal(ones(5))\n results = fill(mcmc_with_warmup(RNG, ℓ, N; reporter = NoProgressReport()), K)\n @test size(stack_posterior_matrices(results)) == (N, K, D)\n @test size(pool_posterior_matrices(results)) == (D, N * K)\nend\n\n# @testset \"tuner framework\" begin\n# s = StepsizeTuner(10)\n# @test length(s) == 10\n# @test repr(s) == \"Stepsize tuner, 10 samples\"\n# c = StepsizeCovTuner(19, 7.0)\n# @test length(c) == 19\n# @test repr(c) ==\n# \"Stepsize and covariance tuner, 19 samples, regularization 7.0\"\n# b = bracketed_doubling_tuner() # testing the defaults\n# @test b isa TunerSequence\n# @test b.tuners == (StepsizeTuner(75), # init\n# StepsizeCovTuner(25, 5.0), # doubling each step\n# StepsizeCovTuner(50, 5.0),\n# StepsizeCovTuner(100, 5.0),\n# StepsizeCovTuner(200, 5.0),\n# StepsizeCovTuner(400, 5.0),\n# StepsizeTuner(50)) # terminate\n# @test repr(b) ==\n# \"\"\"\n# Sequence of 7 tuners, 900 total samples\n# Stepsize tuner, 75 samples\n# Stepsize and covariance tuner, 25 samples, regularization 5.0\n# Stepsize and covariance tuner, 50 samples, regularization 5.0\n# Stepsize and covariance tuner, 100 samples, regularization 5.0\n# Stepsize and covariance tuner, 200 samples, regularization 5.0\n# Stepsize and covariance tuner, 400 samples, regularization 5.0\n# Stepsize tuner, 50 samples\"\"\"\n# end\n" }, { "path": "test/test_stepsize.jl", "content": "#####\n##### Stepsize and adaptation tests\n#####\n\nusing DynamicHMC: find_initial_stepsize, InitialStepsizeSearch, DualAveraging,\n initial_adaptation_state, adapt_stepsize, current_ϵ, final_ϵ, FixedStepsize,\n local_log_acceptance_ratio\n\n@testset \"stepsize general rootfinding\" begin\n A = ϵ -> -ϵ*3.0\n params = InitialStepsizeSearch()\n # parameters consistency\n @test_throws ArgumentError InitialStepsizeSearch(; log_threshold = NaN) # not finite\n @test_throws ArgumentError InitialStepsizeSearch(; log_threshold = 1.0) # too large\n @test_throws ArgumentError InitialStepsizeSearch(; initial_ϵ = -0.5) # not > 0\n @test_throws ArgumentError InitialStepsizeSearch(; maxiter_crossing = 2) # too small\n # crossing\n ϵ = find_initial_stepsize(params, A)\n @test A(ϵ) > params.log_threshold > A(params.initial_ϵ)\n let params = InitialStepsizeSearch(; initial_ϵ = 0.01)\n ϵ = find_initial_stepsize(params, A)\n @test A(ϵ) < params.log_threshold < A(params.initial_ϵ)\n end\n @test_throws DynamicHMCError find_initial_stepsize(params, ϵ -> 1) # constant\nend\n\n\"\"\"\n$(SIGNATURES)\n\nA parametric random acceptance rate that depends on the stepsize. For unit\ntesting acceptance rate tuning.\n\"\"\"\ndummy_acceptance_rate(ϵ, σ = 0.05) = min(1/ϵ * exp(randn(RNG)*σ - σ^2/2), 1)\n\nmean_dummy_acceptance_rate(ϵ, σ = 0.05) = mean(dummy_acceptance_rate(ϵ, σ) for _ in 1:10000)\n\n@testset \"dual averaging far\" begin\n ϵ₀ = 100.0 # way off\n δ = 0.65\n dual_averaging = DualAveraging(; δ = δ)\n A = initial_adaptation_state(dual_averaging, ϵ₀)\n @test A.logϵ̄ == 0 # ϵ̄₀ = 1 in Gelman and Hoffman (2014)\n @test A.m == 1\n @test A.H̄ == 0\n for _ in 1:500\n A = adapt_stepsize(dual_averaging, A, dummy_acceptance_rate(current_ϵ(A)))\n end\n @test mean_dummy_acceptance_rate(final_ϵ(A)) ≈ δ atol = 0.02\nend\n\n@testset \"dual averaging close\" begin\n ϵ₀ = 2.0\n δ = 0.65\n dual_averaging = DualAveraging(; δ = δ)\n A = initial_adaptation_state(dual_averaging, ϵ₀)\n for _ in 1:2000\n A = adapt_stepsize(dual_averaging, A, dummy_acceptance_rate(current_ϵ(A)))\n end\n @test mean_dummy_acceptance_rate(final_ϵ(A)) ≈ δ atol = 0.01\nend\n\n@testset \"dual averaging far and noisy\" begin\n ϵ₀ = 20.0\n δ = 0.65\n dual_averaging = DualAveraging(; δ = δ)\n A = initial_adaptation_state(dual_averaging, ϵ₀)\n for _ in 1:10000\n A = adapt_stepsize(dual_averaging, A, dummy_acceptance_rate(current_ϵ(A), 2.0))\n end\n @test mean_dummy_acceptance_rate(final_ϵ(A), 2.0) ≈ δ atol = 0.04\nend\n\n@testset \"fixed stepsize sanity checks\" begin\n fs = FixedStepsize()\n ϵ = 1.0\n A = initial_adaptation_state(fs, ϵ)\n @test A == adapt_stepsize(fs, A, ϵ)\n @test current_ϵ(A) == ϵ\n @test final_ϵ(A) == ϵ\nend\n\n@testset \"find reasonable stepsize - random H, z\" begin\n p = InitialStepsizeSearch()\n _bkt(A, ϵ, C) = (A(ϵ) - p.log_threshold) * (A(ϵ * C) - p.log_threshold) ≤ 0\n for _ in 1:100\n (; H, z) = rand_Hz(rand(RNG, 3:5))\n A = local_log_acceptance_ratio(H, z)\n ϵ = find_initial_stepsize(p, A)\n @test _bkt(A, ϵ, 0.5) || _bkt(A, ϵ, 2.0)\n end\nend\n\n@testset \"error for non-finite initial density\" begin\n p = InitialStepsizeSearch()\n (; H, z) = rand_Hz(2)\n z = DynamicHMC.PhasePoint(z.Q, [NaN, NaN])\n @test_throws DynamicHMCError find_initial_stepsize(p, local_log_acceptance_ratio(H, z))\nend\n" }, { "path": "test/test_trees.jl", "content": "using DynamicHMC: Directions, next_direction, biased_progressive_logprob2, adjacent_tree,\n sample_trajectory, InvalidTree, is_turning\n\n####\n#### test directions mechanism\n####\n\n@testset \"directions\" begin\n directions = Directions(0b110101)\n is_forwards = Vector{Bool}()\n for i in 1:6\n is_forward, directions = next_direction(directions)\n push!(is_forwards, is_forward)\n end\n @test collect(is_forwards) == [true, false, true, false, true, true]\n @test rand(RNG, Directions).flags isa UInt32\nend\n\n####\n#### dummy trajectory for unit testing\n####\n\n\"\"\"\nTrajectory type that is easy to inspect and reason about, for unit testing.\n\nThe field `visited` keeps track of visited nodes, can be reset with `empty!`.\n\"\"\"\nstruct DummyTrajectory{L,T,D}\n \"Positions that trigger turning.\"\n turning::T\n \"Positions that trigger/mimic divergence.\"\n divergent::D\n \"Log density.\"\n ℓ::L\n \"Visited positions.\"\n visited::Vector{Int}\nend\n\nfunction DummyTrajectory(ℓ; turning = Set(Int[]), divergent = Set(Int[]))\n DummyTrajectory(turning, divergent, ℓ, Int[])\nend\n\nBase.empty!(trajectory) = empty(trajectory.visited)\n\nDynamicHMC.move(::DummyTrajectory, z, is_forward) = z + (is_forward ? 1 : -1)\n\nconst DUMMY_TURN_STATISTICS = Tuple{Bool,UnitRange{Int64}}\n\nfunction DynamicHMC.is_turning(::DummyTrajectory, τ::DUMMY_TURN_STATISTICS)\n turn_flag, positions = τ\n @test length(positions) > 1 # not called on a leaf\n turn_flag\nend\n\nfunction DynamicHMC.combine_turn_statistics(::DummyTrajectory,\n τ₁::DUMMY_TURN_STATISTICS,\n τ₂::DUMMY_TURN_STATISTICS)\n turn_flag1, positions1 = τ₁\n turn_flag2, positions2 = τ₂\n @test last(positions1) + 1 == first(positions2) # adjacency and order\n (turn_flag1 && turn_flag2, first(positions1):last(positions2))\nend\n\nfunction DynamicHMC.combine_visited_statistics(::DummyTrajectory, v₁, v₂)\n a1, s1 = v₁\n a2, s2 = v₂\n (a1 + a2, s1 + s2)\nend\n\nfunction DynamicHMC.combine_proposals(_, ::DummyTrajectory, zeta1, zeta2, logprob2, is_forward)\n lp2 = logprob2 > 0 ? 0.0 : logprob2\n lp1 = logprob2 > 0 ? oftype(lp2, -Inf) : log1mexp(lp2)\n if !is_forward\n # exchange so that we can test for adjacency, and join as a UnitRange\n zeta1, zeta2 = zeta2, zeta1\n lp1, lp2 = lp2, lp1\n end\n z1, p1 = zeta1\n z2, p2 = zeta2\n @test last(z1) + 1 == first(z2) # adjacency and order\n (first(z1):last(z2), vcat(p1 .+ lp1, p2 .+ lp2))\nend\n\nfunction DynamicHMC.calculate_logprob2(::DummyTrajectory, is_doubling, ω₁, ω₂, ω)\n biased_progressive_logprob2(is_doubling, ω₁, ω₂, ω)\nend\n\nfunction DynamicHMC.leaf(trajectory::DummyTrajectory, z, is_initial)\n (; turning, divergent, ℓ, visited) = trajectory\n d = z ∈ divergent\n is_initial && @argcheck !d # don't start with divergent\n Δ = ℓ(z)\n v = is_initial ? (0.0, 0) : (min(exp(Δ), 1), 1)\n !is_initial && push!(visited, z) # save position\n if d\n nothing, v\n else\n (((z:z, [0.0]), # ζ = nodes, log prob. within tree\n Δ, # ω = Δ for leaf\n (z ∈ turning, z:z)), # τ\n v)\n end\nend\n\n\"A log density for testing.\"\ntestℓ(z) = -abs2(z - 3) * 0.1\n\n\"Total acceptance rate of `ℓ` over `z`\"\ntestA(ℓ, z) = sum(min.(exp.(ℓ.(z)), 1))\n\n\"sum of acceptance rates for trajectory.\"\ntestA(trajectory::DummyTrajectory) = testA(trajectory.ℓ, trajectory.visited)\n\n@testset \"dummy adjacent tree full\" begin\n trajectory = DummyTrajectory(testℓ)\n (ζ, ω, τ, z′, i′), v = adjacent_tree(nothing, trajectory, 0, 0, 2, true)\n @test first(ζ) == 1:4\n @test sum(exp, last(ζ)) ≈ 1\n @test trajectory.visited == 1:4\n @test !is_turning(trajectory, τ)\n @test v[1] ≈ testA(trajectory)\n @test v[2] == 4\n @test z′ = i′ == 4\nend\n\n@testset \"dummy adjacent tree turning\" begin\n trajectory = DummyTrajectory(testℓ; turning = 5:7)\n t, v = adjacent_tree(nothing, trajectory, 0, 0, 3, true)\n @test trajectory.visited == 1:6 # [5,6] is turning\n @test t == InvalidTree(5, 6)\n @test v[1] == testA(trajectory)\n @test v[2] == 6\nend\n\n@testset \"dummy adjacent tree divergent\" begin\n trajectory = DummyTrajectory(testℓ; divergent = 5:7)\n t, v = adjacent_tree(nothing, trajectory, 0, 0, 3, true)\n @test trajectory.visited == 1:5 # 5 is divergent\n @test t == InvalidTree(5)\n @test v[1] ≈ testA(testℓ, 1:5)\n @test v[2] == 5\nend\n\n@testset \"dummy adjacent tree full backward\" begin\n trajectory = DummyTrajectory(testℓ)\n (ζ, ω, τ, z′, i′), v = adjacent_tree(nothing, trajectory, 0, 0, 3, false)\n @test first(ζ) == -8:-1\n @test sum(exp, last(ζ)) ≈ 1\n @test trajectory.visited == -(1:8)\n @test !is_turning(trajectory, τ)\n @test v[1] ≈ testA(testℓ, -(1:8))\n @test v[2] == 8\n @test z′ == i′ == -8\nend\n\n@testset \"dummy sampled tree\" begin\n trajectory = DummyTrajectory(testℓ)\n ζ, v, termination, depth = sample_trajectory(nothing, trajectory, 0, 3, Directions(0b101))\n @test trajectory.visited == [1, -1, -2, 2, 3, 4, 5]\n @test first(ζ) == -2:5\n @test sum(exp, last(ζ)) ≈ 1\n @test termination == DynamicHMC.REACHED_MAX_DEPTH\n @test v[1] ≈ testA(trajectory)\n @test v[2] == 7 # initial node does not participate in acceptance rate\nend\n\n####\n#### Detailed balance tests\n####\n\n\"An accumulator for log probabilities associated with a position on a trajectory.\"\nempty_accumulator() = Dict{Int,Float64}()\n\n\"Add log probabilities `πs` at positions `zs` into `accumulator`.\"\nfunction add_log_probabilities!(accumulator, zs, πs)\n for (z, π) in zip(zs, πs)\n accumulator[z] = haskey(accumulator, z) ? logaddexp(accumulator[z], π) : π\n end\n accumulator\nend\n\n\"Normalize an accumulator by depth.\"\nfunction normalize_accumulator(accumulator, depth)\n D = log(0.5) * depth\n Dict((k => v + D) for (k, v) in pairs(accumulator))\nend\n\n\"\"\"\nAn accumulator with the probability of visiting nodes for all trees with `depth`, strarting\nfrom `z`, on `trajectory`.\n\"\"\"\nfunction visited_log_probabilities(trajectory, z, depth)\n accumulator = empty_accumulator()\n for flags in 0:(2^depth - 1)\n ζ = first(sample_trajectory(nothing, trajectory, z, depth, Directions(UInt32(flags))))\n add_log_probabilities!(accumulator, ζ...)\n end\n normalize_accumulator(accumulator, depth)\nend\n\n\"\"\"\nThe probability of visiting node `z′` for all trees with `depth`, strarting from `z`, on\n`trajectory`.\n\"\"\"\nfunction transition_log_probability(trajectory, z, z′, depth)\n p = -Inf\n for flags in 0:(2^depth - 1)\n zs, πs = first(sample_trajectory(nothing, trajectory, z, depth,\n Directions(UInt32(flags))))\n ix = findfirst(isequal(z′), zs)\n if ix ≠ nothing\n p = logaddexp(p, πs[ix])\n end\n end\n p + depth * log(0.5)\nend\n\n@testset \"transition calculations consistency check\" begin\n trajectory = DummyTrajectory(testℓ)\n depth = 5\n z = 9\n for (z′, π) in pairs(visited_log_probabilities(trajectory, z, depth))\n @test π ≈ transition_log_probability(trajectory, z, z′, depth)\n end\nend\n\n\"\"\"\n$(SIGNATURES)\n\nFor all transitions from `z`, test the detailed balance condition, ie\n\n``ℙ(z) ℙ(j ∣ z) = ℙ(j) ℙ(z ∣ j)``\n\nwhere ``ℙ(z) = exp(ℓ(z))`` and the transition probabilities ``ℙ(⋅∣⋅)`` are\ncalculated using `visited_log_probabilities` and `transition_log_probability`.\n\n(We use logs for more accurate calculation.)\n\"\"\"\nfunction test_detailed_balance(trajectory, z, depth; atol = √eps())\n (; ℓ) = trajectory\n ℓz = ℓ(z)\n for (z′, π) in pairs(visited_log_probabilities(trajectory, z, depth))\n π′ = transition_log_probability(trajectory, z′, z, depth)\n @test (π + ℓz) ≈ (π′ + ℓ(z′)) atol = atol\n end\nend\n\n@testset \"detailed balance\" begin\n for max_depth in 1:5\n test_detailed_balance(DummyTrajectory(testℓ), 0, max_depth)\n end\n for max_depth in 1:5\n test_detailed_balance(DummyTrajectory(testℓ; turning = 1:2), 3, max_depth)\n end\n for max_depth in 1:6\n test_detailed_balance(DummyTrajectory(testℓ; divergent = 10:11), 3, max_depth)\n end\n for max_depth in 1:6\n test_detailed_balance(DummyTrajectory(testℓ; divergent = 10:12, turning = -3:-2), 3,\n max_depth)\n end\nend\n" }, { "path": "test/utilities.jl", "content": "\"\"\"\n$(SIGNATURES)\n\nRandom positive definite matrix of size `n` x `n` (for testing).\n\"\"\"\nfunction rand_Σ(::Type{Symmetric}, n)\n A = randn(RNG, n, n)\n Symmetric(A'*A .+ 0.01)\nend\n\nrand_Σ(::Type{Diagonal}, n) = Diagonal(randn(RNG, n).^2 .+ 0.01)\n\nrand_Σ(n::Int) = rand_Σ(Symmetric, n)\n\n\"\"\"\n$(SIGNATURES)\n\nSimulated mean and covariance of `N` values from `f()`.\n\"\"\"\nfunction simulated_meancov(f, N)\n s = f()\n K = length(s)\n x = similar(s, (N, K))\n for i in 1:N\n x[i, :] = f()\n end\n m, C = mean_and_cov(x, 1)\n vec(m), C\nend\n\n@testset \"simulated meancov\" begin\n μ = [2, 1.2]\n D = [2.0, 0.7]\n m, C = simulated_meancov(()-> randn(RNG, 2) .* D .+ μ, 10000)\n @test m ≈ μ atol = 0.05 rtol = 0.1\n @test C ≈ Diagonal(abs2.(D)) atol = 0.05 rtol = 0.1\nend\n\n###\n### Hamiltonian test helper functions\n###\n\n\"\"\"\n$(SIGNATURES)\n\nReturn a reasonable estimate for the largest stable stepsize (which may not be\nstable, but is a good starting point for finding that).\n\n`q` is assumed to be normal with variance `Σ`. `κ` is the kinetic energy.\n\nUsing the transformation ``p̃ = W⁻¹ p``, the kinetic energy is\n\n``p'M⁻¹p = p'W⁻¹'W⁻¹p/2=p̃'p̃/2``\n\nTransforming to ``q̃=W'q``, the variance of which becomes ``W' Σ W``. Return the\nsquare root of its smallest eigenvalue, following Neal (2011, p 136).\n\nWhen ``Σ⁻¹=M=WW'``, this the variance of `q̃` is ``W' Σ W=W' W'⁻¹W⁻¹W=I``, and\nthus decorrelates the density perfectly.\n\"\"\"\nfind_stable_ϵ(κ::GaussianKineticEnergy, Σ) = √eigmin(κ.W'*Σ*κ.W)\n\n\"Multivariate normal with `Σ = LL'`.\"\nmultivariate_normal(μ, L) = (shift(μ) ∘ linear(L))(StandardMultivariateNormal(length(μ)))\n\n\"Multivariate normal with diagonal `Σ` (constant `v` variance).\"\nmultivariate_normal(μ, v::Real = 1) = multivariate_normal(μ, I(length(μ)) * v)\n\n\"\"\"\n$(SIGNATURES)\n\nA `NamedTuple` that contains\n\n- a random `K`-element vector `μ`\n\n- a random `K × K` covariance matrix `Σ`,\n\n- a random Hamiltonian `H` with `ℓ` corresponding to a multivariate normal with `μ`, `Σ`,\n and a random Gaussian kinetic energy (unrelated to `ℓ`).\n\n- a random phasepoint `z`.\n\nUseful for testing.\n\"\"\"\nfunction rand_Hz(K)\n μ = randn(K)\n Σ = rand_Σ(K)\n L = cholesky(Σ).L\n κ = GaussianKineticEnergy(inv(rand_Σ(Diagonal, K)))\n ℓ = multivariate_normal(μ, L)\n H = Hamiltonian(κ, ℓ)\n q = rand(RNG, ℓ)\n p = rand_p(RNG, κ)\n z = PhasePoint(evaluate_ℓ(H.ℓ, q), rand_p(RNG, κ))\n (μ = μ, Σ = Σ, H = H, z = z)\nend\n" } ]