Function to sparsify a matrix given a specific block size

Question

Problem Statement

I am trying to write a function that would sparsify a matrix given a target sparsity and an argument called block_shape which defines the minimum size of zeros block in the matrix. The target doesn't have to be met perfectly, but as close as possible.

For example, given the following arguments,

>>> matrix = [
  [1, 1, 1, 1],
  [1, 1, 1, 1],
  [1, 1, 1, 1],
  [1, 1, 1, 1]
]
>>> target = 0.5
>>> block_shape = (2, 2)

valid outputs of 50% sparsity could be

>>> sparse_matrix = sparsify(matrix, target, block_shape)
>>> sparse_matrix
[
  [1, 1, 0, 0],
  [1, 1, 0, 0],
  [0, 0, 1, 1],
  [0, 0, 1, 1]
]

>>> sparse_matrix = sparsify(matrix, target, block_shape)
>>> sparse_matrix
[
  [1, 0, 0, 1],
  [1, 0, 0, 1],
  [0, 0, 1, 1],
  [0, 0, 1, 1]
]

Note that there could be multiple valid sparsified versions of the input. The only criteris is to get to the target as much as possible. One of the constraints is that only the zeros of shape block_size are considered to be sparse.

For example, the matrix below has a sparsity level of 0%, given the arguments

>>> sparse_matrix = sparsify(matrix, target, block_shape)
>>> sparse_matrix
[
  [1, 0, 0, 1],
  [1, 1, 0, 0],
  [0, 1, 1, 1],
  [0, 0, 0, 0]
]

What I have so far

Currently, I have the following piece of code

import numpy as np

def sparsify(matrix, target, block_shape=None):
    if block_shape is None or block_shape == 1 or block_shape == (1,) or block_shape == (1, 1):
        # 1x1 is just bernoulli with p=target
        probs = np.random.uniform(size=matrix.shape)
        mask = np.zeros(matrix.shape)
        mask[probs >= target] = 1.0
    else:
        if isinstance(block_shape, int):
            block_shape = (block_shape, block_shape)
        if len(block_shape) == 1:
            block_shape = (block_shape[0], block_shape[0])
        mask = np.ones(matrix.shape)
        rows, cols = matrix.shape
        for row in range(rows):
            for col in range(cols):
                submask = mask[row:row+block_shape[0], col:col+block_shape[1]]
                if submask.shape != block_shape:
                    # we don't care about the edges, cannot partially sparsify
                    continue
                if (submask == 0).any():
                    # If current (row, col) is already in the sparsified area, skip
                    continue
                prob = np.random.random()
                if prob < target:
                    submask[:, :] = np.zeros(submask.shape)
    return matrix * mask, mask

The problem with the code above is that it does not match the target if the block size is not (1, 1)

>>> matrix = np.random.randn(100, 100)
>>> matrix, mask = sparsify(matrix, target=0.5, block_shape=(2, 2))

>>> print((matrix == 0).mean())
0.73
>>> print((mask == 0).mean())
0.73

Reason for discrepancy (I think)

I am not sure why I am not getting the target I expect, but I think it has something to do with the fact that I check the probability of every element, instead of the block as a whole. However, I have skipping conditions in my code, so I thought that should cover it

Edits

Edit 1 -- additional examples

Just giving some more examples.

Example 1: Given different block size

>>> sparse_matrix = sparsify(matrix, 0.25, (3, 3))
>>> sparse_matrix
[
  [0, 0, 0, 1],
  [0, 0, 0, 1],
  [0, 0, 0, 1],
  [1, 1, 1, 1]
]

The example above is a valid sparse matrix, although the level of sparsity is not 25%, another valid result could be a matrix of all 1's.

Example 2: Given a different block size and target

>>> sparse_matrix = sparsify(matrix, 0.6, (1, 2))
>>> sparse_matrix
[
  [0, 0, 0, 0],
  [1, 0, 0, 1],
  [0, 0, 1, 1],
  [1, 1, 0, 0]
]

Notice that all zeroes can be put in blocks of (1, 2), and the sparsity level = 60%

Edit 2 -- forgot a constraint

Another constraint that I forgot to mention, but tried incorporating into my code is that the zero blocks must be non-overlapping.

Example 1: The result below is NOT valid

>>> sparse_matrix = sparsify(matrix, 0.5, (2, 2))
>>> sparse_matrix
[
  [0, 0, 1, 1],
  [0, 0, 0, 1],
  [1, 0, 0, 1],
  [1, 1, 1, 1]
]

Although the blocks starting at index (0, 0) and (1, 1) have valid zero-shapes, the result does not meet the requirements. The reason is that only one of those blocks can be considered valid. if we label the zero blocks as z0 and z1, here is what this matrix is:

[
  [z0, z0,  1, 1],
  [z0, z0, z1, 1],
  [ 1, z1, z1, 1],
  [ 1,  1,  1, 1]
]

element at (1, 1) can be treated as belonging to z0 or z1. That means that there is only one sparse block, which makes the level of sparsity at 25% (not ~44%).

user_12829979 · Answer

The probability of becoming 0 is not all equal.

For example: block_shape (2, 2), matrix(0, 0) becoming 0 has probability of target since the loop only passes through once. matrix(1, 0) has probability more than target since the loop passes it twice. similarly, matrix(1, 1) has probability more than (1, 0) because the loop sees it four times at (0, 0), (1, 0), (0, 1), (1, 1).

This also happens in the middle of the matrix due to prior loop operations.

So the main variable affecting the result is the block_shape.

I've been fiddling around for a bit and here's an alternative way using while loop instead of for loop. Simulating through until you reach target probability within err. You just need to watch out for inf loop due to too small err.

import numpy as np

def sparsify(matrix, target, block_shape=None):
    if block_shape is None or block_shape == 1 or block_shape == (1,) or block_shape == (1, 1):
        # 1x1 is just bernoulli with p=target
        probs = np.random.uniform(size=matrix.shape)
        mask = np.zeros(matrix.shape)
        mask[probs >= target] = 1.0
    else:
        if isinstance(block_shape, int):
            block_shape = (block_shape, block_shape)
        if len(block_shape) == 1:
            block_shape = (block_shape[0], block_shape[0])
        mask = np.ones(matrix.shape)
        rows, cols = matrix.shape

        # vars for probability check
        total = float(rows * cols)
        zero_cnt= total - np.count_nonzero(matrix)
        err = 0.005  # .5%

        # simulate until we reach target probability range
        while not target - err < (zero_cnt/ total) < target + err:

            # pick a random point in the matrix
            row = np.random.randint(rows)
            col = np.random.randint(cols)

            # submask = mask[row:row + block_shape[0], col:col + block_shape[1]]
            submask = matrix[row:row + block_shape[0], col:col + block_shape[1]]

            if submask.shape != block_shape:
                # we don't care about the edges, cannot partially sparsify
                continue
            if (submask == 0).any():
                # If current (row, col) is already in the sparsified area, skip
                continue

            # need more 0s to reach target probability range
            if zero_cnt/ total < target - err:  
                matrix[row:row + block_shape[0], col:col + block_shape[1]] = 0

            # need more 1s to reach target probability range
            else:                               
                matrix[row:row + block_shape[0], col:col + block_shape[1]] = 1

            # update 0 count
            zero_cnt = total - np.count_nonzero(matrix)

    return matrix * mask, mask

note.

Didn't check for any optimization or code refactoring.
Didn't use the mask var. Worked on the matrix directly.

matrix = np.ones((100, 100))
matrix, mask = sparsify(matrix, target=0.5, block_shape=(2, 2))
print((matrix == 0).mean())
# prints somewhere between target - err and target + err
# likely to see a lower value in the range since we're counting up (0s)