
<!DOCTYPE html>

<html>
  <head>
    <meta charset="utf-8" />
    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
    <title>Using the SimulatorIMNN &#8212; Information Maximising Neural Network 0.3dev documentation</title>
    <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />
    <link rel="stylesheet" href="../../../../_static/alabaster.css" type="text/css" />
    <link rel="stylesheet" type="text/css" href="../../../../_static/autoclasstoc.css" />
    <script id="documentation_options" data-url_root="../../../../" src="../../../../_static/documentation_options.js"></script>
    <script src="../../../../_static/jquery.js"></script>
    <script src="../../../../_static/underscore.js"></script>
    <script src="../../../../_static/doctools.js"></script>
    <script async="async" src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.7/latest.js?config=TeX-AMS-MML_HTMLorMML"></script>
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" />
    <link rel="next" title="Using the GradientIMNN" href="../GradientIMNN/GradientIMNN.html" />
    <link rel="prev" title="Using specific subclasses" href="../../subclasses.html" />
   
  <link rel="stylesheet" href="../../../../_static/custom.css" type="text/css" />
  
  
  <meta name="viewport" content="width=device-width, initial-scale=0.9, maximum-scale=0.9" />

  </head><body>
  

    <div class="document">
      <div class="documentwrapper">
        <div class="bodywrapper">
          

          <div class="body" role="main">
            
  <div class="section" id="using-the-simulatorimnn">
<h1>Using the SimulatorIMNN<a class="headerlink" href="#using-the-simulatorimnn" title="Permalink to this headline">¶</a></h1>
<p>With a generative model built in <code class="docutils literal notranslate"><span class="pre">jax</span></code> we can generate simulations
on-the-fly to fit an IMNN. This is the best way to prevent against
overfitting due to limited variability in datasets which may introduce
spurious features which are not informative about the model parameters.</p>
<p>For this example we are going to summaries the unknown mean,
<span class="math notranslate nohighlight">\(\mu\)</span>, and variance, <span class="math notranslate nohighlight">\(\Sigma\)</span>, of <span class="math notranslate nohighlight">\(n_{\bf d}=10\)</span> data
points of two 1D random Gaussian field,
<span class="math notranslate nohighlight">\({\bf d}=\{d_i\sim\mathcal{N}(\mu,\Sigma)|i\in[1, n_{\bf d}]\}\)</span>.
This is an interesting problem since we know the likelihood
analytically, but it is non-Gaussian</p>
<div class="math notranslate nohighlight">
\[\mathcal{L}({\bf d}|\mu,\Sigma) = \prod_i^{n_{\bf d}}\frac{1}{\sqrt{2\pi|\Sigma|}}\exp\left[-\frac{1}{2}\frac{(d_i-\mu)^2}{\Sigma}\right]\]</div>
<p>As well as knowing the likelihood for this problem, we also know what
sufficient statistics describe the mean and variance of the data - they
are the mean and the variance</p>
<div class="math notranslate nohighlight">
\[\frac{1}{n_{\bf d}}\sum_i^{n_{\bf d}}d_i = \mu\textrm{  and  }\frac{1}{n_{\bf d}-1}\sum_i^{n_{\bf d}}(d_i-\mu)^2=\Sigma\]</div>
<p>What makes this an interesting problem for the IMNN is the fact that the
sufficient statistic for the variance is non-linear, i.e. it is a sum of
the square of the data, and so linear methods like MOPED would be lossy
in terms of information.</p>
<p>We can calculate the Fisher information by taking the negative second
derivative of the likelihood taking the expectation by inserting the
relations for the sufficient statistics, i.e. and examining at the
fiducial parameter values</p>
<div class="math notranslate nohighlight">
\[\begin{split}{\bf F}_{\alpha\beta} = -\left.\left(\begin{array}{cc}\displaystyle-\frac{n_{\bf d}}{\Sigma}&amp;0\\0&amp;\displaystyle-\frac{n_{\bf d}}{2\Sigma^2}\end{array}\right)\right|_{\Sigma=\Sigma^{\textrm{fid}}}.\end{split}\]</div>
<p>Choosing fiducial parameter values of <span class="math notranslate nohighlight">\(\mu^\textrm{fid}=0\)</span> and
<span class="math notranslate nohighlight">\(\Sigma^\textrm{fid}=1\)</span> we find that the determinant of the Fisher
information matrix is <span class="math notranslate nohighlight">\(|{\bf F}_{\alpha\beta}|=50\)</span>.</p>
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="kn">from</span> <span class="nn">imnn</span> <span class="kn">import</span> <span class="n">SimulatorIMNN</span>

<span class="kn">import</span> <span class="nn">jax</span>
<span class="kn">import</span> <span class="nn">jax.numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="kn">from</span> <span class="nn">jax.experimental</span> <span class="kn">import</span> <span class="n">stax</span><span class="p">,</span> <span class="n">optimizers</span>
</pre></div>
</div>
<p>We’re going to use 1000 summary vectors, with a length of two, at a time
to make an estimate of the covariance of network outputs and the
derivative of the mean of the network outputs with respect to the two
model parameters.</p>
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">n_s</span> <span class="o">=</span> <span class="mi">1000</span>
<span class="n">n_d</span> <span class="o">=</span> <span class="n">n_s</span>

<span class="n">n_params</span> <span class="o">=</span> <span class="mi">2</span>
<span class="n">n_summaries</span> <span class="o">=</span> <span class="n">n_params</span>

<span class="n">input_shape</span> <span class="o">=</span> <span class="p">(</span><span class="mi">10</span><span class="p">,)</span>
</pre></div>
</div>
<p>These are going to be generated on the fly using:</p>
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">simulator</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">θ</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">θ</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">+</span> <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">normal</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">shape</span><span class="o">=</span><span class="n">input_shape</span><span class="p">)</span> <span class="o">*</span> <span class="n">np</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span><span class="n">θ</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>
</pre></div>
</div>
<p>Our fiducial parameter values are <span class="math notranslate nohighlight">\(\mu^\textrm{fid}=0\)</span> and
<span class="math notranslate nohighlight">\(\Sigma^\textrm{fid}=1\)</span></p>
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">θ_fid</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([</span><span class="mf">0.</span><span class="p">,</span> <span class="mf">1.</span><span class="p">])</span>
</pre></div>
</div>
<p>For initialising the neural network and for generating the simulations
on the fly we need two random number generators:</p>
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">rng</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">PRNGKey</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
<span class="n">rng</span><span class="p">,</span> <span class="n">model_key</span><span class="p">,</span> <span class="n">fit_key</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">rng</span><span class="p">,</span> <span class="n">num</span><span class="o">=</span><span class="mi">3</span><span class="p">)</span>
</pre></div>
</div>
<p>We’re going to use <code class="docutils literal notranslate"><span class="pre">jax</span></code>’s stax module to build a simple network
with three hidden layers each with 128 neurons and which are activated
by leaky relu before outputting the two summaries. The optimiser will be
a <code class="docutils literal notranslate"><span class="pre">jax</span></code> Adam optimiser with a step size of 0.001.</p>
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">model</span> <span class="o">=</span> <span class="n">stax</span><span class="o">.</span><span class="n">serial</span><span class="p">(</span>
    <span class="n">stax</span><span class="o">.</span><span class="n">Dense</span><span class="p">(</span><span class="mi">128</span><span class="p">),</span>
    <span class="n">stax</span><span class="o">.</span><span class="n">LeakyRelu</span><span class="p">,</span>
    <span class="n">stax</span><span class="o">.</span><span class="n">Dense</span><span class="p">(</span><span class="mi">128</span><span class="p">),</span>
    <span class="n">stax</span><span class="o">.</span><span class="n">LeakyRelu</span><span class="p">,</span>
    <span class="n">stax</span><span class="o">.</span><span class="n">Dense</span><span class="p">(</span><span class="mi">128</span><span class="p">),</span>
    <span class="n">stax</span><span class="o">.</span><span class="n">LeakyRelu</span><span class="p">,</span>
    <span class="n">stax</span><span class="o">.</span><span class="n">Dense</span><span class="p">(</span><span class="n">n_summaries</span><span class="p">))</span>
<span class="n">optimiser</span> <span class="o">=</span> <span class="n">optimizers</span><span class="o">.</span><span class="n">adam</span><span class="p">(</span><span class="n">step_size</span><span class="o">=</span><span class="mf">1e-3</span><span class="p">)</span>
</pre></div>
</div>
<p>The SimulatorIMNN can now be initialised setting up the network and the
fitting routine (as well as the plotting function)</p>
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">imnn</span> <span class="o">=</span> <span class="n">SimulatorIMNN</span><span class="p">(</span>
    <span class="n">n_s</span><span class="o">=</span><span class="n">n_s</span><span class="p">,</span> <span class="n">n_d</span><span class="o">=</span><span class="n">n_d</span><span class="p">,</span> <span class="n">n_params</span><span class="o">=</span><span class="n">n_params</span><span class="p">,</span> <span class="n">n_summaries</span><span class="o">=</span><span class="n">n_summaries</span><span class="p">,</span>
    <span class="n">input_shape</span><span class="o">=</span><span class="n">input_shape</span><span class="p">,</span> <span class="n">θ_fid</span><span class="o">=</span><span class="n">θ_fid</span><span class="p">,</span> <span class="n">model</span><span class="o">=</span><span class="n">model</span><span class="p">,</span>
    <span class="n">optimiser</span><span class="o">=</span><span class="n">optimiser</span><span class="p">,</span> <span class="n">key_or_state</span><span class="o">=</span><span class="n">model_key</span><span class="p">,</span>
    <span class="n">simulator</span><span class="o">=</span><span class="n">simulator</span><span class="p">)</span>
</pre></div>
</div>
<p>To set the scale of the regularisation we use a coupling strength
<span class="math notranslate nohighlight">\(\lambda\)</span> whose value should mean that the determinant of the
difference between the covariance of network outputs and the identity
matrix is larger than the expected initial value of the determinant of
the Fisher information matrix from the network. How close to the
identity matrix the covariance should be is set by <span class="math notranslate nohighlight">\(\epsilon\)</span>.
These parameters should not be very important, but they will help with
convergence time.</p>
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">λ</span> <span class="o">=</span> <span class="mf">10.</span>
<span class="n">ϵ</span> <span class="o">=</span> <span class="mf">0.1</span>
</pre></div>
</div>
<p>Fitting can then be done simply by calling:</p>
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">imnn</span><span class="o">.</span><span class="n">fit</span><span class="p">(</span><span class="n">λ</span><span class="p">,</span> <span class="n">ϵ</span><span class="p">,</span> <span class="n">rng</span><span class="o">=</span><span class="n">fit_key</span><span class="p">,</span> <span class="n">print_rate</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">best</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
</pre></div>
</div>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span>0it [00:00, ?it/s]
</pre></div>
</div>
<p>Here we have included a <code class="docutils literal notranslate"><span class="pre">print_rate</span></code> for a progress bar, but leaving
this out will massively reduce fitting time (at the expense of not
knowing how many iterations have been run). The IMNN is run (by default)
for a maximum of <code class="docutils literal notranslate"><span class="pre">max_iterations</span> <span class="pre">=</span> <span class="pre">100000</span></code> iterations, but with early
stopping which can turn on after <code class="docutils literal notranslate"><span class="pre">min_iterations</span> <span class="pre">=</span> <span class="pre">100</span></code> iterations and
after <code class="docutils literal notranslate"><span class="pre">patience</span> <span class="pre">=</span> <span class="pre">100</span></code> iterations where the maximum determinant of the
Fisher information matrix has not increased. Note we have chosen to set
<code class="docutils literal notranslate"><span class="pre">imnn.w</span></code> are set to the values of the network parameters at the final
iteration via <code class="docutils literal notranslate"><span class="pre">best</span> <span class="pre">=</span> <span class="pre">False</span></code>, otherwise the network parameters which
obtained the highest value of the determinant of the Fisher information
matrix would be set.</p>
<p>To continue training one can simply generate a new random seed (very
important) and rerun fit</p>
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">rng</span><span class="p">,</span> <span class="n">another_fit_key</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">rng</span><span class="p">)</span>
<span class="n">imnn</span><span class="o">.</span><span class="n">fit</span><span class="p">(</span><span class="n">λ</span><span class="p">,</span> <span class="n">ϵ</span><span class="p">,</span> <span class="n">rng</span><span class="o">=</span><span class="n">another_fit_key</span><span class="p">,</span> <span class="n">print_rate</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">best</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
</pre></div>
</div>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span>0it [00:00, ?it/s]
</pre></div>
</div>
<p>To visualise the fitting history we can plot the results:</p>
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="n">imnn</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">expected_detF</span><span class="o">=</span><span class="mi">50</span><span class="p">);</span>
</pre></div>
</div>
<img alt="../../../../_images/output_21_0.png" src="../../../../_images/output_21_0.png" />
</div>


          </div>
          
        </div>
      </div>
      <div class="sphinxsidebar" role="navigation" aria-label="main navigation">
        <div class="sphinxsidebarwrapper">
<h1 class="logo"><a href="../../../../index.html">Information Maximising Neural Network</a></h1>








<h3>Navigation</h3>
<p class="caption"><span class="caption-text">Contents:</span></p>
<ul class="current">
<li class="toctree-l1"><a class="reference internal" href="../../../details.html">IMNN: Information maximising neural networks</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../install.html">Installing</a></li>
<li class="toctree-l1 current"><a class="reference internal" href="../../../examples.html">Examples</a><ul class="current">
<li class="toctree-l2 current"><a class="reference internal" href="../../subclasses.html">Using specific subclasses</a></li>
<li class="toctree-l2"><a class="reference internal" href="../../lfi.html">Using the likelihood free inference tools</a></li>
<li class="toctree-l2"><a class="reference internal" href="../../2d_field_inference/2d_field_inference.html">Inferring cosmology power spectrum parameters from 2D fields</a></li>
<li class="toctree-l2"><a class="reference internal" href="../../2d_field_inference/2d_field_inference.html#cosmological-parameter-inference-of-log-normal-fields">Cosmological parameter inference of log normal fields</a></li>
<li class="toctree-l2"><a class="reference internal" href="../../additive_noise/additive_noise.html">Constructing an IMNN with on-the-fly additive noise</a></li>
<li class="toctree-l2"><a class="reference internal" href="../../imnn_vs_regression/imnn_vs_regression.html">IMNN vs regression networks</a></li>
<li class="toctree-l2"><a class="reference internal" href="../../../examples.html#imnn">IMNN</a></li>
</ul>
</li>
<li class="toctree-l1"><a class="reference internal" href="../../../modules.html">Specific IMNN classes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../parents.html">Parent classes</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../tf.data.Datasets.html">Using tf.data.Datasets for loading</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../lfi.html">Likelihood-free inference modules</a></li>
</ul>

<div class="relations">
<h3>Related Topics</h3>
<ul>
  <li><a href="../../../../index.html">Documentation overview</a><ul>
  <li><a href="../../../examples.html">Examples</a><ul>
  <li><a href="../../subclasses.html">Using specific subclasses</a><ul>
      <li>Previous: <a href="../../subclasses.html" title="previous chapter">Using specific subclasses</a></li>
      <li>Next: <a href="../GradientIMNN/GradientIMNN.html" title="next chapter">Using the GradientIMNN</a></li>
  </ul></li>
  </ul></li>
  </ul></li>
</ul>
</div>
<div id="searchbox" style="display: none" role="search">
  <h3 id="searchlabel">Quick search</h3>
    <div class="searchformwrapper">
    <form class="search" action="../../../../search.html" method="get">
      <input type="text" name="q" aria-labelledby="searchlabel" />
      <input type="submit" value="Go" />
    </form>
    </div>
</div>
<script>$('#searchbox').show(0);</script>








        </div>
      </div>
      <div class="clearer"></div>
    </div>
    <div class="footer">
      &copy;2021, Tom Charnock.
      
      |
      Powered by <a href="http://sphinx-doc.org/">Sphinx 3.5.4</a>
      &amp; <a href="https://github.com/bitprophet/alabaster">Alabaster 0.7.12</a>
      
      |
      <a href="../../../../_sources/pages/examples/subclasses/SimulatorIMNN/SimulatorIMNN.rst.txt"
          rel="nofollow">Page source</a>
    </div>

    

    
  </body>
</html>