Is there a good way of iterating through several lists(where sometimes they are single valued) in python

Question

So basically my question relates to 'zip' (or izip), and this question which was asked before....

Is there a better way to iterate over two lists, getting one element from each list for each iteration?

If i have two variables - where they either are a 1d array of values length n, or are a single value, how do i loop through them so that I get n values returned.

'zip' kindof does what I want - except that when I pass in a single value, and an array it complains.

I have an example of what I'm aiming for below - basically i have a c function that does a more efficient calculation than python. I want it to act like some of the numpy functions - that deal ok with mixtures of arrays and scalars, so i wrote a python wrapper for it. However - like I say 'zip' fails. I guess in principle I can do some testing of the input s and write a different statement for each variation of scalars and arrays - but it seems like python should have something more clever.... ;) Any advice?

"""
    Example of zip problems.
"""

import numpy as np
import time

def cfun(a, b) :
    """
        Pretending to be c function which doesn't deal with arrays
    """
    if not np.isscalar(a)   or  not np.isscalar(b)  :

        raise Exception('c is freaking out')
    else :

        return a+b

def pyfun(a, b) :
    """
        Python Wrappper - to deal with arrays input
    """

    if not np.isscalar(a)   or  not np.isscalar(b) :
        return np.array([cfun(a_i,b_i) for a_i, b_i in zip(a,b)])

    else :

        return cfun(a, b)

    return cfun(a,b)


a = np.array([1,2])
b= np.array([1,2])
print pyfun(a, b)

a = [1,2]
b = 1
print pyfun(a, b)

edit :

Many thanks everyone for the suggestions everyone. Think i have to go for np.braodcast for the solution - since it seems the simplest from my perspective.....

I changed your `'''` to `"""` so SO's renderer would like it. As a side note, `"""` is actually preferred to `'''` I think — mgilson, Mar 01 '13 at 22:27
[PEP 257](http://www.python.org/dev/peps/pep-0257/) describes docstrings. It says... `For consistency, always use """triple double quotes""" around docstrings. Use r"""raw triple double quotes""" if you use any backslashes in your docstrings. For Unicode docstrings, use u"""Unicode triple-quoted strings"""` — Mark Hildreth, Mar 01 '13 at 23:00
@MarkHildreth -- PEP 257. I knew it was in a PEP somewhere, but when I didn't find it in PEP8, I began to wonder if I was making it up. — mgilson, Mar 02 '13 at 01:20

score 4 · Accepted Answer · answered Mar 01 '13 at 23:49

4

Since you use numpy, you don't need zip() to iterate several arrays and scalars. You can use numpy.broadcast():

In [5]:

list(np.broadcast([1,2,3], 10))

Out[5]:

[(1, 10), (2, 10), (3, 10)]

In [6]:

list(np.broadcast([1,2,3], [10, 20, 30]))

Out[6]:

[(1, 10), (2, 20), (3, 30)]

In [8]:

list(np.broadcast([1,2,3], 100, [10, 20, 30]))

Out[8]:

[(1, 100, 10), (2, 100, 20), (3, 100, 30)]

answered Mar 01 '13 at 23:49

HYRY

94,853
25
187
187

I ran into this after a generator I previously used no longer worked after a `numpy` upgrade. This is really useful, but it seems like the way the question is stated and titled obscures finding this answer. – ryanjdillon Jun 05 '15 at 13:46

score 1 · Answer 2 · answered Mar 01 '13 at 23:03

If you want to force broadcasting, you can use numpy.lib.stride_tricks.broadcast_arrays. Reusing your cfun:

def pyfun(a, b) :
    if not (np.isscalar(a) and np.isscalar(b)) :
        a_bcast, b_bcast = np.lib.stride_tricks.broadcast_arrays(a, b)
        return np.array([cfun(j, k) for j, k in zip(a_bcast, b_bcast)])
    return cfun(a, b)

And now:

>>> pyfun(5, 6)
11
>>> pyfun(5, [6, 7, 8])
array([11, 12, 13])
>>> pyfun([3, 4, 5], [6, 7, 8])
array([ 9, 11, 13])

For your particular application there is probably no advantage over Rob's pure python thing, since your function is still running in a python loop.

Robᵩ · Answer 3 · 2013-03-01T22:57:33.250

A decorator that optinally converts each of the arguments to a sequence might help. Here is the ordinary python (not numpy) version:

# TESTED
def listify(f):
  def dolistify(*args):
    from collections import Iterable
    return f(*(a if isinstance(a, Iterable) else (a,) for a in args))
  return dolistify

@listify
def foo(a,b):
  print a, b

foo( (1,2), (3,4) )
foo( 1, [3,4] )
foo( 1, 2 )

So, in your example we need to use not np.isscalar as the predicate and np.array as the modifier. Because of the decorator, pyfun always receives an array.

#UNTESTED
def listify(f):
  def dolistify(*args):
    from collections import Iterable
    return f(*(np.array([a]) if np.isscalar(a) else a for a in args))
  return dolistify

@listify
def pyfun(a, b) :
    """
        Python Wrappper - to deal with arrays input
    """

    return np.array([cfun(a_i,b_i) for a_i, b_i in zip(a,b)])

Or maybe you could apply the same idea to zip:

#UNTESTED
def MyZip(*args):
  return zip(np.array([a]) if np.isscalar(a) else a for a in args)

def pyfun(a, b) :
    """
        Python Wrappper - to deal with arrays input
    """

    return np.array([cfun(a_i,b_i) for a_i, b_i in MyZip(a,b)])

Is there a good way of iterating through several lists(where sometimes they are single valued) in python

3 Answers3