4

I have a function that depends on several variables, let's say y=f(x1,x2,x3,x4). If each of the variables is a simple number, then the result should be a plain number. If one of the variables is an array, I need the result to be also an array. And so on: if two of the variables are arrays, I need the result to be a 2-dimensional array. Etc.

Example:

def f(x1,x2,x3,x4):
    y=x1*x2/(x3+x4)
    return y

x1=1.0
x2=2.0
x3=3.0
x4=4.0
f(x1,x2,x3,x4)
# should give 2.0/7.0 = 0.2857...

x3=array([1.0,2.0,3.0,4.0,5.0])
f(x1,x2,x3,x4)
# should give a one-dimensional array with shape (5,)

x4=array([10.0,20.0,30.0,40.0,50.0,60.0,70.0])
f(x1,x2,x3,x4)
# should give a two-dimensional array with shape (5,7)

How to do it? (To be as clear as possible for a non-Python reader of my program?)

ali_m
  • 71,714
  • 23
  • 223
  • 298
Amenhotep
  • 920
  • 1
  • 13
  • 18

2 Answers2

2

The proper way to do this is to pass in the properly shaped data. If you want a 2d result, you should pass in 2d arrays. This can be accomplished by np.newaxis.

import numpy as np

def f(x1,x2,x3,x4):
    y = x1*x2/(x3+x4)
    return y

x1 = 1.0
x2 = 2.0
x3 = 3.0
x4 = 4.0
print f(x1,x2,x3,x4)

x3 = np.array([1.0,2.0,3.0,4.0,5.0])
print f(x1,x2,x3,x4)

x3 = x3[:, np.newaxis]
x4 = np.array([10.0,20.0,30.0,40.0,50.0,60.0,70.0])
x4 = x4[np.newaxis, :]
print f(x1,x2,x3,x4)

Of course, the way your question is posed, it's a little ambiguous why you should expect to get an array shaped (5, 7) and not an array shaped (7, 5).

mgilson
  • 300,191
  • 65
  • 633
  • 696
  • I understand. Please give me an example of all of the variables being arrays. Then I have to add... how many newaxis to each? Thanks – Amenhotep Dec 20 '13 at 19:27
  • Shape (5,7) or (7,5) is OK, as long as I get to understand the rules (which I don't, yet :o). – Amenhotep Dec 20 '13 at 19:29
  • @Amenhotep -- Have a look at numpy's [broadcasting documentation](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html) That explains the rules. In summary, if `x` is a 1d array with shape `(6,)`, then `x[:, np.newaxis]` will give an array of shape `(6, 1)` and `x[np.newaxis, :]` gives an array of shape `(1, 6)`. Now when you do an arithmetic operation on a `(5, 1)` and `(1, 7)` the result is a `(5, 7)`, etc. – mgilson Dec 20 '13 at 19:35
  • Yes, this part I understood. Could you please give an example with three variables being arrays? Thank you. – Amenhotep Dec 20 '13 at 19:38
  • `x[:, np.newaxis, np.newaxis] + y[np.newaxis, :, np.newaxis] + z[np.newaxis, np.newaxis, :]` is one example where `x`, `y` and `z` are all 1D arrays. – mgilson Dec 20 '13 at 19:41
0

You might look at using @when to define the function appropriately.

For example:

# Case for single values
@when(int, int, int, int)
def f(x1,x2,x3,x4):
    y = x1*x2/(x3+x4)
    return y

# Case for 1D arrays
@when(list, list, list, list)
def f(x1,x2,x3,x4):
    y = []
    for i in range(len(x1)):
        y.append(x1[i]*x2[i]/(x3[i]+x4[i]))
    return y
Cory Kramer
  • 114,268
  • 16
  • 167
  • 218
  • Interesting. But if I want to cover all the cases, I need to write many @when's, right? For example when x1 and x3 are numbers and x2 and x4 are arrays. Can't it detect automatically all these combinations? – Amenhotep Dec 20 '13 at 19:32
  • Isn't the result a one dimensional array? I need it to be multi-dimensional (if more than one input is an array). – Amenhotep Dec 20 '13 at 19:34
  • What's `when`? Google didn't help. – user2357112 Dec 20 '13 at 19:42
  • It is from the package PEAK-rules: https://pypi.python.org/pypi/PEAK-Rules. It is for making generic functions, among other things. – Cory Kramer Dec 20 '13 at 19:59