N-dimensional diagonal tensor calculation

ndiag.py

# convert from N-D coordinates to 1-D index                                                                                         
def ravel(coord,stride):
    idx = 0
    # loop and multiply each coordinate by the stride                                                                               
    for i in range(len(coord)):
        idx += coord[i] * stride[i]
    return idx

# convert from a 1-D index to N-D coordinates                                                                                       
def unravel(idx,stride):
    # prepare coordinates vector                                                                                                    
    ret = [0] * len(stride)
    # loop and divide index by stride, saving remainder for next dimension                                                          
    for i in range(len(stride)):
        ret[i] = idx / stride[i]
        idx -= ret[i] * stride[i]
    # convert vector to tuple for indexing                                                                                          
    return tuple(ret)

# copy from a to b along the given stride                                                                                           
def copy(a,b,diag_stride):
    # compute the word-aligned strides (dividing by bytes-per-word)                                                                 
    a_stride = np.divide(a.strides,a.dtype.itemsize)
    b_stride = np.divide(b.strides,b.dtype.itemsize)
    # loop and convert an index in b-space to a-space                                                                               
    for b_idx in range(b.size):
        # convert an index in b-space to coordinates                                                                                
        b_coord = unravel(b_idx,b_stride)
        # convert b-space coordinates given the diagonal stride to a-space index                                                    
        a_idx = ravel(b_coord,diag_stride)
        # convert a-space index to a-space coordinates                                                                              
        a_coord = unravel(a_idx,a_stride)
        # assign value from a to b                                                                                                  
        b[b_coord] = a[a_coord]
        
# calculate the diagonal of the given tensor                                                                                        
# arr - input tensor                                                                                                                
# axis0 - the first axis to compute the diagonal over                                                                               
# axis1 - the second axis to compute the diagonal over                                                                              
# returns - diagonal tensor                                                                                                         
def diag(arr,axis0=0,axis1=1):
    # check that there are at least 2 dimensions, there is no 1D diagonal                                                           
    ndim = len(arr.shape)
    if ndim < 2:
        print('error! requries 2+ dimensions')
        return None

    # make sure axis0 is less than axis1, if not swap them                                                                          
    if axis0 < axis1:
        a0 = axis0
        a1 = axis1
    elif axis0 > axis1:
        a0 = axis1
        a1 = axis0
    else:
        print('error! cannot calculate diagonal along the same axes: axis0=%d  axis1=%d' % (axis0,axis1))
        return None

    # get the size of the two axes that we will calculate the diagonal along                                                        
    dim0 = arr.shape[a0]
    dim1 = arr.shape[a1]

    # get the stride of indexing along each axis                                                                                    
    stride0 = arr.strides[a0] / arr.dtype.itemsize
    stride1 = arr.strides[a1] / arr.dtype.itemsize

    # determine the diagonal length                                                                                                 
    diag_size = dim0 if dim0 < dim1 else dim1

    # calculate output diagonal tensor shape                                                                                        
    ret_shape = [0] * (ndim-1)
    ret_strides = [0] * (ndim-1)
    idx = 0

    # loop over all the dimensions in the input tensor and add all axes to the output tensor that are not the ones we're calculating the diagonal along                                                                                                                
    for i in range(ndim):
        if i != a0 and i != a1:
            ret_shape[idx] = arr.shape[i]
            ret_strides[idx] = arr.strides[i] / arr.dtype.itemsize
            idx += 1
    # add the diagonal shape as the last axis in the output tensor                                                                  
    ret_shape[ndim-2] = diag_size
    ret_strides[ndim-2] = stride0 + stride1

    # construct the output tensor                                                                                                   
    ret = np.zeros(ret_shape,dtype=arr.dtype)

    # copy the values along the diagonal from the input tensor to the output tensor                                                 
    copy(arr,ret,ret_strides)

    # return the diagonal tensor                                                                                                    
    return ret