Package: numpyro-ext
Extensions for NumPyroThis library includes a miscellaneous set of helper functions, custom
distributions, and other utilities that I find useful when usingNumPyroin my work.InstallationSince NumPyro, and hence this library, are built on top of JAX, it's typically
good practice to start by installing JAX followingthe installation
instructions. Then, you can
install this library using pip:python-mpipinstallnumpyro-extUsageSince this README is checked usingdoctest, let's start by importing some
common modules that we'll need in all our examples:>>>importjax>>>importjax.numpyasjnp>>>importnumpyro>>>importnumpyro_extDistributionsThe tradition is to importnumpyro_ext.distributionsasdistxto
differentiate fromnumpyro.distributions, which is imported asdist:>>>fromnumpyroimportdistributionsasdist>>>fromnumpyro_extimportdistributionsasdistx>>>key=jax.random.PRNGKey(0)AngleA uniform distribution over angles in radians. The actual sampling is performed
in the two-dimensional vector space proportional to(sin(theta), cos(theta))so that the sampler doesn't see a discontinuity at pi.>>>angle=distx.Angle()>>>print(angle.sample(key,(2,3)))[[0.4...][2.4...]]UnitDiskA uniform distribution over two-dimensional points within the disk of radius 1.
This means that the sum over squares of the last dimension of a random variable
generated from this distribution will always be less than 1.>>>unit_disk=distx.UnitDisk()>>>u=unit_disk.sample(key,(5,))>>>print(jnp.sum(u**2,axis=-1))[0.07...]NoncentralChi2Anon-central chi-squared
distribution.
To use this distribution, you'll need to install the optionaltensorflow-probabilitydependency.>>>ncx2=distx.NoncentralChi2(df=3,nc=2.)>>>print(ncx2.sample(key,(5,)))[2.19...]MarginalizedLinearThe marginalized product of two (possibly multivariate) normal distributions
with a linear relationship between them. The mathematical details of these
models are discussed in detail inthis note,
and this distribution implements the math presented there, in a computationally
efficient way, assuming that the number of marginalized parameters is small
compared to the size of the dataset.The following example shows a particularly simple example of a
fully-marginalized model for fitting a line to data:>>>defmodel(x,y=None):...design_matrix=jnp.vander(x,2)...prior=dist.Normal(0.0,1.0)...data=dist.Normal(0.0,2.0)...numpyro.sample(..."y",...distx.MarginalizedLinear(design_matrix,prior,data),...obs=y...)...Things get a little more interesting when the design matrix and/or the
distributions are functions of non-linear parameters. For example, if we want to
find the period of a sinusoidal signal, also fitting for some unknown excess
measurement uncertainty (often called "jitter") we can use the following model:>>>defmodel(x,y_err,y=None):...period=numpyro.sample("period",dist.Uniform(1.0,250.0))...ln_jitter=numpyro.sample("ln_jitter",dist.Normal(0.0,2.0))...design_matrix=jnp.stack(...[...jnp.sin(2*jnp.pi*x/period),...jnp.cos(2*jnp.pi*x/period),...jnp.ones_like(x),...],...axis=-1,...)...prior=dist.Normal(0.0,10.0).expand([3])...data=dist.Normal(0.0,jnp.sqrt(y_err**2+jnp.exp(2*ln_jitter)))...numpyro.sample(..."y",...distx.MarginalizedLinear(design_matrix,prior,data),...obs=y...)...>>>x=jnp.linspace(-1.0,1.0,5)>>>samples=numpyro.infer.Predictive(model,num_samples=2)(key,x,0.1)>>>print(samples["period"])[......]>>>print(samples["y"])[[.........][.........]]It's often useful to also track conditional samples of the marginalized
parameters during inference. The conditional distribution can be accessed using
theconditionalmethod onMarginalizedLinear:>>>x=jnp.linspace(-1.0,1.0,5)>>>y=jnp.sin(x)# just some fake data>>>design_matrix=jnp.vander(x,2)>>>prior=dist.Normal(0.0,1.0)>>>data=dist.Normal(0.0,2.0)>>>marg=distx.MarginalizedLinear(design_matrix,prior,data)>>>cond=marg.conditional(y)>>>print(type(cond).__name__)MultivariateNormal>>>print(cond.sample(key,(3,)))[[...][...][...]]OptimizationThe inference lore is a little mixed on the benefits of optimization as an
initialization tool for MCMC, but I find that at least in a lot of astronomy
applications, an initial optimization can make a huge difference in performance.
Even if you don't want to use the optimization results as an initialization, it
can still sometimes be useful to numerically search for the maximuma
posterioriparameters for your model. However, the NumPyro interface for these
types of optimization isn't terribly user-friendly, so this library provides
some helpers to make it a little more straightforward.By default, this optimization uses the wrappers of scipy's optimization routines
provided by theJAXoptlibrary, so you'll
need to install JAXopt:python-mpipinstalljaxoptbefore running these examples.The following example shows a simple optimization of a model with a single
parameter:>>>fromnumpyro_extimportoptimasoptimx>>>>>>defmodel(y=None):...x=numpyro.sample("x",dist.Normal(0.0,1.0))...numpyro.sample("y",dist.Normal(x,2.0),obs=y)...>>>soln=optimx.optimize(model)(key,y=0.5)By default, the optimization starts from a prior sample, but you can provide
custom initial coordinates as follows:>>>soln=optimx.optimize(model,start={"x":12.3})(key,y=0.5)Similarly, if you only want to optimize a subset of the parameters, you can
provide a list of parameters to target:>>>soln=optimx.optimize(model,sites=["x"])(key,y=0.5)Information matrix computationThe Fisher information matrix for models with Gaussian likelihoods isstraightforward to
compute,
and this library provides a helper function for automating this computation:>>>fromnumpyro_extimportinformation>>>>>>defmodel(x,y=None):...a=numpyro.sample("a",dist.Normal(0.0,1.0))...b=numpyro.sample("b",dist.Normal(0.0,1.0))...log_alpha=numpyro.sample("log_alpha",dist.Normal(0.0,1.0))...cov=jnp.exp(log_alpha-0.5*(x[:,None]-x[None,:])**2)...cov+=0.1*jnp.eye(len(x))...numpyro.sample(..."y",...dist.MultivariateNormal(loc=a*x+b,covariance_matrix=cov),...obs=y,...)...>>>x=jnp.linspace(-1.0,1.0,5)>>>y=jnp.sin(x)# the input data just needs to have the right shape>>>params={"a":0.5,"b":-0.2,"log_alpha":-0.5}>>>info=information(model)(params,x,y=y)>>>print(info){'a':{'a':...,'b':...'log_alpha':...},'b':...}The returned information matrix is a nested dictionary of dictionaries, indexed
by pairs of parameter names, where the values are the corresponding blocks of
the information matrix.