Test reproduce grackle graph (fig. 8 in Grackle: a Chemistry and Cooling Library for Astrophysics)

pyswiftsim/pyswiftsim/cooling.pxd

@@ -11,14 +11,22 @@ cdef extern from "parser.h":
 
 cdef extern from "units.h":
     struct unit_system:
-        pass
+        double UnitMass_in_cgs;
+        double UnitLength_in_cgs;
+        double UnitTime_in_cgs;
+        double UnitCurrent_in_cgs;
+        double UnitTemperature_in_cgs;
     void units_init(unit_system* us, const swift_params* params,
                     const char* category)
     void units_print_backend(const unit_system *us)
 
-    
+cdef extern from "adiabatic_index.h":
+    float hydro_gamma
+
 cdef extern from "physical_constants.h":
     struct phys_const:
+        double const_boltzmann_k
+        double const_proton_mass
         pass
     void phys_const_init(unit_system* us, phys_const* internal_const);
 

pyswiftsim/pyswiftsim/cooling.pyx

@@ -9,6 +9,7 @@ from pyswiftsim.part_def cimport hydro_space
 
 from pyswiftsim cimport equation_of_state
 from pyswiftsim.equation_of_state cimport gas_entropy_from_internal_energy
+from pyswiftsim.equation_of_state cimport gas_soundspeed_from_internal_energy
 
 from libc.stdlib cimport malloc, free
 from libc.math cimport M_PI
@@ -17,6 +18,9 @@ import numpy as np
 
 ctypedef np.float_t DTYPE_t
 
+_hydro_gamma = hydro_gamma
+
+
 def read_params(str filename):
     # transform python to C
     tmp = filename.encode(u"ascii")
@@ -27,18 +31,18 @@ def read_params(str filename):
     # init params
     cdef swift_params *params = <swift_params*> malloc(sizeof(swift_params));
     parser_read_file(cfilename, params)
-    parser_print_params(params)
+    #parser_print_params(params)
  
     
     return pointer_create(
         params,
-        pointer.swift_params
+        pointer.type_swift_params
     )
 
 def init_units(Pointer p_params):
 
     # deal with inputs
-    pointer_check(p_params, pointer.swift_params)
+    pointer_check(p_params, pointer.type_swift_params)
     cdef swift_params *params = <swift_params*> p_params._data
     
     # init units
@@ -46,7 +50,7 @@ def init_units(Pointer p_params):
     cdef const char* category = "InternalUnitSystem"
     units_init(units, params, category)
 
-    units_print_backend(units)
+    #units_print_backend(units)
 
     # init constants
     cdef phys_const *constants = <phys_const*> malloc(sizeof(phys_const));
@@ -54,33 +58,33 @@ def init_units(Pointer p_params):
 
     # return
     d = {}
-    d["constants"] = pointer_create(constants, pointer.phys_const)
-    d["units"] = pointer_create(units, pointer.unit_system)
+    d["constants"] = pointer_create(constants, pointer.type_phys_const)
+    d["units"] = pointer_create(units, pointer.type_unit_system)
     return d
 
 def pycooling_init(Pointer p_params,
                    Pointer p_units,
                    Pointer p_constants):
     # deal with inputs
-    pointer_check(p_params, pointer.swift_params)
+    pointer_check(p_params, pointer.type_swift_params)
     cdef swift_params *params = <swift_params*> p_params._data
 
-    pointer_check(p_units, pointer.unit_system)
+    pointer_check(p_units, pointer.type_unit_system)
     cdef unit_system *units = <unit_system*> p_units._data
 
-    pointer_check(p_constants, pointer.phys_const)
+    pointer_check(p_constants, pointer.type_phys_const)
     cdef phys_const *constants = <phys_const*> p_constants._data
 
-    units_print_backend(units)
+    #units_print_backend(units)
     # init cooling
     cdef cooling_function_data *cooling = <cooling_function_data*> malloc(sizeof(cooling_function_data));
     cooling_init(params, units, constants, cooling)
 
     # print results
-    cooling_print_backend(cooling)
+    #cooling_print_backend(cooling)
 
     # store and return results
-    return pointer_create(cooling, pointer.cooling_function_data)
+    return pointer_create(cooling, pointer.type_cooling_function_data)
 
 
 def pycooling_rate(Pointer p_units,
@@ -97,15 +101,15 @@ def pycooling_rate(Pointer p_units,
     
     # deal with inputs
     # p_units
-    pointer_check(p_units, pointer.unit_system)
+    pointer_check(p_units, pointer.type_unit_system)
     cdef unit_system *units = <unit_system*> p_units._data
 
     # p_cooling
-    pointer_check(p_cooling, pointer.cooling_function_data)
+    pointer_check(p_cooling, pointer.type_cooling_function_data)
     cdef cooling_function_data *cooling = <cooling_function_data*> p_cooling._data
 
     # p_constants
-    pointer_check(p_constants, pointer.phys_const)
+    pointer_check(p_constants, pointer.type_phys_const)
     cdef phys_const *constants = <phys_const*> p_constants._data
 
     # initialize particles
@@ -133,3 +137,19 @@ def pycooling_rate(Pointer p_units,
 
         
     return rate
+
+
+def soundspeed_from_internal_energy(np.ndarray [DTYPE_t, ndim=1] rho,
+                                    np.ndarray [DTYPE_t, ndim=1] u):
+    
+
+    if rho.shape[0] != u.shape[0]:
+        raise Exception("Rho and u should be of the same size!")
+    
+    cdef int N = rho.shape[0]
+    cdef np.ndarray[DTYPE_t, ndim=1] cs = np.empty_like(rho)
+
+    for i in range(N):
+        cs[i] = gas_soundspeed_from_internal_energy(rho[i], u[i])
+
+    return cs

pyswiftsim/pyswiftsim/equation_of_state.pxd

 cdef extern from "equation_of_state.h":
     cdef float gas_entropy_from_internal_energy(float density, float u)
+    cdef float gas_soundspeed_from_entropy(float density, float entropy)
+    cdef float gas_soundspeed_from_internal_energy(float density, float u)

pyswiftsim/pyswiftsim/pointer.pxd

@@ -3,11 +3,11 @@ cdef extern from "clocks.h":
 
     
 cdef enum pointer_type:
-    swift_params
-    unit_system
-    phys_const
-    cooling_function_data
-    count
+    type_swift_params
+    type_unit_system
+    type_phys_const
+    type_cooling_function_data
+    type_count
 
 cdef const char **pointer_name
 

pyswiftsim/pyswiftsim/pointer.pyx

 from libc.stdlib cimport free
+from cooling cimport unit_system, phys_const
 
 clocks_set_cpufreq(0); # estimate automatically cpufreq and init time
 
@@ -27,6 +28,15 @@ cdef class Pointer:
     def __del__(self):
         free(self._data)
 
+    def __getattr__(self, value):
+        data_type = self._data_type
+        if data_type == type_unit_system:
+            return pointer_unit_system_getattr(self, value)
+        if data_type == type_phys_const:
+            return pointer_phys_const_getattr(self, value)
+        else:
+            raise Exception("Data type not implemented")
+
 
 # method to call when creating a pointer
 cdef pointer_create(void *params, int data_type):
@@ -43,3 +53,30 @@ cdef pointer_check(Pointer ptr, int data_type):
         raise Exception("Value not set for pointer type %s" % data_name)
     
 
+cdef pointer_unit_system_getattr(Pointer ptr, str attr):
+    pointer_check(ptr, type_unit_system)
+    cdef unit_system us = (<unit_system*> ptr._data)[0]
+
+    if attr == "UnitMass_in_cgs":
+        return us.UnitMass_in_cgs
+    elif attr == "UnitLength_in_cgs":
+        return us.UnitLength_in_cgs
+    elif attr == "UnitTime_in_cgs":
+        return us.UnitTime_in_cgs
+    elif attr == "UnitCurrent_in_cgs":
+        return us.UnitCurrent_in_cgs
+    elif attr == "UnitTemperature_in_cgs":
+        return us.UnitTemperature_in_cgs
+    else:
+        raise Exception("unit_system does not contain %s" % attr)
+
+cdef pointer_phys_const_getattr(Pointer ptr, str attr):
+    pointer_check(ptr, type_phys_const)
+    cdef phys_const consts = (<phys_const*> ptr._data)[0]
+
+    if attr == "const_boltzmann_k":
+        return consts.const_boltzmann_k
+    if attr == "const_proton_mass":
+        return consts.const_proton_mass
+    else:
+        raise Exception("phys_const does not contain %s" % attr)

pyswiftsim/test/test_cooling.py

+#!/usr/bin/env python3
+from pyswiftsim import pointer
+from pyswiftsim import cooling
+
+import astropy.units
+
+import matplotlib.pyplot as plt
+
+import numpy as np
+
+# number of particles
+N  = 100
+# params file
+params = "/home/loikki/swift_test/cooling_box/coolingBox.yml"
+
+# read params
+params = cooling.read_params(params)
+#print(params)
+
+# init units
+d = cooling.init_units(params)
+units = d["units"]
+consts = d["constants"]
+
+# init cooling
+cooling_data = cooling.pycooling_init(params, units, consts)
+
+# Init variables
+rho = np.logspace(-6, 4, N).astype(np.float) # in cm-3
+rho *= consts.const_proton_mass * units.UnitLength_in_cgs**3
+T = np.logspace(1,9, N).astype(np.float) # in K
+
+rho, T = np.meshgrid(rho, T)
+shape = rho.shape
+rho = rho.reshape(np.prod(shape))
+T = T.reshape(np.prod(shape))
+
+# u = k_b T / (gamma - 1) mu m_p
+mu = 2
+u = consts.const_boltzmann_k * T / ((cooling._hydro_gamma - 1.) * mu * consts.const_proton_mass)
+
+# compute cooling
+rate = cooling.pycooling_rate(units, cooling_data, consts, rho, u)
+
+cs = cooling.soundspeed_from_internal_energy(rho, u)
+
+rho = rho.reshape(shape) / (units.UnitLength_in_cgs**3 * consts.const_proton_mass)
+T = T.reshape(shape)
+# du / dt  [e / s] = [kg m2 / s3]
+rate = rate.reshape(shape) * units.UnitMass_in_cgs * units.UnitLength_in_cgs**2 / units.UnitTime_in_cgs**3
+cs = cs.reshape(shape) * units.UnitLength_in_cgs / units.UnitTime_in_cgs
+
+L_cool = -rate * cs / astropy.units.kpc.to("cm")
+
+plt.title("Cooling Length")
+plt.contourf(rho, T, np.log(L_cool))
+plt.xlabel("Density [$cm^{-3}$]")
+plt.ylabel("Temperature [K]")
+plt.colorbar()
+
+ax = plt.gca()
+ax.set_xscale('log')
+ax.set_yscale('log')
+
+plt.show()