For loop heading depending on if statement (Python)

Question

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()

It would really help if you had gave more info about what "meat of the for loop" does. I suspect, because you're using numpy, some sort of [broadcasting](http://www.scipy.org/EricsBroadcastingDoc) might the the right tool for this situation. — Bi Rico, Apr 05 '13 at 19:19
the meat of the for loop does a trigonometric interpolation on the arrays in xarray (x values) corresponding to the t values in tarray. And essentially I want to be able to say "hey, xarray has 5 chunks of x data, and tarray only has one, use that same tarray value for all 5 xarrays" or "hey, xarray has n chunks, and so does tarray, so lets use xarray[i] with tarray[i] from i = 0,1,...,n" — KevinShaffer, Apr 05 '13 at 19:21

score 4 · Answer 1 · answered Apr 05 '13 at 19:22

4

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.

answered Apr 05 '13 at 19:22

Aswin Murugesh

10,831
10
40
69

Good answer, however I want to keep this function free if possible, as it's just a file that i will be putting into a package to call as needed. – KevinShaffer Apr 05 '13 at 19:28
so the for loop should be a file output statement only? – Aswin Murugesh Apr 05 '13 at 19:35

user545424 · Answer 2 · 2013-04-05T19:45:28.583

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.

score 2 · Accepted Answer · answered Apr 05 '13 at 19:49

2

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

answered Apr 05 '13 at 19:49

vivek

2,807
5
31
44

score 0 · Answer 4 · edited May 23 '17 at 10:32

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.

For loop heading depending on if statement (Python)

4 Answers4