text:
```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

```

ext: ipynb

max_stars_repo_path: scratch/fields/128x128 - demo.ipynb

max_stars_repo_name: tlmakinen/kosmo-kompress

max_issues_repo_path: scratch/fields/128x128 - demo.ipynb

max_issues_repo_name: tlmakinen/FieldIMNNs

max_forks_repo_path: scratch/fields/128x128 - demo.ipynb

max_forks_repo_name: tlmakinen/FieldIMNNs


text:
Energy Spectrum of Prompt Fission Neutrons
=======================

The energy spectrum of the prompt fission neutrons, which we'll call as the
prompt fission spectrum for simplicity, represents the probability of neutrons
of different energies to be born in a fission event — a very important
parameter.

This quantity, usually denoted by $\chi_p$, is a probability distribution.
So let's check some simple approaches to it.

```python
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
import sympy as sym
from sympy.plotting import plot

mpl.style.use('seaborn')
```

```python
def neutron_pop(k):
    """
    Gived the ratio of neutron pupulation in generation i, N_i, by the initial neutron population N_0 given the number of generations.
    """
    ni_by_n0 = k**i
    return ni_by_n0

```

```python
i = sym.symbols('i')
crit_pop = neutron_pop(1.0)
super_pop = neutron_pop(1.18)
sub_pop = neutron_pop(0.9)

gen = 10

population = plot(sub_pop, crit_pop, super_pop, (i, 0, gen),
                  adaptive = False,
                  nb_of_points = 500,
                  size=(5,3),
                  xlabel='generation (-)',
                  ylabel='$N_i/N_0$ (-)',
                  axis_center=(-0.1,0.7),
                  show=False)
```

```python
x0, y0 = population._series[0].get_points()
x1, y1 = population._series[1].get_points()
x2, y2 = population._series[2].get_points()

fig, ax = plt.subplots(figsize=(5,3))
ax.plot(x0, y0, label='k=0.90')
ax.plot(x1, y1, label='k=1.00')
ax.plot(x2, y2, label='k=1.18')
ax.set_xlabel('generation')
ax.set_ylabel('$N_i/N_0$')
ax.xaxis.set_major_locator(MaxNLocator(integer=True))

ax.legend()

plt.savefig(f'k_pop_i{gen}.pdf', bbox_inches='tight')
```

```python

```


text:
```python
import requests
import pandas as pd
from time import time, mktime
import datetime
import numpy as np
import json
from math import cos, sqrt
from operator import itemgetter
from functools import reduce
import seaborn as sns
import matplotlib.pyplot as plt
```

```python
# defining route and date for get requests
route_id = '1'
date = '2020-06-01'

# route and location info from api, stops from Labs22
stops_url = 'https://raw.githubusercontent.com/Lambda-School-Labs/sfmta-data-analysis-ds/master/deprecated_assets/datasets/route_info.csv'
schedules_url = 'http://sfmta-ds.eba-hqpuyrup.us-east-1.elasticbeanstalk.com/get-route-info' # needs to be updated with api url when live
daily_url = 'http://sfmta-ds.eba-hqpuyrup.us-east-1.elasticbeanstalk.com/daily-general-json'

stops = pd.read_csv(stops_url)
sched_json = requests.get(schedules_url, params={'route_id': route_id,
                                                 'day': date}).json()
daily_json = requests.get(daily_url, params={'day': date}).json()
```

# First Look At Calculating Bunches & Gaps

## Helper Functions

```python
def fcc_projection(loc1, loc2):
    """
    function to apply FCC recommended formulae
    for calculating distances on earth projected to a plane

    significantly faster computationally, negligible loss in accuracy

    Args:
    loc1 - a tuple of lat/lon
    loc2 - a tuple of lat/lon
    """
    lat1, lat2 = loc1[0], loc2[0]
    lon1, lon2 = loc1[1], loc2[1]

    mean_lat = (lat1+lat2)/2
    delta_lat = lat2 - lat1
    delta_lon = lon2 - lon1

    k1 = 111.13209 - 0.56605*cos(2*mean_lat) + .0012*cos(4*mean_lat)
    k2 = 111.41513*cos(mean_lat) - 0.09455*cos(3*mean_lat) + 0.00012*cos(5*mean_lat)

    distance = sqrt((k1*delta_lat)**2 + (k2*delta_lon)**2)

    return distance
```

```python
def get_dwell(df):
    dwell_count = 0
    dwell_totals = []

    for x in df['kph']:
        if x <= 0:
          dwell_count += 1
          dwell_totals.append(dwell_count)
        elif x > 0:
          dwell_totals.append(0)
          dwell_count = 0

    df['dwell'] = [dwell_totals[x] for x in range(len(df))]

    return df
```

```python
def adjust_timestamp(df):
    times = df.timestamp.values
    ages = df.age.values

    df['adjusted_timestamp'] = [pd.Timestamp(times[x]) -
                                pd.Timedelta(seconds=ages[x])
                                for x in range(len(df.timestamp))]

    return df
```

```python
def assign_stops(df, stops):
    """
    applies basic wrangling function
    calculates nearest stop from reported location in km
    returns dataframe with reported location,
    nearest stop (coords and name), and distance between

    implements FCC projection formulae for calculating distance

    Args:
    df - dataframe of transit data, requires 'latitude', 'longitude' columns
    stops - datafram of stops data, requires 'lat', 'lon', 'title' columns
    """

    stop_lats = stops['lat'].values
    stop_lons = stops['lon'].values

    reported_lats = df['latitude'].values
    reported_lons = df['longitude'].values

    stop_points = [{'latitude': stop_lats[x], 'longitude': stop_lons[x]}
                 for x in range(len(stops))]

    reported_points = [{'latitude': reported_lats[x],
                      'longitude': reported_lons[x]}
                     for x in range(len(df))]

    radius = .1524

    stop_point_tuples = [tuple(stop_points[x].values())
                       for x in range(len(stop_points))]

    reported_point_tuples = [tuple(reported_points[x].values())
                           for x in range(len(reported_points))]

    df['reported_location'] = reported_point_tuples

    distances = [{x: fcc_projection(location, x)
                 for x in stop_point_tuples}
                 for location in reported_point_tuples]

    distances_sorted = [{k: v for k, v in sorted(distances[x].items(),
                                                 key=itemgetter(1))}
                       for x in range(len(distances))]

    point_stops = [next(iter(distances_sorted[x].items()))
                   if next(iter(distances_sorted[x].items()))[1] <= radius
                   else None
                   for x in range(len(distances_sorted))]

    stop_tuples = list(zip(stops.lat, stops.lon))
    stop_titles = [stops.title.iloc[stop_tuples.index(stop[0])]
                   if stop != None
                   else None
                   for stop in point_stops]

    df['nearest_stop'] = [x[0] if x != None else None for x in point_stops]
    df['distance_in_km'] = [x[1] if x != None else None for x in point_stops]

    df['title'] = stop_titles

    df = df.drop(columns=['age', 'rid', 'latitude', 'longitude', 'timestamp',
                          'heading'])

    return df
```

```python
def bin_from_timestamp(timestamp, bins):
    diffs = [abs(timestamp - x) for x in bins]
    return bins[diffs.index(min(diffs))]
```

## Function To Calculate B&Gs

```python
def find_disruptions(stops, daily_json, schedules_json):

    start = time()

    # pulling relevant info from schedules_json
    route = schedules_json['route']
    intervals = schedules_json['intervals']
    stop_list = schedules_json['stops']

    # filtering by route
    locations = pd.DataFrame(data=daily_json).dropna()
    locations = locations[locations.rid.eq(route)]

    # unifying direction tags for later merging
    locations['direction'] = ['Inbound' if '_I' in x
                              else 'Outbound' if '_O' in x
                              else None
                              for x in locations.direction.values]

    # filtering by route and current active stops
    stops = stops[stops.route_id.eq(route)]
    stops = stops[stops.tag.isin(stop_list)]
    stops = stops.rename(columns={'dir': 'direction'})

    print(f'Prep Complete: {time()-start} Seconds')

    # adjusting timestamps
    locations = adjust_timestamp(locations)

    print(f'Timestamps Adjusted: {time()-start} Seconds')

    # calculating dwell time for filtering
    locations = get_dwell(locations)

    print(f'Dwell Calculated: {time()-start} Seconds')

    # calculating nearest stops and distances
    locations = assign_stops(locations, stops)

    print(f'Stops Assigned: {time()-start} Seconds')

    # merging to assign titles to stops and unify tags
    locations = locations.merge(stops[stops.route_id.eq(route)],
                                how='left', on=['direction', 'title'])

    locations = locations.drop(columns=[
                                        'route_id', 'lat',
                                        'lon', 'stopId']).dropna()

    # dropping all reports with dwell != 0
    locations = locations[locations.dwell.eq(0)
                          ].reset_index(drop=True
                                        ).sort_values('adjusted_timestamp')

    print(f'Stops Filtered: {time()-start} Seconds')

    # separating by stop to calculate service deltas and disruptions
    stop_dfs = []
    for x in locations.tag.unique():

        df = locations[locations.tag.eq(x)].copy().reset_index(drop=True)

        df['window'] = [round((locations.loc[x, 'distance_in_km'] /
                                      np.mean(df.kph))*3600, 2)
                              for x in range(len(df))]

        df['service_deltas'] = [pd.Timestamp(df.loc[x, 'adjusted_timestamp'])
                                            - pd.Timestamp(df
                                            .loc[x-1,  'adjusted_timestamp'])
                                            if x != 0
                                            else None
                                            for x in range(len(df))]

        df['adjusted_service_deltas'] = [(df.loc[x, 'service_deltas'].total_seconds()
                                        - pd.Timedelta(seconds=df.loc[x, 'window']).total_seconds(),
                                        df.loc[x, 'service_deltas'].total_seconds()
                                        + pd.Timedelta(seconds=df.loc[x, 'window']).total_seconds())
                                        for x in range(len(df))]

        df['service_disruptions'] = [1 if (x[0]/60) < .2 * intervals['mode']
                                            else -1 if (x[1]/60) > 1.5 * intervals['mode']
                                            else 0
                                            for x in df['adjusted_service_deltas']]

        # dropping instances of multiple identical reports
        df.drop_duplicates(subset=['vid', 'adjusted_timestamp'],
                          keep='first', inplace=True)

        # creating new dataframe of timestamp / stop: service_disruptions
        df = pd.DataFrame(data={'adjusted_timestamp': df.adjusted_timestamp.values,
                                f'{x}': df.service_disruptions.values})

        stop_dfs.append(df)

    print(f'Services Analyzed: {time()-start} Seconds')

    # merging all stop-level dataframes
    # result is df of columns: stops, rows: timestamps, cells: B&Gs
    daily_service = reduce(lambda x, y: pd.merge(x, y,
                                               on='adjusted_timestamp',
                                               how='outer', sort=True),
                                               stop_dfs)

    print(f'Master DataFrame Compiled: {time()-start} Seconds')

    # Forward fill of service
    daily_service.fillna(method='ffill', inplace=True, axis=0)

    print(f'Naive Service Imputation Complete: {time()-start} Seconds')

    # creating a simple filter for calculating metrics
    stop_columns = daily_service.copy().columns

    # calculating basic metrics
    daily_service['on_time'] = daily_service[stop_columns].eq(0).sum(axis=1)
    daily_service['bunches'] = daily_service[stop_columns].eq(1).sum(axis=1)
    daily_service['gaps'] = daily_service[stop_columns].eq(-1).sum(axis=1)
    daily_service['service_disruptions'] = daily_service[['bunches',
                                                            'gaps']].sum(axis=1)

    daily_service['coverage'] = [round((x['on_time'] + x['bunches']) /
                                           (x['bunches'] + x['on_time']
                                            + x['gaps']
                                            ) * 100, 2)
                                  for _, x in daily_service.iterrows()]

    daily_service['on_time_%'] = [round(x['on_time'] /
                                            (x['on_time'] +
                                            x['service_disruptions']
                                             ) * 100, 2)
                                  for _, x in daily_service.iterrows()]

    average_coverage = np.mean(daily_service.coverage)
    average_on_time = np.mean(daily_service['on_time_%'])
    total_bunches = daily_service.bunches.sum()
    total_gaps = daily_service.gaps.sum()
    total_disruptions = daily_service.service_disruptions.sum()
    total_on_time = daily_service.on_time.sum()
    total_services = (daily_service.on_time.sum() +
                            daily_service.service_disruptions.sum())

    print(f'Metrics Calculated - Total Time: {time()-start} Seconds')

    return {'dataframe': daily_service,
            'metrics': {'coverage': average_coverage,
                        'on_time_%': average_on_time,
                        'totals': {'bunches': total_bunches,
                                   'gaps': total_gaps,
                                   'disruptions': total_disruptions,
                                   'on_time': total_on_time,
                                   'services': total_services}}}
```

```python
output = find_disruptions(stops, daily_json, sched_json)
daily_service = output['dataframe']
```

    Prep Complete: 0.6126761436462402 Seconds
    Timestamps Adjusted: 0.9255528450012207 Seconds
    Dwell Calculated: 0.9312446117401123 Seconds
    Stops Assigned: 1.7302463054656982 Seconds
    Stops Filtered: 1.7505877017974854 Seconds
    Services Analyzed: 2.103181838989258 Seconds
    Master DataFrame Compiled: 2.187445640563965 Seconds
    Naive Service Imputation Complete: 2.189103603363037 Seconds
    Metrics Calculated - Total Time: 2.478125810623169 Seconds

```python
output['metrics']
```

    {'coverage': 44.532984031936294,
     'on_time_%': 33.92160678642728,
     'totals': {'bunches': 1799,
      'disruptions': 11224,
      'gaps': 9425,
      'on_time': 5555,
      'services': 16779}}

```python
daily_service.head(5)
```

      adjusted_timestamp
      4026.0
      4028.0
      4015.0
      4027.0
      3825.0
      6314.0
      6312.0
      3893.0
      3848.0
      6296.0
      34015.0
      6311.0
      4277.0
      3892.0
      34277.0
      3824.0
      6295.0
      on_time
      bunches
      gaps
      service_disruptions
      coverage
      on_time_%

      0
      2020-06-01 03:42:13
      0.0
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      1
      0
      0
      0
      100.0
      100.0

      1
      2020-06-01 03:47:30
      0.0
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      1
      0
      0
      0
      100.0
      100.0

      2
      2020-06-01 03:49:56
      0.0
      0.0
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      2
      0
      0
      0
      100.0
      100.0

      3
      2020-06-01 03:51:59
      0.0
      0.0
      0.0
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      3
      0
      0
      0
      100.0
      100.0

      4
      2020-06-01 04:02:27
      0.0
      0.0
      0.0
      0.0
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      NaN
      4
      0
      0
      0
      100.0
      100.0

```python
fig, ax = plt.subplots()

sns.lineplot(x=daily_service.index, y=daily_service['on_time_%'])
sns.lineplot(x=daily_service.index, y=daily_service.coverage)

ax.legend(['On-Time %', 'Coverage %'])
ax.grid()
fig.set_size_inches(20, 10)
plt.title('On-Time % & Coverage', {'fontsize': 20})
plt.ylabel('Percentage')
plt.xlabel('Time')
plt.show()

fig, ax = plt.subplots()

sns.lineplot(x=daily_service.index, y=daily_service.bunches)
sns.lineplot(x=daily_service.index, y=daily_service.gaps)

ax.legend(['Bunches', 'Gaps'])
ax.grid()
fig.set_size_inches(20, 10)
plt.title('Bunches & Gaps', {'fontsize': 20})
plt.ylabel('Disruptions')
plt.xlabel('Time')
plt.show()

fig, ax = plt.subplots()

sns.lineplot(x=daily_service.index, y=daily_service.on_time)
sns.lineplot(x=daily_service.index, y=daily_service.service_disruptions)

ax.legend(['On-Time', 'Disruptions'])
ax.grid()
fig.set_size_inches(20, 10)
plt.title('On-Time Services & Service Disruptions', {'fontsize': 20})
plt.ylabel('Services')
plt.xlabel('Time')
plt.show();
```

```python
sns.pairplot(data=daily_service[['on_time_%', 'on_time', 'bunches',
                               'gaps', 'service_disruptions', 'coverage']]);
```

```python
sns.jointplot(daily_service.bunches,
              daily_service.gaps,
              kind='kde')

sns.jointplot(daily_service.gaps,
              daily_service['on_time_%'],
              kind='kde')

sns.jointplot(daily_service.service_disruptions,
              daily_service['on_time_%'],
              kind='kde');
```

# Results

Right off the bat, it's important to note that this is far from a perfect, precise approach. I've made a lot of assumptions and simplifications for the sake of time and ease, namely:
- Dwell-time reported at a stop shouldn't be used to determine bunching/\
gapping; we're concerned with headway *between* busses, not necessarily the progress of a single bus

- Rather than trying to calculate exactly when a bus was at a stop, I've based my bunches and gaps off of a *window* of time calculated as the time since a stop was last serviced $T$ plus or minus the reported distance from the stop $\Delta$ divided by the average kph for that route $v$

\begin{equation}
  w = T \pm (\Delta \div v)
\end{equation}

- If the lower bound of $w$ is below 20% of the common scheduled headway $h$, that stop registers a bunch ($s=1$) at that time; else if the upper bound of $w$ is above 150% of $h$, that stop registers a gap ($s=-1$) at that time; if stop isn't bunched or gapped we register an "on-time" service ($s=0$).

\begin{equation}
  s=\begin{cases}
    1, & \text{if $w<0.2 \cdot h$}.\\
    -1, & \text{if $w>1.5 \cdot h$}.\\
    0, & \text{otherwise.}
  \end{cases}
\end{equation}

- State is calculated for every service at a given stop and assumed to persist until a subsequent service alters that state

- Rather than dealing with a largely empty dataframe, we are dealing only with timestamps for which there was a recorded service for at least one stop in the line; every row contains at least one "true" observation alongside the "assumed" observations, and our number of observations is limited by the number of "true" observations across unique timestamps

- Stops are (temporarily) being pulled from the Labs22 data for my own ease and speed in writing; code should plug-and-play with similarly formatted but up-to-date data

I'm calculating coverage at a given timestamp $C$ as on-time services $s_0$ at that time plus bunched services $s_1$ at that time, divided by total services $s_t$ at that time

\begin{equation}
  C = \frac{s_0 + s_1}{s_t}
\end{equation}

And on-time frequency at a given timestamp $f$ simply as on-time services $s_0$ divided by total services $s_t$

\begin{equation}
  f = \frac{s_0}{s_t}
\end{equation}

The end result of all of this is that we do have actual metrics and visualizations that at least *approximate* the actual health of the system.

By no means would I imagine this to particularly accurate, but it should be good enough to provide more and better insight than a random guess.

**All models are wrong, but some are useful.**

¯\\_(ツ)_/¯

max_stars_repo_name: austiezr/sfmta-data-analysis-ds

max_issues_repo_name: austiezr/sfmta-data-analysis-ds

max_forks_repo_name: austiezr/sfmta-data-analysis-ds


text:
# Método de Newton
J.J
---
\begin{equation}
x_{k+1} = x_{k} - \frac{f(x_{k})}{f'(x_{k})}
\end{equation}

Ejemplo: $f(x) = x^{2} + 3x +1$

```python
import numpy as np
from sympy import *
from sympy.utilities.lambdify import lambdify
import matplotlib.pyplot as plt
init_printing(use_unicode=True)
```

```python
#calcular la derivada de f
x = symbols('x')
funcion = 1 + 3*x + x**2 #escribir la función aqui
dfuncion = diff(funcion, x)
print(str(dfuncion))
```

    2*x + 3

```python
f = lambdify(x, funcion)
df = lambdify((x), dfuncion)
```

```python
e = 0.0001 #error
maxit = 1000000 #iteraciones máximas
x0 = 2. #punto inicial
```

```python
def NewtonR(x0, func = f, dfunc = df, error = e, iterations = maxit):
    it = 0
    while (abs(f(x0)) > e) and (it < maxit):
        it += 1
        x0 = x0 - f(x0)/df(x0)
    return x0
```

```python
sol = NewtonR(x0)
print(sol)
```

    -0.3819660098414816

```python
X = np.linspace(-4, 2, 100)
plt.plot(X, f(X))
plt.plot(sol,f(sol),'ro')
plt.show()
```

```python

```

max_stars_repo_path: Ec. No lineales/Metodo_Newton.ipynb

max_stars_repo_name: JosueJuarez/M-todos-Num-ricos

max_issues_repo_path: Ec. No lineales/Metodo_Newton.ipynb

max_issues_repo_name: JosueJuarez/M-todos-Num-ricos

max_forks_repo_path: Ec. No lineales/Metodo_Newton.ipynb

max_forks_repo_name: JosueJuarez/M-todos-Num-ricos