diff --git a/sbmy_control/analysis/power_spectrum.py b/sbmy_control/analysis/power_spectrum.py index d97e24f..7a92370 100644 --- a/sbmy_control/analysis/power_spectrum.py +++ b/sbmy_control/analysis/power_spectrum.py @@ -1,7 +1,7 @@ import numpy as np import os -nfirst_kmodes = 10 +kmin = 1e-1 kmax = 2e0 Nk = 50 AliasingCorr=False @@ -24,17 +24,12 @@ def crop_field(field, Ncrop): field.L2 = field.N2*d2 -def get_power_spectrum(field, nfirst_kmodes=nfirst_kmodes, kmax=kmax, Nk=Nk, G=None): +def get_power_spectrum(field, kmin=kmin, kmax=kmax, Nk=Nk, G=None): from pysbmy.power import PowerSpectrum from pysbmy.fft import FourierGrid from pysbmy.correlations import get_autocorrelation if G is None: - default_k_modes = PowerSpectrum(field.L0,field.L1,field.L2,int(field.N0),int(field.N1),int(field.N2),).FourierGrid.k_modes - if kmax is None: - kmax = default_k_modes[-1] - if kmax > default_k_modes[-1]: - raise ValueError(f"kmax ({kmax}) is larger than the maximum k mode available in the field ({default_k_modes[-1]}).") G = FourierGrid( field.L0, field.L1, @@ -42,8 +37,8 @@ def get_power_spectrum(field, nfirst_kmodes=nfirst_kmodes, kmax=kmax, Nk=Nk, G=N int(field.N0), int(field.N1), int(field.N2), - k_modes=np.concat([default_k_modes[:nfirst_kmodes],np.logspace( - np.log10(default_k_modes[nfirst_kmodes]), + k_modes=np.concat([PowerSpectrum(field.L0,field.L1,field.L2,int(field.N0),int(field.N1),int(field.N2),).FourierGrid.k_modes[:10],np.logspace( + np.log10(kmin), np.log10(kmax), Nk, )]), @@ -56,18 +51,13 @@ def get_power_spectrum(field, nfirst_kmodes=nfirst_kmodes, kmax=kmax, Nk=Nk, G=N return G, k, Pk -def get_cross_correlations(field_A, field_B, nfirst_kmodes=nfirst_kmodes, kmax=kmax, Nk=Nk, G=None): +def get_cross_correlations(field_A, field_B, kmin=kmin, kmax=kmax, Nk=Nk, G=None): from pysbmy.power import PowerSpectrum from pysbmy.fft import FourierGrid from pysbmy.correlations import get_crosscorrelation if G is None: - default_k_modes = PowerSpectrum(field_A.L0,field_A.L1,field_A.L2,int(field_A.N0),int(field_A.N1),int(field_A.N2),).FourierGrid.k_modes - if kmax is None: - kmax = default_k_modes[-1] - if kmax > default_k_modes[-1]: - raise ValueError(f"kmax ({kmax}) is larger than the maximum k mode available in the field ({default_k_modes[-1]}).") G = FourierGrid( field_A.L0, field_A.L1, @@ -75,8 +65,8 @@ def get_cross_correlations(field_A, field_B, nfirst_kmodes=nfirst_kmodes, kmax=k int(field_A.N0), int(field_A.N1), int(field_A.N2), - k_modes=np.concat([default_k_modes[:nfirst_kmodes],np.logspace( - np.log10(default_k_modes[nfirst_kmodes]), + k_modes=np.concat([PowerSpectrum(field_A.L0,field_A.L1,field_A.L2,int(field_A.N0),int(field_A.N1),int(field_A.N2),).FourierGrid.k_modes[:10],np.logspace( + np.log10(kmin), np.log10(kmax), Nk, )]), @@ -114,7 +104,7 @@ def plot_power_spectra(filenames, yticks = np.linspace(0.9,1.1,11), bound1=0.01, bound2=0.02, - nfirst_kmodes=nfirst_kmodes, + kmin=kmin, kmax=kmax, Nk=Nk, figsize=(8,4), @@ -149,7 +139,7 @@ def plot_power_spectra(filenames, color=colors[i], linestyle=linestyles[i], marker=markers[i],), - power_args=dict(nfirst_kmodes=nfirst_kmodes, + power_args=dict(kmin=kmin, kmax=kmax, Nk=Nk), ) @@ -188,7 +178,7 @@ def plot_cross_correlations(filenames_A, yticks = np.linspace(0.99,1.001,12), bound1=0.001, bound2=0.002, - nfirst_kmodes=nfirst_kmodes, + kmin=kmin, kmax=kmax, Nk=Nk, figsize=(8,4), @@ -227,7 +217,7 @@ def plot_cross_correlations(filenames_A, color=colors[i], linestyle=linestyles[i], marker=markers[i],), - power_args=dict(nfirst_kmodes=nfirst_kmodes, + power_args=dict(kmin=kmin, kmax=kmax, Nk=Nk), ) @@ -293,15 +283,8 @@ def console_main(): parser.add_argument('-yrp', '--ylim_power', type=float, nargs=2, default=[0.9,1.1], help='Y-axis limits.') parser.add_argument('-yrc', '--ylim_corr', type=float, nargs=2, default=[0.99,1.001], help='Y-axis limits.') parser.add_argument('--crop', type=int, default=None, help='Remove the outter N pixels of the fields.') - parser.add_argument('--nfirst_kmodes', type=float, default=10, help='First k modes from all available k modes to be used in the power spectrum before the geomspace part.') - parser.add_argument('--kmax', type=float, default=None, help='Maximum k value for the power spectrum geomspace part. If None, it will be set to the maximum k mode available in the field.') - parser.add_argument('--Nk', type=int, default=50, help='Number of k values for the power spectrum geomspace part.') args = parser.parse_args() - - nfirst_kmodes = args.nfirst_kmodes - kmax = args.kmax - Nk = args.Nk if not args.power_spectrum and not args.cross_correlation: print('You must choose between power_spectrum and cross_correlation.') @@ -315,7 +298,7 @@ def console_main(): from pysbmy.field import read_field F_ref = read_field(args.directory+args.reference) crop_field(F_ref, args.crop) - G, _, Pk_ref = get_power_spectrum(F_ref, nfirst_kmodes=nfirst_kmodes, kmax=kmax, Nk=Nk) + G, _, Pk_ref = get_power_spectrum(F_ref, kmin=kmin, kmax=kmax, Nk=Nk) else: Pk_ref = None G = None @@ -342,7 +325,7 @@ def console_main(): yticks = yticks_power, bound1=0.01, bound2=0.02, - nfirst_kmodes=nfirst_kmodes, + kmin=kmin, kmax=kmax, Nk=Nk, ax=axes[0], @@ -361,7 +344,7 @@ def console_main(): yticks = yticks_corr, bound1=0.001, bound2=0.002, - nfirst_kmodes=nfirst_kmodes, + kmin=kmin, kmax=kmax, Nk=Nk, ax=axes[1], @@ -388,7 +371,7 @@ def console_main(): yticks = yticks_power, bound1=0.01, bound2=0.02, - nfirst_kmodes=nfirst_kmodes, + kmin=kmin, kmax=kmax, Nk=Nk, Ncrop=args.crop, @@ -409,7 +392,7 @@ def console_main(): yticks = yticks_corr, bound1=0.001, bound2=0.002, - nfirst_kmodes=nfirst_kmodes, + kmin=kmin, kmax=kmax, Nk=Nk, Ncrop=args.crop, diff --git a/sbmy_control/low_level.py b/sbmy_control/low_level.py index 26054bf..5a32cdd 100644 --- a/sbmy_control/low_level.py +++ b/sbmy_control/low_level.py @@ -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) diff --git a/sbmy_control/monofonic.py b/sbmy_control/monofonic.py index 861f14b..134b55b 100644 --- a/sbmy_control/monofonic.py +++ b/sbmy_control/monofonic.py @@ -32,7 +32,7 @@ def main_monofonic(parsed_args): print_message("Running monofonic in slurm mode.", 1, "monofonic", verbose=parsed_args.verbose) slurm_dict=parse_arguments_slurm(parsed_args) main_dict=parse_arguments_main(parsed_args) - slurm_script = slurm_dict["scripts"]+f"monofonic_{main_dict['simname']}.sh" + slurm_script = slurm_dict["scripts"]+"monofonic.sh" if not isfile(slurm_script) or parsed_args.force: print_message(f"SLURM script {slurm_script} does not exist (or forced). Creating it.", 2, "monofonic", verbose=parsed_args.verbose) diff --git a/sbmy_control/parameters_card.py b/sbmy_control/parameters_card.py index bbdc1a6..5384adb 100644 --- a/sbmy_control/parameters_card.py +++ b/sbmy_control/parameters_card.py @@ -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, diff --git a/sbmy_control/scripts/convert_snapshot_to_density.py b/sbmy_control/scripts/convert_snapshot_to_density.py index f3935d9..869ee90 100644 --- a/sbmy_control/scripts/convert_snapshot_to_density.py +++ b/sbmy_control/scripts/convert_snapshot_to_density.py @@ -1,7 +1,9 @@ +from pysbmy.density import get_density_pm_snapshot +from pysbmy.snapshot import read_snapshot import argparse -def convert_snapshot_to_density(snapshot_path, output_path, N=None, corner=(0.0, 0.0, 0.0), time=1.0): +def convert_snapshot_to_density(snapshot_path, output_path, N=None, corner=(0.0, 0.0, 0.0)): """ Convert a snapshot to a density field. @@ -16,9 +18,6 @@ def convert_snapshot_to_density(snapshot_path, output_path, N=None, corner=(0.0, corner : tuple of float Corner of the box (x, y, z). """ - from pysbmy.density import get_density_pm_snapshot - from pysbmy.snapshot import read_snapshot - # Read the snapshot print("Reading snapshot...") snap = read_snapshot(snapshot_path) @@ -29,7 +28,6 @@ def convert_snapshot_to_density(snapshot_path, output_path, N=None, corner=(0.0, # Calculate density print("Calculating density...") F = get_density_pm_snapshot(snap, N, N, N, corner[0], corner[1], corner[2]) - F.time = time # Set the time for the field # Write density to file print("Writing density...") @@ -38,7 +36,7 @@ def convert_snapshot_to_density(snapshot_path, output_path, N=None, corner=(0.0, print("Done.") -def convert_snapshots_to_density(snapshot_base_path, output_path, N=None, corner=(0.0, 0.0, 0.0), time=1.0): +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. @@ -53,57 +51,33 @@ def convert_snapshots_to_density(snapshot_base_path, output_path, N=None, corner corner : tuple of float Corner of the box (x, y, z). """ - from pysbmy.density import get_density_pm, density_to_delta - from pysbmy.snapshot import read_snapshot - from pysbmy.field import Field - import numpy as np - # 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 - L0 = L1 = L2 = None - d0 = d1 = d2 = 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 - - if L0 is None: - # Get the grid parameters from the snapshot - L0 = snap.L0 - L1 = snap.L1 - L2 = snap.L2 - d0 = L0 / N - d1 = L1 / N - d2 = L2 / N - else: - # Ensure the grid parameters match the first snapshot - if not (snap.L0 == L0 and snap.L1 == L1 and snap.L2 == L2): - raise ValueError(f"All snapshots must have the same grid parameters. Got {snap.L0}, {snap.L1}, {snap.L2} but expected {L0}, {L1}, {L2}.") - - if A is None: - # Initialize the density field - A = np.zeros((N, N, N), dtype=np.float32) # Calculate density - print("Calculating density for this snapshot...") - A = get_density_pm(snap.pos, A, d0=d0, d1=d1, d2=d2) + 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 - # Convert density to delta field - print("Converting density to delta field...") - A = density_to_delta(A) - - # Create a Field object - F = Field(L0=L0,L1=L1,L2=L2,corner0=corner[0],corner1=corner[1],corner2=corner[2],rank=1,N0=N,N1=N,N2=N,time=time,data=A) + F.data = A # Combine densities from all snapshots # Write density to file print("Writing density...") F.write(output_path) diff --git a/sbmy_control/simbelmyne.py b/sbmy_control/simbelmyne.py index ac2a6ba..a106ae9 100644 --- a/sbmy_control/simbelmyne.py +++ b/sbmy_control/simbelmyne.py @@ -16,7 +16,7 @@ def main_simbelmyne(parsed_args): if isfile(log_file): # Remove the preexisting log file to allow for the progress_bar to be run normally from os import remove - remove(log_file) + oremove(log_file) command_args = ["simbelmyne", paramfile, log_file]