maubin/sbmy_control
0
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

better k_modes for power spectra

Browse source
This commit is contained in:
Mayeul Aubin 2025-07-01 15:15:35 +02:00
parent bf76ce01e9
commit 28fb27ede2
1 changed files with 33 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

49
sbmy_control/analysis/power_spectrum.py
View file

@ -1,7 +1,7 @@
import numpy as np import numpy as np
import os import os
kmin = 1e-1 nfirst_kmodes = 10
kmax = 2e0 kmax = 2e0
Nk = 50 Nk = 50
AliasingCorr=False AliasingCorr=False
@ -24,12 +24,17 @@ def crop_field(field, Ncrop):
field.L2 = field.N2*d2 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.power import PowerSpectrum
from pysbmy.fft import FourierGrid from pysbmy.fft import FourierGrid
from pysbmy.correlations import get_autocorrelation from pysbmy.correlations import get_autocorrelation
if G is None: 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( G = FourierGrid(
field.L0, field.L0,
field.L1, field.L1,
@ -37,8 +42,8 @@ def get_power_spectrum(field, kmin=kmin, kmax=kmax, Nk=Nk, G=None):
int(field.N0), int(field.N0),
int(field.N1), int(field.N1),
int(field.N2), 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( k_modes=np.concat([default_k_modes[:nfirst_kmodes],np.logspace(
np.log10(kmin), np.log10(default_k_modes[nfirst_kmodes]),
np.log10(kmax), np.log10(kmax),
Nk, Nk,
)]), )]),
@ -51,13 +56,18 @@ def get_power_spectrum(field, kmin=kmin, kmax=kmax, Nk=Nk, G=None):
return G, k, Pk 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.power import PowerSpectrum
from pysbmy.fft import FourierGrid from pysbmy.fft import FourierGrid
from pysbmy.correlations import get_crosscorrelation from pysbmy.correlations import get_crosscorrelation
if G is None: 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( G = FourierGrid(
field_A.L0, field_A.L0,
field_A.L1, 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.N0),
int(field_A.N1), int(field_A.N1),
int(field_A.N2), 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( k_modes=np.concat([default_k_modes[:nfirst_kmodes],np.logspace(
np.log10(kmin), np.log10(default_k_modes[nfirst_kmodes]),
np.log10(kmax), np.log10(kmax),
Nk, Nk,
)]), )]),
@ -104,7 +114,7 @@ def plot_power_spectra(filenames,
yticks = np.linspace(0.9,1.1,11), yticks = np.linspace(0.9,1.1,11),
bound1=0.01, bound1=0.01,
bound2=0.02, bound2=0.02,
kmin=kmin, nfirst_kmodes=nfirst_kmodes,
kmax=kmax, kmax=kmax,
Nk=Nk, Nk=Nk,
figsize=(8,4), figsize=(8,4),
@ -139,7 +149,7 @@ def plot_power_spectra(filenames,
color=colors[i], color=colors[i],
linestyle=linestyles[i], linestyle=linestyles[i],
marker=markers[i],), marker=markers[i],),
power_args=dict(kmin=kmin, power_args=dict(nfirst_kmodes=nfirst_kmodes,
kmax=kmax, kmax=kmax,
Nk=Nk), Nk=Nk),
) )
@ -178,7 +188,7 @@ def plot_cross_correlations(filenames_A,
yticks = np.linspace(0.99,1.001,12), yticks = np.linspace(0.99,1.001,12),
bound1=0.001, bound1=0.001,
bound2=0.002, bound2=0.002,
kmin=kmin, nfirst_kmodes=nfirst_kmodes,
kmax=kmax, kmax=kmax,
Nk=Nk, Nk=Nk,
figsize=(8,4), figsize=(8,4),
@ -217,7 +227,7 @@ def plot_cross_correlations(filenames_A,
color=colors[i], color=colors[i],
linestyle=linestyles[i], linestyle=linestyles[i],
marker=markers[i],), marker=markers[i],),
power_args=dict(kmin=kmin, power_args=dict(nfirst_kmodes=nfirst_kmodes,
kmax=kmax, kmax=kmax,
Nk=Nk), Nk=Nk),
) )
@ -283,9 +293,16 @@ 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('-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('-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('--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() 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: if not args.power_spectrum and not args.cross_correlation:
print('You must choose between power_spectrum and cross_correlation.') print('You must choose between power_spectrum and cross_correlation.')
exit() exit()
@ -298,7 +315,7 @@ def console_main():
from pysbmy.field import read_field from pysbmy.field import read_field
F_ref = read_field(args.directory+args.reference) F_ref = read_field(args.directory+args.reference)
crop_field(F_ref, args.crop) 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: else:
Pk_ref = None Pk_ref = None
G = None G = None
@ -325,7 +342,7 @@ def console_main():
yticks = yticks_power, yticks = yticks_power,
bound1=0.01, bound1=0.01,
bound2=0.02, bound2=0.02,
kmin=kmin, nfirst_kmodes=nfirst_kmodes,
kmax=kmax, kmax=kmax,
Nk=Nk, Nk=Nk,
ax=axes[0], ax=axes[0],
@ -344,7 +361,7 @@ def console_main():
yticks = yticks_corr, yticks = yticks_corr,
bound1=0.001, bound1=0.001,
bound2=0.002, bound2=0.002,
kmin=kmin, nfirst_kmodes=nfirst_kmodes,
kmax=kmax, kmax=kmax,
Nk=Nk, Nk=Nk,
ax=axes[1], ax=axes[1],
@ -371,7 +388,7 @@ def console_main():
yticks = yticks_power, yticks = yticks_power,
bound1=0.01, bound1=0.01,
bound2=0.02, bound2=0.02,
kmin=kmin, nfirst_kmodes=nfirst_kmodes,
kmax=kmax, kmax=kmax,
Nk=Nk, Nk=Nk,
Ncrop=args.crop, Ncrop=args.crop,
@ -392,7 +409,7 @@ def console_main():
yticks = yticks_corr, yticks = yticks_corr,
bound1=0.001, bound1=0.001,
bound2=0.002, bound2=0.002,
kmin=kmin, nfirst_kmodes=nfirst_kmodes,
kmax=kmax, kmax=kmax,
Nk=Nk, Nk=Nk,
Ncrop=args.crop, Ncrop=args.crop,