Numpy 2d and 1d array to latex bmatrix

Question

I'm looking for a clean way to migrate numpy arrays to latex bmatrix. It should work for both 2d arrays and horizontal and vertical 1d array.

Example

A = array([[12, 5, 2],
           [20, 4, 8],
           [ 2, 4, 3],
           [ 7, 1,10]])

print A              #2d array
print A[0]           #horizontal array
print A[:,0, None]   #vertical array

array_to_bmatrix(A)
array_to_bmatrix(A[0])
array_to_bmatrix(A[:,0, None])

Out:

[[12  5  2]
 [20  4  8]
 [ 2  4  3]
 [ 7  1 10]]

[12  5  2]

[[12]
 [20]
 [ 2]
 [ 7]]

\begin{bmatrix} 
 12.000 & 5.000 & 2.000 & \\
 20.000 & 4.000 & 8.000 & \\
 2.000 & 4.000 & 3.000 & \\
 7.000 & 1.000 & 10.000 & \\
\end{bmatrix}

\begin{bmatrix} 
 12.000 & 5.000 & 2.000
\end{bmatrix}

\begin{bmatrix} 
 12.000 & \\
 20.000 & \\
 2.000 & \\
 7.000 & \\
\end{bmatrix}

Attempt of solution

def array_to_bmatrix(array):
    begin = '\\begin{bmatrix} \n'
    data = ''
    for line in array:        
        if line.size == 1:
            data = data + ' %.3f &'%line
            data = data + r' \\'
            data = data + '\n'
            continue
        for element in line:
            data = data + ' %.3f &'%element

        data = data + r' \\'
        data = data + '\n'
    end = '\end{bmatrix}'
    print begin + data + end  

This solution works for vertical and 2d arrays, however it outputs horizontal arrays as vertical ones.

array_to_bmatrix(A[0])

Out:

\begin{bmatrix} 
 12.000 & \\
 5.000 & \\
 2.000 & \\
\end{bmatrix}

Answer 1

The __str__ method of the numpy array already does most of the formatting for you. Let's exploit that;

import numpy as np

def bmatrix(a):
    """Returns a LaTeX bmatrix

    :a: numpy array
    :returns: LaTeX bmatrix as a string
    """
    if len(a.shape) > 2:
        raise ValueError('bmatrix can at most display two dimensions')
    lines = str(a).replace('[', '').replace(']', '').splitlines()
    rv = [r'\begin{bmatrix}']
    rv += ['  ' + ' & '.join(l.split()) + r'\\' for l in lines]
    rv +=  [r'\end{bmatrix}']
    return '\n'.join(rv)

A = np.array([[12, 5, 2], [20, 4, 8], [ 2, 4, 3], [ 7, 1, 10]])
print bmatrix(A) + '\n'

B = np.array([[1.2], [3.7], [0.2]])
print bmatrix(B) + '\n'

C = np.array([1.2, 9.3, 0.6, -2.1])
print bmatrix(C) + '\n'

This returns:

\begin{bmatrix}
  12 & 5 & 2\\
  20 & 4 & 8\\
  2 & 4 & 3\\
  7 & 1 & 10\\
\end{bmatrix}

\begin{bmatrix}
  1.2\\
  3.7\\
  0.2\\
\end{bmatrix}

\begin{bmatrix}
  1.2 & 9.3 & 0.6 & -2.1\\
\end{bmatrix}