<commit_before># Numpy + Dask implementation of CASA image data access

from __future__ import print_function, absolute_import, division

import os
from math import ceil, floor
import uuid
import numpy as np

import dask.array

from .casa_low_level_io import getdminfo

__all__ = ['casa_image_dask_reader']


class MemmapWrapper:
    """
    A wrapper class for dask that opens memmap only when __getitem__ is called,
    which prevents issues with too many open files if opening a memmap for all
    chunks at the same time.
    """

    def __init__(self, filename, **kwargs):
        self._filename = filename
        self._kwargs = kwargs
        self.shape = kwargs['shape'][::-1]
        self.dtype = kwargs['dtype']
        self.ndim = len(self.shape)

    def __getitem__(self, item):
        # We open a file manually and return an in-memory copy of the array
        # otherwise the file doesn't get closed properly.
        with open(self._filename) as f:
            return np.memmap(f, mode='readonly', order='F', **self._kwargs).T[item].copy()


class MaskWrapper:
    """
    A wrapper class for dask that opens binary masks from file and returns
    a chunk from it on-the-fly.
    """

    def __init__(self, filename, offset, count, shape):
        self._filename = filename
        self._offset = offset
        self._count = count
        self.shape = shape[::-1]
        self.dtype = np.bool_
        self.ndim = len(self.shape)

    def __getitem__(self, item):
        # The offset is in the final bit array - but fromfile needs to operate
        # by reading in uint8, so we need to make sure we align what we read
        # in to the bytes.
        start = floor(self._offset / 8)
        end = ceil((self._offset + self._count) / 8)
        array_uint8 = np.fromfile(self._filename, dtype=np.uint8,
                                  offset=start, count=end - start)
        array_bits = np.unpackbits(array_uint8, bitorder='little')
        chunk = array_bits[self._offset - start * 8:self._offset + self._count - start * 8]
        return chunk.astype(np.bool_).reshape(self.shape[::-1], order='F').T[item]


def from_array_fast(arrays, asarray=False, lock=False):
    """
    This is a more efficient alternative to doing::

        [dask.array.from_array(array) for array in arrays]

    that avoids a lot of the overhead in from_array by using the Array
    initializer directly.
    """
    slices = tuple(slice(0, size) for size in arrays[0].shape)
    chunk = tuple((size,) for size in arrays[0].shape)
    meta = np.zeros((0,), dtype=arrays[0].dtype)
    dask_arrays = []
    for array in arrays:
        name1 = str(uuid.uuid4())
        name2 = str(uuid.uuid4())
        dsk = {(name1,) + (0,) * array.ndim: (dask.array.core.getter, name2,
                                              slices, asarray, lock),
               name2: array}
        dask_arrays.append(dask.array.Array(dsk, name1, chunk, meta=meta, dtype=array.dtype))
    return dask_arrays<commit_msg>Convert type after slicing for performance<commit_after># Numpy + Dask implementation of CASA image data access

from __future__ import print_function, absolute_import, division

import os
from math import ceil, floor
import uuid
import numpy as np

import dask.array

from .casa_low_level_io import getdminfo

__all__ = ['casa_image_dask_reader']


class MemmapWrapper:
    """
    A wrapper class for dask that opens memmap only when __getitem__ is called,
    which prevents issues with too many open files if opening a memmap for all
    chunks at the same time.
    """

    def __init__(self, filename, **kwargs):
        self._filename = filename
        self._kwargs = kwargs
        self.shape = kwargs['shape'][::-1]
        self.dtype = kwargs['dtype']
        self.ndim = len(self.shape)

    def __getitem__(self, item):
        # We open a file manually and return an in-memory copy of the array
        # otherwise the file doesn't get closed properly.
        with open(self._filename) as f:
            return np.memmap(f, mode='readonly', order='F', **self._kwargs).T[item].copy()


class MaskWrapper:
    """
    A wrapper class for dask that opens binary masks from file and returns
    a chunk from it on-the-fly.
    """

    def __init__(self, filename, offset, count, shape):
        self._filename = filename
        self._offset = offset
        self._count = count
        self.shape = shape[::-1]
        self.dtype = np.bool_
        self.ndim = len(self.shape)

    def __getitem__(self, item):
        # The offset is in the final bit array - but fromfile needs to operate
        # by reading in uint8, so we need to make sure we align what we read
        # in to the bytes.
        start = floor(self._offset / 8)
        end = ceil((self._offset + self._count) / 8)
        array_uint8 = np.fromfile(self._filename, dtype=np.uint8,
                                  offset=start, count=end - start)
        array_bits = np.unpackbits(array_uint8, bitorder='little')
        chunk = array_bits[self._offset - start * 8:self._offset + self._count - start * 8]
        return chunk.reshape(self.shape[::-1], order='F').T[item].astype(np.bool_)


def from_array_fast(arrays, asarray=False, lock=False):
    """
    This is a more efficient alternative to doing::

        [dask.array.from_array(array) for array in arrays]

    that avoids a lot of the overhead in from_array by using the Array
    initializer directly.
    """
    slices = tuple(slice(0, size) for size in arrays[0].shape)
    chunk = tuple((size,) for size in arrays[0].shape)
    meta = np.zeros((0,), dtype=arrays[0].dtype)
    dask_arrays = []
    for array in arrays:
        name1 = str(uuid.uuid4())
        name2 = str(uuid.uuid4())
        dsk = {(name1,) + (0,) * array.ndim: (dask.array.core.getter, name2,
                                              slices, asarray, lock),
               name2: array}
        dask_arrays.append(dask.array.Array(dsk, name1, chunk, meta=meta, dtype=array.dtype))
    return dask_arrays
{'fraction_non_alphanumeric': 0.10336259731654796, 'fraction_numerical': 0.008282259400364419, 'mean_word_length': 2.5351288056206087, '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': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7633606', 'n_tokens_mistral': 1836, 'n_tokens_neox': 1799, 'n_words': 584}