Source code for africanus.model.spectral.dask

# -*- coding: utf-8 -*-


from africanus.model.spectral.spec_model import (
    spectral_model as np_spectral_model,
    SPECTRAL_MODEL_DOC,
)
from africanus.util.requirements import requires_optional

try:
    import dask.array as da
except ImportError as e:
    opt_import_error = e
else:
    opt_import_error = None


def spectral_model_wrapper(stokes, spi, ref_freq, frequencies, base=None):
    return np_spectral_model(stokes, spi[0], ref_freq, frequencies, base=base)


[docs] @requires_optional("dask.array", opt_import_error) def spectral_model(stokes, spi, ref_freq, frequencies, base=0): if len(spi.chunks[1]) != 1: raise ValueError("Chunking along the spi dimension unsupported") pol_dim = () if stokes.ndim == 1 else ("pol",) return da.blockwise( spectral_model_wrapper, ( "source", "chan", ) + pol_dim, stokes, ("source",) + pol_dim, spi, ("source", "spi") + pol_dim, ref_freq, ("source",), frequencies, ("chan",), base=base, dtype=stokes.dtype, )
try: spectral_model.__doc__ = SPECTRAL_MODEL_DOC.substitute( array_type=":class:`dask.array.Array`" ) except AttributeError: pass