Source code for ps_eor.sphcube

"""Frequency-dependent sky data in HEALPix and spherical harmonics.

This module is the spherical counterpart of :mod:`ps_eor.datacube`.
:class:`SphImageCube` stores HEALPix maps shaped
``(n_freqs, 12 * nside**2)``, while :class:`SphDataCube` stores their complex
spherical-harmonic coefficients as ``(n_freqs, n_modes)`` with matching
``(ell, m)`` coordinates. In a harmonic cube, the transverse spatial
coordinate used by the power-spectrum code is ``u = ell / (2 pi)``.

A Cartesian visibility cube can be reprojected to HEALPix and transformed to
spherical harmonics when a large field of view makes a flat-sky treatment
inappropriate::

    from ps_eor import sphcube

    alm_cube = sphcube.SphDataCube.from_cartcube(
        cart_cube,
        nside=128,
        lmax=600,
    )
    alm_cube.filter_lm(lmin=20, lmax=500)
    alm_cube.save('spherical-visibilities.h5')

The transformation can be reversed for inspection. ``mask_corrected=True``
divides by the recorded sky window where its response is sufficiently large::

    sky_cube = alm_cube.image(
        mask_corrected=True,
        mask_correction_threshold=.1,
    )
    sky_cube.plot(fmhz='med')

For multiple observations on the same harmonic grid,
:class:`SphDataCubeCombiner` forms a weighted mean over their common frequency
channels. Most ``filter_*`` methods modify a spherical cube in place; copy it
first when the unfiltered cube must be retained.
"""



import copy
from collections import defaultdict

import healpy as hp
import matplotlib.pyplot as plt
import numpy as np
import tables
from astropy.coordinates import ICRS
from scipy.interpolate import interp1d
from scipy.signal import get_window

from . import datacube, psutil, sphutil




class SphMetaData:
    """Observation metadata associated with a spherical data cube.

    Metadata may describe several accumulated observations. Per-observation
    pointing, duration, window, and weight fields are stored as lists, while
    the HEALPix geometry, integration time, coordinate system, and frequency
    width are shared.
    """

    def __init__(self, metadata):
        self.data = defaultdict(list, metadata)

    def __add__(self, other):
        assert self.data['INT_TIME'] == other.data['INT_TIME']
        assert self.data['COORD_SYS'] == other.data['COORD_SYS']

        new = self.copy()
        for key in ['MJD_OBS', 'RADEC_OBS', 'TIME_OBS', 'WIN_FCT', 'WIN_FCT_FWHM', 'WEIGHTS']:
            new.data[key].extend(other.get(key))

        return new

    @property
    def nside(self):
        """HEALPix ``nside``."""
        return self.get('NSIDE')

    @property
    def int_time(self):
        """Integration time in seconds."""
        return self.get('INT_TIME')

    @property
    def total_time(self):
        """Stored total time, or the sum of observation durations, in seconds."""
        if 'TOTAL_TIME' in self.data:
            return self.data['TOTAL_TIME']
        return np.sum(self.get('TIME_OBS'))

    @property
    def freq_width(self):
        """Frequency-channel width in Hz."""
        return self.get('FREQ_WIDT')

    @property
    def win_fct(self):
        """Combined :class:`SphWindowFunction` reconstructed from the metadata."""
        return SphWindowFunction.from_meta(self)



    @staticmethod
    def new_obs(mjd_obs, radec_obs, time_obs, int_time, nside, win_fct, fwhm, freq_width,
                coord_sys='ICRS', weight=1):
        """Create metadata for one observation.

        Args:
            mjd_obs: Observation reference time in MJD.
            radec_obs: Pointing ``(RA, Dec)`` in degrees.
            time_obs: Observation duration in seconds.
            int_time: Integration time in seconds.
            nside: HEALPix resolution.
            win_fct: Window name, or ``None`` for a unit mask.
            fwhm: Window full-width at half-maximum in radians.
            freq_width: Frequency-channel width in Hz.
            coord_sys: Astropy coordinate-system identifier.
            weight: Relative observation weight.

        Returns:
            SphMetaData: Metadata containing the new observation.
        """
        data = defaultdict(list)
        data['INT_TIME'] = int_time
        data['NSIDE'] = nside
        data['COORD_SYS'] = coord_sys
        data['FREQ_WIDT'] = freq_width
        data['ORIGIN'] = f'ps_eor v_{datacube.__version__}'
        new = SphMetaData(data)
        new.add_obs(mjd_obs, radec_obs, time_obs, win_fct, fwhm, weight)

        return new




    def add_obs(self, mjd_obs, radec_obs, time_obs, win_fct, fwhm, weight=1):
        """Append one observation.

        Args:
            mjd_obs: Observation reference time in MJD.
            radec_obs: Pointing ``(RA, Dec)`` in degrees.
            time_obs: Observation duration in seconds.
            win_fct: Window name, or ``None``.
            fwhm: Window FWHM in radians.
            weight: Relative observation weight.
        """
        self.data['MJD_OBS'].append(mjd_obs)
        self.data['RADEC_OBS'].append(radec_obs)
        self.data['TIME_OBS'].append(time_obs)
        self.data['WIN_FCT'].append(win_fct)
        self.data['WIN_FCT_FWHM'].append(fwhm)
        self.data['WEIGHTS'].append(weight)




    def iter_obs(self, keys):
        """Yield the requested fields for each observation.

        Args:
            keys: Ordered metadata keys to return.

        Yields:
            list: Values corresponding to ``keys`` for one observation.
        """
        for i in range(len(self.data['MJD_OBS'])):
            yield [self.data[key][i] for key in keys]




    def get(self, key):
        """Return metadata field ``key``."""
        return self.data[key]




    def add_weight(self, weight):
        """Multiply every observation weight by ``weight``."""
        self.data['WEIGHTS'] = (weight * np.array(self.data['WEIGHTS'])).tolist()




    def items(self):
        """Return the metadata key-value pairs."""
        return list(self.data.items())




    def copy(self):
        """Return a deep copy of the metadata."""
        return SphMetaData(copy.deepcopy(self.data))






class NoMask(datacube.Mask):
    """Spherical mask that leaves every pixel unchanged."""

    def __init__(self):
        datacube.Mask.__init__(self, [self])



    def generate(self, meta_data):
        """Generate a unit HEALPix mask at ``meta_data.nside``."""
        return np.ones(hp.nside2npix(meta_data.get('NSIDE')))






class NoPrimaryBeam(datacube.BasePrimaryBeam, NoMask):
    """Spherical primary-beam model with unit response."""

    def __init__(self):
        datacube.BasePrimaryBeam.__init__(self, [self])





class SphWindowFunction(datacube.Mask):
    """Azimuthally symmetric sky window centred on a pointing."""

    def __init__(self, name, ra_dec, fwhm, weight=1, primary_beam=None):
        """Configure a spherical sky window.

        Args:
            name: Window name accepted by :func:`scipy.signal.get_window`.
            ra_dec: Window centre ``(RA, Dec)`` in degrees.
            fwhm: Full-width at half-maximum in radians.
            weight: Multiplicative window weight.
            primary_beam: Optional :class:`~ps_eor.datacube.BasePrimaryBeam`
                multiplied into the window.
        """
        self.name = name
        self.ra_dec = ra_dec
        self.fwhm = fwhm
        self.weight = weight
        self.primary_beam = primary_beam
        datacube.Mask.__init__(self, [self])



    @staticmethod
    def from_meta(meta, primary_beam=None):
        """Reconstruct the combined window recorded in spherical metadata.

        Args:
            meta: :class:`SphMetaData` containing observation windows.
            primary_beam: Optional beam to include in every window.

        Returns:
            ps_eor.datacube.Mask: Sum of the per-observation windows.
        """
        win_fcts = []
        for win_fct, win_fct_fwhm, ra_dec, weight in meta.iter_obs(['WIN_FCT', 'WIN_FCT_FWHM',
                                                                    'RADEC_OBS', 'WEIGHTS']):
            if win_fct is None:
                win_fcts.append(NoMask())
            else:
                name = datacube.WindowFunction.parse_winfct_str(win_fct.split('_')[0])
                win_fcts.append(SphWindowFunction(name, ra_dec, win_fct_fwhm, weight,
                                                  primary_beam=primary_beam))
        return datacube.Mask(win_fcts)




    def generate_window(self, nside, pb_fwhm=None, oversample=2):
        """Generate the HEALPix window map.

        Args:
            nside: HEALPix resolution.
            pb_fwhm: Optional primary-beam FWHM in radians.
            oversample: Radial window-sampling factor.

        Returns:
            ndarray: Window map with ``12 * nside**2`` pixels.
        """
        thetas, _phis = hp.pix2ang(nside, np.arange(hp.nside2npix(nside)))
        n = psutil.get_next_even(oversample * (np.pi / 2.) / hp.nside2resol(nside))
        yp = get_window(self.name, n)[n // 2:]
        xp = np.linspace(0, 0.5 * self.fwhm, len(yp))
        win_map = interp1d(xp, yp, fill_value=0, bounds_error=False)(thetas)
        if self.primary_beam is not None and pb_fwhm is not None:
            pb_map = psutil.get_beam(thetas, self.primary_beam.beam_type, pb_fwhm, n_sidelobe=None)
            win_map = win_map * pb_map

        ra, dec = self.ra_dec
        win_map = hp.Rotator(rot=(180 + ra, -dec + 90), inv=True).rotate_map_pixel(win_map)

        return self.weight * win_map




    def generate(self, meta_data, freq=None):
        """Generate the window for the supplied metadata and frequency.

        Args:
            meta_data: Spherical metadata defining ``nside``.
            freq: Frequency in Hz, used for a frequency-dependent primary beam.

        Returns:
            ndarray: HEALPix window map.
        """
        pb_fwhm = self.primary_beam.get_fwhm(freq=freq) if self.primary_beam is not None else None
        return self.generate_window(meta_data.get('NSIDE'), pb_fwhm=pb_fwhm)






class SphDataCube(datacube.DataCube):
    """Frequency cube of spherical-harmonic coefficients.

    The data shape is ``(n_freqs, n_modes)``. Each column is identified by the
    corresponding entries of ``ll`` and ``mm``.
    """

    def __init__(self, alm_cube, ll, mm, freqs, meta, weights=None):
        """Create a spherical-harmonic cube.

        Args:
            alm_cube: Complex coefficients shaped ``(n_freqs, n_modes)``.
            ll, mm: Harmonic coordinates of length ``n_modes``.
            freqs: Frequencies in Hz.
            meta: :class:`SphMetaData` describing the maps and observations.
            weights: Optional aligned weight cube.
        """
        self.ll = ll
        self.mm = mm
        self.ru = self.ll / (2 * np.pi)
        self.meta = meta

        datacube.DataCube.__init__(self, alm_cube, freqs, weights=weights)



    def get_unique_xy(self):
        """Return per-mode keys used to match ``(ell, m)`` between cubes."""
        return self.ll + 1e-6 * self.mm




    @staticmethod
    def from_cartcube(cube, nside, lmax, reproject_order='bilinear'):
        """Reproject a Cartesian cube and transform it to spherical harmonics.

        The Cartesian cube is imaged, reprojected onto an ICRS HEALPix grid,
        and transformed independently at every frequency. Pixels outside the
        Cartesian footprint are set to zero.

        Args:
            cube: :class:`~ps_eor.datacube.CartDataCube` to transform.
            nside: Output HEALPix resolution.
            lmax: Maximum spherical-harmonic degree.
            reproject_order: Interpolation order accepted by
                :func:`reproject.reproject_to_healpix`.

        Returns:
            SphDataCube: Harmonic cube with observation metadata derived from
            the Cartesian input.
        """
        import reproject

        fwcs = cube.meta.wcs.copy()
        fwcs = fwcs.dropaxis(2)
        fwcs = fwcs.dropaxis(2)

        img_cube = cube.regrid().image().data.real
        hmaps = []
        pr = psutil.progress_report(img_cube.shape[0])
        for i, img in enumerate(img_cube):
            pr(i)
            hmaps.append(reproject.reproject_to_healpix((img, fwcs), ICRS(), nside=nside,
                                                        order=reproject_order)[0])
        hmaps = np.array(hmaps)
        hmaps[np.isnan(hmaps)] = 0

        alms = np.array([hp.map2alm(hmap, int(lmax)) for hmap in hmaps])
        ll, mm = sphutil.get_lm(int(lmax))

        meta = SphMetaData.new_obs(
                cube.meta.to_header()['MJD-OBS'], fwcs.wcs.crval, cube.meta.total_time,
                cube.meta.int_time, nside, cube.meta.get('PEWINFCT', 'boxcar'), cube.meta.theta_fov,
                cube.meta.freq_width)

        return SphDataCube(alms, ll, mm, cube.freqs, meta)




    def new_with_data(self, data, weights=None, freqs=None):
        """Create a cube on the same harmonic grid with new data.

        Args:
            data: Replacement array shaped ``(n_freqs, n_modes)``.
            weights: Optional replacement weights.
            freqs: Optional replacement frequencies in Hz.

        Returns:
            SphDataCube: The new cube.
        """
        if freqs is None:
            freqs = self.freqs

        return SphDataCube(data, self.ll, self.mm, freqs, self.meta, weights=weights)




    def filter_outliers(self, outliers):
        """Remove frequency channels selected by ``outliers`` in place."""
        self.freqs = self.freqs[~outliers]
        self.data = self.data[~outliers]




    def filter_uv_from_index(self, idx_uv):
        """Keep the spatial modes selected by ``idx_uv`` in place."""
        datacube.DataCube.filter_uv_from_index(self, idx_uv)
        self.mm = self.mm[idx_uv]
        self.ll = self.ll[idx_uv]
        self.ru = self.ru[idx_uv]




    def filter_uvrange(self, umin, umax):
        """Keep ``ell / (2 pi)`` within ``[umin, umax]`` in place."""
        idx_uv = (self.ru >= umin) & (self.ru <= umax)
        self.filter_uv_from_index(idx_uv)




    def filter_lm(self, lmin, lmax):
        """Keep modes with ``ell`` within ``[lmin, lmax]`` in place."""
        idx_uv = (self.ll >= lmin) & (self.ll <= lmax)
        self.filter_uv_from_index(idx_uv)




    def filter_m_theta_max(self, theta_max):
        """Keep azimuthal modes supported within ``theta_max`` in place.

        Args:
            theta_max: Maximum angular distance in radians.
        """
        idx_uv = self.mm < np.clip(np.sin(theta_max) * self.ll, 1, max(self.ll))
        self.filter_uv_from_index(idx_uv)




    @staticmethod
    def load(filename):
        """Load a spherical-harmonic cube from HDF5.

        Args:
            filename: Input HDF5 filename.

        Returns:
            SphDataCube: The restored coefficients and metadata.
        """
        with tables.open_file(filename, 'r') as h5_file:
            alm_cube = h5_file.root.alm_cube.data.read()
            freqs = h5_file.root.alm_cube.freqs.read()
            ll = h5_file.root.alm_cube.ll.read()
            mm = h5_file.root.alm_cube.mm.read()

            attrs = h5_file.root.alm_cube.data.attrs
            meta = SphMetaData({k: psutil.safe_decode_bytes(attrs[k])
                                     for k in attrs._f_list() if k[0].isupper()})

            # Older files did not store channel width.
            if 'FREQ_WIDT' not in meta.data:
                meta.data = psutil.robust_freq_width(freqs)

        return SphDataCube(alm_cube, ll, mm, freqs, meta)




    def save(self, filename):
        """Save the spherical-harmonic cube to HDF5.

        Args:
            filename: Destination HDF5 filename.
        """
        with tables.open_file(filename, 'w') as h5_file:
            group = h5_file.create_group("/", 'alm_cube', 'Alm cube (n_freqs, n_modes')
            h5_file.create_array(group, 'data', self.data, "Spherical harmonics (K)")
            h5_file.create_array(group, 'freqs', self.freqs, "Frequencies (Hz)")
            h5_file.create_array(group, 'll', self.ll, "l mode")
            h5_file.create_array(group, 'mm', self.mm, "m mode")

            for key, value in self.meta.items():
                h5_file.root.alm_cube.data.attrs[key] = value




    def plot_lm(self, fmhz='med', action_fct=None, ax=None, title=None, **kargs):
        """Plot real coefficients on the ``(ell, m)`` plane.

        Args:
            fmhz: Frequency in MHz, or ``'med'`` or ``'first'``. Ignored when
                ``action_fct`` is supplied.
            action_fct: Optional reduction called as
                ``action_fct(data, axis=0)``.
            ax: Matplotlib axes; a new one is created by default.
            title: Optional axes title.
            **kargs: Arguments passed to :meth:`matplotlib.axes.Axes.imshow`.
        """
        if ax is None:
            _fig, ax = plt.subplots()

        alm = self.data

        if action_fct is None:
            if fmhz == 'med':
                fmhz = self.freqs[0] + (self.freqs[-1] - self.freqs[0]) / 2.
            elif fmhz == 'first':
                fmhz = self.freqs[0]
            elif fmhz == 'first':
                fmhz = self.freqs[-1]

            i = np.argmin(abs(self.freqs - fmhz * 1e6))
            alm = alm[i]
        else:
            alm = action_fct(alm, axis=0)

        alm_map = sphutil.get_lm_map(alm, self.ll, self.mm)

        cbs = psutil.ColorbarSetting(psutil.ColorbarInnerPosition(location=2, height="80%", pad=1))

        extent = (min(self.ll), max(self.ll), min(self.mm), max(self.mm))

        im_mappable = ax.imshow(alm_map.real, extent=extent, aspect='auto', **kargs)
        cbs.add_colorbar(im_mappable, ax)

        ax.set_xlabel('l')
        ax.set_ylabel('m')

        if title is not None:
            ax.set_title(title)




    def regrid(self):
        """Return this cube; spherical coefficients need no UV regridding."""
        return self




    def image(self, mask_corrected=False, mask_correction_threshold=0.1):
        """Transform the coefficients to a HEALPix image cube.

        Args:
            mask_corrected: Divide by the recorded window where it is reliable.
            mask_correction_threshold: Correct pixels above this fraction of
                the peak window response and set the others to zero.

        Returns:
            SphImageCube: Frequency-dependent HEALPix maps.
        """
        data = np.asarray(self.data, order='C')
        hmaps = np.array([sphutil.fast_alm2map(k, self.ll, self.mm, self.meta.nside) for k in data])

        if mask_corrected:
            mask = self.meta.win_fct.generate(self.meta)
            th = mask_correction_threshold
            hmaps[:, mask > th * mask.max()] = hmaps[:, mask > th * mask.max()] / mask[None,
                    mask > th * mask.max()]
            hmaps[:, mask < th * mask.max()] = 0

        return SphImageCube(hmaps, self.freqs, self.meta)






class SphImageCube(datacube.ImageCube):
    """Frequency cube of HEALPix sky maps.

    The data shape is ``(n_freqs, 12 * nside**2)``.
    """

    def __init__(self, hmaps, freqs, meta):
        """Create a HEALPix image cube.

        Args:
            hmaps: HEALPix maps shaped ``(n_freqs, n_pixels)``.
            freqs: Frequencies in Hz.
            meta: :class:`SphMetaData` with the matching ``nside``.
        """
        datacube.ImageCube.__init__(self, hmaps, freqs, meta)



    def apply_window_function(self, win_fct, add_to_meta=True):
        """Apply a sky window in place.

        Args:
            win_fct: :class:`SphWindowFunction` to multiply into every map.
            add_to_meta: Record the window as another observation in the
                metadata.
        """
        win_mask = win_fct.generate(self.meta)
        self.data = self.data * win_mask
        if add_to_meta:
            if self.meta.data['WIN_FCT'] == [None] * len(self.meta.data['WIN_FCT']):
                self.meta = SphMetaData.new_obs(
                        0, win_fct.ra_dec, 0, self.meta.int_time, self.meta.nside, win_fct.name,
                        win_fct.fwhm, self.meta.freq_width, weight=1)
            else:
                self.meta.add_obs(0, win_fct.ra_dec, 0, win_fct.name, win_fct.fwhm, weight=1)




    def ft(self, umin, umax):
        """Transform the maps to spherical harmonics over a UV range.

        The transform uses ``lmax = int(2 pi umax)`` and then retains modes
        whose ``ell / (2 pi)`` lies within the requested interval.

        Args:
            umin, umax: Spatial-frequency limits in wavelengths.

        Returns:
            SphDataCube: The selected harmonic coefficients.
        """
        lmax = int(2 * np.pi * umax)
        alms = np.array([hp.map2alm(k, lmax) for k in self.data], order='C')
        ll, mm = sphutil.get_lm(lmax)

        cube = SphDataCube(alms, ll, mm, self.freqs, self.meta)
        cube.filter_uvrange(umin, umax)

        return cube




    def plot(self, fmhz='med', action_fct=None, dpar=None, dmer=None, title='',
             vmax=None, vmin=None, ax=None, auto_scale_quantiles=None, coord='CG', **kargs):
        """Plot one frequency map or a reduction over frequency.

        Args:
            fmhz: Frequency in MHz, or ``'med'`` or ``'first'``. Ignored when
                ``action_fct`` is supplied.
            action_fct: Optional reduction called as
                ``action_fct(data, axis=0)``.
            dpar, dmer: HEALPix graticule spacing.
            title: Plot title.
            vmax, vmin: Explicit color limits.
            ax: Matplotlib axes; a new one is created by default.
            auto_scale_quantiles: Optional ``(low, high)`` quantiles used as
                color limits.
            coord: HEALPix coordinate conversion passed to
                :func:`healpy.mollview`.
            **kargs: Additional arguments passed to :func:`healpy.mollview`.
        """
        if ax is None:
            _fig, ax = plt.subplots()

        d = self.data

        if action_fct is None:
            if fmhz == 'med':
                fmhz = self.freqs[0] + (self.freqs[-1] - self.freqs[0]) / 2.
            elif fmhz == 'first':
                fmhz = self.freqs[0]
            elif fmhz == 'first':
                fmhz = self.freqs[-1]

            i = np.argmin(abs(self.freqs - fmhz * 1e6))
            d = d[i]
        else:
            d = action_fct(d, axis=0)

        if vmin is not None:
            kargs['min'] = vmin
        if vmax is not None:
            kargs['max'] = vmax

        if auto_scale_quantiles is not None:
            kargs['min'] = np.quantile(d, auto_scale_quantiles[0])
            kargs['max'] = np.quantile(d, auto_scale_quantiles[1])

        plt.sca(ax)
        hp.mollview(d, hold=True, title=title, coord=coord, **kargs)
        hp.graticule(dpar=dpar, dmer=dmer, coord='C', verbose=False)
        hp.graticule(dpar=dpar, dmer=dmer, coord='G', alpha=0.5, c=psutil.blue, verbose=False)




    def plot_slice(self, ax=None, min_dec=-20, npix=1000):
        """Plot a declination-frequency slice through the pointing centre.

        Args:
            ax: Matplotlib axes; a new one is created by default.
            min_dec: Lowest declination in degrees.
            npix: Number of samples along each half of the great-circle slice.
        """
        if ax is None:
            _fig, ax = plt.subplots()

        ra, dec = np.radians(np.array(self.meta.data['RADEC_OBS']).mean(axis=0))
        if ra < 0:
            ra = 2 * np.pi + ra

        decs = np.radians(np.linspace(min_dec, 90, npix))

        ras = ra * np.ones_like(decs)
        thetas, phis = sphutil.radec2thetaphi(ras, decs)
        idx1 = hp.ang2pix(128, thetas, phis)

        ras = (ra + np.pi) * np.ones_like(decs)
        thetas, phis = sphutil.radec2thetaphi(ras, decs)
        idx2 = hp.ang2pix(128, thetas, phis)[::-1]

        extent = [min_dec, 180 - min_dec, self.freqs[0] * 1e-6, self.freqs[-1] * 1e-6]

        im = ax.imshow(np.hstack([self.data[:, idx1], self.data[:, idx2]]), aspect='auto',
                       extent=extent)
        ax.axvline(np.degrees(dec), c=psutil.black, ls='--')
        ax.set_ylabel('Freqs (MHz)')
        ax.set_xlabel(f'DEC (RA = {np.degrees(ra):.1f} deg) [deg]')

        cbs = psutil.ColorbarSetting(psutil.ColorbarOutterPosition(width='3%'))
        cbs.add_colorbar(im, ax)




    def copy(self):
        """Return a deep copy of the map data and metadata."""
        return SphImageCube(self.data.copy(), self.freqs, self.meta.copy())






class SphDataCubeCombiner:
    """Accumulate weighted spherical cubes on their common frequencies."""

    def __init__(self):
        self.data = None
        self.freqs = None
        self.weights = None
        self.ll = None
        self.mm = None
        self.meta = None



    def get_inter_idx(self, freqs1, freqs2):
        """Return masks selecting the common frequencies of two arrays."""
        f1 = datacube._fmhz(freqs1)
        f2 = datacube._fmhz(freqs2)
        freqs = np.intersect1d(f1, f2)
        return np.isin(f1, freqs), np.isin(f2, freqs)




    def add(self, cube, weight):
        """Add a weighted cube to the accumulator.

        Only channels common to all cubes accumulated so far are retained.
        Cubes are expected to share the same ``(ell, m)`` grid.

        Args:
            cube: :class:`SphDataCube` to accumulate.
            weight: Scalar weight applied to the complete observation.
        """
        if self.data is None:
            self.data = weight * cube.data
            self.freqs = cube.freqs.copy()
            self.weights = [weight]
            self.meta = cube.meta.copy()
            self.meta.add_weight(weight)
            self.ll = cube.ll.copy()
            self.mm = cube.mm.copy()
        else:
            idx1, idx2 = self.get_inter_idx(self.freqs, cube.freqs)
            self.data = self.data[idx1]
            self.freqs = self.freqs[idx1]
            self.data += weight * cube.data[idx2]
            self.weights.append(weight)

            cube_meta = cube.meta.copy()
            cube_meta.add_weight(weight)
            self.meta = self.meta + cube_meta




    def get(self):
        """Return the weighted mean of the accumulated cubes.

        Returns:
            SphDataCube: Combined data on the common frequency and harmonic
            grid.
        """

        sum_weight = float(np.sum(self.weights))
        data = self.data.copy() / sum_weight
        meta = self.meta.copy()
        meta.add_weight(1 / sum_weight)

        return SphDataCube(data, self.ll, self.mm, self.freqs.copy(), meta)