From ea04587d56e4647f23e975f572e6ccb233732149 Mon Sep 17 00:00:00 2001 From: maubin Date: Wed, 5 Mar 2025 16:13:35 +0000 Subject: [PATCH] Upload files to "/" --- ICs.py | 184 +++++++++++++++++++++++++++++++++++++++++++++++ __init__.py | Bin 0 -> 1024 bytes args_main.py | 66 +++++++++++++++++ cosmo_params.py | 66 +++++++++++++++++ low_level.py | 185 ++++++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 501 insertions(+) create mode 100644 ICs.py create mode 100644 __init__.py create mode 100644 args_main.py create mode 100644 cosmo_params.py create mode 100644 low_level.py diff --git a/ICs.py b/ICs.py new file mode 100644 index 0000000..a50b2fa --- /dev/null +++ b/ICs.py @@ -0,0 +1,184 @@ + +def main_ICs(parsed_args): + """ + Main function for the initial conditions generator. + """ + from args_main import parse_arguments_main + from parameters_card import parse_arguments_card + from low_level import print_starting_module, print_message, print_ending_module + from os.path import isfile + + main_dict = parse_arguments_main(parsed_args) + card_dict, _ = parse_arguments_card(parsed_args) + + print_starting_module("ICs", verbose=parsed_args.verbose) + + match main_dict["ICs_gen"]: + + case "ext": + if not isfile(card_dict["ICs"]): + raise FileNotFoundError(f"External initial conditions file {card_dict["ICs"]} not found.") + if not card_dict["ICsMode"] == 2: + raise ValueError(f"ICs_gen is in ext mode, but ICsMode is not 2: ICsMode={card_dict["ICsMode"]}") + print_message(f"External initial conditions file {card_dict["ICs"]} found.", 1, "ICs", verbose=parsed_args.verbose) + + case "sbmy": + if card_dict["ICsMode"] == 0: + print_message("ICsMode is 0 so the initial conditions will be generated during the sbmy run.", 1, "ICs", verbose=parsed_args.verbose) + elif card_dict["ICsMode"] == 1: + print_message("ICsMode is 1 so the initial conditions will be generated using the sbmy ICs generator.", 1, "ICs", verbose=parsed_args.verbose) + ICs_sbmy(parsed_args) + elif card_dict["ICsMode"] == 2: + raise ValueError("ICsMode is 2 but ICs_gen is sbmy (not ext).") + else: + raise ValueError(f"ICsMode {card_dict["ICsMode"]} not recognized.") + + case "monofonic": + if card_dict["ICsMode"] == 1 or card_dict["ICsMode"] == 0: + raise ValueError("ICsMode is 1 or 2 but ICs_gen is monofonic.") + elif card_dict["ICsMode"] == 2: + print_message("ICsMode is 2 so the initial conditions will be generated using the monofonic ICs generator.", 1, "ICs", verbose=parsed_args.verbose) + ICs_monofonic(parsed_args) + else: + raise ValueError(f"ICsMode {card_dict["ICsMode"]} not recognized.") + + case _: + raise ValueError(f"ICs_gen {main_dict["ICs_gen"]} not recognized.") + + print_ending_module("ICs", verbose=parsed_args.verbose) + + + + + +def ICs_monofonic(parsed_args): + from monofonic import main_monofonic + main_monofonic(parsed_args) + + +def ICs_sbmy(parsed_args): + from pysbmy.power import PowerSpectrum + from low_level import print_starting_module, print_message + from os.path import isfile + from parameters_card import parse_arguments_card + + print_starting_module("sbmy IC", verbose=parsed_args.verbose) + + card_dict, _ = parse_arguments_card(parsed_args) + + power_spectrum_file = card_dict["InputPowerSpectrum"] + + if not isfile(power_spectrum_file) or parsed_args.force: + print_message(f"Power spectrum file {power_spectrum_file} does not exist or forced. Creating it", 1, "sbmy IC", verbose=parsed_args.verbose) + create_sbmy_power_spectrum_file(parsed_args, card_dict, power_spectrum_file) + print_message(f"Power spectrum file written to {power_spectrum_file}", 2, "sbmy IC", verbose=parsed_args.verbose) + else: + print_message(f"Power spectrum file {power_spectrum_file} exists.", 1, "sbmy IC", verbose=parsed_args.verbose) + + white_noise_field_file = card_dict["InputWhiteNoise"] + + if not isfile(white_noise_field_file) or parsed_args.force: + print_message(f"White noise field file {white_noise_field_file} does not exist or forced. Creating it", 1, "sbmy IC", verbose=parsed_args.verbose) + create_sbmy_white_noise_field(parsed_args, card_dict, white_noise_field_file) + print_message(f"White noise field file written to {white_noise_field_file}", 2, "sbmy IC", verbose=parsed_args.verbose) + else: + print_message(f"White noise field file {white_noise_field_file} exists.", 1, "sbmy IC", verbose=parsed_args.verbose) + + + + + + +def create_sbmy_power_spectrum_file(parsed_args, card_dict, power_spectrum_file): + from cosmo_params import parse_arguments_cosmo + from pysbmy.power import PowerSpectrum + + cosmo_dict = parse_arguments_cosmo(parsed_args) + + if parsed_args.verbose < 2: + from io import BytesIO + from low_level import stdout_redirector, stderr_redirector + f = BytesIO() + g = BytesIO() + with stdout_redirector(f): + with stderr_redirector(g): + Pk = PowerSpectrum(card_dict["L"], + card_dict["L"], + card_dict["L"], + card_dict["N_LPT_mesh"], + card_dict["N_LPT_mesh"], + card_dict["N_LPT_mesh"], + cosmo_dict, + ) + Pk.write(power_spectrum_file) + g.close() + f.close() + + else: + Pk = PowerSpectrum(card_dict["L"], + card_dict["L"], + card_dict["L"], + card_dict["N_LPT_mesh"], + card_dict["N_LPT_mesh"], + card_dict["N_LPT_mesh"], + cosmo_dict, + ) + Pk.write(power_spectrum_file) + + +def create_sbmy_white_noise_field(parsed_args, card_dict, white_noise_field_file): + import numpy as np + from gc import collect + from low_level import print_message + from pysbmy.field import BaseField + + print_message(f"Seed: {parsed_args.seed}", 3, "sbmy IC", verbose=parsed_args.verbose) + rng = np.random.default_rng(parsed_args.seed) + print_message(f"Numpy seed {str(rng.bit_generator.state)}", 3, "sbmy IC", verbose=parsed_args.verbose) + data = rng.standard_normal(size=card_dict["N_LPT_mesh"]**3) + wn = BaseField(card_dict["L"], + card_dict["L"], + card_dict["L"], + card_dict["corner"][0], + card_dict["corner"][1], + card_dict["corner"][2], + 1, + card_dict["N_LPT_mesh"], + card_dict["N_LPT_mesh"], + card_dict["N_LPT_mesh"], + data) + del data + collect() + + if parsed_args.verbose < 2: + from io import BytesIO + from low_level import stdout_redirector, stderr_redirector + f = BytesIO() + g = BytesIO() + with stdout_redirector(f): + with stderr_redirector(g): + wn.write(white_noise_field_file) + g.close() + f.close() + else: + wn.write(white_noise_field_file) + + +if __name__ == "__main__": + from argparse import ArgumentParser + from args_main import register_arguments_main + from parameters_card import register_arguments_card + from cosmo_params import register_arguments_cosmo + from parameters_monofonic import register_arguments_monofonic + from slurm_submission import register_arguments_slurm + + parser = ArgumentParser(description="Generate initial conditions for a Simbelmyne simulation.") + # TODO: reduce the volume of arguments + register_arguments_main(parser) + register_arguments_monofonic(parser) + register_arguments_slurm(parser) + register_arguments_card(parser) + register_arguments_cosmo(parser) + parsed_args = parser.parse_args() + + main_ICs(parsed_args) \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..06d7405020018ddf3cacee90fd4af10487da3d20 GIT binary patch literal 1024 ScmZQz7zLvtFd70QH3R?z00031 literal 0 HcmV?d00001 diff --git a/args_main.py b/args_main.py new file mode 100644 index 0000000..4c60e31 --- /dev/null +++ b/args_main.py @@ -0,0 +1,66 @@ +from argparse import ArgumentParser + +def register_arguments_main(parser:ArgumentParser): + """ + Register the arguments for the main parameters. + """ + parser.add_argument("-d", "--directory", type=str, default="./", help="Main directory where the output will be saved (if other dir and filenames are not specified).") + parser.add_argument("-pd", "--paramdir", type=str, default=None, help="Directory where the parameters are saved (default is -d/params).") + parser.add_argument("-rd", "--resultdir", type=str, default=None, help="Directory where the results are saved (default is -d/results).") + parser.add_argument("-wd", "--workdir", type=str, default=None, help="Directory where the work is done (default is -d/work).") + parser.add_argument("-ld", "--logdir", type=str, default=None, help="Directory where the logs are saved (default is -d/logs).") + parser.add_argument("-n", "--simname", type=str, default=None, help="Name of the simulation (for outputs and parameter card).") + parser.add_argument("-M", "--mode", type=str, default="tCOLA", help="Mode of the program: ICs, PM, LPT, tCOLA, sCOLA, pre_sCOLA, post_sCOLA, alltCOLA, allsCOLA.") + parser.add_argument("-ic","--ICs_gen", type=str, default="ext", help="Initial conditions generator: ext, sbmy, monofonic") + parser.add_argument("-v", "--verbose", type=int, default=1, help="Verbosity level (0: no output, 1: basic output, 2: subprograms output).") + parser.add_argument("--seed", type=int, default=None, help="Seed for the random number generator.") + parser.add_argument("--nthreads", type=int, default=None, help="Number of threads to use.") + parser.add_argument("-E","--execution", type=str, default="local", help="Execution mode: local, slurm.") + parser.add_argument("-F","--force", action="store_true", help="Force the regeneration of all files (parameters and outputs).") + +def parse_arguments_main(parsed_args): + """ + Parse the arguments for the main parameters. + """ + from pathlib import Path + main_dict = dict( + directory=parsed_args.directory, + paramdir=parsed_args.paramdir, + resultdir=parsed_args.resultdir, + workdir=parsed_args.workdir, + logdir=parsed_args.logdir, + mode=parsed_args.mode, + verbose=parsed_args.verbose, + simname=parsed_args.simname, + ICs_gen=parsed_args.ICs_gen, + seed=parsed_args.seed, + nthreads=parsed_args.nthreads, + execution=parsed_args.execution, + force=parsed_args.force, + ) + + ## Now for all the directories, if they are not specified, they are set using the main directory + if main_dict["paramdir"] is None: + main_dict["paramdir"] = main_dict["directory"]+"params/" + if main_dict["resultdir"] is None: + main_dict["resultdir"] = main_dict["directory"]+"results/" + if main_dict["workdir"] is None: + main_dict["workdir"] = main_dict["directory"]+"work/" + if main_dict["logdir"] is None: + main_dict["logdir"] = main_dict["directory"]+"logs/" + if main_dict["simname"] is None: + main_dict["simname"] = main_dict["mode"] + if main_dict["seed"] is None: + import numpy as np + main_dict["seed"] = np.random.randint(0, 2**32) + if main_dict["nthreads"] is None: + import os + main_dict["nthreads"] = os.cpu_count() + + # Create missing directories + Path(main_dict["paramdir"]).mkdir(parents=True, exist_ok=True) + Path(main_dict["resultdir"]).mkdir(parents=True, exist_ok=True) + Path(main_dict["workdir"]).mkdir(parents=True, exist_ok=True) + Path(main_dict["logdir"]).mkdir(parents=True, exist_ok=True) + + return main_dict \ No newline at end of file diff --git a/cosmo_params.py b/cosmo_params.py new file mode 100644 index 0000000..8ca1636 --- /dev/null +++ b/cosmo_params.py @@ -0,0 +1,66 @@ +from argparse import ArgumentParser + +def register_arguments_cosmo(parser:ArgumentParser): + """ + Register the arguments for the cosmological parameters. + """ + parser.add_argument("--h", type=float, default=0.6732, help="Hubble constant.") + parser.add_argument("-Om","--Omega_m", type=float, default=0.302, help="Matter density parameter.") + parser.add_argument("-Ob","--Omega_b", type=float, default=0.049, help="Baryon density parameter.") + parser.add_argument("-Oq","--Omega_q", type=float, default=0.6842, help="Dark energy density parameter.") + parser.add_argument("--n_s", type=float, default=0.968, help="Spectral index.") + parser.add_argument("-s8","--sigma8", type=float, default=0.815, help="RMS amplitude of the density fluctuations.") + parser.add_argument("--A_s", type=float, default=2.148752e-09, help="Amplitude of the primordial power spectrum.") + parser.add_argument("--Tcmb", type=float, default=2.7255, help="CMB temperature.") + parser.add_argument("--k_p", type=float, default=0.05, help="Pivot scale.") + parser.add_argument("--N_ur", type=float, default=2.046, help="Effective number of ultra-relativistic species.") + parser.add_argument("--m_nu1", type=float, default=0.06, help="Mass of the first neutrino species.") + parser.add_argument("--m_nu2", type=float, default=0.0, help="Mass of the second neutrino species.") + parser.add_argument("--m_nu3", type=float, default=0.0, help="Mass of the third neutrino species.") + parser.add_argument("--w_0", type=float, default=-1.0, help="Dark energy equation of state parameter.") + parser.add_argument("--w_a", type=float, default=0.0, help="Dark energy equation of state parameter.") + parser.add_argument("--fnl", type=float, default=100.0, help="Local non-Gaussianity parameter.") + parser.add_argument("--gnl", type=float, default=0.0, help="Equilateral non-Gaussianity parameter.") + + + +def parse_arguments_cosmo(parsed_args): + """ + Parse the arguments for the cosmological parameters. + """ + cosmo_dict = dict( + ## Common parameters + h=parsed_args.h, + Omega_m=parsed_args.Omega_m, + Omega_b=parsed_args.Omega_b, + Omega_q=parsed_args.Omega_q, + Omega_k=0.0, + n_s=parsed_args.n_s, + sigma8=parsed_args.sigma8, + w_0=parsed_args.w_0, + w_a=parsed_args.w_a, + A_s=parsed_args.A_s, + + ## monofonic ICs additional parameters + Tcmb=parsed_args.Tcmb, + k_p=parsed_args.k_p, + N_ur=parsed_args.N_ur, + m_nu1=parsed_args.m_nu1, + m_nu2=parsed_args.m_nu2, + m_nu3=parsed_args.m_nu3, + fnl=parsed_args.fnl, + gnl=parsed_args.gnl, + + ## sbmy ICs additional parameters + Omega_r=0.0, + k_max=10.0, + tau_reio=0.06, + WhichSpectrum="EH", # EH or class + w0_fld=parsed_args.w_0, + wa_fld=parsed_args.w_a, + ) + + return cosmo_dict + +def z2a(z): + return 1.0/(1.0+z) \ No newline at end of file diff --git a/low_level.py b/low_level.py new file mode 100644 index 0000000..f807c31 --- /dev/null +++ b/low_level.py @@ -0,0 +1,185 @@ +#!/usr/bin/env python +# ------------------------------------------------------------------------------------- +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, version 3. +# +# 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. +# +# The text of the license is located in the root directory of the source package. +# ------------------------------------------------------------------------------------- + +__author__ = "Tristan Hoellinger, Mayeul Aubin" +__version__ = "0.1" +__date__ = "2024" +__license__ = "GPLv3" + +""" +Tools to deal with low-level operations e.g. redirecting C stdout. +""" + +from contextlib import contextmanager +import ctypes +import io +import os, sys +import tempfile + +libc = ctypes.CDLL(None) +c_stdout = ctypes.c_void_p.in_dll(libc, "stdout") + + +# Taken from: +# https://eli.thegreenplace.net/2015/redirecting-all-kinds-of-stdout-in-python/ +@contextmanager +def stdout_redirector(stream): + """A context manager that redirects stdout to the given stream. For instance, this + could be used to redirect C code stdout to None (to avoid cluttering the terminal eg + when using tqdm). + + Args: + stream (file-like object): The stream to which stdout should be redirected. + Example: + >>> with stdout_redirector(stream): + >>> print("Hello world!") # Will be printed to stream instead of stdout. + """ + # The original fd stdout points to. Usually 1 on POSIX systems. + original_stdout_fd = sys.stdout.fileno() + + def _redirect_stdout(to_fd): + """Redirect stdout to the given file descriptor.""" + # Flush the C-level buffer stdout + libc.fflush(c_stdout) + # Flush and close sys.stdout - also closes the file descriptor (fd) + sys.stdout.close() + # Make original_stdout_fd point to the same file as to_fd + os.dup2(to_fd, original_stdout_fd) + # Create a new sys.stdout that points to the redirected fd + sys.stdout = io.TextIOWrapper(os.fdopen(original_stdout_fd, "wb")) + + # Save a copy of the original stdout fd in saved_stdout_fd + saved_stdout_fd = os.dup(original_stdout_fd) + try: + # Create a temporary file and redirect stdout to it + tfile = tempfile.TemporaryFile(mode="w+b") + _redirect_stdout(tfile.fileno()) + # Yield to caller, then redirect stdout back to the saved fd + yield + _redirect_stdout(saved_stdout_fd) + # Copy contents of temporary file to the given stream + tfile.flush() + tfile.seek(0, io.SEEK_SET) + stream.write(tfile.read()) + finally: + tfile.close() + os.close(saved_stdout_fd) + + +# Adapted from: +# https://eli.thegreenplace.net/2015/redirecting-all-kinds-of-stdout-in-python/ +c_stderr = ctypes.c_void_p.in_dll(libc, "stderr") + + +@contextmanager +def stderr_redirector(stream): + """A context manager that redirects stderr to the given stream. + For instance, this could be used to redirect C code stderr to None (to avoid + cluttering the terminal eg when using tqdm). + WARNING: this should be used with CAUTION as it redirects all stderr, so NONE OF THE + ERRORS WILL BE DISPLAYED, no matter the severity. + + Args: + stream (file-like object): The stream to which stdout should be redirected. + """ + # The original fd stdout points to. Usually 1 on POSIX systems. + original_stderr_fd = sys.stderr.fileno() + + def _redirect_stderr(to_fd): + """Redirect stderr to the given file descriptor.""" + # Flush the C-level buffer stderr + libc.fflush(c_stderr) + # Flush and close sys.stderr - also closes the file descriptor (fd) + sys.stderr.close() + # Make original_stderr_fd point to the same file as to_fd + os.dup2(to_fd, original_stderr_fd) + # Create a new sys.stderr that points to the redirected fd + sys.stderr = io.TextIOWrapper(os.fdopen(original_stderr_fd, "wb")) + + # Save a copy of the original stdout fd in saved_stdout_fd + saved_stderr_fd = os.dup(original_stderr_fd) + try: + # Create a temporary file and redirect stdout to it + tfile = tempfile.TemporaryFile(mode="w+b") + _redirect_stderr(tfile.fileno()) + # Yield to caller, then redirect stdout back to the saved fd + yield + _redirect_stderr(saved_stderr_fd) + # Copy contents of temporary file to the given stream + tfile.flush() + tfile.seek(0, io.SEEK_SET) + stream.write(tfile.read()) + finally: + tfile.close() + os.close(saved_stderr_fd) + + +def print_level(level:int, module:str): + """ + Generate a string with the current time, the module name and the level of the message. + """ + from datetime import datetime + + max_len_module = 10 + + date = datetime.now() + out="" + out+=date.strftime("%H:%M:%S") + out+=" | " + out+=module.upper().ljust(max_len_module) + out+=" | " + out+=">"*level + out+=" " + + return out + + +def print_message(message:str, level:int, module:str, verbose:int = 1): + """ + Print a message with a given level and module name. + """ + if verbose >= 1: + print(print_level(level, module)+message) + + +def print_header(verbose:int=1): + """ + Print the header of the program. + """ + if verbose >= 1: + from datetime import datetime + date = datetime.now() + width=40 + program_name="SBMY CONTROL" + print("#"*width) + print("# "+program_name.center(width-4)+" #") + print("# "+date.strftime("%Y-%m-%d %H:%M:%S").center(width-4)+" #") + print("#"*width) + print("") + + +def print_starting_module(module:str, verbose:int=1): + """ + Print the starting message of a module. + """ + if verbose >= 1: + print(print_level(0, module)+f"Starting {module} module.") + + +def print_ending_module(module:str, verbose:int=1): + """ + Print the ending message of a module. + """ + if verbose >= 1: + print(print_level(0, module)+f"Ending {module} module.") \ No newline at end of file