Is there an efficient way to pass "all" as a numpy index?-CodePudding

I have code that generates a boolean array that acts as a mask on numpy arrays, along the lines of:

def func():
    a    = numpy.arange(10)
    mask = a % 2 == 0
    return a[mask]

Now, I need to separate this into a case where the mask is created, and one where it is not created and all values are used instead. This could be achieved as follows:

def func(use_mask):
    a    = numpy.arange(10)
    if use_mask:
        mask = a % 2 == 0
    else:
        mask = numpy.ones(10, dtype=bool)
    return a[mask]

However, this becomes extremely wasteful for large arrays, since an equally large boolean array must first be created.

My question is thus: Is there something I can pass as an "index" to recreate the behavior of such an everywhere-true array?

Systematically changing occurrences of a[mask] to something else involving some indexing magic etc. is a valid solution, but just avoiding the masking entirely via an expanded case distinction or something else that changes the structure of the code is not desired, as it would impair readability and maintainability (see next paragraph).

For the sake of completeness, here's what I'm currently considering doing, though this makes the code messier and less streamlined since it expands the if/else beyond where it technically needs to be (in reality, the mask is used more than once, hence every occurrence would need to be contained within the case distinction; I used f1 and f2 as examples here):

def func(use_mask):
    a    = numpy.arange(10)
    if use_mask:
        mask = a % 2 == 0
        r   = f1(a[mask])
        q   = f2(a[mask], r)
        return q
    else:
        r   = f1(a)
        q   = f2(a, r)
        return q

CodePudding user response：

Recall that a[:] returns the contents of a (even if a is multidimensional). We cannot store the : in the mask variable, but we can use a slice object equivalently:

def func(use_mask):
    a = numpy.arange(10)
    if use_mask:
        mask = a % 2 == 0
    else:
        mask = slice(None)
    return a[mask]

This does not use any memory to create the index array. I'm not sure what the CPU usage of the a[slice(None)] operation is, though.