Add python tools for neutrinos

doc/RTD/source/Neutrinos/index.rst

@@ -32,7 +32,8 @@ neutrino mass specified in the cosmology and generate new velocities
 based on the homogeneous (unperturbed) Fermi-Dirac distribution. In
 this case, placeholder neutrino particles should be provided in the
 initial conditions with arbitrary masses and velocities, distributed
-uniformly in the box.
+uniformly in the box. Placeholders can be spawned with the python
+script ``tools/spawn_neutrinos.py``.
 
 Relativistic Drift
 ------------------
@@ -113,7 +114,8 @@ neutrinos respectively. It is recommended to store the units of the wavenumbers
 as an attribute at "Units/Unit length in cgs (U_L)". The ``fixed_bg_density``
 flag determines whether the linear response scales as :math:`\Omega_\nu(a)`
 or the present-day value :math:`\Omega_{\nu,0}`, either of which may be
-appropriate depending on the particle initial conditions.
+appropriate depending on the particle initial conditions. An HDF5 file
+can be generated using classy with the script ``tools/create_perturb_file.py``.
 
 The linear response mode currently only supports degenerate mass models
 with a single neutrino transfer function.

tools/create_perturb_file.py

+#!/usr/bin/env python3
+
+# Utility to create an HDF5 file with cosmological transfer functions using
+# CLASS (requires python packages classy and h5py).
+
+import numpy as np
+import h5py
+from classy import Class
+
+#Filename of the perturbations file to be create
+fname = "perturbations.hdf5"
+
+#Open the file
+f = h5py.File("perturbations.hdf5", mode="w")
+
+#Cosmological parameters
+h = 0.681
+Omega_b = 0.0486
+Omega_cdm = 0.2560110606
+A_s = 2.0993736148e-09
+n_s = 0.967
+
+#Neutrino and radiation parameters
+T_cmb = 2.7255
+T_ncdm = 0.71611
+N_ur = 2.0308
+N_ncdm = 1
+deg_ncdm = [1]
+m_ncdm = [0.06]
+
+#Maximum wavenumber and redshift
+kmax = 30.
+zmax = 1e3
+amin = 1.0 / (zmax + 1)
+
+#CLASS output distance unit
+Mpc_cgs = 3.085677581282e24
+
+#CLASS parameters
+params = {
+	"h": h,
+	"Omega_b": Omega_b,
+	"Omega_cdm": Omega_cdm,
+	"T_cmb": T_cmb,
+	"N_ncdm": N_ncdm,
+	"N_ur": N_ur,
+	"T_ncdm": T_ncdm,
+	"deg_ncdm": "".join(str(x)+"," for x in deg_ncdm)[:-1],
+	"m_ncdm": "".join(str(x)+"," for x in m_ncdm)[:-1],
+	"A_s": A_s,
+	"n_s": n_s,
+	"output": "dTk, vTk",
+	"z_max_pk": zmax,
+	"P_k_max_1/Mpc": kmax,
+	"reio_parametrization": "reio_none",
+	"YHe": "BBN",
+	"k_output_values": kmax,
+	"k_per_decade_for_pk": 100}
+
+print("Running CLASS")
+
+#Run CLASS
+model = Class()
+model.set(params)
+model.compute()
+
+#Extract wavenumbers and prepare redshifts
+k = model.get_transfer(0)["k (h/Mpc)"] * h
+a = np.exp(np.arange(0, np.log(amin), -0.01))
+z = 1.0 / a - 1.0
+nk = len(k)
+nz = len(z)
+
+print("We have", nk, "wavenumbers and", nz, "redshifts")
+
+keys = model.get_transfer(0).keys()
+
+print("Available transfer functions:")
+print(keys)
+
+#Prepare dictionary
+pt = {}
+for key in keys:
+	pt[key] = np.zeros((nz, nk))
+
+#Extract transfer functions
+for i in range(nz):
+	pti = model.get_transfer(z[i])
+	for key in pti:
+		pt[key][i,:] = pti[key]
+
+#Export the perturbations file
+f.create_group("Functions")
+f["Redshifts"] = z
+f["Wavenumbers"] = k
+f.create_group("Units")
+f["Units"].attrs["Unit length in cgs (U_L)"] = Mpc_cgs
+
+#Write the perturbations
+for key in keys:
+	f["Functions/" + key.replace("/", "\\")] = pt[key]
+
+#Close the file
+f.close()
+
+print("Done.")

tools/spawn_neutrinos.py

+#!/usr/bin/env python3
+
+# Utility to spawn placeholder neutrinos in an existing initial conditions file.
+# The neutrinos will be distributed uniformly in space and with with zero mass
+# and velocity.
+
+import h5py
+import numpy as np
+import sys
+
+# This script does not need to be changed for different particle numbers.
+# Usage: ./spawn_neutrinos.py filename
+
+# Constants
+Mpc_cgs = 3.08567758e+24
+Default_N_nu_per100Mpc = 72 # 72^3 neutrinos for a (100 Mpc)^3 box
+Default_nr_neutrinos_per_Mpc3 = (Default_N_nu_per100Mpc / 100.)**3
+
+# Read command line arguments
+if len(sys.argv) <= 1 or sys.argv[1] == "--help" or sys.argv[1] == "-h":
+    print("Usage: ./spawn_neutrinos.py filename (use -h to show this message)")
+    exit(0)
+
+# Open the hdf5 file
+fname = sys.argv[1]
+f = h5py.File(fname, "r+")
+print("Operating on '" + fname + "'")
+print("")
+
+# Check the unit system
+if "Units" in f.keys() and "Unit length in cgs (U_L)" in f["Units"].attrs.keys():
+    Length_Unit = f["Units"].attrs["Unit length in cgs (U_L)"] / Mpc_cgs # Mpc
+else:
+    Length_Unit = 1.0 # Mpc
+
+# Extract the box dimensions and volume
+L = f["Header"].attrs["BoxSize"] / Length_Unit # Mpc
+V = L**3 if np.isscalar(L) else np.product(L) # Mpc^3
+if not np.isscalar(L) and len(L) != 3:
+    raise ValueError("Box dimensions are not cubic")
+
+# Check that the file does not have any neutrinos
+nparts = f["Header"].attrs["NumPart_Total"]
+while len(nparts) < 6:
+    nparts = np.append(nparts, 0)
+if nparts[6] != 0 or "PartType6" in f.keys():
+    raise IOError("This file already has neutrinos.")
+
+# Compute the default number of neutrinos (round to nearest cubic number)
+Default_N_nu = round((Default_nr_neutrinos_per_Mpc3 * V)**(1./3.))
+Default_Nr_neutrinos = int(Default_N_nu**3)
+
+print("The box dimensions are " + str(L) + " Mpc.")
+print("The default number of neutrinos is " +
+      "%g" % Default_N_nu_per100Mpc + "^3 per (100 Mpc)^3.")
+print("The default number of neutrinos is " +
+      "%g" % Default_N_nu + "^3 = " + str(Default_Nr_neutrinos) + ".")
+print("")
+
+#Request the number of neutrino particles to be spawned (with default option)
+Nr_neutrinos = int(input("Enter the number of neutrinos (default " +
+                         "%d" % Default_Nr_neutrinos + "): ")
+                      or "%d" % Default_Nr_neutrinos)
+
+nparts[6] = Nr_neutrinos
+
+print("")
+print("The number of particles per type will be:")
+print("{:25s}: {:12d}".format("Gas", nparts[0]))
+print("{:25s}: {:12d}".format("Dark Matter", nparts[1]))
+print("{:25s}: {:12d}".format("Background Dark Matter", nparts[2]))
+print("{:25s}: {:12d}".format("Sinks", nparts[3]))
+print("{:25s}: {:12d}".format("Stars", nparts[4]))
+print("{:25s}: {:12d}".format("Black Holes", nparts[5]))
+print("{:25s}: {:12d}".format("Neutrinos", nparts[6]))
+print("")
+
+firstID = int(nparts[0:6].sum() + 1)
+print("The first particle ID of the first neutrino will be: " + str(firstID))
+print("")
+
+confirm = input("Enter y to confirm: ")
+if (confirm != "y"):
+    print("Not confirmed. Done for now.")
+    exit(0)
+
+print("")
+print("Generating particle data...")
+
+# Generate particle data
+x = np.random.uniform(0, L, (Nr_neutrinos, 3)) * Length_Unit
+v = np.zeros((Nr_neutrinos, 3))
+m = np.zeros(Nr_neutrinos)
+ids = np.arange(firstID, firstID + Nr_neutrinos)
+
+print("Writing particle data...")
+
+# Store the particle data
+f.create_group("PartType6")
+f["PartType6/Coordinates"] = x
+f["PartType6/Velocities"] = v
+f["PartType6/Masses"] = m
+f["PartType6/ParticleIDs"] = ids
+
+# Update the header
+f["Header"].attrs["NumPart_Total"] = nparts
+
+print("All done.")
+
+# And close
+f.close()