diff --git a/README b/README
deleted file mode 120000
index 602d87dfd9405502dcd988b9dab5bb724dc4471d..0000000000000000000000000000000000000000
--- a/README
+++ /dev/null
@@ -1 +0,0 @@
-Readme.md
\ No newline at end of file
diff --git a/docs/wrapper.rst b/docs/RST/libcooling.rst
similarity index 68%
rename from docs/wrapper.rst
rename to docs/RST/libcooling.rst
index ed4f2febcf5bf875090e8dabfbb6c720629ddc28..b83bf4409a206c2f0e27c83b51bb0a1fc7b2382a 100644
--- a/docs/wrapper.rst
+++ b/docs/RST/libcooling.rst
@@ -1,7 +1,7 @@
-Wrapper
-=======
+Libcooling
+==========
 
-.. automodule:: pyswiftsim.wrapper
+.. automodule:: pyswiftsim.libcooling
     :members:
     :undoc-members: 
     :private-members:
diff --git a/docs/pyswiftsim.rst b/docs/RST/pyswiftsim.rst
similarity index 100%
rename from docs/pyswiftsim.rst
rename to docs/RST/pyswiftsim.rst
diff --git a/docs/conf.py b/docs/conf.py
index aa7b47b0fa36a714271edc205433577513fdff66..8d40b6e2f01cd0c989474cbfab4425561875bd84 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -82,7 +82,7 @@ pygments_style = None
 # The theme to use for HTML and HTML Help pages.  See the documentation for
 # a list of builtin themes.
 #
-html_theme = 'alabaster'
+html_theme = 'sphinx_rtd_theme'
 
 # Theme options are theme-specific and customize the look and feel of a theme
 # further.  For a list of options available for each theme, see the
@@ -199,3 +199,5 @@ todo_include_todos = True
 # -- Options for numpy extension ----------------------------------------------
 
 numpydoc_use_plots = True
+
+numpydoc_xref_param_type = True
diff --git a/docs/index.rst b/docs/index.rst
index 9c79e721d38fe05c2ccbba38238a25fd08a9ea35..d11523e49bcf892c03137a516830300e753817cc 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -6,12 +6,15 @@
 Welcome to PySWIFTsim's documentation!
 ======================================
 
+PySWIFTsim is a python wrapper around the gravity and SPH solver SWIFT.
+The aim is to provide a python tool that benefits from HPC code and allows to directly use some feature of SWIFT in order to either test it or analyze a simulation.
+
 .. toctree::
    :maxdepth: 2
-   :caption: Contents:
 
-   pyswiftsim.rst
-   wrapper.rst
+   RST/pyswiftsim.rst
+   RST/libcooling.rst
+   RST/examples.rst
 
 
 
diff --git a/pyswiftsim/__init__.py b/pyswiftsim/__init__.py
index 526dbce334dd759ae3af7299c78404f02c44fa8d..8c6a5748ae78c94a83aee7c1d808247ac8fd102e 100644
--- a/pyswiftsim/__init__.py
+++ b/pyswiftsim/__init__.py
@@ -1,19 +1,30 @@
-"""
-This file is part of PYSWIFT.
-Copyright (c) 2018 Loic Hausammann (loic.hausammann@epfl.ch)
+# ########################################################################
+# This file is part of PYSWIFT.
+# Copyright (c) 2018 Loic Hausammann (loic.hausammann@epfl.ch)
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+# ########################################################################
 
-This program is free software: you can redistribute it and/or modify
-it under the terms of the GNU Lesser General Public License as published
-by the Free Software Foundation, either version 3 of the License, or
-(at your option) any later version.
+import os
+import urllib.request
 
-This program is distributed in the hope that it will be useful,
-but WITHOUT ANY WARRANTY; without even the implied warranty of
-MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-GNU General Public License for more details.
 
-You should have received a copy of the GNU Lesser General Public License
-along with this program.  If not, see <http://www.gnu.org/licenses/>.
-"""
+def downloadCoolingTable():
+    url = "http://virgodb.cosma.dur.ac.uk/"
+    url += "swift-webstorage/CoolingTables/CloudyData_UVB=HM2012.h5"
+    filename = "CloudyData_UVB=HM2012.h5"
 
-from pyswiftsim import *
+    if not os.path.isfile(filename):
+        # Download the file from `url` and save it locally under `file_name`:
+        urllib.request.urlretrieve(url, filename)
diff --git a/setup.py b/setup.py
index 919adee2849f4fc5dcc55e8dd88d25a6b2ccf337..1c76bc94e8463de8a2bf4af34c8a9c61a2505365 100644
--- a/setup.py
+++ b/setup.py
@@ -118,7 +118,7 @@ data_files = []
 ##############
 
 c_src = glob("src/*.c")
-ext_modules = Extension("pyswiftsim.wrapper",
+ext_modules = Extension("pyswiftsim.libcooling",
                         c_src,
                         include_dirs=include,
                         libraries=lib,
diff --git a/src/cooling_wrapper.c b/src/cooling_wrapper.c
index 8b1d399bf2ecda8de40b7e58e38375c41e1f75e0..47dd1ccb52857d1794719eaac9ba2bdcedded1c2 100644
--- a/src/cooling_wrapper.c
+++ b/src/cooling_wrapper.c
@@ -19,6 +19,48 @@
 #include "cooling_wrapper.h"
 #include "pyswiftsim_tools.h"
 
+/* TODO Remove this */
+struct cooling_function_data cooling;
+int initialized;
+
+
+/**
+ * @brief Set the cooling element fractions
+ *
+ * @param xp The #xpart to set
+ * @param frac The numpy array containing the fractions (id, element)
+ * @param idx The id (in frac) of the particle to set
+ */
+void pycooling_set_fractions(struct xpart *xp, PyArrayObject* frac, const int idx) {
+#ifdef COOLING_GRACKLE
+  struct cooling_xpart_data *data = &xp->cooling_data;
+  data->metal_frac = *(float*)PyArray_GETPTR2(frac, idx, 0);
+
+#ifdef COOLING_GRACKLE
+#if COOLING_GRACKLE_MODE > 0
+  data->HI_frac = *(float*)PyArray_GETPTR2(frac, idx, 1);
+  data->HII_frac = *(float*)PyArray_GETPTR2(frac, idx, 2);
+  data->HeI_frac = *(float*)PyArray_GETPTR2(frac, idx, 3);
+  data->HeII_frac = *(float*)PyArray_GETPTR2(frac, idx, 4);
+  data->HeIII_frac = *(float*)PyArray_GETPTR2(frac, idx, 5);
+  data->e_frac = *(float*)PyArray_GETPTR2(frac, idx, 6);
+#endif // COOLING_GRACKLE_MODE
+#if COOLING_GRACKLE_MODE > 1
+  data->HM_frac = *(float*)PyArray_GETPTR2(frac, idx, 7);
+  data->H2I_frac = *(float*)PyArray_GETPTR2(frac, idx, 8);
+  data->H2II_frac = *(float*)PyArray_GETPTR2(frac, idx, 9);
+#endif // COOLING_GRACKLE_MODE
+#if COOLING_GRACKLE_MODE > 2
+  data->DI_frac = *(float*)PyArray_GETPTR2(frac, idx, 10);
+  data->DII_frac = *(float*)PyArray_GETPTR2(frac, idx, 11);
+  data->HDI_frac = *(float*)PyArray_GETPTR2(frac, idx, 12);
+#endif // COOLING_GRACKLE_MODE
+#endif // COOLING_GRACKLE
+
+#else
+  error("Not implemented");
+#endif
+}
 
 /**
  * @brief Compute the cooling rate
@@ -50,19 +92,19 @@ PyObject* pycooling_rate(PyObject* self, PyObject* args) {
 
   /* parse argument */
   if (!PyArg_ParseTuple(
-        args, "siOO|fO", &filename, &with_cosmo,
-	&rho, &energy, &dt, &fractions))
+        args, "sOO|ifO", &filename, &rho, &energy,
+	&with_cosmo, &dt, &fractions))
     return NULL;
 
   /* check numpy array */
-  if (pytools_check_array(energy, 1, NPY_FLOAT, "Energy") != SUCCESS)
+  if (pytools_check_array(energy, 1, NPY_FLOAT, "Energy") != SWIFT_SUCCESS)
     return NULL;
 
-  if (pytools_check_array(rho, 1, NPY_FLOAT, "Density") != SUCCESS)
+  if (pytools_check_array(rho, 1, NPY_FLOAT, "Density") != SWIFT_SUCCESS)
     return NULL;
 
   if (fractions != NULL &&
-      pytools_check_array(fractions, 2, NPY_FLOAT, "Fractions") != SUCCESS)
+      pytools_check_array(fractions, 2, NPY_FLOAT, "Fractions") != SWIFT_SUCCESS)
     return NULL;
 
   if (PyArray_DIM(energy, 0) != PyArray_DIM(rho, 0))
@@ -75,24 +117,34 @@ PyObject* pycooling_rate(PyObject* self, PyObject* args) {
   size_t N = PyArray_DIM(energy, 0);
 
 
-  /* Initialize everything */
+  /* Initialize parser */
   struct swift_params params;
   parser_read_file(filename, &params);
 
+  /* Init unit system */
   struct unit_system us;
   units_init_from_params(&us, &params, "InternalUnitSystem");
 
+  /* Init physical constant */
   struct phys_const pconst;
   phys_const_init(&us, &params, &pconst);
 
-  struct cooling_function_data cooling;
-
   /* init cooling */
-  cooling_init_backend(&params, &us, &pconst, &cooling);
-
+  if (initialized == 0) {
+    // struct cooling_function_data cooling;
+    cooling_init_backend(&params, &us, &pconst, &cooling);
+    initialized = 1;
+  }
+  else if (initialized == 1) {
+    message("WARNING: not reinitializing the cooling, need to be fixed");
+    initialized = 2;
+  }
+
+  /* Init chemistry */
   struct chemistry_global_data chemistry;
   chemistry_init_backend(&params, &us, &pconst, &chemistry);
 
+  /* Init cosmo */
   struct cosmology cosmo;
 
   if (with_cosmo)
@@ -100,10 +152,18 @@ PyObject* pycooling_rate(PyObject* self, PyObject* args) {
   else
     cosmology_init_no_cosmo(&cosmo);
 
+  /* Init hydro prop */
+  struct hydro_props hydro_props;
+  hydro_props_init(&hydro_props, &pconst, &us, &params);
+
+  /* Init entropy floor */
+  struct entropy_floor_properties floor_props;
+  entropy_floor_init(&floor_props, &pconst, &us,
+		     &hydro_props, &params);
+
   struct part p;
   struct xpart xp;
 
-  pyswift_part_set_to_zero(&p, &xp);
   hydro_first_init_part(&p, &xp);
     
   /* return object */
@@ -115,27 +175,28 @@ PyObject* pycooling_rate(PyObject* self, PyObject* args) {
       /* set particle data */
       p.rho = *(float*) PyArray_GETPTR1(rho, i);
       float u = *(float*) PyArray_GETPTR1(energy, i);
-      p.entropy = gas_entropy_from_internal_energy(p.rho, u);
+      hydro_set_physical_internal_energy(&p, &xp, &cosmo, u);
 
       if (fractions != NULL)
 	pycooling_set_fractions(&xp, fractions, i);
       else
 	cooling_first_init_part(&pconst, &us, &cosmo, &cooling, &p, &xp);
 
-      chemistry_first_init_part(&pconst, &us, &cosmo, &chemistry, &p, &xp);
-	
+      chemistry_first_init_part(&pconst, &us, &cosmo, &chemistry, &p, &xp);	
+
+      hydro_set_physical_internal_energy_dt(&p, &cosmo, 0);
+
       /* compute cooling rate */
-#ifdef COOLING_GRACKLE
+      cooling_cool_part(&pconst, &us, &cosmo, &hydro_props,
+			&floor_props, &cooling, &p, &xp, dt, dt);
       float *tmp = PyArray_GETPTR1(rate, i);
-      *tmp = cooling_new_energy(&pconst, &us, &cosmo, &cooling, &p, &xp, dt);
-      *tmp = *tmp - hydro_get_physical_internal_energy(&p, &xp, &cosmo);
-      *tmp /= dt;
-#else
-      error("Not implemented");
-      //*tmp = cooling_rate(pconst, us, cosmo, cooling, &p, &xp);
-#endif
+      *tmp = hydro_get_physical_internal_energy_dt(&p, &cosmo);
+
     }
 
+  /* Cleanup */
+  cosmology_clean(&cosmo);
+
   return (PyObject *) rate;
   
 }
diff --git a/src/cooling_wrapper.h b/src/cooling_wrapper.h
index 86841cff6fcfed69a26a4fa89e42bd0ab9d1ab0e..bd9017f8abb4c31084db4ace60fa52372f943006 100644
--- a/src/cooling_wrapper.h
+++ b/src/cooling_wrapper.h
@@ -21,6 +21,11 @@
 
 #include "pyswiftsim_includes.h"
 
+/* A few global variables */
+extern struct cooling_function_data cooling;
+extern int initialized;
+
+
 PyObject* pycooling_rate(PyObject* self, PyObject* args);
 
 #endif // __PYSWIFTSIM_COOLING_H__
diff --git a/src/wrapper.c b/src/libcooling.c
similarity index 67%
rename from src/wrapper.c
rename to src/libcooling.c
index af15b99499ce9bd9ca7ea085da4ad418cf9bc856..6248d437d0ad9519b033c81acf8f80760f54280b 100644
--- a/src/wrapper.c
+++ b/src/libcooling.c
@@ -22,30 +22,38 @@
 #include <math.h>
 #include <numpy/arrayobject.h>
 
-
 /* definition of the method table */      
       
-static PyMethodDef wrapper_methods[] = {
+static PyMethodDef libcooling_methods[] = {
 
   {"coolingRate", pycooling_rate, METH_VARARGS,
    "Compute the cooling rate.\n\n"
    "Parameters\n"
    "----------\n\n"
-   "pyconst: swift physical constant\n"
-   "pyus: swift unit system\n"
-   "cooling: swift cooling structure\n"
-   "rho: np.array\n"
-   "\t Mass density in pyus units\n"
-   "energy: np.array\n"
-   "\t Internal energy in pyus units\n"
+   "filename: str\n"
+   "\t Parameter file\n"
+   "rho: array\n"
+   "\t Mass density in internal units\n"
+   "energy: array\n"
+   "\t Internal energy in internal units\n"
    "dt: float, optional\n"
-   "\t Time step in pyus units\n"
-   "fractions: np.array, optional\n"
+   "\t Time step in internal units\n"
+   "with_cosmo: int, optional\n"
+   "\t Use cosmology?\n"
+   "fractions: array, optional\n"
    "\t Fraction of each cooling element (including metals)\n"
+   "\n"
    "Returns\n"
-   "-------\n\n"
-   "rate: np.array\n"
+   "-------\n"
+   "rate: array\n"
    "\t Cooling rate\n"
+   "\n"
+   "Examples\n"
+   "--------\n"
+   ">>> rho = np.array([1])\n"
+   ">>> u = np.array([1000])\n"
+   ">>> filename = 'params.yml'\n"
+   ">>> rate = coolingRate(filename, rho, u)\n"
   },
 
 
@@ -54,12 +62,12 @@ static PyMethodDef wrapper_methods[] = {
       
       
 
-static struct PyModuleDef wrapper_cmodule = {
+static struct PyModuleDef libcooling_cmodule = {
   PyModuleDef_HEAD_INIT,
-  "wrapper",
-  "Wrapper around the SPH cosmological simulation code SWIFT",
+  "libcooling",
+  "Wrapper around the cooling of SWIFT",
   -1,
-  wrapper_methods,
+  libcooling_methods,
   NULL, /* m_slots */
   NULL, /* m_traverse */
   NULL, /* m_clear */
@@ -67,7 +75,7 @@ static struct PyModuleDef wrapper_cmodule = {
 };
 
 
-PyMODINIT_FUNC PyInit_libpyswiftsim(void)
+PyMODINIT_FUNC PyInit_libcooling(void)
 {
   PyObject *m;
 
@@ -76,7 +84,7 @@ PyMODINIT_FUNC PyInit_libpyswiftsim(void)
   
   import_array();
   Py_Initialize();
-  m = PyModule_Create(&wrapper_cmodule);
+  m = PyModule_Create(&libcooling_cmodule);
 
   return m;
 }
diff --git a/src/pyswiftsim_tools.c b/src/pyswiftsim_tools.c
index 4fd232542387b913e7604f99cb95eb18ae3f8e70..e9c8f3d3196cbca65265e7f9d62073d902c31529 100644
--- a/src/pyswiftsim_tools.c
+++ b/src/pyswiftsim_tools.c
@@ -35,21 +35,21 @@ int pytools_check_array(PyArrayObject *obj, int dim, int type, const char* name)
   /* check if array */
   if (!PyArray_Check(obj))
     {
-      pyerror("Expecting a numpy array for %s", name);
+      error("Expecting a numpy array for %s", name);
     }
 
   /* check if required dim */
   if (PyArray_NDIM(obj) != dim)
     {
-      pyerror("Array should be a %i dimensional object for %s", dim, name);
+      error("Array should be a %i dimensional object for %s", dim, name);
     }
 
   /* check data type */
   if (PyArray_TYPE(obj) != type)
     {
-      pyerror("Wrong array type for %s", name);
+      error("Wrong array type for %s", name);
     }
 
-  return SUCCESS;
+  return SWIFT_SUCCESS;
 
 }
diff --git a/src/pyswiftsim_tools.h b/src/pyswiftsim_tools.h
index e3433480c7b65b4aa51566d1af37c52cb9682f63..567314f028f17bf34e395ea77292763d1373447e 100644
--- a/src/pyswiftsim_tools.h
+++ b/src/pyswiftsim_tools.h
@@ -36,8 +36,8 @@
 
 /* error code in pyswiftsim */
 enum error_code {
-  FAIL = 0, // ensure NULL == FAIL
-  SUCCESS,
+  SWIFT_FAIL,
+  SWIFT_SUCCESS,
 };
 
 
diff --git a/test/getCoolingTable.sh b/test/getCoolingTable.sh
deleted file mode 100644
index 809958bfc236e5ab0f7be924c62789fa13b3ac29..0000000000000000000000000000000000000000
--- a/test/getCoolingTable.sh
+++ /dev/null
@@ -1,2 +0,0 @@
-#!/bin/bash
-wget http://virgodb.cosma.dur.ac.uk/swift-webstorage/CoolingTables/CloudyData_UVB=HM2012.h5
diff --git a/test/test_cooling.py b/test/test_cooling.py
index 8623c513f2e45cbabdd0d229e1f449f81bf52225..94e3a46917f74659c8b8502f2922d316f14d5c41 100644
--- a/test/test_cooling.py
+++ b/test/test_cooling.py
@@ -1,6 +1,24 @@
 #!/usr/bin/env python3
+# ########################################################################
+# This file is part of PYSWIFT.
+# Copyright (c) 2018 Loic Hausammann (loic.hausammann@epfl.ch)
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+# ########################################################################
 
-from pyswiftsim.wrapper import coolingRate
+from pyswiftsim.libcooling import coolingRate
+from pyswiftsim import downloadCoolingTable
 import numpy as np
 from astropy import units, constants
 import yaml
@@ -10,6 +28,9 @@ filename = "test_cooling.yml"
 # adiabatic index
 gamma = 5. / 3.
 
+# download cooling table
+downloadCoolingTable()
+
 # Read parameter file
 f = open(filename, "r")
 data = yaml.load(f)
@@ -29,29 +50,35 @@ u_t = u_l / u_v
 u_e = u_m * u_l**2 / u_t**2
 
 # boltzman constant
-k_B = constants.k_B.to("{} J / K".format(u_e))
+k_B = constants.k_B.to("erg / K") * units.K / (units.erg * u_e)
+k_B = float(k_B)
 # proton mass
-m_p = constants.m_p.to("{} g".format(u_l))
+m_p = constants.m_p.to("g") / (units.g * u_m)
+m_p = float(m_p)
 
 
 def internalEnergy(T):
-    u = k_B * T
-    u /= (gamma - 1.) * m_p
+    u = (k_B / m_p) * T
+    u /= (gamma - 1.)
     return u
 
 
 # generates data
+# time step in Myr
+dt = 1e-3
+dt = float(units.Myr.to('s')) / u_t
+
 # density in atom / cm3
 density = np.array([1.], dtype=np.float32)
-density *= float(constants.m_p.to("g") / units.g) / u_m
-density *= u_l**3
+density *= m_p * u_l**3
 
 # Temperature in K
-T = np.array([100], dtype=np.float32)
+T = np.array([1e4], dtype=np.float32)
 u = internalEnergy(T)
 
 # compute the cooling rate
-rate = coolingRate(filename, density, u)
+with_cosmo = False
+rate = coolingRate(filename, density, u, with_cosmo, dt)
 
 # go back to cgs
 rate /= u_e / u_t
diff --git a/test/test_cooling.yml b/test/test_cooling.yml
index 067bf3acb227d754a856b488cb15fa4aaaa1ab48..a30e9edf8a5fa99564e288e5029c038af1fb2b32 100644
--- a/test/test_cooling.yml
+++ b/test/test_cooling.yml
@@ -1,8 +1,8 @@
-# Define the system of units to use internally.
+# Define the system of units to use internally. 
 InternalUnitSystem:
-  UnitMass_in_cgs:     1000   # Grams
-  UnitLength_in_cgs:   1   # Centimeters
-  UnitVelocity_in_cgs: 1   # Centimeters per second
+  UnitMass_in_cgs:     1.989e43   # Grams
+  UnitLength_in_cgs:   3.085e21   # Centimeters
+  UnitVelocity_in_cgs: 1e5   # Centimeters per second
   UnitCurrent_in_cgs:  1   # Amperes
   UnitTemp_in_cgs:     1   # Kelvin
 
@@ -18,3 +18,8 @@ GrackleCooling:
   OutputMode: 0                         # (optional) Write in output corresponding primordial chemistry mode
   MaxSteps: 10000                       # (optional) Max number of step when computing the initial composition
   ConvergenceLimit: 1e-2                # (optional) Convergence threshold (relative) for initial composition
+
+# Parameters for the hydrodynamics scheme
+SPH:
+  resolution_eta:        1.2348   # Target smoothing length in units of the mean inter-particle separation (1.2348 == 48Ngbs with the cubic spline kernel).
+  CFL_condition:         0.1      # Courant-Friedrich-Levy condition for time integration.