<commit_before>"""
Module with classes corresponding to VTK file types.
"""
import functools
import numpy as np

from . import writer
from .data_array import DataArray
from .unstructured import CellType, check_connectivity


def _make_3darray(points):
    "Complete missing coordinates to get 3d points"
    if points.shape[1] == 3:
        return points

    points_3d = np.zeros((points.shape[0], 3))
    for i in range(min(points.shape[1], 3)):
        points_3d[:, i] = points[:, i]
    return points_3d


def write_manager(cls):
    "Make a class into a context manager that writes on exit"
    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        if type is None:
            # Only write if no error
            self.write()
        return True

    cls.__enter__ = __enter__
    cls.__exit__ = __exit__
    return cls


@write_manager
class VTKFile:
    """Generic VTK file"""

    def __init__(self, filename, filetype, **kwargs):
        """
        Create a generic VTK file

        :param filename: name of file or file handle
        :param filetype: VTK format of file
        :param **kwargs: parameters forwarded to Writer
        """
        self.filetype = filetype
        self.filename = filename
        self.writer = writer.Writer(filetype, **kwargs)

        # Center piece
        self.piece = self.writer.registerPiece()

        # Registering data elements
        self.point_data = self.piece.register('PointData')
        self.cell_data = self.piece.register('CellData')
        self.field_data = self.writer.registerComponent(
            'FieldData', self.writer.data_node
        )

    def addPointData(self, data_array, vtk_format='binary'):
        """
        Add a DataArray instance to the PointData section of the file

        :param data_array: DataArray instance
        :param vtk_format: data format. Can be:
            - ``'ascii'``: data is written with numpy.savetxt in-place
            - ``'binary'``: data is written in base64 (with possible
            compression) in-place
            - ``'append'``: data is written in base64 (with possible
            compression) in the ``AppendData`` section
        """
        self.point_data.registerDataArray(data_array, vtk_format)
        return self

    def addCellData(self, data_array, vtk_format='binary'):
        """
        Add a DataArray instance to the PointData section of the file

        Arguments are identical to `addPointData`.
        """
        self.cell_data.registerDataArray(data_array, vtk_format)
        return self

    def addFieldData(self, data_array, vtk_format='binary'):
        """
        Add a DataArray instance to the FieldData section of the file.

        Arguments are identical to `addPointData`.

        The FieldData section contains arrays that are not related to
        the domain geometry. Their shape is therefore freeform.
        """
        self.field_data.registerDataArray(data_array, vtk_format)
        return self

    def write(self):
        """
        Write XML to file.
        """
        self.writer.registerAppend()
        self.writer.write(self.filename)


class ImageData(VTKFile):
    """VTK Image data (coordinates given by a range and constant spacing)"""

    def __init__(self, filename, ranges, points, **kwargs):
        """
        Init an ImageData file (regular orthogonal grid)

        :param filename: name of file or file handle
        :param ranges: list of pairs for coordinate ranges
        :param points: list of number of points
        """
        VTKFile.__init__(self, filename, 'ImageData', **kwargs)

        # Computing spacings
        spacings = [(x[1] - x[0]) / (n - 1) for x, n in zip(ranges, points)]

        # Filling in missing coordinates
        for _ in range(len(points), 3):
            points.append(1)

        # Setting extents, spacing and origin
        extent = functools.reduce(
            lambda x, y: x + "0 {} ".format(y-1), points, "")
        spacings = functools.reduce(
            lambda x, y: x + "{} ".format(y), spacings, "")
        origins = functools.reduce(
            lambda x, y: x + "{} ".format(y[0]), ranges, "")

        self.writer.setDataNodeAttributes({
            'WholeExtent': extent,
            'Spacing': spacings,
            'Origin': origins
        })

        self.piece.setAttributes({
            "Extent": extent
        })


class RectilinearGrid(VTKFile):
    """VTK Rectilinear grid (coordinates are given by 3 seperate ranges)"""

    def __init__(self, filename, coordinates, offsets=None, **kwargs):
        """
        Init an RectilinearGrid file (irregular orthogonal grid)

        :param filename: name of file or file handle
        :param coordinates: list of coordinates for each direction
        """
        VTKFile.__init__(self, filename, 'RectilinearGrid', **kwargs)

        # Checking that we actually have a list or tuple
        if isinstance(coordinates, np.ndarray):
            coordinates = [coordinates]

        self.coordinates = list(coordinates)

        # Filling in missing coordinates
        for _ in range(len(self.coordinates), 3):
            self.coordinates.append(np.array([0.]))

        # Setting data extent
        extent = []

        for coord in self.coordinates:
            if coord.ndim != 1:
                raise ValueError(
                    'Coordinate array should have only one dimension'
                    + ' (has {})'.format(coord.ndim))
            extent.append(coord.size-1)

        if offsets is None:
            offsets = [0] * len(extent)
        elif len(offsets) == len(coordinates):
            offsets = list(offsets)
            offsets += [0] * (len(extent) - len(offsets))
        else:
            raise ValueError(
                'Size of offsets should '
                'match domain dimension {}'.format(len(coordinates)))

        def fold_extent(acc, couple):
            offset, extent = couple

            if offset != 0:
                offset -= 1  # pragma: no cover
            return acc + "{} {} ".format(offset, offset+extent)

        # Create extent string with offsets
        self.extent = functools.reduce(fold_extent, zip(offsets, extent), "")

        self.writer.setDataNodeAttributes({
            "WholeExtent": self.extent<commit_msg>added offsets as explicit keyword in RectilinearGrid<commit_after>"""
Module with classes corresponding to VTK file types.
"""
import functools
import numpy as np

from . import writer
from .data_array import DataArray
from .unstructured import CellType, check_connectivity


def _make_3darray(points):
    "Complete missing coordinates to get 3d points"
    if points.shape[1] == 3:
        return points

    points_3d = np.zeros((points.shape[0], 3))
    for i in range(min(points.shape[1], 3)):
        points_3d[:, i] = points[:, i]
    return points_3d


def write_manager(cls):
    "Make a class into a context manager that writes on exit"
    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        if type is None:
            # Only write if no error
            self.write()
        return True

    cls.__enter__ = __enter__
    cls.__exit__ = __exit__
    return cls


@write_manager
class VTKFile:
    """Generic VTK file"""

    def __init__(self, filename, filetype, **kwargs):
        """
        Create a generic VTK file

        :param filename: name of file or file handle
        :param filetype: VTK format of file
        :param **kwargs: parameters forwarded to Writer
        """
        self.filetype = filetype
        self.filename = filename
        self.writer = writer.Writer(filetype, **kwargs)

        # Center piece
        self.piece = self.writer.registerPiece()

        # Registering data elements
        self.point_data = self.piece.register('PointData')
        self.cell_data = self.piece.register('CellData')
        self.field_data = self.writer.registerComponent(
            'FieldData', self.writer.data_node
        )

    def addPointData(self, data_array, vtk_format='binary'):
        """
        Add a DataArray instance to the PointData section of the file

        :param data_array: DataArray instance
        :param vtk_format: data format. Can be:
            - ``'ascii'``: data is written with numpy.savetxt in-place
            - ``'binary'``: data is written in base64 (with possible
            compression) in-place
            - ``'append'``: data is written in base64 (with possible
            compression) in the ``AppendData`` section
        """
        self.point_data.registerDataArray(data_array, vtk_format)
        return self

    def addCellData(self, data_array, vtk_format='binary'):
        """
        Add a DataArray instance to the PointData section of the file

        Arguments are identical to `addPointData`.
        """
        self.cell_data.registerDataArray(data_array, vtk_format)
        return self

    def addFieldData(self, data_array, vtk_format='binary'):
        """
        Add a DataArray instance to the FieldData section of the file.

        Arguments are identical to `addPointData`.

        The FieldData section contains arrays that are not related to
        the domain geometry. Their shape is therefore freeform.
        """
        self.field_data.registerDataArray(data_array, vtk_format)
        return self

    def write(self):
        """
        Write XML to file.
        """
        self.writer.registerAppend()
        self.writer.write(self.filename)


class ImageData(VTKFile):
    """VTK Image data (coordinates given by a range and constant spacing)"""

    def __init__(self, filename, ranges, points, offsets=None, **kwargs):
        """
        Init an ImageData file (regular orthogonal grid)

        :param filename: name of file or file handle
        :param ranges: list of pairs for coordinate ranges
        :param points: list of number of points
        """
        VTKFile.__init__(self, filename, 'ImageData', **kwargs)

        # Computing spacings
        spacings = [(x[1] - x[0]) / (n - 1) for x, n in zip(ranges, points)]

        # Filling in missing coordinates
        for _ in range(len(points), 3):
            points.append(1)

        # Setting extents, spacing and origin
        extent = functools.reduce(
            lambda x, y: x + "0 {} ".format(y-1), points, "")
        spacings = functools.reduce(
            lambda x, y: x + "{} ".format(y), spacings, "")
        origins = functools.reduce(
            lambda x, y: x + "{} ".format(y[0]), ranges, "")

        self.writer.setDataNodeAttributes({
            'WholeExtent': extent,
            'Spacing': spacings,
            'Origin': origins
        })

        self.piece.setAttributes({
            "Extent": extent
        })


class RectilinearGrid(VTKFile):
    """VTK Rectilinear grid (coordinates are given by 3 seperate ranges)"""

    def __init__(self, filename, coordinates, offsets=None, **kwargs):
        """
        Init an RectilinearGrid file (irregular orthogonal grid)

        :param filename: name of file or file handle
        :param coordinates: list of coordinates for each direction
        """
        VTKFile.__init__(self, filename, 'RectilinearGrid', **kwargs)

        # Checking that we actually have a list or tuple
        if isinstance(coordinates, np.ndarray):
            coordinates = [coordinates]

        self.coordinates = list(coordinates)

        # Filling in missing coordinates
        for _ in range(len(self.coordinates), 3):
            self.coordinates.append(np.array([0.]))

        # Setting data extent
        extent = []

        for coord in self.coordinates:
            if coord.ndim != 1:
                raise ValueError(
                    'Coordinate array should have only one dimension'
                    + ' (has {})'.format(coord.ndim))
            extent.append(coord.size-1)

        if offsets is None:
            offsets = [0] * len(extent)
        elif len(offsets) == len(coordinates):
            offsets = list(offsets)
            offsets += [0] * (len(extent) - len(offsets))
        else:
            raise ValueError(
                'Size of offsets should '
                'match domain dimension {}'.format(len(coordinates)))

        def fold_extent(acc, couple):
            offset, extent = couple
            offset -= offset != 0
            return acc + "{} {} ".format(offset, offset+extent)

        # Create extent string with offsets
        self.extent = functools.reduce(fold_extent, zip(offsets, extent), "")

        self.writer.setDataNodeAttributes({
            "WholeExtent": self.extent
{'fraction_non_alphanumeric': 0.09914947235785163, 'fraction_numerical': 0.004961411245865491, 'mean_word_length': 2.195520885757423, 'pattern_counts': {'":': 4, '<': 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': '3169856', 'n_tokens_mistral': 3560, 'n_tokens_neox': 3404, 'n_words': 1103}