diff --git a/sbmy_control/analysis/power_spectrum.py b/sbmy_control/analysis/power_spectrum.py index 7a92370..d97e24f 100644 --- a/sbmy_control/analysis/power_spectrum.py +++ b/sbmy_control/analysis/power_spectrum.py @@ -1,7 +1,7 @@ import numpy as np import os -kmin = 1e-1 +nfirst_kmodes = 10 kmax = 2e0 Nk = 50 AliasingCorr=False @@ -24,12 +24,17 @@ def crop_field(field, Ncrop): field.L2 = field.N2*d2 -def get_power_spectrum(field, kmin=kmin, kmax=kmax, Nk=Nk, G=None): +def get_power_spectrum(field, nfirst_kmodes=nfirst_kmodes, 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: + default_k_modes = PowerSpectrum(field.L0,field.L1,field.L2,int(field.N0),int(field.N1),int(field.N2),).FourierGrid.k_modes + if kmax is None: + kmax = default_k_modes[-1] + if kmax > default_k_modes[-1]: + raise ValueError(f"kmax ({kmax}) is larger than the maximum k mode available in the field ({default_k_modes[-1]}).") G = FourierGrid( field.L0, field.L1, @@ -37,8 +42,8 @@ def get_power_spectrum(field, kmin=kmin, kmax=kmax, Nk=Nk, G=None): int(field.N0), int(field.N1), int(field.N2), - k_modes=np.concat([PowerSpectrum(field.L0,field.L1,field.L2,int(field.N0),int(field.N1),int(field.N2),).FourierGrid.k_modes[:10],np.logspace( - np.log10(kmin), + k_modes=np.concat([default_k_modes[:nfirst_kmodes],np.logspace( + np.log10(default_k_modes[nfirst_kmodes]), np.log10(kmax), Nk, )]), @@ -51,13 +56,18 @@ def get_power_spectrum(field, kmin=kmin, kmax=kmax, Nk=Nk, G=None): return G, k, Pk -def get_cross_correlations(field_A, field_B, kmin=kmin, kmax=kmax, Nk=Nk, G=None): +def get_cross_correlations(field_A, field_B, nfirst_kmodes=nfirst_kmodes, 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: + default_k_modes = PowerSpectrum(field_A.L0,field_A.L1,field_A.L2,int(field_A.N0),int(field_A.N1),int(field_A.N2),).FourierGrid.k_modes + if kmax is None: + kmax = default_k_modes[-1] + if kmax > default_k_modes[-1]: + raise ValueError(f"kmax ({kmax}) is larger than the maximum k mode available in the field ({default_k_modes[-1]}).") G = FourierGrid( field_A.L0, field_A.L1, @@ -65,8 +75,8 @@ def get_cross_correlations(field_A, field_B, kmin=kmin, kmax=kmax, Nk=Nk, G=None int(field_A.N0), int(field_A.N1), int(field_A.N2), - k_modes=np.concat([PowerSpectrum(field_A.L0,field_A.L1,field_A.L2,int(field_A.N0),int(field_A.N1),int(field_A.N2),).FourierGrid.k_modes[:10],np.logspace( - np.log10(kmin), + k_modes=np.concat([default_k_modes[:nfirst_kmodes],np.logspace( + np.log10(default_k_modes[nfirst_kmodes]), np.log10(kmax), Nk, )]), @@ -104,7 +114,7 @@ def plot_power_spectra(filenames, yticks = np.linspace(0.9,1.1,11), bound1=0.01, bound2=0.02, - kmin=kmin, + nfirst_kmodes=nfirst_kmodes, kmax=kmax, Nk=Nk, figsize=(8,4), @@ -139,7 +149,7 @@ def plot_power_spectra(filenames, color=colors[i], linestyle=linestyles[i], marker=markers[i],), - power_args=dict(kmin=kmin, + power_args=dict(nfirst_kmodes=nfirst_kmodes, kmax=kmax, Nk=Nk), ) @@ -178,7 +188,7 @@ def plot_cross_correlations(filenames_A, yticks = np.linspace(0.99,1.001,12), bound1=0.001, bound2=0.002, - kmin=kmin, + nfirst_kmodes=nfirst_kmodes, kmax=kmax, Nk=Nk, figsize=(8,4), @@ -217,7 +227,7 @@ def plot_cross_correlations(filenames_A, color=colors[i], linestyle=linestyles[i], marker=markers[i],), - power_args=dict(kmin=kmin, + power_args=dict(nfirst_kmodes=nfirst_kmodes, kmax=kmax, Nk=Nk), ) @@ -283,8 +293,15 @@ def console_main(): 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.') parser.add_argument('--crop', type=int, default=None, help='Remove the outter N pixels of the fields.') + parser.add_argument('--nfirst_kmodes', type=float, default=10, help='First k modes from all available k modes to be used in the power spectrum before the geomspace part.') + parser.add_argument('--kmax', type=float, default=None, help='Maximum k value for the power spectrum geomspace part. If None, it will be set to the maximum k mode available in the field.') + parser.add_argument('--Nk', type=int, default=50, help='Number of k values for the power spectrum geomspace part.') args = parser.parse_args() + + nfirst_kmodes = args.nfirst_kmodes + kmax = args.kmax + Nk = args.Nk if not args.power_spectrum and not args.cross_correlation: print('You must choose between power_spectrum and cross_correlation.') @@ -298,7 +315,7 @@ def console_main(): from pysbmy.field import read_field F_ref = read_field(args.directory+args.reference) crop_field(F_ref, args.crop) - G, _, Pk_ref = get_power_spectrum(F_ref, kmin=kmin, kmax=kmax, Nk=Nk) + G, _, Pk_ref = get_power_spectrum(F_ref, nfirst_kmodes=nfirst_kmodes, kmax=kmax, Nk=Nk) else: Pk_ref = None G = None @@ -325,7 +342,7 @@ def console_main(): yticks = yticks_power, bound1=0.01, bound2=0.02, - kmin=kmin, + nfirst_kmodes=nfirst_kmodes, kmax=kmax, Nk=Nk, ax=axes[0], @@ -344,7 +361,7 @@ def console_main(): yticks = yticks_corr, bound1=0.001, bound2=0.002, - kmin=kmin, + nfirst_kmodes=nfirst_kmodes, kmax=kmax, Nk=Nk, ax=axes[1], @@ -371,7 +388,7 @@ def console_main(): yticks = yticks_power, bound1=0.01, bound2=0.02, - kmin=kmin, + nfirst_kmodes=nfirst_kmodes, kmax=kmax, Nk=Nk, Ncrop=args.crop, @@ -392,7 +409,7 @@ def console_main(): yticks = yticks_corr, bound1=0.001, bound2=0.002, - kmin=kmin, + nfirst_kmodes=nfirst_kmodes, kmax=kmax, Nk=Nk, Ncrop=args.crop,