Update cooling wrapper with OMP

setup.py

 #!/usr/bin/env python3
 
-descr = """
-Wrapper around the SPH cosmological simulation code SWIFT
-"""
-
 from setuptools import setup, find_packages, Extension
 import sys
 import os
 from glob import glob
 import numpy
 
+descr = """
+Wrapper around the SPH cosmological simulation code SWIFT
+"""
+
+with_omp = True
 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",
+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",
+    "-fopenmp"
 ]
 
+lflags = [
+    "-fopenmp"
+    ]
+
+
 # deal with arguments
 def parseCmdLine(arg, store=False):
     ret = False
@@ -37,11 +45,12 @@ def parseCmdLine(arg, store=False):
         sys.argv.remove(arg)
 
     return ret
-    
+
+
 swift_path = parseCmdLine("--with-swift", store=True)
 
 # python lib dependency
-install_requires=["numpy"]
+install_requires = ["numpy"]
 
 
 def getValueFromMakefile(swift_root, value):
@@ -51,7 +60,7 @@ def getValueFromMakefile(swift_root, value):
     with open(makefile, "r") as f:
         for line in f.readlines():
             if value == line[:N]:
-                return line[N:-1] # remove \n
+                return line[N:-1]  # remove \n
 
     raise ValueError("Value %s not found in Makefile" % value)
 
@@ -80,9 +89,10 @@ if swift_path:
     include.append(grackle_inc)
 
 # C libraries
-lib = ["m",
-       "swiftsim",
-       "hdf5",
+lib = [
+    "m",
+    "swiftsim",
+    "hdf5",
 ]
 
 lib_dir = []
@@ -90,7 +100,7 @@ lib_dir = []
 if swift_path:
     lib_dir.append(swift_path + "/src/.libs")
     lib_dir.append(hdf5_root + "/lib")
-       
+
 #  src files
 c_src = []
 
@@ -99,7 +109,7 @@ data_files = []
 
 
 ##############
-## C ext
+#  C ext
 ##############
 
 c_src = glob("src/*.c")
@@ -108,45 +118,41 @@ ext_modules = Extension("pyswiftsim.wrapper",
                         include_dirs=include,
                         libraries=lib,
                         library_dirs=lib_dir,
-                        extra_compile_args=cflags)
+                        extra_compile_args=cflags,
+                        extra_link_args=lflags)
 
 ext_modules = [ext_modules]
-    
+
 ##############
-## data
+#  data
 ##############
 
 data_files = []
 
 ##############
-## scripts
+#  scripts
 ##############
 
 list_scripts = []
 
 ##############
-## Setup
+#  Setup
 ##############
 
 setup(
-    name         = "pyswiftsim",
-    version      = "0.1",
-    author       = "Hausammann Loic",
-    author_email = "loic.hausammann@epfl.ch",
-    description  = descr,
-    license      = "GPLv3",
-    keywords     = "nbody sph simulation hpc",
-    url          = "",
-
-    packages         = find_packages(),
-
-    data_files       = data_files,
-
-    scripts          = list_scripts,
-
-    install_requires = install_requires,
-
-    dependency_links = dependency_links,
-    
-    ext_modules      = ext_modules,
+    name="pyswiftsim",
+    version="0.1",
+    author="Hausammann Loic",
+    author_email="loic.hausammann@epfl.ch",
+    description=descr,
+    license="GPLv3",
+    keywords="nbody sph simulation hpc",
+    url="",
+
+    packages=find_packages(),
+    data_files=data_files,
+    scripts=list_scripts,
+    install_requires=install_requires,
+    dependency_links=dependency_links,
+    ext_modules=ext_modules,
 )

src/config_wrapper.h

@@ -16,13 +16,17 @@ PyObject* config_get_cooling() {
 
   /* grackle */
 #elif defined(COOLING_GRACKLE)
-#ifdef CONFIG_BFLOAT_4
-  cooling_name = "grackle_float";
-#elif CONFIG_BFLOAT_8
-  cooling_name = "grackle_double";
+#if COOLING_GRACKLE_MODE == 0
+  cooling_name = "grackle";
+#elif COOLING_GRACKLE_MODE == 1
+  cooling_name = "grackle1";
+#elif COOLING_GRACKLE_MODE == 2
+  cooling_name = "grackle2";  
+#elif COOLING_GRACKLE_MODE == 3
+  cooling_name = "grackle3";
 #else
-  error("Impossible to config for grackle");
-#endif // CONFIG_BFLOAT
+  error("Grackle mode unknown");
+#endif // COOLING_GRACKLE_MODE
 #endif // COOLING_GRACKLE
   
   return PyUnicode_FromString(cooling_name);

src/cooling_wrapper.c

 #include "pyswiftsim_tools.h"
 #include "cooling_wrapper.h"
-
-
+#include <omp.h>
+
+
+/**
+ * @brief Initialize the cooling
+ *
+ * args is expecting pyswiftsim classes in the following order:
+ * SwiftParams, UnitSystem and PhysConst. 
+ *
+ * @param self calling object
+ * @param args arguments
+ * @return CoolingFunctionData
+ */
 PyObject* pycooling_init(PyObject* self, PyObject* args) {
   PyObject *pyparams;
   PyObject *pyus;
@@ -35,6 +46,52 @@ PyObject* pycooling_init(PyObject* self, PyObject* args) {
   return pycooling;
 }
 
+/**
+ * @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) {
+  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
+
+}
+  
+/**
+ * @brief Compute cooling rate
+ *
+ * args is expecting pyswiftsim classes in the following order: 
+ * PhysConst, UnitSystem and CoolingFunctionData.
+ * Then two numpy arrays (density and specific energy) and an optional
+ * float for the time step
+ *
+ * @param self calling object
+ * @param args arguments
+ * @return cooling rate
+ */
 PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) {
   import_array();
   
@@ -44,35 +101,40 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) {
 
   PyArrayObject *rho;
   PyArrayObject *energy;
+  PyArrayObject *fractions = NULL;
 
   float dt = 1e-3;
 
   /* parse argument */
   if (!PyArg_ParseTuple(args,
-			"OOOOO|f",
+			"OOOOO|fOO",
 			&pypconst,
 			&pyus,
 			&pycooling,
 			&rho,
 			&energy,
-			&dt))
+			&dt,
+			&fractions
+			))
     return NULL;
 
   /* check numpy array */
   if (pytools_check_array(energy, 1, NPY_FLOAT) != SUCCESS)
-    {
-      return NULL;
-    }
+    return NULL;
 
   if (pytools_check_array(rho, 1, NPY_FLOAT) != SUCCESS)
-    {
-      return NULL;
-    }
+    return NULL;
+
+  if (fractions != NULL &&
+      pytools_check_array(fractions, 2, NPY_FLOAT) != SUCCESS)
+    return NULL;
 
   if (PyArray_DIM(energy, 0) != PyArray_DIM(rho, 0))
-    {
-      pyerror("Density and energy should have the same dimension");
-    }
+    pyerror("Density and energy should have the same dimension");
+
+  if (fractions != NULL &&
+      PyArray_DIM(fractions, 0) != PyArray_DIM(rho,0))
+    pyerror("Fractions should have the same first dimension than the others");
 
   size_t N = PyArray_DIM(energy, 0);
 
@@ -101,7 +163,11 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) {
   /* return object */
   PyArrayObject *rate = PyArray_SimpleNew(PyArray_NDIM(energy), PyArray_DIMS(energy), NPY_FLOAT);
 
+  /* Release GIL */
+  Py_BEGIN_ALLOW_THREADS;
+  
   /* loop over all particles */
+#pragma omp for
   for(size_t i = 0; i < N; i++)
     {
       /* set particle data */
@@ -109,8 +175,11 @@ 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);
-
+      if (fractions != NULL)
+	pycooling_set_fractions(&xp, fractions, i);
+      else
+	cooling_first_init_part(&p, &xp, cooling);
+	
       /* compute cooling rate */
       float *tmp = PyArray_GETPTR1(rate, i);
 #ifdef COOLING_GRACKLE
@@ -120,6 +189,9 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) {
 #endif
     }
 
+  /* Acquire GIL */
+  Py_END_ALLOW_THREADS;
+
   return rate;
   
 }

src/cooling_wrapper.h

@@ -3,30 +3,8 @@
 
 #include "pyswiftsim_tools.h"
 
-/**
- * @brief Initialize the cooling
- *
- * args is expecting pyswiftsim classes in the following order:
- * SwiftParams, UnitSystem and PhysConst. 
- *
- * @param self calling object
- * @param args arguments
- * @return CoolingFunctionData
- */
 PyObject* pycooling_init(PyObject* self, PyObject* args);
 
-/**
- * @brief Compute cooling rate
- *
- * args is expecting pyswiftsim classes in the following order: 
- * PhysConst, UnitSystem and CoolingFunctionData.
- * Then two numpy arrays (density and specific energy) and an optional
- * float for the time step
- *
- * @param self calling object
- * @param args arguments
- * @return cooling rate
- */
 PyArrayObject* pycooling_rate(PyObject* self, PyObject* args);
 
 #endif // __PYSWIFTSIM_COOLING_H__

src/wrapper.c

@@ -30,7 +30,21 @@ static PyMethodDef wrapper_methods[] = {
    "Initialize cooling."},
 
   {"coolingRate", pycooling_rate, METH_VARARGS,
-   "Compute the cooling rate."},
+   "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"
+   "dt: float, optional\n"
+   "\t Time step in pyus units\n"
+   "fractions: np.array, optional\n"
+   "\t Fraction of each cooling element (including metals)\n"
+  },
 
   {"configGetCooling", config_get_cooling, METH_VARARGS,
    "Get the cooling type."},