[FIXED] How do I vectorize a function in numpy with some fixed parameters?

Issue

I have written a code for approximating a function with the Bernstein polynomials ( https://en.wikipedia.org/wiki/Bernstein_polynomial )

at

https://github.com/pdenapo/metodos-numericos/blob/master/python/bernstein.py

I have a function that gives the polynomial approximating f as bernstein(f, n, p) (where f is the function that I want to approximate, n is the degree and p the point where it is evaluated.

def bernstein(f, n, p):
    return np.sum(
        [f(k / n) * st.binom.pmf(k, n, p) for k in np.arange(0, n + 1)])

Now I want to generate a plot of this function where f and n es fixed, and p runs though a vector generated by np.arrange So I am vectorizing the function in the following way:

 bernstein3 = lambda x: bernstein(f, 3, x)
 bernstein3 = np.vectorize(bernstein3)
 y3 = bernstein3(x)
 plt.plot(x, y3, 'green', label='$B_3$')

It works. But I guess there must be some more elegant, or perhaps more pythonic way of doing this. Any suggestions? Many thanks

Solution

Since SciPy statistic functions are vectorized, your bernstein function can be modified in a straightforward manner to work that way:

import numpy as np
import scipy.stats

def bernstein(f, n, p):
    # Vector of k values
    k = np.arange(n + 1)
    # Add a broadcasting dimension to p
    pd = np.expand_dims(p, -1)
    # Compute approximation
    return np.sum(f(k / n) * scipy.stats.binom.pmf(k, n, pd), -1)

It would be used simply as this:

import numpy as np
import matplotlib.pyplot as plt

def f(x):
    return np.abs(1 / 2 - x)

x = np.linspace(0, 1, 100)
y = f(x)
plt.plot(x, y, 'blue', label='f(x)')
y_approx = bernstein(f, 10, x)
plt.plot(x, y_approx, 'orange', label='f_approx(x)')
plt.show()

Answered By - jdehesa