Array protocol and __array__ in NumPy

Introduction

The array protocol lets objects tell NumPy how to convert themselves into arrays. The __array__ method is a simple way to make your object behave like a NumPy array.

You have a custom data object and want to use NumPy functions on it easily.

You want to convert your object to a NumPy array without writing extra code.

You want your object to work smoothly with libraries that expect NumPy arrays.

You want to speed up data processing by using NumPy's fast operations on your data.

Syntax

NumPy

class YourClass:
    def __array__(self, dtype=None):
        # Return a NumPy array representation of your data
        import numpy as np
        data = ...  # your data here
        if dtype is not None:
            return np.array(data, dtype=dtype)
        else:
            return np.array(data)

The __array__ method should return a NumPy array.

The optional dtype argument lets NumPy request a specific data type.

Examples

Converts a simple list inside a custom object to a NumPy array.

NumPy

import numpy as np

class SimpleData:
    def __init__(self, values):
        self.values = values
    def __array__(self, dtype=None):
        if dtype is not None:
            return np.array(self.values, dtype=dtype)
        return np.array(self.values)

obj = SimpleData([1, 2, 3])
print(np.array(obj))

Handles the edge case where the data is empty, returning an empty array.

NumPy

import numpy as np

class EmptyData:
    def __array__(self, dtype=None):
        return np.array([])

obj = EmptyData()
print(np.array(obj))

Handles the case with only one element, returning a one-element array.

NumPy

import numpy as np

class SingleValueData:
    def __init__(self, value):
        self.value = value
    def __array__(self, dtype=None):
        return np.array([self.value], dtype=dtype) if dtype is not None else np.array([self.value])

obj = SingleValueData(42)
print(np.array(obj))

Converts a list with values at the beginning and end to a NumPy array.

NumPy

import numpy as np

class DataAtEdges:
    def __init__(self, values):
        self.values = values
    def __array__(self, dtype=None):
        return np.array(self.values, dtype=dtype) if dtype is not None else np.array(self.values)

obj = DataAtEdges([0, 1, 2, 3, 4])
print(np.array(obj))

Sample Program

This program shows how a custom object with a __array__ method can be converted to a NumPy array automatically. It prints the original data, the converted array, and calculates the mean using NumPy.

NumPy

import numpy as np

class CustomData:
    def __init__(self, data_list):
        self.data_list = data_list
    def __array__(self, dtype=None):
        # Convert internal list to numpy array
        if dtype is not None:
            return np.array(self.data_list, dtype=dtype)
        return np.array(self.data_list)

# Create an object with some data
custom_object = CustomData([10, 20, 30, 40])

# Print the object data before conversion
print('Original data in object:', custom_object.data_list)

# Convert to numpy array using the array protocol
numpy_array = np.array(custom_object)
print('Converted to numpy array:', numpy_array)

# Use numpy functions on the converted array
mean_value = np.mean(custom_object)
print('Mean of data:', mean_value)

Important Notes

The __array__ method lets NumPy know how to convert your object to an array quickly.

Time complexity depends on the size of your data because it copies data into a new array.

Space complexity is also proportional to the data size since a new array is created.

A common mistake is forgetting to handle the optional dtype argument, which can cause errors when NumPy requests a specific type.

Use __array__ when you want seamless integration with NumPy functions. For simple conversions, you might also use np.array(your_data) directly if your data is already array-like.

Summary

The __array__ method allows custom objects to behave like NumPy arrays.

It helps NumPy convert your data automatically for fast numerical operations.

Remember to handle the optional dtype argument for flexibility.