diff --git a/docs/RST/examples/cooling.rst b/docs/RST/examples/cooling.rst
index 1cc6f4db5c44493cdf7792754084760d5b89954b..3f06576eef49c6208d0ce74caaad27d8f3f8c122 100644
--- a/docs/RST/examples/cooling.rst
+++ b/docs/RST/examples/cooling.rst
@@ -14,6 +14,10 @@ In the first mode, the code generates particles at different temperatures but sa
.. image:: https://obswww.unige.ch/~lhausamm/PySWIFTsim/examples/cooling_rate.png
In the second mode, the code generates a grid of particles at different density and temperatures and then computes the cooling.
+I suppose that in [Smith2016]_, they use the non equilibrium mode of Grackle to generate this graph.
+It requires a careful initialization of the element fractions which is not easy due to the non monotonic behavior of the cooling function.
+For simplicity, only Grackle with the equilibrium mode is shown in this documentation.
+
.. image:: https://obswww.unige.ch/~lhausamm/PySWIFTsim/examples/cooling_rate_2d.png
diff --git a/examples/cooling_box/cooling.yml b/examples/cooling_box/cooling.yml
index 415f22ce7fadc74f042ad91d30079aa8ce0ffbea..66fa4b67f5245eb33a07fae082c870ea3b47535f 100644
--- a/examples/cooling_box/cooling.yml
+++ b/examples/cooling_box/cooling.yml
@@ -16,7 +16,7 @@ GrackleCooling:
CloudyTable: CloudyData_UVB=HM2012.h5 # Name of the Cloudy Table
WithUVbackground: 1 # Enable or not the UV background
Redshift: 0 # Redshift to use (-1 means time based redshift)
- WithMetalCooling: 1 # Enable or not the metal cooling
+ WithMetalCooling: 0 # Enable or not the metal cooling
ProvideVolumetricHeatingRates: 0 # User provide volumetric heating rates
ProvideSpecificHeatingRates: 0 # User provide specific heating rates
SelfShieldingMethod: 0 # Grackle (<= 3) or Gear self shielding method
diff --git a/examples/cooling_box/plot_cooling.py b/examples/cooling_box/plot_cooling.py
index 13c2e4e514fd14b4f749a8a7d7f39a1e19414088..6d8485ec05c2f292d82c0772c08db97681bec21a 100644
--- a/examples/cooling_box/plot_cooling.py
+++ b/examples/cooling_box/plot_cooling.py
@@ -39,13 +39,13 @@ with_cosmo = False
# particle data
# in atom/cc
-part_density = np.array([1.])
-part_temperature = np.array([1e4]) # in K
+part_density = np.array([0.1])
+part_temperature = np.array([1e6]) # in K
gamma = 5./3.
# time in Myr
-t_end = 1.
-N_time = 100000
+t_end = 1e2
+N_time = 1000
# SCRIPT
@@ -86,23 +86,40 @@ def plot_solution(energy, data, us):
unit_length = float(us["UnitLength_in_cgs"])
unit_vel = float(us["UnitVelocity_in_cgs"])
unit_mass = float(us["UnitMass_in_cgs"])
+ unit_temp = float(us["UnitTemp_in_cgs"])
unit_time = unit_length / unit_vel
# change units => cgs
rho *= unit_mass / unit_length**3
- energy *= unit_length**2 / unit_time**2
+ m_p = constants.m_p.to("g") / (unit_mass * units.g)
+ rho = part_density * unit_length**3 * m_p
+ d["density"] = rho.to("")
- time *= unit_time
+ unit_energy = unit_mass * unit_length**2
+ unit_energy /= unit_time**2
+ k_B = constants.k_B.to("erg / K") / (unit_energy / unit_temp)
+ k_B *= units.K / units.erg
+ T = unit_temp * energy / k_B
+ T *= (gamma - 1.) * m_p
+
+ myr = units.s / units.Myr
+ time *= unit_time * float(myr.to(""))
# do plot
plt.figure()
date = strftime("%d %b %Y", gmtime())
plt.title("Date: {}".format(date))
- plt.plot(time, energy)
- plt.xlabel("Time [s]")
- plt.ylabel("Energy [erg]")
+ plt.plot(time, T)
+ plt.xlabel("Time [Myr]")
+ plt.ylabel("Temperature [K]")
+
+ ax = plt.gca()
+ ax.set_xscale("log")
+ ax.set_yscale("log")
+ ax.set_xlim([1e-2, 1e2])
+
plt.show()
diff --git a/examples/cooling_rate/cooling.yml b/examples/cooling_rate/cooling.yml
index 166fa0401a5870b941e503152d9843a8e2417e69..8fa0e720004ba4101d496108db29ef32f3722ef2 100644
--- a/examples/cooling_rate/cooling.yml
+++ b/examples/cooling_rate/cooling.yml
@@ -33,7 +33,7 @@ GrackleCooling:
ProvideVolumetricHeatingRates: 0 # User provide volumetric heating rates
ProvideSpecificHeatingRates: 0 # User provide specific heating rates
SelfShieldingMethod: 0 # Grackle (<= 3) or Gear self shielding method
- ConvergenceLimit: 1e-4
+ ConvergenceLimit: 1e-2
MaxSteps: 20000
Omega: 0.8
diff --git a/examples/cooling_rate/plot_cooling.py b/examples/cooling_rate/plot_cooling.py
index 55d21ee1790177aa1755a7cc2815cdb26414502a..8ea0e9ac66756031d151f467edf1bf48cf5c5a2c 100644
--- a/examples/cooling_rate/plot_cooling.py
+++ b/examples/cooling_rate/plot_cooling.py
@@ -26,12 +26,16 @@ import matplotlib.pyplot as plt
from astropy import units, constants
import yaml
from time import strftime, gmtime
+import sys
plt.rc('text', usetex=True)
downloadCoolingTable()
# PARAMETERS
+
+colormap = "RdBu"
+
# cosmology
with_cosmo = False
@@ -42,7 +46,7 @@ filename = "cooling.yml"
h_mass_frac = 0.76
# density in atom / cm3
-N_rho = 100
+N_rho = int(sys.argv[-1])
if N_rho == 1:
density = np.array([1.])
else:
@@ -169,7 +173,8 @@ def plot_solution(rate, data, us):
N_levels = 100
levels = np.linspace(_min, _max, N_levels)
plt.figure()
- plt.contourf(rho, T, cooling_length, levels)
+ plt.contourf(rho, T, cooling_length, levels,
+ cmap=plt.get_cmap(colormap))
plt.xlabel("Density [atom/cm3]")
plt.ylabel("Temperature [K]")
diff --git a/examples/cooling_rate/run.sh b/examples/cooling_rate/run.sh
index ca4ed090d7161d18ce520af16a81caaaebd04cf3..ce3d510430474c3986b463878c3edff2ba02d212 100644
--- a/examples/cooling_rate/run.sh
+++ b/examples/cooling_rate/run.sh
@@ -1,3 +1,4 @@
#!/bin/bash
-./plot_cooling.py
+./plot_cooling.py 1
+./plot_cooling.py 100
diff --git a/scripts/pyswift_cooling_rate b/scripts/pyswift_cooling_rate
new file mode 100644
index 0000000000000000000000000000000000000000..8b08db13d6e81a30b1cfbb89d2114e1909f44b4f
--- /dev/null
+++ b/scripts/pyswift_cooling_rate
@@ -0,0 +1,276 @@
+#!/usr/bin/env python3
+# ########################################################################
+# This file is part of PYSWIFT.
+# Copyright (c) 2018 Loic Hausammann (loic.hausammann@epfl.ch)
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Lesser General Public License as published
+# by the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this program. If not, see <http://www.gnu.org/licenses/>.
+# ########################################################################
+
+from pyswiftsim.libcooling import coolingRate
+from pyswiftsim import downloadCoolingTable
+
+from copy import deepcopy
+import numpy as np
+import matplotlib.pyplot as plt
+from astropy import units, constants
+import yaml
+from time import strftime, gmtime
+from argparse import ArgumentParser
+import os
+
+plt.rc('text', usetex=True)
+
+# PARAMETERS
+
+colormap = "RdBu"
+
+# cosmology
+with_cosmo = False
+
+
+def parseArguments():
+ parser = ArgumentParser(description="Plot the cooling rate")
+
+ parser.add_argument(
+ "--download_table", dest="table",
+ action="store_true",
+ help="Download the default cooling table")
+
+ parser.add_argument(
+ "--params", dest="params",
+ type=str, default=None,
+ help="SWIFT parameters file")
+
+ parser.add_argument(
+ "--T_min", dest="Tmin",
+ type=float, default=10.,
+ help="Minimal temperature to plot")
+
+ parser.add_argument(
+ "--T_max", dest="Tmax",
+ type=float, default=1e9,
+ help="Maximal temperature to plot")
+
+ parser.add_argument(
+ "--N_T", dest="N_T",
+ type=int, default=1000,
+ help="Number of points for the temperature")
+
+ parser.add_argument(
+ "--rho_min", dest="rho_min",
+ type=float, default=1e-6,
+ help="Minimal density to plot")
+
+ parser.add_argument(
+ "--rho_max", dest="rho_max",
+ type=float, default=1e4,
+ help="Maximal density to plot (only used if more than 1 point)")
+
+ parser.add_argument(
+ "--N_rho", dest="N_rho",
+ type=int, default=1,
+ help="Number of points for the density")
+
+ parser.add_argument(
+ "--gamma", dest="gamma",
+ type=float, default=5. / 3.,
+ help="Adiabatic index")
+
+ parser.add_argument(
+ "--h_frac", dest="h_mass_frac",
+ type=float, default=0.76,
+ help="hydogen mass fraction")
+
+ args = parser.parse_args()
+
+ return args
+
+
+def mu(T, h_mass_frac):
+ # define constants
+ T_transition = 1e4
+ mu_neutral = 4. / (1. + 3. * h_mass_frac)
+ mu_ion = 4. / (8. - 5. * (1 - h_mass_frac))
+
+ # Find correct mu
+ ind = T > T_transition
+ mu = np.zeros(T.shape)
+ mu[ind] = mu_ion
+ mu[~ind] = mu_neutral
+ return mu
+
+
+# SCRIPT
+def generate_initial_condition(us, opt):
+ print("Generating default initial conditions")
+ # get units
+ unit_length = float(us["UnitLength_in_cgs"])
+ unit_vel = float(us["UnitVelocity_in_cgs"])
+ unit_mass = float(us["UnitMass_in_cgs"])
+ unit_temp = float(us["UnitTemp_in_cgs"])
+
+ # density in atom / cm3
+ N_rho = opt.N_rho
+ if N_rho == 1:
+ density = np.array([opt.rho_min])
+ else:
+ rho_min = np.log10(opt.rho_min)
+ rho_max = np.log10(opt.rho_max)
+ density = np.logspace(rho_min, rho_max, N_rho)
+
+ # temperature in K
+ N_T = opt.N_T
+ T_min = np.log10(opt.Tmin)
+ T_max = np.log10(opt.Tmax)
+ temperature = np.logspace(T_min, T_max, N_T)
+
+ # time step in s
+ dt = units.Myr * 1e-8
+ dt = dt.to("s") / units.s
+
+ d = {}
+ # generate grid
+ rho, T = np.meshgrid(density, temperature)
+ rho = deepcopy(rho.reshape(-1))
+ T = T.reshape(-1)
+ d["temperature"] = T
+
+ # Deal with units
+ m_p = constants.m_p.to("g") / (unit_mass * units.g)
+ rho = rho * unit_length**3 * m_p
+ d["density"] = rho.to("")
+
+ unit_time = unit_length / unit_vel
+ unit_energy = unit_mass * unit_length**2
+ unit_energy /= unit_time**2
+ k_B = constants.k_B.to("erg / K") / (unit_energy / unit_temp)
+ k_B *= units.K / units.erg
+ energy = k_B * T / (unit_temp * mu(T, opt.h_mass_frac))
+ energy /= (opt.gamma - 1.) * m_p
+ d["energy"] = energy.to("")
+
+ time_step = dt / unit_time
+ d["time_step"] = time_step
+
+ return d
+
+
+def getParser(filename):
+ with open(filename, "r") as stream:
+ stream = yaml.load(stream)
+ return stream
+
+
+def plot_solution(rate, data, us, opt):
+ print("Plotting solution")
+ energy = data["energy"]
+ rho = data["density"]
+ T = data["temperature"]
+
+ T *= float(us["UnitTemp_in_cgs"])
+
+ # change units => cgs
+ u_m = float(us["UnitMass_in_cgs"])
+ u_l = float(us["UnitLength_in_cgs"])
+ u_v = float(us["UnitVelocity_in_cgs"])
+ u_t = u_l / u_v
+
+ rho *= u_m / u_l**3
+
+ energy *= u_v**2
+
+ rate *= u_v**2 / u_t
+
+ # lambda cooling
+ lam = rate * rho
+
+ N_rho = opt.N_rho
+ N_T = opt.N_T
+
+ # do plot
+ if N_rho == 1:
+ # plot Lambda vs T
+ plt.figure()
+ plt.loglog(T, np.abs(lam),
+ label="SWIFT")
+ plt.xlabel("Temperature [K]")
+ plt.ylabel("$\\Lambda$ [erg s$^{-1}$ cm$^{3}$]")
+
+ else:
+ m_p = constants.m_p.to("g") / units.g
+ rho /= m_p
+
+ shape = [N_rho, N_T]
+ cooling_time = energy / rate
+ cooling_length = np.sqrt(opt.gamma * (opt.gamma-1.) * energy)
+ cooling_length *= 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,
+ cmap=plt.get_cmap(colormap))
+ 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)
+ cbar.ax.set_ylabel("Log( Cooling Length / kpc )")
+
+ date = strftime("%d %b %Y", gmtime())
+ plt.title("Date: {}".format(date))
+ plt.show()
+
+
+if __name__ == "__main__":
+
+ opt = parseArguments()
+
+ if opt.table:
+ downloadCoolingTable()
+
+ if not os.path.isfile(opt.params):
+ raise Exception(
+ "The parameter file '{}' does not exist".format(opt.params))
+
+ parser = getParser(opt.params)
+ us = parser["InternalUnitSystem"]
+
+ d = generate_initial_condition(us, opt)
+
+ # du / dt
+ print("Computing cooling...")
+
+ rate = coolingRate(opt.params,
+ d["density"].astype(np.float32),
+ d["energy"].astype(np.float32),
+ with_cosmo, d["time_step"])
+
+ plot_solution(rate, d, us, opt)
diff --git a/setup.py b/setup.py
index 1c76bc94e8463de8a2bf4af34c8a9c61a2505365..833eb7af74f1d914984ccee2d1389d4a8780a1e5 100644
--- a/setup.py
+++ b/setup.py
@@ -118,15 +118,22 @@ data_files = []
##############
c_src = glob("src/*.c")
-ext_modules = Extension("pyswiftsim.libcooling",
- c_src,
- include_dirs=include,
- libraries=lib,
- library_dirs=lib_dir,
- extra_compile_args=cflags,
- extra_link_args=lflags)
-
-ext_modules = [ext_modules]
+ext_modules = []
+ext_modules.append(Extension("pyswiftsim.libcooling",
+ c_src,
+ include_dirs=include,
+ libraries=lib,
+ library_dirs=lib_dir,
+ extra_compile_args=cflags,
+ extra_link_args=lflags))
+
+ext_modules.append(Extension("pyswiftsim.libstellar",
+ c_src,
+ include_dirs=include,
+ libraries=lib,
+ library_dirs=lib_dir,
+ extra_compile_args=cflags,
+ extra_link_args=lflags))
##############
# data
@@ -138,7 +145,7 @@ data_files = []
# scripts
##############
-list_scripts = []
+list_scripts = glob("scripts/*")
##############
# Setup
diff --git a/src/cooling_wrapper.c b/src/cooling_wrapper.c
index 47dd1ccb52857d1794719eaac9ba2bdcedded1c2..dc60a7ae57c7ee507d35178687f867d4a5e68b8d 100644
--- a/src/cooling_wrapper.c
+++ b/src/cooling_wrapper.c
@@ -117,44 +117,14 @@ PyObject* pycooling_rate(PyObject* self, PyObject* args) {
size_t N = PyArray_DIM(energy, 0);
- /* Initialize parser */
- struct swift_params params;
- parser_read_file(filename, ¶ms);
-
- /* Init unit system */
- struct unit_system us;
- units_init_from_params(&us, ¶ms, "InternalUnitSystem");
-
- /* Init physical constant */
- struct phys_const pconst;
- phys_const_init(&us, ¶ms, &pconst);
-
- /* init cooling */
- if (initialized == 0) {
- // struct cooling_function_data cooling;
- cooling_init_backend(¶ms, &us, &pconst, &cooling);
- initialized = 1;
- }
- else if (initialized == 1) {
- message("WARNING: not reinitializing the cooling, need to be fixed");
- initialized = 2;
- }
+ PYSWIFTSIM_INIT_STRUCTS();
+
+ PYSWIFTSIM_INIT_COOLING();
/* Init chemistry */
struct chemistry_global_data chemistry;
chemistry_init_backend(¶ms, &us, &pconst, &chemistry);
- /* Init cosmo */
- struct cosmology cosmo;
-
- if (with_cosmo)
- cosmology_init(¶ms, &us, &pconst, &cosmo);
- else
- cosmology_init_no_cosmo(&cosmo);
-
- /* Init hydro prop */
- struct hydro_props hydro_props;
- hydro_props_init(&hydro_props, &pconst, &us, ¶ms);
/* Init entropy floor */
struct entropy_floor_properties floor_props;
@@ -177,12 +147,17 @@ PyObject* pycooling_rate(PyObject* self, PyObject* args) {
float u = *(float*) PyArray_GETPTR1(energy, i);
hydro_set_physical_internal_energy(&p, &xp, &cosmo, u);
+ chemistry_first_init_part(&pconst, &us, &cosmo, &chemistry, &p, &xp);
+
if (fractions != NULL)
pycooling_set_fractions(&xp, fractions, i);
- else
- cooling_first_init_part(&pconst, &us, &cosmo, &cooling, &p, &xp);
+ else {
+ /* Need to set the mass in order to estimate the element fractions */
+ hydro_set_mass(&p, p.rho);
+ p.h = 1;
- chemistry_first_init_part(&pconst, &us, &cosmo, &chemistry, &p, &xp);
+ cooling_first_init_part(&pconst, &us, &cosmo, &cooling, &p, &xp);
+ }
hydro_set_physical_internal_energy_dt(&p, &cosmo, 0);
diff --git a/src/cooling_wrapper.h b/src/cooling_wrapper.h
index bd9017f8abb4c31084db4ace60fa52372f943006..e0b97a9c9f70d9033dfaac54ae75bb061f8dfb35 100644
--- a/src/cooling_wrapper.h
+++ b/src/cooling_wrapper.h
@@ -28,5 +28,20 @@ extern int initialized;
PyObject* pycooling_rate(PyObject* self, PyObject* args);
+#define PYSWIFTSIM_INIT_COOLING() \
+ \
+ /* init cooling */ \
+ if (initialized == 0) { \
+ /* struct cooling_function_data cooling; */ \
+ cooling_init_backend(¶ms, &us, &pconst, &cooling); \
+ initialized = 1; \
+ } \
+ else if (initialized == 1) { \
+ message("WARNING: not reinitializing the cooling, need to be fixed"); \
+ initialized = 2; \
+ } \
+ \
+
+
#endif // __PYSWIFTSIM_COOLING_H__
diff --git a/src/libstellar.c b/src/libstellar.c
new file mode 100644
index 0000000000000000000000000000000000000000..a42feb1a341d7fa25da44bbb00d8cc0cdf7e8f1e
--- /dev/null
+++ b/src/libstellar.c
@@ -0,0 +1,72 @@
+/*******************************************************************************
+ * This file is part of PYSWIFTSIM.
+ * Copyright (c) 2018 loic hausammann (loic.hausammann@epfl.ch)
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published
+ * by the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program. If not, see <http://www.gnu.org/licenses/>.
+ *
+ ******************************************************************************/
+#include "stellar_evolution_wrapper.h"
+
+#include <Python.h>
+#include <math.h>
+#include <numpy/arrayobject.h>
+
+/* definition of the method table */
+
+static PyMethodDef libstellar_methods[] = {
+
+ {"initialMassFunction", pyinitial_mass_function, METH_VARARGS,
+ "TODO\n"
+ },
+
+ {"lifetimeFromMass", pylifetime_from_mass, METH_VARARGS,
+ "TODO\n"
+ },
+
+ {"massFromLifetime", pymass_from_lifetime, METH_VARARGS,
+ "TODO\n"
+ },
+
+
+ {NULL, NULL, 0, NULL} /* Sentinel */
+};
+
+
+
+static struct PyModuleDef libstellar_cmodule = {
+ PyModuleDef_HEAD_INIT,
+ "libstellar",
+ "Wrapper around the stellar physics of SWIFT",
+ -1,
+ libstellar_methods,
+ NULL, /* m_slots */
+ NULL, /* m_traverse */
+ NULL, /* m_clear */
+ NULL, /* m_free */
+};
+
+
+PyMODINIT_FUNC PyInit_libstellar(void)
+{
+ PyObject *m;
+
+ /* set time for swift */
+ clocks_set_cpufreq(0);
+
+ import_array();
+ Py_Initialize();
+ m = PyModule_Create(&libstellar_cmodule);
+
+ return m;
+}
diff --git a/src/pyswiftsim_tools.h b/src/pyswiftsim_tools.h
index 567314f028f17bf34e395ea77292763d1373447e..6b75075c21bc10bd6399cce584e128c83adb54ee 100644
--- a/src/pyswiftsim_tools.h
+++ b/src/pyswiftsim_tools.h
@@ -44,4 +44,32 @@ enum error_code {
int pytools_check_array(PyArrayObject *obj, int dim, int type, const char* name);
+#define PYSWIFTSIM_INIT_STRUCTS() \
+ \
+ /* Initialize parser */ \
+ struct swift_params params; \
+ parser_read_file(filename, ¶ms); \
+ \
+ /* Init unit system */ \
+ struct unit_system us; \
+ units_init_from_params(&us, ¶ms, "InternalUnitSystem"); \
+ \
+ /* Init physical constant */ \
+ struct phys_const pconst; \
+ phys_const_init(&us, ¶ms, &pconst); \
+ \
+ /* Init cosmo */ \
+ struct cosmology cosmo; \
+ \
+ if (with_cosmo) \
+ cosmology_init(¶ms, &us, &pconst, &cosmo); \
+ else \
+ cosmology_init_no_cosmo(&cosmo); \
+ \
+ /* Init hydro prop */ \
+ struct hydro_props hydro_props; \
+ hydro_props_init(&hydro_props, &pconst, &us, ¶ms); \
+
+
+
#endif // __PYSWIFTSIM_TOOLS_H__
diff --git a/src/stellar_evolution_wrapper.c b/src/stellar_evolution_wrapper.c
new file mode 100644
index 0000000000000000000000000000000000000000..a4639ea4910cc79a72eda2251e8f20400544c28c
--- /dev/null
+++ b/src/stellar_evolution_wrapper.c
@@ -0,0 +1,191 @@
+/*******************************************************************************
+ * This file is part of PYSWIFTSIM.
+ * Copyright (c) 2018 loic hausammann (loic.hausammann@epfl.ch)
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published
+ * by the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program. If not, see <http://www.gnu.org/licenses/>.
+ *
+ ******************************************************************************/
+#include "stellar_evolution_wrapper.h"
+#include "pyswiftsim_tools.h"
+
+/**
+ * @brief Compute the initial mass function.
+ *
+ * args is expecting:
+ * TODO
+ *
+ * @param self calling object
+ * @param args arguments
+ * @return cooling rate
+ */
+PyObject* pyinitial_mass_function(PyObject* self, PyObject* args) {
+ import_array();
+
+
+ char *filename;
+
+ PyArrayObject *mass;
+
+ /* parse argument */
+ if (!PyArg_ParseTuple(
+ args, "sO", &filename, &mass))
+ return NULL;
+
+ /* check numpy array */
+ if (pytools_check_array(mass, 1, NPY_FLOAT, "Mass") != SWIFT_SUCCESS)
+ return NULL;
+
+ size_t N = PyArray_DIM(mass, 0);
+
+ int with_cosmo = 0;
+ PYSWIFTSIM_INIT_STRUCTS();
+
+ /* Init stellar physics */
+ struct feedback_props fp;
+ feedback_props_init(&fp, &pconst, &hydro_props, &us, ¶ms, &cosmo);
+
+ /* return object */
+ PyArrayObject *imf = (PyArrayObject *)PyArray_SimpleNew(PyArray_NDIM(mass), PyArray_DIMS(mass), NPY_FLOAT);
+
+ for(size_t i = 0; i < N; i++) {
+ float m = *(float*) PyArray_GETPTR1(mass, i);
+ float *imf_i = (float*) PyArray_GETPTR1(imf, i);
+
+
+ *imf_i = stellar_evolution_get_imf(&fp.stellar_model.imf, m);
+ }
+
+ return (PyObject *) imf;
+
+}
+
+/**
+ * @brief Compute the lifetime of a star.
+ *
+ * args is expecting:
+ * TODO
+ *
+ * @param self calling object
+ * @param args arguments
+ * @return cooling rate
+ */
+PyObject* pylifetime_from_mass(PyObject* self, PyObject* args) {
+ import_array();
+
+ char *filename;
+
+ PyArrayObject *mass;
+ PyArrayObject *metallicity;
+
+ /* parse argument */
+ if (!PyArg_ParseTuple(
+ args, "sOO", &filename, &mass, &metallicity))
+ return NULL;
+
+ /* check numpy array */
+ if (pytools_check_array(mass, 1, NPY_FLOAT, "Mass") != SWIFT_SUCCESS)
+ return NULL;
+ if (pytools_check_array(metallicity, 1, NPY_FLOAT, "Metallicity") != SWIFT_SUCCESS)
+ return NULL;
+
+ if (PyArray_DIM(metallicity, 0) != PyArray_DIM(mass, 0))
+ error("Mass and metallicity should have the same dimension");
+
+
+ size_t N = PyArray_DIM(mass, 0);
+
+ int with_cosmo = 0;
+
+ PYSWIFTSIM_INIT_STRUCTS();
+
+ /* Init stellar physics */
+ struct feedback_props fp;
+ feedback_props_init(&fp, &pconst, &hydro_props, &us, ¶ms, &cosmo);
+
+ /* return object */
+ PyArrayObject *lifetime = (PyArrayObject *)PyArray_SimpleNew(PyArray_NDIM(mass), PyArray_DIMS(mass), NPY_FLOAT);
+
+ for(size_t i = 0; i < N; i++) {
+ float m = *(float*) PyArray_GETPTR1(mass, i);
+ float z = *(float*) PyArray_GETPTR1(metallicity, i);
+ float *tau = (float*) PyArray_GETPTR1(lifetime, i);
+
+
+ *tau = stellar_evolution_get_log_lifetime_from_mass(&fp.stellar_model.imf, log10(m), z);
+ *tau = pow(10, *tau);
+ }
+
+ return (PyObject *) lifetime;
+
+}
+
+/**
+ * @brief Compute the mass of a star with a given lifetime.
+ *
+ * args is expecting:
+ * TODO
+ *
+ * @param self calling object
+ * @param args arguments
+ * @return cooling rate
+ */
+PyObject* pymass_from_lifetime(PyObject* self, PyObject* args) {
+ import_array();
+
+ char *filename;
+
+ PyArrayObject *time;
+ PyArrayObject *metallicity;
+
+ /* parse argument */
+ if (!PyArg_ParseTuple(
+ args, "sOO", &filename, &time, &metallicity))
+ return NULL;
+
+ /* check numpy array */
+ if (pytools_check_array(time, 1, NPY_FLOAT, "Lifetime") != SWIFT_SUCCESS)
+ return NULL;
+ if (pytools_check_array(metallicity, 1, NPY_FLOAT, "Metallicity") != SWIFT_SUCCESS)
+ return NULL;
+
+ if (PyArray_DIM(metallicity, 0) != PyArray_DIM(time, 0))
+ error("Lifetime and metallicity should have the same dimension");
+
+
+ size_t N = PyArray_DIM(time, 0);
+
+ int with_cosmo = 0;
+
+ PYSWIFTSIM_INIT_STRUCTS();
+
+ /* Init stellar physics */
+ struct feedback_props fp;
+ feedback_props_init(&fp, &pconst, &hydro_props, &us, ¶ms, &cosmo);
+
+ /* return object */
+ PyArrayObject *mass = (PyArrayObject *)PyArray_SimpleNew(PyArray_NDIM(time), PyArray_DIMS(time), NPY_FLOAT);
+
+ for(size_t i = 0; i < N; i++) {
+ float tau = *(float*) PyArray_GETPTR1(time, i);
+ float z = *(float*) PyArray_GETPTR1(metallicity, i);
+ float *m = (float*) PyArray_GETPTR1(mass, i);
+
+
+ *m = stellar_evolution_get_log_mass_from_lifetime(&fp.stellar_model.imf, log10(tau), z);
+ *m = pow(10, *m);
+ }
+
+ return (PyObject *) mass;
+
+}
diff --git a/src/stellar_evolution_wrapper.h b/src/stellar_evolution_wrapper.h
new file mode 100644
index 0000000000000000000000000000000000000000..a76b04d39b1733a8667cc9435859cfeb9fbfbeea
--- /dev/null
+++ b/src/stellar_evolution_wrapper.h
@@ -0,0 +1,29 @@
+/*******************************************************************************
+ * This file is part of PYSWIFTSIM.
+ * Copyright (c) 2018 loic hausammann (loic.hausammann@epfl.ch)
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published
+ * by the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program. If not, see <http://www.gnu.org/licenses/>.
+ *
+ ******************************************************************************/
+#ifndef __PYSWIFTSIM_STELLAR_H__
+#define __PYSWIFTSIM_STELLAR_H__
+
+#include "pyswiftsim_includes.h"
+
+PyObject* pyinitial_mass_function(PyObject* self, PyObject* args);
+PyObject* pylifetime_from_mass(PyObject* self, PyObject* args);
+PyObject* pymass_from_lifetime(PyObject* self, PyObject* args);
+
+#endif // __PYSWIFTSIM_STELLAR_H__
+