Generating a matrix with each row having normalized weights

Question

I just asked this question over Stack Over Flow on how to improve my code and reposting it here as someone on Stack Overflow recommended this platform.

I have written two python functions and they are correct, but I believe this is not the efficient way to write them, and I am wondering if some can give their input on how to improve them.

The first function is the following which generates an odd number of weights that sum up to one.

import numpy as np
# the argument of the following function has to an non-negative and odd 
# integer 
def weights(n_w):

    M = int((n_w+1)/2)        
    j = np.arange(-M+1, M,1)    
    w = 2.**(-j**2) 
    w = w/w.sum()
    return w   

If I now type

weights(5)

I get five weights, which add up to 1.

Out[442]: array([0.02941176, 0.23529412, 0.47058824, 0.23529412, 0.02941176])

The first and last entry of the output array is the same and similarly second and the fourth are also the same.

So, I can generate any odd number of normalized weights using the function above.

I have used this function to write the following function which generates an n by n square matrix with each row having an odd number of weights and zeros.

def weight_matrix(n_w,n):

    w_matrix = np.zeros((n,n))       
    M = int((n_w+1)/2)
    W_main = weights(n_w)     

    for m in range(1,M-1):
        w = weights(2*m+1)
        w_matrix[m,:int(len(w))] = w

    w_matrix += w_matrix[::-1,::-1]
    w_matrix[0,0], w_matrix[-1,-1] = [1,1]

    nn = 0

    for m in range(M-1,n-M+1):
        w = W_main
        w_matrix[m,nn:nn+n_w] = w
        nn  += 1

    return w_matrix 

An example of this function is

weight_matrix(5,7)

where the second argument specifies the dimension of the output matrix and the first argument decides the maximum number weights to be generated.

Following is the output:

array([[1.        , 0.        , 0.        , 0.        , 0.        ,
    0.        , 0.        ],
   [0.25      , 0.5       , 0.25      , 0.        , 0.        ,
    0.        , 0.        ],
   [0.02941176, 0.23529412, 0.47058824, 0.23529412, 0.02941176,
    0.        , 0.        ],
   [0.        , 0.02941176, 0.23529412, 0.47058824, 0.23529412,
    0.02941176, 0.        ],
   [0.        , 0.        , 0.02941176, 0.23529412, 0.47058824,
    0.23529412, 0.02941176],
   [0.        , 0.        , 0.        , 0.        , 0.25      ,
    0.5       , 0.25      ],
   [0.        , 0.        , 0.        , 0.        , 0.        ,
    0.        , 1.        ]])

Each row adds up to 1

The first element of the first row is 1 which is the only non-zero element of it.

Similarly, the last element of the last row is 1 which is the only non-zero element of it.

The second row and the second to the last have the same three non-zero elements but at different positions.

From row 3 to the last row have the same five non-zero elements but at different positions.

These five non-zero elements shift along the diagonal of the matrix.

The second function gives me control over what size matrix to generate and how long I shall increase the number of weights.

Any suggestion to improve these functions will be appreciated.

Answer 1

• Type-hint your function arguments
• Your weight_matrix, when vectorised, won't be able to rely on weights
• Rather than int((n_w+1)/2), use (n_w + 1)//2
• Add some unit tests based on your known criteria and output data
• It is possible to vectorise with no for-loops, which is usually preferred but I have not done detailed performance analysis. I wouldn't be able to perform such analysis unless I knew how often this function was called and with what sizes.
• Your function fails for me for all parameters that I tried that were not (5, 7), so your index manipulation needs to be improved. If you need to impose that arguments should be odd, you should do so at the top of the function and raise if those conditions are not met.
• Prefer n_weights instead of the abbreviated n_w

The following example vectorisation ignores weights, and uses modified upper triangular indices to zero out the appropriate elements.

import numpy as np


def weights(n_weights: int) -> np.ndarray:
    M = (n_weights + 1) // 2
    j = np.arange(1 - M, M)
    w = 2.**(-j**2)
    return w / w.sum()


def weight_matrix(n_weights: int, n: int) -> np.ndarray:
    j = np.arange(n)[:, np.newaxis]
    w = 2. ** -(j.T - j)**2

    # This is a diagonally dominant matrix, but not a diagonal matrix. Moreover,
    # the zero regions do not align with diagonals. A first cut of the zero-
    # region indices comes from the upper triangle with k=1.
    i, j = np.triu_indices(n=n, k=1)
    zeros = j > np.min(
        np.vstack((2*i, i + n_weights//2)),
        axis=0,
    )
    i = i[zeros]
    j = j[zeros]
    w[i, j] = 0
    w[-1-i, -1-j] = 0

    return w / np.sum(w, axis=1, keepdims=True)


def test_symmetry(w_matrix: np.ndarray) -> None:

    # Each row adds up to 1
    row_sums = np.sum(w_matrix, axis=1)
    assert np.all(np.isclose(
        1, row_sums, rtol=0, atol=1e-15,
    ))

    # The number of non-zeros per row is 2n + 1 to the middle
    nonzeros = np.sum(
        np.logical_not(np.isclose(
            w_matrix, 0, rtol=0, atol=1e-10,
        )), axis=1,
    )
    half_1 = nonzeros[:len(nonzeros)//2]
    half_2 = nonzeros[len(nonzeros)//2+1:]
    assert np.all(half_1 == half_2[::-1])


def test() -> None:
    # Test with known output
    assert np.all(np.isclose(
        weights(5),
        np.array((0.02941176, 0.23529412, 0.47058824, 0.23529412, 0.02941176)),
        rtol=0, atol=1e-8,
    ))

    # Test for symmetry
    w = weights(1000)
    assert w.shape == (999,)
    assert np.all(np.isclose(w[:499], w[:499:-1], rtol=0, atol=1e-16))

    actual = weight_matrix(5, 7)
    expected = np.array([
        [1.00000000, 0.00000000, 0.00000000, 0.00000000, 0.00000000, 0.00000000, 0.00000000],
        [0.25000000, 0.50000000, 0.25000000, 0.00000000, 0.00000000, 0.00000000, 0.00000000],
        [0.02941176, 0.23529412, 0.47058824, 0.23529412, 0.02941176, 0.00000000, 0.00000000],
        [0.00000000, 0.02941176, 0.23529412, 0.47058824, 0.23529412, 0.02941176, 0.00000000],
        [0.00000000, 0.00000000, 0.02941176, 0.23529412, 0.47058824, 0.23529412, 0.02941176],
        [0.00000000, 0.00000000, 0.00000000, 0.00000000, 0.25000000, 0.50000000, 0.25000000],
        [0.00000000, 0.00000000, 0.00000000, 0.00000000, 0.00000000, 0.00000000, 1.00000000],
    ])
    assert np.all(np.isclose(actual, expected, rtol=0, atol=1e-8))

    test_symmetry(actual)
    test_symmetry(weight_matrix(169, 189))


if __name__ == '__main__':
    test()