Is there a way of specifying which argument QuTiP's parallel_map should iterate over?

Question

QuTiP's function parallel_map provides the possibility to compute the value of a given function for several values of its argument in parallel. All the examples show cases where the first positional argument is varied, like the following:

def testFunc1(a, b):
    return a, b

from qutip import parallel_map
parallel_map(testFunc1, (1, 2, 3), task_args=(4,))

This returns [(1, 4), (2, 4), (3, 4)]. Now I'm wondering if it's also possible to have a fixed value for a and a tuple for b. According to the documentation task_args can also be a dictionary, so I tried

parallel_map(testFunc1, (1, 2, 3), task_args={'a': 4})
parallel_map(testFunc1, (1, 2, 3), task_args={'a': 4, 'b': (1, 2, 3)})

but this results in TypeError: can only concatenate tuple (not "dict") to tuple.
When I try

parallel_map(testFunc1, b=(1, 2, 3), task_args={'a': 4})

I get TypeError: parallel_map() missing 1 required positional argument: 'values'.

Does somebody know how to use parallel_map for the n-th positional argument (without writing a function wrapper for each n)?

Answer 1

Looking into the source code of parallel_map reveals why it only works for the first argument of a function:

async_res = [pool.apply_async(task, (value,) + task_args, task_kwargs, _update_progress_bar)
             for value in values]

In this line the parallel processes are created. The function task gets a tuple representing all positional arguments, which is created from 1 element of the values and all the other task_args. Since the element of values is at the first position of the combined tuple ((value,) + task_args) it is always the first argument of the function which varies between its parallel instances. The TypeError: can only concatenate tuple (not "dict") to tuple if a dictionary is used for task_args comes from the fact that the +-operator is only overloaded for (tuple, tuple), not for (tuple, dict).

So in the end there is no way around building a wrapper. Either one for each particular n or a generic one like:

def wrapper(varyingArgument, moreArgs):
    n = moreArgs[0]
    orderedArgs = moreArgs[1:n] + (varyingArgument,) + moreArgs[n:]
    return testFunc1(*orderedArgs)

Calling this with

parallel_map(wrapper, (1,2,3), task_args=((2, 4),))

returns

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

Answer 2

Q : "how to use parallel_map for the n-th positional argument _{(w/o writing a function wrapper for each n)}?"

Avoid creating problems, where they are not, and place whatever your n-th placed externally filled iterable in your def-ed function call-signature right into the parallel_map()-expected ( as documented ) and iteration-processing compliant tuple :

#          (           )-------------------------- parallel_map() expected TUPLE
#your Fun( (  a    vv--)-------------------)----------------your FED-IN ITERABLE
testFunc1( ( 'a', 'b1' ), 'of-no-interest' ) --> (('a', 'b1'), 'of-no-interest')
testFunc1( ( 'a', 'b2' ), 'of-no-interest' ) --> (('a', 'b2'), 'of-no-interest')
testFunc1( ( 'a', 'b3' ), 'of-no-interest' ) --> (('a', 'b2'), 'of-no-interest')

_{"Do you mean something like parallel_map(testFunc1, [(4, 1), (4, 2), (4, 3)], task_args=('of-no-interest',))? Here b always has the value 'of-no-interest'. – A. P. 2 hours ago"}

No,
the example is a clear path to unload a and any and all n-th user-side code FED-IN iterable(s), right as was required above.

def testFun2( a, b ):
    return [ item for item in tuple( a ) ], b

show the way, a call to :

testFun2( ( 'a', 'b', "c", None, 42, -3.14159, "The-N-th-ITERABLE" ),
          'not-important-one-(of-no-interest-HERE-in-solving-N-th-iterable-for-parallel_star()-calls)...'
           )

delivers -->

(['a', 'b', 'c', None, 42, -3.14159, "The-N-th-ITERABLE"], 'not-important-one-(of-no-interest-HERE-in-solving-N-th-iterable-for-parallel_star()-calls)...')

Exactly meeting both a) your wish to have free-hands for any N-th iterable, not just the first positional one and also the b) the very what the call-signature of the parallel_map() expects and was documented to do so :

parallel_map( testFun2,                                        # TASK         Callable
              ( <_USER-SIDE_GENERATOR_FEED-IN_TUPLEsOfPARs_> ),# TASK VALUE(s)Array / List
              any_other_wished2have_call-signature_parameters, # TASK_ARGS    Dict
              ...,                                             # TASK_KWARGS  Dict
              ...                                              # call  KWARGS Dict
              )