Python - 创建一个 class 作为二维 numpy 数组的索引

Question

我有一些索引代码 class TwoDimIndex。我想用它来索引像 arr[idx] 这样的 numpy 二维数组。下面是 class.

的代码

import numpy as np

class TwoDimIndex:
    def __init__(self, i1, i2):
        self.i1 = i1
        self.i2 = i2
        pass

    def __index__(self):
        return self.i1, self.i2 
        # return np.array([self.i1, self.i2]) - Tried this too
        
    def __eq__(self, other):
        return self.i1 == other
        
    def __int__(self):
        return self.i1
        
    def __hash__(self):
        return hash(self.__int__())

# Don't edit code of this function please
def useful_func():
    idx = TwoDimIndex(1, 1) # Can edit create instance
    arr_two_dim = np.array([[0, 1], [2, 3]])
    print(idx.__index__() == (1, 1))
    print(arr_two_dim[1, 1], arr_two_dim[idx.__index__()]) # Success
    print(arr_two_dim[idx]) # Crash here
    # But I want this code work!!!

useful_func()

Class TwoDimIndex 用作索引，例如 arr[TwoDimIndex()].

Jupyter Notebook 和 Python 3.8 的所有代码。但是执行这段代码时出现错误。

IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices

有什么方法可以使 class TwoDimIndex 的实例成为 numpy 二维数组索引？

Answer 1

我找到了从元组继承的简短解决方案。

class TwoDimIndex(tuple):
    def __init__(self, tup):
        self.tup = tuple(tup)
        pass

    def __index__(self):
        return self.tup

def useful_func():
    idx = TwoDimIndex([1, 1]) # Edit at create class
    arr_two_dim = np.array([[0, 1], [2, 3]])
    print(idx.__index__() == (1, 1))
    print(arr_two_dim[1, 1], arr_two_dim[idx.__index__()]) # Success
    print(arr_two_dim[idx]) # Success now

useful_func()

不确定这是正确的方法，但它有效。

Python - 创建一个 class 作为二维 numpy 数组的索引

Python - create a class as index for 2-dimension numpy array

python-3.x

numpy

class

indexing

jupyter