Offset values of tensor by different amounts

1

Hi everyone,

I have a challenging problem. At least I find it challenging!
For an ecology problem, I have a grid (matrix) representing my environment, and values on each cell representing the local population.
Every now and then, the population disperses: Each individual in each cell moves left or right (not up or down at this early stage) by a distance that is sampled from a normal distribution. I would like to vectorize the dispersion operation, at least for one iteration (i.e. for one individual per occupied cell).

For instance, let’s assume that before the dispersion, I have the grid below:

|  6  8 10  7  5  0  0  |
|  5  7  9 11  8  7  3  |
|  4  6  6  8  5  4  1  |

After one iteration, let’s say that one individual in the first row moved +2 (to the right), one individual in the second row by -1 (to the left) and two individuals in the third row moved by -3 and +1 respectively. These movements could be represented by a matrix as below, where each value represents the distance travelled by the first individual at that position (obviousy this would have to be repeated for the other individuals in each position):

|  0  0  0  0  2  0  0  |
|  0 -1  0  0  0  0  0  |
|  0  0  0 -3  0  1  0  |

This is what the new environment could look like:

|  6  8 10  7  4  0  1  |
|  6  6  9 11  8  7  3  |
|  5  6  6  7  5  3  2  |

Is there a way of providing such a matrix of offsets and modifying the original matrix without going though nested for loops? Note that the distance matrix approach is what I have in mind but there might be better ways of achieving this.

Thanks in advance!

2

This way:

import torch

population = torch.tensor([
                   [6, 8, 10,  7, 5, 0, 0],
                   [5, 7,  9, 11, 8, 7, 3],
                   [4, 6,  6,  8, 5, 4, 1]
             ]
)

moves = torch.tensor([
                   [0, 0,  0,  0, 2, 0, 0],
                   [0, 0, -1,  0, 0, 0, 0],
                   [0, 0,  0, -3, 0, 1, 0]
             ]
)

right = (moves > 0).nonzero(as_tuple=True)
left = (moves < 0).nonzero(as_tuple=True)

mask = torch.zeros_like(population)

mask[right] = -1
mask[left] = -1

mask[right[0], right[1] + moves[right]] = 1
mask[left[0], left[1] + moves[left]] = 1

population += mask

print(population)
1 Like
3

Brilliant, thanks so much!

1 Like