# Path: katdal/chunkstore.py
# class PlaceholderChunk:
#     """Chunk returned to indicate missing data."""
# 
#     def __init__(self, shape, dtype):
#         self.shape = shape
#         self.dtype = np.dtype(dtype)
# 
#     def __getitem__(self, index):
#         # Create an array with a zero-sized dtype, so that it takes no memory
#         dummy = np.empty(self.shape, dtype=[])
#         new_shape = dummy[index].shape
#         return PlaceholderChunk(new_shape, self.dtype)
# 
# Path: katdal/flags.py
# DATA_LOST = 1 << DATA_LOST_BIT
# 
# Path: katdal/van_vleck.py
# def autocorr_lookup_table(levels, size=4000):
#     """Lookup table that corrects complex autocorrelation quantisation effects.
# 
#     This maps the variance of a quantised complex voltage signal to the variance
#     of the unquantised signal under the assumption that the signal is proper
#     (circularly-symmetric) complex normally distributed.
# 
#     Parameters
#     ----------
#     levels : sequence of float
#         Quantisation levels for real and imaginary components of voltage signal
#     size : int, optional
#         Size of lookup table
# 
#     Returns
#     -------
#     quantised_autocorr_table, true_autocorr_table : array of float, shape (`size`,)
#         Lookup table associating quantised autocorrelations and unquantised
#         autocorrelations (i.e. power/variance of complex signals)
#     """
#     # Terminology:
#     # x = Proper complex normal voltage signal (zero-mean)
#     # rxx = Power (variance) *per* real/imag component of unquantised / true x
#     # sxx = Power (variance) *per* real/imag component of quantised x
#     abs_levels = np.abs(levels)
#     sxx_min_nonzero = abs_levels[abs_levels > 0].min() ** 2
#     sxx_max = abs_levels.max() ** 2
#     # Sweep across range of true power values, placing more table entries at tricky lower end
#     rxx_grid = np.r_[np.logspace(-2.4, 0, size // 2, endpoint=False),
#                      np.logspace(0, np.log10(sxx_max / sxx_min_nonzero) + 8, size - 2 - size // 2)]
#     # Shift the table to place inflection point at minimum non-zero sxx
#     rxx_grid *= sxx_min_nonzero
#     # Map true power to expected quantised power
#     sxx_mean = _squared_quant_norm0_mean(levels, rxx_grid)
#     # Extend quantised power values to its maximum range
#     sxx_table = np.r_[0., sxx_mean, sxx_max]
#     # Replace asymptotic with linear decay at bottom end, and clip unbounded growth at top end
#     rxx_table = np.r_[0., rxx_grid, rxx_grid[-1]]
#     # The factor 2 converts power per real/imag component to power/variance of complex signal
#     return 2. * sxx_table, 2. * rxx_table


import itertools
import logging
import dask.array as da
import numba
import numpy as np
import toolz
from dask.array.rechunk import intersect_chunks
from dask.highlevelgraph import HighLevelGraph
from .chunkstore import PlaceholderChunk
from .flags import DATA_LOST
from .van_vleck import autocorr_lookup_table
    """Correct autocorrelations for quantisation effects (Van Vleck correction).

    This is a first-order correction that only adjusts the mean autocorrelations,
    which in turn affects the autocorrelation and crosscorrelation weights.
    A complete correction would also adjust the mean crosscorrelations, and
    further improve the weight estimates based on Bayesian statistics.

    Parameters
    ----------
    vis : :class:`dask.array.Array` of complex64, shape (*T*, *F*, *B*)
        Complex visibility data as function of time, frequency, correlation product
    corrprods : sequence of 2-tuples or ndarray, containing str
        Input labels of the correlation products, used to find autocorrelations
    levels : sequence of float, optional
        Quantisation levels of real/imag components of complex digital signal
        entering correlator (defaults to MeerKAT F-engine output levels)

    Returns
    -------
    corrected_vis : :class:`dask.array.Array` of complex64, shape (*T*, *F*, *B*)
        Complex visibility data with autocorrelations corrected for quantisation
    """
    assert len(corrprods) == vis.shape[2]
    # Ensure that we have only a single chunk on the baseline axis.
    if len(vis.chunks[2]) > 1:
        vis = vis.rechunk({2: vis.shape[2]})
    auto_indices, _, _ = corrprod_to_autocorr(corrprods)
    if levels is None:
        # 255-level "8-bit" output of MeerKAT F-engine requantiser
        levels = np.arange(-127., 128.)

# Path: katdal/chunkstore.py
# class PlaceholderChunk:
#     """Chunk returned to indicate missing data."""
# 
#     def __init__(self, shape, dtype):
#         self.shape = shape
#         self.dtype = np.dtype(dtype)
# 
#     def __getitem__(self, index):
#         # Create an array with a zero-sized dtype, so that it takes no memory
#         dummy = np.empty(self.shape, dtype=[])
#         new_shape = dummy[index].shape
#         return PlaceholderChunk(new_shape, self.dtype)
# 
# Path: katdal/flags.py
# DATA_LOST = 1 << DATA_LOST_BIT
# 
# Path: katdal/van_vleck.py
# def autocorr_lookup_table(levels, size=4000):
#     """Lookup table that corrects complex autocorrelation quantisation effects.
# 
#     This maps the variance of a quantised complex voltage signal to the variance
#     of the unquantised signal under the assumption that the signal is proper
#     (circularly-symmetric) complex normally distributed.
# 
#     Parameters
#     ----------
#     levels : sequence of float
#         Quantisation levels for real and imaginary components of voltage signal
#     size : int, optional
#         Size of lookup table
# 
#     Returns
#     -------
#     quantised_autocorr_table, true_autocorr_table : array of float, shape (`size`,)
#         Lookup table associating quantised autocorrelations and unquantised
#         autocorrelations (i.e. power/variance of complex signals)
#     """
#     # Terminology:
#     # x = Proper complex normal voltage signal (zero-mean)
#     # rxx = Power (variance) *per* real/imag component of unquantised / true x
#     # sxx = Power (variance) *per* real/imag component of quantised x
#     abs_levels = np.abs(levels)
#     sxx_min_nonzero = abs_levels[abs_levels > 0].min() ** 2
#     sxx_max = abs_levels.max() ** 2
#     # Sweep across range of true power values, placing more table entries at tricky lower end
#     rxx_grid = np.r_[np.logspace(-2.4, 0, size // 2, endpoint=False),
#                      np.logspace(0, np.log10(sxx_max / sxx_min_nonzero) + 8, size - 2 - size // 2)]
#     # Shift the table to place inflection point at minimum non-zero sxx
#     rxx_grid *= sxx_min_nonzero
#     # Map true power to expected quantised power
#     sxx_mean = _squared_quant_norm0_mean(levels, rxx_grid)
#     # Extend quantised power values to its maximum range
#     sxx_table = np.r_[0., sxx_mean, sxx_max]
#     # Replace asymptotic with linear decay at bottom end, and clip unbounded growth at top end
#     rxx_table = np.r_[0., rxx_grid, rxx_grid[-1]]
#     # The factor 2 converts power per real/imag component to power/variance of complex signal
#     return 2. * sxx_table, 2. * rxx_table


# Path: katdal/vis_flags_weights.py
import itertools
import logging
import dask.array as da
import numba
import numpy as np
import toolz
from dask.array.rechunk import intersect_chunks
from dask.highlevelgraph import HighLevelGraph
from .chunkstore import PlaceholderChunk
from .flags import DATA_LOST
from .van_vleck import autocorr_lookup_table

    """Correct autocorrelations for quantisation effects (Van Vleck correction).

    This is a first-order correction that only adjusts the mean autocorrelations,
    which in turn affects the autocorrelation and crosscorrelation weights.
    A complete correction would also adjust the mean crosscorrelations, and
    further improve the weight estimates based on Bayesian statistics.

    Parameters
    ----------
    vis : :class:`dask.array.Array` of complex64, shape (*T*, *F*, *B*)
        Complex visibility data as function of time, frequency, correlation product
    corrprods : sequence of 2-tuples or ndarray, containing str
        Input labels of the correlation products, used to find autocorrelations
    levels : sequence of float, optional
        Quantisation levels of real/imag components of complex digital signal
        entering correlator (defaults to MeerKAT F-engine output levels)

    Returns
    -------
    corrected_vis : :class:`dask.array.Array` of complex64, shape (*T*, *F*, *B*)
        Complex visibility data with autocorrelations corrected for quantisation
    """
    assert len(corrprods) == vis.shape[2]
    # Ensure that we have only a single chunk on the baseline axis.
    if len(vis.chunks[2]) > 1:
        vis = vis.rechunk({2: vis.shape[2]})
    auto_indices, _, _ = corrprod_to_autocorr(corrprods)
    if levels is None:
        # 255-level "8-bit" output of MeerKAT F-engine requantiser
        levels = np.arange(-127., 128.)

    quantised_autocorr_table, true_autocorr_table = autocorr_lookup_table(levels)