Update cooling example and repository structure

test/cooling.yml → examples/cooling_rate/cooling.yml

@@ -37,10 +37,14 @@ InitialConditions:
 
 # 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
+  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
+  ProvideVolumetricHeatingRates: 0 # User provide volumetric heating rates
+  ProvideSpecificHeatingRates: 0 # User provide specific heating rates
+  SelfShieldingMethod: 0 # Grackle (<= 3) or Gear self shielding method
+  OutputMode: 0 # Write in output corresponding primordial chemistry mode
 
 Gravity:
   eta:          0.025    # Constant dimensionless multiplier for time integration.

examples/cooling_rate/generate_grackle_data.py

+#!/usr/bin/env python
+########################################################################
+#
+# Cooling rate data generation. This code is a modified version
+# of a code developed by the Grackle Team.
+#
+#
+# Copyright (c) 2013-2016, Grackle Development Team.
+#
+# Distributed under the terms of the Enzo Public Licence.
+#
+# The full license is in the file LICENSE, distributed with this
+# software.
+########################################################################
+
+"""
+Generate (or update) a hdf5 file containing
+the grackle data for each primordial chemistry
+"""
+
+import numpy as np
+from h5py import File
+from copy import deepcopy
+
+from pygrackle import \
+    chemistry_data, \
+    setup_fluid_container
+
+from pygrackle.utilities.physical_constants import \
+    mass_hydrogen_cgs, \
+    sec_per_Myr
+
+filename = "grackle.hdf5"
+
+
+def generate_data(primordial_chemistry):
+    current_redshift = 0.
+
+    # Set solver parameters
+    my_chemistry = chemistry_data()
+    my_chemistry.use_grackle = 1
+    my_chemistry.with_radiative_cooling = 1
+    my_chemistry.primordial_chemistry = primordial_chemistry
+    my_chemistry.metal_cooling = 1
+    my_chemistry.UVbackground = 1
+    my_chemistry.self_shielding_method = 0
+    my_chemistry.H2_self_shielding = 0
+    my_chemistry.grackle_data_file = "CloudyData_UVB=HM2012.h5"
+
+    # Set units
+    my_chemistry.comoving_coordinates = 0  # proper units
+    my_chemistry.a_units = 1.0
+    my_chemistry.a_value = 1.0 / (1.0 + current_redshift) / \
+        my_chemistry.a_units
+    my_chemistry.length_units = 3.085e21
+    my_chemistry.density_units = 1.989e43 / my_chemistry.length_units**3
+    my_chemistry.velocity_units = 20725573.785998672
+    my_chemistry.time_units = my_chemistry.length_units / \
+        my_chemistry.velocity_units
+
+    density = mass_hydrogen_cgs
+    dt = 1e-8 * sec_per_Myr / my_chemistry.time_units
+    # Call convenience function for setting up a fluid container.
+    # This container holds the solver parameters, units, and fields.
+    temperature = np.logspace(1, 9, 1000)
+    fc = setup_fluid_container(my_chemistry,
+                               metal_mass_fraction=0.01295,
+                               temperature=temperature,
+                               density=density,
+                               converge=True)
+
+    f = File(filename, "a")
+    data_name = "PrimordialChemistry%i" % primordial_chemistry
+    if data_name in f:
+        print "Updating Dataset"
+        data = f[data_name]
+        f.pop(data_name)
+    else:
+        print "Creating Dataset"
+
+    data = f.create_group(data_name)
+
+    # Units
+    data.create_group("Units")
+    tmp = data["Units"]
+    tmp.attrs["Length"] = my_chemistry.length_units
+    tmp.attrs["Density"] = my_chemistry.density_units
+    tmp.attrs["Velocity"] = my_chemistry.velocity_units
+    tmp.attrs["Time"] = my_chemistry.time_units
+
+    # Parameters
+    data.create_group("Params")
+    tmp = data["Params"]
+    tmp.attrs["MetalCooling"] = my_chemistry.metal_cooling
+    tmp.attrs["UVBackground"] = my_chemistry.UVbackground
+    tmp.attrs["SelfShieldingMethod"] = my_chemistry.self_shielding_method
+    tmp.attrs["TimeStep"] = dt
+
+    # Inputs
+    data.create_group("Input")
+    tmp = data["Input"]
+    # energy
+    dset = tmp.create_dataset("Energy", fc["energy"].shape,
+                              dtype=fc["energy"].dtype)
+    dset[:] = fc["energy"]
+    # density
+    dset = tmp.create_dataset("Density", fc["density"].shape,
+                              dtype=fc["density"].dtype)
+    dset[:] = fc["density"]
+    # Temperautre
+    dset = tmp.create_dataset("Temperature", temperature.shape,
+                              dtype=temperature.dtype)
+    dset[:] = temperature
+
+    old = deepcopy(fc)
+    fc.solve_chemistry(dt)
+    rate = (fc["energy"] - old["energy"]) / dt
+
+    data.create_group("Output")
+    tmp = data["Output"]
+    # Rate
+    dset = tmp.create_dataset("Rate", rate.shape,
+                              dtype=rate.dtype)
+    dset[:] = rate
+    # Energy
+    dset = tmp.create_dataset("Energy", fc["energy"].shape,
+                              dtype=fc["energy"].dtype)
+    dset[:] = fc["energy"]
+
+    f.close()
+
+
+if __name__ == "__main__":
+    for i in range(4):
+        print "Computing Primordial Chemistry %i" % i
+        generate_data(i)

examples/cooling_rate/run.sh

+#!/bin/bash
+
+# Get the Grackle cooling table
+if [ ! -e CloudyData_UVB=HM2012.h5 ]
+then
+    echo "Fetching the Cloudy tables required by Grackle..."
+    ./getCoolingTable.sh
+fi
+
+# Generate Grackle data if not present
+if [ ! -e grackle.hdf5 ]
+then
+    echo "Generating Grackle Data..."
+    ./generate_grackle_data.py
+fi
+
+./test_cooling.py

examples/cooling_rate/test_cooling.py

+#!/usr/bin/env python3
+
+from pyswiftsim import wrapper
+
+from copy import deepcopy
+import numpy as np
+import matplotlib.pyplot as plt
+from astropy import units
+from os.path import isfile
+from h5py import File
+
+plt.rc('text', usetex=True)
+
+# PARAMETERS
+
+# grackle primordial chemistry
+primordial_chemistry = 0
+
+# reference data
+grackle_filename = "grackle.hdf5"
+
+# swift params filename
+filename = "cooling.yml"
+
+# if grackle_filename does not exist
+# use following values
+
+# density in atom / cm3
+N_rho = 1
+# with N_rho > 1, the code is not implemented to deal
+# with the reference
+if N_rho == 1:
+    default_density = np.array([1.])
+else:
+    default_density = np.logspace(-3, 1, N_rho)
+
+# temperature in K
+N_T = 1000
+default_temperature = np.logspace(1, 9, N_T)
+
+# time step in s
+default_dt = units.Myr * 1e-8
+default_dt = default_dt.to("s") / units.s
+
+# adiabatic index
+gamma = 5. / 3.
+
+# SCRIPT
+
+
+def generate_default_initial_condition(us, pconst):
+    print("Generating default initial conditions")
+    d = {}
+    # generate grid
+    rho, T = np.meshgrid(default_density, default_temperature)
+    rho = deepcopy(rho.reshape(-1))
+    T = T.reshape(-1)
+    d["temperature"] = T
+
+    # Deal with units
+    rho *= us.UnitLength_in_cgs**3 * pconst.const_proton_mass
+    d["density"] = rho
+
+    energy = pconst.const_boltzmann_k * T / us.UnitTemperature_in_cgs
+    energy /= (gamma - 1.) * pconst.const_proton_mass
+    d["energy"] = energy
+
+    dt = default_dt / us.UnitTime_in_cgs
+    d["time_step"] = dt
+
+    return d
+
+
+def read_grackle_data(filename, us, primordial_chemistry):
+    print("Reading initial conditions")
+    f = File(filename, "r")
+    data = f["PrimordialChemistry%i" % primordial_chemistry]
+
+    # read units
+    tmp = data["Units"].attrs
+
+    u_len = tmp["Length"] / us.UnitLength_in_cgs
+    u_den = tmp["Density"] * us.UnitLength_in_cgs**3 / us.UnitMass_in_cgs
+    u_time = tmp["Time"] / us.UnitTime_in_cgs
+
+    # read input
+    tmp = data["Input"]
+
+    energy = tmp["Energy"][:] * u_len**2 / u_time**2
+    d["energy"] = energy
+
+    T = tmp["Temperature"][:] / us.UnitTemperature_in_cgs
+    d["temperature"] = T
+
+    density = tmp["Density"][:] * u_den
+    d["density"] = density
+
+    dt = data["Params"].attrs["TimeStep"] * u_time
+    d["time_step"] = dt
+
+    # read output
+    tmp = data["Output"]
+
+    energy = tmp["Energy"][:] * u_len**2 / u_time**2
+    d["out_energy"] = energy
+
+    rate = tmp["Rate"][:] * u_len**2 / (u_time**3)
+    d["rate"] = rate
+
+    f.close()
+    return d
+
+
+def initialize_swift(filename):
+    d = {}
+    # parse swift params
+    print("Reading parameters")
+    params = wrapper.parserReadFile(filename)
+    d["params"] = params
+    # init units
+    print("Initialization of the unit system")
+    us, pconst = wrapper.unitSystemInit(params)
+    d["us"] = us
+    d["pconst"] = pconst
+    # init cooling
+    print("Initialization of the cooling")
+    cooling = wrapper.coolingInit(params, us, pconst)
+    d["cooling"] = cooling
+    return d
+
+
+def plot_solution(rate, data, us):
+    energy = data["energy"]
+    rho = data["density"]
+    T = data["temperature"]
+
+    ref = False
+    if "rate" in data:
+        ref = True
+        grackle_rate = data["rate"]
+
+    # change units => cgs
+    rho *= us.UnitMass_in_cgs / us.UnitLength_in_cgs**3
+
+    T *= us.UnitTemperature_in_cgs
+
+    energy *= us.UnitLength_in_cgs**2 / us.UnitTime_in_cgs**2
+
+    rate *= us.UnitLength_in_cgs**2 / us.UnitTime_in_cgs**3
+
+    if ref:
+        grackle_rate *= us.UnitLength_in_cgs**2 / us.UnitTime_in_cgs**3
+
+    # lambda cooling
+    lam = rate * rho
+
+    if ref:
+        grackle_lam = grackle_rate * rho
+
+    # do plot
+    if N_rho == 1:
+        # plot Lambda vs T
+        plt.figure()
+        plt.loglog(T, np.abs(lam), label="SWIFT")
+        if ref:
+            # plot reference
+            plt.loglog(T, np.abs(grackle_lam), label="Ref.")
+            plt.legend()
+            plt.xlabel("Temperature [K]")
+            plt.ylabel("$\\Lambda$ [erg s$^{-1}$ cm$^{3}$]")
+
+            # plot error vs T
+            plt.figure()
+            plt.plot(T, (lam - grackle_lam) / grackle_lam)
+            plt.gca().set_xscale("log")
+            plt.xlabel("Temperature [K]")
+            plt.ylabel(
+                r"$(\Lambda - \Lambda_\textrm{ref}) / \Lambda_\textrm{ref}$")
+
+        plt.show()
+
+    else:
+        shape = [N_rho, N_T]
+        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()
+
+
+if __name__ == "__main__":
+
+    d = initialize_swift(filename)
+    pconst = d["pconst"]
+    us = d["us"]
+    params = d["params"]
+    cooling = d["cooling"]
+
+    if isfile(grackle_filename):
+        d = read_grackle_data(filename, us, primordial_chemistry)
+    else:
+        d = generate_default_initial_condition(us, pconst)
+
+    # du / dt
+    print("Computing cooling...")
+    rate = wrapper.coolingRate(pconst, us, cooling,
+                               d["density"].astype(np.float32),
+                               d["energy"].astype(np.float32),
+                               d["time_step"])
+
+    plot_solution(rate, d, us)

test/test_cooling.py

-#!/usr/bin/env python3
-
-from pyswiftsim import wrapper
-
-from copy import deepcopy
-import numpy as np
-import matplotlib.pyplot as plt
-from astropy import units
-from os.path import isfile
-
-#
-# parameters
-#
-
-# adiabatic index
-gamma = 5. / 3.
-
-# swift params filename
-filename = "test/cooling.yml"
-
-# number of points
-N_rho = 1
-N_T = 1000
-
-# 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)
-
-# time step
-dt = units.Myr * 1e-8
-dt = dt.to("s") / units.s
-
-
-# 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
-print("Reading parameters")
-params = wrapper.parserReadFile(filename)
-# init units
-print("Initialization of the unit system")
-us, pconst = wrapper.unitSystemInit(params)
-# init cooling
-print("Initialization of the cooling")
-cooling = wrapper.coolingInit(params, us, pconst)
-
-
-#
-# Deal with units
-#
-
-# change units of rho and T
-# rho
-rho *= us.UnitLength_in_cgs**3 * pconst.const_proton_mass
-
-# specific energy
-# check if reference solution is given
-ref = False
-if N_rho == 1 and isfile("energy.npy") and isfile("rate.npy"):
-    ref = True
-    print("Using reference solution")
-    energy = np.load("energy.npy")
-else:
-    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
-print("Computing cooling...")
-rate = wrapper.coolingRate(pconst, us, cooling,
-                           rho.astype(np.float32),
-                           energy.astype(np.float32),
-                           dt)
-
-if ref:
-    rate_ref = np.load("rate.npy")
-
-print("Computing done")
-
-#
-# 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 ref:
-    rate_ref *= us.UnitLength_in_cgs**2 / us.UnitTime_in_cgs**3
-
-if N_rho == 1 or N_T == 1:
-    lambda_ = rate * rho * us.UnitMass_in_cgs / us.UnitLength_in_cgs**3
-    if ref:
-        lambda_ref = rate_ref * rho * us.UnitMass_in_cgs \
-            / us.UnitLength_in_cgs**3
-
-    if ref:
-        plt.figure()
-        plt.plot(T, (lambda_ - lambda_ref) / lambda_ref)
-        plt.gca().set_xscale("log")
-
-    if N_rho == 1:
-        plt.figure()
-        plt.loglog(T, np.abs(lambda_), label="SWIFT")
-        if ref:
-            plt.loglog(T, np.abs(lambda_ref), label="Ref.")
-            plt.legend()
-        plt.xlabel("Temperature [K]")
-        plt.ylabel("$\\Lambda$ [erg s$^{-1}$ cm$^{3}$]")
-
-    else:
-        plt.figure()
-        plt.loglog(rho, np.abs(lambda_))
-        plt.xlabel("Density [atom / cm3]")
-        plt.ylabel("$\\Lambda$ [erg s$^{-1}$ cm$^{3}$]")
-
-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()