diff --git a/pyswiftsim/structure.py b/pyswiftsim/structure.py
index f42e65e8a1ef0014a6183d68ddb9393a05867121..2e89c0c5e6a3427d0ec38b0ede0d8903160e4376 100644
--- a/pyswiftsim/structure.py
+++ b/pyswiftsim/structure.py
@@ -283,28 +283,45 @@ class UnitSystem(SwiftStruct):
         super().__init__(self.struct_format, data, parent)
 
 
+######################################################################
+#                                                                    #
+#                     chemistry_part_data                            #
+#                                                                    #
+######################################################################
+class ChemistryPartData(SwiftStruct):
+    _format = "f"
+    _name = [
+        "he_density"
+    ]
+
+    def __init__(self, data, parent=None):
+        super().__init__(self.struct_format, data, parent)
+
 ######################################################################
 #                                                                    #
 #                          Part                                      #
 #                                                                    #
 ######################################################################
 class Part(SwiftStruct):
-    _format = "qP3d3f3ffffffN7f4c"
+    _format = "qP3d3f3ffffffN7f4{chemistry}sc".format(
+        chemistry=struct.calcsize(ChemistryPartData._format)
+    )
     _name = [
-            "id",
-            "gpart",
-            "pos",
-            "vel",
-            "a_hydro",
-            "h",
-            "mass",
-            "rho",
-            "entropy",
-            "entropy_dt",
-            "last_offset",
-            "density",
-            "time_bin"
-        ]
+        "id",
+        "gpart",
+        "pos",
+        "vel",
+        "a_hydro",
+        "h",
+        "mass",
+        "rho",
+        "entropy",
+        "entropy_dt",
+        "last_offset",
+        "density",
+        "chemistry_data",
+        "time_bin"
+    ]
 
     def __init__(self, data, parent=None):
         super().__init__(self.struct_format, data, parent)
@@ -416,7 +433,7 @@ class PhysConst(SwiftStruct):
         super().__init__(self.struct_format, data, parent)
 
 
-class GrackleUnits(SwiftStruct):
+class GrackleCodeUnits(SwiftStruct):
     _format = "idddddd"
     _name = [
         "comoving_coordinates",
@@ -493,17 +510,17 @@ class CoolingFunctionData(SwiftStruct):
             "cooling_tstep_mult"
         ]
     elif cooling_type == "grackle":
-        _format = "200cidd{units}s{chemistry}s".format(
-            units=struct.calcsize(GrackleUnits._format),
+        _format = "200cidd{code_units}s{chemistry}s".format(
+            code_units=struct.calcsize(GrackleCodeUnits._format),
             chemistry=struct.calcsize(GrackleChemistryData._format)
         )
         _name = [
-            "cloudy_table",
-            "uv_background",
-            "redshift",
-            "density_self_shielding",
-            "units",
-            "chemistry"
+            "GrackleCloudyTable",
+            "UVbackground",
+            "GrackleRedshift",
+            "GrackleHSShieldingDensityThreshold",
+            "code_units",
+            "chemistry",
         ]
 
         @property
@@ -514,7 +531,7 @@ class CoolingFunctionData(SwiftStruct):
             }
 
             units = {
-                "class": GrackleUnits,
+                "class": GrackleCodeUnits,
                 "size": 1
             }
             return {
diff --git a/setup.py b/setup.py
index 3efb6897a93e83a649044df5aab442b22b33769d..b87b13a808514f6cf666837b5570337f8a721e28 100644
--- a/setup.py
+++ b/setup.py
@@ -12,6 +12,17 @@ import numpy
 
 os.environ["CC"] = "mpicc"
 
+cflags = ["-Werror",
+          "-Wall",
+          "-Wextra",
+          # disables some warnings due to python
+          "-Wno-unused-parameter",
+          "-Wno-strict-prototypes",
+          "-Wno-unused-function",
+          "-Wno-incompatible-pointer-types",
+          "-Wno-missing-field-initializers",
+]
+
 # deal with arguments
 def parseCmdLine(arg, store=False):
     ret = False
@@ -96,7 +107,8 @@ ext_modules = Extension("pyswiftsim.wrapper",
                         c_src,
                         include_dirs=include,
                         libraries=lib,
-                        library_dirs=lib_dir)
+                        library_dirs=lib_dir,
+                        extra_compile_args=cflags)
 
 ext_modules = [ext_modules]
     
diff --git a/src/config_wrapper.h b/src/config_wrapper.h
index bc84738cd80c204ab003f74da0539bd6a19156c6..6b9b2a8e752bba7826b699d76050e946c57c4f1c 100644
--- a/src/config_wrapper.h
+++ b/src/config_wrapper.h
@@ -3,8 +3,6 @@
 
 #include "pyswiftsim_tools.h"
 
-#include <config.h>
-
 /**
  * @brief give the cooling type name
  *
diff --git a/src/cooling_wrapper.c b/src/cooling_wrapper.c
index 59fd431e911caf87677e2ea826e6d3f30ec3f235..cb23170b544e455b181f8007277b7d55ace4ffa8 100644
--- a/src/cooling_wrapper.c
+++ b/src/cooling_wrapper.c
@@ -1,13 +1,6 @@
 #include "pyswiftsim_tools.h"
 #include "cooling_wrapper.h"
 
-#include <cooling.h>
-#include <cooling_struct.h>
-#include <equation_of_state.h>
-
-
-#include <numpy/arrayobject.h>
-
 
 PyObject* pycooling_init(PyObject* self, PyObject* args) {
   PyObject *pyparams;
@@ -43,7 +36,7 @@ PyObject* pycooling_init(PyObject* self, PyObject* args) {
 }
 
 PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) {
-  IMPORT_ARRAY();
+  import_array();
   
   PyObject *pycooling;
   PyObject *pyus;
@@ -98,8 +91,10 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) {
     return NULL;
 
   struct part p;
+  struct xpart xp;
 
 #ifdef COOLING_GRACKLE
+  /* set grackle_data */
   grackle_data = &cooling->chemistry;
 #endif
 
@@ -114,12 +109,14 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) {
       float u = *(float*) PyArray_GETPTR1(energy, i);
       p.entropy = gas_entropy_from_internal_energy(p.rho, u);
 
+      cooling_first_init_part(&p, &xp, cooling);
+
       /* compute cooling rate */
       float *tmp = PyArray_GETPTR1(rate, i);
 #ifdef COOLING_GRACKLE
-      *tmp = cooling_rate(pconst, us, cooling, &p, dt);
+      *tmp = cooling_rate(pconst, us, cooling, &p, &xp, dt);
 #else
-      *tmp = cooling_rate(pconst, us, cooling, &p);
+      *tmp = cooling_rate(pconst, us, cooling, &p, &xp);
 #endif
     }
 
diff --git a/src/parser_wrapper.c b/src/parser_wrapper.c
index ef164f58389ef5d3d2ecbce0de2e4aec081b2736..3bf033f208428b2480f796ccdf415f7a09f8a2b7 100644
--- a/src/parser_wrapper.c
+++ b/src/parser_wrapper.c
@@ -1,7 +1,6 @@
 #include "parser_wrapper.h"
 #include "pyswiftsim_tools.h"
 
-#include <parser.h>
 
 PyObject* pyparser_read_file(PyObject *self, PyObject *args)
 {
diff --git a/src/part_wrapper.c b/src/part_wrapper.c
index a4765eecae376cd000fb4c4850bee4b420cd8b0a..76d4512ae235c880ecd52d4c993f58f7e8ac1dbd 100644
--- a/src/part_wrapper.c
+++ b/src/part_wrapper.c
@@ -1,8 +1,5 @@
 #include "pyswiftsim_tools.h"
 
-#include <part.h>
-#include <hydro.h>
-
 #include <Python.h>
 #include <stdlib.h>
 #include <string.h>
diff --git a/src/part_wrapper.h b/src/part_wrapper.h
index 58a426fa48a57c6c88dc591df01754ffe8ea8e91..5875f59da06cc9aba5ca9605a3ac4b377cc1a673 100644
--- a/src/part_wrapper.h
+++ b/src/part_wrapper.h
@@ -2,7 +2,6 @@
 #define __PYSWIFTSIM_PART_H__
 
 #include <Python.h>
-#include <part.h>
 
 /**
  * @brief Test function for the struct
diff --git a/src/pyswiftsim_tools.c b/src/pyswiftsim_tools.c
index 41814598be72886abeb35ca6a750fb59b3748a1e..b8673de8a40465ef1806f847c9381b04bbd8dcff 100644
--- a/src/pyswiftsim_tools.c
+++ b/src/pyswiftsim_tools.c
@@ -1,15 +1,5 @@
 #include "pyswiftsim_tools.h"
 
-/* include swift */
-#include <part.h>
-#include <units.h>
-#include <parser.h>
-#include <physical_constants.h>
-#include <cooling_struct.h>
-
-#include <Python.h>
-#include <numpy/arrayobject.h>
-
 
 const size_t class_size[class_count] = {
   sizeof(struct unit_system),
@@ -19,7 +9,7 @@ const size_t class_size[class_count] = {
   sizeof(struct cooling_function_data)
 };
   
-const char *class_name[class_count] = {
+char *class_name[class_count] = {
   "UnitSystem",
   "Part",
   "SwiftParams",
@@ -69,7 +59,6 @@ PyObject* pytools_return(void *p, int class)
   size_t nber_bytes;
 
   char module_name[STRING_SIZE] = "pyswiftsim.structure";
-  char *class_pyname;
 
   /* check class */
   if (class >= class_count)
@@ -77,7 +66,7 @@ PyObject* pytools_return(void *p, int class)
 
   /* get class information */
   nber_bytes = class_size[class];
-  class_pyname = class_name[class];
+  char *class_pyname = class_name[class];
 
   /* import python class */
   python_class = pytools_import(module_name, class_pyname);
@@ -126,8 +115,8 @@ char* pytools_get_type_name(PyObject *obj)
     }
 
   /* transform to C */
-  size_t size;
-  char *name = PyUnicode_AsUTF8AndSize(recv, size);
+  Py_ssize_t size;
+  char *name = PyUnicode_AsUTF8AndSize(recv, &size);
   Py_DECREF(recv);
 
   if (name == NULL)
@@ -155,9 +144,9 @@ char* pytools_construct(PyObject* obj, int class)
   PyObject *pyclass = pytools_import(module_name, class_pyname);
 
   /* check if classes correspond */
-  int test = !PyObject_IsInstance(obj, pyclass);
+  int test = PyObject_IsInstance(obj, pyclass);
   Py_DECREF(pyclass);
-  if (test)
+  if (!test)
     {
       char *recv = pytools_get_type_name(obj);
       if (recv == NULL)
@@ -181,8 +170,7 @@ char* pytools_construct(PyObject* obj, int class)
 
 int pytools_check_array(PyArrayObject *obj, int dim, int type)
 {
-  /* ensure to have numpy arrays */
-  IMPORT_ARRAY();
+  IMPORT_ARRAY1(1);
 
   /* check if array */
   if (!PyArray_Check(obj))
diff --git a/src/pyswiftsim_tools.h b/src/pyswiftsim_tools.h
index 9d03af9f77ade6f241bcf93c2252464fc9ae0e10..65cb3df581181723e2df87a3d5ff42f79a27059d 100644
--- a/src/pyswiftsim_tools.h
+++ b/src/pyswiftsim_tools.h
@@ -6,18 +6,12 @@
 #include <Python.h>
 #include <numpy/arrayobject.h>
 
+#undef assert
+#include <swift.h>
+
 #define DIM 3
 #define STRING_SIZE 200
 
-/* Check if arrays are imported and if not import them */
-#define IMPORT_ARRAY()				\
-  ({						\
-    if (PyArray_API == NULL)			\
-      {						\
-	import_array();				\
-      }						\
-  })
-
 /* Set the error message for python (still need to return NULL) */
 #define pyerror(s, ...)						\
   ({								\
@@ -37,6 +31,17 @@
 	    __FUNCTION__, __LINE__, ##__VA_ARGS__);		\
   })
 
+#define IMPORT_ARRAY1(ret)			\
+  {						\
+    if (PyArray_API == NULL)			\
+      {						\
+	import_array1(ret);			\
+      }						\
+  }
+
+#define IMPORT_ARRAY() IMPORT_ARRAY1(NUMPY_IMPORT_ARRAY_RETVAL);
+	
+
 /* Enum swift classes */
 enum swift_class {
   class_unit_system,
@@ -50,7 +55,7 @@ enum swift_class {
 /* size of each structure in enum swift_class */
 extern const size_t class_size[];
 /* name of each Python class representing a swift class */
-extern const char *class_name[];
+extern char *class_name[];
 
 /* error code in pyswiftsim */
 enum error_code {
diff --git a/src/units_wrapper.c b/src/units_wrapper.c
index aeda170ac6f6e74e55ce10e0c11c637ae3621778..fe3a23a0cac0ed011cbaf9417d77c0d933e4721c 100644
--- a/src/units_wrapper.c
+++ b/src/units_wrapper.c
@@ -1,9 +1,5 @@
 #include "pyswiftsim_tools.h"
 
-#include <units.h>
-#include <parser.h>
-#include <physical_constants.h>
-
 #include <Python.h>
 #include <stdlib.h>
 #include <string.h>
@@ -38,7 +34,7 @@ PyObject* pyunit_system_init(PyObject *self, PyObject *args)
   if (!PyArg_ParseTuple(args, "O", &parser))
     return NULL;
 
-  struct swift_params *params = pytools_construct(parser, class_swift_params);
+  struct swift_params *params = (struct swift_params*) pytools_construct(parser, class_swift_params);
 
   if (params == NULL)
     return NULL;
diff --git a/src/wrapper.c b/src/wrapper.c
index 27fc59183559641a6b0d3fbfa1e59011b16eae92..12e7ef2b12653c9648e91bdf1c727cddb6993395 100644
--- a/src/wrapper.c
+++ b/src/wrapper.c
@@ -5,7 +5,6 @@
 #include "pyswiftsim_tools.h"
 #include "config_wrapper.h"
 
-
 #include <Python.h>
 #include <math.h>
 #include <numpy/arrayobject.h>
diff --git a/test/test_cooling.py b/test/test_cooling.py
index 43148ce9a7e875bf9a747afc8e7847d605a3c845..f8bb2aeca7e0c7a2e9e71cf209681711c4aff66a 100644
--- a/test/test_cooling.py
+++ b/test/test_cooling.py
@@ -19,8 +19,8 @@ gamma = 5. / 3.
 filename = "test/cooling.yml"
 
 # number of points
-N_rho = 300
-N_T = 300
+N_rho = 1
+N_T = 100
 
 # density in atom / cm3
 if N_rho == 1:
@@ -49,10 +49,13 @@ T = T.reshape(-1)
 #
 
 # parse swift params
+print("Reading parameters")
 params = wrapper.parserReadFile(filename)
 # init units
+print("Initialization of the unit system")
 us, pconst = wrapper.unitSystemInit(params)
 # init cooling
+print("Initialization of the cooling")
 cooling = wrapper.coolingInit(params, us, pconst)
 
 
@@ -74,10 +77,12 @@ dt /= us.UnitTime_in_cgs
 #
 
 # du / dt
+print("Computing cooling...")
 rate = wrapper.coolingRate(pconst, us, cooling,
                            rho.astype(np.float32),
                            energy.astype(np.float32),
                            dt)
+print("Computing done")
 
 #
 # plot