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

powerspectrum

Browse source
This commit is contained in:
Mayeul Aubin 2025-03-24 17:34:01 +01:00
parent 0447279fcc
commit fa54f87866
1 changed files with 99 additions and 68 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

111
analysis/power_spectrum.py
View file

@ -5,12 +5,12 @@ kmax = 2e0
Nk = 50 Nk = 50
AliasingCorr=False 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.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:
G = FourierGrid( G = FourierGrid(
field.L0, field.L0,
field.L1, field.L1,
@ -29,15 +29,16 @@ def get_power_spectrum(field, kmin=kmin, kmax=kmax, Nk=Nk):
Pk, _ = get_autocorrelation(field, G, AliasingCorr) Pk, _ = get_autocorrelation(field, G, AliasingCorr)
Pk = Pk[1:] 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.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:
G = FourierGrid( G = FourierGrid(
field_A.L0, field_A.L0,
field_A.L1, field_A.L1,
@ -56,21 +57,21 @@ def get_cross_correlations(field_A, field_B, kmin=kmin, kmax=kmax, Nk=Nk):
_, _, Rks, _ = get_crosscorrelation(field_A, field_B, G, AliasingCorr) _, _, Rks, _ = get_crosscorrelation(field_A, field_B, G, AliasingCorr)
Rks = Rks[1:] 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: if Pk_ref is not None:
ax.plot(k, Pk/Pk_ref-1, **plot_args) ax.plot(k, Pk/Pk_ref, **plot_args)
else: else:
ax.plot(k, Pk, **plot_args) 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) ax.plot(k, Rks, **plot_args)
return ax return ax, G, k, Rks
def plot_power_spectra(filenames, def plot_power_spectra(filenames,
@ -79,6 +80,7 @@ def plot_power_spectra(filenames,
linestyles=None, linestyles=None,
markers=None, markers=None,
Pk_ref=None, Pk_ref=None,
G=None,
ylims=[0.9,1.1], ylims=[0.9,1.1],
yticks = np.linspace(0.9,1.1,11), yticks = np.linspace(0.9,1.1,11),
bound1=0.01, bound1=0.01,
@ -108,9 +110,10 @@ def plot_power_spectra(filenames,
for i, filename in enumerate(filenames): for i, filename in enumerate(filenames):
field = read_field(filename) field = read_field(filename)
add_power_spectrum_to_plot(ax=ax, _, G, k, _ = add_power_spectrum_to_plot(ax=ax,
field=field, field=field,
Pk_ref=Pk_ref, Pk_ref=Pk_ref,
G=G,
plot_args=dict(label=labels[i], plot_args=dict(label=labels[i],
color=colors[i], color=colors[i],
linestyle=linestyles[i], linestyle=linestyles[i],
@ -120,15 +123,15 @@ def plot_power_spectra(filenames,
Nk=Nk), Nk=Nk),
) )
ax.set_xscale('log') ax.set_xscale('log')
ax.set_xlim(kmin, kmax) ax.set_xlim(k.min(),k[-2])
if ylims is not None: if ylims is not None:
ax.set_ylim(ylims) ax.set_ylim(ylims)
if yticks is not None: if yticks is not None:
ax.set_yticks(yticks) 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: 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: else:
ax.set_ylabel('$P(k)$') ax.set_ylabel('$P(k)$')
@ -145,6 +148,7 @@ def plot_power_spectra(filenames,
def plot_cross_correlations(filenames_A, def plot_cross_correlations(filenames_A,
filename_B, filename_B,
G=None,
labels=None, labels=None,
colors=None, colors=None,
linestyles=None, linestyles=None,
@ -180,9 +184,10 @@ def plot_cross_correlations(filenames_A,
for i, filename_A in enumerate(filenames_A): for i, filename_A in enumerate(filenames_A):
field_A = read_field(filename_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_A=field_A,
field_B=field_B, field_B=field_B,
G=G,
plot_args=dict(label=labels[i], plot_args=dict(label=labels[i],
color=colors[i], color=colors[i],
linestyle=linestyles[i], linestyle=linestyles[i],
@ -192,12 +197,12 @@ def plot_cross_correlations(filenames_A,
Nk=Nk), Nk=Nk),
) )
ax.set_xscale('log') ax.set_xscale('log')
ax.set_xlim(kmin, kmax) ax.set_xlim(k.min(), k[-2])
if ylims is not None: if ylims is not None:
ax.set_ylim(ylims) ax.set_ylim(ylims)
if yticks is not None: if yticks is not None:
ax.set_yticks(yticks) 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)$') ax.set_ylabel('$R(k)$')
if bound1 is not None: if bound1 is not None:
@ -211,6 +216,28 @@ def plot_cross_correlations(filenames_A,
return fig, ax 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__": if __name__ == "__main__":
from argparse import ArgumentParser from argparse import ArgumentParser
parser = ArgumentParser(description='Plot power spectra of fields') 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('-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('-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('-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() args = parser.parse_args()
@ -235,12 +264,18 @@ if __name__ == "__main__":
if args.reference is not None: if args.reference is not None:
from pysbmy.field import read_field 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: else:
Pk_ref = None Pk_ref = None
G = None
filenames = [args.directory+f for f in args.filenames] 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: if args.power_spectrum and args.cross_correlation:
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
fig, axes = plt.subplots(2, 1, figsize=(8,8)) fig, axes = plt.subplots(2, 1, figsize=(8,8))
@ -250,14 +285,11 @@ if __name__ == "__main__":
linestyles=args.linestyles, linestyles=args.linestyles,
markers=args.markers, markers=args.markers,
Pk_ref=Pk_ref, 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, kmin=kmin,
kmax=kmax, kmax=kmax,
Nk=Nk, Nk=Nk,
@ -266,18 +298,15 @@ if __name__ == "__main__":
plot_cross_correlations(filenames_A=filenames, plot_cross_correlations(filenames_A=filenames,
filename_B=args.directory+args.reference, filename_B=args.directory+args.reference,
G=G,
labels=args.labels, labels=args.labels,
colors=args.colors, colors=args.colors,
linestyles=args.linestyles, linestyles=args.linestyles,
markers=args.markers, 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, kmin=kmin,
kmax=kmax, kmax=kmax,
Nk=Nk, Nk=Nk,
@ -298,8 +327,9 @@ if __name__ == "__main__":
linestyles=args.linestyles, linestyles=args.linestyles,
markers=args.markers, markers=args.markers,
Pk_ref=Pk_ref, 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, bound1=0.01,
bound2=0.02, bound2=0.02,
kmin=kmin, kmin=kmin,
@ -312,12 +342,13 @@ if __name__ == "__main__":
elif args.cross_correlation: elif args.cross_correlation:
fig, ax = plot_cross_correlations(filenames_A=filenames, fig, ax = plot_cross_correlations(filenames_A=filenames,
filename_B=args.reference, filename_B=args.reference,
G=G,
labels=args.labels, labels=args.labels,
colors=args.colors, colors=args.colors,
linestyles=args.linestyles, linestyles=args.linestyles,
markers=args.markers, 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,
kmin=kmin, kmin=kmin,