Commit eb17f0c4 authored by alexandre beelen's avatar alexandre beelen
Browse files

feat(generator) Add gen_psffield to create syntetic fields from catalogs

parent d42cd047
Pipeline #2602 canceled with stage
......@@ -7,17 +7,19 @@ Basic Usage of powspec
import matplotlib.pyplot as plt
# sphinx_gallery_thumbnail_path = '_static/demo.png'
import astropy.units as u
from astropy.convolution import Gaussian2DKernel
from astropy.visualization import quantity_support
import numpy as np
from powspec.powspec import power_spectral_density
from powspec.utils.generator import gen_pkfield
from powspec.utils.generator import gen_pkfield, gen_psffield
quantity_support()
# %%
# Create fake images
# ------------------
# Create fake Gaussian field images
# ---------------------------------
#
# Create a list of fake images with different P(k)
#
......@@ -61,3 +63,36 @@ for i, (image, powspec, alpha) in enumerate(zip(images, powspecs, alphas)):
ax.imshow(image.value, origin="lower")
plt.show()
# %%
# Create PSF field image
# ----------------------
#
# Create a fake catalog of sources
n_pix = 1024
n_sources = 1024*5
positions = np.random.uniform(0, n_pix, size=(2, n_sources))
fluxes = np.random.uniform(1, 10, n_sources)
image = gen_psffield(positions, fluxes, n_pix, kernel=Gaussian2DKernel(10), factor=4) * u.Jy / u.beam
# %%
# Compute P(k)
# ------------
#
# Compute power spectra of each images
#
powspec, k = power_spectral_density(image, res=res, range='tight-log', bins=auto)
k_mid = np.mean(u.Quantity([k[1:], k[:-1]]), axis=0)
# %%
# Plots
# -----
fig, axes = plt.subplots(ncols=2)
axes[0].imshow(image.value, origin='lower')
axes[1].loglog(k_mid, powspec)
fig.show()
import numpy as np
from copy import deepcopy
from astropy.convolution import Kernel2D, convolve_fft
from astropy.convolution.kernels import _round_up_to_odd_integer
__all__ = ["Pk", "gen_pkfield"]
......@@ -9,8 +13,24 @@ def Pk(k, alpha=-11.0 / 3, fknee=1):
def gen_pkfield(npix=32, alpha=-11.0 / 3, fknee=1, res=1):
"""Generate a 2D square map with P(k) field"""
"""Generate a gaussian field with (k/k_0)^alpha law.
Parameters
----------
npix : int, optional
number of pixels for the map, by default 32
alpha : [float], optional
the power law index, by default -11.0/3
fknee : float or ~astropy.units.Quantity, optional
the knee frequency in 1/res unit, where P(k) = 1, by default 1
res : float or ~astropy.units.Quantity, optional
size of a pixel, by default 1
Returns
-------
data : ndarray, shape(n_pix, n_pix)
the requested gaussian field
"""
ufreq = np.fft.fftfreq(npix, d=res)
kfreq = np.sqrt(ufreq[:, np.newaxis] ** 2 + ufreq ** 2)
......@@ -22,3 +42,67 @@ def gen_pkfield(npix=32, alpha=-11.0 / 3, fknee=1, res=1):
fft_img = np.sqrt(psd) * (np.cos(pha) + 1j * np.sin(pha))
return np.real(np.fft.ifft2(fft_img)) * npix / res ** 2
def gen_psffield(positions, fluxes=None, shape=(32, 32), kernel=None, factor=None):
"""Generate a point spread function field given a catalog of point source.
Parameters
----------
positions : array_like, shape (2, M)
x, y positions in pixel coordinates
fluxes : array_like, shape (M,)
corresponding peak fluxes
shape : int or [int, int], optional
the output image shape
kernel : ~astropy.convolution.Kernel2D, optional
the 2D kernel to be used for the PSF
factor : [int], optional
a overpixelization factor used for the projection before smoothing, by default None
Returns
-------
array : ndarray, shape(nx, ny)
The corresponding map
"""
if factor is None:
factor = 1
if fluxes is None:
fluxes = np.ones(positions.shape[1])
if isinstance(shape, (int, np.int)):
shape = [shape, shape]
_shape = np.array(shape) * factor
_positions = (np.asarray(positions) + 0.5) * factor - 0.5
if kernel is not None:
# Upscale the kernel with factor
kernel = deepcopy(kernel)
for param in ["x_stddev", "y_stddev", "width", "radius", "radius_in"]:
if param in kernel._model.param_names:
getattr(kernel._model, param).value *= factor
Kernel2D.__init__(
kernel,
x_size=_round_up_to_odd_integer(kernel.shape[1] * factor),
y_size=_round_up_to_odd_integer(kernel.shape[0] * factor),
),
# Range are maximum bins edges
hist2d_kwd = {"bins": _shape, "range": ((-0.5, _shape[0] - 0.5), (-0.5, _shape[1] - 0.5))}
# reverse the _positions because it needs to be y x
array = np.histogram2d(*_positions[::-1], weights=fluxes, **hist2d_kwd)[0]
# Remove nan if present
array[np.isnan(array)] = 0
if kernel is not None:
kernel.normalize("peak")
array = convolve_fft(array, kernel, normalize_kernel=False, boundary="wrap") / factor ** 2
# Average rebinning onto the input shape
array = array.reshape((shape[0], factor, shape[1], factor)).sum(-1).sum(1)
return array
......@@ -9,7 +9,7 @@ def test_Pk():
assert np.all(Pk(np.arange(10), alpha=0, fknee=1) == np.ones(10))
def gen_pkfield():
def test_gen_pkfield():
from ..generator import gen_pkfield
np.random.seed(42)
......@@ -18,3 +18,112 @@ def gen_pkfield():
assert pkfield.shape == (256, 256)
npt.assert_almost_equal(np.mean(pkfield), 0)
npt.assert_almost_equal(np.std(pkfield), 1, decimal=3)
def test_gen_psffield():
from ..generator import gen_psffield
from astropy.convolution import Gaussian2DKernel
np.random.seed(42)
n_pix = 64
n_sources = 3
sigma = 2
positions = np.random.uniform(n_pix // 5, n_pix * 4 // 5, size=(2, n_sources))
m = gen_psffield(positions, shape=n_pix, kernel=None, factor=None)
assert np.sum(m) == n_sources
shape = (n_pix, n_pix)
m = gen_psffield(positions, shape=shape, kernel=None, factor=None)
assert np.sum(m) == n_sources
m_f = gen_psffield(positions, shape=shape, kernel=None, factor=10)
npt.assert_equal(m, m_f)
m = gen_psffield(positions, shape=shape, kernel=Gaussian2DKernel(sigma), factor=None)
npt.assert_almost_equal(np.log10(n_sources * 2 * np.pi * sigma ** 2), np.log10(m.sum()), decimal=4)
m = gen_psffield(positions, shape=shape, kernel=Gaussian2DKernel(sigma), factor=10)
npt.assert_almost_equal(np.log10(n_sources * 2 * np.pi * sigma ** 2), np.log10(m.sum()), decimal=4)
def pixelated_gaussian(shape, center, sigma):
"""x, y bin edges"""
from scipy.special import erf
x = np.linspace(-0.5, shape[1] - 0.5, shape[1] + 1)
y = np.linspace(-0.5, shape[0] - 0.5, shape[0] + 1)
x, y = np.meshgrid(x, y)
x1 = x[:-1, :-1]
x2 = x[1:, 1:]
y1 = y[:-1, :-1]
y2 = y[1:, 1:]
m = (
1
/ (4 * (x2 - x1) * (y2 - y1))
* (erf((x2 - center[0]) / (np.sqrt(2) * sigma[0])) - erf((x1 - center[0]) / (np.sqrt(2) * sigma[0])))
* (erf((y2 - center[1]) / (np.sqrt(2) * sigma[1])) - erf((y1 - center[1]) / (np.sqrt(2) * sigma[1])))
* (2 * np.pi * sigma[0] * sigma[1])
)
return m
m_pix = np.sum([pixelated_gaussian(m.shape, center, (sigma, sigma)) for center in positions.T], axis=0)
# Should be better
npt.assert_almost_equal(m_pix, m, decimal=2)
assert np.all(np.abs(m_pix - m) < 1e-1)
fluxes = [1, 2, 3]
m_f = gen_psffield(positions, fluxes, shape=shape, kernel=Gaussian2DKernel(sigma), factor=10)
npt.assert_almost_equal(np.log10(2 * np.pi * sigma ** 2 * np.sum(fluxes)), np.log10(m_f.sum()), decimal=4)
def main():
import matplotlib.pyplot as plt
from astropy.convolution import Gaussian2DKernel
plt.ion()
def pixelated_gaussian(shape, center, sigma):
"""x, y bin edges"""
from scipy.special import erf
x = np.linspace(-0.5, shape[1] - 0.5, shape[1] + 1)
y = np.linspace(-0.5, shape[0] - 0.5, shape[0] + 1)
x, y = np.meshgrid(x, y)
x1 = x[:-1, :-1]
x2 = x[1:, 1:]
y1 = y[:-1, :-1]
y2 = y[1:, 1:]
m = (
1
/ (4 * (x2 - x1) * (y2 - y1))
* (erf((x2 - center[0]) / (np.sqrt(2) * sigma[0])) - erf((x1 - center[0]) / (np.sqrt(2) * sigma[0])))
* (erf((y2 - center[1]) / (np.sqrt(2) * sigma[1])) - erf((y1 - center[1]) / (np.sqrt(2) * sigma[1])))
* (2 * np.pi * sigma[0] * sigma[1])
)
return m
def centered_gaussian(shape, center, sigma):
"""x, y bin edges"""
x = np.linspace(0, shape[1] - 1, shape[1])
y = np.linspace(0, shape[0] - 1, shape[0])
x, y = np.meshgrid(x, y)
m = np.exp(-((x - center[0]) ** 2) / (2 * sigma[0] ** 2) - (y - center[1]) ** 2 / (2 * sigma[1] ** 2))
return m
n_pix = 11
n_sources = 3
np.random.seed(0)
positions = np.random.uniform(1, n_pix - 1, size=(2, n_sources))
shape = (n_pix, n_pix)
test = np.sum([centered_gaussian(shape, center, (1, 1)) for center in positions.T], axis=0)
pix_test = np.sum([pixelated_gaussian(shape, center, (1, 1)) for center in positions.T], axis=0)
m = gen_psffield(positions, shape, kernel=Gaussian2DKernel(1), factor=4)
plt.clf()
plt.imshow(m, extent=(-0.5, shape[0] - 0.5, -0.5, shape[1] - 0.5), origin="lower")
plt.scatter(*positions)
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment