<commit_before>import numpy as np


class CoefficientArray:
    def __init__(self, dtype=np.complex, ctype=np.matrix):
        """
        dtype -- number type
        ctype -- container type (default np.matrix)
        """
        self.dtype = dtype
        self.ctype = ctype
        self._data = {}

    def __getitem__(self, key):
        """
        key -- (k, ispn)
        """

        # # return as view, not required
        # return self._data[key][:]

        # return as view
        return self._data[key]

    def __setitem__(self, key, item):
        """
        """
        if key in self._data:
            x = self._data[key]
            # make sure shapes don't change
            try:
                # view, no copy needed
                x[:] = self.ctype(item, copy=False)
            except TypeError:
                # not a view, make a copy
                x = self.ctype(item, copy=True)
        else:
            self._data[key] = self.ctype(item, dtype=self.dtype, copy=True)

    def sum(self, **kwargs):
        """
        """
        from mpi4py import MPI
        loc_sum = np.array(sum([np.sum(v) for _, v in self.items()]), dtype=np.complex128)
        rcvBuf = np.array(0.0, dtype=np.complex128)

        MPI.COMM_WORLD.Allreduce(
            [loc_sum, MPI.DOUBLE_COMPLEX],
            [rcvBuf, MPI.DOUBLE_COMPLEX],
                op=MPI.SUM)
        return np.asscalar(rcvBuf)

    def __mul__(self, other):
        """
        Returns a new object of type type(self)
        """
        out = type(self)(dtype=self.dtype)
        if isinstance(other, CoefficientArray):
            for key in other._data.keys():
                out[key] = np.einsum('ij,ij->ij', self._data[key], other._data[key])
        elif np.isscalar(other):
            for key in self._data.keys():
                out[key] = self._data[key] * other
        else:
            raise TypeError('wrong type')
        return out

    def __matmul__(self, other):
        """
        TODO
        """
        # TODO: complex | double -> complex, double | double -> double
        out = type(self)(dtype=np.complex)
        assert(self.ctype is np.matrix)
        if isinstance(other, CoefficientArray):
            for key in other._data.keys():
                out[key] = self._data[key] * other._data[key]
        elif np.isscalar(other):
            for key in self._data.keys():
                out[key] = self._data[key] * other
        else:
            raise TypeError('wrong type')
        return out

    def __add__(self, other):
        """
        """
        out = type(self)(dtype=self.dtype)
        if isinstance(other, CoefficientArray):
            for key in other._data.keys():
                out[key] = self._data[key] + other._data[key]
        elif np.isscalar(other):
            for key in self._data.keys():
                out[key] = self._data[key] + other
        return out

    def __neg__(self):
        """

        """
        out = type(self)(dtype=self.dtype)
        for key, val in self._data.items():
            out[key] = -val
        return out

    def abs(self):
        """
        """
        out = type(self)(dtype=self.dtype)
        for key in self._data.keys():
            out[key] = np.abs(self._data[key])
        return out

    def keys(self):
        return self._data.keys()

    def __sub__(self, other):
        """
        """
        return self.__add__(-1 * other)

    def conjugate(self):
        """
        """
        out = type(self)(dtype=self.dtype)
        for key, val in self._data.items():
            out[key] = np.conj(val)
        return out<commit_msg>python module: docstrings<commit_after>import numpy as np


class CoefficientArray:
    def __init__(self, dtype=np.complex, ctype=np.matrix):
        """
        dtype -- number type
        ctype -- container type (default np.matrix)
        """
        self.dtype = dtype
        self.ctype = ctype
        self._data = {}

    def __getitem__(self, key):
        """
        key -- (k, ispn)
        """

        # # return as view, not required
        # return self._data[key][:]

        # return as view
        return self._data[key]

    def __setitem__(self, key, item):
        """
        """
        if key in self._data:
            x = self._data[key]
            # make sure shapes don't change
            try:
                # view, no copy needed
                x[:] = self.ctype(item, copy=False)
            except TypeError:
                # not a view, make a copy
                x = self.ctype(item, copy=True)
        else:
            self._data[key] = self.ctype(item, dtype=self.dtype, copy=True)

    def sum(self, **kwargs):
        """
        """
        from mpi4py import MPI
        loc_sum = np.array(sum([np.sum(v) for _, v in self.items()]), dtype=np.complex128)
        rcvBuf = np.array(0.0, dtype=np.complex128)

        MPI.COMM_WORLD.Allreduce(
            [loc_sum, MPI.DOUBLE_COMPLEX],
            [rcvBuf, MPI.DOUBLE_COMPLEX],
                op=MPI.SUM)
        return np.asscalar(rcvBuf)

    def __mul__(self, other):
        """
        Returns a new object of type type(self)
        """
        out = type(self)(dtype=self.dtype)
        if isinstance(other, CoefficientArray):
            for key in other._data.keys():
                out[key] = np.einsum('ij,ij->ij', self._data[key], other._data[key])
        elif np.isscalar(other):
            for key in self._data.keys():
                out[key] = self._data[key] * other
        else:
            raise TypeError('wrong type')
        return out

    def __matmul__(self, other):
        """
        TODO
        """
        # TODO: complex | double -> complex, double | double -> double
        out = type(self)(dtype=np.complex)
        assert(self.ctype is np.matrix)
        if isinstance(other, CoefficientArray):
            for key in other._data.keys():
                out[key] = self._data[key] * other._data[key]
        elif np.isscalar(other):
            for key in self._data.keys():
                out[key] = self._data[key] * other
        else:
            raise TypeError('wrong type')
        return out

    def __add__(self, other):
        out = type(self)(dtype=self.dtype)
        if isinstance(other, CoefficientArray):
            for key in other._data.keys():
                out[key] = self._data[key] + other._data[key]
        elif np.isscalar(other):
            for key in self._data.keys():
                out[key] = self._data[key] + other
        return out

    def __neg__(self):
        out = type(self)(dtype=self.dtype)
        for key, val in self._data.items():
            out[key] = -val
        return out

    def abs(self):
        out = type(self)(dtype=self.dtype)
        for key in self._data.keys():
            out[key] = np.abs(self._data[key])
        return out

    def keys(self):
        return self._data.keys()

    def __sub__(self, other):
        return self.__add__(-1 * other)

    def conjugate(self):
        out = type(self)(dtype=self.dtype)
        for key, val in self._data.items():
            out[key] = np.conj(val)
        return out
{'fraction_non_alphanumeric': 0.12672521957340024, 'fraction_numerical': 0.0027882336539802035, 'mean_word_length': 1.6239941477688369, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, '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': '5095755', 'n_tokens_mistral': 2196, 'n_tokens_neox': 2134, 'n_words': 517}