"""Classical foreground fitting and subtraction methods.
All fitters implement the same small interface: call ``run(data_cube,
noise_cube)`` and receive a :class:`FitterResult` containing the foreground
model in :attr:`FitterResult.fit` and the residual data in
:attr:`FitterResult.sub`. Both outputs retain the input cube geometry, so the
residual can be passed directly to flagging or power-spectrum estimation.
A common deterministic fit models each visibility with a smooth function of
frequency. The noise cube supplies the per-frequency noise scale used by the
fit::
from ps_eor import fgfit
fitter = fgfit.PolyForegroundFit(
deg=3,
fit_type='power_poly',
)
result = fitter.run(data_cube, noise_cube)
foreground = result.fit
residual = result.sub
Valid polynomial families are ``'poly'``, ``'bernstein'``, ``'power_poly'``,
and ``'power_bernstein'``. The power variants are often useful for smooth
foreground spectra.
For a component-based separation, :class:`PcaForegroundFit` fits the real and
imaginary data independently and also exposes its component representation::
result = fgfit.PcaForegroundFit(n_cmpt=3).run(data_cube, noise_cube)
first_component = result.inverse_transform(0)
These fitters provide inexpensive deterministic baselines. Probabilistic
foreground separation and posterior sampling are available through
:mod:`ps_eor.ml_gpr`.
"""
import numpy as np
from . import datacube, fitutil
[docs]
class FitterResult:
"""Foreground model and residual cubes.
Attributes:
fit: Fitted foreground model.
sub: Input data minus the model.
"""
def __init__(self, cube_fit, cube_sub):
self.fit = cube_fit
self.sub = cube_sub
[docs]
class MixForegroundResult(FitterResult):
"""A fitted model with its component mixing representation."""
def __init__(self, cube_fit, cube_sub, cube_mix, inv_trans):
FitterResult.__init__(self, cube_fit, cube_sub)
self.mix = cube_mix
self.inv_trans = inv_trans
[docs]
def get_component(self, n):
"""The ``n``-th fitted component as a one-channel cube.
Returns:
CartDataCube: component ``n`` of the mixing cube.
"""
return self.mix.get_freq(n)
[docs]
class AbstractForegroundFitter:
"""Interface implemented by foreground fitters."""
[docs]
def run(self, data_cube, data_cube_noise):
"""Fit the foregrounds in ``data_cube`` and separate them from the data.
Args:
data_cube: the data to fit (a :class:`~ps_eor.datacube.CartDataCube`).
data_cube_noise: a noise cube of the same geometry, used to weight
the fit and set thresholds.
Returns:
FitterResult: the foreground model (``.fit``) and residual (``.sub``).
"""
raise NotImplementedError()
[docs]
class NoAction(AbstractForegroundFitter):
"""A no-op fitter: a zero model, leaving the data untouched."""
[docs]
def run(self, i_cube, v_cube):
"""Return a zero foreground model, keeping ``i_cube`` as the residual.
Returns:
FitterResult: ``.fit`` is zeros, ``.sub`` is ``i_cube`` unchanged.
"""
return FitterResult(i_cube.new_with_data(np.zeros_like(i_cube.data)), i_cube)
[docs]
class GmcaForegroundFit(AbstractForegroundFitter):
"""Foreground separation with Generalized Morphological Component Analysis."""
def __init__(self, n_cmpt, mints=0, do_wave_transform=False, do_poly_fit=0):
"""Configure the GMCA model.
Args:
n_cmpt: Number of GMCA components.
mints: Final threshold in units of the median absolute deviation.
do_wave_transform: Transform the frequency axis before GMCA.
do_poly_fit: Polynomial degree applied to the GMCA reconstruction;
zero disables this step.
"""
self.n_cmpt = n_cmpt
self.mints = mints
self.do_wave_transform = do_wave_transform
self.do_poly_fit = do_poly_fit
[docs]
def run(self, data_cube, data_cube_noise):
"""Separate foregrounds with GMCA (see :meth:`AbstractForegroundFitter.run`).
Returns:
FitterResult: the GMCA model (``.fit``) and residual (``.sub``).
"""
w = 1
if isinstance(data_cube, datacube.DataCube) and data_cube.weights is not None:
w = data_cube.weights.data
y = data_cube.data * w
y_fit = fitutil.alm_gmca_fit(y, self.n_cmpt, data_cube_noise.data,
self.mints, self.do_wave_transform, self.do_poly_fit)
cube_fit = data_cube.new_with_data(y_fit / w)
cube_sub = data_cube.new_with_data((y - y_fit) / w)
cube_sub.set_weights(data_cube.weights)
return FitterResult(cube_fit, cube_sub)
[docs]
class PcaForegroundFit(AbstractForegroundFitter):
"""Foreground separation using independent PCA fits to real and imaginary data."""
def __init__(self, n_cmpt):
"""Configure a fit with ``n_cmpt`` principal components."""
self.n_cmpt = n_cmpt
[docs]
def run(self, data_cube, data_cube_noise):
"""Separate foregrounds with PCA (see :meth:`AbstractForegroundFitter.run`).
Returns:
MixForegroundResult: the model (``.fit``), residual (``.sub``) and the
component mixing cube (``.mix`` / :meth:`MixForegroundResult.get_component`).
"""
y = data_cube.data
y_fit, y_mix, inv_trans = fitutil.alm_pca_fit(y, self.n_cmpt, return_mix=True)
cube_fit = data_cube.new_with_data(y_fit)
cube_sub = data_cube.new_with_data(y - y_fit)
cube_mix = data_cube.new_with_data(y_mix, freqs=np.arange(self.n_cmpt))
cube_sub.set_weights(data_cube.weights)
return MixForegroundResult(cube_fit, cube_sub, cube_mix, inv_trans)
[docs]
class PolyForegroundFit(AbstractForegroundFitter):
"""Fit a smooth spectral model independently to each visibility."""
def __init__(self, deg, fit_type):
"""Configure the spectral model.
Args:
deg: Polynomial degree.
fit_type: ``'poly'``, ``'bernstein'``, ``'power_poly'``, or
``'power_bernstein'``.
"""
self.deg = deg
self.fit_type = fit_type
self.fit_fct = fitutil.get_fit_fct(self.fit_type)
[docs]
def run(self, data_cube, data_cube_noise):
"""Fit a smooth spectral model per visibility (see
:meth:`AbstractForegroundFitter.run`).
Returns:
FitterResult: the smooth model (``.fit``) and residual (``.sub``).
"""
x = data_cube.freqs / 1e6
y = data_cube.data
y_v = data_cube_noise.data
noiserms = np.std(y_v, axis=1)
y_fit = fitutil.alm_poly_fit(x, y, noiserms, self.deg, self.fit_fct)
cube_fit = data_cube.new_with_data(y_fit)
cube_sub = data_cube.new_with_data(y - y_fit)
cube_sub.set_weights(data_cube.weights)
return FitterResult(cube_fit, cube_sub)