Numpy: Elegant solution to calculating output matrix of function with indefinite number of variables

Question

I have been searching all weekend for a more elegant (read: no manual for-loop programming) solution to the following problem:

Lets assume we have a custom function f() with an indefinite number of inputs. For simplicity, let us start with two:

def f(x,y):
    return x + y

Now I pass an array with the correct number of variables to this function:

x = np.array([x0, x1, x2, ..., xn])
y = np.array([y0, y1, y2, ..., yn])

The answer I look for is:

z = np.array([[x0 + y0, x0 + y1, x0 + y2, ..., x0 + yn],
              [x1 + y0, x1 + y1, x1 + y2, ..., x1 + yn],
              [x2 + y0, x2 + y1, x2 + y2, ..., x2 + yn],
              ...])

So, in summary, I am looking for a function, that I can pass another custom function to, which then calculates all possible combinations, without me having to program a ridiculous number of for-loops.

Please help me, hive mind!

Edit 1: the custom function may be of arbitrary complexity. From my real world problem, this is one example:

def f(x, y):
    return 1 - (x/2)**y*binom(y, y/2)

Edit 2: the accepted answer works as intended. The answer of Dishin H Goyani linking to stackoverflow.com/a/32742943/6075699 produces the same result using a different path.

Thanks everyone! Stackoverflow rules!

Outside of NumPy, you could probably make use of [`itertools.product`](https://docs.python.org/3/library/itertools.html#itertools.product). — 9769953, Feb 17 '20 at 08:53
You might want to look at numpy array broadcasting - https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html — Axeon Thra, Feb 17 '20 at 09:19
@DishinHGoyani: This is brilliant! Works like a charm for a two dimensional problem! — olivherbst, Feb 17 '20 at 09:19

score 1 · Answer 1 · edited Feb 17 '20 at 09:03

1

It seems like you're looking for an application of the dot product:

In [1]: a = np.array([1, 2, 3, 4, 5, 6, 7])
In [2]: b = np.array([1, 2, 3, 4, 5, 6, 7])
In [3]: np.dot(np.matrix(a).T, np.matrix(b))
Out[3]: 
 matrix([[ 1,  2,  3,  4,  5,  6,  7],
         [ 2,  4,  6,  8, 10, 12, 14],
         [ 3,  6,  9, 12, 15, 18, 21],
         [ 4,  8, 12, 16, 20, 24, 28],
         [ 5, 10, 15, 20, 25, 30, 35],
         [ 6, 12, 18, 24, 30, 36, 42],
         [ 7, 14, 21, 28, 35, 42, 49]])

This only works, if you convert your array into something that has 2 dimensions (like the matrix object) of which one is a singleton dimension. Then transpose one of them and the dot product can give you exactly what you need.

edited Feb 17 '20 at 09:03

Brown Bear

19,655
10
58
76

answered Feb 17 '20 at 08:57

meetaig

913
10
26

I don't see, how this would work for a custom function with higher complexity? Please see my edit! – olivherbst Feb 17 '20 at 09:09
Even your first cell is wrong, unless I'm misreading how you intend to flatten that out. Not the first element should be 1 + 1 according to OP. – kabanus Feb 17 '20 at 09:11

kabanus · Accepted Answer · 2020-02-17T09:24:35.730

Mesh grid can help you create all the pairs, than just sum them. As a bonus it's extensible to more dimensions:

>>> y = np.arange(1,5)
>>> x = np.arange(6,10)
>>> x
array([6, 7, 8, 9])
>>> y
array([1, 2, 3, 4])
>>> sum(np.meshgrid(x,y))
array([[ 7,  8,  9, 10],
       [ 8,  9, 10, 11],
       [ 9, 10, 11, 12],
       [10, 11, 12, 13]])

To put this in a function taking an unknown number of arrays:

def meshSum(*arrays):
    return sum(np.meshgrid(*arrays))

An example with another array:

>>> z = np.arange(11,15)
>>> def meshSum(*arrays):
...     return sum(np.meshgrid(*arrays))
...
>>> meshSum(x,y,z)
array([[[18, 19, 20, 21],
        [19, 20, 21, 22],
        [20, 21, 22, 23],
        [21, 22, 23, 24]],

       [[19, 20, 21, 22],
        [20, 21, 22, 23],
        [21, 22, 23, 24],
        [22, 23, 24, 25]],

       [[20, 21, 22, 23],
        [21, 22, 23, 24],
        [22, 23, 24, 25],
        [23, 24, 25, 26]],

       [[21, 22, 23, 24],
        [22, 23, 24, 25],
        [23, 24, 25, 26],
        [24, 25, 26, 27]]])

Following your edit, for an arbitrary operation on the mesh

def meshOperation(f, *arrays):
    return f((*np.meshgrid(*arrays))

where f must take either a *args argument or a number of arguments equal to len(arrays), so meshOperation(f, x, y) is valid in your final example.

After edit: This definitely works for 2D problems, for 3D I will try next. Weirdly, it switches the input order around, such that you have to transpose the output array. — olivherbst, Feb 17 '20 at 09:40
@olivherbst weird and expected are a matter of your taste and experience. Use `.T` in the result to transpose it. — kabanus, Feb 17 '20 at 09:41

score 0 · Answer 3 · answered Feb 17 '20 at 09:25

I suggest numpy array broadcasting.

Example:

import numpy as np

# initial arrays
x = np.arange(1, 15, 3)
y = np.arange(1, 6) + 100

# get them to 2d
x2 = np.atleast_2d(x)
y2 = np.atleast_2d(y).T  #y should be vertical

# simple stuff
print("sum:\n", x2 + y2)

# complicated stuff
print("complicated:\n", x2/(1+y2) + np.exp(-y2/(1+x2)))

Numpy: Elegant solution to calculating output matrix of function with indefinite number of variables

3 Answers3