Merge branch 'main' into tests/GravPotOutput

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
 
-
+    if G is None:
-    G = FourierGrid(
+        G = FourierGrid(
-        field.L0,
+            field.L0,
-        field.L1,
+            field.L1,
-        field.L2,
+            field.L2,
-        field.N0,
+            field.N0,
-        field.N1,
+            field.N1,
-        field.N2,
+            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(
+            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(kmin),
-            np.log10(kmax),
+                np.log10(kmax),
-            Nk,
+                Nk,
-        )]),
+            )]),
-        kmax=kmax,
+            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(
+    if G is None:
-        field_A.L0,
+        G = FourierGrid(
-        field_A.L1,
+            field_A.L0,
-        field_A.L2,
+            field_A.L1,
-        field_A.N0,
+            field_A.L2,
-        field_A.N1,
+            field_A.N0,
-        field_A.N2,
+            field_A.N1,
-        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(
+            field_A.N2,
-            np.log10(kmin),
+            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(kmax),
+                np.log10(kmin),
-            Nk,
+                np.log10(kmax),
-        )]),
+                Nk,
-        kmax=kmax,
+            )]),
-    )
+            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={}):
+def add_power_spectrum_to_plot(ax, field, Pk_ref=None, G=None, plot_args={}, power_args={}):
-    k, Pk = get_power_spectrum(field, **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={}):
+def add_cross_correlations_to_plot(ax, field_A, field_B, G=None, plot_args={}, power_args={}):
-    k, Rks = get_cross_correlations(field_A, field_B, **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],
+                           G=G,
-                        #    yticks = np.linspace(0.9,1.1,11),
+                           ylims=ylims_power,
-                        #    bound1=0.01,
+                           yticks = yticks_power,
-                        #    bound2=0.02,
+                           bound1=0.01,
-                           ylims=None,
+                           bound2=0.02,
-                           yticks = None,
-                           bound1=None,
-                           bound2=None,
                            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],
+                                ylims=ylims_corr,
-                                # yticks = np.linspace(0.99,1.001,12),
+                                yticks = yticks_corr,
-                                # bound1=0.001,
+                                bound1=0.001,
-                                # bound2=0.002,
+                                bound2=0.002,
-                                ylims=None,
-                                yticks = None,
-                                bound1=None,
-                                bound2=None,
                                 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],
+                                    G=G,
-                                     yticks = np.linspace(0.9,1.1,11),
+                                    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],
+                                        ylims=ylims_corr,
-                                          yticks = np.linspace(0.99,1.001,12),
+                                        yticks = yticks_corr,
                                           bound1=0.001,
                                           bound2=0.002,
                                           kmin=kmin,

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, 
+                                reference=None, 
-                                titles,
+                                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
+    if reference is not None:
-    for i, data in enumerate(data_list):
+        # Plot the data itself
-        im = axes[0, i].imshow(data, cmap=cmap, origin='lower', vmin=vmin, vmax=vmax)
+        for i, data in enumerate(data_list):
-        axes[0, i].set_title(titles[i], fontsize=fs_titles)
+            im = axes[0, i].imshow(data, cmap=cmap, origin='lower', vmin=vmin, vmax=vmax)
-    fig.colorbar(im, ax=axes[0, :], orientation='vertical')
+            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
+        # Plot the data compared to the reference
-    for i, data in enumerate(data_list):
+        for i, data in enumerate(data_list):
-        im = axes[1, i].imshow(data - reference, cmap=cmap, origin='lower', vmin=vmin_diff, vmax=vmax_diff)
+            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_title(f'{titles[i]} - Reference', fontsize=fs_titles)
-    fig.colorbar(im, ax=axes[1, :], orientation='vertical')
+            axes[1, i].set_xticks(ticks[i])
-    
+            axes[1, i].set_yticks(ticks[i])
-    # Add the score on the plots
+            axes[1, i].set_xticklabels(tick_labels[i])
-    for i, data in enumerate(data_list):
+            axes[1, i].set_yticklabels(tick_labels[i])
-        axes[1, i].text(0.5, 0.9, f"Score: {score(data, reference):.2e}", fontsize=10, transform=axes[1, i].transAxes, color='white')
+            axes[1, i].set_xlabel('Mpc/h')
-    # plt.tight_layout()
+        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:

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)
+    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,
+    }

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:

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