Reputation: 858
I want to write a conditional for loop, where if an array is of length 1, use one for loop heading, and if this len(array) > 1 and == len(array2), use a different for loop heading, and if neither of these conditions are true, quit with an error of my choice. The real issue is that I don't want to have this if statement, and then the for loops, when the for loops actually are IDENTICAL, except for the heading, and rather long, so doubling the code seems like a waste.
Is there a nice way to do this where I only have to have the meat of the for loop written once?
Note: xarray and tarray are multi-d numpy arrays ie) xarray = array([[1,2,3],[4,5,6]])
The snippit of code looks as such:
if len(tarray) > 1 and len(xarray) == len(tarray):
for x,ts in zip(xarray,tarray):
#stuff
if len(tarray) == 1:
for x in xarray:
#same stuff as above for loop
else:
print 'Dimension Mismatch -- Quitting:'
quit()
Upvotes: 0
Views: 360
Reputation: 40982
In order to change the for
loop header in run time you can use exec
.
codeToRun = "";
if len(tarray) > 1 and len(xarray) == len(tarray):
codeToRun += """for x,ts in zip(xarray,tarray): """
elif len(tarray) == 1:
codeToRun += """for x in xarray:"""
else:
codeToRun += """print 'Dimension Mismatch -- Quitting:'
quit()"""
exec codeToRun + """# rest of code """
see How do I execute a string containing Python code in Python? and consider to use eval
.
Upvotes: 0
Reputation: 2867
This should work:
if len(tarray) >= 1:
res = zip(tarray, xarray) if len(tarray) == len(xarray) else xarray
else:
# Error message
for each in res:
# Do some stuff
Upvotes: 2
Reputation: 16189
If I understand your problem correctly, before the for loop you should be able to do something like:
try:
xx, tt = np.broadcast_arrays(xarray,tarray)
except ValueError:
# raised if two arrays cannot be broadcast together
for x, t in zip(xx,tt):
# do stuff
So, for example:
>>> xarray = np.arange(6).reshape((2,3))
>>> tarray = np.atleast_2d(np.arange(3))
>>> xarray
array([[0, 1, 2],
[3, 4, 5]])
>>> tarray
array([[0, 1, 2]])
>>> xx, tt = np.broadcast_arrays(xarray,tarray)
>>> xx
array([[0, 1, 2],
[3, 4, 5]])
>>> tt
array([[0, 1, 2],
[0, 1, 2]])
>>> tarray = np.arange(6,12).reshape((2,3))
>>> tarray
array([[ 6, 7, 8],
[ 9, 10, 11]])
>>> xx, tt = np.broadcast_arrays(xarray,tarray)
>>> tt
array([[ 6, 7, 8],
[ 9, 10, 11]])
I would also suggest doing some checking beforehand that both the xarray
and tarray
are 2D.
Upvotes: 2
Reputation: 11070
If the contents of the for loop are the same, then you can call a function from both the for loops, containing the common code.
And the if statements, checking is necessary if you want the conditions to be met.
Upvotes: 4