[FIXED] calculate matrix without using nested if and for loops?

Issue

I'm trying to make a matrix A (200, 200) where A[i][j] = A[j][i] follows Bernoulli distribution of probability p if the element is in a same cluster of index(i.e. 0-49, 50-99, 100-149, 150-199) and probability q if not.

So far I tried to implement it straightforwardly like below,

import numpy as np
from scipy.stats import bernoulli
def link_weight(p=0.05, q=0.04):
    A = np.zeros((200, 200))
    for i in range(len(A)):
        for j in range(len(A)):
            if i == j:
                A[i][j] = 0
            elif i < j:
                if i < 50:
                    if j < 50:
                        A[i][j] = bernoulli.rvs(p)
                        A[j][i] = bernoulli.rvs(p)
                    else:
                        A[i][j] = bernoulli.rvs(q)
                        A[j][i] = bernoulli.rvs(q)
                if 50 <= i <= 99:
                    if 50 <= j <= 99:
                        A[i][j] = bernoulli.rvs(p)
                        A[j][i] = bernoulli.rvs(p)
                    else:
                        A[i][j] = bernoulli.rvs(q)
                        A[j][i] = bernoulli.rvs(q)
                if 100 <= i <= 149:
                    if 100 <= j <= 149:
                        A[i][j] = bernoulli.rvs(p)
                        A[j][i] = bernoulli.rvs(p)
                    else:
                        A[i][j] = bernoulli.rvs(q)
                        A[j][i] = bernoulli.rvs(q)
                if 150 <= i <= 199:
                    if 150 <= j <= 199:
                        A[i][j] = bernoulli.rvs(p)
                        A[j][i] = bernoulli.rvs(p)
                    else:
                        A[i][j] = bernoulli.rvs(q)
                        A[j][i] = bernoulli.rvs(q)
    return A

But I wanna make it more readable and faster. How can I prove this code by avoiding nested if/for loops?

Solution

Numpy indexing is really powerful so you could use that:

import numpy as np
from scipy.stats import bernoulli

def link_weight(p=0.05, q=0.04):
    # Create bernouli array
    A = bernoulli.rvs(q, size=(200, 200))
    # Randomize values p=p
    A[0:50, 0:50] = bernoulli.rvs(p, size=(50,50))
    A[50:100, 50:100] = bernoulli.rvs(p, size=(50,50))
    A[100:150, 100:150] = bernoulli.rvs(p, size=(50,50))
    A[150:200, 150:200] = bernoulli.rvs(p, size=(50,50))
    # Make symmetric
    A = A + A.T - np.diag(A.diagonal())
    return A

print(link_weight())

• Symmertic matrix: https://stackoverflow.com/a/10806947/14536215
• Indexing: https://numpy.org/doc/stable/reference/arrays.indexing.html

Edit: Change the "symmetrization" to use this: https://stackoverflow.com/a/2573982/14536215

Answered By - Tzane