Have numpy.concatenate return proper subclass rather than plain ndarray

Question

I have a numpy array subclass, and I'd like to be able to concatenate them.

import numpy as np

class BreakfastArray(np.ndarray):
    def __new__(cls, n=1):
        dtypes=[("waffles", int), ("eggs", int)]
        obj = np.zeros(n, dtype=dtypes).view(cls)
        return obj
        
b1 = BreakfastArray(n=1)
b2 = BreakfastArray(n=2)
con_b1b2 = np.concatenate([b1, b2])

print(b1.__class__, con_b1b2.__class__)

this outputs <class '__main__.BreakfastArray'> <class 'numpy.ndarray'>, but I'd like the concatenated array to also be a BreakfastArray class. It looks like I probably need to add a __array_finalize__ method, but I can't figure out the right way to do it.

Barmar · Accepted Answer · 2024-12-14 01:03:22Z

1

Create con_bib2 as a new BreakfastArray, then use out=con_bib2 when concatenating to store into it.

con_bib2 = BreakfastArray(n=3)
np.concatenate([b1, b2], out=con_bib2)

answered Dec 14, 2024 at 1:03

Barmar

789k57 gold badges555 silver badges669 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

simon · Accepted Answer · 2024-12-16 15:00:41Z

Short answer:

If you want to be pragmatic, use Barmar's approach.
If you are insisting, then override __array_function__ rather than __array_finalize__.

Long answer:

Once you provide your own implementation of BreakfastArray.__array_finalize__, such as the following:

    def __array_finalize__(self, obj):
        print("finalize", f"self: {type(self)}", f"obj: {type(obj)}")
        return super().__array_finalize__(obj)

… then you will realize that with your current code, it is only called twice, namely once for each instance creation (b1, b2) – but not for the concatenation.

What will be called, however, is __array_function__. We can directly follow the example from its documentation and adapt it to your needs:

import numpy as np

HANDLED_FUNCTIONS = {}

class BreakfastArray(np.ndarray):
    def __new__(cls, n=1):
        dtypes=[("waffles", int), ("eggs", int)]
        obj = np.zeros(n, dtype=dtypes).view(cls)
        return obj
    
    def __array_function__(self, func, types, args, kwargs):
        if func not in HANDLED_FUNCTIONS:
            return NotImplemented
        if not all(issubclass(t, BreakfastArray) for t in types):
            return NotImplemented
        return HANDLED_FUNCTIONS[func](*args, **kwargs)

def implements(numpy_function):
    def decorator(func):
        HANDLED_FUNCTIONS[numpy_function] = func
        return func
    return decorator

@implements(np.concatenate)
def concatenate(arrays):
    result = BreakfastArray(n=sum(len(a) for a in arrays))
    return np.concatenate([np.asarray(a) for a in arrays], out=result)
        
b1 = BreakfastArray(n=1)
b2 = BreakfastArray(n=2)
con_b1b2 = np.concatenate([b1, b2])

Provide a dictionary HANDLED_FUNCTIONS to hold custom implementations of Numpy functions for BreakfastArray.
Provide a decorator implements to store the implementations into the dictionary of custom implementations.
Provide own implementation of __array_function__ to access the dictionary of custom implementations.
Provide own implementation of concatenate to be stored in the dictionary of custom implementations. Note that here I call np.asarray() on the arrays to be concatenated, to avoid an infinite recursion into our own concatenate implementation; also, note that the current implementation does not support axis or out parameters.

If we only ever want to override concatenate, this can be simplified to:

import numpy as np

class BreakfastArray(np.ndarray):
    def __new__(cls, n=1):
        dtypes=[("waffles", int), ("eggs", int)]
        obj = np.zeros(n, dtype=dtypes).view(cls)
        return obj
    
    def __array_function__(self, func, types, args, kwargs):
        if func != np.concatenate or not all(issubclass(t, BreakfastArray) for t in types):
            return NotImplemented
        arrays = args[0]
        result = BreakfastArray(n=sum(len(a) for a in arrays))
        return np.concatenate([np.asarray(a) for a in arrays], out=result)
        
b1 = BreakfastArray(n=1)
b2 = BreakfastArray(n=2)
con_b1b2 = np.concatenate([b1, b2])

I.P. Freeley · Accepted Answer · 2024-12-16 21:53:25Z

Expanding simon's solution, this is what I settled on so other numpy functions fall-back to standard ndarray (so, numpy.unique(b2["waffles"]) works as expected). Also a slight change to concatenate so it will work for any subclasses as well.

import numpy as np

HANDLED_FUNCTIONS = {}

class BreakfastArray(np.ndarray):
    def __new__(cls, *args, n=1, **kwargs):
        dtypes=[("waffles", int), ("eggs", int)]
        obj = np.zeros(n, dtype=dtypes).view(cls)
        return obj

    def __array_function__(self, func, types, args, kwargs):
        # If we want "standard numpy behavior",
        # convert any BreakfastArray to ndarray views
        if func not in HANDLED_FUNCTIONS:
            new_args = []
            for arg in args:
                if issubclass(arg.__class__, BreakfastArray):
                    new_args.append(arg.view(np.ndarray))
                else:
                    new_args.append(arg)
            return func(*new_args, **kwargs)
        if not all(issubclass(t, BreakfastArray) for t in types):
            return NotImplemented
        return HANDLED_FUNCTIONS[func](*args, **kwargs)

def implements(numpy_function):
    def decorator(func):
        HANDLED_FUNCTIONS[numpy_function] = func
        return func
    return decorator

@implements(np.concatenate)
def concatenate(arrays):
    result = arrays[0].__class__(n=sum(len(a) for a in arrays))
    return np.concatenate([np.asarray(a) for a in arrays], out=result)

Collectives™ on Stack Overflow

Have numpy.concatenate return proper subclass rather than plain ndarray

3 Answers 3

Comments

Short answer:

Long answer:

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

Comments

Short answer:

Long answer:

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related