183 lines
7.5 KiB
Python
183 lines
7.5 KiB
Python
|
|
|
|
def create_parameters_dict(ID,ts_file,cosmo_params):
|
|
"""
|
|
Create a dictionnary of paramters for a given ID for the sbmy_control program.
|
|
"""
|
|
|
|
params=dict(directory="/home/aubin/data/PotentialBCs/ML_dataset/",
|
|
simname=ID,
|
|
mode="alltCOLA",
|
|
ICs_gen="monofonic",
|
|
verbose=1,
|
|
nthreads=128,
|
|
execution="slurm",
|
|
N_particles=768,
|
|
L=480.0,
|
|
ICs="/home/aubin/data/PotentialBCs/ML_dataset/initial_conditions/ICs_"+ID+"_DM_delta.h5",
|
|
WriteGravPot=True,
|
|
OutputGravitationalPotentialBase="/home/aubin/data/PotentialBCs/ML_dataset/gravitational_potential/gravpot_"+ID,
|
|
MeshGravPot=768,
|
|
WriteDensity=True,
|
|
OutputDensityBase="/home/aubin/data/PotentialBCs/ML_dataset/density/density_"+ID,
|
|
MeshDensity=768,
|
|
WriteReferenceFrame=True,
|
|
OutputMomentaBase="/home/aubin/data/PotentialBCs/ML_dataset/momenta/momenta_"+ID+"_",
|
|
monofonic_config="/home/aubin/data/PotentialBCs/ML_dataset/monofonic_config/monofonic_config_"+ID+".conf",
|
|
monofonic_output="/home/aubin/data/PotentialBCs/ML_dataset/initial_conditions/ICs_"+ID+"_",
|
|
TimeSteppingFileName=ts_file,
|
|
Omega_m=cosmo_params["Omega_m"],
|
|
Omega_b=cosmo_params["Omega_b"],
|
|
Omega_q=cosmo_params["Omega_q"],
|
|
h=cosmo_params["h"],
|
|
n_s=cosmo_params["n_s"],
|
|
sigma8=cosmo_params["sigma8"],
|
|
Tcmb=cosmo_params["Tcmb"],
|
|
k_p=cosmo_params["k_p"],
|
|
N_ur=cosmo_params["N_ur"],
|
|
m_nu1=cosmo_params["m_nu1"],
|
|
m_nu2=cosmo_params["m_nu2"],
|
|
m_nu3=cosmo_params["m_nu3"],
|
|
w_0=cosmo_params["w_0"],
|
|
w_a=cosmo_params["w_a"],
|
|
fnl=cosmo_params["fnl"],
|
|
gnl=cosmo_params["gnl"],
|
|
)
|
|
|
|
return params
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import numpy as np
|
|
import pandas as pd
|
|
from pysbmy.timestepping import StandardTimeStepping
|
|
from pysbmy.cosmology import full_growthfactors_solver
|
|
from sbmy_control.parameters_monofonic import get_config_from_dict, create_monofonic_config
|
|
|
|
# Read the cosmological parameter from the CSV file
|
|
w0wa_set_file = "/home/aubin/PotentialBCs/Planck2018/w0wa_set.csv"
|
|
w0wa_set = pd.read_csv(w0wa_set_file)
|
|
|
|
# Extract the IDs from the DataFrame
|
|
IDs = w0wa_set['Run'].tolist()
|
|
|
|
# Define the timestepping configurations (TimeStepDistribution, nsteps, name)
|
|
ts_configs=[(0,10,'Ta1'),
|
|
(0,20,'Ta2'),
|
|
(3,10,'TD1'),
|
|
(3,20,'TD2'),]
|
|
|
|
# Define the additional cosmological parameters
|
|
cosmo_defaults_add = {
|
|
"Omega_k":0.0,
|
|
"Omega_r": 0.0,
|
|
"Tcmb": 2.7255,
|
|
"k_p": 0.05,
|
|
"N_ur": 2.046,
|
|
"m_nu1": 0.06,
|
|
"m_nu2": 0.0,
|
|
"m_nu3": 0.0,
|
|
"fnl": 0.0,
|
|
"gnl": 0.0,
|
|
"k_max":10.0,
|
|
"tau_reio":0.06,
|
|
"WhichSpectrum":"EH",
|
|
}
|
|
|
|
# Iterate over the IDs
|
|
for k, ID in enumerate(IDs):
|
|
if k<6: # Skip the first 6 IDs, already done
|
|
continue
|
|
print(f"Processing ID {k+1}/{len(IDs)}: {ID}")
|
|
|
|
# Get the cosmological parameters for the current ID
|
|
cosmo_params = w0wa_set[w0wa_set['Run'] == ID].iloc[0].to_dict()
|
|
# Add the additional cosmological parameters
|
|
cosmo_params.update(cosmo_defaults_add)
|
|
|
|
cosmo_params["w_0"] = cosmo_params["w0"]
|
|
cosmo_params["w_a"] = cosmo_params["wa"]
|
|
cosmo_params["w0_fld"] = cosmo_params["w0"]
|
|
cosmo_params["wa_fld"] = cosmo_params["wa"]
|
|
|
|
# Create the timestepping
|
|
print(f"> Creating timestepping with configuration {ts_configs[k % len(ts_configs)]} for ID {ID}")
|
|
this_ts_config = ts_configs[k % len(ts_configs)]
|
|
TS = StandardTimeStepping(
|
|
ai= 1/(1+19.0),
|
|
af = 1.,
|
|
integrator=2,
|
|
TimeStepDistribution=this_ts_config[0],
|
|
nsteps=this_ts_config[1],
|
|
cosmo=cosmo_params
|
|
)
|
|
|
|
# Write the timestepping
|
|
print(f"> Writing timestepping to file for ID {ID}")
|
|
ts_file = f"/home/aubin/data/PotentialBCs/ML_dataset/params/ts_{ID}_{this_ts_config[2]}.h5"
|
|
TS.write(ts_file)
|
|
|
|
# For each force evaluation, get the following values: a, D1, D2, which can be obtained with
|
|
print(f"> Calculating growth factors for ID {ID}")
|
|
a_force = TS.aStart[TS.operations == 2]
|
|
D_array = full_growthfactors_solver(a_force, cosmo_params)
|
|
D1_force = D_array[:, 0]
|
|
D2_force = D_array[:, 2]
|
|
|
|
# For each force evaluation, save the following values to a file: ID, nforce, a, D1, D2, h, Omega_m, sigma8
|
|
print(f"> Saving the cosmological parameters and time values for ID {ID}")
|
|
for nforce in range(len(a_force)):
|
|
output_file = f"/home/aubin/data/PotentialBCs/ML_dataset/cosmo_and_time/cosmo_and_time_parameters_{ID}_nforce{nforce}.txt"
|
|
with open(output_file, 'w') as f:
|
|
f.write(f"ID: {ID}\n")
|
|
f.write(f"nforce: {nforce}\n")
|
|
f.write(f"a: {a_force[nforce]}\n")
|
|
f.write(f"D1: {D1_force[nforce]}\n")
|
|
f.write(f"D2: {D2_force[nforce]}\n")
|
|
f.write(f"h: {cosmo_params['h']}\n")
|
|
f.write(f"Omega_m: {cosmo_params['Omega_m']}\n")
|
|
f.write(f"sigma8: {cosmo_params['sigma8']}\n")
|
|
|
|
# Create the configuration file for monofonic
|
|
print(f"> Creating the configuration file for monofonic for ID {ID}")
|
|
monofonic_config_file = f"/home/aubin/data/PotentialBCs/ML_dataset/monofonic_config/monofonic_config_{ID}.conf"
|
|
monofonic_dict = {
|
|
"config" : monofonic_config_file,
|
|
"output" : f"/home/aubin/data/PotentialBCs/ML_dataset/initial_conditions/ICs_{ID}_",
|
|
"gridres": 768,
|
|
"boxlength": 480.0,
|
|
"seed" : np.random.randint(0, 2**32),
|
|
"ParameterSet":None,
|
|
"nthreads":128,
|
|
**cosmo_params,
|
|
}
|
|
create_monofonic_config(monofonic_config_file, get_config_from_dict(monofonic_dict))
|
|
|
|
# Run the sbmy_control command with the parameters for the current ID
|
|
print(f"> Running sbmy_control for ID {ID}")
|
|
params = create_parameters_dict(ID,ts_file,cosmo_params)
|
|
|
|
import subprocess
|
|
|
|
command_args = ["sbmy_control"]
|
|
for key, value in params.items():
|
|
command_args.append(f"--{key}")
|
|
command_args.append(f"{value}")
|
|
|
|
subprocess.run(command_args)
|
|
|
|
# Remove the initial conditions particles files that are unnecessary
|
|
print(f"> Removing unnecessary initial conditions files for ID {ID}")
|
|
import os
|
|
import glob
|
|
ic_particles_files = glob.glob(f"/home/aubin/data/PotentialBCs/ML_dataset/initial_conditions/ICs_{ID}_particles*")
|
|
for file in ic_particles_files:
|
|
os.remove(file)
|
|
|
|
# Remove monofonic slurm script
|
|
monofonic_slurm_script = f"/home/aubin/data/PotentialBCs/ML_dataset/slurm_scripts/monofonic.sh"
|
|
if os.path.isfile(monofonic_slurm_script):
|
|
print(f"> Removing monofonic slurm script for ID {ID}")
|
|
os.remove(monofonic_slurm_script)
|
|
|
|
print(f"> Finished processing ID {ID}\n\n\n")
|