From a0d11dce96b1e5729bde5ef11386bf928257bb50 Mon Sep 17 00:00:00 2001
From: lhausamm <loic_hausammann@hotmail.com>
Date: Fri, 8 Dec 2017 10:38:51 +0100
Subject: [PATCH] Cooling is working and able to reproduce grackle paper graphs
---
pyswiftsim/structure.py | 126 ++++++++++++++++++--------------------
setup.py | 37 +++++++----
src/cooling_wrapper.c | 95 ++++++++++++++++++++++++++---
src/cooling_wrapper.h | 2 +-
src/part_wrapper.c | 2 +-
src/pyswiftsim_tools.c | 26 ++++++++
src/pyswiftsim_tools.h | 17 ++++++
src/units_wrapper.c | 2 +-
src/wrapper.c | 9 ++-
test/cooling.yml | 51 ++++++++++++++++
test/test_cooling.py | 132 +++++++++++++++++++++++++++++++++++++---
11 files changed, 400 insertions(+), 99 deletions(-)
create mode 100644 test/cooling.yml
diff --git a/pyswiftsim/structure.py b/pyswiftsim/structure.py
index c183a04..6197a22 100644
--- a/pyswiftsim/structure.py
+++ b/pyswiftsim/structure.py
@@ -1,3 +1,5 @@
+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)
diff --git a/setup.py b/setup.py
index 285ff2f..3efb689 100644
--- a/setup.py
+++ b/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 = []
diff --git a/src/cooling_wrapper.c b/src/cooling_wrapper.c
index ce09541..b81e13f 100644
--- a/src/cooling_wrapper.c
+++ b/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;
+
}
diff --git a/src/cooling_wrapper.h b/src/cooling_wrapper.h
index 55dcec6..95717e8 100644
--- a/src/cooling_wrapper.h
+++ b/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__
diff --git a/src/part_wrapper.c b/src/part_wrapper.c
index 742d1a7..0fd3220 100644
--- a/src/part_wrapper.c
+++ b/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;
}
+
diff --git a/src/pyswiftsim_tools.c b/src/pyswiftsim_tools.c
index a01ede6..3cbad7f 100644
--- a/src/pyswiftsim_tools.c
+++ b/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;
+
+}
diff --git a/src/pyswiftsim_tools.h b/src/pyswiftsim_tools.h
index fb2a366..75f2737 100644
--- a/src/pyswiftsim_tools.h
+++ b/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__
diff --git a/src/units_wrapper.c b/src/units_wrapper.c
index 47ada02..1858c2b 100644
--- a/src/units_wrapper.c
+++ b/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);
}
diff --git a/src/wrapper.c b/src/wrapper.c
index 8c4b0ba..7ef79ae 100644
--- a/src/wrapper.c
+++ b/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)
diff --git a/test/cooling.yml b/test/cooling.yml
new file mode 100644
index 0000000..96555e6
--- /dev/null
+++ b/test/cooling.yml
@@ -0,0 +1,51 @@
+# 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).
diff --git a/test/test_cooling.py b/test/test_cooling.py
index 616b6af..e2a8c6f 100644
--- a/test/test_cooling.py
+++ b/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()
--
GitLab