KevinShaffer
KevinShaffer

Reputation: 858

For loop heading depending on if statement (Python)

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

Answers (4)

0x90
0x90

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

vivek
vivek

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

user545424
user545424

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

Aswin Murugesh
Aswin Murugesh

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

Related Questions