From 9d6f8f9de20295ac079ee239dd4b8702222d5573 Mon Sep 17 00:00:00 2001 From: Mayeul Aubin Date: Mon, 17 Mar 2025 11:37:30 +0100 Subject: [PATCH] added analysis --- analysis/__init__.py | 0 analysis/power_spectrum.py | 335 +++++++++++++++++++++++++++++++++++++ analysis/slices.py | 114 +++++++++++++ 3 files changed, 449 insertions(+) create mode 100644 analysis/__init__.py create mode 100644 analysis/power_spectrum.py create mode 100644 analysis/slices.py diff --git a/analysis/__init__.py b/analysis/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/analysis/power_spectrum.py b/analysis/power_spectrum.py new file mode 100644 index 0000000..53fdc4f --- /dev/null +++ b/analysis/power_spectrum.py @@ -0,0 +1,335 @@ +import numpy as np + +kmin = 1e-1 +kmax = 2e0 +Nk = 50 +AliasingCorr=False + +def get_power_spectrum(field, kmin=kmin, kmax=kmax, Nk=Nk): + 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, + ) + k = G.k_modes[1:] + Pk, _ = get_autocorrelation(field, G, AliasingCorr) + Pk = Pk[1:] + + return k, Pk + + +def get_cross_correlations(field_A, field_B, kmin=kmin, kmax=kmax, Nk=Nk): + 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, + ) + k = G.k_modes[1:] + _, _, Rks, _ = get_crosscorrelation(field_A, field_B, G, AliasingCorr) + Rks = Rks[1:] + + return 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) + if Pk_ref is not None: + ax.plot(k, Pk/Pk_ref-1, **plot_args) + else: + ax.plot(k, Pk, **plot_args) + return ax + +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) + ax.plot(k, Rks, **plot_args) + return ax + + +def plot_power_spectra(filenames, + labels=None, + colors=None, + linestyles=None, + markers=None, + Pk_ref=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) + add_power_spectrum_to_plot(ax=ax, + field=field, + Pk_ref=Pk_ref, + 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(kmin, kmax) + 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}$]') + + if Pk_ref is not None: + ax.set_ylabel(r'$P(k)/P_\mathrm{ref}(k)-1$') + 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, + 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) + add_cross_correlations_to_plot(ax=ax, + field_A=field_A, + field_B=field_B, + 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(kmin, kmax) + 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_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 + + +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.') + + 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 + Pk_ref = get_power_spectrum(read_field(args.directory+args.reference), kmin=kmin, kmax=kmax, Nk=Nk)[1] + else: + Pk_ref = None + + filenames = [args.directory+f for f in args.filenames] + + 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, + # 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, + kmin=kmin, + kmax=kmax, + Nk=Nk, + ax=axes[0], + fig=fig) + + plot_cross_correlations(filenames_A=filenames, + filename_B=args.directory+args.reference, + 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, + 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, + 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) + 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, + 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, + 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') + + diff --git a/analysis/slices.py b/analysis/slices.py new file mode 100644 index 0000000..0456e7b --- /dev/null +++ b/analysis/slices.py @@ -0,0 +1,114 @@ +import numpy as np + +fs = 18 +fs_titles = fs -4 + +def plot_imshow_with_reference( data_list, + reference, + titles, + vmin=None, + vmax=None, + cmap='viridis'): + """ + Plot the imshow of a list of 2D arrays with two rows: one for the data itself, + one for the data compared to a reference. Each row will have a common colorbar. + + Parameters: + - data_list: list of 2D arrays to be plotted + - reference: 2D array to be used as reference for comparison + - titles: list of titles for each subplot + - cmap: colormap to be used for plotting + """ + import matplotlib.pyplot as plt + + if titles is None: + titles = [None for f in data_list] + + 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)) + + 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) + 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') + + # 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() + + return fig, axes + + +if __name__ == "__main__": + from argparse import ArgumentParser + parser = ArgumentParser(description='Comparisons of fields slices.') + + parser.add_argument('-a','--axis', type=int, default=0, help='Axis along which the slices will be taken.') + parser.add_argument('-i','--index', type=int, default=None, help='Index of the slice along the axis.') + 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', '--cmap', type=str, default='viridis', help='Colormap to be used for plotting.') + parser.add_argument('-vmin', type=float, default=None, help='Minimum value for the colorbar.') + parser.add_argument('-vmax', type=float, default=None, help='Maximum value for the colorbar.') + 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.') + + args = parser.parse_args() + + from pysbmy.field import read_field + + ref_field = read_field(args.directory+args.reference) + fields = [read_field(args.directory+f) for f in args.filenames] + + if args.index is None: + index = ref_field.N0//2 + else: + index=args.index + + match args.axis: + case 0 | 'x': + reference = ref_field.data[index,:,:] + fields = [f.data[index,:,:] for f in fields] + case 1 | 'y': + reference = ref_field.data[:,index,:] + fields = [f.data[:,index,:] for f in fields] + case 2 | 'z': + reference = ref_field.data[:,:,index] + 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) + 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.suptitle(args.title) + + if args.output is not None: + fig.savefig(args.output) + else: + fig.savefig(args.directory+'slices.png') \ No newline at end of file