Package: pymc3-ext
PyMC3 ExtrasThis library include various experimental or otherwise special purpose extras for use with PyMC3 that have been extracted from theexoplanetproject.
The most widely useful component is probably the custom tuning functions for the PyMC3 NUTS sampler that isdescribed below, but it also includes some helper functions fornon-linear optimizationandsome custom distributions.You'll find the usage instructions below and automatically generated tutorial notebooks on thenotebooksbranch on GitHub.InstallationYou'll need a Python installation (tested on versions 3.7 and 3.8) and it can often be best to install PyMC3 usingcondaso that it can handle all the details of compiler setup.
This step isoptional, but I would normally create a clean conda environment for projects that use PyMC3:# Optionalcondacreate-nname-of-my-projectpython=3.8pymc3
condaactivatename-of-my-projectThe easiest way to install this package is usingpip:python-mpipinstall-Upymc3-extThis will also update the dependencies like PyMC3, which is probably what you want because this is only tested on recent versions of both of those packages.NUTS tuningThe main function that thepymc3-extlibrary provides is a drop-in replacement for thepm.samplefunction in PyMC3:importpymc3_extaspmxwithmodel:trace=pmx.sample()This adjusts the defaults (dense mass matrix adaptation and higher target acceptance fraction, for example) to be more in line with the needs of users in astrophysics, but it also provides some other features that might improve sampling performance.The main features are described below, but more details can be found in theSampling tutorial notebook.Tuning scheduleOne main difference betweenpmx.sampleand the standard one implemented in PyMC3 is the tuning schedule.
In PyMC3, the default tuning schedule is very simple:(a)update the step size parameter every step, and(b)update the mass matrix everyNsteps.
In reality the procedure is slightly more convoluted, but this is the gist.pmx.sampleinstead uses a "fast" and "slow" update schedule with doubling adaptation windows that is nearly identical tothe scheme used by the Stan project.
We have found that this can significantly outperform the default algorithm and rarely leads to worse performance.Parameter groupsSince thispmx.samplefunction defaults to estimating a dense mass matrix, the memory and computational requirements (as well as the numerical error) for the sampler can grow significantly for problems with a large number of parameters.
This library includes support for grouping parameters that are known to have covariances so that the mass matrix is block diagonal rather than dense.
For example, in the following snippet, the parameterxis known to be correlated, while the parameteryonly requires a diagonal mass matrix:importpymc3aspmimportpymc3_extaspmxwithpm.Model():# `L` is the Cholesky factor of a covariance matrix with off-diagonal elementsx=pm.MvNormal("x",mu=np.zeros(L.shape[0]),chol=L,shape=L.shape[0])y=pm.Normal("y",shape=5)trace=pmx.sample(parameter_groups=[[x],pmx.sampling.ParameterGroup([y],"diag"),],)Variable target acceptance fraction (experimental)Thepmx.samplefunction also includes support for an experimental feature where the target acceptance fraction is adjusted throughout the tuning phase.
This can be useful when the early warm up windows are slow to run because a high target acceptance fraction will require tiny step sizes that aren't well suited for searching for the typical set.
This feature can be used by providing theinitial_acceptparameter (usually set to something like0.5) topmx.sample.
This will be the target acceptance fraction at the beginning and it will be increased totarget_accept(0.9by default) throught the tuning phase.OptimizationWhen PyMC3 added a warning to thepm.find_MAPfunction, we implemented a custom non-linear optimization framework inexoplanetbecause it is often useful to be able to optimize (at least) some parameters when initializing the sampler for many problems in astrophysics (and probably elsewhere).
Whilepm.find_MAPno longer complains, thepymc3_ext.optimizefunction is included here for backwards compatibility even though it should have similar behavior topm.find_MAP.
To use this function, you'll do something like the following:importpymc3_extaspmxwithmodel:soln=pmx.optimize(vars=[var1,var2])soln=pmx.optimize(start=soln,vars=[var3])DistributionsMost of the custom distributions in this library are there to make working with periodic parameters (like angles) easier.
All of these reparameterizations could be implemented manually without too much trouble, but it can be useful to have them in a more compact form.
Here is a list of the included distributions and a short description:pmx.UnitVector: A vector where the sum of squares is fixed to unity. For a multidimensional shape, the normalization is performed along the last dimension.pmx.UnitDisk: Two dimensional parameters constrianed to live within the unit disk. This will be useful when you have an angle and a magnitude that must be in the range zero to one (for example, an eccentricity vector for a bound orbit). This distribution is constrained such that the sum of squares along the zeroth axis will always be less than one. Note that the shape of this distribution must be two in the zeroth axis.pmx.Angle: An angle constrained to be in the range -pi to pi. The actual sampling is performed in the two dimensional vector space(sin(theta), cos(theta))so that the sampler doesn't see a discontinuity at pi. As a technical detail, the performance of this distribution can be affected using theregularizationparameter which helps deal with pathelogical geometries introduced when this parameter is well/poorly constrained. The default value (10.0) was selected as a reasonable default choice, but you might get better performance by adjusting this.pmx.Periodic: An extension topmx.Anglethat supports arbitrary upper and lower bounds for the allowed range.pmx.UnitUniform: This distribution is equivalent topm.Uniform(lower=0, upper=1), but it can be more numerically stable in some cases.LicenseCopyright 2020 Dan Foreman-Mackey and contributors.pymc3-ext is free software made available under the MIT License. For details see the LICENSE file.