diff --git a/analysis/power_spectrum.py b/analysis/power_spectrum.py index 53fdc4f..64dacb5 100644 --- a/analysis/power_spectrum.py +++ b/analysis/power_spectrum.py @@ -5,72 +5,73 @@ kmax = 2e0 Nk = 50 AliasingCorr=False -def get_power_spectrum(field, kmin=kmin, kmax=kmax, Nk=Nk): +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 - - G = FourierGrid( - field.L0, - field.L1, - field.L2, - field.N0, - field.N1, - field.N2, - k_modes=np.concat([PowerSpectrum(field.L0,field.L1,field.L2,field.N0,field.N1,field.N2,).FourierGrid.k_modes[:10],np.logspace( - np.log10(kmin), - np.log10(kmax), - Nk, - )]), - kmax=kmax, - ) + if G is None: + G = FourierGrid( + field.L0, + field.L1, + field.L2, + field.N0, + field.N1, + field.N2, + k_modes=np.concat([PowerSpectrum(field.L0,field.L1,field.L2,field.N0,field.N1,field.N2,).FourierGrid.k_modes[:10],np.logspace( + np.log10(kmin), + np.log10(kmax), + Nk, + )]), + kmax=kmax, + ) k = G.k_modes[1:] Pk, _ = get_autocorrelation(field, G, AliasingCorr) Pk = Pk[1:] - return k, Pk + return G, k, Pk -def get_cross_correlations(field_A, field_B, kmin=kmin, kmax=kmax, Nk=Nk): +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 - G = FourierGrid( - field_A.L0, - field_A.L1, - field_A.L2, - field_A.N0, - field_A.N1, - field_A.N2, - k_modes=np.concat([PowerSpectrum(field_A.L0,field_A.L1,field_A.L2,field_A.N0,field_A.N1,field_A.N2,).FourierGrid.k_modes[:10],np.logspace( - np.log10(kmin), - np.log10(kmax), - Nk, - )]), - kmax=kmax, - ) + if G is None: + G = FourierGrid( + field_A.L0, + field_A.L1, + field_A.L2, + field_A.N0, + field_A.N1, + field_A.N2, + k_modes=np.concat([PowerSpectrum(field_A.L0,field_A.L1,field_A.L2,field_A.N0,field_A.N1,field_A.N2,).FourierGrid.k_modes[:10],np.logspace( + np.log10(kmin), + np.log10(kmax), + Nk, + )]), + kmax=kmax, + ) k = G.k_modes[1:] _, _, Rks, _ = get_crosscorrelation(field_A, field_B, G, AliasingCorr) Rks = Rks[1:] - return k, Rks + return G, k, Rks -def add_power_spectrum_to_plot(ax, field, Pk_ref=None, plot_args={}, power_args={}): - k, Pk = get_power_spectrum(field, **power_args) +def add_power_spectrum_to_plot(ax, field, Pk_ref=None, G=None, plot_args={}, power_args={}): + G, k, Pk = get_power_spectrum(field, G=G, **power_args) if Pk_ref is not None: - ax.plot(k, Pk/Pk_ref-1, **plot_args) + ax.plot(k, Pk/Pk_ref, **plot_args) else: ax.plot(k, Pk, **plot_args) - return ax + return ax, G, k, Pk -def add_cross_correlations_to_plot(ax, field_A, field_B, plot_args={}, power_args={}): - k, Rks = get_cross_correlations(field_A, field_B, **power_args) +def add_cross_correlations_to_plot(ax, field_A, field_B, G=None, plot_args={}, power_args={}): + G, k, Rks = get_cross_correlations(field_A, field_B, G=G, **power_args) ax.plot(k, Rks, **plot_args) - return ax + return ax, G, k, Rks def plot_power_spectra(filenames, @@ -79,6 +80,7 @@ def plot_power_spectra(filenames, linestyles=None, markers=None, Pk_ref=None, + G=None, ylims=[0.9,1.1], yticks = np.linspace(0.9,1.1,11), bound1=0.01, @@ -108,9 +110,10 @@ def plot_power_spectra(filenames, for i, filename in enumerate(filenames): field = read_field(filename) - add_power_spectrum_to_plot(ax=ax, + _, G, k, _ = add_power_spectrum_to_plot(ax=ax, field=field, Pk_ref=Pk_ref, + G=G, plot_args=dict(label=labels[i], color=colors[i], linestyle=linestyles[i], @@ -120,15 +123,15 @@ def plot_power_spectra(filenames, Nk=Nk), ) ax.set_xscale('log') - ax.set_xlim(kmin, kmax) + ax.set_xlim(k.min(),k[-2]) if ylims is not None: ax.set_ylim(ylims) if yticks is not None: ax.set_yticks(yticks) - ax.set_xlabel(r'$k$ [$h/\mathrm{Mpc}$]') + ax.set_xlabel(r'$k$ [$h/\mathrm{Mpc}$]', labelpad=-10) if Pk_ref is not None: - ax.set_ylabel(r'$P(k)/P_\mathrm{ref}(k)-1$') + ax.set_ylabel(r'$P(k)/P_\mathrm{ref}(k)$') else: ax.set_ylabel('$P(k)$') @@ -145,6 +148,7 @@ def plot_power_spectra(filenames, def plot_cross_correlations(filenames_A, filename_B, + G=None, labels=None, colors=None, linestyles=None, @@ -180,9 +184,10 @@ def plot_cross_correlations(filenames_A, for i, filename_A in enumerate(filenames_A): field_A = read_field(filename_A) - add_cross_correlations_to_plot(ax=ax, + _, G, k, _ = add_cross_correlations_to_plot(ax=ax, field_A=field_A, field_B=field_B, + G=G, plot_args=dict(label=labels[i], color=colors[i], linestyle=linestyles[i], @@ -192,12 +197,12 @@ def plot_cross_correlations(filenames_A, Nk=Nk), ) ax.set_xscale('log') - ax.set_xlim(kmin, kmax) + ax.set_xlim(k.min(), k[-2]) if ylims is not None: ax.set_ylim(ylims) if yticks is not None: ax.set_yticks(yticks) - ax.set_xlabel(r'$k$ [$h/\mathrm{Mpc}$]') + ax.set_xlabel(r'$k$ [$h/\mathrm{Mpc}$]', labelpad=-10) ax.set_ylabel('$R(k)$') if bound1 is not None: @@ -211,6 +216,28 @@ def plot_cross_correlations(filenames_A, return fig, ax +def get_ylims_and_yticks(ylims): + + if ylims[0] == ylims[1]: + ylims = None + yticks = None + else: + diff_ylims = ylims[1] - ylims[0] + factor = 1 + while diff_ylims<5.: + diff_ylims *= 10 + factor *= 10 + if diff_ylims<12.: + yticks = np.linspace(int(ylims[0]*factor)/factor,int(factor*ylims[1])/factor, int(diff_ylims)+1) + else: + while diff_ylims>12.: + diff_ylims /= 2. + factor /= 2. + yticks = np.linspace(int(ylims[0]*factor)/factor,int(factor*ylims[1])/factor, int(diff_ylims)+1) + + return ylims, yticks + + if __name__ == "__main__": from argparse import ArgumentParser parser = ArgumentParser(description='Plot power spectra of fields') @@ -226,6 +253,8 @@ if __name__ == "__main__": parser.add_argument('-ls', '--linestyles', type=str, nargs='+', default=None, help='Linestyles for each field.') parser.add_argument('-m', '--markers', type=str, nargs='+', default=None, help='Markers for each field.') parser.add_argument('-t','--title', type=str, default=None, help='Title of the plot.') + 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.') args = parser.parse_args() @@ -235,12 +264,18 @@ if __name__ == "__main__": if args.reference is not None: from pysbmy.field import read_field - Pk_ref = get_power_spectrum(read_field(args.directory+args.reference), kmin=kmin, kmax=kmax, Nk=Nk)[1] + G, _, Pk_ref = get_power_spectrum(read_field(args.directory+args.reference), kmin=kmin, kmax=kmax, Nk=Nk) else: Pk_ref = None + G = None + filenames = [args.directory+f for f in args.filenames] + ylims_power, yticks_power = get_ylims_and_yticks(args.ylim_power) + ylims_corr, yticks_corr = get_ylims_and_yticks(args.ylim_corr) + + if args.power_spectrum and args.cross_correlation: import matplotlib.pyplot as plt fig, axes = plt.subplots(2, 1, figsize=(8,8)) @@ -250,14 +285,11 @@ if __name__ == "__main__": linestyles=args.linestyles, markers=args.markers, Pk_ref=Pk_ref, - # ylims=[0.9,1.1], - # yticks = np.linspace(0.9,1.1,11), - # bound1=0.01, - # bound2=0.02, - ylims=None, - yticks = None, - bound1=None, - bound2=None, + G=G, + ylims=ylims_power, + yticks = yticks_power, + bound1=0.01, + bound2=0.02, kmin=kmin, kmax=kmax, Nk=Nk, @@ -266,18 +298,15 @@ if __name__ == "__main__": plot_cross_correlations(filenames_A=filenames, filename_B=args.directory+args.reference, + G=G, labels=args.labels, colors=args.colors, linestyles=args.linestyles, markers=args.markers, - # ylims=[0.99, 1.001], - # yticks = np.linspace(0.99,1.001,12), - # bound1=0.001, - # bound2=0.002, - ylims=None, - yticks = None, - bound1=None, - bound2=None, + ylims=ylims_corr, + yticks = yticks_corr, + bound1=0.001, + bound2=0.002, kmin=kmin, kmax=kmax, Nk=Nk, @@ -298,8 +327,9 @@ if __name__ == "__main__": linestyles=args.linestyles, markers=args.markers, Pk_ref=Pk_ref, - ylims=[0.9,1.1], - yticks = np.linspace(0.9,1.1,11), + G=G, + ylims=ylims_power, + yticks = yticks_power, bound1=0.01, bound2=0.02, kmin=kmin, @@ -312,12 +342,13 @@ if __name__ == "__main__": elif args.cross_correlation: fig, ax = plot_cross_correlations(filenames_A=filenames, filename_B=args.reference, + G=G, labels=args.labels, colors=args.colors, linestyles=args.linestyles, markers=args.markers, - ylims=[0.99, 1.001], - yticks = np.linspace(0.99,1.001,12), + ylims=ylims_corr, + yticks = yticks_corr, bound1=0.001, bound2=0.002, kmin=kmin, diff --git a/analysis/slices.py b/analysis/slices.py index 0456e7b..8b260cc 100644 --- a/analysis/slices.py +++ b/analysis/slices.py @@ -1,13 +1,18 @@ import numpy as np +import sys +sys.path.append('/home/aubin/Simbelmyne/sbmy_control/') +from cosmo_params import register_arguments_cosmo, parse_arguments_cosmo + fs = 18 fs_titles = fs -4 def plot_imshow_with_reference( data_list, - reference, - titles, + reference=None, + titles=None, vmin=None, - vmax=None, + vmax=None, + L=None, cmap='viridis'): """ Plot the imshow of a list of 2D arrays with two rows: one for the data itself, @@ -23,38 +28,70 @@ def plot_imshow_with_reference( data_list, if titles is None: titles = [None for f in data_list] + + if L is None: + L = [len(data) for data in data_list] + elif isinstance(L, int) or isinstance(L, float): + L = [L for data in data_list] + + sep = 10 if L[0] < 50 else 20 if L[0] < 200 else 100 + ticks = [np.arange(0, l+1, sep)*len(dat)/l for l, dat in zip(L,data_list)] + tick_labels = [np.arange(0, l+1, sep) for l in L] def score(data, reference): return np.linalg.norm(data-reference)/np.linalg.norm(reference) n = len(data_list) - fig, axes = plt.subplots(2, n, figsize=(5 * n, 10)) + fig, axes = plt.subplots(1 if reference is None else 2, n, figsize=(5 * n, 5 if reference is None else 5*2), dpi=max(500, data_list[0].shape[0]//2)) if vmin is None or vmax is None: vmin = min(np.quantile(data,0.01) for data in data_list) vmax = max(np.quantile(data,0.99) for data in data_list) + + if reference is not None: vmin_diff = min(np.quantile((data-reference),0.01) for data in data_list) vmax_diff = max(np.quantile((data-reference),0.99) for data in data_list) else: vmin_diff = vmin vmax_diff = vmax - # Plot the data itself - for i, data in enumerate(data_list): - im = axes[0, i].imshow(data, cmap=cmap, origin='lower', vmin=vmin, vmax=vmax) - axes[0, i].set_title(titles[i], fontsize=fs_titles) - fig.colorbar(im, ax=axes[0, :], orientation='vertical') + if reference is not None: + # Plot the data itself + for i, data in enumerate(data_list): + im = axes[0, i].imshow(data, cmap=cmap, origin='lower', vmin=vmin, vmax=vmax) + axes[0, i].set_title(titles[i], fontsize=fs_titles) + axes[0, i].set_xticks(ticks[i]) + axes[0, i].set_yticks(ticks[i]) + axes[0, i].set_xticklabels(tick_labels[i]) + axes[0, i].set_yticklabels(tick_labels[i]) + axes[0, i].set_xlabel('Mpc/h') + fig.colorbar(im, ax=axes[0, :], orientation='vertical') - # Plot the data compared to the reference - for i, data in enumerate(data_list): - im = axes[1, i].imshow(data - reference, cmap=cmap, origin='lower', vmin=vmin_diff, vmax=vmax_diff) - axes[1, i].set_title(f'{titles[i]} - Reference', fontsize=fs_titles) - fig.colorbar(im, ax=axes[1, :], orientation='vertical') - - # Add the score on the plots - for i, data in enumerate(data_list): - axes[1, i].text(0.5, 0.9, f"Score: {score(data, reference):.2e}", fontsize=10, transform=axes[1, i].transAxes, color='white') - # plt.tight_layout() + # Plot the data compared to the reference + for i, data in enumerate(data_list): + im = axes[1, i].imshow(data - reference, cmap=cmap, origin='lower', vmin=vmin_diff, vmax=vmax_diff) + axes[1, i].set_title(f'{titles[i]} - Reference', fontsize=fs_titles) + axes[1, i].set_xticks(ticks[i]) + axes[1, i].set_yticks(ticks[i]) + axes[1, i].set_xticklabels(tick_labels[i]) + axes[1, i].set_yticklabels(tick_labels[i]) + axes[1, i].set_xlabel('Mpc/h') + fig.colorbar(im, ax=axes[1, :], orientation='vertical') + + # Add the score on the plots + for i, data in enumerate(data_list): + axes[1, i].text(0.5, 0.9, f"Score: {score(data, reference):.2e}", fontsize=10, transform=axes[1, i].transAxes, color='white') + # plt.tight_layout() + else: + for i, data in enumerate(data_list): + im = axes[i].imshow(data, cmap=cmap, origin='lower', vmin=vmin, vmax=vmax) + axes[i].set_title(titles[i], fontsize=fs_titles) + axes[i].set_xticks(ticks[i]) + axes[i].set_yticks(ticks[i]) + axes[i].set_xticklabels(tick_labels[i]) + axes[i].set_yticklabels(tick_labels[i]) + axes[i].set_xlabel('Mpc/h') + fig.colorbar(im, ax=axes[:], orientation='vertical') return fig, axes @@ -76,36 +113,48 @@ if __name__ == "__main__": parser.add_argument('-t', '--title', type=str, default=None, help='Title for the plot.') parser.add_argument('-log','--log_scale', action='store_true', help='Use log scale for the data.') + register_arguments_cosmo(parser) + args = parser.parse_args() from pysbmy.field import read_field + from pysbmy.cosmology import d_plus - ref_field = read_field(args.directory+args.reference) + ref_field = read_field(args.directory+args.reference) if args.reference is not None else None fields = [read_field(args.directory+f) for f in args.filenames] if args.index is None: - index = ref_field.N0//2 + index = fields[0].N0//2 else: index=args.index + + # args.labels=[f"a={f.time:.2f}" for f in fields] + L = [f.L0 for f in fields] match args.axis: case 0 | 'x': - reference = ref_field.data[index,:,:] + reference = ref_field.data[index,:,:] if ref_field is not None else None fields = [f.data[index,:,:] for f in fields] + # reference = ref_field.data[index,:,:]/d_plus(1e-3,ref_field.time,parse_arguments_cosmo(args)) + # fields = [f.data[index,:,:]/d_plus(1e-3,f.time,parse_arguments_cosmo(args)) for f in fields] + # reference = ref_field.data[index,:,:]/d_plus(1e-3,0.05,parse_arguments_cosmo(args)) + # fields = [f.data[index,:,:]/d_plus(1e-3,time,parse_arguments_cosmo(args)) for f,time in zip(fields,[0.05, 1.0])] case 1 | 'y': - reference = ref_field.data[:,index,:] + reference = ref_field.data[:,index,:] if ref_field is not None else None fields = [f.data[:,index,:] for f in fields] case 2 | 'z': - reference = ref_field.data[:,:,index] + reference = ref_field.data[:,:,index] if ref_field is not None else None fields = [f.data[:,:,index] for f in fields] case _: raise ValueError(f"Wrong axis provided : {args.axis}") if args.log_scale: - reference = np.log10(2.+reference) + reference = np.log10(2.+reference) if ref_field is not None else None fields = [np.log10(2.+f) for f in fields] - fig, axes = plot_imshow_with_reference(fields,reference,args.labels, vmin=args.vmin, vmax=args.vmax,cmap=args.cmap) + + + fig, axes = plot_imshow_with_reference(fields,reference,args.labels, vmin=args.vmin, vmax=args.vmax,cmap=args.cmap, L=L) fig.suptitle(args.title) if args.output is not None: diff --git a/cosmo_params.py b/cosmo_params.py index 8ca1636..14e48dc 100644 --- a/cosmo_params.py +++ b/cosmo_params.py @@ -63,4 +63,32 @@ def parse_arguments_cosmo(parsed_args): return cosmo_dict def z2a(z): - return 1.0/(1.0+z) \ No newline at end of file + return 1.0/(1.0+z) + + +cosmo_defaults = { + "h": 0.6732, + "Omega_m": 0.302, + "Omega_b": 0.049, + "Omega_q": 0.6842, + "Omega_k":0.0, + "Omega_r": 0.0, + "n_s": 0.968, + "sigma8": 0.815, + "A_s": 2.148752e-09, + "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": 100.0, + "gnl": 0.0, + "k_max":10.0, + "tau_reio":0.06, + "WhichSpectrum":"EH", + "w0_fld":-1.0, + "wa_fld":0.0, + } \ No newline at end of file diff --git a/parameters_card.py b/parameters_card.py index fe1768c..082c297 100644 --- a/parameters_card.py +++ b/parameters_card.py @@ -207,9 +207,9 @@ def parse_arguments_card(parsed_args): if card_dict["OutputSnapshotsBase"] is None: card_dict["OutputSnapshotsBase"] = main_dict["resultdir"]+"particles_"+main_dict["simname"] if card_dict["OutputFinalSnapshot"] is None: - card_dict["OutputFinalSnapshot"] = main_dict["resultdir"]+ligthcone_prefix+"final_particles_"+main_dict["simname"]+".gadget3" + card_dict["OutputFinalSnapshot"] = main_dict["resultdir"]+ligthcone_prefix+"final_particles_"+main_dict["simname"]+("_lc" if card_dict["GenerateLightcone"] else "")+".gadget3" if card_dict["OutputFinalDensity"] is None: - card_dict["OutputFinalDensity"] = main_dict["resultdir"]+ligthcone_prefix+"final_density_"+main_dict["simname"]+".h5" + card_dict["OutputFinalDensity"] = main_dict["resultdir"]+ligthcone_prefix+"final_density_"+main_dict["simname"]+("_lc" if card_dict["GenerateLightcone"] else "")+".h5" if card_dict["OutputTilesBase"] is None: card_dict["OutputTilesBase"] = main_dict["workdir"]+main_dict["simname"]+"_tile" if card_dict["OutputLPTPotential1"] is None: diff --git a/scola.py b/scola.py index a6156ad..ff341ab 100644 --- a/scola.py +++ b/scola.py @@ -13,6 +13,11 @@ def main_scola(parsed_args): nboxes_tot = int(parsed_args.N_tiles**3) + if have_all_tiles(parsed_args) and not parsed_args.force: + print_message("All tiles already exist. Use -F to overwrite.", 1, "scola", verbose=parsed_args.verbose) + print_ending_module("scola", verbose=parsed_args.verbose) + return + if parsed_args.execution == "local": from parameters_card import parse_arguments_card from tqdm import tqdm