Package: imnn-tf
Optimising a neural network to maximise the Fisher information provides us with a function able to massively compress data without losing information about parameters of interest. This function can then be used for likelihood-free inference.The module here provides both the routines for fitting a neural network by maximising the Fisher information as well as a few methods for performing likelihood-free inference and approximate Bayesian computation.Specifically, the neural network takes some data,, and maps it to a compressed summary,, wherecan have the same dimensionality as that of the parameter space, rather than the data space, potentially without losing any information. To do so we maximise the Fisher information of the summary statistics provided by the neural network, and in doing so, find a functional form of the optimal compression.To train the neural network a batch of simulationscreated at a fiducial parameter valuefor training (and another for validation). These simulations are compressed by the neural network to obtain some statistic, i.e. the output of the neural network. We can use these to calculate the covariance,, of the compressed summaries. The sensitivity to model parameters uses the derivative of the simulation. This can be provided analytically or numercially usingcreated above the fiducial parameter valueandcreated below the fiducial parameter valueThe simulations are compressed using the network and used to find mean of the summariesIf the derivative of the simulations with respect to the parameters can be calculated analytically (or via autograd, etc.) then that can be used directly using the chain rule since the derivative of the network outputs with respect to the network input can be calculated easilyWe then useandto calculate the Fisher informationSince any linear rescaling of the summaries is also a summary, when maximising the Fisher information we set their scale usingwhereis a regularisation term whose strength is dictated bywithas a strength andas a rate parameter which can be determined from a closeness condition on the Frobenius norm of the difference between the covariance (and inverse covariance) from the identity matrix.When using this code please citeTom Charnock, Guilhem Lavaux and Benjamin D. WandeltAutomatic physical inference with information maximizing neural networks.Physical Review D. 97 083004 (2018)doi:10.1103/PhysRevD.97.083004arXiv:1802.03537The code in the paper can be downloaded as v1 or v1.1 of the code kept on zenodo:https://doi.org/10.5281/zenodo.1175196The code can be installed using:pip install imnn-tfor:git clone https://bitbucket.org/tomcharnock/imnn-tf.git
cd imnn-tf
python3 setup.py installModulesAvailable are modules for fitting the IMNN inIMNNand for doing likelihood-free inference inLFI. Examples of how to use these modules are available in theexamplesdirectory.IMNNThe basic call for the IMNN isimnn=IMNN.IMNN(n_s,n_d,n_params,n_summaries,θ_fid,δθ,input_shape,fiducial,derivative,validation_fiducial,validation_derivative,{model},{optimiser},{save},{load},{weights},{directory},{filename},{at_once},{map_fn},{check_shape},{verbose},{dtype},{itype})wheren_s- number of simulations to calculate covariance of fiducial simulationsn_d- number of simulations to calculate mean of the derivative simulationsn_params- number of parameters in the modeln_summaries- number of summaries to compress the data toθ_fid- fiducial parameter valuesδθ- parameter differences for numerical derivativeinput_shape- shape of a single simulationfiducial- a numpy array of fiducial simulations, generative function of individual fiducial simulations, or list of TFRecord filenamesderivative-  a numpy array of derivative simulations, generative function of individual derivative simulations, or list of TFRecord filenamesvalidation_fiducial- a numpy array of fiducial simulations, generative function of individual fiducial simulations, or list of TFRecord filenamesvalidation_derivative-  a numpy array of derivative simulations, generative function of individual derivative simulations, or list of TFRecord filenamesmodel- akeras-like model (optional if loading)optimiser- akeras-like optimiser (optional if loading)save- boolean describing whether to save or not (needdirectoryandfilenameifsave=True)load- boolean describing whether to load model or not (needdirectory,`filenameand optionallyweightsifload=True)weights- string with name of file of saved weightsdirectory- string with directory to load or save modelfilename- string with filename to load or save modelat_once- number of simulations to process with model at once (should ben_sif memory is large enough)map_fn- function to preprocess data (Noneif no preprocessing)check_shape- boolean describing whether to check shape of simulation on initialisationverbose- boolean to turn on and off descriptive write outdtype- TensorFlow float type (defaulttf.float32)itype- TensorFlow int type (defaulttf.int32)Fitimnn.fit({n_iterations},{λ},{ϵ},{reset},{patience},{min_iterations},{checkpoint},{tqdm_notebook},{weight_file})wheren_iterations- number of iterations to run the fitting for (can beNonewhen using patience)λ- strength of the regularisationϵ- distance of covariance (and inverse covariance) from the identityreset- boolean describing whether to reset weights and start training from scratchpatience- number of iterations of decreasing Fisher information of the validation set before stoppingmin_iterations- number of iterations before early stopping turns oncheckpoint- number of iterations between model saving (default turned off)tqdm_notebook-Trueif using a Jupyter notebook and False otherwise #TODO - make automatic (might already be implemented)weight_file- string with filename to save model weights toOnce trained, statistics are saved in a history dictionary attributeimnn.history"det_F"- determinant of the Fisher information of the training set"det_C"- determinant of the covariance matrix of the training set"det_Cinv"- determinant of the inverse covariance matrix of the training set"dμ_dθ"- derivative of the mean of the training set summaries"reg"- value of the regularisation"r"- value of the dynamic strength of the regularisation"val_det_F"- determinant of the Fisher information of the validation set"val_det_C"- determinant of the covariance matrix of the validation set"val_det_Cinv"- determinant of the inverse covariance matrix of the validation set"val_dμ_dθ"- derivative of the mean of the validation set summariesPlotimnn.plot({regulariser},{known_det_fisher},{figsize})whereregulariser- boolean describing whether to plot the regularisation historyknown_det_fisher- value of the determinant of the target Fisher information if already knownfigsize- tuple with the size of the figure if not defaultEstimate parametersGaussian estimates of the parameter values can be obtained from the network by runningimnn.get_estimate(input_data)where`input_datais data input to the network (shapeNone + input_shape). Note that if you want to make estimates without initialising the IMNN (once trained), the model can be loaded, along with the saved data during fit. For an IMNN saved withdirectory="model"andfilename=modelthen an estimator can be made usingestimator_parameters=np.load("model/model/estimator.npz")Finv=estimator_parameters["Finv"]θ_fid=estimator_parameters["θ_fid"]dμ_dθ=estimator_parameters["dμ_dθ"]Cinv=estimator_parameters["Cinv"]μ=estimator_parameters["μ"]@tf.function:defestimator(data):returntf.add(θ_fid,tf.einsum("ij,jk,kl,ml->mi",Finv,dμ_dθ,Cinv,model(data)-μ))ordefestimator(data):returnθ_fid+np.einsum("ij,jk,kl,ml->mi",Finv,dμ_dθ,Cinv,model(data)-μ)Training and validation data formatThe data must have the correct shape. For a single simulation with shapeinput_shapethen a fiducial data array must have a shape offiducial.shape=(n_s,)+input_shapeThe derivatives need to have a shape ofderivative.shape=(n_d,2,n_params)+input_shapewherederivative[:,0,...]is the lower part of the numerical derivative andderivative[:,1,...]is the upper part of the numerical derivative andderivative[:,:, i,...]labels theith parameter.If the data won’t fit in memory then we can load data via a generative functiondeffiducial_loader(seed):yieldfiducial[seed],seeddefderivative_loader(seed,derivative,parameter):yieldderivative[seed,derivative,parameter](seed,derivative,parameter)The function yields a single simulation at for each call labelled with the seed index (seedin range0ton_s) for the fiducial loader. The derivative loader yields a single simulation at a given seed, given upper or lower derivative and given parameter index (seedin range0ton_d,derivativein range0to1, andparameterin range0ton_params). In the above functions,fiducialandderivativeare some way of grabbing the data - it could be reading from file or from memory, etc. This has quite a bit of overhead and so it would be preferred to save the data as a TFRecord format. Instructions on how to do this for ingestion by the IMNN is available in theexamples/TFRecords.ipynbandexamples/IMNN - TFRecords.ipynbtutorials.Network model and optimiserThe IMNN is based onkeras-like network and optimisers, so an example could bemodel=tf.keras.Sequential([tf.keras.Input(shape=input_shape),tf.keras.layers.Dense(128),LeakyReLU(0.01),tf.keras.layers.Dense(128),LeakyReLU(0.01),tf.keras.layers.Dense(n_summaries),])opt=tf.keras.optimizers.Adam()Make sure to choose this network sensibly so that it best pulls the information from the data.LFIThe LFI module provides a Gaussian approximation to the posterior, a simple approximation Bayesian computation (ABC) implementation and a population Monte Carlo (PMC). These work with any estimator and not just with the IMNN.Gaussian approximation of the posteriorThe Gaussian approximation takes the inverse Fisher information as the variance of a Gaussian posterior (as implied by the Cramer-Rao bound) whose mean is at the estimate value.GA=LFI.GaussianApproximation(target_data,prior,Fisher,get_estimate,{labels})wheretarget_data- as many pieces of data to be inferred (target_data.shape = (None,) + input_shape)prior- the prior distribution which can be sampled from and whose probability can be evaluated with anevent_shapeof at least[1](suggested to use a TensorFlow Probability distribution)Fisher- Fisher information matrix (imnn.For otherwise for non-IMNN)get_estimate- function providing estimate of then_paramsmodel parameters from the data (imnn.get_estimateor otherwise for non-IMNN)labels- list of strings for labelling plotsPlot Fisher informationThe inverse Fisher information can plotted usingGA.plot_Fisher({figsize})Gaussian approximation to the likelihood and posteriorThe Gaussian approximation to the likelihood (prob) and the posterior (and their logarithms) can be obtained usingGA.log_prob({grid},{gridsize})GA.prob({grid},{gridsize})GA.log_posterior({grid},{gridsize})GA.posterior({grid},{gridsize})wheregrid- a set of parameters or an array of parameter or a meshgrid of parameter to evaluate the likelihood or posterior at (ifNonegridsize takes over)gridsize- a tuple of lengthn_paramswith the size of the meshgrid to make #TODO might crash ifGA.prior.low=-np.inffor any parameter orGA.prior.high=np.inffor any parameter. This defaults to20for every parameter ifgrid=Noneandgridsizeis not providedPlotting posteriorThe posterior can be plotted usingGA.plot({grid},{gridsize},**kwargs)where**kwargsare a variety ofmatplotlibarguments.Approximate Bayesian computation (ABC)The ABC draws parameter values from the prior and makes simulations at these points. These simulations are then summarised and then the distance between these estimates and the estimate of the target data can be calculated. Estimates within some small ϵ-ball around the target estimate are approximately samples from the posterior. Note that the larger the value of ϵ, the worse the approximation to the posterior.Note that a simulator of the data is needed. The simulator must be a functiondefsimulator(parameters,seed,simulator_args):returnsimulationwhereseedis a random number generator andsimulator_argsis a dict of arguments. Theseedandsimulator_argsare only for setting up the simulator - the function used in the ABC (and PMC) call must only take an array of parameters and return an array of simulations made at those parameter values. The function can call external codes, submit jobs on a cluster, etc. as long as the simulations are returned in the same order as the passed parameter array.The ABC can be initialised usingABC=LFI.ApproximateBayesianComputation(target_data,prior,Fisher,get_estimate,simulator,{labels})wheretarget_data- as many pieces of data to be inferred (target_data.shape = (None,) + input_shape)prior- the prior distribution which can be sampled from and whose probability can be evaluated with anevent_shapeof at least[1](suggested to use a TensorFlow Probability distribution)Fisher- Fisher information matrix (imnn.For otherwise for non-IMNN)get_estimate- function providing estimate of then_paramsmodel parameters from the data (imnn.get_estimateor otherwise for non-IMNN)simulator- function taking array of parameter values and returning simulations made at those valueslabels- list of strings for labelling plotsObtaining samplesThe ABC can be run usingABC(draws,{at_once},{save_sims})orABC.ABC(draws,{at_once},{save_sims},{PMC},{update})wheredraws- the number of simulations to make (or parameter values to make the simulations ifPMC=True)at_once- boolean describing whether to process (and make) all simulations at once or notsave_sims- string with the filename to save the sims (as a.npy) if providedPMC- boolean describing whetherdrawsis a number of simulations ordrawsis an array of parameter values to make simulations atupdate- boolean describing whether to update the ABC attributes onces the ABC is run or notOnce this is run the parameters, estimates, differences from the estimate and the target and the distance from the target are found asABC.parametersABC.estimatesABC.differencesABC.distancesAcception and rejection of samplesABConly runs the simulations and calculates the estimate distances but doesn’t do the accept and reject step within the ϵ-ball. This is done usingABC.accept_reject(ϵ)whereϵ`- a float describing the radius of the ϵ-ballOnce this is run more attributes are filledABC.num_accepted- number of accepted samplesABC.num_rejected- number of rejected samplesABC.num_draws- total number of samples doneABC.accepted_parametersABC.accepted_differencesABC.accepted_estimatesABC.accepted_distancesABC.rejected_parametersABC.rejected_differencesABC.rejected_estimatesABC.rejected_distancesAutomatic rejection samplerTo get a certain number of draws within a chosen ϵ-ball one can runABC.get_min_accepted(ϵ,accepted,{min_draws},{at_once},{save_sims},{tqdm_notebook})whereϵ- a float describing the radius of the ϵ-ballaccepted- the number of samples to be accepted within the ϵ-ballmin_draws- how many simulations to do at a time iteratively until enough simulations are acceptedat_once- boolean describing whether to process (and make) all simulations at once or notsave_sims- string with the filename to save the sims (as a.npy) if providedtqdm_notebook- True if using a Jupyter notebook and False otherwise #TODO - make automatic (might already be implemented)Histogrammed posteriorThe posterior is approximated by histogramming the accepted samples from the ABC (and acception/rejection) and can be calculated usingABC.posterior({bins},{ranges},{ϵ},{draws},{accepted},{at_once},{save_sims},{tqdm_notebook})wherebins- number of bins in the histogram defining the posteriorranges- minimum and maximum values for each parameter in the histogramOptionally any of the parameters forABC.ABC(...),ABC.accept_reject(...), and/orABC.get_min_accepted(...)can be passed toABC.posterior(...)to run the ABC when calling posterior rather than calling the sampling step first.Plot plosteriorThe posterior can be plotted usingABC.plot({smoothing},{bins},{ranges},{ϵ},{draws},{accepted},{at_once},{save_sims},{tqdm_notebook},**kwargs)wheresmoothing- the pixel range of a Gaussian smoothing of the histogram for plotting (smoothing causes inflation of the posterior)Optionally any of the parameters forABC.ABC(...),ABC.accept_reject(...), and/orABC.get_min_accepted(...)can be passed toABC.plot(...)to run the ABC when making the plot rather than calling the sampling step first.matplotlibparameters can also be passed for the plotting routine.Plot samplesThe samples can also be plotted usingABC.scatter_plot({axes},{rejected},{ϵ},{draws},{accepted},{at_once},{save_sims},{tqdm_notebook},**kwargs)whereaxes- either"parameter_estimate","parameter_parameter", or"estimate_estimate"for plotting the estimates against the parameters, or the parameters against the parameters or the estimates against the estimates (the last two are good for diagnostics such as the completeness of the sampling from the prior and the shape and correlation of the estimation function)rejected- a number between0and1describing the fraction of the rejected samples to plot (there are often orders of magnitude more samples rejected and so it makes sense to plot fewer, if they are to be plotted at all)Optionally any of the parameters forABC.ABC(...),ABC.accept_reject(...), and/orABC.get_min_accepted(...)can be passed toABC.scatter_plot(...)to run the ABC when making the plot rather than calling the sampling step first.matplotlibparameters can also be passed for the plotting routine.Running with saved simulationsIf simulations have already been run and we want to perform a simple ABC on them then we can set the simulator to return the saved simulations,saved_simulations->(?)+ input_shape, and pass the corresponding saved parameters,saved_parameters->(?,n_params)`, used to make the simulation.simulator=lambda_:saved_simulationsABC=LFI.ApproximateBayesianComputation(target_data,prior,Fisher,get_estimate,simulator,{labels})ABC(draws=saved_parameters,predrawn=True,save_sims=None,{at_once})ABC.accept_reject(ϵ)Population Monte Carlo (PMC)Whilst we can obtain approximate posteriors using ABC, the rejection rate is very high because we sample always from the prior. Population Monte Carlo (PMC) uses statistics of the population of samples to propose new parameter values, so each new simulation is more likely to be accepted. This prevents us needing to define an ϵ parameter to define the acceptance distance. Instead we start with a population from the prior and iteratively move samples inwards. Once it becomes difficult to move the population any more, i.e. the number of attempts to accept a parameter becomes very large, then the distribution is seen to be a stable approximation to the posterior.The whole module works very similarly toABCwith a few changes in arguments.PMC=LFI.PopulationMonteCarlo(target_data,prior,Fisher,get_estimate,simulator,{labels})wheretarget_data- as many pieces of data to be inferred (target_data.shape = (None,) + input_shape)prior- the prior distribution which can be sampled from and whose probability can be evaluated with anevent_shapeof at least[1](suggested to use a TensorFlow Probability distribution)Fisher- Fisher information matrix (imnn.For otherwise for non-IMNN)get_estimate- function providing estimate of then_paramsmodel parameters from the data (imnn.get_estimateor otherwise for non-IMNN)simulator- function taking array of parameter values and returning simulations made at those valueslabels- list of strings for labelling plotsObtaining accepted samplesThe PMC can be run by callingPMC(draws,initial_draws,criterion,{percentile},{at_once},{save_sims},{tqdm_notebook})orPMC.PMC(draws,initial_draws,criterion,{percentile},{at_once},{save_sims},{tqdm_notebook})wheredraws- number of samples from the posteriorinitial_draws- number of samples from the prior to start the PMC (must be equal to or greater than the number of draws from the posteriorcriterion- the stopping condition, the fraction of times samples are accepted in any one iteration of the PMC (when this is small then many samples are not accepted into the population, suggesting a stationary distribution)percentile- the percentage of points which are considered the in the main sample (making this small moves more samples at once, but with reduced statistics from the population, default set to 75%, it takes longer to run (but may be cheaper in number of simulations) if set to a high value orNone)at_once- boolean describing whether to process (and make) all simulations at once or notsave_sims- string with the filename to save the sims (as a.npy) if providedtqdm_notebook-Trueif using a Jupyter notebook and False otherwise #TODO - make automatic (might already be implemented)Histogrammed posteriorThe posterior is approximated by histogramming the accepted samples from the PMC and can be calculated usingPMC.posterior({bins},{ranges},{draws},{initial_draws},{criterion},{percentile},{at_once},{save_sims},{tqdm_notebook})wherebins- number of bins in the histogram defining the posteriorranges- minimum and maximum values for each parameter in the histogramOptionally any of the parameters forPMC.PMC(...)can be passed toPMC.posterior(...)`to run the PMC when calling posterior rather than calling the sampling step first.Plot posteriorThe posterior can be plotted usingPMC.plot({smoothing},{bins},{ranges},{draws},{initial_draws},{criterion},{percentile},{at_once},{save_sims},{tqdm_notebook},**kwargs)wheresmoothing- the pixel range of a Gaussian smoothing of the histogram for plotting (smoothing causes inflation of the posterior)Optionally any of the parameters forPMC.PMC(...)can be passed toPMC.plot(...)to run the PMC when making the plot rather than calling the sampling step first.matplotlibparameters can also be passed for the plotting routine.Plot samplesThe samples can also be plotted usingPMC.scatter_plot({axes},{draws},{initial_draws},{criterion},{percentile},{at_once},{save_sims},{tqdm_notebook},**kwargs)whereaxes- either"parameter_estimate","parameter_parameter", or"estimate_estimate"for plotting the estimates against the parameters, or the parameters against the parameters or the estimates against the estimates (the last two are good for diagnostics such as the completeness of the sampling from the prior and the shape and correlation of the estimation function)Optionally any of the parameters forPMC.PMC(...)can be passed toPMC.scatter_plot(...)to run the PMC when making the plot rather than calling the sampling step first.matplotlibparameters can also be passed for the plotting routine.TODOThe module is under constant development, and progress can be checked in thedevbranch. Current additions to the IMNN includePut back summary support
- Previous versions of the IMNN had the ability to pass arbitrary summaries along with network summaries. This is useful because it can be a suggestion of how much information is gained over other summarising functions (such as the two point statistics, etc.)
- Need to accept array, generative function and TFRecords with summaries and split covariance between summaries and network outputs for regularisationJAX implementation of all routines
- This is under private development currentlyDocstrings written for LFIWrite unit tests