pythonnumpywhile-loopvectorizationdo-while

Error while vectorizing a function containing while loop


I have the follwing code:

import numpy as np
from numpy import random

@np.vectorize
def simulation(_):
    numRolls = 0
    diff = 0
    prevRoll = -999999
    while True:
        roll = random.randint(1, 7)
        print(prevRoll, roll)
        numRolls += 1
        diff = np.abs(roll - prevRoll)
        if diff == 1:
            break
        else:
            prevRoll = roll
    
    return numRolls

If I remove the decorator, the code works perfectly fine. But using the np.vectorize, it seems the while loop continues to run even if it hits break, and then when it hits break again, the loop stops, which seems really odd. Can anyone tell me what the problem is and what the solution would be? I need to have the function vectorized to find the mean.


Solution

  • You have the answer in the docs

    If otypes is not specified, then a call to the function with the first argument will be used to determine the number of outputs. The results of this call will be cached if cache is True to prevent calling the function twice. However, to implement the cache, the original function must be wrapped which will slow down subsequent calls, so only do this if your function is expensive.

    You need to set otypes with the return types

    @np.vectorize(otypes=[int])
    def simulation(_):
        ...