Select Git revision
plot_cooling.py
plot_cooling.py 8.71 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 = 3
# 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
cosmo = wrapper.cosmologyInit(params, us, pconst)
d["cosmo"] = cosmo
# 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"]
cosmo = d["cosmo"]
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,
cosmo,
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,
cosmo,
cooling,
chemistry,
rho,
ene,
dt,
frac.astype(np.float32))
plot_solution(energy, d, us)