How to support modified data interpretations in NumPy ndarrays?

I am trying to write a Python 3 class that stores some data in a NumPy np.ndarray. However, I want my class to also contain a piece of information about how to interpret the data values.

For example, let's assume the dtype of the ndarray is np.float32, but there is also a "color" that modifies the meaning of those floating-point values. So, if I want to add a red number and a blue number, I must first convert both numbers to magenta in order to legally add their underlying _data arrays. The result of the addition will then have _color = "magenta".

This is just a toy example. In reality, the "color" is not a string (it's better to think of it as an integer), the "color" of the result is mathematically determined from the "color" of the two inputs, and the conversion between any two "colors" is mathematically defined.

class MyClass:
    
    def __init__(self, data : np.ndarray, color : str):
        self._data = data
        self._color = color
    
    
    # Example: Adding red numbers and blue numbers produces magenta numbers
    def convert(self, other_color):
        if self._color == "red" and other_color == "blue":
            return MyClass(10*self._data, "magenta")
        elif self._color == "blue" and other_color == "red":
            return MyClass(self._data/10, "magenta")
    
    
    def __add__(self, other):
        if other._color == self._color:
            # If the colors match, then just add the data values
            return MyClass(self._data + other._data, self._color)
        else:
            # If the colors don't match, then convert to the output color before adding
            new_self = self.convert(other._color)
            new_other = other.convert(self._color)
            return new_self + new_other

My problem is that the _color information lives alongside the _data. So, I can't seem to define sensible indexing behavior for my class:

• If I define __getitem__ to return self._data[i], then the _color information is lost.
• If I define __getitem__ to return MyClass(self._data[i], self._color) then I'm creating a new object that contains a scalar number. This will cause plenty of problems (for example, I can legally index that_object[i], leading to certain error.
• If I define __getitem__ to return MyClass(self._data[i:i+1], self._color) then I'm indexing an array to get an array, which leads to plenty of other problems. For example, my_object[i] = my_object[i] looks sensible, but would throw an error.

I then started thinking that what I really want is a different dtype for each different "color". That way, the indexed value would have the "color" information encoded for free in the dtype... but I don't know how to implement that.

The theoretical total number of "colors" is likely to be roughly 100,000. However, fewer than 100 would be used in any single script execution. So, I guess it may be possible to maintain a list/dictionary/? of the used "colors" and how they map to dynamically generated classes ... but Python tends to quietly convert types in ways I don't expect, so that is probably not the right path to go down.

All I know is that I don't want to store the "color" alongside every data value. The data arrays can be ~billions of entries, with one "color" for all entries.

How can I keep track of this "color" information, while also having a usable class?

It's not tenable to define every dunder (__add__, etc.) It's also probably not tenable to inherit from np.ndarray, which is what it would take for a compatible derived class.

You can get away with a thin wrapper:

from typing import NamedTuple, Sequence, Any, Callable

import numpy as np

Colour = int
RED: Colour = 0x0000FF
MAGENTA: Colour = 0xFF00FF
BLUE: Colour = 0xFF0000


def common_colour(colours: Sequence[Colour]) -> Colour:
    # magic happens here
    return sum(colours)


class ColouredArray(NamedTuple):
    colour: Colour
    data: np.ndarray

    def __str__(self) -> str:
        return f'({self.colour}) {self.data}'

    def convert(self, new_colour: Colour) -> 'ColouredArray':
        return ColouredArray(
            # magic happens here
            data=self.data * new_colour/self.colour,
            colour=new_colour,
        )


def all_common(arrays: Sequence['ColouredArray']) -> tuple['ColouredArray']:
    new_colour = common_colour([a.colour for a in arrays])
    return tuple(
        array.convert(new_colour) for array in arrays
    )


def call_common(method: Callable, *args, **kwargs) -> tuple[Colour, Any]:
    new_colour = common_colour([
        arg.colour
        for arg in (*args, *kwargs.values())
        if isinstance(arg, ColouredArray)
    ])
    return new_colour, method(
        *(
            arg.convert(new_colour).data if isinstance(arg, ColouredArray) else arg
            for arg in args
        ),
        **{
            k: arg.convert(new_colour).data if isinstance(arg, ColouredArray) else arg
            for k, arg in kwargs.items()
        },
    )


y = ColouredArray(*call_common(
    np.interp,
    x=ColouredArray(RED, np.arange(5)),
    xp=ColouredArray(RED, np.arange(1, 30, 2)),
    fp=ColouredArray(BLUE, np.arange(11, 40, 2)),
    left=-1,
))
print(y)

In this example, the tuple from call_common is unpacked directly into the ColouredArray constructor because interp returns a single np.ndarray. In other situations such as the 4-tuple returned from numpy.linalg.lstsq it would be up to the caller to unpack and reconstruct a coloured array as needed.