Creating a 4 dimensional tensor mask in numpy

Question

While working on a research problem, I encountered the following rank 4 tensor mask:

M[i,j,k,l] = delta(i+j,k+l)

Where delta(x,y) = 1 iff x==y, else 0 (the kronecker delta)

For example M(2,3,1,4) = 1 since (2+3=1+4) but M(1,3,2,4)= 0 since (1+3 is not 2+4)

Is there a vectorized way to create this using numpy? If I used for loops, I have to go through N^4 iterations which for my current value of N is already too much. Also there's an issue of memory.

Latex does not work here. Better share a small example and the expected output — yatu, Apr 12 '20 at 17:26

score 0 · Answer 1 · answered Dec 25 '20 at 19:06

Adapting the suggestion from https://stackoverflow.com/a/18860925/2062663, you could try using Numpy's meshgrid as follows:

import numpy as np


def get_mask(shape):
    i, j, k, l = np.meshgrid(*map(np.arange, shape),
                             sparse=1,
                             indexing='ij')
    return i + j == k + l

Here are a few test cases:

def test_1():
    expected = np.array([[[[1]]]])
    actual = get_mask((1, 1, 1, 1))
    assert actual == expected


def test_2():
    expected = np.array([[[[1, 0], [0, 0]],
                          [[0, 1], [1, 0]]],
                         [[[0, 1], [1, 0]],
                          [[0, 0], [0, 1]]]])

    actual = get_mask((2, 2, 2, 2))
    assert (actual == expected).all()


def test_4():
    actual = get_mask((4, 4, 4, 4))
    assert actual[1, 2, 0, 3] == 1
    assert actual[1, 1, 2, 0] == 1
    assert actual[0, 2, 1, 3] == 0
    assert actual[3, 2, 0, 3] == 0

Creating a 4 dimensional tensor mask in numpy

1 Answers1