# This program is free software: you can redistribute it and/or modify # it under the terms of the GNU General Public License as published by # the Free Software Foundation, either version 3 of the License, or # (at your option) any later version. # # This program is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # GNU General Public License for more details. # # You should have received a copy of the GNU General Public License # along with this program. If not, see . # # Copyright(C) 2013-2021 Max-Planck-Society # Authors: Reimar Leike, Philipp Arras, Philipp Frank # # NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik. ############################################################################### # Variational Inference (VI) # # This script demonstrates how MGVI and GeoVI work for an inference problem # with only two real quantities of interest. This enables us to plot the # posterior probability density as two-dimensional plot. The approximate # posterior samples are contrasted with the maximum-a-posterior (MAP) solution # together with samples drawn with the Laplace method. This method uses the # local curvature at the MAP solution as inverse covariance of a Gaussian # probability density. ############################################################################### import numpy as np import pylab as plt from matplotlib.colors import LogNorm import nifty7 as ift def main(): dom = ift.UnstructuredDomain(1) scale = 10 a = ift.FieldAdapter(dom, 'a') b = ift.FieldAdapter(dom, 'b') lh = (a.adjoint @ a).scale(scale) + (b.adjoint @ b).scale(-1.35*2).exp() lh = ift.VariableCovarianceGaussianEnergy(dom, 'a', 'b', np.float64) @ lh icsamp = ift.AbsDeltaEnergyController(deltaE=0.1, iteration_limit=2) ham = ift.StandardHamiltonian(lh, icsamp) x_limits = [-8/scale, 8/scale] x_limits_scaled = [-8, 8] y_limits = [-4, 4] x = np.linspace(*x_limits, num=401) y = np.linspace(*y_limits, num=401) xx, yy = np.meshgrid(x, y, indexing='ij') def np_ham(x, y): prior = x**2 + y**2 mean = x*scale lcov = 1.35*2*y lh = .5*(mean**2*np.exp(-lcov) + lcov) return lh + prior z = np.exp(-1*np_ham(xx, yy)) plt.plot(y, np.sum(z, axis=0)) plt.xlabel('y') plt.ylabel('unnormalized pdf') plt.title('Marginal density') plt.pause(2.0) plt.close() plt.plot(x*scale, np.sum(z, axis=1)) plt.xlabel('x') plt.ylabel('unnormalized pdf') plt.title('Marginal density') plt.pause(2.0) plt.close() mapx = xx[z==np.max(z)] mapy = yy[z==np.max(z)] meanx = (xx*z).sum()/z.sum() meany = (yy*z).sum()/z.sum() n_samples = 100 minimizer = ift.NewtonCG( ift.GradientNormController(iteration_limit=2, name='Mini')) IC = ift.StochasticAbsDeltaEnergyController(0.1, iteration_limit=20, name='advi') stochastic_minimizer_mf = ift.ADVIOptimizer(IC, eta = 0.5) stochastic_minimizer_fc = ift.ADVIOptimizer(IC, eta = 0.5) posmg = posgeo = posmf = posfc = ift.from_random(ham.domain, 'normal') fc = ift.FullCovarianceVI(posfc, ham, 10, False, initial_sig=0.01) mf = ift.MeanFieldVI(posmf, ham, 10, False, initial_sig=0.01) fig, axs = plt.subplots(2, 2, figsize=[12, 8]) axs = axs.flatten() def update_plot(runs): for axx, (nn, kl, pp, sam) in zip(axs,runs): axx.clear() axx.imshow(z.T, origin='lower', norm=LogNorm(vmin=1e-3, vmax=np.max(z)), cmap='gist_earth_r', extent=x_limits_scaled + y_limits) xs, ys = [], [] if sam: samples = (samp + pp for samp in kl.samples) else: samples = (kl.draw_sample() for _ in range(n_samples)) mx, my = 0., 0. for samp in samples: a = samp.val['a'] xs.append(a) mx += a b = samp.val['b'] ys.append(b) my += b mx /= n_samples my /= n_samples axx.scatter(np.array(xs)*scale, np.array(ys), label=f'{nn} samples') axx.scatter(mx*scale, my, label=f'{nn} mean') axx.scatter(mapx*scale, mapy, label = 'MAP') axx.scatter(meanx*scale, meany, label = 'Posterior mean') axx.set_title(nn) axx.set_xlim(x_limits_scaled) axx.set_ylim(y_limits) axx.legend(loc='lower right') axs[0].xaxis.set_visible(False) axs[1].xaxis.set_visible(False) axs[1].yaxis.set_visible(False) axs[2].set_xlabel('x') axs[2].set_ylabel('y') axs[3].yaxis.set_visible(False) axs[3].set_xlabel('x') plt.tight_layout() plt.draw() plt.pause(2.0) for ii in range(15): if ii % 2 == 0: # Resample mgkl = ift.MetricGaussianKL(posmg, ham, n_samples, False) mini_samp = ift.NewtonCG(ift.AbsDeltaEnergyController(1E-8, iteration_limit=5)) geokl = ift.GeoMetricKL(posgeo, ham, n_samples, mini_samp, False) runs = (("MGVI", mgkl, posmg, True), ("GeoVI", geokl, posgeo, True), ("MeanfieldVI", mf, posmf, False), ("FullCovarianceVI", fc, posfc, False)) update_plot(runs) mgkl, _ = minimizer(mgkl) geokl, _ = minimizer(geokl) mf.minimize(stochastic_minimizer_mf) fc.minimize(stochastic_minimizer_fc) posmg = mgkl.position posgeo = geokl.position posmf = mf.mean posfc = fc.mean runs = (("MGVI", mgkl, posmg, True), ("GeoVI", geokl, posgeo, True), ("MeanfieldVI", mf, posmf, False), ("FullCovarianceVI", fc, posfc, False)) update_plot(runs) ift.logger.info('Finished') # Uncomment the following line in order to leave the plots open # plt.show() if __name__ == '__main__': main()