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"], A_s=cosmo_params["A_s"], 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 planck_set_file = "/home/aubin/PotentialBCs/Planck2018/Planck_set.csv" planck_set = pd.read_csv(planck_set_file) # Extract the IDs from the DataFrame IDs = planck_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, "w_0": -1.0, "w_a": 0.0, "fnl": 0.0, "gnl": 0.0, "k_max":10.0, "tau_reio":0.06, "WhichSpectrum":"EH", "w0_fld":-1.0, "wa_fld":0.0, } # Iterate over the IDs for k, ID in enumerate(IDs): print(f"Processing ID {k+1}/{len(IDs)}: {ID}") if k < 4: continue # Already done for the first 4 IDs # Get the cosmological parameters for the current ID cosmo_params = planck_set[planck_set['Run'] == ID].iloc[0].to_dict() # Add the additional cosmological parameters cosmo_params.update(cosmo_defaults_add) # Multiply A_s by 1e-9 to convert it to the correct units cosmo_params["A_s"] *= 1e-9 # 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")