added analysis

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')
+    
+

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')