"""Numerical helpers for classical foreground fitting."""
import functools
import numpy as np
from scipy.special import comb
from sklearn.decomposition import PCA
from . import pscart, psutil
[docs]
def inv_pca(X, Y, pca, mode):
"""Reconstruct data from a fitted PCA model.
Args:
X: the original data (kept for interface symmetry; unused).
Y: the PCA coefficients.
pca: the fitted sklearn ``PCA`` object.
mode: reconstruct from this single component, or ``None`` for all.
Returns:
ndarray: the reconstruction.
"""
if mode is None:
return pca.inverse_transform(Y[:, :])
idx = slice(mode, mode + 1) if mode is not None else slice(None)
s = np.sqrt(pca.explained_variance_[idx, np.newaxis]) if pca.whiten else 1
return np.dot(Y[:, idx], s * pca.components_[idx, :]) + pca.mean_[:]
[docs]
def inv_pca_complex(X_real, X_imag, Y_real, Y_imag, pca_real, pca_imag, n_mode):
"""Reconstruct complex data from independent real / imaginary PCA models.
Args:
X_real, X_imag: the original real / imaginary data.
Y_real, Y_imag: their PCA coefficients.
pca_real, pca_imag: the fitted ``PCA`` objects.
n_mode: reconstruct from this single component, or ``None`` for all.
Returns:
ndarray: the complex reconstruction.
"""
Xrec_real = inv_pca(X_real, Y_real, pca_real, n_mode)
Xrec_imag = inv_pca(X_imag, Y_imag, pca_imag, n_mode)
return Xrec_real + 1j * Xrec_imag
[docs]
def alm_pca_fit(alm, n_cmpt, verbose=True, return_mix=False):
"""Fit ``n_cmpt`` PCA components to complex frequency-dependent data (real and
imaginary parts fitted independently).
Args:
alm: complex data, shape ``(n_freqs, n_modes)``.
n_cmpt: number of principal components.
verbose (bool): print the explained-variance ratio.
return_mix (bool): also return the component mixing and the inverse
transform.
Returns:
ndarray: the fitted (smooth) data; or ``(fit, mix, inv_trans)`` when
``return_mix``.
"""
X_real = alm.T.real
X_imag = alm.T.imag
pca_real = PCA(n_components=n_cmpt, whiten=True)
pca_imag = PCA(n_components=n_cmpt, whiten=True)
Y_real = pca_real.fit_transform(X_real)
Y_imag = pca_imag.fit_transform(X_imag)
if verbose:
print('PCA: Percentage of variance explained:', pca_real.explained_variance_ratio_)
inv_trans = functools.partial(inv_pca_complex, X_real, X_imag, Y_real, Y_imag, pca_real,
pca_imag)
alm_pca_fitted = inv_trans(None).T
if return_mix:
return alm_pca_fitted, (Y_real + 1j * Y_imag).T, inv_trans
return alm_pca_fitted
[docs]
def gmca_fit(X, n_cmpt, mints=0, do_wave_transform=False, do_poly_fit=0, X_n=None):
"""Fit ``n_cmpt`` GMCA components to a real-valued array.
Requires the optional ``pyGMCA`` and ``libwise`` packages.
Args:
X: real data, shape ``(n_freqs, n_samples)``.
n_cmpt: number of GMCA components.
mints: final threshold, in units of the median absolute deviation.
do_wave_transform (bool): wavelet-transform the frequency axis first.
do_poly_fit (int): if > 0, smooth the reconstruction with a Bernstein
polynomial of this degree.
X_n: optional noise array (kept for interface symmetry; unused).
Returns:
ndarray: the GMCA reconstruction, same shape as ``X``.
"""
from libwise import wtutils
from pyGMCA.bss.amca import pyAMCA as pam
n_scales = 6
if do_wave_transform:
Xw = np.hstack(wtutils.wavedec(X, 'b1', n_scales, dec=wtutils.uiwt, axis=0,
boundary='symm'))
else:
Xw = X
S, A = pam.AMCA(Xw, n_cmpt, mints=mints)
if do_wave_transform:
Y = wtutils.waverec(np.split(np.dot(A, S).real, n_scales + 1, axis=1),
'b1', rec=wtutils.uiwt_inv, axis=0, boundary='symm')
else:
Y = np.dot(A, S).real
if do_poly_fit > 0:
Y = alm_poly_fit(np.arange(X.shape[0]), Y, np.ones(X.shape[0]), do_poly_fit,
fit_fct=bernstein_fit)
return Y
[docs]
def alm_gmca_fit(alm, n_cmpt, alm_n, mints=0, do_wave_transform=False, do_poly_fit=0):
"""GMCA fit applied independently to the real and imaginary parts of ``alm``.
Args:
alm: complex data, shape ``(n_freqs, n_modes)``.
n_cmpt: number of GMCA components.
alm_n: noise array passed through to :func:`gmca_fit`.
mints, do_wave_transform, do_poly_fit: see :func:`gmca_fit`.
Returns:
ndarray: the complex GMCA reconstruction.
"""
Yr = gmca_fit(alm.real, n_cmpt, mints=mints, do_wave_transform=do_wave_transform,
do_poly_fit=do_poly_fit, X_n=alm_n)
Yi = gmca_fit(alm.imag, n_cmpt, mints=mints, do_wave_transform=do_wave_transform,
do_poly_fit=do_poly_fit, X_n=alm_n)
return Yr + 1j * Yi
[docs]
def bernstein_poly(i, n, x):
"""The ``i``-th Bernstein basis polynomial of degree ``n``, evaluated at ``x``.
Returns:
ndarray: the basis values at ``x``.
"""
return comb(n, i) * (x ** (n - i)) * (1 - x)**i
[docs]
def poly_fit(x, y, noiserms, deg, min_deg=0, full_cov=False):
"""Noise-weighted least-squares polynomial fit of ``y`` against ``x``.
Args:
x, y: the data points.
noiserms: per-point 1-sigma noise, used as inverse-variance weights.
deg: maximum polynomial degree.
min_deg: minimum degree (drop lower-order terms).
full_cov (bool): return the full model covariance instead of
per-coefficient standard deviations.
Returns:
tuple: ``(coeffs, model, cov)`` -- the fitted coefficients, the model at
``x``, and either the full model covariance (``full_cov``) or the
per-coefficient standard deviations.
"""
C_Dinv = np.diagflat(1 / noiserms ** 2)
A = np.vstack([x ** k for k in np.arange(min_deg, deg + 1)]).T
lhs = np.dot(np.dot(A.T, C_Dinv), A)
rhs = np.dot(np.dot(A.T, C_Dinv), y)
s = np.linalg.solve(lhs, rhs)
y_s = np.dot(A, s)
if full_cov:
cov_sigma = np.dot(np.dot(A, np.linalg.inv(lhs)), A.T)
else:
cov_sigma = np.sqrt(np.diag(np.linalg.inv(lhs)))
return s, y_s, cov_sigma
[docs]
def bernstein_fit(x, y, noiserms, deg, full_cov=False):
"""Noise-weighted least-squares fit of a Bernstein-polynomial expansion.
Same arguments and return as :func:`poly_fit`, using a Bernstein basis of
degree ``deg`` (``x`` should lie in ``[0, 1]``).
Returns:
tuple: ``(coeffs, model, cov)`` -- see :func:`poly_fit`.
"""
ber_basis = []
for j in range(deg):
for i in range(j + 1):
ber_basis.append(bernstein_poly(i, j, x))
C_Dinv = np.diagflat(1 / noiserms ** 2)
A = np.vstack(ber_basis).T
lhs = np.dot(np.dot(A.T, C_Dinv), A)
rhs = np.dot(np.dot(A.T, C_Dinv), y)
s = np.linalg.solve(lhs, rhs)
if full_cov:
cov_sigma = np.dot(np.dot(A, np.linalg.pinv(lhs)), A.T)
else:
cov_sigma = np.sqrt(np.diag(np.linalg.inv(lhs)))
y_s = np.dot(A, s)
return s, y_s, cov_sigma
[docs]
def powerlaw_fit(x, y, noiserms, deg, min_deg=0, bernstein=False, offset=True,
output_parameters=False):
"""Fit a smooth spectrum as a polynomial in log-log space (a generalised
power law).
Args:
x, y: the data (e.g. frequency and amplitude).
noiserms: per-point 1-sigma noise.
deg: polynomial degree in log-log space.
min_deg: minimum degree.
bernstein (bool): use a Bernstein basis instead of a plain polynomial.
offset (bool): add an offset so the log is defined for non-positive ``y``.
output_parameters (bool): also return the fitted coefficients.
Returns:
ndarray: the model evaluated at ``x``; or ``(model, coeffs)`` when
``output_parameters``.
"""
if x[0] > 0:
x = x / x[0]
sgn = np.sign(np.mean(y))
y_m = y * sgn
offset = -10 * min(np.min(y), -0.05) if offset else 0
y_m += offset
if bernstein:
s, y_s, _cov_sigma = bernstein_fit(np.log(x), np.log(y_m), noiserms, deg)
else:
s, y_s, _cov_sigma = poly_fit(np.log(x), np.log(y_m), noiserms, deg, min_deg=min_deg)
y_s = sgn * (np.exp(y_s) - offset)
if output_parameters:
return y_s, s
return y_s
[docs]
def powerlaw_fit_bernstein(x, y, noiserms, deg):
"""A :func:`powerlaw_fit` using a Bernstein basis (``bernstein=True``).
Returns:
ndarray: the model evaluated at ``x``.
"""
return powerlaw_fit(x, y, noiserms, deg, bernstein=True)
[docs]
def get_fit_fct(fit_type):
"""The fit function named by ``fit_type``.
Args:
fit_type (str): one of ``'poly'``, ``'bernstein'``, ``'power_poly'``,
``'power_bernstein'``.
Returns:
callable: the matching fit function.
"""
if fit_type == 'poly':
return poly_fit
elif fit_type == 'bernstein':
return bernstein_fit
elif fit_type == 'power_poly':
return powerlaw_fit
elif fit_type == 'power_bernstein':
return powerlaw_fit_bernstein
else:
print(f"Error: fit_type '{fit_type}' invalid")
[docs]
def alm_poly_fit(freqs, alm, noiserms, deg, fit_fct=powerlaw_fit, **kargs):
"""Fit a spectral model to every column of ``alm`` at once.
Args:
freqs: the spectral axis (e.g. frequency in MHz).
alm: data, shape ``(n_freqs, n_cols)``.
noiserms: per-channel 1-sigma noise.
deg: model degree.
fit_fct: the fit function to apply (default :func:`powerlaw_fit`).
**kargs: forwarded to ``fit_fct``.
Returns:
ndarray: the fitted model, same shape as ``alm``.
"""
res = fit_fct(freqs, alm, noiserms, deg, **kargs)
if len(res) == 3:
res = res[1]
return res
[docs]
def fit_cl_cube_poly(cube, du=10, poly_deg=4, log=True):
"""Fit a smooth polynomial to a cube's angular power spectrum P(k_per).
Args:
cube: the data cube.
du: uv-bin width, in wavelengths.
poly_deg: polynomial degree.
log (bool): fit in log-log space.
Returns:
tuple: ``(fit_fct, uv_mean, ps, ps_err)`` -- the fitted function of
baseline length, the bin centres, and the measured power and its error.
"""
uv_bins = np.arange(cube.ru.min(), cube.ru.max(), du)
uv_mean = np.array([(a + b) / 2 for (a, b) in psutil.pairwise(uv_bins)])
w = cube.weights.data if cube.weights is not None else None
ps, ps_err, _, _ = pscart.get_cross_power_spectra(
cube.data, cube.data, cube.uu, cube.vv, 2 * np.pi * uv_mean, weight_cube=w,
uniform_bins=True)
ps = ps.real.mean(axis=0)
ps_err = ps_err.real.mean(axis=0)
if log:
fit_fct = psutil.polynomial_fit_log(uv_mean, ps, ps_err, poly_deg)
else:
fit_fct = psutil.polynomial_fit(uv_mean, ps, ps_err, poly_deg)
return fit_fct, uv_mean, ps, ps_err