So i was trying this and I find it really unfeasible. I am not that aware about smart ways to do the following. Can somebody help ? Also inputs of lists are quite big.
This task was to create an image from the values generated by me.
center_star contains list of [x,y] pairs which are centers of various point like objects.
1800 value represents that image to be generated is of 1800x1800 pixel.
Sigma variable has value 2 by default.
final=[[0]*1800]*1800
for i in range(len(center_stars)):
xi=center_stars[i][0]
yi=center_stars[i][1]
print(i)
for j in range(1800):
for k in range(1800):
final[j][k]+=gauss_amplitude[i]*(math.e**((-1*((xi-j)**2+(yi-k)**2))/2*sigma*sigma))
Is there a smarter way to save time using some of the numpy operation and execute this piece of code in less time?
Something like that:
import math
import numpy as np
N = 1800
final = np.empty((N, N))
final1 = np.empty((N, N))
j = np.arange(N)
k = np.arange(N)
jj, kk = np.meshgrid(j, k)
sigma = 2.
s = 0.5 / (sigma * sigma)
for i in range(len(center_stars)):
xi = center_stars[i][0]
yi = center_stars[i][1]
final += gauss_amplitude[i] * np.exp(- ((xi - jj.T)**2 + (yi - kk.T)**2) * s)
# Code below is for comparison:
for j in range(N):
for k in range(N):
final1[j][k]+=gauss_amplitude[i] * (math.e** (-((xi-j)**2+(yi-k)**2)/(2*sigma*sigma)))
Besides, I assume you missed brackets around 2*sigma*sigma
you can try to compact your code like this:
Gauss=lambda i,j,k,xi,yi:gauss_amplitude[i]*(math.e**((-((xi-j)**2+(yi-k)**2))/(2*sigma*sigma)))
final=[[Gauss(i,j,k,x[0],x[1]) for j in range(1800) for k in range(1800)] for i,x in enumerate(center_starts)]
If your sigmas are all the same, you can achieve this with out any loop by using
scipy.signal.convolve2d.
import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import convolve2d
from scipy.stats import multivariate_normal
sigma = 3
width = 200 # smaller than yours so you can see the single pixels
n_stars = 50
# draw some random stars
star_row = np.random.randint(0, width, n_stars)
star_col = np.random.randint(0, width, n_stars)
star_amplitude = np.random.normal(50, 10, n_stars)
# assign amplitudes to center pixel of stars
amplitudes = np.zeros((width, width))
amplitudes[star_row, star_col] = star_amplitude
# create 2d gaussian kernel
row = col = np.arange(-4 * sigma, 4 * sigma + 1)
grid = np.stack(np.meshgrid(row, col)).T
kernel = multivariate_normal(
[0, 0],
[[sigma**2, 0], [0, sigma**2]]
).pdf(grid)
kernel /= kernel.sum()
# convolve with 2d gaussian
final = convolve2d(amplitudes, kernel, mode='same')
fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(9, 3))
img = ax1.imshow(amplitudes)
fig.colorbar(img, ax=ax1)
ax1.set_label('Before Convolution')
img = ax2.imshow(kernel)
fig.colorbar(img, ax=ax2)
ax2.set_label('Convolution Kernel')
img = ax3.imshow(final)
fig.colorbar(img, ax=ax3)
ax3.set_label('After Convolution')
fig.tight_layout()
fig.savefig('conv2d.png', dpi=300)
Result:
If the sigmas differ, you get a way with a single loop over the possible sigmas.
Related
I am applying numpy.fft.fft function to a simple 'top-hat' function, i.e. all values are 0 except over a given x-range where they are 1.
The problem arises when I apply a fast fourier transform to the data. It outputs what I am expecting, which is a Gaussian profile, but it looks like instead of being centred on 0 the peaks are at the edges. I am expecting the peak to be central. It's as if the output was split down the centre, and mirrored.
Is there a way of centering this output data other than manually flipping the output array? Why is it doing this?
Here's my full code:
import numpy as np
from numpy.fft import fft, ifft
import matplotlib.pyplot as plt
def find_nearest(array, value):
""" Finds the nearest array value index for a particular value."""
array = np.asarray(array)
idx = (np.abs(array - value)).argmin()
return idx
aperture_length = 81.92e-3 # in metres
start = - aperture_length / 2
end = aperture_length / 2
step = 0.02e-3
x = np.arange(start, end, step)
E = np.zeros(len(x))
slit_length = 1e-3
ind1 = find_nearest(x, -slit_length/2)
ind2 = find_nearest(x, slit_length/2)
for i in range(len(E)):
if i <= ind2 and i >= ind1:
E[i] = 1
# Compute FFT
trans = fft(E, norm='ortho')
fig = plt.figure(dpi=150)
ax1 = plt.subplot(1, 1, 1)
ax1.plot(x, E)
ax1.set_title(r'E vs. x')
plt.figure(dpi=150)
ax2 = plt.subplot(1, 1, 1)
ax2.plot(x, trans)
# ax2.set_xlim(-0.02, 0.02)
ax2.set_title(r'FFT of E')
plt.tight_layout()
Overall, I want to calculate a fourier transform of a given data set and filter out some of the frequencies with the biggest absolute values. So:
1) Given a data array D with accompanying times t, 2) find the k biggest fourier coefficients and 3) remove those coefficients from the data, in order to filter out certain signals from the original data.
Something goes wrong in the end when plotting the filtered data set over the given times. I'm not exactly sure, where the error is. The final 'filtered data' plot doesn't look even slightly smoothened and somehow changes its position compared with the original data. Is my code completely bad?
Part 1):
n = 1000
limit_low = 0
limit_high = 0.48
D = np.random.normal(0, 0.5, n) + np.abs(np.random.normal(0, 2, n) * np.sin(np.linspace(0, 3*np.pi, n))) + np.sin(np.linspace(0, 5*np.pi, n))**2 + np.sin(np.linspace(1, 6*np.pi, n))**2
scaling = (limit_high - limit_low) / (max(D) - min(D))
D = D * scaling
D = D + (limit_low - min(D)) # given data
t = linspace(0,D.size-1,D.size) # times
Part 2):
from numpy import linspace
import numpy as np
from scipy import fft, ifft
D_f = fft.fft(D) # fft of D-dataset
#---extract the k biggest coefficients out of D_f ---
k = 15
I, bigvals = [], []
for i in np.argsort(-D_f):
if D_f[i] not in bigvals:
bigvals.append(D_f[i])
I.append(i)
if len(I) == k:
break
bigcofs = np.zeros(len(D_f))
bigcofs[I] = D_f[I] # array with only zeros in in except for the k maximal coefficients
Part 3):
D_filter = fft.ifft(bigcofs)
D_new = D - D_filter
p1=plt.plot(t,D,'r')
p2=plt.plot(t,D_new,'b');
plt.legend((p1[0], p2[0]), ('original data', 'filtered data'))
I appreciate your help, thanks in advance.
There were two issues I noticed:
you likely want the components with largest absolute value, so np.argsort(-np.abs(D_f)) instead of np.argsort(-D_f).
More subtly: bigcofs = np.zeros(len(D_f)) is of type float64 and was discarding the imaginary part at the line bigcofs[I] = D_f[I]. You can fix this with bigcofs = np.zeros(len(D_f), dtype=complex)
I've improved your code slightly below to get desired results:
import numpy as np
from scipy import fft, ifft
import matplotlib.pyplot as plt
n = 1000
limit_low = 0
limit_high = 0.48
N_THRESH = 10
D = 0.5*np.random.normal(0, 0.5, n) + 0.5*np.abs(np.random.normal(0, 2, n) * np.sin(np.linspace(0, 3*np.pi, n))) + np.sin(np.linspace(0, 5*np.pi, n))**2 + np.sin(np.linspace(1, 6*np.pi, n))**2
scaling = (limit_high - limit_low) / (max(D) - min(D))
D = D * scaling
D = D + (limit_low - min(D)) # given data
t = np.linspace(0,D.size-1,D.size) # times
# transformed data
D_fft = fft.fft(D)
# Create boolean mask for N largest indices
idx_sorted = np.argsort(-np.abs(D_fft))
idx = idx_sorted[0:N_THRESH]
mask = np.zeros(D_fft.shape).astype(bool)
mask[idx] = True
# Split fft above, below N_THRESH points:
D_below = D_fft.copy()
D_below[mask] = 0
D_above = D_fft.copy()
D_above[~mask] = 0
#inverse separated functions
D_above = fft.ifft(D_above)
D_below = fft.ifft(D_below)
# plot
plt.ion()
f, (ax1, ax2, ax3) = plt.subplots(3,1)
l1, = ax1.plot(t, D, c="r", label="original")
l2, = ax2.plot(t, D_above, c="g", label="top {} coeff. signal".format(N_THRESH))
l3, = ax3.plot(t, D_below, c="b", label="remaining signal")
f.legend(handles=[l1,l2,l3])
plt.show()
Good day to you fellow programmer !
Today I would like to do something that I believe is tricky. I have a very large 2D array called tac that basically contains time curve values and a file containing a tuple of coordinates called coor which contains information on where to place these curves in a 3D array. What this set of variables represents is actually a 4D array: the first 3 dimensions represent space dimensions and the fourth is time. The whole thing is stored as is to avoid storing an immense amount of zeros.
I would like to apply, for each time (in other words, each values in the 4th dimension), a gaussian kernel to this set of data. I was able to generate this kernel and to perform the convolution quite easily for a fixed standard deviation for the whole array using scipy.ndimage.convolve. The kernel was created using scipy.signal.gaussian. Here is a brief example of the principle where tac_4d contains the 4D array (stores a lot of data I know... but one problem at the time):
def gaussian_kernel_3d(radius, sigma):
num = 2 * radius + 1
kernel_1d = signal.gaussian(num, std=sigma).reshape(num, 1)
kernel_2d = np.outer(kernel_1d, kernel_1d)
kernel_3d = np.outer(kernel_1d, kernel_2d).reshape(num, num, num)
kernel_3d = np.expand_dims(kernel_3d, -1)
return kernel_3d
g = gaussian_kernel_3d(1, .5)
cag = nd.convolve(tac_4d, g, mode='constant', cval=0.0)
The trick is now to convolve the array with a kernel which standard deviation is different for each SPACE coordinate. In other words, I would have a 3D array std containing standard deviations for each coordinate of the array.
It seems https://github.com/sheliak/varconvolve is the code needed to take care of this problem. However I don't really understand how to use it and quite frankly, I would prefer to come up with a genuine solution. Do you guys see a way to solve this problem?
Thanks in advance !
EDIT
Here is what I hope can be considered MCVE
import numpy as np
from scipy import signal
from scipy import ndimage as nd
def gaussian_kernel_2d(radius, sigma):
num = 2 * radius + 1
kernel_1d = signal.gaussian(num, std=sigma).reshape(num, 1)
kernel_2d = np.outer(kernel_1d, kernel_1d)
return kernel_2d
def gaussian_kernel_3d(radius, sigma):
num = 2 * radius + 1
kernel_1d = signal.gaussian(num, std=sigma).reshape(num, 1)
kernel_2d = np.outer(kernel_1d, kernel_1d)
kernel_3d = np.outer(kernel_1d, kernel_2d).reshape(num, num, num)
kernel_3d = np.expand_dims(kernel_3d, -1)
return kernel_3d
np.random.seed(0)
number_of_tac = 150
time_samples = 915
z, y, x = 100, 150, 100
voxel_number = x * y * z
# TACs in the right order
tac = np.random.uniform(0, 4, time_samples * number_of_tac).reshape(number_of_tac, time_samples)
arr = np.array([0] * (voxel_number - number_of_tac) + [1] * number_of_tac)
np.random.shuffle(arr)
arr = arr.reshape(z, y, x)
coor = np.where(arr != 0) # non-empty voxel
# Algorithm to replace TAC in 3D space
nnz = np.zeros(arr.shape)
nnz[coor] = 1
tac_4d = np.zeros((x, y, z, time_samples))
tac_4d[np.where(nnz == 1)] = tac
# 3D convolution for all time
# TODO: find a way to make standard deviation change for each voxel
g = gaussian_kernel_3d(1, 1) # 3D kernel of std = 1
v = np.random.uniform(0, 1, x * y * z).reshape(z, y, x) # 3D array of std
cag = nd.convolve(tac_4d, g, mode='constant', cval=0.0) # convolution
Essentially, you have a 4D dataset, shape (nx, ny, nz, nt) that is sparse in (nx, ny, nz) and dense in the nt axis. If (i, j, k) are coordinates of nonzero points in the sparse dimensions, you want to convolve with a Gaussian 3D kernel that has a sigma that depends on (i, j, k).
For example, if there are nonzero points at [1, 2, 5] and [1, 4, 5] with corresponding sigmas 0.1 and 1.0, then the output at coordinates [1, 3, 5] is affected mostly by the [1, 4, 5] point because that one has the largest point spread.
Your question is ambiguous; it could also mean that point [1, 3, 5] has a its own associated sigma, for example 0.5, and pulls data from the two adjacent points with equal weight. I will assume the first definition (sigma values associated with input points, not with output points).
Because the operation is not a true convolution, there is no fast FFT-based method to do the entire operation in one operation. Instead, you have to loop over the sigma values. Fortunately, your example has only 150 nonzero points, so the loop is not too expensive.
Here is an implementation. I keep the data in sparse representation as long as possible.
import scipy.signal
import numpy as np
def kernel3d(mm, sigma):
"""Return (mm, mm, mm) shaped, normalized kernel."""
g1 = scipy.signal.gaussian(mm, std=sigma)
g3 = g1.reshape(mm, 1, 1) * g1.reshape(1, mm, 1) * g1.reshape(1, 1, mm)
return g3 * (1/g3.sum())
np.random.seed(1)
s = 2 # scaling factor (original problem: s=10)
nx, ny, nz, nt, nnz = 10*s, 11*s, 12*s, 91*s, 15*s
# select nnz random voxels to fill with time series data
randint = np.random.randint
tseries = {} # key: (i, j, k) tuple; value: time series data, shape (nt,)
for _ in range(nnz):
while True:
ijk = (randint(nx), randint(ny), randint(nz))
if ijk not in tseries:
tseries[ijk] = np.random.uniform(0, 1, size=nt)
break
ijks = np.array(list(tseries.keys())) # shape (nnz, 3)
# sigmas: key: (i, j, k) tuple; value: standard deviation
sigmas = { k: np.random.uniform(0, 2) for k in tseries.keys() }
# output will be stored as dense array, padded to avoid edge issues
# with convolution.
m = 5 # padding size
cag_4dp = np.zeros((nx+2*m, ny+2*m, nz+2*m, nt))
mm = 2*m + 1 # kernel width
for (i, j, k), tdata in tseries.items():
kernel = kernel3d(mm, sigmas[(i, j, k)]).reshape(mm, mm, mm, 1)
# convolution of one voxel by kernel is trivial.
# slice4d_c has shape (mm, mm, mm, nt).
slice4d_c = kernel * tdata
cag_4dp[i:i+mm, j:j+mm, k:k+mm, :] += slice4d_c
cag_4d = cag_4dp[m:-m, m:-m, m:-m, :]
#%%
import matplotlib.pyplot as plt
fig, axs = plt.subplots(2, 2, tight_layout=True)
plt.close('all')
# find a few planes
#ks = np.where(np.any(cag_4d != 0, axis=(0, 1,3)))[0]
ks = ijks[:4, 2]
for ax, k in zip(axs.ravel(), ks):
ax.imshow(cag_4d[:, :, k, nt//2].T)
ax.set_title(f'Voxel [:, :, {k}] at time {nt//2}')
fig.show()
for ijk, sigma in sigmas.items():
print(f'{ijk}: sigma={sigma:.2f}')
I have a not-quite linear gradient at some angle to the horizontal as an image. Here's some toy data:
g = np.ones((5,20))
for x in range(g.shape[0]):
for y in range(g.shape[1]):
g[x,y] += (x+y)*0.1+(y*0.01)
I want to essentially correct the skew in the gradient so that it is horizontal, i.e. the gradient increases to the right and all vertical slices are constant.
This will of course produce a parallelogram with a larger x-axis than the input image. Returning a masked Numpy array would be ideal. Here's a (terrible) cartoon to quickly illustrate.
Any idea how to achieve this? Thanks!
You can intepolate to determine the skewness and interpolate again to correct it.
import numpy as np
from scipy.ndimage.interpolation import map_coordinates
m, n = g.shape
j_shift = np.interp(g[:,0], g[0,:], np.arange(n))
pad = int(np.max(j_shift))
i, j = np.indices((m, n + pad))
z = map_coordinates(g, [i, j - j_shift[:,None]], cval=np.nan)
This works on the example image, but you have to do some additional checks to make it function on other gradients. It does not work on gradients that are nonlinear in the x-direction though. Demo:
Full script:
import numpy as np
from scipy.ndimage.interpolation import map_coordinates
def fix(g):
x = 1 if g[0,0] < g[0,-1] else -1
y = 1 if g[0,0] < g[-1,0] else -1
g = g[::y,::x]
m, n = g.shape
j_shift = np.interp(g[:,0], g[0,:], np.arange(n))
pad = int(np.max(j_shift))
i, j = np.indices((m, n + pad))
z = map_coordinates(g, [i, j - j_shift[:,None]], cval=np.nan)
return z[::y,::x]
import matplotlib.pyplot as plt
i, j = np.indices((50,100))
g = 0.01*i**2 + j
plt.figure(figsize=(6,5))
plt.subplot(211)
plt.imshow(g[::-1], interpolation='none')
plt.title('original')
plt.subplot(212)
plt.imshow(fix(g[::-1]), interpolation='none')
plt.title('fixed')
plt.tight_layout()
I'm new to python, and I have this code for calculating the potential inside a 1x1 box using fourier series, but a part of it is going way too slow (marked in the code below).
If someone could help me with this, I suspect I could've done something with the numpy library, but I'm not that familiar with it.
import matplotlib.pyplot as plt
import pylab
import sys
from matplotlib import rc
rc('text', usetex=False)
rc('font', family = 'serif')
#One of the boundary conditions for the potential.
def func1(x,n):
V_c = 1
V_0 = V_c * np.sin(n*np.pi*x)
return V_0*np.sin(n*np.pi*x)
#To calculate the potential inside a box:
def v(x,y):
n = 1;
sum = 0;
nmax = 20;
while n < nmax:
[C_n, err] = quad(func1, 0, 1, args=(n), );
sum = sum + 2*(C_n/np.sinh(np.pi*n)*np.sin(n*np.pi*x)*np.sinh(n*np.pi*y));
n = n + 1;
return sum;
def main(argv):
x_axis = np.linspace(0,1,100)
y_axis = np.linspace(0,1,100)
V_0 = np.zeros(100)
V_1 = np.zeros(100)
n = 4;
#Plotter for V0 = v_c * sin () x
for i in range(100):
V_0[i] = V_0_1(i/100, n)
plt.plot(x_axis, V_0)
plt.xlabel('x/L')
plt.ylabel('V_0')
plt.title('V_0(x) = sin(m*pi*x/L), n = 4')
plt.show()
#Plot for V_0 = V_c(1-(x-1/2)^4)
for i in range(100):
V_1[i] = V_0_2(i/100)
plt.figure()
plt.plot(x_axis, V_1)
plt.xlabel('x/L')
plt.ylabel('V_0')
plt.title('V_0(x) = 1- (x/L - 1/2)^4)')
#plt.legend()
plt.show()
#Plot V(x/L,y/L) on the boundary:
V_0_Y = np.zeros(100)
V_1_Y = np.zeros(100)
V_X_0 = np.zeros(100)
V_X_1 = np.zeros(100)
for i in range(100):
V_0_Y[i] = v(0, i/100)
V_1_Y[i] = v(1, i/100)
V_X_0[i] = v(i/100, 0)
V_X_1[i] = v(i/100, 1)
# V(x/L = 0, y/L):
plt.figure()
plt.plot(x_axis, V_0_Y)
plt.title('V(x/L = 0, y/L)')
plt.show()
# V(x/L = 1, y/L):
plt.figure()
plt.plot(x_axis, V_1_Y)
plt.title('V(x/L = 1, y/L)')
plt.show()
# V(x/L, y/L = 0):
plt.figure()
plt.plot(x_axis, V_X_0)
plt.title('V(x/L, y/L = 0)')
plt.show()
# V(x/L, y/L = 1):
plt.figure()
plt.plot(x_axis, V_X_1)
plt.title('V(x/L, y/L = 1)')
plt.show()
#Plot V(x,y)
#######
# This is where the code is way too slow, it takes like 10 minutes when n in v(x,y) is 20.
#######
V = np.zeros(10000).reshape((100,100))
for i in range(100):
for j in range(100):
V[i,j] = v(j/100, i/100)
plt.figure()
plt.contour(x_axis, y_axis, V, 50)
plt.savefig('V_1')
plt.show()
if __name__ == "__main__":
main(sys.argv[1:])
You can find how to use FFT/DFT in this document :
Discretized continuous Fourier transform with numpy
Also, regarding your V matrix, there are many ways to improve the execution speed. One is to make sure you use Python 3, or xrange() instead of range() if you a are still in Python 2.. I usually put these lines in my Python code, to allow it to run evenly wether I use Python 3. or 2.*
# Don't want to generate huge lists in memory... use standard range for Python 3.*
range = xrange if isinstance(range(2),
list) else range
Then, instead of re-computing j/100 and i/100, you can precompute these values and put them in an array; knowing that a division is much more costly than a multiplication ! Something like :
ratios = np.arange(100) / 100
V = np.zeros(10000).reshape((100,100))
j = 0
while j < 100:
i = 0
while i < 100:
V[i,j] = v(values[j], values[i])
i += 1
j += 1
Well, anyway, this is rather cosmetic and will not save your life; and you still need to call the function v()...
Then, you can use weave :
http://docs.scipy.org/doc/scipy-0.14.0/reference/tutorial/weave.html
Or write all your pure computation/loop code in C, compile it and generate a module which you can call from Python.
You should look into numpy's broadcasting tricks and vectorization (several references, one of the first good links that pops up is from Matlab but it is just as applicable to numpy - can anyone recommend me a good numpy link in the comments that I might point other users to in the future?).
What I saw in your code (once you remove all the unnecessary bits like plots and unused functions), is that you are essentially doing this:
from __future__ import division
from scipy.integrate import quad
import numpy as np
import matplotlib.pyplot as plt
def func1(x,n):
return 1*np.sin(n*np.pi*x)**2
def v(x,y):
n = 1;
sum = 0;
nmax = 20;
while n < nmax:
[C_n, err] = quad(func1, 0, 1, args=(n), );
sum = sum + 2*(C_n/np.sinh(np.pi*n)*np.sin(n*np.pi*x)*np.sinh(n*np.pi*y));
n = n + 1;
return sum;
def main():
x_axis = np.linspace(0,1,100)
y_axis = np.linspace(0,1,100)
#######
# This is where the code is way too slow, it takes like 10 minutes when n in v(x,y) is 20.
#######
V = np.zeros(10000).reshape((100,100))
for i in range(100):
for j in range(100):
V[i,j] = v(j/100, i/100)
plt.figure()
plt.contour(x_axis, y_axis, V, 50)
plt.show()
if __name__ == "__main__":
main()
If you look carefully (you could use a profiler too), you'll see that you're integrating your function func1 (which I'll rename into the integrand) about 20 times for each element in the 100x100 array V. However, the integrand doesn't change! So you can already bring it out of your loop. If you do that, and use broadcasting tricks, you could end up with something like this:
import numpy as np
from scipy.integrate import quad
import matplotlib.pyplot as plt
def integrand(x,n):
return 1*np.sin(n*np.pi*x)**2
sine_order = np.arange(1,20).reshape(-1,1,1) # Make an array along the third dimension
integration_results = np.empty_like(sine_order, dtype=np.float)
for enu, order in enumerate(sine_order):
integration_results[enu] = quad(integrand, 0, 1, args=(order,))[0]
y,x = np.ogrid[0:1:.01, 0:1:.01]
term = integration_results / np.sinh(np.pi * sine_order) * np.sin(sine_order * np.pi * x) * np.sinh(sine_order * np.pi * y)
# This is the key: you have a 3D matrix here and with this summation,
# you're basically squashing the entire 3D structure into a flat, 2D
# representation. This 'squashing' is done by means of a sum.
V = 2*np.sum(term, axis=0)
x_axis = np.linspace(0,1,100)
y_axis = np.linspace(0,1,100)
plt.figure()
plt.contour(x_axis, y_axis, V, 50)
plt.show()
which runs in less than a second on my system.
Broadcasting becomes much more understandable if you take pen&paper and draw out the vectors that you are "broadcasting" as if you were constructing a building, from basic Tetris-blocks.
These two versions are functionally the same, but one is completely vectorized, while the other uses python for-loops. As a new user to python and numpy, I definitely recommend reading through the broadcasting basics. Good luck!