Calling a function containing an array of x values

Question

I want to call a function in a program, that has the same format as the following, but where the x values are in the form of an array of shape = (426, 240). Can someone help with this?

The function is:

def f(x):
    if x < 0:
        return -2*x
    else :
        return -x

    x = np.arange(-100, 100, 1)

    plt.plot(x, list(map(f, x)), 'b-')  # for python3

    #plt.show()

The part of the code that calls the function would look like this:

def nucleation_and_motion_in_G_gradient_fluid_2D(writer, args, R=60):
    dx = 2*R / args.height
    x = (np.arange(args.width) - args.width // 2) * dx
    y = (np.arange(args.height) - args.height // 2) * dx
    x, y = np.meshgrid(x, y, indexing='ij')

def source_G(t):
    center = np.exp(-0.5*(t-5)**2) * 10
    gradient = (1+np.tanh(t-30)) * 0.0003
    piecewise_1 = f(x) # ***function f(x) called here***

    return -( 
        np.exp(-0.5*(x*x + y*y)) #+ np.exp(-0.5*((x)**2 + y*y))
    ) * center + piecewise_1 * gradient   # piecewise function test

Main code here.

I already know the code works for a trapezoid function in combination with the x array, as follows:

(code requires: from scipy import signal)

def trapezoid_signal(x, width=2., slope=1., amp=10., offs=1):
        a = slope * width * signal.sawtooth(2 * np.pi * 1/10 * x/width - 0.8, width=0.5)/4.
        a[a>amp/2.] = amp/2.
        a[a<-amp/2.] = -amp/2.
        return a + amp/2. + offs

def source_G(t):
    center = np.exp(-0.5*(t-5)**2) * 10
    gradient = (1+np.tanh(t-30)) * 0.0003
    trapezoid = trapezoid_signal(x, width=40, slope=5, amp=50)

    return -( 
        np.exp(-0.5*(x**2 + y**2)) 
    ) * center + trapezoid * gradient # one soliton particle in 2 dimensions of xy with z axis as concentration potential

Answer 1

If you want to make this

def f(x):
    if x < 0:
        return -2*x
    else :
        return -x

compatible with vectorization, you can use the following (very common) trick:

def f(x):
    neg = x < 0
    return neg * (-2 * x) + (1 - neg) * -x

It works!

>>> f(np.arange(-5, 5))
array([10,  8,  6,  4,  2,  0, -1, -2, -3, -4])