Cooling is working and able to reproduce grackle paper graphs

pyswiftsim/structure.py

+from pyswiftsim import wrapper
+
 import struct
 import numpy
 from ctypes import *
@@ -52,14 +54,14 @@ class SwiftStruct(struct.Struct):
 
         Should be implemented for each sub class
         """
-        raise NotImplementedError("SwiftStruct should not be used")
+        return self.__class__._name
 
     @property
     def struct_format(self):
         """
         String containing the data type of each attribute
         """
-        raise NotImplementedError("SwiftStruct should not be used")
+        return self.__class__._format
 
     @property
     def struct_substruct(self):
@@ -218,13 +220,18 @@ class SwiftStruct(struct.Struct):
 #                                                                    #
 ######################################################################
 class ArrayStruct(SwiftStruct):
+    _name = [
+        "array_data"
+    ]
+        
     def __init__(self, struct_format, data, parent, name):
         super().__init__(struct_format, data, parent)
-        self._name = name
         self._format = struct_format
+        self._name = name
 
     def __getitem__(self, ii):
         data = self.unpack(self.data)
+        data = self._clean(data)
         return data[ii]
     
     def __setitem__(self, ii, value):
@@ -232,20 +239,24 @@ class ArrayStruct(SwiftStruct):
         data[ii] = value
         setattr(self.parent, self._name, data)
 
+    def _clean(self, data):
+        if self._format[-1] == "c":
+            data = b"".join(data)
+            # find end of txt character
+            n = data.index(b"\x00")
+            data = data[:n].decode("utf-8")
+
+        return data
+
     def __str__(self):
         data = self.unpack(self.data)
+        data = self._clean(data)
         return str(data)
         
         
     def getArray(self):
         return self.unpack(self.data)
 
-    @property
-    def struct_name(self):
-        return [
-            "array_data"
-        ]
-
     @property
     def struct_format(self):
         return self._format
@@ -257,24 +268,18 @@ class ArrayStruct(SwiftStruct):
 #                                                                    #
 ######################################################################
 class UnitSystem(SwiftStruct):
+    _format = "ddddd"
+    _name = [
+        "UnitMass_in_cgs",
+        "UnitLength_in_cgs",
+        "UnitTime_in_cgs",
+        "UnitCurrent_in_cgs",
+        "UnitTemperature_in_cgs",
+    ]
 
     def __init__(self, data, parent=None):
         super().__init__(self.struct_format, data, parent)
 
-    @property
-    def struct_format(self):
-        return "ddddd"
-
-    @property
-    def struct_name(self):
-        return  [
-            "UnitMass_in_cgs",
-            "UnitLength_in_cgs",
-            "UnitTime_in_cgs",
-            "UnitCurrent_in_cgs",
-            "UnitTemperature_in_cgs",
-        ]
-
 
 ######################################################################
 #                                                                    #
@@ -282,18 +287,8 @@ class UnitSystem(SwiftStruct):
 #                                                                    #
 ######################################################################
 class Part(SwiftStruct):
-    def __init__(self, data, parent=None):
-        super().__init__(self.struct_format, data, parent)
-
-    @property
-    def struct_format(self):
-        print("ERROR, need to fix density/time_bin")
-        return "qP3d3f3ffffffN7f4c"
-
-
-    @property
-    def struct_name(self):
-        return  [
+    _format = "qP3d3f3ffffffN7f4c"
+    _name = [
             "id",
             "gpart",
             "pos",
@@ -309,6 +304,11 @@ class Part(SwiftStruct):
             "time_bin"
         ]
 
+    def __init__(self, data, parent=None):
+        super().__init__(self.struct_format, data, parent)
+        print("ERROR, need to fix density/time_bin")
+        
+    
 ######################################################################
 #                                                                    #
 #                       Parameter                                    #
@@ -327,14 +327,6 @@ class Parameter(SwiftStruct):
 
     def __init__(self, data, parent=None):
         super().__init__(self.struct_format, data, parent)
-
-    @property
-    def struct_format(self):
-        return Parameter._format
-
-    @property
-    def struct_name(self):
-        return Parameter._name
     
 ######################################################################
 #                                                                    #
@@ -352,15 +344,6 @@ class Section(SwiftStruct):
     def __init__(self, data, parent=None):
         super().__init__(self.struct_format, data, parent)
 
-    @property
-    def struct_format(self):
-        return Section._format
-
-
-    @property
-    def struct_name(self):
-        return Section._name 
-
 ######################################################################
 #                                                                    #
 #                       SwiftParams                                  #
@@ -384,15 +367,6 @@ class SwiftParams(SwiftStruct):
     def __init__(self, data, parent=None):
         super().__init__(self.struct_format, data, parent)
 
-    @property
-    def struct_format(self):
-        return SwiftParams._format
-
-
-    @property
-    def struct_name(self):
-        return SwiftParams._name
-
     @property
     def struct_substruct(self):
         sec = {
@@ -439,11 +413,29 @@ class PhysConst(SwiftStruct):
     def __init__(self, data, parent=None):
         super().__init__(self.struct_format, data, parent)
 
-    @property
-    def struct_format(self):
-        return PhysConst._format
 
+class CoolingFunctionData(SwiftStruct):
+    cooling_type = wrapper.configGetCooling()
+    if cooling_type == "const_lambda":
+        _format = "dffff"
+        _name = [
+            "lambda",
+            "hydrogen_mass_abundance",
+            "mean_molecular_weight",
+            "min_energy",
+            "cooling_tstep_mult"
+        ]
+    elif cooling_type == "grackle":
+        _format = "200cidd"
+        _name = [
+            "GrackleCloudyTable",
+            "UVbackground",
+            "GrackleRedshift",
+            "GrackleHSShieldingDensityThreshold"
+        ]
+    else:
+        raise ValueError(
+            "Cooling Type %s not implemented" % cooling_type)
 
-    @property
-    def struct_name(self):
-        return PhysConst._name
+    def __init__(self, data, parent=None):
+        super().__init__(self.struct_format, data, parent)

setup.py

@@ -13,13 +13,21 @@ import numpy
 os.environ["CC"] = "mpicc"
 
 # deal with arguments
-swift_path = None
-swift_arg = "--with-swift"
-if swift_arg in sys.argv:
-    i = sys.argv.index(swift_arg)
-    swift_path = sys.argv[i+1]
-    sys.argv.remove(swift_arg)
-    sys.argv.remove(swift_path)
+def parseCmdLine(arg, store=False):
+    ret = False
+    if arg in sys.argv:
+        if store:
+            i = sys.argv.index(arg)
+            ret = sys.argv[i+1]
+            sys.argv.remove(ret)
+        else:
+            ret = True
+
+        sys.argv.remove(arg)
+
+    return ret
+    
+swift_path = parseCmdLine("--with-swift", store=True)
 
 # python lib dependency
 install_requires=["numpy"]
@@ -37,7 +45,7 @@ def getValueFromMakefile(swift_root, value):
     raise ValueError("Value %s not found in Makefile" % value)
 
 
-if swift_path is not None:
+if swift_path:
     hdf5_root = getValueFromMakefile(swift_path, "H5CC = ")
 
 # link to non pypi dependencies
@@ -46,7 +54,7 @@ dependency_links = []
 # C includes
 include = [numpy.get_include()]
 
-if swift_path is not None:
+if swift_path:
     # swift
     include.extend([
         swift_path + "/src",
@@ -56,19 +64,22 @@ if swift_path is not None:
     # hdf5
     include.append(hdf5_root + "/include")
 
+    # grackle
+    grackle_inc = getValueFromMakefile(swift_path, "GRACKLE_INCS = -I")
+    include.append(grackle_inc)
+
 # C libraries
 lib = ["m",
        "swiftsim",
-       "hdf5"
+       "hdf5",
 ]
 
 lib_dir = []
 
-if swift_path is not None:
+if swift_path:
     lib_dir.append(swift_path + "/src/.libs")
     lib_dir.append(hdf5_root + "/lib")
-
-    
+       
 #  src files
 c_src = []
 

src/cooling_wrapper.c

@@ -3,24 +3,29 @@
 
 #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;
-  PyObject* pyus;
-  PyObject* pypconst;
+  PyObject *pyparams;
+  PyObject *pyus;
+  PyObject *pypconst;
 
   if (!PyArg_ParseTuple(args, "OOO", &pyparams, &pyus, &pypconst))
       return NULL;
 
-  struct swift_params *params = pytools_construct(pyparams, class_swift_params);
+  struct swift_params *params = (struct swift_params*) pytools_construct(pyparams, class_swift_params);
   if (params == NULL)
     return NULL;
 
-  struct unit_system *us = pytools_construct(pyus, class_unit_system);
+  struct unit_system *us = (struct unit_system*) pytools_construct(pyus, class_unit_system);
   if (us == NULL)
     return NULL;
   
-  struct phys_const *pconst = pytools_construct(pypconst, class_phys_const);
+  struct phys_const *pconst = (struct phys_const*) pytools_construct(pypconst, class_phys_const);
   if (pconst == NULL)
     return NULL;
 
@@ -34,6 +39,80 @@ PyObject* pycooling_init(PyObject* self, PyObject* args) {
   return pycooling;
 }
 
-PyObject* pycooling_rate(PyObject* self, PyObject* args) {
-  pyerror("Not implemented");
+PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) {
+  IMPORT_ARRAY();
+  
+  PyObject *pycooling;
+  PyObject *pyus;
+  PyObject *pypconst;
+
+  PyArrayObject *rho;
+  PyArrayObject *energy;
+
+  float dt = 1e-3;
+
+  if (!PyArg_ParseTuple(args,
+			"OOOOO|f",
+			&pypconst,
+			&pyus,
+			&pycooling,
+			&rho,
+			&energy,
+			&dt))
+    return NULL;
+
+  if (pytools_check_array(energy, 1, NPY_FLOAT) != SUCCESS)
+    {
+      return NULL;
+    }
+
+  if (pytools_check_array(rho, 1, NPY_FLOAT) != SUCCESS)
+    {
+      return NULL;
+    }
+
+  if (PyArray_DIM(energy, 0) != PyArray_DIM(rho, 0))
+    {
+      pyerror("Density and energy should have the same dimension");
+    }
+
+    size_t N = PyArray_DIM(energy, 0);
+
+  struct cooling_function_data *cooling = (struct cooling_function_data*) pytools_construct(pycooling, class_cooling_function_data);
+  if (cooling == NULL)
+    return NULL;
+
+  struct unit_system *us = (struct unit_system*) pytools_construct(pyus, class_unit_system);
+  if (us == NULL)
+    return NULL;
+
+  struct phys_const *pconst = (struct phys_const*) pytools_construct(pypconst, class_phys_const);
+  if (pconst == NULL)
+    return NULL;
+
+  struct part p;
+
+#ifdef COOLING_GRACKLE
+  grackle_data.grackle_data_file = cooling->GrackleCloudyTable;
+#endif
+
+  PyArrayObject *rate = PyArray_NewLikeArray(energy, NPY_ANYORDER, NULL, 1);
+
+  for(size_t i = 0; i < N; i++)
+    {
+      p.rho = *(float*) PyArray_GETPTR1(rho, i);
+      float u = *(float*) PyArray_GETPTR1(energy, i);
+      p.entropy = gas_entropy_from_internal_energy(p.rho, u);
+
+      float *tmp = PyArray_GETPTR1(rate, i);
+
+#ifdef COOLING_GRACKLE
+      *tmp = cooling_rate(pconst, us, cooling, &p, dt);
+#else
+      *tmp = cooling_rate(pconst, us, cooling, &p);
+#endif
+    }
+
+  return rate;
+  
 }

src/cooling_wrapper.h

@@ -5,7 +5,7 @@
 
 PyObject* pycooling_init(PyObject* self, PyObject* args);
 
-PyObject* pycooling_rate(PyObject* self, PyObject* args);
+PyArrayObject* pycooling_rate(PyObject* self, PyObject* args);
 
 #endif // __PYSWIFTSIM_COOLING_H__
 

src/part_wrapper.c

@@ -26,7 +26,6 @@ PyObject* pypart_test_struct(PyObject *self, PyObject *args)
   p->rho = 12.;
   p->entropy = 13;
   p->entropy_dt = 14;
-  p->last_offset = 15;
   
   hydro_init_part(p, NULL);
 
@@ -36,3 +35,4 @@ PyObject* pypart_test_struct(PyObject *self, PyObject *args)
   
   return object;
 }
+

src/pyswiftsim_tools.c

@@ -8,6 +8,8 @@
 #include <cooling_struct.h>
 
 #include <Python.h>
+#include <numpy/arrayobject.h>
+
 
 const size_t class_size[class_count] = {
   sizeof(struct unit_system),
@@ -164,3 +166,27 @@ char* pytools_construct(PyObject* obj, int class)
   Py_DECREF(data);
   return ret;
 }
+
+
+int pytools_check_array(PyArrayObject *obj, int dim, int type)
+{
+  IMPORT_ARRAY();
+  
+  if (!PyArray_Check(obj))
+    {
+      pyerror("Expecting a numpy array");
+    }
+  
+  if (PyArray_NDIM(obj) != dim)
+    {
+      pyerror("Array should be a %i dimensional object", dim);
+    }
+
+  if (PyArray_TYPE(obj) != type)
+    {
+      pyerror("Wrong array type");
+    }
+
+  return SUCCESS;
+
+}

src/pyswiftsim_tools.h

@@ -4,10 +4,20 @@
 #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
 
 #include <Python.h>
+#include <numpy/arrayobject.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, ...)						\
   ({								\
@@ -19,6 +29,12 @@
     return FAIL;						\
   })
 
+#define pydebug(s, ...)						\
+  ({								\
+    PyErr_Print();						\
+    fprintf(stderr, "\n%s:%s():%i: Test " s "\n", __FILE__,	\
+	    __FUNCTION__, __LINE__, ##__VA_ARGS__);		\
+  })
 
 enum class {
   class_unit_system,
@@ -45,4 +61,5 @@ PyObject* pytools_import(char* module, char* object_name);
 
 char* pytools_get_type_name(PyObject *obj);
 
+int pytools_check_array(PyArrayObject *obj, int dim, int type);
 #endif // __PYSWIFTSIM_TOOLS_H__

src/units_wrapper.c

@@ -52,7 +52,7 @@ PyObject* pyunit_system_init(PyObject *self, PyObject *args)
   
   PyObject *pypconst = pytools_return(&pconst, class_phys_const);
   if (pypconst == NULL)
-    return NULL;
+    return NULL;  
 
   return PyTuple_Pack(2, pyus, pypconst);
 }

src/wrapper.c

@@ -3,6 +3,7 @@
 #include "parser_wrapper.h"
 #include "cooling_wrapper.h"
 #include "pyswiftsim_tools.h"
+#include "config_wrapper.h"
 
 
 #include <Python.h>
@@ -29,6 +30,12 @@ static PyMethodDef wrapper_methods[] = {
   {"coolingInit", pycooling_init, METH_VARARGS,
    "Initialize cooling."},
 
+  {"coolingRate", pycooling_rate, METH_VARARGS,
+   "Compute the cooling rate."},
+
+  {"configGetCooling", config_get_cooling, METH_VARARGS,
+   "Get the cooling type."},
+  
   {NULL, NULL, 0, NULL}        /* Sentinel */
 };      
       
@@ -47,8 +54,8 @@ PyMODINIT_FUNC PyInit_wrapper(void)
 {
   PyObject *m;
 
-  Py_Initialize();
   import_array();
+  Py_Initialize();
   m = PyModule_Create(&wrapper_cmodule);
 
   if (m == NULL)

test/cooling.yml

+# Define the system of units to use internally. 
+InternalUnitSystem:
+  UnitMass_in_cgs:     1.989e43   # Grams
+  UnitLength_in_cgs:   3.085e21   # Centimeters
+  UnitVelocity_in_cgs: 20725573.785998672   # Centimeters per second
+  UnitCurrent_in_cgs:  1   # Amperes
+  UnitTemp_in_cgs:     1   # Kelvin
+
+# Parameters governing the time integration
+TimeIntegration:
+  time_begin: 0.    # The starting time of the simulation (in internal units).
+  time_end:   5e-2  # The end time of the simulation (in internal units).
+  dt_min:     1e-7  # The minimal time-step size of the simulation (in internal units).
+  dt_max:     1e-4  # The maximal time-step size of the simulation (in internal units).
+
+# Parameters governing the snapshots
+Snapshots:
+  basename:            sedov # Common part of the name of output files
+  time_first:          0.    # Time of the first output (in internal units)
+  delta_time:          1e-2  # Time difference between consecutive outputs (in internal units)
+
+# Parameters governing the conserved quantities statistics
+Statistics:
+  delta_time:          1e-3 # Time between statistics output
+  tasks_per_cell:      1000
+  
+# Parameters for the hydrodynamics scheme
+SPH:
+  resolution_eta:        1.2348   # Target smoothing length in units of the mean inter-particle separation (1.2348 == 48Ngbs with the cubic spline kernel).
+  CFL_condition:         0.1      # Courant-Friedrich-Levy condition for time integration.
+
+# Parameters related to the initial conditions
+InitialConditions:
+  file_name:  ./sedov.hdf5          # The file to read
+  h_scaling:           3.33
+
+
+# Cooling with Grackle 2.0
+GrackleCooling:
+  GrackleCloudyTable: CloudyData_UVB=HM2012.h5 # Name of the Cloudy Table
+  UVbackground: 1 # Enable or not the UV background
+  GrackleRedshift: 0 # Redshift to use (-1 means time based redshift)
+  GrackleHSShieldingDensityThreshold: 1.1708e-26 # self shielding threshold in internal system of units
+
+Gravity:
+  eta:          0.025    # Constant dimensionless multiplier for time integration.
+  theta:        0.7      # Opening angle (Multipole acceptance criterion)
+  epsilon:      0.1      # Softening length (in internal units).
+  a_smooth:     1.25     # (Optional) Smoothing scale in top-level cell sizes to smooth the long-range forces over (this is the default value).
+  r_cut_max:    4.5      # (Optional) Cut-off in number of top-level cells beyond which no FMM forces are computed (this is the default value).
+  r_cut_min:    0.1      # (Optional) Cut-off in number of top-level cells below which no truncation of FMM forces are performed (this is the default value).

test/test_cooling.py

@@ -3,16 +3,134 @@
 from pyswiftsim import wrapper
 from pyswiftsim import structure
 
-filename = "/home/loikki/swift_test/cooling_sedov/sedov.yml"
+from copy import deepcopy
+import numpy as np
+import matplotlib.pyplot as plt
+from astropy import units
 
-params = wrapper.parserReadFile(filename)
+#
+# parameters
+#
 
-us, pconst = wrapper.unitSystemInit(params)
+# adiabatic index
+gamma = 5. / 3.
+
+# swift params filename
+filename = "test/cooling.yml"
+
+# number of points
+N_rho = 3000
+N_T = 3000
+
+# density in atom / cm3
+if N_rho == 1:
+    rho = np.array([1.])
+else:
+    #rho = np.array([1.])
+    rho = np.logspace(-6, 4, N_rho)
+
+# temperature in K
+T = np.logspace(1, 9, N_T)
 
-#print(us)
-#print(pconst)
+# time step
+dt = units.Myr * 1e-5
+dt = dt.to("s") / units.s
 
-print(type(params), type(us), type(pconst))
+
+
+# generate grid
+rho, T = np.meshgrid(rho, T)
+shape = rho.shape
+rho = deepcopy(rho.reshape(-1))
+T = T.reshape(-1)
+
+#
+# swift init
+#
+
+# parse swift params
+params = wrapper.parserReadFile(filename)
+# init units
+us, pconst = wrapper.unitSystemInit(params)
+# init cooling
 cooling = wrapper.coolingInit(params, us, pconst)
 
-print(cooling)
+
+#
+# Deal with units
+#
+
+# change units of rho and T
+# rho
+rho *= us.UnitLength_in_cgs**3 * pconst.const_proton_mass
+# specific energy
+energy = pconst.const_boltzmann_k * T / us.UnitTemperature_in_cgs
+energy /= (gamma - 1.) * pconst.const_proton_mass
+# time step
+dt /= us.UnitTime_in_cgs
+
+#
+# compute rate
+#
+
+# du / dt
+rate = wrapper.coolingRate(pconst, us, cooling,
+                           rho.astype(np.float32),
+                           energy.astype(np.float32),
+                           dt)
+
+#
+# plot
+#
+
+
+# change units => cgs
+energy *= us.UnitLength_in_cgs**2 / us.UnitTime_in_cgs**2
+
+rate *= us.UnitLength_in_cgs**2 / us.UnitTime_in_cgs**3
+    
+if N_rho == 1 or N_T == 1:
+    plt.figure()
+    n = rho / (pconst.const_proton_mass * us.UnitLength_in_cgs**3)
+    proton_mass_in_cgs = pconst.const_proton_mass * us.UnitMass_in_cgs
+    lambda_ = rate * proton_mass_in_cgs / (n)
+    
+    if N_rho == 1:
+        plt.loglog(T, np.abs(lambda_))
+        plt.xlabel("Temperature [K]")
+    else:
+        plt.loglog(rho, np.abs(lambda_))
+        plt.xlabel("Density [atom / cm3]")
+    plt.ylabel("Rate")
+    
+else:
+    cooling_time = energy / rate
+    cooling_length = np.sqrt(gamma * (gamma-1.) * energy) * cooling_time
+
+    cooling_length = np.log10(np.abs(cooling_length) / units.kpc.to('cm'))
+    
+    # reshape
+    rho = rho.reshape(shape)
+    T = T.reshape(shape)
+    energy = energy.reshape(shape)
+    cooling_length = cooling_length.reshape(shape)
+
+    _min = -7
+    _max = 7
+    N_levels = 100
+    levels = np.linspace(_min, _max, N_levels)
+    plt.figure()
+    plt.contourf(rho, T, cooling_length, levels)
+    plt.xlabel("Density [atom/cm3]")
+    plt.ylabel("Temperature [K]")
+
+    ax = plt.gca()
+    ax.set_xscale("log")
+    ax.set_yscale("log")
+    
+    cbar = plt.colorbar()
+    tc = np.arange(_min, _max, 1.5)
+    cbar.set_ticks(tc)
+    cbar.set_ticklabels(tc)
+
+plt.show()