[PYTHON MATRIX-SIG] Diagonals and traces

[Konrad Hinsen]

I have made an attempt at producing a more useful version and
diagonal() (and added the trivial trace()), which I herewith
submit to the scrutiny of the high court ;-)

Features: It takes two axis numbers (defaulting to 0 and 1) plus an
offset (defaulting to zero) for picking the parallels. Its return
value has a rank of one less then the input, and its first axis
corresponds to the diagonal, i.e. it equals
array([a[0,0], a[1,1], ...]), assuming the default axis values.

This version is incompatible with the current one in two respects:

1) In case you want diagonal-parallels, you have to specify the
   axis arguments.

2) For arrays of rank > 2, you have to be aware of the different
   axis defaults (-2 and -1 in Numeric, 0 and 1 in my version)
   and the different return value (diagonal along axis -1 in
   Numeric, along 0 in my version). This choice is more consistent
   with the other structural operations, which default to axis=0.

And here's the code:

------------------------------------------------------------
def diagonal(a, axis1=0, axis2=1, offset=0):
    a = array(a)
    if axis2 < axis1: axis1, axis2 = axis2, axis1
    if axis2 > 1:
	new_axes = range(len(a.shape))
	del new_axes[axis2]; del new_axes[axis1]
	new_axes[0:0] = [axis1, axis2]
	a = transpose(a, new_axes)
    s = a.shape
    n1 = s[0]
    n2 = s[1]
    n = n1*n2
    s = (n,) + s[2:]
    a = reshape(a, s)
    if offset < 0: offset = n2-offset-1
    return take(a, arange(offset, n, n2+1), 0)

def trace(a, axis1=0, axis2=1, offset=0):
    return add.reduce(diagonal(a, axis1, axis2, offset))
------------------------------------------------------------

Konrad.
-- 
-------------------------------------------------------------------------------
Konrad Hinsen                          | E-Mail: hinsen@ibs.ibs.fr
Laboratoire de Dynamique Moleculaire   | Tel.: +33-76.88.99.28
Institut de Biologie Structurale       | Fax:  +33-76.88.54.94
41, av. des Martyrs                    | Deutsch/Esperanto/English/
38027 Grenoble Cedex 1, France         | Nederlands/Francais
-------------------------------------------------------------------------------

=================
MATRIX-SIG  - SIG on Matrix Math for Python

send messages to: matrix-sig@python.org
administrivia to: matrix-sig-request@python.org
=================