Add cooling_rate scripts, start implementing stellar physics

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
 

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

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()
 
 

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
 

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]")
 

examples/cooling_rate/run.sh

 #!/bin/bash
 
-./plot_cooling.py
+./plot_cooling.py 1
+./plot_cooling.py 100

scripts/pyswift_cooling_rate

+#!/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)

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

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, &params);
-
-  /* Init unit system */
-  struct unit_system us;
-  units_init_from_params(&us, &params, "InternalUnitSystem");
-
-  /* Init physical constant */
-  struct phys_const pconst;
-  phys_const_init(&us, &params, &pconst);
-
-  /* init cooling */
-  if (initialized == 0) {
-    // struct cooling_function_data cooling;
-    cooling_init_backend(&params, &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(&params, &us, &pconst, &chemistry);
 
-  /* Init cosmo */
-  struct cosmology cosmo;
-
-  if (with_cosmo)
-    cosmology_init(&params, &us, &pconst, &cosmo);
-  else
-    cosmology_init_no_cosmo(&cosmo);
-
-  /* Init hydro prop */
-  struct hydro_props hydro_props;
-  hydro_props_init(&hydro_props, &pconst, &us, &params);
 
   /* 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);
 

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(&params, &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__
 

src/libstellar.c

+/*******************************************************************************
+ * 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;
+}

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, &params);				\
+								\
+  /* Init unit system */					\
+  struct unit_system us;					\
+  units_init_from_params(&us, &params, "InternalUnitSystem");	\
+								\
+  /* Init physical constant */					\
+  struct phys_const pconst;					\
+  phys_const_init(&us, &params, &pconst);			\
+								\
+  /* Init cosmo */						\
+  struct cosmology cosmo;					\
+								\
+  if (with_cosmo)						\
+    cosmology_init(&params, &us, &pconst, &cosmo);		\
+  else								\
+    cosmology_init_no_cosmo(&cosmo);				\
+								\
+  /* Init hydro prop */						\
+  struct hydro_props hydro_props;				\
+  hydro_props_init(&hydro_props, &pconst, &us, &params);	\
+  
+
+
 #endif // __PYSWIFTSIM_TOOLS_H__

src/stellar_evolution_wrapper.c

+/*******************************************************************************
+ * 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, &params, &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, &params, &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, &params, &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;
+  
+}

src/stellar_evolution_wrapper.h

+/*******************************************************************************
+ * 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__
+