path: root/numpy/core/defmatrix.py

__all__ = ['matrix', 'bmat', 'mat', 'asmatrix']

import sys
import numeric as N
from numeric import concatenate, isscalar, binary_repr

# make translation table
_table = [None]*256
for k in range(256):
    _table[k] = chr(k)
_table = ''.join(_table)

_numchars = '0123456789.-+jeEL'
_todelete = []
for k in _table:
    if k not in _numchars:
        _todelete.append(k)
_todelete = ''.join(_todelete)
del k

def _eval(astr):
    return eval(astr.translate(_table,_todelete))

def _convert_from_string(data):
    rows = data.split(';')
    newdata = []
    count = 0
    for row in rows:
        trow = row.split(',')
        newrow = []
        for col in trow:
            temp = col.split()
            newrow.extend(map(_eval,temp))
        if count == 0:
            Ncols = len(newrow)
        elif len(newrow) != Ncols:
            raise ValueError, "Rows not the same size."
        count += 1
        newdata.append(newrow)
    return newdata

def asmatrix(data, dtype=None):
    """ Returns 'data' as a matrix.  Unlike matrix(), no copy is performed
    if 'data' is already a matrix or array.  Equivalent to:
    matrix(data, copy=False)
    """
    return matrix(data, dtype=dtype, copy=False)



class matrix(N.ndarray):
    __array_priority__ = 10.0
    def __new__(subtype, data, dtype=None, copy=True):
        if isinstance(data, matrix):
            dtype2 = data.dtype
            if (dtype is None):
                dtype = dtype2
            if (dtype2 == dtype) and (not copy):
                return data
            return data.astype(dtype)

        if isinstance(data, N.ndarray):
            if dtype is None:
                intype = data.dtype
            else:
                intype = N.dtype(dtype)
            new = data.view(subtype)
            if intype != data.dtype:
                return new.astype(intype)
            if copy: return new.copy()
            else: return new

        if isinstance(data, str):
            data = _convert_from_string(data)

        # now convert data to an array
        arr = N.array(data, dtype=dtype, copy=copy)
        ndim = arr.ndim
        shape = arr.shape
        if (ndim > 2):
            raise ValueError, "matrix must be 2-dimensional"
        elif ndim == 0:
            shape = (1,1)
        elif ndim == 1:
            shape = (1,shape[0])

        order = False
        if (ndim == 2) and arr.flags.fortran:
            order = True

        if not (order or arr.flags.contiguous):
            arr = arr.copy()

        ret = N.ndarray.__new__(subtype, shape, arr.dtype,
                                buffer=arr,
                                order=order)
        return ret

    def __array_finalize__(self, obj):
        self._getitem = False
        if (isinstance(obj, matrix) and obj._getitem): return
        ndim = self.ndim
        if (ndim == 2):
            return
        if (ndim > 2):
            newshape = tuple([x for x in self.shape if x > 1])
            ndim = len(newshape)
            if ndim == 2:
                self.shape = newshape
                return
            elif (ndim > 2):
                raise ValueError, "shape too large to be a matrix."
        else:
            newshape = self.shape
        if ndim == 0:
            self.shape = (1,1)
        elif ndim == 1:
            self.shape = (1,newshape[0])
        return

    def __getitem__(self, index):
        self._getitem = True
        try:
            out = N.ndarray.__getitem__(self, index)
        finally:
            self._getitem = False

        if not isinstance(out, N.ndarray):
            return out

        if out.ndim == 0:
            return out[()]
        if out.ndim == 1:
            sh = out.shape[0]
            # Determine when we should have a column array
            try:
                n = len(index)
            except:
                n = 0
            if n > 1 and isscalar(index[1]):
                out.shape = (sh,1)
            else:
                out.shape = (1,sh)
        return out

    def _get_truendim(self):
        shp = self.shape
        truend = 0
        for val in shp:
            if (val > 1): truend += 1
        return truend


    def __mul__(self, other):
        if isinstance(other,(N.ndarray, list, tuple)) :
            # This promotes 1-D vectors to row vectors
            return N.dot(self, asmatrix(other))
        if N.isscalar(other) or not hasattr(other, '__rmul__') :
            return N.dot(self, other)
        return NotImplemented

    def __rmul__(self, other):
        return N.dot(other, self)

    def __imul__(self, other):
        self[:] = self * other
        return self

    def __pow__(self, other):
        shape = self.shape
        if len(shape) != 2 or shape[0] != shape[1]:
            raise TypeError, "matrix is not square"
        if type(other) in (type(1), type(1L)):
            if other==0:
                return matrix(N.identity(shape[0]))
            if other<0:
                x = self.I
                other=-other
            else:
                x=self
            result = x
            if other <= 3:
                while(other>1):
                    result=result*x
                    other=other-1
                return result
            # binary decomposition to reduce the number of Matrix
            #  Multiplies for other > 3.
            beta = binary_repr(other)
            t = len(beta)
            Z,q = x.copy(),0
            while beta[t-q-1] == '0':
                Z *= Z
                q += 1
            result = Z.copy()
            for k in range(q+1,t):
                Z *= Z
                if beta[t-k-1] == '1':
                    result *= Z
            return result
        else:
            raise TypeError, "exponent must be an integer"

    def __rpow__(self, other):
        return NotImplemented

    def __repr__(self):
        s = repr(self.__array__()).replace('array', 'matrix')
        # now, 'matrix' has 6 letters, and 'array' 5, so the columns don't
        # line up anymore. We need to add a space.
        l = s.splitlines()
        for i in range(1, len(l)):
            if l[i]:
                l[i] = ' ' + l[i]
        return '\n'.join(l)

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

    def _align(self, axis):
        """A convenience function for operations that need to preserve axis
        orientation.
        """
        if axis is None:
            return self[0,0]
        elif axis==0:
            return self
        elif axis==1:
            return self.transpose()
        else:
            raise ValueError, "unsupported axis"

    # To preserve orientation of result...
    def sum(self, axis=None, dtype=None, out=None):
        """Sum the matrix over the given axis.  If the axis is None, sum
        over all dimensions.  This preserves the orientation of the
        result as a row or column.
        """
        return N.ndarray.sum(self, axis, dtype, out)._align(axis)

    def mean(self, axis=None, out=None):
        """Compute the mean along the specified axis.

        Returns the average of the array elements.  The average is taken over
        the flattened array by default, otherwise over the specified axis.

        :Parameters:

            axis : integer
                Axis along which the means are computed. The default is
                to compute the standard deviation of the flattened array.

            dtype : type
                Type to use in computing the means. For arrays of integer type
                the default is float32, for arrays of float types it is the
                same as the array type.

            out : ndarray
                Alternative output array in which to place the result. It must
                have the same shape as the expected output but the type will be
                cast if necessary.

        :Returns:

            mean : The return type varies, see above.
                A new array holding the result is returned unless out is
                specified, in which case a reference to out is returned.

        :SeeAlso:

            - var : variance
            - std : standard deviation

        Notes
        -----

            The mean is the sum of the elements along the axis divided by the
            number of elements.

        """
        return N.ndarray.mean(self, axis, out)._align(axis)

    def std(self, axis=None, dtype=None, out=None):
        """Compute the standard deviation along the specified axis.

        Returns the standard deviation of the array elements, a measure of the
        spread of a distribution. The standard deviation is computed for the
        flattened array by default, otherwise over the specified axis.

        :Parameters:

            axis : integer
                Axis along which the standard deviation is computed. The
                default is to compute the standard deviation of the flattened
                array.

            dtype : type
                Type to use in computing the standard deviation. For arrays of
                integer type the default is float32, for arrays of float types
                it is the same as the array type.

            out : ndarray
                Alternative output array in which to place the result. It must
                have the same shape as the expected output but the type will be
                cast if necessary.

        :Returns:

            standard deviation : The return type varies, see above.
                A new array holding the result is returned unless out is
                specified, in which case a reference to out is returned.

        :SeeAlso:

            - var : variance
            - mean : average

        Notes
        -----

          The standard deviation is the square root of the average of the
          squared deviations from the mean, i.e. var = sqrt(mean((x -
          x.mean())**2)).  The computed standard deviation is biased, i.e., the
          mean is computed by dividing by the number of elements, N, rather
          than by N-1.

        """
        return N.ndarray.std(self, axis, dtype, out)._align(axis)

    def var(self, axis=None, dtype=None, out=None):
        """Compute the variance along the specified axis.

        Returns the variance of the array elements, a measure of the spread of
        a distribution.  The variance is computed for the flattened array by
        default, otherwise over the specified axis.

        :Parameters:

            axis : integer
                Axis along which the variance is computed. The default is to
                compute the variance of the flattened array.

            dtype : type
                Type to use in computing the variance. For arrays of integer
                type the default is float32, for arrays of float types it is
                the same as the array type.

            out : ndarray
                Alternative output array in which to place the result. It must
                have the same shape as the expected output but the type will be
                cast if necessary.

        :Returns:

            variance : depends, see above
                A new array holding the result is returned unless out is
                specified, in which case a reference to out is returned.

        :SeeAlso:

            - std : standard deviation
            - mean : average

        Notes
        -----

          The variance is the average of the squared deviations from the mean,
          i.e.  var = mean((x - x.mean())**2).  The computed variance is
          biased, i.e., the mean is computed by dividing by the number of
          elements, N, rather than by N-1.

        """
        return N.ndarray.var(self, axis, dtype, out)._align(axis)

    def prod(self, axis=None, dtype=None, out=None):
        return N.ndarray.prod(self, axis, dtype, out)._align(axis)

    def any(self, axis=None, out=None):
        return N.ndarray.any(self, axis, out)._align(axis)

    def all(self, axis=None, out=None):
        return N.ndarray.all(self, axis, out)._align(axis)

    def max(self, axis=None, out=None):
        return N.ndarray.max(self, axis, out)._align(axis)

    def argmax(self, axis=None, out=None):
        return N.ndarray.argmax(self, axis, out)._align(axis)

    def min(self, axis=None, out=None):
        return N.ndarray.min(self, axis, out)._align(axis)

    def argmin(self, axis=None, out=None):
        return N.ndarray.argmin(self, axis, out)._align(axis)

    def ptp(self, axis=None, out=None):
        return N.ndarray.ptp(self, axis, out)._align(axis)

    # Needed becase tolist method expects a[i]
    #  to have dimension a.ndim-1
    def tolist(self):
        return self.__array__().tolist()

    def getI(self):
        M,N = self.shape
        if M == N:
            from numpy.dual import inv as func
        else:
            from numpy.dual import pinv as func
        return asmatrix(func(self))

    def getA(self):
        return self.__array__()

    def getA1(self):
        return self.__array__().ravel()

    def getT(self):
        return self.transpose()

    def getH(self):
        if issubclass(self.dtype.type, N.complexfloating):
            return self.transpose().conjugate()
        else:
            return self.transpose()

    T = property(getT, None, doc="transpose")
    A = property(getA, None, doc="base array")
    A1 = property(getA1, None, doc="1-d base array")
    H = property(getH, None, doc="hermitian (conjugate) transpose")
    I = property(getI, None, doc="inverse")

def _from_string(str,gdict,ldict):
    rows = str.split(';')
    rowtup = []
    for row in rows:
        trow = row.split(',')
        newrow = []
        for x in trow:
            newrow.extend(x.split())
        trow = newrow
        coltup = []
        for col in trow:
            col = col.strip()
            try:
                thismat = ldict[col]
            except KeyError:
                try:
                    thismat = gdict[col]
                except KeyError:
                    raise KeyError, "%s not found" % (col,)

            coltup.append(thismat)
        rowtup.append(concatenate(coltup,axis=-1))
    return concatenate(rowtup,axis=0)


def bmat(obj, ldict=None, gdict=None):
    """Build a matrix object from string, nested sequence, or array.

    Ex:  F = bmat('A, B; C, D')
         F = bmat([[A,B],[C,D]])
         F = bmat(r_[c_[A,B],c_[C,D]])

        all produce the same Matrix Object    [ A  B ]
                                              [ C  D ]

        if A, B, C, and D are appropriately shaped 2-d arrays.
    """
    if isinstance(obj, str):
        if gdict is None:
            # get previous frame
            frame = sys._getframe().f_back
            glob_dict = frame.f_globals
            loc_dict = frame.f_locals
        else:
            glob_dict = gdict
            loc_dict = ldict

        return matrix(_from_string(obj, glob_dict, loc_dict))

    if isinstance(obj, (tuple, list)):
        # [[A,B],[C,D]]
        arr_rows = []
        for row in obj:
            if isinstance(row, N.ndarray):  # not 2-d
                return matrix(concatenate(obj,axis=-1))
            else:
                arr_rows.append(concatenate(row,axis=-1))
        return matrix(concatenate(arr_rows,axis=0))
    if isinstance(obj, N.ndarray):
        return matrix(obj)

mat = asmatrix