[FIXED] How to change the font color of the specified elements in Numpy

Issue

I want to know how to change the font color of the specified elements to mark them in the numpy array. I referenced the answer to this question: Coloring entries in an Matrix/2D-numpy array? I write the following two functions to do this. For instance, in the following example, I want to change the color of a[0,1], a[1,2], a[2,0]. In the trans(a, j, *lst) function, a is the input matrix, j is the color number, and *lst is the index of the element to be modified.

import numpy as np
a = np.array([[0,0,0],
              [0,0,0],
              [23.7135132,0,0],
              [32.22954112,0,0],
              [71.01941585,31.27384681,0]])
print(a,'\n')

def trans(a,j,*lst):
    num_row = a.shape[0]
    num_column = a.shape[1]
    add = ''
    row = []
    column = []
    # record the row index and column index of each input (in *lst)
    for k in range(len(lst)):
        row.append(lst[k][0])
        column.append(lst[k][1])
        
    # Generate the same content as 'a' and change the color
    for m in range(num_row):
        for n in range(num_column):
            for k in range(len(lst)):
                # the element whose color needs to be modified
                if m == row[k] and n == column[k]:
                    add += "\033[3{}m{}\033[0m ".format(j, a[row[k],column[k]])
                    break
                # the element whose color does not need to be modified
                if k == len(lst)-1:
                    add += "\033[3{}m{}\033[0m ".format(0, a[m,n])
    # There is a space at the end of "\033[3{}m{}\033[0m "
    add = add.strip(' ')
    add = add.split(' ')
    add = np.array(add)
    add = add.reshape(num_row,num_column)
    return add

def show(a, row_sep=" "):
    n, m = a.shape
    fmt_str = "\n".join([row_sep.join(["{}"]*m)]*n)
    return fmt_str.format(*a.ravel())

print(show(trans(a,1,[0,1],[1,2],[2,0])))

But the results are as follows.

As you can see, the elements cannot be aligned, so I want to know how to automatically align these elements, just like the display of numpy array. Or is there any other way to directly modify the font color of numpy array? I just want to make some special elements more conspicuous.

Solution

Here's a very general version that wraps np.array2string(). It works with any dtype and/or custom formatters, and the output should look identical to np.array2string(). Instead of trying to parse the formatted string, it uses its own custom formatter to add escape sequences around the target elements. It also uses a str subclass so the length calculations for line wrapping ignore the escape sequences.

import re
import numpy as np

class EscapeString(str):
    """A string that excludes SGR escape sequences from its length."""
    def __len__(self):
        return len(re.sub(r"\033\[[\d:;]*m", "", self))
    def __add__(self, other):
        return EscapeString(super().__add__(other))
    def __radd__(self, other):
        return EscapeString(str.__add__(other, self))

def color_array(arr, color, *lst, **kwargs):
    """Takes the same keyword arguments as np.array2string()."""
    # adjust some kwarg name differences from _make_options_dict()
    names = {"max_line_width": "linewidth", "suppress_small": "suppress"}
    options_kwargs = {names.get(k, k): v for k, v in kwargs.items() if v is not None}
    # this comes from arrayprint.array2string()
    overrides = np.core.arrayprint._make_options_dict(**options_kwargs)
    options = np.get_printoptions()
    options.update(overrides)
    # from arrayprint._array2string()
    format_function = np.core.arrayprint._get_format_function(arr, **options)

    # convert input index lists to tuples
    target_indices = set(map(tuple, lst))

    def color_formatter(i):
        # convert flat index to coordinates
        idx = np.unravel_index(i, arr.shape)
        s = format_function(arr[idx])
        if idx in target_indices:
            return EscapeString(f"\033[{30+color}m{s}\033[0m")
        return s

    # array of indices into the flat array
    indices = np.arange(arr.size).reshape(arr.shape)
    kwargs["formatter"] = {"int": color_formatter}
    return np.array2string(indices, **kwargs)

Example:

np.set_printoptions(threshold=99, linewidth=50)
a = np.arange(100).reshape(5, 20)
print(a)
print(np.array2string(a, threshold=a.size))
print()
print(color_array(a, 1, [0, 1], [1, 3], [0, 12], [4, 17]))
print(color_array(a, 1, [0, 1], [1, 3], [0, 12], [4, 17], threshold=a.size))

Answered By - yut23