Package: nifty8
NIFTy - Numerical Information Field TheoryNIFTyproject homepage:ift.pages.mpcdf.de/nifty| Found a bug?github.com/nifty-ppl/nifty/issues| Need help?github.com/nifty-ppl/nifty/discussionsSummaryDescriptionNIFTy, "NumericalInformationFieldTheory", is a Bayesian imaging library.
It is designed to infer the million to billion dimensional posterior distribution in the image space from noisy input data.
At the core of NIFTy lies a set of powerful Gaussian Process (GP) models and accurate Variational Inference (VI) algorithms.Gaussian ProcessesOne standard tool from the NIFTy toolbox are the structured GP models.
These models usually rely on the harmonic domain being easily accessible, e.g. for pixels spaced on a regular Cartesian grid, the natural choice to represent a stationary kernel is the Fourier domain.
An example, initializing a non-parameteric GP prior for a $128 \times 128$ space with unit volume is shown in the following.fromnifty8importreasjftdims=(128,128)cfm=jft.CorrelatedFieldMaker("cf")cfm.set_amplitude_total_offset(offset_mean=2,offset_std=(1e-1,3e-2))cfm.add_fluctuations(# Axis over which the kernle is defineddims,distances=tuple(1.0/dfordindims),fluctuations=(1.0,5e-1),loglogavgslope=(-3.0,2e-1),flexibility=(1e0,2e-1),asperity=(5e-1,5e-2),prefix="ax1",non_parametric_kind="power",)correlated_field=cfm.finalize()# forward model for a GP priorNot all problems are well described by regularly spaced pixels.
For more complicated pixel spacings, NIFTy features Iterative Charted Refinement, a GP model for arbitrarily deformed spaces.
This model exploits nearest neighbor relations on various coarsening of the discretized modeled space and runs very efficiently on GPUs.
For one dimensional problems with arbitrarily spaced pixel, NIFTy also implements multiple flavors of Gauss-Markov processes.Building Up Complex ModelsModels are rarely just a GP prior.
Commonly, a model contains at least several non-linearities that transform the GP prior or combine it with other random variables.
For building more complex models, NIFTy provides aModelclass that offers a somewhat familiar object-oriented design yet is fully JAX compatible and functional under the hood.
The following code showcases such a model that builds up a slightly more involved model using the objects from the previous example.fromjaximportnumpyasjnpclassForward(jft.Model):def__init__(self,correlated_field):self._cf=correlated_field# Track a method with which a random input for the model. This is not# strictly required but is usually handy when building deep models.super().__init__(init=correlated_field.init)def__call__(self,x):# NOTE, any kind of masking of the output, non-linear and linear# transformation could be carried out here. Models can also combined and# nested in any way and form.returnjnp.exp(self._cf(x))forward=Forward(correlated_field)data=jnp.load("data.npy")lh=jft.Poissonian(data).amend(forward)All GP models in NIFTy as well as all likelihoods are models and their attributes are exposed to JAX, meaning JAX understands what it means if a computation involvesselfor other models.
In other words,correlated_field,forward, andlhfrom the code snippets shown here are all so-called pytrees in JAX and, e.g., the following is valid codejax.jit(lambda l, x: l(x))(lh, x0)withx0some arbitrarily chosen valid input tolh.
Inspired byequinox, individual attributes of the class can be marked as non-static or static viadataclass.field(metadata=dict(static=...))for the purpose of compiling.
Depending on the value, JAX will either treat the attribute as unknown placeholder or as known concrete attribute and potentially inline it during compiles.
This mechanism is extensively used in likelihoods to avoid inlining large constants such as the data and avoiding expensive re-compiles whenever possible.Variational InferenceNIFTy is built for models with millions to billions of degrees of freedom.
To probe the posterior efficiently and accurately, NIFTy relies on VI.
At the core of the VI methods lie an alternating procedure in which we switch between optimizing the Kullback–Leibler divergence for a specific shape of the variational posterior and updating the shape of the variational posterior.A typical minimization with NIFTy is shown in the following.
It retrieves six independent, antithetically mirrored samples from the approximate posterior via 25 iterations of alternating between optimization and sample adaption.
The final result is stored in thesamplesvariable.
A convenient one-shot wrapper for the below isjft.optimize_kl.
By virtue of all modeling tools in NIFTy being written in JAX, it is also possible to combine NIFTy tools withblackjaxor any other posterior sampler in the JAX ecosystem.fromjaximportrandomkey=random.PRNGKey(42)key,sk=random.split(key,2)# NIFTy is agnostic w.r.t. the type of input it gets as long as it supports core# arithmetic properties. Tell NIFTy to treat our parameter dictionary as a# vector.samples=jft.Samples(pos=jft.Vector(lh.init(sk)),samples=None,keys=None)delta=1e-4absdelta=delta*jft.size(samples.pos)opt_vi=jft.OptimizeVI(lh,n_total_iterations=25)opt_vi_st=opt_vi.init_state(key,# Typically on the order of 2-12n_samples=lambdai:1ifi<2else(2ifi<4else6),# Arguments for the conjugate gradient method used to drawing samples from# an implicit covariance matrixdraw_linear_kwargs=dict(cg_name="SL",cg_kwargs=dict(absdelta=absdelta/10.0,maxiter=100)),# Arguements for the minimizer in the nonlinear updating of the samplesnonlinearly_update_kwargs=dict(minimize_kwargs=dict(name="SN",xtol=delta,cg_kwargs=dict(name=None),maxiter=5)),# Arguments for the minimizer of the KL-divergence cost potentialkl_kwargs=dict(minimize_kwargs=dict(name="M",xtol=delta,maxiter=35)),sample_mode=lambdai:"nonlinear_resample"ifi<3else"nonlinear_update",)foriinrange(opt_vi.n_total_iterations):print(f"Iteration{i+1:04d}")# Continuously updates the samples of the approximate posterior distributionsamples,opt_vi_st=opt_vi.update(samples,opt_vi_st)print(opt_vi.get_status_message(samples,opt_vi_st))InstallationIf you only want to use NIFTy in your projects, but not change its source code, the easiest way to install NIFTy is via pip:pip install --user 'nifty8[re]'The line above installs the optional JAX backend termed NIFTy.re in addition to the numpy-based NIFTy.If you might want to adapt the NIFTy source code, we suggest installing NIFTy as editable python package using the following commands:git clone -b NIFTy_8 https://gitlab.mpcdf.mpg.de/ift/nifty.git
cd nifty
pip install --user --editable .Building the DocumentationNIFTy's documentation is generated viaSphinxand is available online atift.pages.mpcdf.de/nifty.To build the documentation locally, run:sudo apt-get install dvipng jupyter-nbconvert texlive-latex-base texlive-latex-extra
pip install --user sphinx jupytext pydata-sphinx-theme myst-parser
cd <nifty_directory>
bash docs/generate.shTo view the documentation, opendocs/build/index.htmlin your browser.Note: Make sure that you reinstall nifty after each change since sphinx imports nifty from the Python path.Run the testsTo run the tests, additional packages are required:sudo apt-get install python3-pytest-covAfterwards the tests (including a coverage report) can be run using the following command in the repository root:pytest --cov=nifty8 testFirst StepsFor a quick start, you can browse through theinformal introductionor dive into NIFTy by running one of the demonstrations, e.g.:python demos/0_intro.pyLicensing termsMost of NIFTy is licensed under the terms of theGPLv3license with NIFTy.re being a notable exception.
NIFTy.re is licensed under GPL-2.0+ OR BSD-2-Clause.
All of NIFTy is distributedwithout any warranty.Citing NIFTyTo cite the probabilistic programming framework NIFTy, please use the citation provided below.
In addition to citing NIFTy itself, please consider crediting the Gaussian process models you used and the inference machinery.
Seethe corresponding entry on citing NIFTy in the documentationfor further details.@misc{niftyre,
    author       = {{Edenhofer}, Gordian and {Frank}, Philipp and {Roth}, Jakob and {Leike}, Reimar H. and {Guerdi}, Massin and {Scheel-Platz}, Lukas I. and {Guardiani}, Matteo and {Eberle}, Vincent and {Westerkamp}, Margret and {Enßlin}, Torsten A.},
    title        = {{Re-Envisioning Numerical Information Field Theory (NIFTy.re): A Library for Gaussian Processes and Variational Inference}},
    keywords     = {Software},
    year         = 2024,
    howpublished = {in preparation}
}