Skip to content
Snippets Groups Projects
plot_cooling.py 8.42 KiB
#!/usr/bin/env python3
"""
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/>.
"""

from pyswiftsim import wrapper

import numpy as np
import matplotlib.pyplot as plt
from os.path import isfile
from h5py import File
from glob import glob

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

# PARAMETERS

# grackle primordial chemistry (for comparison)
primordial_chemistry = 1

# reference data
grackle_filename = "grackle.hdf5"
compute_equilibrium = True

# swift params filename
filename = "cooling.yml"

# particle data
# in atom/cc if generating
# if reading => same than in grackle.hdf5
part_density = 2.4570967948472597e7
part_temperature = 1e6  # in K
gamma = 5./3.

# default parameters
t_end = 1.
N_time = 100000

# cooling output
swift_output_filename = "swift.hdf5"

# simulations
simulations = {
    "SWIFT": "/home/loikki/data/swiftsim/examples/CoolingBox/coolingBox_*.hdf5"
}
# SCRIPT


def get_fields(primordial_chemistry):
    fields = [
        "metal"
    ]
    if primordial_chemistry > 0:
        fields.extend([
            "HI",
            "HII",
            "HeI",
            "HeII",
            "HeIII",
            "de"])

    if primordial_chemistry > 1:
        fields.extend([
            "HM",
            "H2I",
            "H2II"
            ])

    if primordial_chemistry > 2:
        fields.extend([
            "DI",
            "DII",
            "HDI"
        ])
    return fields


def generate_default_initial_condition(us, pconst):
    print("Generating default initial conditions")
    d = {}
    # generate grid
    d["temperature"] = part_temperature

    # Deal with units
    rho = part_density * us.UnitLength_in_cgs**3 * pconst.const_proton_mass
    d["density"] = [rho]

    energy = pconst.const_boltzmann_k * part_temperature \
        / us.UnitTemperature_in_cgs
    energy /= (gamma - 1.) * pconst.const_proton_mass
    d["energy"] = energy

    d["time"] = np.linspace(0, t_end, N_time)

    return d


def get_io_fields(cooling, chemistry, output):
    a = 1. / (1. + cooling.redshift)
    fields = {
        "MetalCooling": cooling.with_metal_cooling,
        "UVBackground": cooling.with_uv_background,
        "SelfShieldingMethod": cooling.self_shielding_method,
        "MetalMassFraction": chemistry.initial_metallicity,
        "ScaleFactor": a,
        "Temperature": part_temperature,
        "Density": part_density,
    }

    if output:
        fields.update({
            "MaxIter": cooling.max_step,
            "Tolerance": cooling.convergence_limit,
        })
    return fields


def getSimulationIndice(data, cooling, chemistry, output):

    fields = get_io_fields(cooling, chemistry, output)

    ind = 0
    for name in data:
        test = True
        d = data[name]
        if "Params" in d:
            tmp = d["Params"]
            for key in fields.keys():
                check = abs(tmp.attrs[key] - fields[key]) > \
                        1e-3 * tmp.attrs[key]
                if check:
                    test = False
                    break
        if test:
            break
        ind += 1
    if ind == len(data):
        return -1
    else:
        return ind


def read_grackle_data(filename, us, pconst, primordial_chemistry,
                      cooling, chemistry):
    print("Reading initial conditions")
    f = File(filename, "r")
    data = f["PrimordialChemistry%i" % primordial_chemistry]

    i = getSimulationIndice(data, cooling, chemistry, output=False)
    if i < 0:
        raise IOError("Unable to locate corresponding grackle simulation")
    else:
        tmp = list(data.keys())
        data = data[tmp[i]]

    # read units
    tmp = data["Units"].attrs

    u_len = tmp["Length"] / us.UnitLength_in_cgs
    u_time = tmp["Time"] / us.UnitTime_in_cgs
    u_den = tmp["Density"] * us.UnitLength_in_cgs**3 / us.UnitMass_in_cgs

    # read input
    tmp = data["Params"]
    energy = tmp.attrs["Energy"] * u_len**2 / u_time**2
    d["energy"] = energy

    density = tmp.attrs["Density"] * u_den
    d["density"] = density

    # read fractions
    tmp = data["Input"]
    for i in get_fields(primordial_chemistry):
        d[i] = tmp[i][:]

    # read output
    tmp = data["Output"]

    energy = tmp["Energy"][:] * u_len**2 / u_time**2
    d["out_energy"] = energy

    t = tmp["Time"][:] * u_time
    d["time"] = t

    f.close()
    return d


def initialize_swift(filename):
    print("Initialization of SWIFT")
    d = {}

    # parse swift params
    params = wrapper.parserReadFile(filename)
    d["params"] = params

    # init units
    us, pconst = wrapper.unitSystemInit(params)
    d["us"] = us
    d["pconst"] = pconst

    # init cooling
    cooling = wrapper.coolingInit(params, us, pconst)
    d["cooling"] = cooling

    # init chemistry
    chemistry = wrapper.chemistryInit(params, us, pconst)
    d["chemistry"] = chemistry

    return d


def plot_solution(energy, data, us):
    print("Plotting solution")
    rho = data["density"]
    time = data["time"]

    ref = False
    if "out_energy" in data:
        ref = True
        grackle_energy = data["out_energy"]

    # change units => cgs
    rho *= us.UnitMass_in_cgs / us.UnitLength_in_cgs**3

    energy *= us.UnitLength_in_cgs**2 / us.UnitTime_in_cgs**2

    if ref:
        grackle_energy *= us.UnitLength_in_cgs**2 / us.UnitTime_in_cgs**2

    time *= us.UnitTime_in_cgs

    # do plot

    plt.figure()
    plt.plot(time, energy, label="PySWIFTsim, %s" % wrapper.configGetCooling())
    if ref:
        plt.plot(time, grackle_energy,
                 label="Grackle, Prim. Chem. %i" % primordial_chemistry)

    for i in simulations:
        files = glob(simulations[i])
        E = np.zeros([len(files)])
        t = np.zeros([len(files)])
        j = 0
        for f in files:
            f = File(f)
            tmp = f["Units"]
            UL = tmp.attrs["Unit length in cgs (U_L)"]
            UT = tmp.attrs["Unit time in cgs (U_t)"]
            UE = UL**2 / UT**2
            tmp = f["PartType0"]
            E[j] = np.mean(tmp["InternalEnergy"]) * UE
            t[j] = f["Header"].attrs["Time"] * UT
            j += 1
            cooling_type = f["SubgridScheme"].attrs["Cooling Model"]
        plt.plot(t, E, label=i + ": " + cooling_type.decode("utf8"))
    plt.xlabel("Time [s]")
    plt.ylabel("Energy [erg]")
    plt.legend()
    plt.show()


if __name__ == "__main__":

    d = initialize_swift(filename)
    pconst = d["pconst"]
    us = d["us"]
    params = d["params"]
    cooling = d["cooling"]
    chemistry = d["chemistry"]

    if isfile(grackle_filename):
        d = read_grackle_data(grackle_filename, us, pconst,
                              primordial_chemistry, cooling,
                              chemistry)
    else:
        d = generate_default_initial_condition(us, pconst)

    t = d["time"]
    dt = t[1] - t[0]

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

    N = len(t)
    energy = np.zeros(t.shape)
    energy[0] = d["energy"]
    rho = np.array(d["density"], dtype=np.float32)
    for i in range(N-1):
        ene = np.array([energy[i]], dtype=np.float32)
        if not compute_equilibrium:
            energy[i+1] = wrapper.doCooling(pconst, us, cooling,
                                            chemistry,
                                            rho,
                                            ene,
                                            dt)
        else:
            fields = get_fields(primordial_chemistry)
            N = len(fields)
            frac = np.zeros([1, N])
            for j in range(N):
                frac[:, j] = d[fields[j]]
            energy[i+1] = wrapper.doCooling(pconst, us, cooling,
                                            chemistry,
                                            rho,
                                            ene,
                                            dt,
                                            frac.astype(np.float32))

    plot_solution(energy, d, us)