<commit_before>from __future__ import print_function, absolute_import, division

import six
import warnings
import tempfile
from astropy.io import fits
from astropy.wcs import WCS
from astropy import units as u
import numpy as np
from radio_beam import Beam, Beams

from .. import SpectralCube, StokesSpectralCube, BooleanArrayMask, LazyMask, VaryingResolutionSpectralCube
from .. import cube_utils
from .. utils import BeamWarning

# Read and write from a CASA image. This has a few
# complications. First, by default CASA does not return the
# "python order" and so we either have to transpose the cube on
# read or have dueling conventions. Second, CASA often has
# degenerate stokes axes present in unpredictable places (3rd or
# 4th without a clear expectation). We need to replicate these
# when writing but don't want them in memory. By default, try to
# yield the same array in memory that we would get from astropy.


def is_casa_image(input, **kwargs):
    if isinstance(input, six.string_types):
        if input.endswith('.image'):
            return True
    return False


def wcs_casa2astropy(ia, coordsys):
    """
    Convert a casac.coordsys object into an astropy.wcs.WCS object
    """

    # Rather than try and parse the CASA coordsys ourselves, we delegate
    # to CASA by getting it to write out a FITS file and reading back in
    # using WCS

    from casatasks import exportfits

    tmpimagefile = tempfile.mktemp() + '.image'
    tmpfitsfile = tempfile.mktemp() + '.fits'
    ia.newimagefromarray(outfile=tmpimagefile,
                         pixels=np.ones([1] * coordsys.naxes()),
                         csys=coordsys.torecord(), log=False)
    exportfits(tmpimagefile, tmpfitsfile, stokeslast=False)

    return WCS(tmpfitsfile)


def load_casa_image(filename, skipdata=False,
                    skipvalid=False, skipcs=False, **kwargs):
    """
    Load a cube (into memory?) from a CASA image. By default it will transpose
    the cube into a 'python' order and drop degenerate axes. These options can
    be suppressed. The object holds the coordsys object from the image in
    memory.
    """

    try:
        import casatools
        ia = casatools.image()
    except ImportError:
        try:
            from taskinit import iatool
            ia = iatool()
        except ImportError:
            raise ImportError("Could not import CASA (casac) and therefore cannot read CASA .image files")

    # use the ia tool to get the file contents
    try:
        ia.open(filename)
    except AssertionError as ex:
        if 'must be of cReqPath type' in str(ex):
            raise IOError("File {0} not found.  Error was: {1}"
                          .format(filename, str(ex)))
        else:
            raise ex

    # read in the data
    if not skipdata:
        # CASA data are apparently transposed.
        data = ia.getchunk().transpose()

    # CASA stores validity of data as a mask
    if not skipvalid:
        valid = ia.getchunk(getmask=True).transpose()

    # transpose is dealt with within the cube object

    # read in coordinate system object
    casa_cs = ia.coordsys()

    wcs = wcs_casa2astropy(ia, casa_cs)

    unit = ia.brightnessunit()
<commit_msg>create array-like reader for CASA images so we can retrieve just the chunks we need<commit_after>from __future__ import print_function, absolute_import, division

import six
import warnings
import tempfile
from astropy.io import fits
from astropy.wcs import WCS
from astropy import units as u
from astropy.wcs.wcsapi.sliced_low_level_wcs import sanitize_slices
import numpy as np
from radio_beam import Beam, Beams

from .. import SpectralCube, StokesSpectralCube, BooleanArrayMask, LazyMask, VaryingResolutionSpectralCube
from .. import cube_utils
from .. utils import BeamWarning, cached

# Read and write from a CASA image. This has a few
# complications. First, by default CASA does not return the
# "python order" and so we either have to transpose the cube on
# read or have dueling conventions. Second, CASA often has
# degenerate stokes axes present in unpredictable places (3rd or
# 4th without a clear expectation). We need to replicate these
# when writing but don't want them in memory. By default, try to
# yield the same array in memory that we would get from astropy.


def is_casa_image(input, **kwargs):
    if isinstance(input, six.string_types):
        if input.endswith('.image'):
            return True
    return False


def wcs_casa2astropy(ia, coordsys):
    """
    Convert a casac.coordsys object into an astropy.wcs.WCS object
    """

    # Rather than try and parse the CASA coordsys ourselves, we delegate
    # to CASA by getting it to write out a FITS file and reading back in
    # using WCS

    from casatasks import exportfits

    tmpimagefile = tempfile.mktemp() + '.image'
    tmpfitsfile = tempfile.mktemp() + '.fits'
    ia.newimagefromarray(outfile=tmpimagefile,
                         pixels=np.ones([1] * coordsys.naxes()),
                         csys=coordsys.torecord(), log=False)
    exportfits(tmpimagefile, tmpfitsfile, stokeslast=False)

    return WCS(tmpfitsfile)


class ArraylikeCasaData:

    def __init__(self, filename, ia_kwargs={}):

        try:
            import casatools
            self.ia = casatools.image()
        except ImportError:
            try:
                from taskinit import iatool
                self.ia = iatool()
            except ImportError:
                raise ImportError("Could not import CASA (casac) and therefore cannot read CASA .image files")

        # use the ia tool to get the file contents
        try:
            self.ia.open(filename)
        except AssertionError as ex:
            if 'must be of cReqPath type' in str(ex):
                raise IOError("File {0} not found.  Error was: {1}"
                              .format(filename, str(ex)))
            else:
                raise ex

        self.ia_kwargs = ia_kwargs

        self._cache = {}

    @property
    @cached
    def shape(self):
        return tuple(self.ia.shape())

    @property
    @cached
    def ndim(self):
        return len(self.shape)

    @property
    @cached
    def dtype(self):
        return np.dtype(self.ia.pixeltype())


    def __getitem__(self, value):

        value = sanitize_slices(value, self.ndim)

        blc = [(slc.start or -1) for slc in value]
        trc = [(slc.stop or -1) for slc in value]
        inc = [(slc.step or 1) for slc in value]

        data = self.ia.getchunk(blc=blc, trc=trc, inc=inc, **self.ia_kwargs)

        return data.transpose()


def load_casa_image(filename, skipdata=False,
                    skipvalid=False, skipcs=False, **kwargs):
    """
    Load a cube (into memory?) from a CASA image. By default it will transpose
    the cube into a 'python' order and drop degenerate axes. These options can
    be suppressed. The object holds the coordsys object from the image in
    memory.
    """

    try:
        import casatools
        ia = casatools.image()
    except ImportError:
        try:
            from taskinit import iatool
            ia = iatool()
        except ImportError:
            raise ImportError("Could not import CASA (casac) and therefore cannot read CASA .image files")

    # use the ia tool to get the file contents
    try:
        ia.open(filename)
    except AssertionError as ex:
        if 'must be of cReqPath type' in str(ex):
            raise IOError("File {0} not found.  Error was: {1}"
                          .format(filename, str(ex)))
        else:
            raise ex

    # read in the data
    if not skipdata:
        # CASA data are apparently transposed.
        data = ArraylikeCasaData(filename)

    # CASA stores validity of data as a mask
    if not skipvalid:
        valid = ArraylikeCasaData(filename, ia_kwargs={'getmask': True})

    # transpose is dealt with within the cube object

    # read in coordinate system object
    casa_cs = ia.coordsys()

    wcs = wcs_casa2astropy(ia, casa_cs)

    unit = ia.brightnessunit()

{'fraction_non_alphanumeric': 0.0744324525493115, 'fraction_numerical': 0.002357027664061531, 'mean_word_length': 2.5313184406482696, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6118367', 'n_tokens_mistral': 2452, 'n_tokens_neox': 2298, 'n_words': 877}