Compare commits

...
Sign in to create a new pull request.

17 commits

4 changed files with 129 additions and 11 deletions

View file

@ -64,10 +64,14 @@ def stdout_redirector(stream):
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)
@ -237,7 +241,7 @@ def progress_bar_from_logfile(filename:str, desc:str="", verbose:int=1, **kwargs
from tqdm import tqdm
from time import sleep
k=0
limit=600
limit=6000
update_interval=0.2
wait_until_file_exists(filename, verbose=verbose, limit=limit)
current_operation, total_operations = get_progress_from_logfile(filename)

View file

@ -47,6 +47,15 @@ def register_arguments_card(parser:ArgumentParser):
parser.add_argument("--OutputFCsDensity", type=str, default=None, help="Output FCs density file.")
parser.add_argument("--OutputFCsSnapshot", type=str, default=None, help="Output FCs snapshot file.")
parser.add_argument("--OutputRngStateLPT", type=str, default=None, help="Output RNG state file.")
## Tests with phiBCs and density
parser.add_argument("--WriteGravPot", type=bool, default=False, help="Write gravitational potential.")
parser.add_argument("--OutputGravitationalPotentialBase", type=str, default=None, help="Output gravitational potential base.")
parser.add_argument("--MeshGravPot", type=int, default=None, help="Mesh for gravitational potential.")
parser.add_argument("--WriteDensity", type=bool, default=False, help="Write density.")
parser.add_argument("--OutputDensityBase", type=str, default=None, help="Output density base.")
parser.add_argument("--MeshDensity", type=int, default=None, help="Mesh for density.")
parser.add_argument("--LoadPhiBCs", type=bool, default=False, help="Load phiBCs.")
parser.add_argument("--InputPhiBCsBase", type=str, default=None, help="Input phiBCs file base.")
parser.add_argument("--WriteReferenceFrame", type=bool, default=False, help="Write reference frame (COCA).")
parser.add_argument("--OutputMomentaBase", type=str, default=None, help="Output momenta base (COCA).")
parser.add_argument("--ReadReferenceFrame", type=bool, default=False, help="Read reference frame (COCA).")
@ -126,6 +135,16 @@ def parse_arguments_card(parsed_args):
OutputFCsDensity=parsed_args.OutputFCsDensity,
OutputFCsSnapshot=parsed_args.OutputFCsSnapshot,
OutputRngStateLPT=parsed_args.OutputRngStateLPT,
## Tests with phiBCs and density
WriteGravPot=parsed_args.WriteGravPot,
OutputGravitationalPotentialBase=parsed_args.OutputGravitationalPotentialBase,
MeshGravPot=parsed_args.MeshGravPot,
WriteDensity=parsed_args.WriteDensity,
OutputDensityBase=parsed_args.OutputDensityBase,
MeshDensity=parsed_args.MeshDensity,
LoadPhiBCs=parsed_args.LoadPhiBCs,
InputPhiBCsBase=parsed_args.InputPhiBCsBase,
## Cosmological parameters
h=cosmo_dict["h"],
Omega_m=cosmo_dict["Omega_m"],
Omega_b=cosmo_dict["Omega_b"],
@ -211,6 +230,17 @@ def parse_arguments_card(parsed_args):
card_dict["OutputFCsSnapshot"] = main_dict["resultdir"]+"final_particles_"+main_dict["simname"]+".gadget3"
if card_dict["OutputRngStateLPT"] is None:
card_dict["OutputRngStateLPT"] = main_dict["workdir"]+"rng_state.h5"
## Tests with phiBCs and density
if card_dict["OutputGravitationalPotentialBase"] is None:
card_dict["OutputGravitationalPotentialBase"] = main_dict["workdir"]+"gravpot_"+main_dict["simname"]
if card_dict["MeshGravPot"] is None:
card_dict["MeshGravPot"] = card_dict["N_PM_mesh"]
if card_dict["OutputDensityBase"] is None:
card_dict["OutputDensityBase"] = main_dict["workdir"]+"density_"+main_dict["simname"]
if card_dict["MeshDensity"] is None:
card_dict["MeshDensity"] = card_dict["N_PM_mesh"]
if card_dict["InputPhiBCsBase"] is None:
card_dict["InputPhiBCsBase"] = main_dict["workdir"]+"gravpot_tCOLA"
if card_dict["OutputMomentaBase"] is None:
card_dict["OutputMomentaBase"] = main_dict["workdir"]+"momenta_"+main_dict["simname"]+"_"
if card_dict["InputMomentaBase"] is None:
@ -333,6 +363,15 @@ def create_parameter_card_dict(
OutputFCsDensity:str = 'fcs_density.h5',
OutputFCsSnapshot:str = 'fcs_particles.gadget3',
## Tests with phiBCs and density
WriteGravPot:bool = True,
OutputGravitationalPotentialBase:str = 'gravitational_potential.h5',
MeshGravPot:int = 128,
WriteDensity:bool = False,
OutputDensityBase:str = 'density.h5',
MeshDensity:int = 128,
LoadPhiBCs:bool = False,
InputPhiBCsBase:str = 'gravitational_potential.h5',
## COCA parameters
WriteReferenceFrame:bool = False,
OutputMomentaBase:str = 'momenta_',
@ -403,6 +442,17 @@ def create_parameter_card_dict(
OutputLPTPotential1=OutputLPTPotential1,
OutputLPTPotential2=OutputLPTPotential2,
OutputTilesBase=OutputTilesBase,
# Tests with phiBCs and density
WriteGravPot=int(WriteGravPot),
OutputGravitationalPotentialBase=OutputGravitationalPotentialBase,
MeshGravPot=MeshGravPot,
WriteDensity=int(WriteDensity),
OutputDensityBase=OutputDensityBase,
MeshDensity=MeshDensity,
LoadPhiBCs=int(LoadPhiBCs),
InputPhiBCsBase=InputPhiBCsBase,
h=h,
Omega_m=Omega_m,
Omega_b=Omega_b,

View file

@ -95,8 +95,6 @@ def get_config_from_dict(monofonic_dict):
"Omega_L": monofonic_dict["Omega_q"],
"H0": monofonic_dict["h"]*100.,
"n_s": monofonic_dict["n_s"],
"sigma_8": monofonic_dict["sigma8"],
"A_s": monofonic_dict["A_s"],
"Tcmb": monofonic_dict["Tcmb"],
"k_p": monofonic_dict["k_p"],
"N_ur": monofonic_dict["N_ur"],
@ -110,7 +108,14 @@ def get_config_from_dict(monofonic_dict):
"transfer": "CLASS",
"ztarget": 2.5,
"ZeroRadiation": False
}
}
if "A_s" in monofonic_dict.keys() and monofonic_dict["A_s"] is not None:
config["cosmology"]["A_s"] = monofonic_dict["A_s"]
elif "sigma8" in monofonic_dict.keys() and monofonic_dict["sigma8"] is not None:
config["cosmology"]["sigma8"] = monofonic_dict["sigma8"]
else:
raise KeyError("Either A_s or sigma8 must be provided in the cosmology parameters.")
config["random"] = {
"generator": "NGENIC",

View file

@ -34,6 +34,55 @@ def convert_snapshot_to_density(snapshot_path, output_path, N=None, corner=(0.0,
F.write(output_path)
print("Density written to", output_path)
print("Done.")
def convert_snapshots_to_density(snapshot_base_path, output_path, N=None, corner=(0.0, 0.0, 0.0)):
"""
Convert multiple snapshots to density fields.
Parameters
----------
snapshot_base_path : str
Base path for the snapshot files.
output_path : str
Path to the output density file.
N : int
Size of the density field grid (N x N x N).
corner : tuple of float
Corner of the box (x, y, z).
"""
# Get all snapshot files with path "snapshot_base_path*"
import glob
snapshot_files = glob.glob(snapshot_base_path + "*")
if not snapshot_files:
raise FileNotFoundError(f"No snapshot files found at {snapshot_base_path}")
A = None
F = None
for snapshot_path in snapshot_files:
# Read the snapshot
print(f"Reading snapshot {snapshot_path}...")
snap = read_snapshot(snapshot_path)
if N is None:
N = snap.Np0
# Calculate density
print("Calculating density...")
F = get_density_pm_snapshot(snap, N, N, N, corner[0], corner[1], corner[2])
if A is None:
A = F.data
else:
A += F.data
F.data = A # Combine densities from all snapshots
# Write density to file
print("Writing density...")
F.write(output_path)
print("Density written to", output_path)
print("Done.")
def console_main():
@ -70,13 +119,23 @@ def console_main():
args = parser.parse_args()
convert_snapshot_to_density(
snapshot_path=args.snapshot,
output_path=args.output,
N=args.N,
corner=args.corner,
)
if args.snapshot.endswith("h5") or args.snapshot.endswith("hdf5") or args.snapshot.endswith("gadget3"):
# If the snapshot is a single file, convert it to density
convert_snapshot_to_density(
snapshot_path=args.snapshot,
output_path=args.output,
N=args.N,
corner=args.corner,
)
else:
# If the snapshot is a base path, convert all snapshots to density
convert_snapshots_to_density(
snapshot_base_path=args.snapshot,
output_path=args.output,
N=args.N,
corner=args.corner,
)
if __name__ == "__main__":
console_main()
console_main()