diff --git a/setup.py b/setup.py
index b87b13a808514f6cf666837b5570337f8a721e28..6f8712d727e90d1af93e599ed618d67ad68da860 100644
--- a/setup.py
+++ b/setup.py
@@ -1,28 +1,36 @@
#!/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,
)
diff --git a/src/config_wrapper.h b/src/config_wrapper.h
index 3bc7a954e0d8a3fef33088317b767caa9c1d0f14..5a17aeb3c933dc9b36a7884bc11125c7a2b5c9ef 100644
--- a/src/config_wrapper.h
+++ b/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);
diff --git a/src/cooling_wrapper.c b/src/cooling_wrapper.c
index 6047370445aeb18143381f2586068b51384a294c..11f5bdb97a6d694e8f3c396553034d31bc723015 100644
--- a/src/cooling_wrapper.c
+++ b/src/cooling_wrapper.c
@@ -1,7 +1,18 @@
#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;
}
diff --git a/src/cooling_wrapper.h b/src/cooling_wrapper.h
index 4fbdd8abbc746d5d42aa4e32baa707c7770aa602..95717e8ac5d2867b25ddcebe33548fe86ecfa6cc 100644
--- a/src/cooling_wrapper.h
+++ b/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__
diff --git a/src/wrapper.c b/src/wrapper.c
index 12e7ef2b12653c9648e91bdf1c727cddb6993395..46076c5cb627eaed51654bb02107fcaeefcbbfb6 100644
--- a/src/wrapper.c
+++ b/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."},