<commit_before>from bisect import bisect
from copy import deepcopy
import enum
from itertools import product
import logging
from operator import itemgetter
import os
import time
import six

import nibabel as nib
from nibabel.affines import apply_affine
from nibabel.streamlines.array_sequence import ArraySequence
from nibabel.streamlines import detect_format
from nibabel.streamlines.tractogram import (Tractogram,
                                            PerArraySequenceDict,
                                            PerArrayDict)
import numpy as np

from dipy.io.vtk import save_vtk_streamlines, load_vtk_streamlines
from dipy.io.dpy import Dpy


class Space(enum.Enum):
    """ Enum to simplify future change to convention """
    VOX = 'vox'
    VOXMM = 'voxmm'
    RASMM = 'rasmm'


class StatefulTractogram(object):
    """ Class for stateful representation of collections of streamlines
    Object designed to be identical no matter the file format
    (trk, tck, fib, dpy). Facilitate transformation between space and
    data manipulation for each streamline / point.
    """

    def __init__(self, streamlines, reference, space,
                 shifted_origin=False,
                 data_per_point=None, data_per_streamline=None):
        """ Create a strict, state-aware, robust tractogram

        Parameters
        ----------
        streamlines : list or ArraySequence
            Streamlines of the tractogram
        reference : Nifti or Trk filename, Nifti1Image or TrkFile,
            Nifti1Header or trk.header (dict)
            Reference that provides the spatial attribute.
            Typically a nifti-related object from the native diffusion used for
            streamlines generation
        space : string
            Current space in which the streamlines are (vox, voxmm or rasmm)
            Typically after tracking the space is VOX, after nibabel loading
            the space is RASMM
        shifted_origin : bool
            Information on the position of the origin,
            False is Trackvis standard, default (corner of the voxel)
            True is NIFTI standard (center of the voxel)
        data_per_point : dict
            Dictionary in which each key has X items, each items has Y_i items
            X being the number of streamlines
            Y_i being the number of points on streamlines #i
        data_per_streamline : dict
            Dictionary in which each key has X items
            X being the number of streamlines

        Notes
        -----
        Very important to respect the convention, verify that streamlines
        match the reference and are effectively in the right space.

        Any change to the number of streamlines, data_per_point or
        data_per_streamline requires particular verification.

        In a case of manipulation not allowed by this object, use Nibabel
        directly and be careful.
        """
        if data_per_point is None:
            data_per_point = {}

        if data_per_streamline is None:
            data_per_streamline = {}

        self._tractogram = Tractogram(streamlines,
                                      data_per_point=data_per_point,
                                      data_per_streamline=data_per_streamline)

        space_attribute = get_reference_info(reference)
        if space_attribute is None:
            raise TypeError('Reference MUST be one of the following:\n' +
                            'Nifti or Trk filename, Nifti1Image or TrkFile, ' +
                            'Nifti1Header or trk.header (dict)')

        (self._affine, self._dimensions,
         self._voxel_sizes, self._voxel_order) = space_attribute
        self._inv_affine = np.linalg.inv(self._affine)

        if space not in Space:
            raise ValueError('Space MUST be from Space enum, e.g Space.VOX')
        self._space = space

        if not isinstance(shifted_origin, bool):
            raise TypeError('shifted_origin MUST be a boolean')
        self._shifted_origin = shifted_origin
        logging.debug(self)

    def __str__(self):
        """ Generate the string for printing """
        text = 'Affine: \n{}'.format(
            np.array2string(self._affine,
                            formatter={'float_kind': lambda x: "%.6f" % x}))
        text += '\ndimensions: {}'.format(
            np.array2string(self._dimensions))
        text += '\nvoxel_sizes: {}'.format(
            np.array2string(self._voxel_sizes,
                            formatter={'float_kind': lambda x: "%.2f" % x}))
        text += '\nvoxel_order: {}'.format(self._voxel_order)

        text += '\nstreamline_count: {}'.format(self._get_streamline_count())
        text += '\npoint_count: {}'.format(self._get_point_count())
        text += '\ndata_per_streamline keys: {}'.format(
            self.get_data_per_point().keys())
        text += '\ndata_per_point keys: {}'.format(
            self.get_data_per_streamline().keys())

        return text

    def __len__(self):
        """ Define the length of the object """
        return self._get_streamline_count()

    def get_space_attribute(self):
        """ Getter for spatial attribute """
        return self._affine, self._dimensions, self._voxel_sizes, \
            self._voxel_order

    def get_current_space(self):
        """ Getter for the current space """
        return self._space

    def get_current_shift(self):
        """ Getter for shift """
        return self._shifted_origin

    def get_streamlines(self):
        """ Partially safe getter for streamlines """
        return self._tractogram.streamlines

    def get_streamlines_copy(self):
        """ Safe getter for streamlines (for slicing) """
        return deepcopy(list(self._tractogram.streamlines))

    def set_streamlines(self, streamlines):
        """ Modify streamlines. Creating a new object would be less risky.

        Parameters
        ----------
        streamlines : list or ArraySequence (list and deepcopy recommanded)
            Streamlines of the tractogram
        """
        self._tractogram._streamlines = ArraySequence(streamlines)
        self.set_data_per_point(self.get_data_per_point())
        self.set_data_per_streamline(self.get_data_per_streamline())
        logging.warning('Streamlines has been modified')

    def get_data_per_point(self):
        """ Getter for data_per_point """
        return self._tractogram.data_per_point

    def set_data_per_point(self, data):
        """ Modify point data . Creating a new object would be less risky.

        Parameters
        ----------
        data : dict
            Dictionary in which each key has X items, each items has Y_i items
            X being the number of streamlines
            Y_i being the number of points on streamlines #i
        """
        self._tractogram.data_per_point = data
        logging.warning('Data_per_point has been modified')

    def get_data_per_streamline(self):
        """ Getter for data_per_streamline """
        return self._tractogram.data_per_streamline

    def set_data_per_streamline(self, data):
        """ Modify point data . Creating a new object would be less risky.

        Parameters
        ----------
        data : dict
            Dictionary in which each key has X items, each items has Y_i items
            X being the number of streamlines
        """
        self._tractogram.data_per_streamline = data
        logging.warning('Data_per_streamline has been modified')

    def to_vox(self):
        """ Safe function to transform streamlines and update state """
        if self._space == Space.VOXMM:
            self._voxmm_to_vox()
        elif self._space == Space.RASMM:
            self._rasmm_to_vox()

    def to_voxmm(self):
        """ Safe function to transform streamlines and update state """
        if self._space == Space.VOX:
            self._vox_to_voxmm()
        elif self._space == Space.RASMM:
            self._rasmm_to_voxmm()

    def to_rasmm(self):
        """ Safe function to transform streamlines and update state """
        if self._space == Space.VOX:
            self._vox_to_rasmm()
        elif self._space == Space.VOXMM:
            self._voxmm_to_rasmm()

    def to_center(self):
        """ Safe function to shift streamlines so the center of voxel is
        the origin """
        if self._shifted_origin:
            self._shift_voxel_origin()

    def to_corner(self):
        """ Safe function to shift streamlines so the corner of voxel is
        the origin """
        if not self._shifted_origin:
            self._shift_voxel_origin()

    def compute_bounding_box(self):
        """ Compute the bounding box of the streamlines in their current state

        Returns
        -------
        output : ndarray
            8 corners of the XYZ aligned box, all zeros if no streamlines
        """
        if self._tractogram.streamlines.data.size > 0:
            bbox_min = np.min(self._tractogram.streamlines.data, axis=0)
            bbox_max = np.max(self._tractogram.streamlines.data, axis=0)

            return np.asarray(list(product(*zip(bbox_min, bbox_max))))

        return np.zeros((8, 3))

    def is_bbox_in_vox_valid(self):
        """ Verify that the bounding box is valid in voxel space
        Will transform the streamlines for OBB, slow for big tractogram

        Returns
        -------
        output : bool
            Are the streamlines within the volume of the associated reference
        """
        old_space = deepcopy(self.get_current_space())
        old_shift = deepcopy(self.get_current_shift())

        # Do to rotation, equivalent of a OBB must be done
        self.to_vox()
        self.to_corner()
        bbox_corners = deepcopy(self.compute_bounding_box())

        is_valid = True
        if np.any(bbox_corners < 0):
            logging.error('Voxel space values lower than 0.0')
            logging.debug(bbox_corners)
            is_valid = False

        if np.any(bbox_corners[:, 0] > self._dimensions[0]) or \
                np.any(bbox_corners[:, 1] > self._dimensions[1]) or \
                np.any(bbox_corners[:, 2] > self._dimensions[2]):
            logging.error('Voxel space values higher than dimensions')
            logging.debug(bbox_corners)
            is_valid = False

        if old_space == Space.RASMM:
            self.to_rasmm()
        elif old_space == Space.VOXMM:
            self.to_voxmm()

        if not old_shift:
            self.to_center()

        return is_valid

    def remove_invalid_streamlines(self):
        old_space = deepcopy(self.get_current_space())
        old_shift = deepcopy(self.get_current_shift())

        self.to_vox()
        self.to_corner()

        min_condition = np.min(self._tractogram.streamlines.data,
                               axis=1) < 0.0
        max_condition = np.any(self._tractogram.streamlines.data >
                               self._dimensions, axis=1)
        ic_offsets_indices = np.where(np.logical_or(min_condition,
                                                    max_condition))[0]

        ic_indices = []
        for i in ic_offsets_indices:
            ic_indices.append(bisect(
                self._tractogram.streamlines._offsets, i) - 1)

        indices_to_save = np.setdiff1d(np.arange(len(self._tractogram)),
                                       np.array(ic_indices)).astype(int)

        tmp_streamlines = \
            itemgetter(*indices_to_save)(self.get_streamlines_copy())
        tmp_data_per_point = {}
        tmp_data_per_streamline = {}

        for key in self._tractogram.data_per_point:
            tmp_data_per_point[key] = \
                self._tractogram.data_per_point[key][indices_to_save]

        for key in self._tractogram.data_per_streamline:
            tmp_data_per_streamline[key] = \
                self._tractogram.data_per_streamline[key][indices_to_save]

        self._tractogram = Tractogram(tmp_streamlines,
                                      affine_to_rasmm=np.eye(4))

        self._tractogram.data_per_point = tmp_data_per_point
        self._tractogram.data_per_streamline = tmp_data_per_streamline

        if old_space == Space.RASMM:
            self.to_rasmm()
        elif old_space == Space.VOXMM:
            self.to_voxmm()

        if not old_shift:
            self.to_center()<commit_msg>Added docstring for invalid_streamlines function<commit_after>from bisect import bisect
from copy import deepcopy
import enum
from itertools import product
import logging
from operator import itemgetter
import os
import time
import six

import nibabel as nib
from nibabel.affines import apply_affine
from nibabel.streamlines.array_sequence import ArraySequence
from nibabel.streamlines import detect_format
from nibabel.streamlines.tractogram import (Tractogram,
                                            PerArraySequenceDict,
                                            PerArrayDict)
import numpy as np

from dipy.io.vtk import save_vtk_streamlines, load_vtk_streamlines
from dipy.io.dpy import Dpy


class Space(enum.Enum):
    """ Enum to simplify future change to convention """
    VOX = 'vox'
    VOXMM = 'voxmm'
    RASMM = 'rasmm'


class StatefulTractogram(object):
    """ Class for stateful representation of collections of streamlines
    Object designed to be identical no matter the file format
    (trk, tck, fib, dpy). Facilitate transformation between space and
    data manipulation for each streamline / point.
    """

    def __init__(self, streamlines, reference, space,
                 shifted_origin=False,
                 data_per_point=None, data_per_streamline=None):
        """ Create a strict, state-aware, robust tractogram

        Parameters
        ----------
        streamlines : list or ArraySequence
            Streamlines of the tractogram
        reference : Nifti or Trk filename, Nifti1Image or TrkFile,
            Nifti1Header or trk.header (dict)
            Reference that provides the spatial attribute.
            Typically a nifti-related object from the native diffusion used for
            streamlines generation
        space : string
            Current space in which the streamlines are (vox, voxmm or rasmm)
            Typically after tracking the space is VOX, after nibabel loading
            the space is RASMM
        shifted_origin : bool
            Information on the position of the origin,
            False is Trackvis standard, default (corner of the voxel)
            True is NIFTI standard (center of the voxel)
        data_per_point : dict
            Dictionary in which each key has X items, each items has Y_i items
            X being the number of streamlines
            Y_i being the number of points on streamlines #i
        data_per_streamline : dict
            Dictionary in which each key has X items
            X being the number of streamlines

        Notes
        -----
        Very important to respect the convention, verify that streamlines
        match the reference and are effectively in the right space.

        Any change to the number of streamlines, data_per_point or
        data_per_streamline requires particular verification.

        In a case of manipulation not allowed by this object, use Nibabel
        directly and be careful.
        """
        if data_per_point is None:
            data_per_point = {}

        if data_per_streamline is None:
            data_per_streamline = {}

        self._tractogram = Tractogram(streamlines,
                                      data_per_point=data_per_point,
                                      data_per_streamline=data_per_streamline)

        space_attribute = get_reference_info(reference)
        if space_attribute is None:
            raise TypeError('Reference MUST be one of the following:\n' +
                            'Nifti or Trk filename, Nifti1Image or TrkFile, ' +
                            'Nifti1Header or trk.header (dict)')

        (self._affine, self._dimensions,
         self._voxel_sizes, self._voxel_order) = space_attribute
        self._inv_affine = np.linalg.inv(self._affine)

        if space not in Space:
            raise ValueError('Space MUST be from Space enum, e.g Space.VOX')
        self._space = space

        if not isinstance(shifted_origin, bool):
            raise TypeError('shifted_origin MUST be a boolean')
        self._shifted_origin = shifted_origin
        logging.debug(self)

    def __str__(self):
        """ Generate the string for printing """
        text = 'Affine: \n{}'.format(
            np.array2string(self._affine,
                            formatter={'float_kind': lambda x: "%.6f" % x}))
        text += '\ndimensions: {}'.format(
            np.array2string(self._dimensions))
        text += '\nvoxel_sizes: {}'.format(
            np.array2string(self._voxel_sizes,
                            formatter={'float_kind': lambda x: "%.2f" % x}))
        text += '\nvoxel_order: {}'.format(self._voxel_order)

        text += '\nstreamline_count: {}'.format(self._get_streamline_count())
        text += '\npoint_count: {}'.format(self._get_point_count())
        text += '\ndata_per_streamline keys: {}'.format(
            self.get_data_per_point().keys())
        text += '\ndata_per_point keys: {}'.format(
            self.get_data_per_streamline().keys())

        return text

    def __len__(self):
        """ Define the length of the object """
        return self._get_streamline_count()

    def get_space_attribute(self):
        """ Getter for spatial attribute """
        return self._affine, self._dimensions, self._voxel_sizes, \
            self._voxel_order

    def get_current_space(self):
        """ Getter for the current space """
        return self._space

    def get_current_shift(self):
        """ Getter for shift """
        return self._shifted_origin

    def get_streamlines(self):
        """ Partially safe getter for streamlines """
        return self._tractogram.streamlines

    def get_streamlines_copy(self):
        """ Safe getter for streamlines (for slicing) """
        return deepcopy(list(self._tractogram.streamlines))

    def set_streamlines(self, streamlines):
        """ Modify streamlines. Creating a new object would be less risky.

        Parameters
        ----------
        streamlines : list or ArraySequence (list and deepcopy recommanded)
            Streamlines of the tractogram
        """
        self._tractogram._streamlines = ArraySequence(streamlines)
        self.set_data_per_point(self.get_data_per_point())
        self.set_data_per_streamline(self.get_data_per_streamline())
        logging.warning('Streamlines has been modified')

    def get_data_per_point(self):
        """ Getter for data_per_point """
        return self._tractogram.data_per_point

    def set_data_per_point(self, data):
        """ Modify point data . Creating a new object would be less risky.

        Parameters
        ----------
        data : dict
            Dictionary in which each key has X items, each items has Y_i items
            X being the number of streamlines
            Y_i being the number of points on streamlines #i
        """
        self._tractogram.data_per_point = data
        logging.warning('Data_per_point has been modified')

    def get_data_per_streamline(self):
        """ Getter for data_per_streamline """
        return self._tractogram.data_per_streamline

    def set_data_per_streamline(self, data):
        """ Modify point data . Creating a new object would be less risky.

        Parameters
        ----------
        data : dict
            Dictionary in which each key has X items, each items has Y_i items
            X being the number of streamlines
        """
        self._tractogram.data_per_streamline = data
        logging.warning('Data_per_streamline has been modified')

    def to_vox(self):
        """ Safe function to transform streamlines and update state """
        if self._space == Space.VOXMM:
            self._voxmm_to_vox()
        elif self._space == Space.RASMM:
            self._rasmm_to_vox()

    def to_voxmm(self):
        """ Safe function to transform streamlines and update state """
        if self._space == Space.VOX:
            self._vox_to_voxmm()
        elif self._space == Space.RASMM:
            self._rasmm_to_voxmm()

    def to_rasmm(self):
        """ Safe function to transform streamlines and update state """
        if self._space == Space.VOX:
            self._vox_to_rasmm()
        elif self._space == Space.VOXMM:
            self._voxmm_to_rasmm()

    def to_center(self):
        """ Safe function to shift streamlines so the center of voxel is
        the origin """
        if self._shifted_origin:
            self._shift_voxel_origin()

    def to_corner(self):
        """ Safe function to shift streamlines so the corner of voxel is
        the origin """
        if not self._shifted_origin:
            self._shift_voxel_origin()

    def compute_bounding_box(self):
        """ Compute the bounding box of the streamlines in their current state

        Returns
        -------
        output : ndarray
            8 corners of the XYZ aligned box, all zeros if no streamlines
        """
        if self._tractogram.streamlines.data.size > 0:
            bbox_min = np.min(self._tractogram.streamlines.data, axis=0)
            bbox_max = np.max(self._tractogram.streamlines.data, axis=0)

            return np.asarray(list(product(*zip(bbox_min, bbox_max))))

        return np.zeros((8, 3))

    def is_bbox_in_vox_valid(self):
        """ Verify that the bounding box is valid in voxel space
        Will transform the streamlines for OBB, slow for big tractogram

        Returns
        -------
        output : bool
            Are the streamlines within the volume of the associated reference
        """
        old_space = deepcopy(self.get_current_space())
        old_shift = deepcopy(self.get_current_shift())

        # Do to rotation, equivalent of a OBB must be done
        self.to_vox()
        self.to_corner()
        bbox_corners = deepcopy(self.compute_bounding_box())

        is_valid = True
        if np.any(bbox_corners < 0):
            logging.error('Voxel space values lower than 0.0')
            logging.debug(bbox_corners)
            is_valid = False

        if np.any(bbox_corners[:, 0] > self._dimensions[0]) or \
                np.any(bbox_corners[:, 1] > self._dimensions[1]) or \
                np.any(bbox_corners[:, 2] > self._dimensions[2]):
            logging.error('Voxel space values higher than dimensions')
            logging.debug(bbox_corners)
            is_valid = False

        if old_space == Space.RASMM:
            self.to_rasmm()
        elif old_space == Space.VOXMM:
            self.to_voxmm()

        if not old_shift:
            self.to_center()

        return is_valid

    def remove_invalid_streamlines(self):
        """ Remove streamlines with invalid coordinates from the object.
        Will also remove the data_per_point and data_per_streamlines.
        Invalid coordinates are any X,Y,Z values above the reference
        dimensions or below zero         Returns
        -------
        output : tuple
            Tuple of two list, indices_to_remove, indices_to_keep
        """
        old_space = deepcopy(self.get_current_space())
        old_shift = deepcopy(self.get_current_shift())

        self.to_vox()
        self.to_corner()

        min_condition = np.min(self._tractogram.streamlines.data,
                               axis=1) < 0.0
        max_condition = np.any(self._tractogram.streamlines.data >
                               self._dimensions, axis=1)
        ic_offsets_indices = np.where(np.logical_or(min_condition,
                                                    max_condition))[0]

        indices_to_remove = []
        for i in ic_offsets_indices:
            indices_to_remove.append(bisect(
                self._tractogram.streamlines._offsets, i) - 1)

        indices_to_keep = np.setdiff1d(np.arange(len(self._tractogram)),
                                       np.array(indices_to_remove)).astype(int)

        tmp_streamlines = \
            itemgetter(*indices_to_keep)(self.get_streamlines_copy())
        tmp_data_per_point = {}
        tmp_data_per_streamline = {}

        for key in self._tractogram.data_per_point:
            tmp_data_per_point[key] = \
                self._tractogram.data_per_point[key][indices_to_keep]

        for key in self._tractogram.data_per_streamline:
            tmp_data_per_streamline[key] = \
                self._tractogram.data_per_streamline[key][indices_to_keep]

        self._tractogram = Tractogram(tmp_streamlines,
                                      affine_to_rasmm=np.eye(4))

        self._tractogram.data_per_point = tmp_data_per_point
        self._tractogram.data_per_streamline = tmp_data_per_streamline

        if old_space == Space.RASMM:
            self.to_rasmm()
        elif old_space == Space.VOXMM:
            self.to_voxmm()

        if not old_shift:
            self.to_center()

        return indices_to_remove, indices_to_keep
{'fraction_non_alphanumeric': 0.08059464816650148, 'fraction_numerical': 0.0025371655104063427, 'mean_word_length': 2.100921942224954, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 20, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2950422', 'n_tokens_mistral': 7081, 'n_tokens_neox': 6839, 'n_words': 1915}