mirror of
https://github.com/hoellin/selfisys_public.git
synced 2025-06-08 09:01:11 +00:00
Initial commit: v0.1.0
This commit is contained in:
hoellinger
commit
c041384662
172 changed files with 45744 additions and 0 deletions
889
src/selfisys/sbmy_interface.py
Normal file
889
src/selfisys/sbmy_interface.py
Normal file
|
|
@ -0,0 +1,889 @@
|
|||
#!/usr/bin/env python3
|
||||
# ----------------------------------------------------------------------
|
||||
# Copyright (C) 2024 Tristan Hoellinger
|
||||
# Distributed under the GNU General Public License v3.0 (GPLv3).
|
||||
# See the LICENSE file in the root directory for details.
|
||||
# SPDX-License-Identifier: GPL-3.0-or-later
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
__author__ = "Tristan Hoellinger"
|
||||
__version__ = "0.1.0"
|
||||
__date__ = "2024"
|
||||
__license__ = "GPLv3"
|
||||
|
||||
"""Simbelmynë-related functions for the SelfiSys pipeline.
|
||||
"""
|
||||
|
||||
import os
|
||||
import gc
|
||||
from typing import Optional, List, Tuple
|
||||
|
||||
from selfisys.utils.logger import getCustomLogger, INDENT, UNINDENT
|
||||
|
||||
logger = getCustomLogger(__name__)
|
||||
|
||||
|
||||
def get_power_spectrum_from_cosmo(
|
||||
L,
|
||||
size,
|
||||
cosmo,
|
||||
fname_power_spectrum,
|
||||
force=False,
|
||||
):
|
||||
"""
|
||||
Compute a power spectrum from cosmological parameters and save it to
|
||||
disk.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
L : float
|
||||
Size of the simulation box (in Mpc/h).
|
||||
size : int
|
||||
Number of grid points along each axis.
|
||||
cosmo : dict
|
||||
Cosmological parameters (and infrastructure parameters).
|
||||
fname_power_spectrum : str
|
||||
Name (including path) of the power spectrum file to read/write.
|
||||
force : bool, optional
|
||||
If True, forces recomputation even if the file exists. Default
|
||||
is False.
|
||||
|
||||
Raises
|
||||
------
|
||||
OSError
|
||||
If file writing fails or the directory path is invalid.
|
||||
RuntimeError
|
||||
For unexpected issues during power spectrum computation.
|
||||
"""
|
||||
if not os.path.exists(fname_power_spectrum) or force:
|
||||
from pysbmy.power import PowerSpectrum
|
||||
|
||||
try:
|
||||
logger.debug("Computing power spectrum for L=%.2f, size=%d", L, size)
|
||||
P = PowerSpectrum(L, L, L, size, size, size, cosmo)
|
||||
P.write(fname_power_spectrum)
|
||||
logger.debug("Power spectrum written to %s", fname_power_spectrum)
|
||||
except OSError as e:
|
||||
logger.error("File write error at %s: %s", fname_power_spectrum, str(e))
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.critical("Unexpected error in power spectrum computation: %s", str(e))
|
||||
raise RuntimeError("get_power_spectrum_from_cosmo failed.") from e
|
||||
finally:
|
||||
gc.collect()
|
||||
|
||||
|
||||
def compute_Phi(
|
||||
G_ss_path,
|
||||
P_ss_path,
|
||||
g_obj,
|
||||
norm,
|
||||
AliasingCorr=True,
|
||||
verbosity=1,
|
||||
):
|
||||
"""
|
||||
Compute the summary statistics from a field object, based on a
|
||||
provided summary-statistics Fourier grid and baseline spectrum.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
G_ss_path : str
|
||||
Path to the FourierGrid file used for summary-statistics.
|
||||
P_ss_path : str
|
||||
Path to the baseline power spectrum file for normalisation.
|
||||
g_obj : Field
|
||||
Input field object from which to compute summary statistics.
|
||||
norm : ndarray
|
||||
Normalisation constants for the summary statistics.
|
||||
AliasingCorr : bool, optional
|
||||
Whether to apply aliasing correction. Default is True.
|
||||
verbosity : int, optional
|
||||
Verbosity level (0=quiet, 1=normal, 2=debug). Default 1.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Phi : ndarray
|
||||
Vector of summary statistics.
|
||||
|
||||
Raises
|
||||
------
|
||||
OSError
|
||||
If file reading fails at G_ss_path or P_ss_path.
|
||||
RuntimeError
|
||||
If unexpected issues occur during computation.
|
||||
"""
|
||||
from pysbmy.correlations import get_autocorrelation
|
||||
from pysbmy.power import FourierGrid, PowerSpectrum
|
||||
from pysbmy import c_double
|
||||
from io import BytesIO
|
||||
|
||||
try:
|
||||
logger.debug("Reading FourierGrid from %s", G_ss_path)
|
||||
G_ss = FourierGrid.read(G_ss_path)
|
||||
|
||||
if verbosity > 1:
|
||||
Pk, _ = get_autocorrelation(g_obj, G_ss, AliasingCorr=AliasingCorr)
|
||||
else:
|
||||
from selfisys.utils.low_level import stdout_redirector
|
||||
|
||||
f = BytesIO()
|
||||
with stdout_redirector(f):
|
||||
Pk, _ = get_autocorrelation(g_obj, G_ss, AliasingCorr=AliasingCorr)
|
||||
f.close()
|
||||
|
||||
logger.debug("Reading baseline PowerSpectrum from %s", P_ss_path)
|
||||
P_ss = PowerSpectrum.read(P_ss_path)
|
||||
Phi = Pk / (norm * P_ss.powerspectrum)
|
||||
|
||||
del G_ss, P_ss
|
||||
gc.collect()
|
||||
return Phi.astype(c_double)
|
||||
except OSError as e:
|
||||
logger.error("File not found or inaccessible: %s", str(e))
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.critical("Unexpected error in compute_Phi: %s", str(e))
|
||||
raise RuntimeError("compute_Phi failed.") from e
|
||||
finally:
|
||||
gc.collect()
|
||||
|
||||
|
||||
def generate_white_noise_Field(
|
||||
L,
|
||||
size,
|
||||
seedphase,
|
||||
fname_whitenoise,
|
||||
seedname_whitenoise,
|
||||
force_phase=False,
|
||||
):
|
||||
"""
|
||||
Generate a white noise realisation in physical space and write it to
|
||||
disk.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
L : float
|
||||
Size of the simulation box (in Mpc/h).
|
||||
size : int
|
||||
Number of grid points along each axis.
|
||||
seedphase : int or list of int
|
||||
User-provided seed to generate the initial white noise.
|
||||
fname_whitenoise : str
|
||||
File path to write the white noise realisation.
|
||||
seedname_whitenoise : str
|
||||
File path to write the seed state of the RNG.
|
||||
force_phase : bool, optional
|
||||
If True, forces regeneration of the random phases. Default is
|
||||
False.
|
||||
|
||||
Raises
|
||||
------
|
||||
OSError
|
||||
If file writing fails or directory paths are invalid.
|
||||
RuntimeError
|
||||
For unexpected issues.
|
||||
"""
|
||||
if not os.path.exists(fname_whitenoise) or force_phase:
|
||||
import numpy as np
|
||||
from pysbmy.field import BaseField
|
||||
|
||||
try:
|
||||
logger.debug("Generating white noise for L=%.2f, size=%d", L, size)
|
||||
rng = np.random.default_rng(seedphase)
|
||||
|
||||
logger.debug("Saving RNG state to %s", seedname_whitenoise)
|
||||
np.save(seedname_whitenoise, rng.bit_generator.state)
|
||||
with open(seedname_whitenoise + ".txt", "w") as f:
|
||||
f.write(str(rng.bit_generator.state))
|
||||
|
||||
data = rng.standard_normal(size=size**3)
|
||||
wn = BaseField(L, L, L, 0, 0, 0, 1, size, size, size, data)
|
||||
del data
|
||||
|
||||
wn.write(fname_whitenoise)
|
||||
logger.debug("White noise field written to %s", fname_whitenoise)
|
||||
del wn
|
||||
except OSError as e:
|
||||
logger.error("Writing white noise failed at '%s': %s", fname_whitenoise, str(e))
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.critical("Unexpected error in generate_white_noise_Field: %s", str(e))
|
||||
raise RuntimeError("generate_white_noise_Field failed.") from e
|
||||
finally:
|
||||
gc.collect()
|
||||
|
||||
|
||||
def setup_sbmy_parfiles(
|
||||
d,
|
||||
cosmology,
|
||||
file_names,
|
||||
hiddenbox_params,
|
||||
force=False,
|
||||
):
|
||||
"""Set up Simbelmynë parameter file (please refer to the Simbelmynë
|
||||
documentation for more details).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
d : int
|
||||
Index (from 1 to S) specifying a direction in parameter space, 0
|
||||
for the expansion point, or -1 for mock data.
|
||||
cosmology : array, double, dimension=5
|
||||
Cosmological parameters.
|
||||
file_names : dict
|
||||
Dictionary containing the names of the input/output files for
|
||||
the simulation.
|
||||
hiddenbox_params : dict
|
||||
See the `HiddenBox` class for more details.
|
||||
force : bool, optional, default=False
|
||||
If True, forces recompute the simulation parameter files.
|
||||
|
||||
"""
|
||||
from os.path import exists
|
||||
|
||||
fname_simparfile = file_names["fname_simparfile"]
|
||||
fname_power_spectrum = file_names["fname_power_spectrum"]
|
||||
fname_whitenoise = file_names["fname_whitenoise"]
|
||||
fname_outputinitialdensity = file_names["fname_outputinitialdensity"]
|
||||
fnames_outputrealspacedensity = file_names["fnames_outputrealspacedensity"]
|
||||
fnames_outputdensity = file_names["fnames_outputdensity"]
|
||||
fnames_outputLPTdensity = file_names["fnames_outputLPTdensity"]
|
||||
|
||||
Npop = hiddenbox_params["Npop"]
|
||||
Np0 = hiddenbox_params["Np0"]
|
||||
Npm0 = hiddenbox_params["Npm0"]
|
||||
size = hiddenbox_params["size"]
|
||||
L = hiddenbox_params["L"]
|
||||
Ntimesteps = hiddenbox_params["Ntimesteps"]
|
||||
sim_params = hiddenbox_params["sim_params"]
|
||||
eff_redshifts = hiddenbox_params["eff_redshifts"]
|
||||
TimeSteps = hiddenbox_params["TimeSteps"]
|
||||
TimeStepDistribution = hiddenbox_params["TimeStepDistribution"]
|
||||
modified_selfi = hiddenbox_params["modified_selfi"]
|
||||
fsimdir = hiddenbox_params["fsimdir"]
|
||||
|
||||
if not exists(fname_simparfile + "_{}.sbmy".format(Npop)) or force:
|
||||
from pysbmy import param_file
|
||||
from re import search
|
||||
from selfisys.global_parameters import BASEID_OBS
|
||||
|
||||
if TimeSteps is not None and eff_redshifts is None:
|
||||
raise ValueError("TimeSteps must be provided if eff_redshifts is None.")
|
||||
|
||||
regex = r"([a-zA-Z]+)(\d+)?([a-zA-Z]+)?(\d+)?([a-zA-Z]+)?"
|
||||
m = search(regex, sim_params)
|
||||
if m.group(1) == "std":
|
||||
# Single LPT+COLA/PM Simbelmynë data card with linear time
|
||||
# stepping
|
||||
if m.group(2)[0] == "0":
|
||||
RedshiftLPT = float("0." + m.group(2)[1:])
|
||||
else:
|
||||
RedshiftLPT = int(m.group(2))
|
||||
|
||||
RedshiftFCs = 0.0
|
||||
WriteLPTSnapshot = 0
|
||||
WriteLPTDensity = 0
|
||||
match m.group(3):
|
||||
case "RSD":
|
||||
ModulePMCOLA = 0
|
||||
EvolutionMode = 2
|
||||
NumberOfTimeSteps = 0
|
||||
RedshiftFCs = RedshiftLPT
|
||||
NonLinearRSD = 1
|
||||
case "PM":
|
||||
ModulePMCOLA = 1
|
||||
EvolutionMode = 1
|
||||
NumberOfTimeSteps = m.group(4)
|
||||
NonLinearRSD = 0 if (m.group(5) and m.group(5)[:3] == "lin") else 1
|
||||
case "COLA":
|
||||
ModulePMCOLA = 1
|
||||
EvolutionMode = 2
|
||||
NumberOfTimeSteps = m.group(4)
|
||||
NonLinearRSD = 0 if (m.group(5) and m.group(5)[:3] == "lin") else 1
|
||||
case _:
|
||||
raise ValueError("sim_params = {} not valid".format(sim_params))
|
||||
NumberOfTimeSteps = int(m.group(4)) if m.group(4) is not None else 0
|
||||
elif m.group(1) == "custom":
|
||||
# Single LPT+COLA/PM Simbelmynë card with user-provided time
|
||||
# stepping object
|
||||
RedshiftLPT = int(m.group(2))
|
||||
match m.group(3):
|
||||
case None:
|
||||
ModulePMCOLA = 0
|
||||
EvolutionMode = 2
|
||||
case "PM":
|
||||
ModulePMCOLA = 1
|
||||
EvolutionMode = 1
|
||||
NonLinearRSD = 0 if (m.group(5) and m.group(5)[:3] == "lin") else 1
|
||||
case "COLA":
|
||||
ModulePMCOLA = 1
|
||||
EvolutionMode = 2
|
||||
NonLinearRSD = 0 if (m.group(5) and m.group(5)[:3] == "lin") else 1
|
||||
case _:
|
||||
raise ValueError("sim_params = {} not valid".format(sim_params))
|
||||
if TimeStepDistribution is None:
|
||||
raise ValueError("TimeStepDistribution must be provided for 'custom'.")
|
||||
elif m.group(1) == "splitLPT":
|
||||
# Use as many Simbelmynë data cards as there are populations
|
||||
# of galaxies
|
||||
if eff_redshifts is None:
|
||||
raise ValueError("eff_redshifts must be provided for 'splitLPT'.")
|
||||
elif len(eff_redshifts) != Ntimesteps:
|
||||
raise ValueError("len(eff_redshifts) != Ntimesteps")
|
||||
elif m.group(1) == "split":
|
||||
# Use as many Simbelmynë data cards as there are populations
|
||||
# of galaxies
|
||||
if TimeStepDistribution is None:
|
||||
raise ValueError("TimeStepDistribution must be for 'split'.")
|
||||
if eff_redshifts is None:
|
||||
raise ValueError("eff_redshifts must be provided for 'split'.")
|
||||
elif len(eff_redshifts) != Ntimesteps:
|
||||
raise ValueError("len(eff_redshifts) != Ntimesteps")
|
||||
|
||||
RedshiftLPT = int(m.group(2))
|
||||
match m.group(3):
|
||||
case "RSD":
|
||||
ModulePMCOLA = 1
|
||||
EvolutionMode = 2
|
||||
NonLinearRSD = 1
|
||||
case _:
|
||||
raise ValueError("sim_params = {} not valid".format(sim_params))
|
||||
NumberOfTimeSteps = int(m.group(4)) if m.group(4) is not None else 0
|
||||
else:
|
||||
raise ValueError("sim_params = {} not valid" + sim_params)
|
||||
|
||||
if sim_params[-3:] == BASEID_OBS:
|
||||
from selfisys.global_parameters import (
|
||||
h_obs as h,
|
||||
Omega_b_obs as Omega_b,
|
||||
Omega_m_obs as Omega_m,
|
||||
nS_obs as nS,
|
||||
sigma8_obs as sigma8,
|
||||
)
|
||||
else:
|
||||
if modified_selfi:
|
||||
# Treat the cosmological parameters as nuisance
|
||||
# parameters within the hidden box forward model
|
||||
h, Omega_b, Omega_m, nS, sigma8 = cosmology
|
||||
else:
|
||||
# Fix the fiducial cosmology within the hidden box
|
||||
from selfisys.global_parameters import (
|
||||
h_planck as h,
|
||||
Omega_b_planck as Omega_b,
|
||||
Omega_m_planck as Omega_m,
|
||||
nS_planck as nS,
|
||||
sigma8_planck as sigma8,
|
||||
)
|
||||
|
||||
if d < 0: # -1 for mock data, -2 to recompute the observations
|
||||
WriteInitialConditions = 1
|
||||
WriteDensities = 1 # also write real space density fields
|
||||
else: # d=0 for expansion point or d>0 for the gradients
|
||||
WriteInitialConditions = 0
|
||||
WriteDensities = 1 # also write real space density fields
|
||||
|
||||
if m.group(1) == "std":
|
||||
S = param_file( ## Module LPT ##
|
||||
ModuleLPT=1,
|
||||
# Basic setup:
|
||||
Particles=Np0,
|
||||
Mesh=size,
|
||||
BoxSize=L,
|
||||
corner0=0.0,
|
||||
corner1=0.0,
|
||||
corner2=0.0,
|
||||
# Initial conditions:
|
||||
ICsMode=1,
|
||||
WriteICsRngState=0,
|
||||
WriteInitialConditions=WriteInitialConditions,
|
||||
InputWhiteNoise=fname_whitenoise,
|
||||
OutputInitialConditions=fname_outputinitialdensity,
|
||||
# Power spectrum:
|
||||
InputPowerSpectrum=fname_power_spectrum,
|
||||
# Final conditions for LPT:
|
||||
RedshiftLPT=RedshiftLPT,
|
||||
WriteLPTSnapshot=WriteLPTSnapshot,
|
||||
WriteLPTDensity=WriteLPTDensity,
|
||||
OutputLPTDensity=fnames_outputLPTdensity,
|
||||
####################
|
||||
## Module PM/COLA ##
|
||||
####################
|
||||
ModulePMCOLA=ModulePMCOLA,
|
||||
EvolutionMode=EvolutionMode, # 1 for PM, 2 for COLA
|
||||
ParticleMesh=Npm0,
|
||||
NumberOfTimeSteps=NumberOfTimeSteps,
|
||||
# Final snapshot:
|
||||
RedshiftFCs=RedshiftFCs,
|
||||
WriteFinalSnapshot=0,
|
||||
WriteFinalDensity=WriteDensities,
|
||||
OutputFinalDensity=fnames_outputrealspacedensity[0],
|
||||
#########
|
||||
## RSD ##
|
||||
#########
|
||||
ModuleRSD=1,
|
||||
WriteIntermediaryRSD=0,
|
||||
DoNonLinearMapping=NonLinearRSD,
|
||||
WriteRSDensity=1,
|
||||
OutputRSDensity=fnames_outputdensity[0],
|
||||
#############################
|
||||
## Cosmological parameters ##
|
||||
#############################
|
||||
h=h,
|
||||
Omega_q=1.0 - Omega_m,
|
||||
Omega_b=Omega_b,
|
||||
Omega_m=Omega_m,
|
||||
Omega_k=0.0,
|
||||
n_s=nS,
|
||||
sigma8=sigma8,
|
||||
w0_fld=-1.0,
|
||||
wa_fld=0.0,
|
||||
)
|
||||
S.write(fname_simparfile + "_{}.sbmy".format(Npop))
|
||||
elif m.group(1) == "custom":
|
||||
RedshiftFCs = eff_redshifts
|
||||
fname_outputdensity = (
|
||||
fnames_outputdensity[0][: fnames_outputdensity[0].rfind("_")] + ".h5"
|
||||
)
|
||||
S = param_file( ## Module LPT ##
|
||||
ModuleLPT=1,
|
||||
# Basic setup:
|
||||
Particles=Np0,
|
||||
Mesh=size,
|
||||
BoxSize=L,
|
||||
corner0=0.0,
|
||||
corner1=0.0,
|
||||
corner2=0.0,
|
||||
# Initial conditions:
|
||||
ICsMode=1,
|
||||
WriteICsRngState=0,
|
||||
WriteInitialConditions=WriteInitialConditions,
|
||||
InputWhiteNoise=fname_whitenoise,
|
||||
OutputInitialConditions=fname_outputinitialdensity,
|
||||
# Power spectrum:
|
||||
InputPowerSpectrum=fname_power_spectrum,
|
||||
# Final conditions for LPT:
|
||||
RedshiftLPT=RedshiftLPT,
|
||||
WriteLPTSnapshot=0,
|
||||
WriteLPTDensity=0,
|
||||
####################
|
||||
## Module PM/COLA ##
|
||||
####################
|
||||
ModulePMCOLA=ModulePMCOLA,
|
||||
EvolutionMode=EvolutionMode, # 1 for PM, 2 for COLA
|
||||
ParticleMesh=Npm0,
|
||||
OutputKickBase=fsimdir + "/data/cola_kick_",
|
||||
# Final snapshot:
|
||||
RedshiftFCs=RedshiftFCs,
|
||||
WriteFinalSnapshot=0,
|
||||
WriteFinalDensity=0,
|
||||
OutputFinalDensity=fnames_outputrealspacedensity[0],
|
||||
# Intermediate snapshots:
|
||||
WriteSnapshots=0,
|
||||
WriteDensities=WriteDensities,
|
||||
OutputDensitiesBase=fnames_outputrealspacedensity[0][
|
||||
: fnames_outputrealspacedensity[0].rfind("_")
|
||||
]
|
||||
+ "_",
|
||||
OutputDensitiesExt=".h5",
|
||||
############################
|
||||
## Time step distribution ##
|
||||
############################
|
||||
TimeStepDistribution=TimeStepDistribution,
|
||||
ModifiedDiscretization=1, # Modified KD discretisation
|
||||
n_LPT=-2.5, # Exponent for the Ansatz in KD operators
|
||||
#########
|
||||
## RSD ##
|
||||
#########
|
||||
ModuleRSD=1,
|
||||
WriteIntermediaryRSD=1,
|
||||
DoNonLinearMapping=NonLinearRSD,
|
||||
WriteRSDensity=1,
|
||||
OutputRSDensity=fname_outputdensity,
|
||||
#############################
|
||||
## Cosmological parameters ##
|
||||
#############################
|
||||
h=h,
|
||||
Omega_q=1.0 - Omega_m,
|
||||
Omega_b=Omega_b,
|
||||
Omega_m=Omega_m,
|
||||
Omega_k=0.0,
|
||||
n_s=nS,
|
||||
sigma8=sigma8,
|
||||
w0_fld=-1.0,
|
||||
wa_fld=0.0,
|
||||
)
|
||||
S.write(fname_simparfile + "_{}.sbmy".format(Npop))
|
||||
elif m.group(1) == "split":
|
||||
datadir = fsimdir + "/data/"
|
||||
RedshiftFCs = eff_redshifts[0]
|
||||
|
||||
# Write the parameter file for the first simulation
|
||||
S = param_file(
|
||||
################
|
||||
## Module LPT ##
|
||||
################
|
||||
ModuleLPT=1,
|
||||
# Basic setup:
|
||||
Particles=Np0,
|
||||
Mesh=size,
|
||||
BoxSize=L,
|
||||
corner0=0.0,
|
||||
corner1=0.0,
|
||||
corner2=0.0,
|
||||
# Initial conditions:
|
||||
ICsMode=1,
|
||||
WriteICsRngState=0,
|
||||
WriteInitialConditions=WriteInitialConditions,
|
||||
InputWhiteNoise=fname_whitenoise,
|
||||
OutputInitialConditions=fname_outputinitialdensity,
|
||||
# Power spectrum:
|
||||
InputPowerSpectrum=fname_power_spectrum,
|
||||
# Final conditions for LPT:
|
||||
RedshiftLPT=RedshiftLPT,
|
||||
WriteLPTSnapshot=0,
|
||||
WriteLPTDensity=0,
|
||||
####################
|
||||
## Module PM/COLA ##
|
||||
####################
|
||||
ModulePMCOLA=ModulePMCOLA,
|
||||
EvolutionMode=EvolutionMode,
|
||||
ParticleMesh=Npm0,
|
||||
OutputKickBase=datadir + "cola_kick_0_",
|
||||
# Final snapshot:
|
||||
RedshiftFCs=RedshiftFCs,
|
||||
WriteFinalSnapshot=1,
|
||||
OutputFinalSnapshot=datadir + "cola_snapshot_0.gadget3",
|
||||
WriteFinalDensity=1,
|
||||
OutputFinalDensity=fnames_outputrealspacedensity[0],
|
||||
WriteLPTDisplacements=1,
|
||||
OutputPsiLPT1=datadir + "lpt_psi1_0.h5",
|
||||
OutputPsiLPT2=datadir + "lpt_psi2_0.h5",
|
||||
############################
|
||||
## Time step distribution ##
|
||||
############################
|
||||
TimeStepDistribution=TimeStepDistribution[0],
|
||||
ModifiedDiscretization=1,
|
||||
#########
|
||||
## RSD ##
|
||||
#########
|
||||
ModuleRSD=1,
|
||||
WriteIntermediaryRSD=0,
|
||||
DoNonLinearMapping=NonLinearRSD,
|
||||
WriteRSDensity=1,
|
||||
OutputRSDensity=fnames_outputdensity[0],
|
||||
#############################
|
||||
## Cosmological parameters ##
|
||||
#############################
|
||||
h=h,
|
||||
Omega_q=1.0 - Omega_m,
|
||||
Omega_b=Omega_b,
|
||||
Omega_m=Omega_m,
|
||||
Omega_k=0.0,
|
||||
n_s=nS,
|
||||
sigma8=sigma8,
|
||||
w0_fld=-1.0,
|
||||
wa_fld=0.0,
|
||||
)
|
||||
S.write(fname_simparfile + "_pop0.sbmy")
|
||||
|
||||
for i in range(1, Ntimesteps):
|
||||
RedshiftFCs = eff_redshifts[i]
|
||||
|
||||
S = param_file(
|
||||
ModuleLPT=0,
|
||||
# Basic setup:
|
||||
Particles=Np0,
|
||||
Mesh=size,
|
||||
BoxSize=L,
|
||||
corner0=0.0,
|
||||
corner1=0.0,
|
||||
corner2=0.0,
|
||||
InputPsiLPT1=datadir + "lpt_psi1_0.h5",
|
||||
InputPsiLPT2=datadir + "lpt_psi2_0.h5",
|
||||
####################
|
||||
## Module PM/COLA ##
|
||||
####################
|
||||
ModulePMCOLA=ModulePMCOLA,
|
||||
InputPMCOLASnapshot=datadir + "cola_snapshot_{:d}.gadget3".format(i - 1),
|
||||
EvolutionMode=EvolutionMode,
|
||||
ParticleMesh=Npm0,
|
||||
OutputKickBase=datadir + "cola_kick_{:d}_".format(i),
|
||||
# Final snapshot:
|
||||
RedshiftFCs=RedshiftFCs,
|
||||
WriteFinalSnapshot=1,
|
||||
OutputFinalSnapshot=datadir + "cola_snapshot_{:d}.gadget3".format(i),
|
||||
WriteFinalDensity=1,
|
||||
OutputFinalDensity=fnames_outputrealspacedensity[::-1][i],
|
||||
WriteLPTDisplacements=0,
|
||||
############################
|
||||
## Time step distribution ##
|
||||
############################
|
||||
TimeStepDistribution=TimeStepDistribution[i],
|
||||
ModifiedDiscretization=1,
|
||||
#########
|
||||
## RSD ##
|
||||
#########
|
||||
ModuleRSD=1,
|
||||
WriteIntermediaryRSD=0,
|
||||
DoNonLinearMapping=NonLinearRSD,
|
||||
WriteRSDensity=1,
|
||||
OutputRSDensity=fnames_outputdensity[i],
|
||||
#############################
|
||||
## Cosmological parameters ##
|
||||
#############################
|
||||
h=h,
|
||||
Omega_q=1.0 - Omega_m,
|
||||
Omega_b=Omega_b,
|
||||
Omega_m=Omega_m,
|
||||
Omega_k=0.0,
|
||||
n_s=nS,
|
||||
sigma8=sigma8,
|
||||
w0_fld=-1.0,
|
||||
wa_fld=0.0,
|
||||
)
|
||||
S.write(fname_simparfile + "_pop{}.sbmy".format(i))
|
||||
elif m.group(1) == "splitLPT":
|
||||
datadir = fsimdir + "/data/"
|
||||
RedshiftLPT = eff_redshifts[0]
|
||||
RedshiftFCs = eff_redshifts[0]
|
||||
|
||||
# Write the parameter file for the first simulation
|
||||
S = param_file(
|
||||
################
|
||||
## Module LPT ##
|
||||
################
|
||||
ModuleLPT=1,
|
||||
# Basic setup:
|
||||
Particles=Np0,
|
||||
Mesh=size,
|
||||
BoxSize=L,
|
||||
corner0=0.0,
|
||||
corner1=0.0,
|
||||
corner2=0.0,
|
||||
# Initial conditions:
|
||||
ICsMode=1,
|
||||
WriteICsRngState=0,
|
||||
InputWhiteNoise=fname_whitenoise,
|
||||
OutputInitialConditions=fname_outputinitialdensity,
|
||||
InputPowerSpectrum=fname_power_spectrum,
|
||||
# Final conditions for LPT:
|
||||
RedshiftLPT=RedshiftLPT,
|
||||
WriteLPTSnapshot=0,
|
||||
WriteLPTDensity=0,
|
||||
# Final snapshot:
|
||||
RedshiftFCs=RedshiftFCs,
|
||||
WriteFinalDensity=0,
|
||||
WriteLPTDisplacements=0,
|
||||
#########
|
||||
## RSD ##
|
||||
#########
|
||||
ModuleRSD=1,
|
||||
WriteIntermediaryRSD=0,
|
||||
DoNonLinearMapping=NonLinearRSD,
|
||||
WriteRSDensity=1,
|
||||
OutputRSDensity=fnames_outputdensity[0],
|
||||
#############################
|
||||
## Cosmological parameters ##
|
||||
#############################
|
||||
h=h,
|
||||
Omega_q=1.0 - Omega_m,
|
||||
Omega_b=Omega_b,
|
||||
Omega_m=Omega_m,
|
||||
Omega_k=0.0,
|
||||
n_s=nS,
|
||||
sigma8=sigma8,
|
||||
w0_fld=-1.0,
|
||||
wa_fld=0.0,
|
||||
)
|
||||
S.write(fname_simparfile + "_pop0.sbmy")
|
||||
|
||||
for i in range(1, Ntimesteps):
|
||||
RedshiftLPT = eff_redshifts[i]
|
||||
RedshiftFCs = eff_redshifts[i]
|
||||
|
||||
S = param_file(
|
||||
################
|
||||
## Module LPT ##
|
||||
################
|
||||
ModuleLPT=1,
|
||||
# Basic setup:
|
||||
Particles=Np0,
|
||||
Mesh=size,
|
||||
BoxSize=L,
|
||||
corner0=0.0,
|
||||
corner1=0.0,
|
||||
corner2=0.0,
|
||||
# Initial conditions:
|
||||
ICsMode=1,
|
||||
WriteICsRngState=0,
|
||||
InputWhiteNoise=fname_whitenoise,
|
||||
OutputInitialConditions=fname_outputinitialdensity,
|
||||
InputPowerSpectrum=fname_power_spectrum,
|
||||
# Final conditions for LPT:
|
||||
RedshiftLPT=RedshiftLPT,
|
||||
WriteLPTDensity=0,
|
||||
WriteLPTDisplacements=0,
|
||||
#########
|
||||
## RSD ##
|
||||
#########
|
||||
ModuleRSD=1,
|
||||
WriteIntermediaryRSD=0,
|
||||
DoNonLinearMapping=NonLinearRSD,
|
||||
WriteRSDensity=1,
|
||||
OutputRSDensity=fnames_outputdensity[i],
|
||||
#############################
|
||||
## Cosmological parameters ##
|
||||
#############################
|
||||
h=h,
|
||||
Omega_q=1.0 - Omega_m,
|
||||
Omega_b=Omega_b,
|
||||
Omega_m=Omega_m,
|
||||
Omega_k=0.0,
|
||||
n_s=nS,
|
||||
sigma8=sigma8,
|
||||
w0_fld=-1.0,
|
||||
wa_fld=0.0,
|
||||
)
|
||||
S.write(fname_simparfile + "_pop{}.sbmy".format(i))
|
||||
|
||||
|
||||
def handle_time_stepping(
|
||||
aa: List[float],
|
||||
total_steps: int,
|
||||
modeldir: str,
|
||||
figuresdir: str,
|
||||
sim_params: str,
|
||||
force: bool = False,
|
||||
) -> Tuple[Optional[List[int]], Optional[float]]:
|
||||
"""
|
||||
Create and merge individual time-stepping objects.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
aa : list of float
|
||||
List of scale factors in ascending order.
|
||||
total_steps : int
|
||||
Total number of time steps to distribute among the provided
|
||||
scale factors.
|
||||
modeldir : str
|
||||
Directory path to store generated time-stepping files.
|
||||
figuresdir : str
|
||||
Directory path to store time-stepping plots.
|
||||
sim_params : str
|
||||
Simulation parameter string (e.g., "custom", "std", "nograv").
|
||||
force : bool, optional
|
||||
Whether to force recompute the time-stepping files. Default is
|
||||
False.
|
||||
|
||||
Returns
|
||||
-------
|
||||
merged_path : str
|
||||
Path to the merged time-stepping file.
|
||||
indices_steps_cumul : list of int or None
|
||||
Cumulative indices for the distributed steps. Returns None if
|
||||
using splitLPT or if `sim_params` indicates an alternative
|
||||
strategy.
|
||||
eff_redshifts : float or None
|
||||
Effective redshift derived from the final scale factor in
|
||||
'custom' or 'nograv' mode. None otherwise.
|
||||
|
||||
Raises
|
||||
------
|
||||
NotImplementedError
|
||||
If a unsupported time-stepping strategy is used.
|
||||
OSError
|
||||
If file or directory operations fail.
|
||||
RuntimeError
|
||||
If unexpected issues occur during time-stepping setup.
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
from pysbmy.timestepping import StandardTimeStepping, read_timestepping
|
||||
from selfisys.utils.plot_utils import reset_plotting, setup_plotting
|
||||
from selfisys.utils.timestepping import merge_nTS
|
||||
|
||||
logger.info("Evaluating time-stepping strategy: %s", sim_params)
|
||||
|
||||
indices_steps_cumul = None
|
||||
eff_redshifts = None
|
||||
|
||||
isstd = sim_params.startswith("std")
|
||||
splitLPT = sim_params.startswith("splitLPT")
|
||||
|
||||
try:
|
||||
# Case 1: standard approach with distributed steps
|
||||
if not isstd and not splitLPT:
|
||||
reset_plotting() # Revert to default plotting style
|
||||
merged_path = modeldir + "merged.h5"
|
||||
|
||||
# Create time-stepping
|
||||
if not os.path.exists(merged_path) or force:
|
||||
logger.info("Setting up time-stepping...")
|
||||
|
||||
# Distribute steps among the scale factors
|
||||
nsteps = [
|
||||
round((aa[i + 1] - aa[i]) / (aa[-1] - aa[0]) * total_steps)
|
||||
for i in range(len(aa) - 1)
|
||||
]
|
||||
# Adjust the largest gap if rounding caused a mismatch
|
||||
if sum(nsteps) != total_steps:
|
||||
nsteps[nsteps.index(max(nsteps))] += total_steps - sum(nsteps)
|
||||
|
||||
indices_steps_cumul = list(np.cumsum(nsteps) - 1)
|
||||
np.save(modeldir + "indices_steps_cumul.npy", indices_steps_cumul)
|
||||
|
||||
INDENT()
|
||||
logger.diagnostic("Generating individual time-stepping objects...")
|
||||
|
||||
TS_paths = []
|
||||
for i, (ai, af) in enumerate(zip(aa[:-1], aa[1:])):
|
||||
snapshots = np.full((nsteps[i]), False)
|
||||
snapshots[-1] = True # Mark last step as a snapshot
|
||||
TS = StandardTimeStepping(ai, af, snapshots, 0)
|
||||
TS_path = modeldir + f"ts{i+1}.h5"
|
||||
TS.write(str(TS_path))
|
||||
TS_paths.append(TS_path)
|
||||
|
||||
# Ensure the timestepping object are readable and plot
|
||||
for i, path_ts in enumerate(TS_paths):
|
||||
read_timestepping(str(path_ts)).plot(path=str(figuresdir + f"TS{i}.png"))
|
||||
logger.diagnostic("Generating individual time-stepping objects done.")
|
||||
|
||||
logger.diagnostic("Merging time-stepping...")
|
||||
merge_nTS([str(p) for p in TS_paths], merged_path)
|
||||
TS_merged = read_timestepping(merged_path)
|
||||
TS_merged.plot(path=str(figuresdir + "TS_merged.png"))
|
||||
|
||||
# Restore the project's plotting style
|
||||
setup_plotting()
|
||||
logger.diagnostic("Merging time-stepping done.")
|
||||
UNINDENT()
|
||||
logger.info("Setting up time-stepping done.")
|
||||
else:
|
||||
logger.diagnostic("Time-stepping objects already computed.")
|
||||
|
||||
# Evaluate final effective redshift
|
||||
if sim_params.startswith("custom") or sim_params.startswith("nograv"):
|
||||
eff_redshifts = 1 / aa[-1] - 1
|
||||
else:
|
||||
raise NotImplementedError("Time-stepping strategy not yet implemented.")
|
||||
|
||||
# Case 2: splitted
|
||||
elif splitLPT:
|
||||
indices_steps_cumul = [f"pop{i}" for i in range(1, len(aa))]
|
||||
eff_redshifts = [1 / a - 1 for a in aa[1:]]
|
||||
|
||||
# Case 3: other
|
||||
else:
|
||||
logger.diagnostic("Standard time-stepping or no special distribution required.")
|
||||
|
||||
except OSError as e:
|
||||
logger.error("File or directory access error in handle_time_stepping: %s", str(e))
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.critical("An error occurred during time-stepping setup: %s", str(e))
|
||||
raise RuntimeError("Time-stepping setup failed.") from e
|
||||
finally:
|
||||
gc.collect()
|
||||
|
||||
return merged_path, indices_steps_cumul, eff_redshifts
|
||||
Loading…
Add table
Add a link
Reference in a new issue