%pylab inline
# This function appears to be written for scalar values only.
def f(x):
return 3 * x * x - 1
print f(-1.0)
print f(0.0)
print f(1.0)
# But thanks to NumPy it can operate on arrays as well!
import numpy as np
x = np.array([-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0])
print f(x)
# Graphing the result.
import pylab as pl
plot = pl.figure().add_subplot(111)
plot.plot(x, f(x))
plot.grid()
# Use linspace() to make lots of x values for a smoother plot.
x = linspace(-3.0, 3.0, 200)
plot = pl.figure().add_subplot(111)
plot.plot(x, f(x))
plot.grid()
# What if we start with data and want the parabola?
# We can use optimize() to fit a curve to noisy data.
x = numpy.array([-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0])
y = numpy.array([10.3, 5.1, 0.9, 3.6, 8.4, 16.3, 32.7])
def parabola(x, a, b, c):
return a*x*x + b*x + c
from scipy.optimize import curve_fit
(a, b, c), pcov = curve_fit(parabola, x, y)
print a, b, c
# Here is the curve that optimize() figured out,
# plotted next to the points we used for input.
plot = pl.figure().add_subplot(111)
plot.scatter(x, y)
plot.plot(x, parabola(x, a, b, c))