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:

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.

CodePudding user response：

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))