Optimisation of For looping with zip() for handling huge data computations

I am writing a code in python for my research code which is used to handle huge data. I use 3D arrays along with 1D arrays. I use for loops frequently in my code. I would like to optimise the for loops with better options such as zip() and if possible any other better coding idea. I wrote a sample code here, with all features I use in my coding. For example: Call a function with in a loop, 2. multiplication of each element of a 2D code with each element of a 1D array as shown in the code. Kindly guide me.

#*******************************************************************
import numpy as np
import time
import random
#*******************************************************************
Nx = 5    # resolution x-dir
Ny = 5    # resolution y-dir
NL = 7
weights = np.array([0.2,0.5,0.8,1.5,2.8,5.9,1.5]) # sums to 1
ux = np.array([[1, 2,3,4,5], [3, 4,6,8,9], [1, 4,3,6,2], [4, 5,7,4,3], [1, 2,6,5,11]])
print(ux)
#+++++++++++++++++++++++++++++++++++++++++++
def add(a,b):
    c = a+b
    return(c)
#++++++++++++++++++++++++++++++++++++++++++++
# Normal For Looping
def main():
    start = time.time()
    for i in range(Nx):
        for j in range(Ny):
            sum = 0.0
            for k in range(NL):
                a= add(i,j)
                sum= sum + weights[k]*ux[i,j]+a
            ux[i,j]= sum
    end = time.time()
    print('Time:',end-start)
#+++++++++++++++++++++++++++++++++++++++++++
    print(ux)
#
    return 0
#++++++++++++++++++++++
if __name__== "__main__":
  main()
#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

A big thing that stands out, your inner loop over k contains two invariants: ux[i,j] and a. These could be pushed out of the inner loop, reducing the amount of additions and lookups:

def main():
    start = time.time()
    for i in range(Nx):
        for j in range(Ny):
            tot = 0.0
            val = ux[i, j] + i + j

            for k in range(NL):
                tot += weights[k] * val

            ux[i,j]= tot

    end = time.time()
    print('Time:',end-start)

Now that this has happened, we can see that we really can just do weights * val and eliminate the inner loops entirely:

def main():
    for i in range(Nx):
        for j in range(Ny):
            val = ux[i, j] + i + j

            ux[i,j]= (weights * val).sum()

But since the multiplication of weights and val being added is commutative, weights.sum() is a constant, and can just be:

weight = weights.sum()

def main():
    for i in range(Nx):
        for j in range(Ny):
            val = ux[i, j] + i + j
            ux[i,j] = val * weight

This should cut your computation time by a decent amount already.

The last bit uses np.grid to get the indexes of each position in ux:

grid = np.indices((5, 5))

# grid[0] is row indices and grid[1] is column indices
arr = grid[0] + grid[1] + ux

# and finally, the multiplication of each value against
# the weight constant
result = arr * weight

So your whole code boils down to:

def main():
    grid = np.indices((Nx, Ny))
    weight = weights.sum()

    arr = ux + grid[0] + grid[1]

    result = arr * weight

</details>