import numpy as np kmin = 1e-1 kmax = 2e0 Nk = 50 AliasingCorr=False 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: 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 G, k, Pk 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: 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 G, k, Rks 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, **plot_args) else: ax.plot(k, Pk, **plot_args) return ax, G, k, Pk 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, G, k, Rks def plot_power_spectra(filenames, labels=None, colors=None, 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, bound2=0.02, kmin=kmin, kmax=kmax, Nk=Nk, figsize=(8,4), dpi=300, ax=None, fig=None,): import matplotlib.pyplot as plt from pysbmy.field import read_field if fig is None or ax is None: fig, ax = plt.subplots(figsize=figsize, dpi=dpi) if labels is None: labels = [None for f in filenames] if colors is None: colors = [None for f in filenames] if linestyles is None: linestyles = [None for f in filenames] if markers is None: markers = [None for f in filenames] for i, filename in enumerate(filenames): field = read_field(filename) _, 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], marker=markers[i],), power_args=dict(kmin=kmin, kmax=kmax, Nk=Nk), ) ax.set_xscale('log') 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}$]', labelpad=-10) if Pk_ref is not None: ax.set_ylabel(r'$P(k)/P_\mathrm{ref}(k)$') else: ax.set_ylabel('$P(k)$') if bound1 is not None: ax.axhspan(1 - bound1, 1 + bound1, color="grey", alpha=0.2) if bound2 is not None: ax.axhspan(1 - bound2, 1 + bound2, color="grey", alpha=0.1) ax.grid(which="major",alpha=0.5) ax.grid(which="minor",alpha=0.2) return fig, ax def plot_cross_correlations(filenames_A, filename_B, G=None, labels=None, colors=None, linestyles=None, markers=None, ylims=[0.99, 1.001], yticks = np.linspace(0.99,1.001,12), bound1=0.001, bound2=0.002, kmin=kmin, kmax=kmax, Nk=Nk, figsize=(8,4), dpi=300, ax=None, fig=None,): import matplotlib.pyplot as plt from pysbmy.field import read_field if fig is None or ax is None: fig, ax = plt.subplots(figsize=figsize, dpi=dpi) if labels is None: labels = [None for f in filenames_A] if colors is None: colors = [None for f in filenames_A] if linestyles is None: linestyles = [None for f in filenames_A] if markers is None: markers = [None for f in filenames_A] field_B = read_field(filename_B) for i, filename_A in enumerate(filenames_A): field_A = read_field(filename_A) _, 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], marker=markers[i],), power_args=dict(kmin=kmin, kmax=kmax, Nk=Nk), ) ax.set_xscale('log') 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}$]', labelpad=-10) ax.set_ylabel('$R(k)$') if bound1 is not None: ax.axhspan(1 - bound1, 1 + bound1, color="grey", alpha=0.2) if bound2 is not None: ax.axhspan(1 - bound2, 1 + bound2, color="grey", alpha=0.1) ax.grid(which="major",alpha=0.5) ax.grid(which="minor",alpha=0.2) 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') parser.add_argument('-ps', '--power_spectrum', action='store_true', help='Plot power spectra.') parser.add_argument('-cc', '--cross_correlation', action='store_true', help='Plot cross correlations.') parser.add_argument('-d', '--directory', type=str, required=True, help='Directory containing the fields files.') parser.add_argument('-ref', '--reference', type=str, default=None, help='Reference field file.') parser.add_argument('-f', '--filenames', type=str, nargs='+', required=True, help='Field files to be plotted.') parser.add_argument('-o', '--output', type=str, default=None, help='Output plot file name.') parser.add_argument('-l', '--labels', type=str, nargs='+', default=None, help='Labels for each field.') parser.add_argument('-c', '--colors', type=str, nargs='+', default=None, help='Colors for each field.') 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() if not args.power_spectrum and not args.cross_correlation: print('You must choose between power_spectrum and cross_correlation.') exit() if args.reference is not None: from pysbmy.field import read_field 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)) plot_power_spectra(filenames=filenames, labels=args.labels, colors=args.colors, linestyles=args.linestyles, markers=args.markers, Pk_ref=Pk_ref, G=G, ylims=ylims_power, yticks = yticks_power, bound1=0.01, bound2=0.02, kmin=kmin, kmax=kmax, Nk=Nk, ax=axes[0], fig=fig) 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=ylims_corr, yticks = yticks_corr, bound1=0.001, bound2=0.002, kmin=kmin, kmax=kmax, Nk=Nk, ax=axes[1], fig=fig) axes[1].legend(loc='lower left') axes[0].set_title("Power Spectrum") axes[1].set_title("Cross Correlations") if args.title is not None: fig.suptitle(args.title) elif args.power_spectrum: fig, ax = plot_power_spectra(filenames=filenames, labels=args.labels, colors=args.colors, linestyles=args.linestyles, markers=args.markers, Pk_ref=Pk_ref, G=G, ylims=ylims_power, yticks = yticks_power, bound1=0.01, bound2=0.02, kmin=kmin, kmax=kmax, Nk=Nk) ax.legend() if args.title is not None: ax.set_title(args.title) 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=ylims_corr, yticks = yticks_corr, bound1=0.001, bound2=0.002, kmin=kmin, kmax=kmax, Nk=Nk) ax.legend(loc='lower left') if args.title is not None: ax.set_title(args.title) if args.output is not None: fig.savefig(args.output) else: fig.savefig(args.directory+'power_spectrum.png')