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,