Pytorch

README.md

@@ -1,24 +1,48 @@
 # pydelfi
 
-**NOTE:** currently only works with tensorflow <= 1.15; tensorflow 2 update coming soon.
-
 **Density Estimation Likelihood-Free Inference** with neural density estimators and adaptive acquisition of simulations. The implemented methods are described in detail in [Alsing, Charnock, Feeney and Wandelt 2019](https://arxiv.org/abs/1903.00007), and are based closely on [Papamakarios, Sterratt and Murray 2018](https://arxiv.org/pdf/1805.07226.pdf), [Lueckmann et al 2018](https://arxiv.org/abs/1805.09294) and [Alsing, Wandelt and Feeney, 2018](https://academic.oup.com/mnras/article-abstract/477/3/2874/4956055?redirectedFrom=fulltext). Please cite these papers if you use this code!
 
 **Installation:**
 
-The code is in python3 and has the following dependencies:<br>
+The code is in python3. There is a Tensorflow 1 (most stable, see below) and Tensorflow 2 version that can be installed as follows:<br>
+
+**Tensorflow 1 (stable)**
+
+This can be found on the master branch and has the following dependencies:<br>
 [tensorflow](https://www.tensorflow.org) (<=1.15) <br> 
 [getdist](http://getdist.readthedocs.io/en/latest/)<br>
-[emcee](http://dfm.io/emcee/current/)<br>
+[emcee](http://dfm.io/emcee/current/) (>=3.0.2)<br>
 [tqdm](https://github.com/tqdm/tqdm)<br>
 [mpi4py](https://mpi4py.readthedocs.io/en/stable/) (if MPI is required)<br>
 
 You can install the requirements and this package with,
+
 ```
+pip install tensorflow==1.15
 pip install git+https://github.com/justinalsing/pydelfi.git
 ```
+(`tensorflow-gpu==1.15` for GPU acceleration instead of `tensorflow==1.15`)
+
 or alternatively, pip install the requirements and then clone the repo and run `python setup.py install`
 
+**Tensorflow 2**
+
+The Tensorflow 2 version can be found on the `tf2-tom` branch and can be installed as follows. We reccommend you do the install inside a virtual environment to keep version conflicts under control, ie., 
+
+```
+mkdir ~/envs
+virtualenv ~/envs/pydelfi
+source ~/envs/pydelfi/bin/activate
+```
+
+Followed by a pip install of pydelfi:
+
+```
+pip install git+https://github.com/justinalsing/pydelfi.git@tf2-tom
+```
+
+Note: the Mixture Density Networks (MDN) in the tf2 version are currently not performing as well as in the tf1 version (but the Masked Autoregressive Flows are fine). We are getting ot the bottom of this, and also working on expanding the suite of conditional density estimators in a coming update. Watch this space.
+
 **Documentation and tutorials:** 
 
 Once everything is installed, try out either `cosmic_shear.ipynb` or `jla_sne.ipynb` as example templates for how to use the code; plugging in your own simulator and letting pydelfi do it's thing. 

examples/cosmic_shear.ipynb

@@ -6,19 +6,28 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import sys\n",
+    "import tensorflow as tf\n",
     "import numpy as np\n",
     "import matplotlib.pyplot as plt\n",
     "import getdist\n",
     "from getdist import plots, MCSamples\n",
     "import pydelfi.priors as priors\n",
     "import pydelfi.ndes as ndes\n",
     "import pydelfi.delfi as delfi\n",
-    "import tensorflow as tf\n",
+    "import pydelfi.score as score\n",
     "import simulators.cosmic_shear.cosmic_shear as cosmic_shear\n",
     "import pickle\n",
-    "import pydelfi.score as score\n",
-    "tf.logging.set_verbosity(tf.logging.ERROR)\n",
-    "%matplotlib inline"
+    "import tensorflow_probability as tfp\n",
+    "tfd = tfp.distributions"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Set up the simulator\n",
+    "This must have the signature `simulator(parameters, seed, args, batch)` -> `np.array([batch, ndata])`"
    ]
   },
   {
@@ -27,37 +36,45 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "### SET UP THE SIMULATOR ###\n",
-    "\n",
-    "# Set up the tomography simulations\n",
     "CosmicShearSimulator = cosmic_shear.TomographicCosmicShear(pz = pickle.load(open('simulators/cosmic_shear/pz_5bin.pkl', 'rb')),\n",
     "                                                           lmin = 10, lmax = 1000, n_ell_bins = 5, \n",
     "                                                           sigma_e = 0.3, nbar = 30, Area = 15000)\n",
     "\n",
-    "# Simulator function: This must be of the form simulator(theta, seed, args) -> simulated data vector\n",
     "def simulator(theta, seed, simulator_args, batch=1):\n",
     "    return CosmicShearSimulator.simulate(theta, seed)\n",
     "\n",
-    "# Simulator arguments\n",
     "simulator_args = None"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Set up the prior"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
-    "### SET UP THE PRIOR ###\n",
+    "lower = np.array([0, 0.4, 0, 0.4, 0.7]).astype('float32')\n",
+    "upper = np.array([1, 1.2, 0.1, 1.0, 1.3]).astype('float32')\n",
+    "prior_mean = np.array([0.3, 0.8, 0.05, 0.70, 0.96]).astype('float32')\n",
+    "prior_covariance = (np.eye(5)*np.array([0.1, 0.1, 0.05, 0.3, 0.3])**2).astype('float32')\n",
+    "prior_stddev = np.sqrt(np.diag(prior_covariance))\n",
     "\n",
-    "# Define the priors parameters\n",
-    "lower = np.array([0, 0.4, 0, 0.4, 0.7])\n",
-    "upper = np.array([1, 1.2, 0.1, 1.0, 1.3])\n",
-    "prior_mean = np.array([0.3, 0.8, 0.05, 0.70, 0.96])\n",
-    "prior_covariance = np.eye(5)*np.array([0.1, 0.1, 0.05, 0.3, 0.3])**2\n",
-    "\n",
-    "# Prior\n",
-    "prior = priors.TruncatedGaussian(prior_mean, prior_covariance, lower, upper)"
+    "prior = tfd.Blockwise([tfd.TruncatedNormal(loc=prior_mean[i], scale=prior_stddev[i], low=lower[i], high=upper[i]) for i in range(5)])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Set up the compressor\n",
+    "Must have the signature `compressor(data, args)` -> `np.array([n_summaries])`<br>\n",
+    "In this case we are going to do Wishart score compression."
    ]
   },
   {
@@ -66,8 +83,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "### SET UP THE COMPRESSOR ###\n",
-    "\n",
     "# Fiducial parameters\n",
     "theta_fiducial = np.array([0.3, 0.8, 0.05, 0.70, 0.96])\n",
     "\n",
@@ -94,39 +109,85 @@
     "compressor_args = None"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Generate mock data vector"
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
-    "### GENERATE MOCK DATA VECTOR ###\n",
-    "\n",
     "seed = 0\n",
     "data = simulator(theta_fiducial, seed, simulator_args)\n",
     "compressed_data = compressor(data, compressor_args)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Create ensemble of NDEs"
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "WARNING:tensorflow:From /obs/njeffrey/envs/delfi2env/lib/python3.7/site-packages/tensorflow/python/ops/linalg/linear_operator_lower_triangular.py:167: calling LinearOperator.__init__ (from tensorflow.python.ops.linalg.linear_operator) with graph_parents is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Do not pass `graph_parents`.  They will  no longer be used.\n",
+      "WARNING:tensorflow:From /obs/njeffrey/envs/delfi2env/lib/python3.7/site-packages/tensorflow_probability/python/distributions/distribution.py:334: calling TransformedDistribution.__init__ (from tensorflow_probability.python.distributions.transformed_distribution) with batch_shape is deprecated and will be removed after 2020-06-01.\n",
+      "Instructions for updating:\n",
+      "`batch_shape` and `event_shape` args are deprecated. Please use `tfd.Sample`, `tfd.Independent`, and broadcasted parameters of the base distribution instead. For example, replace `tfd.TransformedDistribution(tfd.Normal(0., 1.), tfb.Exp(), batch_shape=[2, 3], event_shape=[4])` with `tfd.TransformedDistrbution(tfd.Sample(tfd.Normal(tf.zeros([2, 3]), 1.),sample_shape=[4]), tfb.Exp())` or `tfd.TransformedDistribution(tfd.Independent(tfd.Normal(tf.zeros([2, 3, 4]), 1.), reinterpreted_batch_ndims=1), tfb.Exp())`.\n",
+      "WARNING:tensorflow:From /obs/njeffrey/envs/delfi2env/lib/python3.7/site-packages/tensorflow_probability/python/distributions/distribution.py:334: calling TransformedDistribution.__init__ (from tensorflow_probability.python.distributions.transformed_distribution) with event_shape is deprecated and will be removed after 2020-06-01.\n",
+      "Instructions for updating:\n",
+      "`batch_shape` and `event_shape` args are deprecated. Please use `tfd.Sample`, `tfd.Independent`, and broadcasted parameters of the base distribution instead. For example, replace `tfd.TransformedDistribution(tfd.Normal(0., 1.), tfb.Exp(), batch_shape=[2, 3], event_shape=[4])` with `tfd.TransformedDistrbution(tfd.Sample(tfd.Normal(tf.zeros([2, 3]), 1.),sample_shape=[4]), tfb.Exp())` or `tfd.TransformedDistribution(tfd.Independent(tfd.Normal(tf.zeros([2, 3, 4]), 1.), reinterpreted_batch_ndims=1), tfb.Exp())`.\n"
+     ]
+    }
+   ],
    "source": [
-    "# Create an ensemble of NDEs\n",
-    "NDEs = [ndes.ConditionalMaskedAutoregressiveFlow(n_parameters=5, n_data=5, n_hiddens=[50,50], n_mades=5, act_fun=tf.tanh, index=0),\n",
-    "        ndes.MixtureDensityNetwork(n_parameters=5, n_data=5, n_components=1, n_hidden=[30,30], activations=[tf.tanh, tf.tanh], index=1),\n",
-    "        ndes.MixtureDensityNetwork(n_parameters=5, n_data=5, n_components=2, n_hidden=[30,30], activations=[tf.tanh, tf.tanh], index=2),\n",
-    "        ndes.MixtureDensityNetwork(n_parameters=5, n_data=5, n_components=3, n_hidden=[30,30], activations=[tf.tanh, tf.tanh], index=3),\n",
-    "        ndes.MixtureDensityNetwork(n_parameters=5, n_data=5, n_components=4, n_hidden=[30,30], activations=[tf.tanh, tf.tanh], index=4),\n",
-    "        ndes.MixtureDensityNetwork(n_parameters=5, n_data=5, n_components=5, n_hidden=[30,30], activations=[tf.tanh, tf.tanh], index=5)]\n",
+    "NDEs = [ndes.ConditionalMaskedAutoregressiveFlow(\n",
+    "            n_parameters=5,\n",
+    "            n_data=5,\n",
+    "            n_mades=5,\n",
+    "            n_hidden=[30,30], \n",
+    "            activation=tf.keras.layers.LeakyReLU(0.01),\n",
+    "            kernel_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=1e-5, seed=None),\n",
+    "            all_layers=True)]\n",
+    "\n",
+    "NDEs += [ndes.MixtureDensityNetwork(\n",
+    "            n_parameters=5,\n",
+    "            n_data=5, \n",
+    "            n_components=i+1,\n",
+    "            n_hidden=[30], \n",
+    "            activation=tf.keras.layers.LeakyReLU(0.01))\n",
+    "        for i in range(5)]\n",
     "\n",
-    "# Create the DELFI object\n",
-    "DelfiEnsemble = delfi.Delfi(compressed_data, prior, NDEs, Finv=Finv, theta_fiducial=theta_fiducial, \n",
-    "                       param_limits = [lower, upper],\n",
-    "                       param_names = ['\\Omega_m', 'S_8', '\\Omega_b', 'h', 'n_s'], \n",
-    "                       results_dir = \"simulators/cosmic_shear/results/\",\n",
-    "                       input_normalization='fisher')"
+    "NDEs += [ndes.SinhArcSinhMADE(\n",
+    "            n_parameters=5,\n",
+    "            n_data=5,\n",
+    "            n_hidden=[64],\n",
+    "            activation=tf.tanh,\n",
+    "            kernel_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=1e-5, seed=None),\n",
+    "            bias_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=1e-5, seed=None)\n",
+    "        )]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Create DELFI object"
    ]
   },
   {
@@ -135,8 +196,42 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Do the Fisher pre-training\n",
-    "DelfiEnsemble.fisher_pretraining()"
+    "DelfiEnsemble = delfi.Delfi(compressed_data, prior, NDEs, \n",
+    "                            Finv=Finv, \n",
+    "                            theta_fiducial=theta_fiducial,\n",
+    "                            param_limits = [lower, upper],\n",
+    "                            param_names=['\\Omega_m', 'S_8', '\\Omega_b', 'h', 'n_s'], \n",
+    "                            results_dir=\"simulators/cosmic_shear/results\",\n",
+    "                            filename=\"cosmic_shear\",\n",
+    "                            optimiser=tf.keras.optimizers.Adam(lr=1e-4),\n",
+    "                            optimiser_arguments=None,\n",
+    "                            dtype=tf.float32,\n",
+    "                            posterior_chain_length=200,\n",
+    "                            nwalkers=500,\n",
+    "                            input_normalization=\"fisher\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Fisher pre-training to initialize NDEs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# DelfiEnsemble.fisher_pretraining(n_batch=5000, epochs=1000, patience=20, plot=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Sequential Neural Likelihood"
    ]
   },
   {
@@ -151,22 +246,15 @@
     "n_populations = 39\n",
     "\n",
     "# Do the SNL training\n",
-    "DelfiEnsemble.sequential_training(simulator, compressor, n_initial, n_batch, n_populations, patience=10, save_intermediate_posteriors=True)"
+    "DelfiEnsemble.sequential_training(simulator, compressor, n_initial, n_batch, n_populations, patience=10, plot=True, save_intermediate_posteriors=True)"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "delfi2env",
    "language": "python",
-   "name": "python3"
+   "name": "delfi2env"
   },
   "language_info": {
    "codemirror_mode": {
@@ -178,7 +266,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.5"
+   "version": "3.7.3"
   }
  },
  "nbformat": 4,