Fix merge

pyswiftsim/structure.py

@@ -531,7 +531,7 @@ class CoolingFunctionData(SwiftStruct):
             }
 
             units = {
-                "class": GrackleUnits,
+                "class": GrackleCodeUnits,
                 "size": 1
             }
             return {

src/cooling_wrapper.c

@@ -94,17 +94,13 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) {
   struct xpart xp;
 
 #ifdef COOLING_GRACKLE
+  /* set grackle_data */
   grackle_data = &cooling->chemistry;
 #endif
 
   /* return object */
   PyArrayObject *rate = PyArray_NewLikeArray(energy, NPY_ANYORDER, NULL, 1);
 
-#ifdef COOLING_GRACKLE
-  /* set grackle_data */
-  grackle_data = &cooling->chemistry;
-#endif
-
   /* loop over all particles */
   for(size_t i = 0; i < N; i++)
     {
@@ -122,7 +118,6 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) {
 #else
       *tmp = cooling_rate(pconst, us, cooling, &p, &xp);
 #endif
-      return rate;
     }
 
   return rate;

test/test_cooling.py

@@ -77,7 +77,6 @@ dt /= us.UnitTime_in_cgs
 #
 
 # du / dt
-print(cooling)
 print("Computing cooling...")
 rate = wrapper.coolingRate(pconst, us, cooling,
                            rho.astype(np.float32),