python-pandas 0.7.0-1 / usr / share / pyshared / pandas / sparse / panel.py

"""
Data structures for sparse float data. Life is made simpler by dealing only with
float64 data
"""

# pylint: disable=E1101,E1103,W0231

import numpy as np

from pandas.core.common import _pickle_array, _unpickle_array, _mut_exclusive
from pandas.core.index import Index, MultiIndex, _ensure_index
from pandas.core.frame import DataFrame
from pandas.core.panel import Panel

from pandas.sparse.frame import SparseDataFrame

from pandas.util.decorators import deprecate

class SparsePanelAxis(object):

    def __init__(self, cache_field, frame_attr):
        self.cache_field = cache_field
        self.frame_attr = frame_attr

    def __get__(self, obj, type=None):
        return getattr(obj, self.cache_field, None)

    def __set__(self, obj, value):
        value = _ensure_index(value)

        if isinstance(value, MultiIndex):
            raise NotImplementedError

        for v in obj._frames.itervalues():
            setattr(v, self.frame_attr, value)

        setattr(obj, self.cache_field, value)


class SparsePanel(Panel):
    """
    Sparse version of Panel

    Parameters
    ----------
    frames : dict of DataFrame objects
    items : array-like
    major_axis : array-like
    minor_axis : array-like
    default_kind : {'block', 'integer'}, default 'block'
        Default sparse kind for converting Series to SparseSeries. Will not
        override SparseSeries passed into constructor
    default_fill_value : float
        Default fill_value for converting Series to SparseSeries. Will not
        override SparseSeries passed in

    Notes
    -----
    """
    ndim = 3

    def __init__(self, frames, items=None, major_axis=None, minor_axis=None,
                 default_fill_value=np.nan, default_kind='block'):
        if isinstance(frames, np.ndarray):
            new_frames = {}
            for item, vals in zip(items, frames):
                new_frames[item] = \
                    SparseDataFrame(vals, index=major_axis,
                                    columns=minor_axis,
                                    default_fill_value=default_fill_value,
                                    default_kind=default_kind)
            frames = new_frames

        assert(isinstance(frames, dict))

        self.default_fill_value = fill_value = default_fill_value
        self.default_kind = kind = default_kind

        # pre-filter, if necessary
        if items is None:
            items = Index(sorted(frames.keys()))
        items = _ensure_index(items)

        (clean_frames,
         major_axis,
         minor_axis) = _convert_frames(frames, major_axis,
                                       minor_axis, kind=kind,
                                       fill_value=fill_value)

        self._frames = clean_frames

        # do we want to fill missing ones?
        for item in items:
            if item not in clean_frames:
                raise Exception('column %s not found in data' % item)

        self._items = items
        self.major_axis = major_axis
        self.minor_axis = minor_axis

    def _consolidate_inplace(self): # pragma: no cover
        # do nothing when DataFrame calls this method
        pass

    def __array_wrap__(self, result):
        return SparsePanel(result, items=self.items,
                           major_axis=self.major_axis,
                           minor_axis=self.minor_axis,
                           default_kind=self.default_kind,
                           default_fill_value=self.default_fill_value)

    @classmethod
    def from_dict(cls, data):
        """
        Analogous to Panel.from_dict
        """
        return SparsePanel(data)

    def to_dense(self):
        """
        Convert SparsePanel to (dense) Panel

        Returns
        -------
        dense : Panel
        """
        return Panel(self.values, self.items, self.major_axis,
                     self.minor_axis)

    @property
    def values(self):
        # return dense values
        return np.array([self._frames[item].values
                         for item in self.items])

    # need a special property for items to make the field assignable

    _items = None
    def _get_items(self):
        return self._items

    def _set_items(self, new_items):
        new_items = _ensure_index(new_items)
        if isinstance(new_items, MultiIndex):
            raise NotImplementedError

        # need to create new frames dict

        old_frame_dict = self._frames
        old_items = self._items
        self._frames = dict((new_k, old_frame_dict[old_k])
                            for new_k, old_k in zip(new_items, old_items))
        self._items = new_items
    items = property(fget=_get_items, fset=_set_items)

    # DataFrame's index
    major_axis = SparsePanelAxis('_major_axis', 'index')

    # DataFrame's columns / "items"
    minor_axis = SparsePanelAxis('_minor_axis', 'columns')

    def _get_item_cache(self, key):
        return self._frames[key]

    def __setitem__(self, key, value):
        if isinstance(value, DataFrame):
            value = value.reindex(index=self.major_axis,
                                  columns=self.minor_axis)
            if not isinstance(value, SparseDataFrame):
                value = value.to_sparse(fill_value=self.default_fill_value,
                                        kind=self.default_kind)
        else:
            raise ValueError('only DataFrame objects can be set currently')

        self._frames[key] = value

        if key not in self.items:
            self._items = Index(list(self.items) + [key])

    def set_value(self, item, major, minor, value):
        """
        Quickly set single value at (item, major, minor) location

        Parameters
        ----------
        item : item label (panel item)
        major : major axis label (panel item row)
        minor : minor axis label (panel item column)
        value : scalar

        Notes
        -----
        This method *always* returns a new object. It is not particularly
        efficient but is provided for API compatibility with Panel

        Returns
        -------
        panel : SparsePanel
        """
        dense = self.to_dense().set_value(item, major, minor, value)
        return dense.to_sparse(kind=self.default_kind,
                               fill_value=self.default_fill_value)

    def __delitem__(self, key):
        loc = self.items.get_loc(key)
        indices = range(loc) + range(loc + 1, len(self.items))
        del self._frames[key]
        self._items = self._items.take(indices)

    def __getstate__(self):
        # pickling
        return (self._frames, _pickle_array(self.items),
                _pickle_array(self.major_axis), _pickle_array(self.minor_axis),
                self.default_fill_value, self.default_kind)

    def __setstate__(self, state):
        frames, items, major, minor, fv, kind = state

        self.default_fill_value = fv
        self.default_kind = kind
        self._items = _unpickle_array(items)
        self._major_axis = _unpickle_array(major)
        self._minor_axis = _unpickle_array(minor)
        self._frames = frames

    def copy(self):
        """
        Make a (shallow) copy of the sparse panel

        Returns
        -------
        copy : SparsePanel
        """
        return SparsePanel(self._frames.copy(), items=self.items,
                           major_axis=self.major_axis,
                           minor_axis=self.minor_axis,
                           default_fill_value=self.default_fill_value,
                           default_kind=self.default_kind)

    def to_frame(self, filter_observations=True):
        """
        Convert SparsePanel to (dense) DataFrame

        Returns
        -------
        frame : DataFrame
        """
        if not filter_observations:
            raise Exception('filter_observations=False not supported for '
                            'SparsePanel.to_long')

        I, N, K = self.shape
        counts = np.zeros(N * K, dtype=int)

        d_values = {}
        d_indexer = {}

        for item in self.items:
            frame = self[item]

            values, major, minor = _stack_sparse_info(frame)

            # values are stacked column-major
            indexer = minor * N + major
            counts.put(indexer, counts.take(indexer) + 1) # cuteness

            d_values[item] = values
            d_indexer[item] = indexer

        # have full set of observations for each item
        mask = counts == I

        # for each item, take mask values at index locations for those sparse
        # values, and use that to select values
        values = np.column_stack([d_values[item][mask.take(d_indexer[item])]
                                  for item in self.items])

        inds, = mask.nonzero()

        # still column major
        major_labels = inds % N
        minor_labels = inds // N

        index = MultiIndex(levels=[self.major_axis, self.minor_axis],
                           labels=[major_labels, minor_labels])

        df = DataFrame(values, index=index, columns=self.items)
        return df.sortlevel(level=0)

    to_long = deprecate('to_long', to_frame)
    toLong = deprecate('toLong', to_frame)

    def reindex(self, major=None, items=None, minor=None, major_axis=None,
                minor_axis=None, copy=False):
        """
        Conform / reshape panel axis labels to new input labels

        Parameters
        ----------
        major : array-like, default None
        items : array-like, default None
        minor : array-like, default None
        copy : boolean, default False
            Copy underlying SparseDataFrame objects

        Returns
        -------
        reindexed : SparsePanel
        """
        major = _mut_exclusive(major, major_axis)
        minor = _mut_exclusive(minor, minor_axis)

        if None == major == items == minor:
            raise ValueError('Must specify at least one axis')

        major = self.major_axis if major is None else major
        minor = self.minor_axis if minor is None else minor

        if items is not None:
            new_frames = {}
            for item in items:
                if item in self._frames:
                    new_frames[item] = self._frames[item]
                else:
                    raise Exception('Reindexing with new items not yet '
                                    'supported')
        else:
            new_frames = self._frames

        if copy:
            new_frames = dict((k, v.copy()) for k, v in new_frames.iteritems())

        return SparsePanel(new_frames, items=items,
                           major_axis=major,
                           minor_axis=minor,
                           default_fill_value=self.default_fill_value,
                           default_kind=self.default_kind)

    def _combine(self, other, func, axis=0):
        if isinstance(other, DataFrame):
            return self._combineFrame(other, func, axis=axis)
        elif isinstance(other, Panel):
            return self._combinePanel(other, func)
        elif np.isscalar(other):
            new_frames = dict((k, func(v, other))
                              for k, v in self.iterkv())
            return self._new_like(new_frames)

    def _combineFrame(self, other, func, axis=0):
        index, columns = self._get_plane_axes(axis)
        axis = self._get_axis_number(axis)

        other = other.reindex(index=index, columns=columns)

        if axis == 0:
            new_values = func(self.values, other.values)
        elif axis == 1:
            new_values = func(self.values.swapaxes(0, 1), other.values.T)
            new_values = new_values.swapaxes(0, 1)
        elif axis == 2:
            new_values = func(self.values.swapaxes(0, 2), other.values)
            new_values = new_values.swapaxes(0, 2)

        # TODO: make faster!
        new_frames = {}
        for item, item_slice in zip(self.items, new_values):
            old_frame = self[item]
            ofv = old_frame.default_fill_value
            ok = old_frame.default_kind
            new_frames[item] = SparseDataFrame(item_slice,
                                               index=self.major_axis,
                                               columns=self.minor_axis,
                                               default_fill_value=ofv,
                                               default_kind=ok)

        return self._new_like(new_frames)

    def _new_like(self, new_frames):
        return SparsePanel(new_frames, self.items, self.major_axis,
                           self.minor_axis,
                           default_fill_value=self.default_fill_value,
                           default_kind=self.default_kind)

    def _combinePanel(self, other, func):
        items = self.items + other.items
        major = self.major_axis + other.major_axis
        minor = self.minor_axis + other.minor_axis

        # could check that everything's the same size, but forget it

        this = self.reindex(items=items, major=major, minor=minor)
        other = other.reindex(items=items, major=major, minor=minor)

        new_frames = {}
        for item in items:
            new_frames[item] = func(this[item], other[item])

        # maybe unnecessary
        new_default_fill = func(self.default_fill_value,
                                other.default_fill_value)

        return SparsePanel(new_frames, items, major, minor,
                           default_fill_value=new_default_fill,
                           default_kind=self.default_kind)

    def major_xs(self, key):
        """
        Return slice of panel along major axis

        Parameters
        ----------
        key : object
            Major axis label

        Returns
        -------
        y : DataFrame
            index -> minor axis, columns -> items
        """
        slices = dict((k, v.xs(key)) for k, v in self.iterkv())
        return DataFrame(slices, index=self.minor_axis, columns=self.items)

    def minor_xs(self, key):
        """
        Return slice of panel along minor axis

        Parameters
        ----------
        key : object
            Minor axis label

        Returns
        -------
        y : SparseDataFrame
            index -> major axis, columns -> items
        """
        slices = dict((k, v[key]) for k, v in self.iterkv())
        return SparseDataFrame(slices, index=self.major_axis,
                               columns=self.items,
                               default_fill_value=self.default_fill_value,
                               default_kind=self.default_kind)

SparseWidePanel = SparsePanel

def _convert_frames(frames, index, columns, fill_value=np.nan, kind='block'):
    from pandas.core.panel import _get_combined_index
    output = {}
    for item, df in frames.iteritems():
        if not isinstance(df, SparseDataFrame):
            df = SparseDataFrame(df, default_kind=kind,
                                 default_fill_value=fill_value)

        output[item] = df

    if index is None:
        all_indexes = [df.index for df in output.values()]
        index = _get_combined_index(all_indexes)
    if columns is None:
        all_columns = [df.columns for df in output.values()]
        columns = _get_combined_index(all_columns)

    index = _ensure_index(index)
    columns = _ensure_index(columns)

    for item, df in output.iteritems():
        if not (df.index.equals(index) and df.columns.equals(columns)):
            output[item] = df.reindex(index=index, columns=columns)

    return output, index, columns


def _stack_sparse_info(frame):
    lengths = [s.sp_index.npoints for _, s in frame.iteritems()]

    # this is pretty fast
    minor_labels = np.repeat(np.arange(len(frame.columns)), lengths)

    inds_to_concat = []
    vals_to_concat = []
    for col in frame.columns:
        series = frame[col]

        if not np.isnan(series.fill_value):
            raise Exception('This routine assumes NaN fill value')

        int_index = series.sp_index.to_int_index()
        inds_to_concat.append(int_index.indices)
        vals_to_concat.append(series.sp_values)

    major_labels = np.concatenate(inds_to_concat)
    sparse_values = np.concatenate(vals_to_concat)

    return sparse_values, major_labels, minor_labels