pythonwolfram-mathematica

Python equivalent of Mathematica's Sow/Reap


Suppose in Mathematica I define the following function:

f[list_] := Map[Prime[Sow[#]] &, list];

which outputs a list of prime numbers, such that if the input list has n at position i, then the output list will contain the nth prime number at position i. For example,

In[2]:= f[{1, 3, 4}]

Out[2]= {2, 5, 7}

Now, if for some reason (debugging, etc...) I want to check what values are being fed into the Prime function. Because of the Sow command in the function, I can do

In[3] := Reap[f[{1, 3, 4}]]

Out[3] := {{2, 5, 7}, {{1, 3, 4}}}

For more details on Sow/Reap, see the Wolfram Documentation. My question is, is there a natural Python equivalent of Mathematica's Sow and Reap functionality? In particular, is there a way to do this kind of thing without explicitly returning extra things from the python function you want to do it to, writing a second python function that is almost the same but returns something extra, or using global variables?


Solution

  • I came up with two ways to implement a rudimentary version of something like this, each with its own limitations. Here's the first version:

    farm = []
    
    def sower(func):
        def wrapped(*args, **kw):
            farm.append([])
            return func(*args, **kw)
        return wrapped
    
    def sow(val):
        farm[-1].append(val)
        return val
    
    def reap(val):
        return val, farm.pop()
    

    You can use it like this (based on one of the examples from the Mathematica doc page):

    >>> @sower
    ... def someSum():
    ...     return sum(sow(x**2 + 1) if (x**2 + 1) % 2 == 0 else x**2 + 1 for x in xrange(1, 11))
    >>> someSum()
    395
    >>> reap(someSum())
    (395, [2, 10, 26, 50, 82])
    

    This has a number of limitations:

    1. Any function that wants to use sow has to be decorated with the sower decorator. This means you can't use sow inside inside inline expressions like list comprehensions the way the Mathematica examples do. You might be able to hack this by inspecting the call stack, but it could get ugly.
    2. Any values that are sown but not reaped get stored in the "farm" forever, so the farm will get bigger and bigger over time.
    3. It doesn't have the "tag" abilities shown in the docs, although that wouldn't be too hard to add.

    Writing this made me think of a simpler implementation with slightly different tradeoffs:

    farm = []
    
    def sow(val):
        if farm:
            farm[-1].append(val)
        return val
    
    def reap(expr):
        farm.append([])
        val = expr()
        return val, farm.pop()
    

    This one you can use like this, which is somewhat more similar to the Mathematica version:

    >>> reap(lambda: sum(sow(x**2 + 1) if (x**2 + 1) % 2 == 0 else x**2 + 1 for x in xrange(1, 11)))
    (395, [2, 10, 26, 50, 82])
    

    This one doesn't require the decorator, and it cleans up reaped values, but it takes a no-argument function as its argument, which requires you to wrap your sowing expression in a function (here done with lambda). Also, this means that all sown values in any function called by the reaped expression will be inserted into the same list, which could result in weird ordering; I can't tell from the Mathematica docs if that's what Mathematica does or what.