python实现MATLAB单指标索引

python实现MATLAB单指标索引，即用数字索引多维数组，MATLAB里A(4), 代表列向量A(:)(4)，下面代码实现这种风格索引，当然是从0开始计数，但用过方法py2matlab, 实现从1开始计数。可从Pip下载:

 pip install pindex

# -*- coding: utf-8 -*-
'''pindex.py

a general method to 

example
    ind = SingleIndex([1,2,3])
    s = ind('helloworld')
    print(s)
    ind.py2matlab()
    s = ind('helloworld')
    print(s)

    ind = Index([1,2])
    ss = ind('helloworld')
    print(ss)

    ind = Index(([1,2],[3,4]))
    try:
        ss = ind(['hello', 'world'])
        print(ss)
    except:
        ind.py2matlab()
        ss = ind(['hello', 'world'])
        print(ss)

'''

import numpy as np


class Arith:
    # interface
    def __init__(self,*args):
        pass

    # override + -
    def __add__(self,other):
        pass

    def __radd__(self,other):
        return self+other

    def __neg__(self):
        pass

    def __pos__(self):
        return self

    def __iadd__(self,other):
        return self+other

    def __sub__(self,other):
        return self+(-other)

    def __isub__(self,other):
        return self-other

    def __rsub__(self,other):
        return -self+other

    #override * /
    def __mul__(self,other):
        pass

    def __rmul__(self,other):
        return self*other

    def __imul__(self,other):
        return self*other

    def __truediv__(self,other):
        pass

    def __itruediv__(self,other):
        return self/other

    def __rtruediv__(self,other):
        pass

    def __irtruediv__(self,other):
        return other/other

    def __pow__(self,other):
        pass

    def __ipow__(self,other):
        return self**other


def ind2tuple(ind, ns):
    '''single index (int) to multi-index under ns shape of array
    example:
>>> ind.ind2tuple(7,(3,4))
(1, 2)
>>> ind.ind2tuple(17,(3,4,3))
(2, 1, 1)'''
    dim = len(ns)
    if dim == 0:
        return ()
    if dim == 1:
        return (ind,)
    else:
        d, r = divmod(ind+1, np.prod(ns[:-1]))
        if r == 0:
            return tuple(ind-1 for ind in ns[:-1]) + (d-1,)
        else:
            return ind2tuple(r-1, ns[:-1]) + (d,)


def ind2iter(ind, N=None):
    # slice to range, int and list to list
    if isinstance(ind, list):
        return ind
    elif isinstance(ind, slice): # slc is a slice
        a = 0 if ind.start is None else ind.start
        b = N if ind.stop is None else ind.stop # b <= N
        return range(a, b, c) if ind.step else range(a, b)
    else:
        return [ind]


def addind(ind1, ind2):
    # ind1, ind2 are int list or range (iter)
    if isinstance(ind2, int):
        if isinstance(ind1, int):
            return ind1 + ind2
        elif isinstance(ind1, slice):
            return slice(ind1.start + ind2, ind1.stop + ind2, ind1.step)
        else:
            return [a+ind2 for a in ind1]
    elif isinstance(ind2, slice):
        if isinstance(ind1, int):
            return addind(ind2, ind1)
        elif isinstance(ind1, slice):
            return slice(ind1.start+ind2.start, ind1.stop+ind2.stop, (ind1.step if ind1.step!=None else 1)+(ind12.step if ind2.step!=None else 1))
        else:
            return list(map(np.add, ind1, ind2iter(ind2)))
    else:
        if isinstance(ind1, list):
            return list(map(np.add, ind1, ind2))
        else:
            return addind(ind2, ind1)


def mulind(ind, num):
    # ind is int list or slice, num is int
    if num==0:
        return 0
    if isinstance(num,int):
        if isinstance(ind, int):
            return ind1*num
        elif isinstance(ind, list):
            if num>0:
                return [num*a for a in ind]
            elif num<0:
                return [num*a for a in ind[::-1]]
        else:
            if num>0:
                return slice(num*ind.start, num*ind.stop, num*ind.step if ind.step!=None else num)
            elif num<0:
                return slice(num*ind.stop, num*ind.start, num*ind.step if ind.step!=None else num)
    elif isinstance(ind, int):
        return ind*num


class SingleIndex(Arith):
    '''single index: int, list, slice'''
    def __init__(self, value, N=None):
        # value: int, list, slice
        self.value = value

    def isa(self, T):
        return isinstance(self.value, T)

    def __add__(self, other):
        return SingleIndex(addind(self.value, other.value if isinstance(other, SingleIndex) else other))

    def __mul__(self, num):
        return SingleIndex(mulind(self.value, num))

    def __neg__(self):
        return -1 * self

    def __str__(self):
        return str(self.value)

    def __repr__(self):
        return 'Single Index: {0}'.format(self.value)

    def isempty(self, N=None):
        if self.isa(int):
            return False
        elif self.isa(list):
            return self.value == []
        else:
            a, b, c = self.value.start, self.value.stop, self.value.step
            if c is None or c > 0:
                if a is None or a == 0:
                    return b is None or -N < b < 0 or b > 0
                elif a > 0:
                    return b is None or b > a or b < 0 and b + N > a
                else:
                    return b is None or b > a + N or b < 0 and b > a
            else:
                if a is None or a == N - 1:
                    return b is None or -N <= b < -1 or b < N - 1
                elif a > 0:
                    return b is None or b < a or b < 0 and b + N < a
                else:
                    return b is None or b < a + N or b < 0 and b < a

    def max(self, N=None):
        if self.isa(int):
            return self.value
        elif self.isa(list):
            return max(self.value)
        else:
            a, b, c = self.value.start, self.value.stop, self.value.step
            if c is None or c > 1:
                b = N - 1 if b is None else b - 1
                if c == 1:
                    return b
                else:
                    r = mod(b - a, c)
                    return b - r
            else:
                return N - 1 if a is None else a


    def min(self, N=None):
        if self.isa(int):
            return self.value
        elif self.isa(list):
            return min(self.value)
        else:
            a, b, c = self.value.start, self.value.stop, self.value.step
            if c is None or c > 1:
                return 0 if a is None else a
            else:
                b = 0 if b is None else b + 1
                if c == -1:
                    return b
                else:
                    r = mod(a - b, -c)
                    return b - r


    def __call__(self, x):
        # self(x) == x[self]
        if isinstance(x, (str, tuple, list)):
            if self.isa((int, slice)):
                return x[self.value]
            else:
                return [x[k] for k in self.value]
        elif isinstance(x, np.ndarray):
            ns = x.shape
            if self.isa(int):
                return x[ind2tuple(self.value, ns)]
            elif self.isa(slice):
                return [x[ind2tuple(k, ns)] for k in ind2iter(self.value)]
            else:
                return [x[ind2tuple(k, ns)] for k in self.value]


    def py2matlab(self):
        # python style to matlab style
        if self.isa(int):
            self.value -= 1
        elif self.isa(slice):
            a, b, c = self.value.start, self.value.stop, self.value.step
            if c is None or c > 0:  # a <= b
                self.value = slice(a-1 if a else None, None if b == 0 else b, c)
            elif c < 0: # a >= b
                self.value = slice(a-1 if a else None, None if b == 1 or b is None else b-2, c)
        else:
            self.value = [v-1 for v in self.value]


    def matlab2py(self):
        # matlab style to python style; inverse of py2matlab
        if self.isa(int):
            self.value += 1
        elif self.isa(slice):
            a, b, c = self.value.start, self.value.stop, self.value.step
            if c is None or c > 0:
                if a is not None:
                    a += 1
            else:
                if a is not None:
                    a += 1
                if b is not None:
                    b += 2
            self.value = slice(a, b, c)
        else:
            self.value = [v+1 for v in self.value]



class Index(Arith, np.ndarray):
    # Multi-index, array of SingleIndex
    def __new__(cls, args=0):
        if isinstance(args, (int, list, slice)):
            return super(Index, cls).__new__(cls, shape=1, dtype=SingleIndex, buffer=np.array(SingleIndex(args)))
        elif isinstance(args, tuple):
            return super(Index, cls).__new__(cls, shape=len(args), dtype=SingleIndex, buffer=np.array([SingleIndex(a) for a in args]))
        else:
            raise TypeError('args should be of int, slice, list(of ints), tuple')


    @property
    def dim(self):
        return len(self)

    def __call__(self, array):
        # self(x) == x[self]
        if len(self) == 1:
            return self[0](array)
        else:
            return np.array([self[1:](a) for a in self[0](array)])

    def __add__(self, other):
        if isinstance(other, int):
            return Index(tuple(a+other for a in self))
        else:
            return Index(map(np.add, self, other))

    def __neg__(self):
        return Index(map(lambda x:-x, self))

    def __mul__(self,other):
        if np.isscalar(other):
            return Index(map(lambda x:x*other, self))
        else:
            return Index(map(np.add, self, other))

    def __str__(self):
        return ','.join(map(str, self))

    def minInd(self):
        # minimal index
        m = tuple(a.min() for a in self)
        return Index(m)

    def maxInd(self):
        # minimal index
        m = tuple(a.max() for a in self)
        return Index(m)

    def min_(self, *others):
        return Index(tuple(min(min(other.minInd()[k] for other in others), self.minInd()[k]) for k in range(self.dim)))

    def max_(self, *others):
        return Index(tuple(max(max(other.maxInd()[k] for other in others), self.maxInd()[k]) for k in range(self.dim)))

    def py2matlab(self):
        for ind in self:
            ind.py2matlab()

    def matlab2py(self):
        for ind in self:
            ind.matlab2py()


def irange(minInd, maxInd):
    return Index(tuple(slice(a, b+1) for a, b in zip(minInd, maxInd)))


if __name__ == '__main__':
    ind = Index(([0,1],[0,1]))
    A = [[1,2,6],[3,4,5]]
    B = ind(A)
    print(B)