Add grackle in structure and improve code

pyswiftsim/structure.py

@@ -283,28 +283,45 @@ class UnitSystem(SwiftStruct):
         super().__init__(self.struct_format, data, parent)
 
 
+######################################################################
+#                                                                    #
+#                     chemistry_part_data                            #
+#                                                                    #
+######################################################################
+class ChemistryPartData(SwiftStruct):
+    _format = "f"
+    _name = [
+        "he_density"
+    ]
+
+    def __init__(self, data, parent=None):
+        super().__init__(self.struct_format, data, parent)
+
 ######################################################################
 #                                                                    #
 #                          Part                                      #
 #                                                                    #
 ######################################################################
 class Part(SwiftStruct):
-    _format = "qP3d3f3ffffffN7f4c"
+    _format = "qP3d3f3ffffffN7f4{chemistry}sc".format(
+        chemistry=struct.calcsize(ChemistryPartData._format)
+    )
     _name = [
-            "id",
-            "gpart",
-            "pos",
-            "vel",
-            "a_hydro",
-            "h",
-            "mass",
-            "rho",
-            "entropy",
-            "entropy_dt",
-            "last_offset",
-            "density",
-            "time_bin"
-        ]
+        "id",
+        "gpart",
+        "pos",
+        "vel",
+        "a_hydro",
+        "h",
+        "mass",
+        "rho",
+        "entropy",
+        "entropy_dt",
+        "last_offset",
+        "density",
+        "chemistry_data",
+        "time_bin"
+    ]
 
     def __init__(self, data, parent=None):
         super().__init__(self.struct_format, data, parent)
@@ -416,7 +433,7 @@ class PhysConst(SwiftStruct):
         super().__init__(self.struct_format, data, parent)
 
 
-class GrackleUnits(SwiftStruct):
+class GrackleCodeUnits(SwiftStruct):
     _format = "idddddd"
     _name = [
         "comoving_coordinates",
@@ -493,17 +510,17 @@ class CoolingFunctionData(SwiftStruct):
             "cooling_tstep_mult"
         ]
     elif cooling_type == "grackle":
-        _format = "200cidd{units}s{chemistry}s".format(
-            units=struct.calcsize(GrackleUnits._format),
+        _format = "200cidd{code_units}s{chemistry}s".format(
+            code_units=struct.calcsize(GrackleCodeUnits._format),
             chemistry=struct.calcsize(GrackleChemistryData._format)
         )
         _name = [
-            "cloudy_table",
-            "uv_background",
-            "redshift",
-            "density_self_shielding",
-            "units",
-            "chemistry"
+            "GrackleCloudyTable",
+            "UVbackground",
+            "GrackleRedshift",
+            "GrackleHSShieldingDensityThreshold",
+            "code_units",
+            "chemistry",
         ]
 
         @property

setup.py

@@ -12,6 +12,13 @@ import numpy
 
 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"]
+
 # deal with arguments
 def parseCmdLine(arg, store=False):
     ret = False

src/config_wrapper.h

@@ -3,8 +3,6 @@
 
 #include "pyswiftsim_tools.h"
 
-#include <config.h>
-
 /**
  * @brief give the cooling type name
  *

src/cooling_wrapper.c

 #include "pyswiftsim_tools.h"
 #include "cooling_wrapper.h"
 
-#include <cooling.h>
-#include <cooling_struct.h>
-#include <equation_of_state.h>
-
-
-#include <numpy/arrayobject.h>
-
 
 PyObject* pycooling_init(PyObject* self, PyObject* args) {
   PyObject *pyparams;
@@ -43,7 +36,7 @@ PyObject* pycooling_init(PyObject* self, PyObject* args) {
 }
 
 PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) {
-  IMPORT_ARRAY();
+  import_array();
   
   PyObject *pycooling;
   PyObject *pyus;
@@ -98,6 +91,7 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) {
     return NULL;
 
   struct part p;
+  struct xpart xp;
 
 #ifdef COOLING_GRACKLE
   grackle_data = &cooling->chemistry;
@@ -106,6 +100,11 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) {
   /* 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++)
     {
@@ -114,13 +113,16 @@ 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);
+
       /* compute cooling rate */
       float *tmp = PyArray_GETPTR1(rate, i);
 #ifdef COOLING_GRACKLE
-      *tmp = cooling_rate(pconst, us, cooling, &p, dt);
+      *tmp = cooling_rate(pconst, us, cooling, &p, &xp, dt);
 #else
-      *tmp = cooling_rate(pconst, us, cooling, &p);
+      *tmp = cooling_rate(pconst, us, cooling, &p, &xp);
 #endif
+      return rate;
     }
 
   return rate;

src/parser_wrapper.c

 #include "parser_wrapper.h"
 #include "pyswiftsim_tools.h"
 
-#include <parser.h>
 
 PyObject* pyparser_read_file(PyObject *self, PyObject *args)
 {

src/part_wrapper.c

 #include "pyswiftsim_tools.h"
 
-#include <part.h>
-#include <hydro.h>
-
 #include <Python.h>
 #include <stdlib.h>
 #include <string.h>

src/part_wrapper.h

@@ -2,7 +2,6 @@
 #define __PYSWIFTSIM_PART_H__
 
 #include <Python.h>
-#include <part.h>
 
 /**
  * @brief Test function for the struct

src/pyswiftsim_tools.c

 #include "pyswiftsim_tools.h"
 
-/* include swift */
-#include <part.h>
-#include <units.h>
-#include <parser.h>
-#include <physical_constants.h>
-#include <cooling_struct.h>
-
-#include <Python.h>
-#include <numpy/arrayobject.h>
-
 
 const size_t class_size[class_count] = {
   sizeof(struct unit_system),
@@ -69,7 +59,6 @@ PyObject* pytools_return(void *p, int class)
   size_t nber_bytes;
 
   char module_name[STRING_SIZE] = "pyswiftsim.structure";
-  char *class_pyname;
 
   /* check class */
   if (class >= class_count)
@@ -77,7 +66,7 @@ PyObject* pytools_return(void *p, int class)
 
   /* get class information */
   nber_bytes = class_size[class];
-  class_pyname = class_name[class];
+  const char *class_pyname = class_name[class];
 
   /* import python class */
   python_class = pytools_import(module_name, class_pyname);
@@ -126,8 +115,8 @@ char* pytools_get_type_name(PyObject *obj)
     }
 
   /* transform to C */
-  size_t size;
-  char *name = PyUnicode_AsUTF8AndSize(recv, size);
+  Py_ssize_t size;
+  char *name = PyUnicode_AsUTF8AndSize(recv, &size);
   Py_DECREF(recv);
 
   if (name == NULL)
@@ -142,7 +131,7 @@ char* pytools_get_type_name(PyObject *obj)
 char* pytools_construct(PyObject* obj, int class)
 {
   char *module_name = "pyswiftsim.structure";
-  char *class_pyname;
+  const char *class_pyname;
 
   /* check python class */
   if (class >= class_count)
@@ -155,9 +144,9 @@ char* pytools_construct(PyObject* obj, int class)
   PyObject *pyclass = pytools_import(module_name, class_pyname);
 
   /* check if classes correspond */
-  int test = !PyObject_IsInstance(obj, pyclass);
+  int test = PyObject_IsInstance(obj, pyclass);
   Py_DECREF(pyclass);
-  if (test)
+  if (!test)
     {
       char *recv = pytools_get_type_name(obj);
       if (recv == NULL)
@@ -181,8 +170,7 @@ char* pytools_construct(PyObject* obj, int class)
 
 int pytools_check_array(PyArrayObject *obj, int dim, int type)
 {
-  /* ensure to have numpy arrays */
-  IMPORT_ARRAY();
+  IMPORT_ARRAY1(1);
 
   /* check if array */
   if (!PyArray_Check(obj))

src/pyswiftsim_tools.h

@@ -6,18 +6,12 @@
 #include <Python.h>
 #include <numpy/arrayobject.h>
 
+#undef assert
+#include <swift.h>
+
 #define DIM 3
 #define STRING_SIZE 200
 
-/* Check if arrays are imported and if not import them */
-#define IMPORT_ARRAY()				\
-  ({						\
-    if (PyArray_API == NULL)			\
-      {						\
-	import_array();				\
-      }						\
-  })
-
 /* Set the error message for python (still need to return NULL) */
 #define pyerror(s, ...)						\
   ({								\
@@ -37,6 +31,17 @@
 	    __FUNCTION__, __LINE__, ##__VA_ARGS__);		\
   })
 
+#define IMPORT_ARRAY1(ret)			\
+  {						\
+    if (PyArray_API == NULL)			\
+      {						\
+	import_array1(ret);			\
+      }						\
+  }
+
+#define IMPORT_ARRAY() IMPORT_ARRAY1(NUMPY_IMPORT_ARRAY_RETVAL);
+	
+
 /* Enum swift classes */
 enum swift_class {
   class_unit_system,

src/units_wrapper.c

 #include "pyswiftsim_tools.h"
 
-#include <units.h>
-#include <parser.h>
-#include <physical_constants.h>
-
 #include <Python.h>
 #include <stdlib.h>
 #include <string.h>
@@ -38,7 +34,7 @@ PyObject* pyunit_system_init(PyObject *self, PyObject *args)
   if (!PyArg_ParseTuple(args, "O", &parser))
     return NULL;
 
-  struct swift_params *params = pytools_construct(parser, class_swift_params);
+  struct swift_params *params = (struct swift_params*) pytools_construct(parser, class_swift_params);
 
   if (params == NULL)
     return NULL;

src/wrapper.c

@@ -5,7 +5,6 @@
 #include "pyswiftsim_tools.h"
 #include "config_wrapper.h"
 
-
 #include <Python.h>
 #include <math.h>
 #include <numpy/arrayobject.h>

test/test_cooling.py

@@ -19,8 +19,8 @@ gamma = 5. / 3.
 filename = "test/cooling.yml"
 
 # number of points
-N_rho = 300
-N_T = 300
+N_rho = 1000
+N_T = 1
 
 # density in atom / cm3
 if N_rho == 1:
@@ -49,10 +49,13 @@ T = T.reshape(-1)
 #
 
 # parse swift params
+print("Reading parameters")
 params = wrapper.parserReadFile(filename)
 # init units
+print("Initialization of the unit system")
 us, pconst = wrapper.unitSystemInit(params)
 # init cooling
+print("Initialization of the cooling")
 cooling = wrapper.coolingInit(params, us, pconst)
 
 
@@ -74,10 +77,13 @@ dt /= us.UnitTime_in_cgs
 #
 
 # du / dt
+print(cooling)
+print("Computing cooling...")
 rate = wrapper.coolingRate(pconst, us, cooling,
                            rho.astype(np.float32),
                            energy.astype(np.float32),
                            dt)
+print("Computing done")
 
 #
 # plot