Replace column by 0 based on probability

Question

How to replace column in the numpy array be certain number based on probability, if it is (1,X,X) shape. I found code to replace rows, but cannot figure out how to modify it, so it is applicable for columns replacement.

grid_example = np.random.rand(1,5,5)
probs = np.random.random((1,5))
grid_example[probs < 0.25] = 0
grid_example

Thanks!

Answer 1

Use:

import numpy as np

rng = np.random.default_rng(42)

grid_example = rng.random((1, 5, 5))
probs = rng.random((1, 5))

grid_example[..., (probs < 0.25).flatten()] = 0
print(grid_example)

Output

[[[0.         0.43887844 0.         0.         0.09417735]
  [0.         0.7611397  0.         0.         0.45038594]
  [0.         0.92676499 0.         0.         0.4434142 ]
  [0.         0.55458479 0.         0.         0.6316644 ]
  [0.         0.35452597 0.         0.         0.7783835 ]]]

The notation [..., (probs < 0.25).flatten()] applies the boolean indexing to the last index. More on the documentation.