### 10 releases (6 breaking)

0.8.0 | Sep 20, 2023 |
---|---|

0.6.0 | Jul 21, 2023 |

0.4.1 | Jan 2, 2023 |

0.4.0 | Nov 12, 2022 |

0.2.1 | Jul 20, 2022 |

#**1** in #nuts

**709** downloads per month

**MIT**license

265KB

3K
SLoC

Sample from posterior distributions using the No U-turn Sampler (NUTS). For details see the original NUTS paper and the more recent introduction.

This crate was developed as a faster replacement of the sampler in PyMC, to be used with the new numba backend of aesara. The python wrapper for this sampler is nutpie.

## Usage

`use` `nuts_rs``::``{`CpuLogpFunc`,` LogpError`,` new_sampler`,` SamplerArgs`,` Chain`,` SampleStats`}``;`
`use` `thiserror``::`Error`;`
`//` Define a function that computes the unnormalized posterior density
`//` and its gradient.
`struct` `PosteriorDensity` `{``}`
`//` The density might fail in a recoverable or non-recoverable manner...
`#``[``derive``(``Debug``,` Error`)``]`
`enum` `PosteriorLogpError` `{``}`
`impl` `LogpError ``for`` ``PosteriorLogpError` `{`
`fn` `is_recoverable``(``&``self``)`` ``->` `bool` `{` `false` `}`
`}`
`impl` `CpuLogpFunc ``for`` ``PosteriorDensity` `{`
`type` `Err` `=` PosteriorLogpError`;`
`//` We define a 10 dimensional normal distribution
`fn` `dim``(``&``self``)`` ``->` `usize` `{` `10` `}`
`//` The normal likelihood with mean 3 and its gradient.
`fn` `logp``(``&``mut` `self`, `position``:` `&`[`f64`], `grad``:` `&``mut` [`f64`]`)`` ``->` `Result``<``f64`, `Self``::``Err``>` `{`
`let` mu `=` `3``f64``;`
`let` logp `=` position
`.``iter``(``)`
`.``copied``(``)`
`.``zip``(`grad`.``iter_mut``(``)``)`
`.``map``(``|``(``x``,` `grad``)``|` `{`
`let` diff `=` x `-` mu`;`
`*`grad `=` `-`diff`;`
`-`diff `*` diff `/` `2``f64`
`}``)`
`.``sum``(``)``;`
`return` `Ok``(`logp`)`
`}`
`}`
`//` We get the default sampler arguments
`let` `mut` sampler_args `=` `SamplerArgs``::`default`(``)``;`
`//` and modify as we like
sampler_args`.`step_size_adapt`.`target `=` `0.``8``;`
sampler_args`.`num_tune `=` `1000``;`
sampler_args`.`maxdepth `=` `3``;` `//` small value just for testing...
sampler_args`.`mass_matrix_adapt`.`store_mass_matrix `=` `true``;`
`//` We instanciate our posterior density function
`let` logp_func `=` PosteriorDensity `{``}``;`
`let` chain `=` `0``;`
`let` seed `=` `42``;`
`let` `mut` sampler `=` `new_sampler``(`logp_func`,` sampler_args`,` chain`,` seed`)``;`
`//` Set to some initial position and start drawing samples.
sampler`.``set_position``(``&``vec!``[``0``f64``;` `10``]``)``.``expect``(``"`Unrecoverable error during init`"``)``;`
`let` `mut` trace `=` `vec!``[``]``;` `//` Collection of all draws
`let` `mut` stats `=` `vec!``[``]``;` `//` Collection of statistics like the acceptance rate for each draw
`for` `_` `in` `0``..``2000` `{`
`let` `(`draw`,` info`)` `=` sampler`.``draw``(``)``.``expect``(``"`Unrecoverable error during sampling`"``)``;`
trace`.``push``(`draw`)``;`
`let` _info_vec `=` info`.``to_vec``(``)``;` `//` We can collect the stats in a Vec
`//` Or get more detailed information about divergences
`if` `let` `Some``(`div_info`)` `=` info`.``divergence_info``(``)` `{`
`println!``(``"`Divergence at position `{:?}``"``,` div_info`.``start_location``(``)``)``;`
`}`
`dbg!``(``&`info`)``;`
stats`.``push``(`info`)``;`
`}`

Sampling several chains in parallel so that samples are accessable as they are generated
is implemented in

.`sample_parallel`

## Implementation details

This crate mostly follows the implementation of NUTS in Stan and PyMC, only tuning of mass matrix and step size differs somewhat.

#### Dependencies

~4–5.5MB

~116K SLoC