diff --git a/examples/Cooling/CoolingBox/coolingBox.yml b/examples/Cooling/CoolingBox/coolingBox.yml
index 97b452e5ea376ab10bba155e0ff3e9acb8ed02bd..853e480cd4ffba4baa659232e6d5f068b4ea2815 100644
--- a/examples/Cooling/CoolingBox/coolingBox.yml
+++ b/examples/Cooling/CoolingBox/coolingBox.yml
@@ -48,10 +48,12 @@ GrackleCooling:
   ProvideVolumetricHeatingRates: 0 # User provide volumetric heating rates
   ProvideSpecificHeatingRates: 0 # User provide specific heating rates
   SelfShieldingMethod: 0 # Grackle (<= 3) or Gear self shielding method
-  OutputMode: 1 # Write in output corresponding primordial chemistry mode
   MaxSteps: 1000
   ConvergenceLimit: 1e-2
   
+GearChemistry:
+  InitialMetallicity: 0.01295
+
 EAGLECooling:
   dir_name:                ./coolingtables/
   H_reion_z:               11.5
@@ -71,5 +73,3 @@ EAGLEChemistry:             # Solar abundances
   init_abundance_Silicon:   6.825874e-4
   init_abundance_Iron:      1.1032152e-3
 
-GearChemistry:
-  InitialMetallicity: 0.01295
diff --git a/examples/Cooling/CoolingBox/energy_plot.py b/examples/Cooling/CoolingBox/energy_plot.py
deleted file mode 100644
index 45f0b4f6b11c3855a919f6a98fd0ca006a887f82..0000000000000000000000000000000000000000
--- a/examples/Cooling/CoolingBox/energy_plot.py
+++ /dev/null
@@ -1,128 +0,0 @@
-import matplotlib
-matplotlib.use("Agg")
-from pylab import *
-import h5py
-
-# Plot parameters
-params = {'axes.labelsize': 10,
-'axes.titlesize': 10,
-'font.size': 12,
-'legend.fontsize': 12,
-'xtick.labelsize': 10,
-'ytick.labelsize': 10,
-'text.usetex': True,
- 'figure.figsize' : (3.15,3.15),
-'figure.subplot.left'    : 0.145,
-'figure.subplot.right'   : 0.99,
-'figure.subplot.bottom'  : 0.11,
-'figure.subplot.top'     : 0.99,
-'figure.subplot.wspace'  : 0.15,
-'figure.subplot.hspace'  : 0.12,
-'lines.markersize' : 6,
-'lines.linewidth' : 3.,
-'text.latex.unicode': True
-}
-rcParams.update(params)
-rc('font',**{'family':'sans-serif','sans-serif':['Times']})
-
-
-import numpy as np
-import h5py as h5
-import sys
-
-# File containing the total energy
-stats_filename = "./energy.txt"
-
-# First snapshot
-snap_filename = "coolingBox_0000.hdf5"
-
-# Some constants in cgs units
-k_b = 1.38E-16 #boltzmann
-m_p = 1.67e-24 #proton mass
-
-# Initial conditions set in makeIC.py
-T_init = 1.0e5
-
-# Read the initial state of the gas
-f = h5.File(snap_filename,'r')
-rho = np.mean(f["/PartType0/Density"])
-pressure = np.mean(f["/PartType0/Pressure"])
-
-# Read the units parameters from the snapshot
-units = f["InternalCodeUnits"]
-unit_mass = units.attrs["Unit mass in cgs (U_M)"]
-unit_length = units.attrs["Unit length in cgs (U_L)"]
-unit_time = units.attrs["Unit time in cgs (U_t)"]
-
-# Read the properties of the cooling function
-parameters = f["Parameters"]
-cooling_lambda = float(parameters.attrs["LambdaCooling:lambda_cgs"])
-min_T = float(parameters.attrs["LambdaCooling:minimum_temperature"])
-mu = float(parameters.attrs["LambdaCooling:mean_molecular_weight"])
-X_H = float(parameters.attrs["LambdaCooling:hydrogen_mass_abundance"])
-
-# Read the adiabatic index
-gamma = float(f["HydroScheme"].attrs["Adiabatic index"])
-
-print "Initial density :", rho
-print "Initial pressure:", pressure
-print "Adiabatic index :", gamma
-
-# Read energy and time arrays
-array = np.genfromtxt(stats_filename,skip_header = 1)
-time = array[:,0]
-total_mass = array[:,1]
-total_energy = array[:,2]
-kinetic_energy = array[:,3]
-internal_energy = array[:,4]
-radiated_energy = array[:,8]
-initial_energy = total_energy[0]
-
-# Conversions to cgs
-rho_cgs = rho * unit_mass / (unit_length)**3
-time_cgs = time * unit_time
-total_energy_cgs = total_energy / total_mass[0] * unit_length**2 / (unit_time)**2
-kinetic_energy_cgs = kinetic_energy / total_mass[0] * unit_length**2 / (unit_time)**2
-internal_energy_cgs = internal_energy / total_mass[0] * unit_length**2 / (unit_time)**2
-radiated_energy_cgs = radiated_energy / total_mass[0] * unit_length**2 / (unit_time)**2  
-
-# Find the energy floor
-u_floor_cgs = k_b * min_T / (mu * m_p * (gamma - 1.))
-
-# Find analytic solution
-initial_energy_cgs = initial_energy/total_mass[0] * unit_length**2 / (unit_time)**2 
-n_H_cgs = X_H * rho_cgs / m_p
-du_dt_cgs = -cooling_lambda * n_H_cgs**2 / rho_cgs
-cooling_time_cgs = (initial_energy_cgs/(-du_dt_cgs))[0]
-analytic_time_cgs = np.linspace(0, cooling_time_cgs * 1.8, 1000)
-u_analytic_cgs = du_dt_cgs*analytic_time_cgs + initial_energy_cgs
-u_analytic_cgs[u_analytic_cgs < u_floor_cgs] = u_floor_cgs
-
-print "Cooling time:", cooling_time_cgs, "[s]"
-
-# Read snapshots
-u_snapshots_cgs = zeros(25)
-t_snapshots_cgs = zeros(25)
-for i in range(25):
-    snap = h5.File("coolingBox_%0.4d.hdf5"%i,'r')
-    u_snapshots_cgs[i] = sum(snap["/PartType0/InternalEnergy"][:] * snap["/PartType0/Masses"][:])  / total_mass[0] * unit_length**2 / (unit_time)**2
-    t_snapshots_cgs[i] = snap["/Header"].attrs["Time"] * unit_time
-
-
-figure()
-plot(time_cgs, total_energy_cgs, 'r-', lw=1.6, label="Gas total energy")
-plot(t_snapshots_cgs, u_snapshots_cgs, 'rD', ms=3)
-plot(time_cgs, radiated_energy_cgs, 'g-', lw=1.6, label="Radiated energy")
-plot(time_cgs, total_energy_cgs + radiated_energy_cgs, 'b-', lw=0.6, label="Gas total + radiated")
-
-plot(analytic_time_cgs, u_analytic_cgs, '--', color='k', alpha=0.8, lw=1.0, label="Analytic solution")
-
-legend(loc="upper right", fontsize=8, frameon=False, handlelength=3, ncol=1)
-xlabel("${\\rm{Time~[s]}}$", labelpad=0)
-ylabel("${\\rm{Energy~[erg]}}$")
-xlim(0, 1.5*cooling_time_cgs)
-ylim(0, 1.5*u_analytic_cgs[0])
-
-savefig("energy.png", dpi=200)
-
-
diff --git a/examples/Cooling/CoolingBox/makeIC.py b/examples/Cooling/CoolingBox/makeIC.py
index f863e174b1fcd404ae178fe324c7a165598b4af0..807ad5378d87891ed878a85c1020d5983474e13d 100644
--- a/examples/Cooling/CoolingBox/makeIC.py
+++ b/examples/Cooling/CoolingBox/makeIC.py
@@ -1,58 +1,69 @@
 ###############################################################################
- # This file is part of SWIFT.
- # Copyright (c) 2016 Stefan Arridge (stefan.arridge@durhama.ac.uk)
- #                    Matthieu Schaller (matthieu.schaller@durham.ac.uk)
- # 
- # 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 file is part of SWIFT.
+# Copyright (c) 2016 Stefan Arridge (stefan.arridge@durhama.ac.uk)
+#                    Matthieu Schaller (matthieu.schaller@durham.ac.uk)
+#
+# 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/>.
+#
+##############################################################################
 
 import h5py
-import sys
-from numpy import *
+import numpy as np
 
 # Generates a SWIFT IC file with a constant density and pressure
 
 # Parameters
-periodic= 1           # 1 For periodic box
-boxSize = 1           # 1 kiloparsec    
-rho = 3.2e3           # Density in code units (3.2e6 is 0.1 hydrogen atoms per cm^3)
-P = 4.5e6             # Pressure in code units (at 10^5K)
-gamma = 5./3.         # Gas adiabatic index
-eta = 1.2349          # 48 ngbs with cubic spline kernel
-fileName = "coolingBox.hdf5" 
-
-#---------------------------------------------------
+periodic = 1  # 1 For periodic box
+boxSize = 1  # 1 kiloparsec
+rho = 3.2e3  # Density in code units (3.2e6 is 0.1 hydrogen atoms per cm^3)
+T = 4e3  # Initial Temperature
+gamma = 5./3.  # Gas adiabatic index
+fileName = "coolingBox.hdf5"
+# ---------------------------------------------------
+
+# defines some constants
+# need to be changed in plotTemperature.py too
+h_frac = 0.76
+mu = 4. / (1. + 3. * h_frac)
+
+m_h_cgs = 1.67e-24
+k_b_cgs = 1.38e-16
+
+# defines units
+unit_length = 3.0857e21  # kpc
+unit_mass = 2.0e33  # solar mass
+unit_time = 3.0857e16  # ~ Gyr
 
 # Read id, position and h from glass
 glass = h5py.File("glassCube_32.hdf5", "r")
 ids = glass["/PartType0/ParticleIDs"][:]
-pos = glass["/PartType0/Coordinates"][:,:] * boxSize
+pos = glass["/PartType0/Coordinates"][:, :] * boxSize
 h = glass["/PartType0/SmoothingLength"][:] * boxSize
 
 # Compute basic properties
-numPart = size(pos) / 3
+numPart = np.size(pos) // 3
 mass = boxSize**3 * rho / numPart
-internalEnergy = P / ((gamma - 1.) * rho)
+internalEnergy = k_b_cgs * T * mu / ((gamma - 1.) * m_h_cgs)
+internalEnergy *= (unit_time / unit_length)**2
 
-#File
-file = h5py.File(fileName, 'w')
+# File
+f = h5py.File(fileName, 'w')
 
 # Header
-grp = file.create_group("/Header")
+grp = f.create_group("/Header")
 grp.attrs["BoxSize"] = boxSize
-grp.attrs["NumPart_Total"] =  [numPart, 0, 0, 0, 0, 0]
+grp.attrs["NumPart_Total"] = [numPart, 0, 0, 0, 0, 0]
 grp.attrs["NumPart_Total_HighWord"] = [0, 0, 0, 0, 0, 0]
 grp.attrs["NumPart_ThisFile"] = [numPart, 0, 0, 0, 0, 0]
 grp.attrs["Time"] = 0.0
@@ -61,43 +72,43 @@ grp.attrs["MassTable"] = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
 grp.attrs["Flag_Entropy_ICs"] = 0
 
 # Runtime parameters
-grp = file.create_group("/RuntimePars")
+grp = f.create_group("/RuntimePars")
 grp.attrs["PeriodicBoundariesOn"] = periodic
 
 # Units
-grp = file.create_group("/Units")
-grp.attrs["Unit length in cgs (U_L)"] = 3.0857e21 
-grp.attrs["Unit mass in cgs (U_M)"] = 2.0e33 
-grp.attrs["Unit time in cgs (U_t)"] = 3.0857e16 
+grp = f.create_group("/Units")
+grp.attrs["Unit length in cgs (U_L)"] = unit_length
+grp.attrs["Unit mass in cgs (U_M)"] = unit_mass
+grp.attrs["Unit time in cgs (U_t)"] = unit_time
 grp.attrs["Unit current in cgs (U_I)"] = 1.
 grp.attrs["Unit temperature in cgs (U_T)"] = 1.
 
-#Particle group
-grp = file.create_group("/PartType0")
+# Particle group
+grp = f.create_group("/PartType0")
 
-v  = zeros((numPart, 3))
+v = np.zeros((numPart, 3))
 ds = grp.create_dataset('Velocities', (numPart, 3), 'f')
 ds[()] = v
 
-m = full((numPart, 1), mass)
-ds = grp.create_dataset('Masses', (numPart,1), 'f')
+m = np.full((numPart, 1), mass)
+ds = grp.create_dataset('Masses', (numPart, 1), 'f')
 ds[()] = m
 
-h = reshape(h, (numPart, 1))
+h = np.reshape(h, (numPart, 1))
 ds = grp.create_dataset('SmoothingLength', (numPart, 1), 'f')
 ds[()] = h
 
-u = full((numPart, 1), internalEnergy)
-ds = grp.create_dataset('InternalEnergy', (numPart,1), 'f')
+u = np.full((numPart, 1), internalEnergy)
+ds = grp.create_dataset('InternalEnergy', (numPart, 1), 'f')
 ds[()] = u
 
-ids = reshape(ids, (numPart, 1))
+ids = np.reshape(ids, (numPart, 1))
 ds = grp.create_dataset('ParticleIDs', (numPart, 1), 'L')
 ds[()] = ids
 
 ds = grp.create_dataset('Coordinates', (numPart, 3), 'd')
 ds[()] = pos
 
-file.close()
+f.close()
 
-print numPart
+print("Initial condition generated")
diff --git a/examples/Cooling/CoolingBox/plotTemperature.py b/examples/Cooling/CoolingBox/plotTemperature.py
new file mode 100644
index 0000000000000000000000000000000000000000..3dabee2c0add5d6370c60a39b745f8a73cf9fc8d
--- /dev/null
+++ b/examples/Cooling/CoolingBox/plotTemperature.py
@@ -0,0 +1,111 @@
+from h5py import File
+import numpy as np
+import matplotlib
+# matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+# Plot parameters
+params = {
+    'axes.labelsize': 10,
+    'axes.titlesize': 10,
+    'font.size': 12,
+    'legend.fontsize': 12,
+    'xtick.labelsize': 10,
+    'ytick.labelsize': 10,
+    'text.usetex': True,
+    'figure.figsize': (3.15, 3.15),
+    'figure.subplot.left': 0.145,
+    'figure.subplot.right': 0.99,
+    'figure.subplot.bottom': 0.11,
+    'figure.subplot.top': 0.99,
+    'figure.subplot.wspace': 0.15,
+    'figure.subplot.hspace': 0.12,
+    'lines.markersize': 6,
+    'lines.linewidth': 3.,
+}
+plt.rcParams.update(params)
+
+
+# Some constants in cgs units
+k_b_cgs = 1.38e-16  # boltzmann
+m_h_cgs = 1.67e-24  # proton mass
+# need to be changed in makeIC.py too
+h_frac = 0.76
+mu = 4. / (1. + 3. * h_frac)
+
+
+# File containing the total energy
+stats_filename = "./energy.txt"
+
+# First snapshot
+snap_filename = "coolingBox_0000.hdf5"
+
+# Read the initial state of the gas
+f = File(snap_filename, 'r')
+
+# Read the units parameters from the snapshot
+units = f["InternalCodeUnits"]
+unit_mass = units.attrs["Unit mass in cgs (U_M)"]
+unit_length = units.attrs["Unit length in cgs (U_L)"]
+unit_time = units.attrs["Unit time in cgs (U_t)"]
+
+# Read the adiabatic index
+gamma = float(f["HydroScheme"].attrs["Adiabatic index"])
+
+
+def Temperature(u):
+    """ Compute the temperature from the internal energy. """
+    u *= (unit_length / unit_time)**2
+    return u * (gamma - 1.) * m_h_cgs / (mu * k_b_cgs)
+
+
+# Read energy and time arrays
+array = np.genfromtxt(stats_filename, skip_header=1)
+time = array[:, 0] * unit_time
+total_mass = array[:, 1]
+total_energy = array[:, 2]
+kinetic_energy = array[:, 3]
+internal_energy = array[:, 4]
+radiated_energy = array[:, 8]
+initial_energy = total_energy[0]
+
+# Conversions to cgs
+total_energy_cgs = total_energy / total_mass[0]
+total_energy_cgs = Temperature(total_energy_cgs)
+
+kinetic_energy_cgs = kinetic_energy / total_mass[0]
+kinetic_energy_cgs = Temperature(kinetic_energy_cgs)
+
+internal_energy_cgs = internal_energy / total_mass[0]
+internal_energy_cgs = Temperature(internal_energy_cgs)
+
+radiated_energy_cgs = radiated_energy / total_mass[0]
+radiated_energy_cgs = Temperature(radiated_energy_cgs)
+
+# Read snapshots
+temp_snap = np.zeros(25)
+time_snap_cgs = np.zeros(25)
+for i in range(25):
+    snap = File("coolingBox_%0.4d.hdf5" % i, 'r')
+    u = snap["/PartType0/InternalEnergy"][:] * snap["/PartType0/Masses"][:]
+    u = sum(u) / total_mass[0]
+    temp_snap[i] = Temperature(u)
+    time_snap_cgs[i] = snap["/Header"].attrs["Time"] * unit_time
+
+
+plt.figure()
+
+Myr_in_yr = 3.15e13
+plt.plot(time, total_energy_cgs, 'r-', lw=1.6, label="Gas total temperature")
+plt.plot(time_snap_cgs, temp_snap, 'rD', ms=3)
+plt.plot(time, radiated_energy_cgs, 'g-', lw=1.6, label="Radiated temperature")
+plt.plot(time, total_energy_cgs + radiated_energy_cgs, 'b-',
+         lw=0.6, label="Gas total + radiated")
+
+plt.legend(loc="upper right", fontsize=8, frameon=False,
+           handlelength=3, ncol=1)
+plt.xlabel("${\\rm{Time~[Myr]}}$", labelpad=0)
+plt.ylabel("${\\rm{Temperature~[K]}}$")
+
+plt.show()
+# plt.savefig("temperature.png", dpi=200)
diff --git a/examples/Cooling/CoolingBox/run.sh b/examples/Cooling/CoolingBox/run.sh
index 8a0d717e49324cdc49301fa0c551d5b7da198d5e..89a6c1acfbb4b88e7b8bf2d3a7d7cfd192b3991b 100755
--- a/examples/Cooling/CoolingBox/run.sh
+++ b/examples/Cooling/CoolingBox/run.sh
@@ -21,7 +21,7 @@ then
 fi
 
 # Run SWIFT
-../../swift --cosmology --hydro --cooling --threads=4 -n 1000 coolingBox.yml
+../../swift --hydro --cooling --threads=4 -n 1000 coolingBox.yml
 
 # Check energy conservation and cooling rate
-python energy_plot.py
+python plotTemperature.py