Search code examples
pythonnumpynumpy-ndarray

Pretty-print indices/coordinates of 2D Numpy array


Does Numpy provide built-in capabilities to print the indices/coordinates of a 2-d Numpy array at its borders?

What I mean is the following: Given, for example, the array

a = np.arange(6).reshape(2, 3).astype(float)

I would like to have a printout as follows:

     0    1    2
0  0.0  1.0  2.0
1  3.0  4.0  5.0

It needn't be exactly like that, but the corresponding information should be conveyed. That is, I want to have a "header column" that shows the row indices (here: 0, 1) and a "header row" that shows the column indices (here: 0, 1, 2).

At present, I simply use Pandas for this purpose (i.e. I convert the Numpy array to a Pandas DataFrame and then print the DataFrame). This is also suggested, for example, in this answer to a related question. However, this would always require Pandas, even if I don't use it for anything else in a given project.

So my question, again, is: can I achieve the same printout only with Numpy (thus without Pandas)? Or, if not, is there a lightweight solution (in the sense of a few lines of code without an additional library) to achieve a similar output?


Solution

  • You can try:

    import numpy as np
    
    def pretty_print(arr, precision=2, sep=1):
        # Convert array contents to list of list of formatted strings
        rows = [[f"{v:.{precision}f}" for v in row] for row in arr]
        # Calculate width of row index's column (maxl) and remaining columns (maxw)
        maxl = len(str(arr.shape[0] - 1))
        maxw = max(len(str(arr.shape[1] - 1)), *(len(v) for row in rows for v in row))
        # Print column indexes and content, pad string entries to appropriate width
        print(" " * maxl, *(f"{i:>{maxw}}" for i in range(0, arr.shape[1])), sep=" " * sep)
        for i, row in enumerate(rows):
            print(f"{i:>{maxl}}", *(f"{v:>{maxw}}" for v in row), sep=" " * sep)
    
    a = np.arange(12).reshape(3, 4).astype(float)
    a[0, 0] = np.inf
    a[1, 1] = -np.inf
    a[2, 3] *= -1
    pretty_print(a, precision=2, sep=1)
    pretty_print(np.ones(shape=(3, 14)), precision=0)
    

    Prints:

           0      1      2      3
    0    inf   1.00   2.00   3.00
    1   4.00   -inf   6.00   7.00
    2   8.00   9.00  10.00 -11.00
    

    and

       0  1  2  3  4  5  6  7  8  9 10 11 12 13
    0  1  1  1  1  1  1  1  1  1  1  1  1  1  1
    1  1  1  1  1  1  1  1  1  1  1  1  1  1  1
    2  1  1  1  1  1  1  1  1  1  1  1  1  1  1