Skip to content
Snippets Groups Projects
Commit dfebccd9 authored by lhausamm's avatar lhausamm
Browse files

Update cooling wrapper with OMP

parent 2204383a
Branches
No related tags found
1 merge request!4Add cooling rate example
Hide whitespace changes
Inline Side-by-side
setup.py
+ 52
− 46
View file @ dfebccd9
#!/usr/bin/env python3 #!/usr/bin/env python3
descr = """
Wrapper around the SPH cosmological simulation code SWIFT
"""
from setuptools import setup, find_packages, Extension from setuptools import setup, find_packages, Extension
import sys import sys
import os import os
from glob import glob from glob import glob
import numpy import numpy
descr = """
Wrapper around the SPH cosmological simulation code SWIFT
"""
with_omp = True
os.environ["CC"] = "mpicc" os.environ["CC"] = "mpicc"
cflags = ["-Werror", cflags = [
"-Wall", "-Werror",
"-Wextra", "-Wall",
# disables some warnings due to python "-Wextra",
"-Wno-unused-parameter", # disables some warnings due to python
"-Wno-strict-prototypes", "-Wno-unused-parameter",
"-Wno-unused-function", "-Wno-strict-prototypes",
"-Wno-incompatible-pointer-types", "-Wno-unused-function",
"-Wno-missing-field-initializers", "-Wno-incompatible-pointer-types",
"-Wno-missing-field-initializers",
"-fopenmp"
] ]
lflags = [
"-fopenmp"
]
# deal with arguments # deal with arguments
def parseCmdLine(arg, store=False): def parseCmdLine(arg, store=False):
ret = False ret = False
...@@ -37,11 +45,12 @@ def parseCmdLine(arg, store=False): ...@@ -37,11 +45,12 @@ def parseCmdLine(arg, store=False):
sys.argv.remove(arg) sys.argv.remove(arg)
return ret return ret
swift_path = parseCmdLine("--with-swift", store=True) swift_path = parseCmdLine("--with-swift", store=True)
# python lib dependency # python lib dependency
install_requires=["numpy"] install_requires = ["numpy"]
def getValueFromMakefile(swift_root, value): def getValueFromMakefile(swift_root, value):
...@@ -51,7 +60,7 @@ def getValueFromMakefile(swift_root, value): ...@@ -51,7 +60,7 @@ def getValueFromMakefile(swift_root, value):
with open(makefile, "r") as f: with open(makefile, "r") as f:
for line in f.readlines(): for line in f.readlines():
if value == line[:N]: 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) raise ValueError("Value %s not found in Makefile" % value)
...@@ -80,9 +89,10 @@ if swift_path: ...@@ -80,9 +89,10 @@ if swift_path:
include.append(grackle_inc) include.append(grackle_inc)
# C libraries # C libraries
lib = ["m", lib = [
"swiftsim", "m",
"hdf5", "swiftsim",
"hdf5",
] ]
lib_dir = [] lib_dir = []
...@@ -90,7 +100,7 @@ lib_dir = [] ...@@ -90,7 +100,7 @@ lib_dir = []
if swift_path: if swift_path:
lib_dir.append(swift_path + "/src/.libs") lib_dir.append(swift_path + "/src/.libs")
lib_dir.append(hdf5_root + "/lib") lib_dir.append(hdf5_root + "/lib")
# src files # src files
c_src = [] c_src = []
...@@ -99,7 +109,7 @@ data_files = [] ...@@ -99,7 +109,7 @@ data_files = []
############## ##############
## C ext # C ext
############## ##############
c_src = glob("src/*.c") c_src = glob("src/*.c")
...@@ -108,45 +118,41 @@ ext_modules = Extension("pyswiftsim.wrapper", ...@@ -108,45 +118,41 @@ ext_modules = Extension("pyswiftsim.wrapper",
include_dirs=include, include_dirs=include,
libraries=lib, libraries=lib,
library_dirs=lib_dir, library_dirs=lib_dir,
extra_compile_args=cflags) extra_compile_args=cflags,
extra_link_args=lflags)
ext_modules = [ext_modules] ext_modules = [ext_modules]
############## ##############
## data # data
############## ##############
data_files = [] data_files = []
############## ##############
## scripts # scripts
############## ##############
list_scripts = [] list_scripts = []
############## ##############
## Setup # Setup
############## ##############
setup( setup(
name = "pyswiftsim", name="pyswiftsim",
version = "0.1", version="0.1",
author = "Hausammann Loic", author="Hausammann Loic",
author_email = "loic.hausammann@epfl.ch", author_email="loic.hausammann@epfl.ch",
description = descr, description=descr,
license = "GPLv3", license="GPLv3",
keywords = "nbody sph simulation hpc", keywords="nbody sph simulation hpc",
url = "", url="",
packages = find_packages(), packages=find_packages(),
data_files=data_files,
data_files = data_files, scripts=list_scripts,
install_requires=install_requires,
scripts = list_scripts, dependency_links=dependency_links,
ext_modules=ext_modules,
install_requires = install_requires,
dependency_links = dependency_links,
ext_modules = ext_modules,
) )
src/config_wrapper.h
+ 10
− 6
View file @ dfebccd9
...@@ -16,13 +16,17 @@ PyObject* config_get_cooling() { ...@@ -16,13 +16,17 @@ PyObject* config_get_cooling() {
/* grackle */ /* grackle */
#elif defined(COOLING_GRACKLE) #elif defined(COOLING_GRACKLE)
#ifdef CONFIG_BFLOAT_4 #if COOLING_GRACKLE_MODE == 0
cooling_name = "grackle_float"; cooling_name = "grackle";
#elif CONFIG_BFLOAT_8 #elif COOLING_GRACKLE_MODE == 1
cooling_name = "grackle_double"; cooling_name = "grackle1";
#elif COOLING_GRACKLE_MODE == 2
cooling_name = "grackle2";
#elif COOLING_GRACKLE_MODE == 3
cooling_name = "grackle3";
#else #else
error("Impossible to config for grackle"); error("Grackle mode unknown");
#endif // CONFIG_BFLOAT #endif // COOLING_GRACKLE_MODE
#endif // COOLING_GRACKLE #endif // COOLING_GRACKLE
return PyUnicode_FromString(cooling_name); return PyUnicode_FromString(cooling_name);
......
src/cooling_wrapper.c
+ 87
− 15
View file @ dfebccd9
#include "pyswiftsim_tools.h" #include "pyswiftsim_tools.h"
#include "cooling_wrapper.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* pycooling_init(PyObject* self, PyObject* args) {
PyObject *pyparams; PyObject *pyparams;
PyObject *pyus; PyObject *pyus;
...@@ -35,6 +46,52 @@ PyObject* pycooling_init(PyObject* self, PyObject* args) { ...@@ -35,6 +46,52 @@ PyObject* pycooling_init(PyObject* self, PyObject* args) {
return pycooling; 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) { PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) {
import_array(); import_array();
...@@ -44,35 +101,40 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) { ...@@ -44,35 +101,40 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) {
PyArrayObject *rho; PyArrayObject *rho;
PyArrayObject *energy; PyArrayObject *energy;
PyArrayObject *fractions = NULL;
float dt = 1e-3; float dt = 1e-3;
/* parse argument */ /* parse argument */
if (!PyArg_ParseTuple(args, if (!PyArg_ParseTuple(args,
"OOOOO|f", "OOOOO|fOO",
&pypconst, &pypconst,
&pyus, &pyus,
&pycooling, &pycooling,
&rho, &rho,
&energy, &energy,
&dt)) &dt,
&fractions
))
return NULL; return NULL;
/* check numpy array */ /* check numpy array */
if (pytools_check_array(energy, 1, NPY_FLOAT) != SUCCESS) if (pytools_check_array(energy, 1, NPY_FLOAT) != SUCCESS)
{ return NULL;
return NULL;
}
if (pytools_check_array(rho, 1, NPY_FLOAT) != SUCCESS) 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)) 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); size_t N = PyArray_DIM(energy, 0);
...@@ -101,7 +163,11 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) { ...@@ -101,7 +163,11 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) {
/* return object */ /* return object */
PyArrayObject *rate = PyArray_SimpleNew(PyArray_NDIM(energy), PyArray_DIMS(energy), NPY_FLOAT); PyArrayObject *rate = PyArray_SimpleNew(PyArray_NDIM(energy), PyArray_DIMS(energy), NPY_FLOAT);
/* Release GIL */
Py_BEGIN_ALLOW_THREADS;
/* loop over all particles */ /* loop over all particles */
#pragma omp for
for(size_t i = 0; i < N; i++) for(size_t i = 0; i < N; i++)
{ {
/* set particle data */ /* set particle data */
...@@ -109,8 +175,11 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) { ...@@ -109,8 +175,11 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) {
float u = *(float*) PyArray_GETPTR1(energy, i); float u = *(float*) PyArray_GETPTR1(energy, i);
p.entropy = gas_entropy_from_internal_energy(p.rho, u); 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 */ /* compute cooling rate */
float *tmp = PyArray_GETPTR1(rate, i); float *tmp = PyArray_GETPTR1(rate, i);
#ifdef COOLING_GRACKLE #ifdef COOLING_GRACKLE
...@@ -120,6 +189,9 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) { ...@@ -120,6 +189,9 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) {
#endif #endif
} }
/* Acquire GIL */
Py_END_ALLOW_THREADS;
return rate; return rate;
} }
src/cooling_wrapper.h
+ 0
− 22
View file @ dfebccd9
...@@ -3,30 +3,8 @@ ...@@ -3,30 +3,8 @@
#include "pyswiftsim_tools.h" #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); 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); PyArrayObject* pycooling_rate(PyObject* self, PyObject* args);
#endif // __PYSWIFTSIM_COOLING_H__ #endif // __PYSWIFTSIM_COOLING_H__
......
src/wrapper.c
+ 15
− 1
View file @ dfebccd9
...@@ -30,7 +30,21 @@ static PyMethodDef wrapper_methods[] = { ...@@ -30,7 +30,21 @@ static PyMethodDef wrapper_methods[] = {
"Initialize cooling."}, "Initialize cooling."},
{"coolingRate", pycooling_rate, METH_VARARGS, {"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, {"configGetCooling", config_get_cooling, METH_VARARGS,
"Get the cooling type."}, "Get the cooling type."},
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment