Accessing a matrix element by matrix[(a, b), c] instead of matrix[a, b, c].

Question

I want to achieve the following:

1. Have a AxBxC matrix (where A,B,C are integers).
2. Access that matrix not as matrix[a, b, c] but as matrix[(a, b), c], this is, I have two variables, var1 = (x, y) and var2 = z and want access my matrix as matrix[var1, var2].

How can this be done? I am using numpy matrix, if it makes any difference.

I know I could use matrix[var1[0], var1[1], var2], but if possible I'd like to know if there is any other more elegant way.

Thanks!

Answer 1

If var1 = (x,y), and var2 = z, you can use

matrix[var1][var2]

Answer 2

I think you can simply subclass the NumPy matrix type, with a new class of your own; and overload the __getitem__() nethod to accept a tuple. Something like this:

class SpecialMatrix(np.matrix):
    def __getitem__(self, arg1, arg2, arg3=None):
        try:
            i, j = arg1
            k = arg2
            assert(arg3 is None)
            x = super(SpecialMatrix, self).__getitem__(i, j, k)
        except TypeError:
            assert(arg3 is not None)
            return super(SpecialMatrix, self).__getitem__(arg1, arg2, arg3)

And do something similar with __setitem__().

I'm not sure if __getitem__() takes multiple arguments like I'm showing here, or if it takes a tuple, or what. I don't have NumPy available as I write this answer, sorry.

EDIT: I re-wrote the example to use super() instead of directly calling the base class. It has been a while since I did anything with subclassing in Python.

EDIT: I just looked at the accepted answer. That's totally the way to do it. I'll leave this up in case anyone finds it educational, but the simple way is best.