Skip to content

Commit c392350

Browse files
authored
Docs (#29)
* Touch up docstrings * Move integrator and step size adaptation to top-level package * Add step sizes and integrators to API docs * Remove step sizes and integration from top-level * STY: black line length 100 * DOC: more documentation * REV: remove python3 * MAINT: refactor stats reshaping * TST: improve tests
1 parent 2f89c79 commit c392350

12 files changed

+309
-143
lines changed

docs/_static/notebooks/quickstart.ipynb

+180-4
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,33 @@
44
"cell_type": "markdown",
55
"metadata": {},
66
"source": [
7-
"# Quickstart"
7+
"# LittleMCMC Quickstart\n",
8+
"\n",
9+
"LittleMCMC is a lightweight and performant implementation of HMC and NUTS in Python, spun out of the PyMC project. In this quickstart tutorial, we will introduce LittleMCMC\n",
10+
"\n",
11+
"## Table of Contents\n",
12+
"\n",
13+
"- [Who should use LittleMCMC?](#Who-should-use-LittleMCMC?)\n",
14+
"- [Sampling](#Sampling)\n",
15+
" - [Inspecting the Output of `lmc.sample`](#Inspecting-the-Output-of-lmc.sample)\n",
16+
"- [Other Modules](#Other-Modules)\n",
17+
"\n",
18+
"## Who should use LittleMCMC?\n",
19+
"\n",
20+
"<div class=\"alert alert-block alert-info\">\n",
21+
"LittleMCMC is a fairly bare bones library with a very niche use case. Most users will probably find that [PyMC3](https://github.com/pymc-devs/pymc3) will satisfy their needs, with better strength of support and quality of documentation.\n",
22+
"</div>\n",
23+
"\n",
24+
"If you:\n",
25+
"\n",
26+
"1. Have model with only continuous parameters,\n",
27+
"1. Are willing to manually \"unconstrain\" all of your model's parameters (if necessary),\n",
28+
"1. Have methods to compute the log probability of the model and its derivative, exposed via a Python callable,\n",
29+
"1. And all you need is an implementation of HMC/NUTS (preferably in Python) to sample from your model\n",
30+
"\n",
31+
"then you should consider using LittleMCMC!\n",
32+
"\n",
33+
"## Sampling"
834
]
935
},
1036
{
@@ -38,10 +64,160 @@
3864
},
3965
{
4066
"cell_type": "code",
41-
"execution_count": null,
67+
"execution_count": 3,
4268
"metadata": {},
43-
"outputs": [],
44-
"source": []
69+
"outputs": [
70+
{
71+
"name": "stderr",
72+
"output_type": "stream",
73+
"text": [
74+
"/Users/george/miniconda3/lib/python3.6/site-packages/ipykernel_launcher.py:2: RuntimeWarning: divide by zero encountered in log\n",
75+
" \n",
76+
"/Users/george/miniconda3/lib/python3.6/site-packages/ipykernel_launcher.py:2: RuntimeWarning: divide by zero encountered in log\n",
77+
" \n"
78+
]
79+
},
80+
{
81+
"name": "stdout",
82+
"output_type": "stream",
83+
"text": [
84+
"\n"
85+
]
86+
}
87+
],
88+
"source": [
89+
"trace, stats, results = lmc.sample(\n",
90+
" logp_dlogp_func=logp_dlogp_func,\n",
91+
" size=1,\n",
92+
" draws=1000,\n",
93+
" tune=500,\n",
94+
" step=lmc.NUTS(logp_dlogp_func=logp_dlogp_func, size=1),\n",
95+
" chains=4,\n",
96+
" cores=4,\n",
97+
" progressbar=\"notebook\"\n",
98+
")"
99+
]
100+
},
101+
{
102+
"cell_type": "markdown",
103+
"metadata": {},
104+
"source": [
105+
"### Inspecting the Output of `lmc.sample`"
106+
]
107+
},
108+
{
109+
"cell_type": "code",
110+
"execution_count": 4,
111+
"metadata": {},
112+
"outputs": [
113+
{
114+
"data": {
115+
"text/plain": [
116+
"array([-0.38331274, -1.76994233, -0.67234733, ..., 0.27817656,\n",
117+
" 0.29250676, 0.42966184])"
118+
]
119+
},
120+
"execution_count": 4,
121+
"metadata": {},
122+
"output_type": "execute_result"
123+
}
124+
],
125+
"source": [
126+
"trace"
127+
]
128+
},
129+
{
130+
"cell_type": "code",
131+
"execution_count": 5,
132+
"metadata": {},
133+
"outputs": [
134+
{
135+
"data": {
136+
"text/plain": [
137+
"(4000,)"
138+
]
139+
},
140+
"execution_count": 5,
141+
"metadata": {},
142+
"output_type": "execute_result"
143+
}
144+
],
145+
"source": [
146+
"trace.shape"
147+
]
148+
},
149+
{
150+
"cell_type": "code",
151+
"execution_count": 6,
152+
"metadata": {},
153+
"outputs": [
154+
{
155+
"data": {
156+
"text/plain": [
157+
"{'depth': array([1, 1, 1, ..., 1, 2, 1]),\n",
158+
" 'step_size': array([0.94586326, 0.94586326, 0.94586326, ..., 2.16938615, 2.16938615,\n",
159+
" 2.16938615]),\n",
160+
" 'tune': array([False, False, False, ..., False, False, False]),\n",
161+
" 'mean_tree_accept': array([1. , 0.43665689, 1. , ..., 0.98765583, 0.72296808,\n",
162+
" 0.97965297]),\n",
163+
" 'step_size_bar': array([1.20597596, 1.20597596, 1.20597596, ..., 1.28614833, 1.28614833,\n",
164+
" 1.28614833]),\n",
165+
" 'tree_size': array([1., 1., 1., ..., 1., 3., 1.]),\n",
166+
" 'diverging': array([False, False, False, ..., False, False, False]),\n",
167+
" 'energy_error': array([-0.25675836, 0.82860753, -0.74393026, ..., 0.01242099,\n",
168+
" 0.00169732, 0.02055688]),\n",
169+
" 'energy': array([1.25393394, 2.56056236, 1.91071276, ..., 0.95981431, 1.76229677,\n",
170+
" 1.02575724]),\n",
171+
" 'max_energy_error': array([-0.25675836, 0.82860753, -0.74393026, ..., 0.01242099,\n",
172+
" 0.56981615, 0.02055688]),\n",
173+
" 'model_logp': array([-0.99240286, -2.48528646, -1.144964 , ..., -0.95762963,\n",
174+
" -0.96171864, -1.01124318])}"
175+
]
176+
},
177+
"execution_count": 6,
178+
"metadata": {},
179+
"output_type": "execute_result"
180+
}
181+
],
182+
"source": [
183+
"stats"
184+
]
185+
},
186+
{
187+
"cell_type": "code",
188+
"execution_count": 7,
189+
"metadata": {},
190+
"outputs": [
191+
{
192+
"data": {
193+
"text/plain": [
194+
"(4000,)"
195+
]
196+
},
197+
"execution_count": 7,
198+
"metadata": {},
199+
"output_type": "execute_result"
200+
}
201+
],
202+
"source": [
203+
"stats[\"diverging\"].shape"
204+
]
205+
},
206+
{
207+
"cell_type": "markdown",
208+
"metadata": {},
209+
"source": [
210+
"## Other Modules\n",
211+
"\n",
212+
"LittleMCMC exposes:\n",
213+
"\n",
214+
"1. Two step methods: Hamiltonian Monte Carlo (HMC) and the No-U-Turn Sampler (NUTS)\n",
215+
"1. Quadpotentials (a.k.a. mass matrices or inverse metrics)\n",
216+
"1. Dual-averaging step size adaptation\n",
217+
"1. Leapfrog integration\n",
218+
"\n",
219+
"Refer to the [API Reference](https://littlemcmc.readthedocs.io/en/latest/api.html) for more information."
220+
]
45221
}
46222
],
47223
"metadata": {

docs/api.rst

+21
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@ Step Methods
2626
HamiltonianMC
2727
NUTS
2828

29+
.. _quadpotentials_api:
30+
2931
Quadpotentials (a.k.a. Mass Matrices)
3032
-------------------------------------
3133

@@ -39,3 +41,22 @@ Quadpotentials (a.k.a. Mass Matrices)
3941
QuadPotentialDiagAdapt
4042
QuadPotentialFullAdapt
4143

44+
.. _step_sizes_api:
45+
46+
Dual Averaging Step Size Adaptation
47+
-----------------------------------
48+
49+
.. autosummary::
50+
:toctree: generated/
51+
52+
step_sizes.DualAverageAdaptation
53+
54+
.. _integrators_api:
55+
56+
Integrators
57+
-----------
58+
59+
.. autosummary::
60+
:toctree: generated/
61+
62+
integration.CpuLeapfrogIntegrator

littlemcmc/base_hmc.py

+27-23
Original file line numberDiff line numberDiff line change
@@ -54,25 +54,34 @@ def __init__(
5454
the log-probability, respectively.
5555
size : int
5656
Total number of parameters. Dimensionality of the output of
57-
`logp_dlogp_func`.
57+
``logp_dlogp_func``.
5858
scaling : 1 or 2-dimensional array-like
5959
Scaling for momentum distribution. 1 dimensional arrays are
60-
interpreted as a matrix diagonal. Only one of `scaling` or
61-
`potential` may be non-None.
60+
interpreted as a matrix diagonal. Only one of ``scaling`` or
61+
``potential`` may be non-None.
6262
is_cov : bool
6363
Treat scaling as a covariance matrix/vector if True, else treat
6464
it as a precision matrix/vector
6565
potential : littlemcmc.quadpotential.Potential, optional
66-
An object that represents the Hamiltonian with methods `velocity`,
67-
`energy`, and `random` methods. Only one of `scaling` or `potential`
68-
may be non-None.
66+
An object that represents the Hamiltonian with methods ``velocity``,
67+
``energy``, and ``random`` methods. Only one of ``scaling`` or
68+
``potential`` may be non-None.
6969
target_accept : float
70+
Adapt the step size such that the average acceptance probability
71+
across the trajectories are close to target_accept. Higher values
72+
for target_accept lead to smaller step sizes. Setting this to higher
73+
values like 0.9 or 0.99 can help with sampling from difficult
74+
posteriors. Valid values are between 0 and 1 (exclusive).
7075
Emax : float
76+
The maximum allowable change in the value of the Hamiltonian. Any
77+
trajectories that result in changes in the value of the Hamiltonian
78+
larger than ``Emax`` will be declared divergent.
7179
adapt_step_size : bool, default=True
7280
If True, performs dual averaging step size adaptation. If False,
73-
`k`, `t0`, `gamma` and `target_accept` are ignored.
81+
``k``, ``t0``, ``gamma`` and ``target_accept`` are ignored.
7482
step_scale : float
75-
Size of steps to take, automatically scaled down by 1 / (size ** 0.25)
83+
Size of steps to take, automatically scaled down by 1 / (``size`` **
84+
0.25).
7685
gamma : float, default .05
7786
k : float, default .75
7887
Parameter for dual averaging for step size adaptation. Values
@@ -81,8 +90,9 @@ def __init__(
8190
t0 : int, default 10
8291
Parameter for dual averaging. Higher values slow initial adaptation.
8392
step_rand : Python callable
84-
Called on step size to randomize, immediately before adapting step
85-
size.
93+
Callback for step size adaptation. Called on the step size at each
94+
iteration immediately before performing dual-averaging step size
95+
adaptation.
8696
"""
8797
self._logp_dlogp_func = logp_dlogp_func
8898
self.adapt_step_size = adapt_step_size
@@ -109,9 +119,7 @@ def __init__(
109119
else:
110120
self.potential = quad_potential(scaling, is_cov)
111121

112-
self.integrator = integration.CpuLeapfrogIntegrator(
113-
self.potential, self._logp_dlogp_func
114-
)
122+
self.integrator = integration.CpuLeapfrogIntegrator(self.potential, self._logp_dlogp_func)
115123
self._step_rand = step_rand
116124
self._warnings: List[SamplerWarning] = []
117125
self._samples_after_tune = 0
@@ -122,9 +130,7 @@ def stop_tuning(self) -> None:
122130
if hasattr(self, "tune"):
123131
self.tune = False
124132

125-
def _hamiltonian_step(
126-
self, start: np.ndarray, p0: np.ndarray, step_size: float
127-
) -> HMCStepData:
133+
def _hamiltonian_step(self, start: np.ndarray, p0: np.ndarray, step_size: float) -> HMCStepData:
128134
"""Compute one Hamiltonian trajectory and return the next state.
129135
130136
Subclasses must overwrite this method and return a `HMCStepData`.
@@ -138,9 +144,7 @@ def _astep(self, q0: np.ndarray):
138144

139145
if not np.isfinite(start.energy):
140146
raise ValueError(
141-
"Bad initial energy: {}. The model might be misspecified.".format(
142-
start.energy
143-
)
147+
"Bad initial energy: {}. The model might be misspecified.".format(start.energy)
144148
)
145149

146150
# Adapt step size
@@ -194,7 +198,9 @@ def warnings(self) -> List[SamplerWarning]:
194198
message = ""
195199
n_divs = self._num_divs_sample
196200
if n_divs and self._samples_after_tune == n_divs:
197-
message = "The chain contains only diverging samples. The model is probably misspecified."
201+
message = (
202+
"The chain contains only diverging samples. The model is probably misspecified."
203+
)
198204
elif n_divs == 1:
199205
message = (
200206
"There was 1 divergence after tuning. Increase "
@@ -207,9 +213,7 @@ def warnings(self) -> List[SamplerWarning]:
207213
)
208214

209215
if message:
210-
warning = SamplerWarning(
211-
WarningType.DIVERGENCES, message, "error", None, None, None
212-
)
216+
warning = SamplerWarning(WarningType.DIVERGENCES, message, "error", None, None, None)
213217
warnings.append(warning)
214218

215219
warnings.extend(self.step_adapt.warnings())

0 commit comments

Comments
 (0)