```python
%load_ext autoreload
%autoreload 2
```


```python
!python -m pip install --upgrade jax==0.2.10 jaxlib==0.1.62+cuda101 -f https://storage.googleapis.com/jax-releases/jax_releases.html
```


```python
cd /mnt/home/tmakinen/repositories/scratch/gitinstall/imnn
```


```python
!git checkout dev;
```

    D	examples/2D sims - IMNN.ipynb
    M	examples/IMNN - JAX.ipynb
    Already on 'dev'
    Your branch is ahead of 'origin/dev' by 3 commits.
      (use "git push" to publish your local commits)



```python
!python -m pip install .
```

    Already up-to-date.



```python
pwd
```




    '/mnt/home/tmakinen/repositories/scratch/dev-imnn/imnn'




```python
!python -m setup.py install
```

    /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/setuptools/dist.py:458: UserWarning: Normalizing '0.3dev' to '0.3.dev0'
      warnings.warn(tmpl.format(**locals()))
    running install
    running bdist_egg
    running egg_info
    writing IMNN.egg-info/PKG-INFO
    writing dependency_links to IMNN.egg-info/dependency_links.txt
    writing requirements to IMNN.egg-info/requires.txt
    writing top-level names to IMNN.egg-info/top_level.txt
    reading manifest file 'IMNN.egg-info/SOURCES.txt'
    writing manifest file 'IMNN.egg-info/SOURCES.txt'
    installing library code to build/bdist.linux-x86_64/egg
    running install_lib
    running build_py
    copying IMNN/experimental/jax/imnn/imnn.py -> build/lib/IMNN/experimental/jax/imnn
    creating build/bdist.linux-x86_64/egg
    creating build/bdist.linux-x86_64/egg/IMNN
    creating build/bdist.linux-x86_64/egg/IMNN/LFI
    copying build/lib/IMNN/LFI/LFI.py -> build/bdist.linux-x86_64/egg/IMNN/LFI
    copying build/lib/IMNN/LFI/__init__.py -> build/bdist.linux-x86_64/egg/IMNN/LFI
    copying build/lib/IMNN/IMNN.py -> build/bdist.linux-x86_64/egg/IMNN
    copying build/lib/IMNN/tfIMNN.py -> build/bdist.linux-x86_64/egg/IMNN
    copying build/lib/IMNN/__init__.py -> build/bdist.linux-x86_64/egg/IMNN
    creating build/bdist.linux-x86_64/egg/IMNN/utils
    copying build/lib/IMNN/utils/__init__.py -> build/bdist.linux-x86_64/egg/IMNN/utils
    copying build/lib/IMNN/utils/utils.py -> build/bdist.linux-x86_64/egg/IMNN/utils
    copying build/lib/IMNN/utils/TFRecords.py -> build/bdist.linux-x86_64/egg/IMNN/utils
    creating build/bdist.linux-x86_64/egg/IMNN/experimental
    copying build/lib/IMNN/experimental/__init__.py -> build/bdist.linux-x86_64/egg/IMNN/experimental
    creating build/bdist.linux-x86_64/egg/IMNN/experimental/jax
    creating build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    copying build/lib/IMNN/experimental/jax/imnn/_tensorflow_imnn.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    copying build/lib/IMNN/experimental/jax/imnn/tensorflow_gradient_imnn.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    copying build/lib/IMNN/experimental/jax/imnn/aggregated_gradient_imnn_test.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    copying build/lib/IMNN/experimental/jax/imnn/gradient_imnn_test.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    copying build/lib/IMNN/experimental/jax/imnn/numerical_gradient_imnn_test.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    copying build/lib/IMNN/experimental/jax/imnn/_imnn.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    copying build/lib/IMNN/experimental/jax/imnn/aggregated_numerical_gradient_imnn_test.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    copying build/lib/IMNN/experimental/jax/imnn/imnn.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    copying build/lib/IMNN/experimental/jax/imnn/aggregated_numerical_gradient_imnn.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    copying build/lib/IMNN/experimental/jax/imnn/gradient_imnn.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    copying build/lib/IMNN/experimental/jax/imnn/__init__.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    copying build/lib/IMNN/experimental/jax/imnn/_imnn_test.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    copying build/lib/IMNN/experimental/jax/imnn/simulator_imnn_test.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    copying build/lib/IMNN/experimental/jax/imnn/aggregated_simulator_imnn.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    copying build/lib/IMNN/experimental/jax/imnn/tensorflow_numerical_gradient_imnn.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    copying build/lib/IMNN/experimental/jax/imnn/simulator_imnn.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    copying build/lib/IMNN/experimental/jax/imnn/_tensorflow_imnn_test.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    copying build/lib/IMNN/experimental/jax/imnn/aggregated_gradient_imnn.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    copying build/lib/IMNN/experimental/jax/imnn/numerical_gradient_imnn.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    copying build/lib/IMNN/experimental/jax/imnn/aggregated_simulator_imnn_test.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn
    creating build/bdist.linux-x86_64/egg/IMNN/experimental/jax/lfi
    copying build/lib/IMNN/experimental/jax/lfi/gaussian_approximation.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/lfi
    copying build/lib/IMNN/experimental/jax/lfi/lfi.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/lfi
    copying build/lib/IMNN/experimental/jax/lfi/__init__.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/lfi
    copying build/lib/IMNN/experimental/jax/lfi/approximate_bayesian_computation.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/lfi
    copying build/lib/IMNN/experimental/jax/__init__.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax
    creating build/bdist.linux-x86_64/egg/IMNN/experimental/jax/utils
    copying build/lib/IMNN/experimental/jax/utils/__init__.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/utils
    copying build/lib/IMNN/experimental/jax/utils/utils.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/utils
    copying build/lib/IMNN/experimental/jax/utils/keras_utils.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/utils
    copying build/lib/IMNN/experimental/jax/utils/container.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/utils
    copying build/lib/IMNN/experimental/jax/utils/jac.py -> build/bdist.linux-x86_64/egg/IMNN/experimental/jax/utils
    creating build/bdist.linux-x86_64/egg/test
    copying build/lib/test/tensorflow_defaults.py -> build/bdist.linux-x86_64/egg/test
    copying build/lib/test/simulator_defaults.py -> build/bdist.linux-x86_64/egg/test
    copying build/lib/test/numerical_gradient_defaults.py -> build/bdist.linux-x86_64/egg/test
    copying build/lib/test/__init__.py -> build/bdist.linux-x86_64/egg/test
    copying build/lib/test/gradient_defaults.py -> build/bdist.linux-x86_64/egg/test
    copying build/lib/test/defaults.py -> build/bdist.linux-x86_64/egg/test
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/LFI/LFI.py to LFI.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/LFI/__init__.py to __init__.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/IMNN.py to IMNN.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/tfIMNN.py to tfIMNN.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/__init__.py to __init__.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/utils/__init__.py to __init__.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/utils/utils.py to utils.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/utils/TFRecords.py to TFRecords.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/__init__.py to __init__.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn/_tensorflow_imnn.py to _tensorflow_imnn.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn/tensorflow_gradient_imnn.py to tensorflow_gradient_imnn.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn/aggregated_gradient_imnn_test.py to aggregated_gradient_imnn_test.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn/gradient_imnn_test.py to gradient_imnn_test.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn/numerical_gradient_imnn_test.py to numerical_gradient_imnn_test.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn/_imnn.py to _imnn.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn/aggregated_numerical_gradient_imnn_test.py to aggregated_numerical_gradient_imnn_test.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn/imnn.py to imnn.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn/aggregated_numerical_gradient_imnn.py to aggregated_numerical_gradient_imnn.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn/gradient_imnn.py to gradient_imnn.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn/__init__.py to __init__.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn/_imnn_test.py to _imnn_test.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn/simulator_imnn_test.py to simulator_imnn_test.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn/aggregated_simulator_imnn.py to aggregated_simulator_imnn.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn/tensorflow_numerical_gradient_imnn.py to tensorflow_numerical_gradient_imnn.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn/simulator_imnn.py to simulator_imnn.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn/_tensorflow_imnn_test.py to _tensorflow_imnn_test.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn/aggregated_gradient_imnn.py to aggregated_gradient_imnn.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn/numerical_gradient_imnn.py to numerical_gradient_imnn.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/imnn/aggregated_simulator_imnn_test.py to aggregated_simulator_imnn_test.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/lfi/gaussian_approximation.py to gaussian_approximation.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/lfi/lfi.py to lfi.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/lfi/__init__.py to __init__.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/lfi/approximate_bayesian_computation.py to approximate_bayesian_computation.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/__init__.py to __init__.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/utils/__init__.py to __init__.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/utils/utils.py to utils.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/utils/keras_utils.py to keras_utils.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/utils/container.py to container.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/IMNN/experimental/jax/utils/jac.py to jac.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/test/tensorflow_defaults.py to tensorflow_defaults.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/test/simulator_defaults.py to simulator_defaults.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/test/numerical_gradient_defaults.py to numerical_gradient_defaults.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/test/__init__.py to __init__.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/test/gradient_defaults.py to gradient_defaults.cpython-36.pyc
    byte-compiling build/bdist.linux-x86_64/egg/test/defaults.py to defaults.cpython-36.pyc
    creating build/bdist.linux-x86_64/egg/EGG-INFO
    copying IMNN.egg-info/PKG-INFO -> build/bdist.linux-x86_64/egg/EGG-INFO
    copying IMNN.egg-info/SOURCES.txt -> build/bdist.linux-x86_64/egg/EGG-INFO
    copying IMNN.egg-info/dependency_links.txt -> build/bdist.linux-x86_64/egg/EGG-INFO
    copying IMNN.egg-info/requires.txt -> build/bdist.linux-x86_64/egg/EGG-INFO
    copying IMNN.egg-info/top_level.txt -> build/bdist.linux-x86_64/egg/EGG-INFO
    zip_safe flag not set; analyzing archive contents...
    creating 'dist/IMNN-0.3.dev0-py3.6.egg' and adding 'build/bdist.linux-x86_64/egg' to it
    removing 'build/bdist.linux-x86_64/egg' (and everything under it)
    Processing IMNN-0.3.dev0-py3.6.egg
    Removing /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/IMNN-0.3.dev0-py3.6.egg
    Copying IMNN-0.3.dev0-py3.6.egg to /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages
    IMNN 0.3.dev0 is already the active version in easy-install.pth
    
    Installed /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/IMNN-0.3.dev0-py3.6.egg
    Processing dependencies for IMNN==0.3.dev0
    Searching for matplotlib==3.3.2
    Best match: matplotlib 3.3.2
    Adding matplotlib 3.3.2 to easy-install.pth file
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages
    Searching for scipy==1.5.3
    Best match: scipy 1.5.3
    Adding scipy 1.5.3 to easy-install.pth file
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages
    Searching for numpy==1.19.4
    Best match: numpy 1.19.4
    Adding numpy 1.19.4 to easy-install.pth file
    Installing f2py script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    Installing f2py3 script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    Installing f2py3.6 script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages
    Searching for tqdm==4.55.1
    Best match: tqdm 4.55.1
    Processing tqdm-4.55.1-py3.6.egg
    tqdm 4.55.1 is already the active version in easy-install.pth
    Installing tqdm script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/tqdm-4.55.1-py3.6.egg
    Searching for jax==0.2.8
    Best match: jax 0.2.8
    Processing jax-0.2.8-py3.6.egg
    jax 0.2.8 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/jax-0.2.8-py3.6.egg
    Searching for jaxlib==0.1.57+cuda101
    Best match: jaxlib 0.1.57+cuda101
    Removing jaxlib 0.1.59 from easy-install.pth file
    Adding jaxlib 0.1.57+cuda101 to easy-install.pth file
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages
    Searching for tensorflow==2.4.0
    Best match: tensorflow 2.4.0
    Processing tensorflow-2.4.0-py3.6-linux-x86_64.egg
    tensorflow 2.4.0 is already the active version in easy-install.pth
    Installing estimator_ckpt_converter script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    Installing import_pb_to_tensorboard script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    Installing saved_model_cli script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    Installing tensorboard script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    Installing tf_upgrade_v2 script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    Installing tflite_convert script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    Installing toco script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    Installing toco_from_protos script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/tensorflow-2.4.0-py3.6-linux-x86_64.egg
    Searching for certifi==2020.12.5
    Best match: certifi 2020.12.5
    Adding certifi 2020.12.5 to easy-install.pth file
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages
    Searching for Pillow==8.1.0
    Best match: Pillow 8.1.0
    Adding Pillow 8.1.0 to easy-install.pth file
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages
    Searching for python-dateutil==2.8.1
    Best match: python-dateutil 2.8.1
    Adding python-dateutil 2.8.1 to easy-install.pth file
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages
    Searching for kiwisolver==1.3.1
    Best match: kiwisolver 1.3.1
    Adding kiwisolver 1.3.1 to easy-install.pth file
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages
    Searching for cycler==0.10.0
    Best match: cycler 0.10.0
    Adding cycler 0.10.0 to easy-install.pth file
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages
    Searching for pyparsing==2.4.7
    Best match: pyparsing 2.4.7
    Adding pyparsing 2.4.7 to easy-install.pth file
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages
    Searching for opt-einsum==3.3.0
    Best match: opt-einsum 3.3.0
    Adding opt-einsum 3.3.0 to easy-install.pth file
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages
    Searching for absl-py==0.11.0
    Best match: absl-py 0.11.0
    Adding absl-py 0.11.0 to easy-install.pth file
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages
    Searching for flatbuffers==1.12
    Best match: flatbuffers 1.12
    Adding flatbuffers 1.12 to easy-install.pth file
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages
    Searching for wrapt==1.12.1
    Best match: wrapt 1.12.1
    Processing wrapt-1.12.1-py3.6-linux-x86_64.egg
    wrapt 1.12.1 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/wrapt-1.12.1-py3.6-linux-x86_64.egg
    Searching for wheel==0.36.2
    Best match: wheel 0.36.2
    Adding wheel 0.36.2 to easy-install.pth file
    Installing wheel script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages
    Searching for typing-extensions==3.7.4.3
    Best match: typing-extensions 3.7.4.3
    Processing typing_extensions-3.7.4.3-py3.6.egg
    typing-extensions 3.7.4.3 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/typing_extensions-3.7.4.3-py3.6.egg
    Searching for termcolor==1.1.0
    Best match: termcolor 1.1.0
    Processing termcolor-1.1.0-py3.6.egg
    termcolor 1.1.0 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/termcolor-1.1.0-py3.6.egg
    Searching for tensorflow-estimator==2.4.0
    Best match: tensorflow-estimator 2.4.0
    Processing tensorflow_estimator-2.4.0-py3.6.egg
    tensorflow-estimator 2.4.0 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/tensorflow_estimator-2.4.0-py3.6.egg
    Searching for tensorboard==2.4.0
    Best match: tensorboard 2.4.0
    Processing tensorboard-2.4.0-py3.6.egg
    tensorboard 2.4.0 is already the active version in easy-install.pth
    Installing tensorboard script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/tensorboard-2.4.0-py3.6.egg
    Searching for six==1.15.0
    Best match: six 1.15.0
    Adding six 1.15.0 to easy-install.pth file
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages
    Searching for protobuf==4.0.0rc2
    Best match: protobuf 4.0.0rc2
    Processing protobuf-4.0.0rc2-py3.6.egg
    protobuf 4.0.0rc2 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/protobuf-4.0.0rc2-py3.6.egg
    Searching for Keras-Preprocessing==1.1.2
    Best match: Keras-Preprocessing 1.1.2
    Processing Keras_Preprocessing-1.1.2-py3.6.egg
    Keras-Preprocessing 1.1.2 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/Keras_Preprocessing-1.1.2-py3.6.egg
    Searching for h5py==2.10.0
    Best match: h5py 2.10.0
    Processing h5py-2.10.0-py3.6-linux-x86_64.egg
    h5py 2.10.0 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/h5py-2.10.0-py3.6-linux-x86_64.egg
    Searching for grpcio==1.32.0
    Best match: grpcio 1.32.0
    Processing grpcio-1.32.0-py3.6-linux-x86_64.egg
    grpcio 1.32.0 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/grpcio-1.32.0-py3.6-linux-x86_64.egg
    Searching for google-pasta==0.2.0
    Best match: google-pasta 0.2.0
    Processing google_pasta-0.2.0-py3.6.egg
    google-pasta 0.2.0 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/google_pasta-0.2.0-py3.6.egg
    Searching for gast==0.3.3
    Best match: gast 0.3.3
    Processing gast-0.3.3-py3.6.egg
    gast 0.3.3 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/gast-0.3.3-py3.6.egg
    Searching for astunparse==1.6.3
    Best match: astunparse 1.6.3
    Processing astunparse-1.6.3-py3.6.egg
    astunparse 1.6.3 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/astunparse-1.6.3-py3.6.egg
    Searching for Werkzeug==1.0.1
    Best match: Werkzeug 1.0.1
    Processing Werkzeug-1.0.1-py3.6.egg
    Werkzeug 1.0.1 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/Werkzeug-1.0.1-py3.6.egg
    Searching for tensorboard-plugin-wit==1.7.0
    Best match: tensorboard-plugin-wit 1.7.0
    Processing tensorboard_plugin_wit-1.7.0-py3.6.egg
    tensorboard-plugin-wit 1.7.0 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/tensorboard_plugin_wit-1.7.0-py3.6.egg
    Searching for setuptools==49.6.0.post20201009
    Best match: setuptools 49.6.0.post20201009
    Adding setuptools 49.6.0.post20201009 to easy-install.pth file
    Installing easy_install script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages
    Searching for requests==2.25.1
    Best match: requests 2.25.1
    Processing requests-2.25.1-py3.6.egg
    requests 2.25.1 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/requests-2.25.1-py3.6.egg
    Searching for Markdown==3.3.3
    Best match: Markdown 3.3.3
    Processing Markdown-3.3.3-py3.6.egg
    Markdown 3.3.3 is already the active version in easy-install.pth
    Installing markdown_py script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/Markdown-3.3.3-py3.6.egg
    Searching for google-auth==1.24.0
    Best match: google-auth 1.24.0
    Processing google_auth-1.24.0-py3.6.egg
    google-auth 1.24.0 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/google_auth-1.24.0-py3.6.egg
    Searching for google-auth-oauthlib==0.4.2
    Best match: google-auth-oauthlib 0.4.2
    Processing google_auth_oauthlib-0.4.2-py3.6.egg
    google-auth-oauthlib 0.4.2 is already the active version in easy-install.pth
    Installing google-oauthlib-tool script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/google_auth_oauthlib-0.4.2-py3.6.egg
    Searching for urllib3==1.26.2
    Best match: urllib3 1.26.2
    Processing urllib3-1.26.2-py3.6.egg
    urllib3 1.26.2 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/urllib3-1.26.2-py3.6.egg
    Searching for idna==2.10
    Best match: idna 2.10
    Processing idna-2.10-py3.6.egg
    idna 2.10 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/idna-2.10-py3.6.egg
    Searching for chardet==4.0.0
    Best match: chardet 4.0.0
    Processing chardet-4.0.0-py3.6.egg
    chardet 4.0.0 is already the active version in easy-install.pth
    Installing chardetect script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/chardet-4.0.0-py3.6.egg
    Searching for importlib-metadata==3.3.0
    Best match: importlib-metadata 3.3.0
    Processing importlib_metadata-3.3.0-py3.6.egg
    importlib-metadata 3.3.0 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/importlib_metadata-3.3.0-py3.6.egg
    Searching for rsa==4.6
    Best match: rsa 4.6
    Processing rsa-4.6-py3.6.egg
    rsa 4.6 is already the active version in easy-install.pth
    Installing pyrsa-decrypt script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    Installing pyrsa-encrypt script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    Installing pyrsa-keygen script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    Installing pyrsa-priv2pub script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    Installing pyrsa-sign script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    Installing pyrsa-verify script to /mnt/home/tmakinen/anaconda3/envs/pyjax/bin
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/rsa-4.6-py3.6.egg
    Searching for pyasn1-modules==0.2.8
    Best match: pyasn1-modules 0.2.8
    Processing pyasn1_modules-0.2.8-py3.6.egg
    pyasn1-modules 0.2.8 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/pyasn1_modules-0.2.8-py3.6.egg
    Searching for cachetools==4.2.0
    Best match: cachetools 4.2.0
    Processing cachetools-4.2.0-py3.6.egg
    cachetools 4.2.0 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/cachetools-4.2.0-py3.6.egg
    Searching for requests-oauthlib==1.3.0
    Best match: requests-oauthlib 1.3.0
    Processing requests_oauthlib-1.3.0-py3.6.egg
    requests-oauthlib 1.3.0 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/requests_oauthlib-1.3.0-py3.6.egg
    Searching for zipp==3.4.0
    Best match: zipp 3.4.0
    Processing zipp-3.4.0-py3.6.egg
    zipp 3.4.0 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/zipp-3.4.0-py3.6.egg
    Searching for pyasn1==0.4.8
    Best match: pyasn1 0.4.8
    Processing pyasn1-0.4.8-py3.6.egg
    pyasn1 0.4.8 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/pyasn1-0.4.8-py3.6.egg
    Searching for oauthlib==3.1.0
    Best match: oauthlib 3.1.0
    Processing oauthlib-3.1.0-py3.6.egg
    oauthlib 3.1.0 is already the active version in easy-install.pth
    
    Using /mnt/home/tmakinen/anaconda3/envs/pyjax/lib/python3.6/site-packages/oauthlib-3.1.0-py3.6.egg
    Finished processing dependencies for IMNN==0.3.dev0
    /mnt/home/tmakinen/anaconda3/envs/pyjax/bin/python: Error while finding module specification for 'setup.py' (AttributeError: module 'setup' has no attribute '__path__')



```python
!XLA_FLAGS=--xla_gpu_cuda_data_dir=/cm/shared/sw/pkg/devel/cuda/10.1.243_418.87.00
from jax.config import config
config.update("jax_debug_nans", True)
config.update('jax_enable_x64', True)

```


```python
import itertools
import math
from functools import partial

import numpy as onp

import jax
print("jax version: ", jax.__version__)
import jax.experimental.optimizers as optimizers
import jax.experimental.stax as stax
import jax.numpy as np
from jax import jit
import matplotlib.pyplot as plt
import scipy

import imnn
print("IMNN version: ", imnn.__version__)

from imnn.imnn import (
    AggregatedGradientIMNN,
    AggregatedNumericalGradientIMNN,
    AggregatedSimulatorIMNN,
    GradientIMNN,
    NumericalGradientIMNN,
    SimulatorIMNN,
)
from imnn.lfi import (
    ApproximateBayesianComputation,
    GaussianApproximation,
)
from imnn.utils import value_and_jacrev, value_and_jacfwd

rng = jax.random.PRNGKey(0)

```

    jax version:  0.2.10
    IMNN version:  0.3dev



```python
!XLA_FLAGS=--xla_cpu_enable_fast_math=false
```


```python
from jax.lib import xla_bridge
print(xla_bridge.get_backend().platform)
```

    gpu


# Model in STAX


```python
n_summaries = 2
n_s = 5000
n_d = 5000
λ = 100.0
ϵ = 0.1
```


```python
# define inception block layer

def InceptBlock2(filters, strides, do_5x5=True, do_3x3=True):
    """InceptNet convolutional striding block.
    filters: tuple: (f1,f2,f3)
    filters1: for conv1x1
    filters2: for conv1x1,conv3x3
    filters3L for conv1x1,conv5x5"""
    
    filters1, filters2, filters3 = filters
    
    conv1x1 = stax.serial(stax.Conv(filters1, (1,1), strides, padding="SAME"))
    
    filters4 = filters2
    conv3x3 = stax.serial(stax.Conv(filters2, (1,1), strides=None, padding="SAME"),
                        stax.Conv(filters4, (3,3), strides, padding="SAME"))
                        
    filters5 = filters3
    conv5x5 = stax.serial(stax.Conv(filters3, (1,1), strides=None, padding="SAME"),
                         stax.Conv(filters5, (5,5), strides, padding="SAME")) 
    
    maxpool = stax.serial(stax.MaxPool((3,3), padding="SAME"),
                         stax.Conv(filters4, (1,1), strides, padding="SAME"))
                            
    if do_3x3:
        if do_5x5:
            return stax.serial(
                  stax.FanOut(4),  # should num=3 or 2 here ?
                  stax.parallel(conv1x1, conv3x3, conv5x5, maxpool),
                  stax.FanInConcat(), 
                  stax.LeakyRelu)
        else:
            return stax.serial(
                  stax.FanOut(3),  # should num=3 or 2 here ?
                  stax.parallel(conv1x1, conv3x3, maxpool),
                  stax.FanInConcat(), 
                  stax.LeakyRelu)
            
    else:
        return stax.serial(
              stax.FanOut(2),  # should num=3 or 2 here ?
              stax.parallel(conv1x1, maxpool),
              stax.FanInConcat(), 
              stax.LeakyRelu)
    
    
def Reshape(newshape):
    """Layer function for a reshape layer."""
    init_fun = lambda rng, input_shape: (newshape,())
    apply_fun = lambda params, inputs, **kwargs: np.reshape(inputs,newshape)
    return init_fun, apply_fun    

```


```python
rng,drop_rng = jax.random.split(rng)
fs = 64 #for 128x128 sims
               

model = stax.serial(
        InceptBlock2((fs,fs,fs), strides=(4,4)),
        InceptBlock2((fs,fs,fs), strides=(4,4)),
        InceptBlock2((fs,fs,fs), strides=(4,4)),
        InceptBlock2((fs,fs,fs), strides=(2,2), do_5x5=False, do_3x3=False),
        stax.Conv(n_summaries, (1,1), strides=(1,1), padding="SAME"),
        stax.Flatten,
        Reshape((n_summaries,))
)
```


```python
optimiser = optimizers.adam(step_size=1e-3)
```

# Random seeds for IMNN


```python
rng, initial_model_key = jax.random.split(rng)
rng, fitting_key = jax.random.split(rng)
```

# Random seeds for ABC


```python
rng, abc_key = jax.random.split(rng)
```

# 2D Gaussian Field Simulator in JAX

Steps to creating $(N \times N)$ 2D Gaussian field for IMNN:
1. Generate a $(N\times N)$ white noise field $\varphi$ such that $\langle \varphi_k \varphi_{-k} \rangle' = 1$
2. Fourier Transform $\varphi$ to real space: $R_{\rm white}(\textbf{x}) \rightarrow R_{\rm white}(\textbf{k})$
    - note: NumPy's DFT Fourier convention is:
    $$\phi_{ab}^{\textbf{k}} = \sum_{c,d = 0}^{N-1} \exp{(-i x_c k_a - i x_d k_b) \phi^{\textbf{x}}_{cd}}$$
    $$\phi_{ab}^{\textbf{x}} = \frac{1}{N^2}\sum_{c,d = 0}^{N-1} \exp{(-i x_c k_a - i x_d k_b) \phi^{\textbf{k}}_{cd}}$$    
3. Scale white noise $R_{\rm white}(\textbf{k})$ by the chosen power spectrum evaluated over a field of $k$ values:
$$ R_P(\textbf{k}) = P^{1/2}(k) R_{\rm white}(\textbf{k}) $$
    - note: here we need to ensure that this array of amplitudes is Hermitian, e.g. $\phi^{* \textbf{k}}_{a(N/2 + b)} = \phi^{\textbf{k}}_{a(N/2 - b)}$. This is accomplished by choosing indexes $k_a = k_b = \frac{2\pi}{N} (0, \dots, N/2, -N/2+1, \dots, -1)$ and then evaluating the square root of the outer product of the meshgrid between the two: $k = \sqrt{k^2_a + k^2_b}$. We can then evaluate $P^{1/2}(k)$.
    
4. Fourier Transform $R_{P}(\textbf{k})$ to real space: $ R_P(\textbf{x}) = \int d^d \tilde{k} e^{i\textbf{k} \cdot \textbf{x}} R_p(\textbf{k}) $:

$$R_{ab}^{\textbf{x}} = \frac{1}{N^2}\sum_{c,d = 0}^{N-1} \exp{(-i x_c k_a - i x_d k_b) R^{\textbf{k}}_{cd}}$$


```python
# SET 32-BiT floats for model !
θ_fid = np.array([1.0, 0.5], dtype=np.float32)
δθ = np.array([0.1, 0.1], dtype=np.float32)
n_params = 2
N = 128
dim = 2
L = 128

field_shape = (N,N)
input_shape = (1,1, N,N)
simulator_args = {"N": N, "L": L, "dim": dim, "shape": field_shape, 'vol_norm': False, "N_scale": True, "squeeze": False}
```

# simulator class for a Powerbox in Jax
Attributes:
- simulator for Gaussian fields from input power spectrum
- analytic Fisher information computation at a given $\theta$


```python
rng,fg_key = jax.random.split(rng)
foregrounds = jax.random.normal(fg_key, (1000, 1,) + simulator_args['shape'])*0


def default_P(k, A, B):
    return A*k**-B


class powerBoxJax:
    def __init__(self, shape, pk=None, k=None):
        
        if pk is None:
            self.pk = default_P
        else:
            self.pk = pk
            
        if k is None:
            self.k = np.sqrt(np.sum(np.array(np.meshgrid(*(
                (np.hstack((np.arange(0, _shape//2 + 1), 
                    np.arange(-_shape//2 + 1, 0))) * 2*np.pi / _shape)**2
                    for _shape in shape))), axis=0))
        else:
            self.k = k
            
        self.shape = shape
        self.N = shape[0]
        
    def simulator(self, rng, θ, simulator_args=simulator_args, add_foregrounds=False):
        def P(k, A=1, B=1):
            return self.pk(k, A, B)

        def fn(key, A, B):
            shape = self.shape #simulator_args["shape"]

            k = self.k

            new_shape = ()
            for _shape in shape:
                if _shape % 2 == 0:
                    new_shape += (_shape+1,)
                else:
                    new_shape += (_shape,)


            key1,key2 = jax.random.split(key)
    
            if add_foregrounds:
                foreground = foregrounds[jax.random.randint(key2, 
                                            minval=0, maxval=1000, shape=())]
            else:
                foreground = 0.

            # L is in length units, like Gpc
            L = simulator_args['L']
            dim = simulator_args['dim']

            if np.isscalar(L):
                L = [L]*int(dim)

            else:
                L = np.array(L)


            V = np.prod(np.array(L))

            scale = V**(1./dim)

            Lk = ()

            _N = 1
            for i,_shape in enumerate(shape):
                _N *= _shape
                Lk += (_shape / L[i],) # 1 / dx


            fft_norm = np.prod(np.array(Lk))

            _dims = len(shape)
            tpl = ()
            for _d in range(_dims):
                tpl += (_d,)

            # POWERBOX IMPLEMENTATION

            mag = jax.random.normal(key1, shape=tuple(N for N in new_shape))
            
            # random phases
            pha = 2 * np.pi * jax.random.uniform(key1, shape=tuple(N for N in new_shape))
            
            # now make hermitian field (reality condition)
            revidx = (slice(None, None, -1),) * len(mag.shape)
            mag = (mag + mag[revidx]) / np.sqrt(2) 
            pha = (pha - pha[revidx]) / 2 + np.pi
            dk = mag * (np.cos(pha) + 1j * np.sin(pha)) # output is complex

            cutidx = (slice(None, -1),) * len(new_shape)

            dk = dk[cutidx]
            powers = np.concatenate((np.zeros(1), 
                        np.sqrt(P(k.flatten()[1:], A=A, B=B)))).reshape(k.shape)

            # normalize power by volume
            if simulator_args['vol_norm']:
                powers = powers/V

            fourier_field = powers * dk

            fourier_field = jax.ops.index_update(
                fourier_field,
                np.zeros(len(shape), dtype=int),
                np.zeros((1,)))
            
            field = np.real(np.fft.ifftn(fourier_field) * fft_norm * V)
            
            if simulator_args["N_scale"]:
                field *= scale

            field = np.expand_dims(field + foreground, (0,))


            if not simulator_args["squeeze"]:
                field = np.expand_dims(field, (0,))

            return np.array(field, dtype='float32')

        shape = self.shape
        A, B = θ


        if A.shape == B.shape:
            if len(A.shape) == 0:
                return fn(rng, A, B)
            else:
                keys = jax.random.split(rng, num=A.shape[0] + 1)
                rng = keys[0]
                keys = keys[1:]
                return jax.vmap(
                    lambda key, A, B: simulator(key, (A, B), simulator_args=simulator_args)
                )(keys, A, B)
        else:
            if len(A.shape) > 0:
                keys = jax.random.split(rng, num=A.shape[0] + 1)
                rng = keys[0]
                keys = keys[1:]
                return jax.vmap(
                    lambda key, A: simulator(key, (A, B), simulator_args=simulator_args)
                )(keys, A)
            elif len(B.shape) > 0:
                keys = jax.random.split(rng, num=B.shape[0])
                return jax.vmap(
                    lambda key, B: simulator(key, (A, B), simulator_args=simulator_args)
                )(keys, B)



    def AnalyticFisher(self, 
               θ, 
               kvec=None,
               N=None
              ):
        """
        Code for computing the Analytic Fisher for a Gaussian
        Field with power spectrum P(k) = Ak^-B
        """
        A,B = θ
        if N is None:
            N = self.N
        
        # we want all UNIQUE fourier modes
        if kvec is None:
            kvec = self.k[1:N//2, 1:N//2]

        pk = lambda k : A*(k**-B)             # P(k) = Ak^(-B)
        p_a = lambda k : k**-B                # deriv w.r.t. A
        p_b = lambda k : -A*(k**-B)*np.log(k) # deriv w.r.t. B


        powers = (pk(kvec.flatten()[:]))#np.concatenate((np.ones(1), 
            

        powera =  (p_a(kvec.flatten()[:])) #np.concatenate((np.zeros(1), 
           

        powerb = (p_b(kvec.flatten()[:])) #np.concatenate((np.zeros(1), 
            


        Cinv = np.diag(1. / powers)          # diagonal inv. covariance
        Ca =  np.diag(powera / 1.)           # C_{,A}
        Cb =  np.diag(powerb / 1.)           # C_{,B}

        Faa = 0.5 * np.trace((Ca @ Cinv @ Ca @ Cinv))
        Fab = 0.5 * np.trace((Ca @ Cinv @ Cb @ Cinv))
        Fba = 0.5 * np.trace((Cb @ Cinv @ Ca @ Cinv))
        Fbb = 0.5 * np.trace((Cb @ Cinv @ Cb @ Cinv))

        return np.array([[Faa, Fab], [Fba, Fbb]])


```

# analytic likelihood class
- return the log-likelihood (computed on a grid) for power spectrum parameters for Gaussian Fields
- takes in a PowerboxJax simulator object `PBJ` and fourier-transformed field target data `Δ_target`
methods:
- `get_likelihood`: return likelihood marginal on a grid
- `plot_contours`: plot likelihood contours on a meshgrid
- `plot_corner`: plot liklihood contours, target crosshairs, and image data on a 2x2 corner plot


```python
class analyticFieldLikelihood:

    def __init__(self,
                 PBJ, 
                 field_shape,
                 Δ_target,
                 prior,
                 k=None,
                 pk=None,
                 gridsize=20, 
                 tiling=2):

        """code for computing a gaussian field's likelihood for power spectrum parameters
            PBJ :         powerBox simulator object
            field_shape : list. shape of field input
            Δ :           array-like. FFT of the real-space field
            prior :       array-like. range over which to compute the likelihood
            k :           array-like. fourier modes over which to compute P(k)
            tiling :      list or int. tiling=2 means likelihood will be computed as 2x2 grid
            gridsize :    how large to make the likelihood surface
        """
        
        
        if k is None:
            self.k = PBJ.k

        if pk is None:
            self.pk = PBJ.pk

        self.field_shape = field_shape
        self.gridsize = gridsize
        if np.isscalar(tiling):
            self.tiling = [tiling]*2
        else:
            self.tiling = tiling    
        #self.tilesize = gridsize // tiling
        self.N = np.sqrt(np.prod(np.array(field_shape)))  # should just be N for NxN grid
        self.prior = prior

        self.k = k
        self.Δ = Δ_target

    def Pk(self, k, A=1, B=0.5):
        return self.pk(k, A, B)


        return np.diag(pk)
    def log_likelihood(self, k, A, B, Δ):
        Δ = Δ.flatten()[:]
        k = k

        dlength = len(k.flatten())

        def fn(_A, _B):

            nrm = np.pad(np.ones(dlength-2)*2, (1,1), constant_values=1.)
            nrm = jax.ops.index_update(
              nrm, np.array([[0],[(dlength-2)]]), np.array([[1],[1]]))
            nrm = 1

            powers = self.Pk(k.flatten()[:], A=_A, B=_B)

            # covariance is P(k)
            C = powers * nrm
            invC = 1./self.Pk(k.flatten()[:], A=_A, B=_B)

            logdetC = np.sum(np.log(C))
            pi2 = np.pi * 2.
            m_half_size = -0.5 * len(Δ)
            exponent = - 0.5 * np.sum(np.conj(Δ) * invC * Δ)
            norm = -0.5 * logdetC + m_half_size*np.log(pi2) 
            return (exponent + norm)
        return jax.vmap(fn)(A, B)

    def get_likelihood(self, return_grid=False, shift=None):
        A_start = self.prior[0][0]
        A_end = self.prior[1][0]
        B_start = self.prior[0][1]
        B_end = self.prior[1][1]

        region_size = [self.gridsize // self.tiling[i] for i in range(len(self.tiling))]

        print("computing likelihood on a %dx%d grid \n \
        in tiles of size %dx%d"%(self.gridsize, self.gridsize, region_size[0], region_size[1]))


        def get_like_region(A0, A1, B0, B1, qsize):
            A_range = np.linspace(A0, A1, qsize)
            B_range = np.linspace(B0, B1, qsize)
            A, B = np.meshgrid(A_range, B_range)

            return (self.log_likelihood(self.k,
                A.ravel(), B.ravel(), self.Δ).reshape(qsize,qsize))



        A_incr = (A_end - A_start) / self.tiling[0]
        B_incr = (B_end - B_start) / self.tiling[1]

        # marks the ends of linspace
        A_starts = [A_start + (i)*A_incr for i in range(self.tiling[0])]
        A_ends = [A_start + (i+1)*A_incr for i in range(self.tiling[0])]
        B_starts = [B_start + (i)*B_incr for i in range(self.tiling[1])]
        B_ends = [B_start + (i+1)*B_incr for i in range(self.tiling[1])]

        _like_cols = []
        for _col in range(self.tiling[0]):
            # slide horizontally in A
            _like_row = []
            for _row in range(self.tiling[1]):
                # slide vertically in B

                _like = get_like_region(A_starts[_row], 
                                        A_ends[_row], 
                                        B_starts[_col], 
                                        B_ends[_col],
                                        region_size[0],
                                        )
                _like_row.append(_like)

            _like_cols.append(np.concatenate(_like_row, axis=1))

        _log_likelihood = np.real(np.concatenate(_like_cols, axis=0))

        if shift is None:
            shift = np.max(_log_likelihood)

        #print('shift', shift)
        print('loglike mean', np.mean(_log_likelihood))

        #_log_likelihood = _log_likelihood - shift

        if return_grid:
            _A_range = np.linspace(self.prior[0,0], self.prior[1,0], self.gridsize)
            _B_range = np.linspace(self.prior[0,0], self.prior[1,0], self.gridsize)
            return (_log_likelihood), _A_range, _B_range

        return (_log_likelihood)



    def plot_contours(self, ax=None, 
                      θ_ref=None, shift=None, 
                      xlabel='A', ylabel='B', 
                      return_like=True):

        _like, _A, _B = self.get_likelihood(return_grid=True, shift=shift)
        _like = scipy.special.softmax(np.real(_like))

        _A, _B = np.meshgrid(_A, _B)

        if ax is None:
            fig,ax = plt.subplots(figsize=(10,10))

        mesh = ax.contourf(_A, _B, _like)
        plt.colorbar(mesh, ax=ax)

        if θ_ref is not None:
            ax.scatter(θ_ref[0], θ_ref[1], zorder=10, marker='+', s=100, color='r')

        ax.set_xlabel('A')
        ax.set_ylabel('B')

        if return_like:
            return _like, ax

        else:
            return ax
        

    def plot_corner(self, ax=None, θ_ref=None, label="Analytic likelihood", 
                    image_data=None, return_like=False):

        _like, _A_range, _B_range = self.get_likelihood(return_grid=True)

        likelihoodA = scipy.special.softmax(np.real(_like)).sum(0) #np.real(likelihood).sum(0)
        likelihoodA /= likelihoodA.sum() * (_A_range[1] -  _A_range[0])
        likelihoodB =  scipy.special.softmax(np.real(_like)).sum(1) #np.real(likelihood).sum(1)
        likelihoodB /= likelihoodB.sum() * (_B_range[1] -  _B_range[0])

        _like = scipy.special.softmax(np.real(_like))
        sorted_marginal = np.sort(_like.flatten())[::-1]
        cdf = np.cumsum(sorted_marginal / sorted_marginal.sum())
        value = []
        for level in [0.997, 0.95, 0.68]:
            this_value = sorted_marginal[np.argmin(np.abs(cdf - level))]
            if len(value) == 0:
                value.append(this_value)
            elif this_value <= value[-1]:
                break
            else:
                value.append(this_value)

        if ax is None:
            fig,ax = plt.subplots(nrows=2, ncols=2)

        ax[1,0].contour(_A_range, _B_range, _like, levels=value, colors='C2', alpha=0.7)
        ax[0,0].plot(_A_range, likelihoodA, color='C2', label=None, alpha=0.7)
        ax[1,1].plot(likelihoodB, _B_range, color='C2', label='loglike', alpha=0.7)
        
        if image_data is not None:
            ax[0,1].imshow(np.squeeze(image_data))
        else:
            ax[0,1].axis("off")

        if θ_ref is not None:
            ax[0,0].axvline(θ_ref[0], linestyle='--', c='k')
            ax[1,0].axvline(θ_ref[0], linestyle='--', c='k')
            ax[1,0].axhline(θ_ref[1], linestyle='--', c='k')
            ax[1,1].axhline(θ_ref[1], linestyle='--', c='k', label=r'$\theta_{\rm ref}$')

        ax[1,0].set_xlabel(r'$A$')
        ax[1,0].set_ylabel(r'$B$')
        
        if return_like:
            return ax,_like

        else:
            return ax
```


```python
PBJ = powerBoxJax(simulator_args['shape'])
simulator = PBJ.simulator
```

## sim and gradient


```python
def simulator_gradient(rng, θ, simulator_args=simulator_args):
    return value_and_jacrev(simulator, argnums=1, allow_int=True, holomorphic=True)(rng, θ, simulator_args=simulator_args)
```


```python
rng, key = jax.random.split(rng)
```


```python
field_shape
```




    (128, 128)




```python
# plot example simulation and derivative
deriv_args = {"N": N, "L": L, "dim": dim, "shape": field_shape, "vol_norm": True, "N_scale": True, "squeeze": False}
simulation, simulation_gradient = value_and_jacfwd(simulator, argnums=1)(rng, θ_fid, simulator_args=deriv_args)

plt.imshow(np.squeeze(simulation[0]), extent=(0,1,0,1))
plt.colorbar()
plt.title('example simulation')
plt.show()

plt.imshow(np.squeeze(simulation_gradient[0].T[0].T), extent=(0,1,0,1))
plt.title('gradient of simulation')
plt.colorbar()
plt.show()
```


```python
def get_simulations(rng, n_s, θ, simulator_args=simulator_args):
    def get_simulator(key):
        return simulator(key, θ, simulator_args=simulator_args)

    keys = jax.random.split(rng, num=n_s)
    return jax.vmap(get_simulator)(np.array(keys))


def get_simulation_gradients(rng, n_s, n_d, θ, simulator_args=simulator_args):
    def get_batch_gradient(key):
        return simulator_gradient(key, θ, simulator_args=simulator_args)

    keys = jax.random.split(rng, num=n_s)
    return jax.vmap(get_batch_gradient)(np.array(keys)[:n_d])
```

# known analytic Fisher information

For a gaussian field, the likelihood is written
$$     \mathcal{L}(\Delta | \theta) = \frac{1}{(2\pi)^{N_p / 2} \det(C)^{1/2}}\exp{\left(-\frac{1}{2} \Delta C^{-1} \Delta \right)}$$

Where $\Delta \in \mathbb{R}^{N_p},\ N_p=N_k=V=N\times N$ is the Fourier Transform of the observed real-space field.
This yields a Fisher information matrix of

$$F_{\alpha \beta} = \langle -\frac{\partial^2 \ln \mathcal{L}}{\partial \lambda_\alpha \partial \lambda_\beta} \rangle= \frac{1}{2} {\rm Tr} (C_{, \alpha} C^{-1} C_{, \beta} C^{-1}) $$
where the covariance is
$$ C_{k_i, k_j} =  P(k_i)\delta_{ij}$$
The associated derivatives for a power law $P(k) = Ak^{-B}$ are
$$\begin{align}
    C_{,A} &= \left( k^{-B} \right)\delta_{ij} \\
    C_{,B} &= \left( -Ak^{-B}\ln k \right) \delta_{ij}
\end{align} $$

We notice that the Fisher information is *only* a function of the power spectrum parameters. It tells us the curvature of the chosen model (likelihood function) at a given $\theta$. The analytic Fisher information is the maximum amount of information we can expect the IMNN to extract from our simulations.


<!-- Alternatively, we can explore a volume integral analytically from the definition of C :
where the Fisher matrix is given by
$$ F_{\alpha \beta} = \sum_k \frac{1}{(\delta P_k)^2} \frac{\partial P_k}{\partial \lambda_\alpha} \frac{\partial P_k}{\partial \lambda_\beta}$$
and the error on $P_k$ is given (for a square, 2D box) as
$$ \delta P_k = \sqrt{\frac{2}{k (\Delta k) V} } \left( P_k + \frac{1}{\bar{n}} \right) $$ -->

<!-- For a gaussian field, the likelihood is written
$$ \ln \mathcal{L}(\theta | \vec{d}) = \ln \mathcal{L}(\theta | \Delta) = \sqrt{\frac{1}{2\pi C}} \exp{\frac{-{\Delta}^2}{2C}}$$

where $\vec{d} = \Delta$ is the overdensity field (in a cosmological context this is the measured temperature or galaxy count in a sky survey). Given that the power spectrum describes the correlations at different scales $k$, we can define the correlation via the power spectrum $C = P(k) = Ak^{-B}$ to compute the log-likelihood. The Fisher information matrix, given as 
$$ F_{\alpha \beta} = \langle - \frac{\partial^2 \ln \mathcal{L}}{\partial \theta_\alpha \partial \theta_\beta} \rangle $$
can then be computed analytically for our likelihood:
$$ F_{\alpha \beta} = \sum_k \frac{1}{(\delta P_k)^2} \frac{\partial P_k}{\partial \theta_\alpha} \frac{\partial P_k}{\partial \theta_\beta} $$

where $\delta P_k = \sqrt{\frac{2}{4\pi \Delta k V k_{\rm tot}^2}} (P_k + \frac{1}{\bar{n}})$ is the error on $P_k$ with survey volume $V$, sampling interval $\Delta k$, and shot noise $1/\bar{n}$. Using the fact that $d\ln P_k = \frac{d P_k}{P_k}$, we can rewrite the sum as an integral:
$$ F_{\alpha \beta} = 2 \pi \left( \frac{V_{\rm eff}}{\lambda^3} \right) \int_{k_{\rm min}}^{k_{\rm max}} d \ln k  \frac{\partial \ln P_k}{\partial \theta_\alpha} \frac{\partial \ln P_k}{\partial \theta_\beta}$$

Where $V_{\rm eff}$ and $\lambda^3$ are our effective windowed survey size and survey extent, respectively (set to 1 for now). Doing the integration explicitly, we obtain the Fisher matrix for parameters $\theta = (A, B)$:
$$ F = 2 \pi \left( \frac{V_{\rm eff}}{\lambda^3} \right) \begin{bmatrix}
    \frac{1}{A^2} \ln (\frac{k_{\rm max}}{k_{\rm min}}) & \frac{1}{2A} ((\ln k_{\rm min})^2 - (\ln k_{\rm max})^2)  \\
     \frac{1}{2A} ((\ln k_{\rm min})^2 - (\ln k_{\rm max})^2) & \frac{1}{3}((\ln k_{\rm max})^3 - (\ln k_{\rm min})^3)   \\
\end{bmatrix}$$
 -->


For our fiducial model with our data vector of size $128^2$, our $\rm det|F|$ reads:


```python
N = simulator_args["N"]
shape = simulator_args["shape"]

kbin = np.sqrt(np.sum(np.array(np.meshgrid(*(
    np.hstack((np.arange(0, _shape//2 + 1), 
        np.arange(-_shape//2 + 1, 0))) *2* np.pi / _shape
        for _shape in shape)))**2, axis=0))

f_expected = PBJ.AnalyticFisher(θ_fid, kvec=None)
print("analytic F(θ_fid): ", f_expected)

detf_expected = np.linalg.det(f_expected)
print("analytic det(F(θ_fid)): ", detf_expected)
```

    analytic F(θ_fid):  [[ 1984.5        -1545.06379524]
     [-1545.06379524  1690.49264978]]
    analytic det(F(θ_fid)):  967560.5321330648



```python
# MAKE SIMULATION
N = simulator_args["N"]
shape = (N,N)

θ_sim = np.array([0.7, 0.8])
simulator_args = {"N": N, "L": L, "dim": dim, "shape": shape, "vol_norm": True, "N_scale": False, "squeeze": True}
simulator_args["shape"] = (N,N)
simkey,rng = jax.random.split(rng)
#sim = np.squeeze(target_data)#
sim = np.squeeze(simulator(simkey, θ_sim, simulator_args=simulator_args))
sim_fft = (np.fft.fft2(sim)) #/ (N**2)
```


```python
# PLOT ANALYTIC POSTERIOR
# IGNORE FIRST FOURIER MODE -- no information here !
gridsize = 100                # for likelihood gridding
Δ = sim_fft[1:N//2, 1:N//2]
k = kbin[1:N//2, 1:N//2]  
prior_range = np.array([[0.1, 0.1], [1.25, 1.25]])
AL = analyticFieldLikelihood(PBJ, shape, Δ, prior_range, k=k, gridsize=gridsize, tiling=[5,5])
```


```python
#plt.style.use('default')
ax = AL.plot_corner(θ_ref=θ_sim, image_data=sim)
```


```python
simulator_args
```




    {'N': 128,
     'L': 128,
     'dim': 2,
     'shape': (128, 128),
     'vol_norm': True,
     'N_scale': False,
     'squeeze': True}




```python
def sanity_check(gridsize=50, num=20):
    likes = []
    likeAs = []
    likeBs = []
    rng1 = jax.random.PRNGKey(13)
    values = []
    θ_target = np.array([0.8, 0.8], dtype=np.float32)

    for t in range(num):
        key, rng1 = jax.random.split(rng1)
        targ = simulator(
            key,
            θ_target)

        gridsize = 50                # for likelihood gridding
        Δ = np.fft.fftn(np.squeeze(targ))[1:N//2, 1:N//2] / N
        k = kbin[1:N//2, 1:N//2]  
        prior_range = np.array([[0.1, 0.1], [1.25, 1.25]])
        AL = analyticFieldLikelihood(PBJ, shape, Δ, prior_range, k=k, 
                                     gridsize=gridsize, tiling=[5,5])

        likelihood,A_range,B_range = AL.get_likelihood(shift=None, return_grid=True)


        _A_range = A_range#*np.exp(shift)
        _B_range = B_range#*np.exp(shift)

        likelihoodA = scipy.special.softmax(np.real(likelihood)).sum(0) #np.real(likelihood).sum(0)
        likelihoodA /= likelihoodA.sum() * (_A_range[1] -  _A_range[0])
        likelihoodB =  scipy.special.softmax(np.real(likelihood)).sum(1) #np.real(likelihood).sum(1)
        likelihoodB /= likelihoodB.sum() * (_B_range[1] -  _B_range[0])

        likelihood = scipy.special.softmax(np.real(likelihood))
        sorted_marginal = np.sort(likelihood.flatten())[::-1]
        cdf = np.cumsum(sorted_marginal / sorted_marginal.sum())
        value = []
        for level in [0.997, 0.95, 0.68]:
            this_value = sorted_marginal[np.argmin(np.abs(cdf - level))]
            if len(value) == 0:
                value.append(this_value)
            elif this_value <= value[-1]:
                break
            else:
                value.append(this_value)
        #fig, ax = plt.subplots(2, 2, figsize=(10, 10))

        #likelihood /= likelihood.sum()

        likes.append(likelihood)
        likeAs.append(likelihoodA)
        likeBs.append(likelihoodB)

        values.append(value)
        
        return likes,likeAs,likeBs,values
```


```python
likes,likeAs,likeBs,values = sanity_check()
fig,ax = plt.subplots(nrows=2, ncols=2)

for l,like in enumerate(likes):
    ax[1,0].contour(A_range, B_range, like, levels=value, colors='#FF8D33', alpha=0.5)
    ax[0, 0].plot(A_range, likeAs[l], color='#FF8D33', label=None, alpha=0.5)
    ax[0, 1].axis("off")
    ax[1, 1].plot(likeBs[l], B_range, color='#FF8D33', label='loglike', alpha=0.5)
    
ax[1,0].scatter(θ_target[0], θ_target[1], marker='+', s=50, color='blue', zorder=20)
ax[0,0].axvline(θ_target[0], linestyle='--', c='k')
ax[1,0].axvline(θ_target[0], linestyle='--', c='k')
ax[1,0].axhline(θ_target[1], linestyle='--', c='k')
ax[1,1].axhline(θ_target[1], linestyle='--', c='k', label=r'$\theta_{\rm target}$')

ax[1,0].set_xlabel(r'$A$')
ax[1,0].set_ylabel(r'$B$')
```

# Initialise IMNN


```python
simulator_args["squeeze"] = False
```


```python
simulator_args['vol_norm'] = True
simulator_args['N_scale'] = True # false 
simulator_args['L'] = L
```


```python
simulator_args
```




    {'N': 128,
     'L': 128,
     'dim': 2,
     'shape': (128, 128),
     'vol_norm': True,
     'N_scale': True,
     'squeeze': False}




```python
IMNN = SimulatorIMNN(
        n_s=5000,
        n_d=5000,
        n_params=n_params,
        n_summaries=n_summaries,
        input_shape=input_shape,
        θ_fid=θ_fid,
        model=model,
        optimiser=optimiser,
        key_or_state=initial_model_key,
        simulator=lambda rng, θ: simulator(rng, θ, simulator_args=simulator_args),
#         devices=[jax.devices()[0]],
#         n_per_device=1000
        )
```

# Fit


```python
# SAVING IMNN ATTRIBUTES
import cloudpickle as pickle
import os

def save_weights(IMNN, folder_name='./model', weights='final'):
    
    # create output directory
    if not os.path.exists(folder_name): 
        os.mkdir(folder_name)
    
    def pckl_me(obj, path):
        with open(path, 'wb') as file_pi:
            pickle.dump(obj, file_pi)
        file_pi.close()
        
    # save IMNN (optimiser) state:
    savestate = jax.experimental.optimizers.unpack_optimizer_state(IMNN.state)
    pckl_me(savestate, os.path.join(folder_name, 'IMNN_state'))
    
    # save weights
    if weights == 'final':
        np.save(os.path.join(folder_name, 'final_w'), IMNN.final_w)
    
    else:
        np.save(os.path.join(folder_name, 'best_w'), IMNN.best_w)
        
    # save initial weights
    np.save(os.path.join(folder_name, 'initial_w'), IMNN.initial_w)
        
    # save training history
    pckl_me(IMNN.history, os.path.join(folder_name, 'history'))
    
    # save important attributes as a dict
    imnn_attributes = {
        'n_s': IMNN.n_s,
        'n_d': IMNN.n_d,
        'input_shape': IMNN.input_shape,
        'n_params' : IMNN.n_params,            
        'n_summaries': IMNN.n_summaries,
        'θ_fid': IMNN.θ_fid,
        'F': IMNN.F,   
        'validate': IMNN.validate,
        'simulate': IMNN.simulate,
        
    }
    
    pckl_me(imnn_attributes, os.path.join(folder_name, 'IMNN_attributes'))
    
    print('saved weights and attributes to the file ', folder_name)
    
    
def load_weights(IMNN, folder_name='./model', weights='final', load_attributes=True):
    
    
    def unpckl_me(path):
        file = open(path, 'rb')
        return pickle.load(file)
    
    # load and assign weights
    if weights=='final':
        weights = np.load(os.path.join(folder_name, 'final_w.npy'),  allow_pickle=True)
        IMNN.final_w = weights
    else:
        weights = np.load(os.path.join(folder_name, 'best_w.npy'),  allow_pickle=True)
        IMNN.best_w = weights
        
    # re-pack and load the optimiser state
    loadstate = unpckl_me(os.path.join(folder_name, 'IMNN_state'))
    IMNN.state = jax.experimental.optimizers.pack_optimizer_state(loadstate)
    
    
    # load history
    IMNN.history = unpckl_me(os.path.join(folder_name, 'history'))
    
    # load important attributes
    if load_attributes:
        IMNN.intial_w =  np.load(os.path.join(folder_name, 'initial_w.npy'), allow_pickle=True)
        attributes = unpckl_me(os.path.join('test_model', 'IMNN_attributes'))
        IMNN.θ_fid = attributes['θ_fid']
        IMNN.n_s = attributes['n_s']
        IMNN.n_d = attributes['n_d']
        IMNN.input_shape = attributes['input_shape']
        
    print('loaded IMNN with these attributes: ', attributes)
    
```


```python
# # test save functions
# save_weights(IMNN, folder_name='./model')

# # test load functions

# # initialize a new imnn with different attributes and then load the old file
# # to overwrite them

# my_new_IMNN = SimIMNN(
#         n_s=300,
#         n_d=100,
#         n_params=n_params,
#         n_summaries=n_summaries,
#         input_shape=input_shape,
#         θ_fid=np.array([1.0,1.0]),
#         key=initial_model_key,
#         model=model,
#         optimiser=optimiser,
#         simulator=lambda rng, θ: simulator(rng, θ, simulator_args=simulator_args),
#         )

# load_weights(my_new_IMNN, folder_name='./model', load_attributes=True)
# my_new_IMNN.set_F_statistics(rng, my_new_IMNN.best_w, my_new_IMNN.θ_fid, my_new_IMNN.n_s, my_new_IMNN.n_d, validate=True)
```


```python
θ_fid
```




    DeviceArray([1. , 0.5], dtype=float32)




```python
%%time
for i in range(1):
    rng,fit_rng = jax.random.split(rng)
    IMNN.fit(λ=10., ϵ=ϵ, rng=fit_rng, min_iterations=500) #for IMNN, IMNN_rng in zip(IMNNs, IMNN_rngs);
    
#save_weights(IMNN, folder_name='./big_incept128')
```

    CPU times: user 12min 57s, sys: 21.6 s, total: 13min 19s
    Wall time: 13min 26s



```python
IMNNs = [IMNN]
```


```python
latexify(fig_width=3.37)
plt.plot(IMNN.history['detF'][:])
plt.plot(np.ones(len(IMNN.history['detF'][:]))*detf_expected, c='k', linestyle='--')
plt.ylim(1e-2, 1e7)
plt.ylabel(r'$\det \textbf{F}$')
plt.xlabel('number of epochs')
plt.yscale('log')
plt.tight_layout()
plt.savefig('/mnt/home/tmakinen/repositories/field-plots/128x128-training.png', dpi=400)
```


```python
np.linalg.det(IMNNs[0].F) #/ (detf_expected)
```




    DeviceArray(878238.75, dtype=float32)




```python
IMNNs[0].F
```




    DeviceArray([[ 2485.2085, -1814.8022],
                 [-1814.8022,  1678.63  ]], dtype=float32)




```python
print('IMNN F:', IMNN.F)

print('IMNN det F:', np.linalg.det(IMNN.F))

print('IMNN F / analytic det F: ', (np.linalg.det(IMNN.F)) / detf_expected)
```

    IMNN F: [[ 2485.2085 -1814.8022]
     [-1814.8022  1678.63  ]]
    IMNN det F: 878238.75
    IMNN F / analytic det F:  0.9076835203931398


# Data for ABC example


```python
class uniform:
    def __init__(self, low, high):
        self.low = np.array(low)
        self.high = np.array(high)
        self.event_shape = [[] for i in range(self.low.shape[0])]
    def sample(self, n=None, seed=None):
        if n is None:
            n = 1
        keys = np.array(jax.random.split(
            seed, 
            num=len(self.event_shape)))
        return jax.vmap(
            lambda key, low, high : jax.random.uniform(
                key, 
                shape=(n,), 
                minval=low, 
                maxval=high))(
            keys, self.low, self.high)

```


```python
prior = uniform([0.6, 0.2], [1.25, 1.20])
```


```python
simulator_args
```




    {'N': 128,
     'L': 128,
     'dim': 2,
     'shape': (128, 128),
     'N_scale': True,
     'vol_norm': True,
     'squeeze': True}




```python
simulator_args = {"N": N, "L": L, "dim": dim, "shape": shape, "N_scale": True, "vol_norm": True, "squeeze": True}

rng, key = jax.random.split(rng)
θ_target = np.array([0.9, 0.6])
target_data = simulator(
    key,
    θ_target,
    simulator_args={**simulator_args, **{'squeeze':False}})
```

# analytic likelihood calculation


```python
target_data = np.load('./128x128-gauss/example-field_theta=%d_%d.npy'%(θ_target[0]*10, θ_target[1]*10))
target_data = np.expand_dims(np.expand_dims(np.expand_dims(target_data,0),0),0)
```


```python
gridsize = 100                # for likelihood gridding
Δ = np.fft.fftn(np.squeeze(target_data))[1:N//2, 1:N//2] / N
k = kbin[1:N//2, 1:N//2]  
prior_range = np.array([[0.1, 0.1], [1.25, 1.25]])
AL = analyticFieldLikelihood(PBJ, shape, Δ, prior_range, k=k, gridsize=gridsize, tiling=[5,5])
```


```python
%%time
%matplotlib inline

ax,like = AL.plot_corner(θ_ref=θ_target, image_data=target_data, return_like=True)
latexify(fig_height=5.37)
plt.show()
```


```python
np.save('./128x128-gauss/example-field_theta=%d_%d'%(θ_target[0]*10, θ_target[1]*10), np.squeeze(target_data))
```

# Gaussian approximation


```python
@jit #partial(jax.jit, static_argnums=0)
def get_estimate(d):
    if len(d.shape) == 1:
        return IMNN.θ_fid + np.einsum(
            "ij,kj,kl,l->i",
            IMNN.invF,
            IMNN.dμ_dθ,
            IMNN.invC,
            IMNN.model(IMNN.best_w, d, rng=rng) - IMNN.μ)
    else:
        return IMNN.θ_fid + np.einsum(
            "ij,kj,kl,ml->mi",
            IMNN.invF,
            IMNN.dμ_dθ,
            IMNN.invC,
            IMNN.model(IMNN.best_w, d, rng=rng) - IMNN.μ)
```


```python
estimates = IMNN.get_estimate(target_data) #[i.get_estimate(target_data) for i in IMNNs];
```


```python
estimates
```




    DeviceArray([[0.88272667, 0.5606261 ]], dtype=float32)




```python
GAs = [GaussianApproximation(IMNN.get_estimate(target_data), IMNN.invF, prior)]
       #GaussianApproximation(get_estimate(target_data), np.linalg.inv(f_expected), prior)]
```


```python
%matplotlib inline
for i, (GA, label) in enumerate(zip(GAs, ['sim IMNN'])):
    if i == 0:
        ax = GA.marginal_plot(
            axis_labels=[r"$A$", r"$B$"], label='on-the-fly IMNN', colours="C{}".format(i)
        )
    else:
        GA.marginal_plot(ax=ax, label='sim IMNN', colours="C{}".format(i), ncol=8)
```

# ABC


```python
{**simulator_args, **{'squeeze':False}}
```




    {'N': 128,
     'L': 128,
     'dim': 2,
     'shape': (128, 128),
     'N_scale': True,
     'vol_norm': True,
     'squeeze': False}




```python
ABC = ApproximateBayesianComputation(
        target_data, prior, 
     lambda A,B : simulator(A,B, simulator_args={**simulator_args, **{'squeeze':False}}), 
                                       IMNN.get_estimate, F=IMNN.F, gridsize=50
    )
```


```python
%%time
rng,abc_key = jax.random.split(rng)
ABC(rng=abc_key,
        n_samples=int(1e3),
        min_accepted=15000,
        max_iterations=50000,
        ϵ=0.05,
        smoothing=0.);
```

    1256 accepted in last  50000 iterations  (50000000 simulations done).
    CPU times: user 10min 56s, sys: 6.43 s, total: 11min 2s
    Wall time: 11min 5s





    (<imnn.utils.container.container at 0x15544cb57220>,
     <imnn.utils.container.container at 0x15544800ce80>,
     <imnn.utils.container.container at 0x15544800c370>)




```python
ABC.parameters.accepted[0].shape
```




    (1279, 2)




```python
ABC.parameters.accepted[0][0]
```




    DeviceArray([0.86439636, 0.55966151], dtype=float64)




```python
#np.save("accepted.npy", ABC.parameters.accepted)
```


```python
#ax = ABC.scatter_summaries(points=ABC.summaries.rejected, colours='red')
ABC.scatter_summaries( colours='blue')
```


```python
likelihood, A_range, B_range = AL.get_likelihood(return_grid=True)

likelihoodA = scipy.special.softmax(np.real(likelihood)).sum(0) #np.real(likelihood).sum(0)
likelihoodA /= likelihoodA.sum() * (A_range[1] -  A_range[0])
likelihoodB =  scipy.special.softmax(np.real(likelihood)).sum(1) #np.real(likelihood).sum(1)
likelihoodB /= likelihoodB.sum() * (B_range[1] -  B_range[0])

likelihood = scipy.special.softmax(np.real(likelihood))
sorted_marginal = np.sort(likelihood.flatten())[::-1]
cdf = np.cumsum(sorted_marginal / sorted_marginal.sum())
value = []
for level in [0.997, 0.95, 0.68]:
    this_value = sorted_marginal[np.argmin(np.abs(cdf - level))]
    if len(value) == 0:
        value.append(this_value)
    elif this_value <= value[-1]:
        break
    else:
        value.append(this_value)
        
```

    computing likelihood on a 100x100 grid 
             in tiles of size 20x20
    shift -4510.754477564908
    loglike mean -5927.549294142041



```python
%matplotlib inline
#plt.style.use('default')
new_colors = [ '#2c0342', '#286d87', '#4fb49d', '#9af486']
#fig = plt.figure(constrained_layout=True, figsize=(3.41*1., 3.41*1.))
#fig,ax = plt.subplots(nrows=2, ncols=2, figsize=(3.37*2, 3.37*2))
#ax = fig.subplots(nrows=2, ncols=2)
fig,ax = plt.subplots(nrows=2, ncols=2, figsize=(3.41*1., 3.41*1.), 
                      gridspec_kw={'height_ratios': [1, 1], 'width_ratios':[1,1]})
latexify(fig_width=3.41, fig_height=3.41)

# just to fiddle with the label
ax[0,0].plot(0.3, 0., color=new_colors[0],
           marker='o', label='ABC')

cmap_reversed = matplotlib.cm.get_cmap('viridis_r')



ax = GAs[0].marginal_plot(ax=ax, colours='#00c133', #new_colors[1], 
                          axis_labels=[r"$A$", r"$B$"], label="Gaussian Approximation", ncol=1,
                            linestyle='dotted')
ax[0,0].legend(framealpha=0.)
ax[0,1].imshow(np.squeeze(target_data), cmap='viridis')


ax[0,0].axvline(θ_target[0], linestyle='--', c='k')
ax[1,0].axvline(θ_target[0], linestyle='--', c='k')
ax[1,0].axhline(θ_target[1], linestyle='--', c='k')
ax[1,1].axhline(θ_target[1], linestyle='--', c='k', label=r'$\theta_{\rm target}$')

ax[1,0].set_xlabel(r'$A$')
ax[1,0].set_ylabel(r'$B$')

ax[0,0].axvline(θ_fid[0], linestyle='--', c='k', alpha=0.4)
ax[1,0].axvline(θ_fid[0], linestyle='--', c='k', alpha=0.4)
ax[1,0].axhline(θ_fid[1], linestyle='--', c='k', alpha=0.4)
ax[1,1].axhline(θ_fid[1], linestyle='--', c='k', alpha=0.4, label=r'$\theta_{\rm fid}$')

ax[1,1].legend(framealpha=0.)


# add in the likelihood estimate
ax[0, 0].plot(A_range, likelihoodA, color='#FF8D33', label='Analytic Likelihood', linestyle='--')
#ax[0,0].legend(framealpha=0.)
ax[0, 1].axis("off")
ax[1, 0].contour(A_range, B_range, likelihood, levels=value, colors='#FF8D33', linestyles='--')
ax[1, 1].plot(likelihoodB, B_range, color='#FF8D33', label='Analytic Likelihood', linestyle='--')

# ax = ABC.scatter_plot(ax=ax,
#                       colours=new_colors[0], 
#                       axis_labels=[r"$A$", r"$B$"], 
#                       s=8,
#                       label='ABC', bbox_to_anchor=None)

# ABC scatter plots
ax[0,0].hist(ABC.parameters.accepted[0][:, 0], color=new_colors[0], histtype='step', density=True)
ax[1,0].scatter(ABC.parameters.accepted[0][:, 0], ABC.parameters.accepted[0][:, 1], s=8, alpha=0.6,
                c=np.log(ABC.distances.accepted[0]), cmap='Purples', edgecolors=None, linewidths=0, marker='.')

ax[1,1].hist(ABC.parameters.accepted[0][:, 1], color=new_colors[0], 
             histtype='step', density=True, orientation='horizontal')


#ax[1,0].legend(framealpha=0.)



ax[0,0].set_xlim(0.55, 1.1)
#ax[0,0].set_ylim(0., 10.1)
ax[1,0].set_xlim(0.55, 1.1)
ax[1,0].set_ylim(0.3, 0.9)
ax[1,1].set_ylim(0.3, 0.9)

ax[0,0].legend(framealpha=0., bbox_to_anchor=(1.08, 1.5), frameon=False)


ax[1,1].set_yticks([])
ax[1,1].set_xticks([])
ax[0,0].set_xticks([])
ax[0,0].set_yticks([])

#ax[0,0].set_ylabel(r'$\mathit{P}(A|\textbf{x})$')
#ax[1,1].set_xlabel(r'$\mathit{P}(B|\textbf{x})$')

plt.subplots_adjust(wspace=0.1, hspace=0.1)
#plt.tight_layout()

plt.savefig('/mnt/home/tmakinen/repositories/field-plots/128x128-inference-contours-3sig-nolab-col.png', dpi=400, bbox_inches='tight')
#plt.subplots_adjust(wspace=0, hspace=0)
plt.show()
```


```python
# save all the contours 

# save all analytic calculations
fname = './128x128-gauss/'
np.save(fname + 'target_field_1', np.squeeze(target_data))
np.save(fname + 'analytic_likelihood', likelihood)
np.save(fname + 'ranges', np.stack((A_range, B_range), axis=0))

# save ABC posterior
np.save("./128x128-gauss/IMNN_accepted.npy", ABC.parameters.accepted)
```


```python
# Create figures in Python that handle LaTeX, and save images to files in my
# preferred formatting. I typically place this code in the root of each of my
# projects, and import using:
# from latexify import *
# which will also run the latexify() function on the import.

# Based on code from https://nipunbatra.github.io/blog/2014/latexify.html

import matplotlib
import matplotlib.pyplot as plt
from math import sqrt

#Back-end to use depends on the system
from matplotlib.backends.backend_pgf import FigureCanvasPgf
matplotlib.backend_bases.register_backend('pdf', FigureCanvasPgf)
# matplotlib.use('pgf')
# from matplotlib.backends.backend_pgf import FigureCanvasPgf
# matplotlib.backend_bases.register_backend('ps', FigureCanvasPgf)

import seaborn as sns
sns.set_style("white")

#my preferred palette. From
#https://seaborn.pydata.org/tutorial/color_palettes.html: "The cubehelix color
#palette system makes sequential palettes with a linear increase or decrease in
#brightness and some variation in hue. This means that the information in your
#colormap will be preserved when converted to black and white (for printing) or
#when viewed by a colorblind individual."

# I typically set the number of colors (below, 8) to the distinct colors I need
# in a given plot, so as to use the full range.
sns.set_palette(sns.color_palette("cubehelix", 8))


# The following is the latexify function. It allows you to create 2 column or 1
# column figures. You may also wish to alter the height or width of the figure.
# The default settings are good for most cases. You may also change the
# parameters such as labelsize and fontsize based on your classfile.
def latexify(fig_width=None, fig_height=None, columns=1):
    """Set up matplotlib's RC params for LaTeX plotting.
    Call this before plotting a figure.
    Parameters
    ----------
    fig_width : float, optional, inches
    fig_height : float,  optional, inches
    columns : {1, 2}
    """

    # code adapted from http://www.scipy.org/Cookbook/Matplotlib/LaTeX_Examples
    # Width and max height in inches for IEEE journals taken from
    # computer.org/cms/Computer.org/Journal%20templates/transactions_art_guide.pdf

    assert(columns in [1, 2])

    if fig_width is None:
        fig_width = 6.9 if columns == 1 else 13.8  # width in inches #3.39

    if fig_height is None:
        golden_mean = (sqrt(5) - 1.0) / 2.0    # Aesthetic ratio
        fig_height = fig_width * golden_mean  # height in inches

    MAX_HEIGHT_INCHES = 16.0
    if fig_height > MAX_HEIGHT_INCHES:
        print(("WARNING: fig_height too large:" + fig_height +
              "so will reduce to" + MAX_HEIGHT_INCHES + "inches."))
        fig_height = MAX_HEIGHT_INCHES

    params = {
            # 'backend': 'ps',
           #   'pgf.rcfonts': False,
           #   'pgf.preamble': ['\\usepackage{gensymb}', '\\usepackage[dvipsnames]{xcolor}'],
           #   "pgf.texsystem": "pdflatex",
             # 'text.latex.preamble': ['\\usepackage{gensymb}', '\\usepackage[dvipsnames]{xcolor}'],
              'text.latex.preamble': '\\usepackage{mathptmx}',
              #values below are useful defaults. individual plot fontsizes are
              #modified as necessary.
              'axes.labelsize': 8,  # fontsize for x and y labels
              'axes.titlesize': 8,
              'font.size': 8,
              'legend.fontsize': 8,
              'xtick.labelsize': 8,
              'ytick.labelsize': 8,
              'text.usetex': True,
              'figure.figsize': [fig_width, fig_height],
              'font.family': 'serif',
              'font.serif': 'Times',
               'lines.linewidth': 1.5,
               'lines.markersize':1,
               'xtick.major.pad' : 2,
               'ytick.major.pad' : 2,
                'axes.xmargin' :  .0,  # x margin.  See `axes.Axes.margins`
                'axes.ymargin' : .0,  # y margin See `axes.Axes.margins`
              }

    matplotlib.rcParams.update(params)

def saveimage(name, fig = plt, extension = 'pdf', folder = 'plots/'):
    sns.despine()

    #Minor ticks off by default in matplotlib
    # plt.minorticks_off()

    #grid being off is the default for seaborn white style, so not needed.
    # plt.grid(False, axis = "x")
    # plt.grid(False, axis = "y")
    
    fig.savefig('{}{}.{}'.format(folder,name, extension), bbox_inches = 'tight')

latexify()
```


```python

```

ipynb
b2aa4432ebc9ce254d07cc4d7ecb3c068327b1b7
Jupyter Notebook
1b363559116136d276b7c085b56a80520a7b0830
Apache-2.0
tlmakinen/FieldIMNNs
scratch/fields/128x128 - demo.ipynb
1b363559116136d276b7c085b56a80520a7b0830
Apache-2.0
tlmakinen/FieldIMNNs
scratch/fields/128x128 - demo.ipynb
2282e20d26dfc523e5705c0e02e0524b639df8a4
Apache-2.0
tlmakinen/kosmo-kompress
scratch/fields/128x128 - demo.ipynb
2021-11-04T12:41:06.000Z
2021-07-16T05:22:15.000Z