plot_cooling.py

#!/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
import pyswiftsim

from copy import deepcopy
import numpy as np
import matplotlib.pyplot as plt
from astropy import units, constants
from time import strftime, gmtime
import sys

plt.rc('text', usetex=True)

pyswiftsim.downloadCoolingTable()

# PARAMETERS

colormap = "RdBu"

# cosmology
with_cosmo = False

# include metals?
with_metals = 1

# hydrogen mass fraction
h_mass_frac = 0.76

# density in atom / cm3
N_rho = int(sys.argv[-1])
if N_rho == 1:
    density = np.array([1.])
else:
    density = np.logspace(-6, 4, N_rho)

# temperature in K
N_T = 1000
temperature = np.logspace(1, 9, N_T)

# time step in s
dt = units.Myr * 1e-8
dt = dt.to("s") / units.s

# adiabatic index
gamma = 5. / 3.


def mu(T):
    # 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):
    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"])

    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))
    energy /= (gamma - 1.) * m_p
    d["energy"] = energy.to("")

    time_step = dt / unit_time
    d["time_step"] = time_step

    return d


def plot_solution(rate, data, us):
    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

    # 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(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,
                     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__":

    overwrite = {
        "GrackleCooling": {
            "WithMetalCooling": with_metals,
        }
    }
    with pyswiftsim.ParameterFile(overwrite) as filename:

        parser = pyswiftsim.parseYamlFile(filename)
        us = parser["InternalUnitSystem"]

        d = generate_initial_condition(us)

        # du / dt
        print("Computing cooling...")

        rate = coolingRate(filename,
                           d["density"].astype(np.float32),
                           d["energy"].astype(np.float32),
                           with_cosmo, d["time_step"])

        plot_solution(rate, d, us)