#!/usr/bin/env python2.7
"""
This python script illustrates how to find a scaling
law that matches data using log-transformed least squares
regression.
We are fitting the log-linear relationship
log(v) = k + m log(u)
which corresponds to the scaling law
m k
v(u) = u e
"""
from numpy import *
from pylab import loglog, show, xlabel, ylabel, savefig, text, subplot, plot
# data file at http://www.math.psu.edu/treluga/450/egg.data
egg_data = loadtxt('egg.data', delimiter=',')
print egg_data
# First, we log-transform our data set
log_data = log(egg_data)/log(10) # use log base 10
# Now, we construct A and b for A x = b.
b = log_data[:, 1:2]
A = hstack([ones((6, 1)), log_data[:, 0:1]])
# solve A x = b for x using gaussian elimination with partial pivoting
x = linalg.solve( A.T.dot(A), A.T.dot(b) )
k, m = float(x[0]), float(x[1]) # float calls convert 0-d arrays to scalars
print (k, m)
## Alternatively, we could use numpy's built-in least-squares
## algorithm. If you look close, you'll see that the answer
## is a little different from the one above. This is because
## the algorithm is using a different calculation method based
## on SVD decomposition of the matrix. SVD is more powerful and
## more robust, but also slower for large problems.
#
(x_alt, residues, a_rank, svd_vals) = linalg.lstsq(A, b)
print (float(x_alt[0]), float(x_alt[1]))
####################
#
# Now, we display our results
#
####################
z = 10**linspace(0, 6, 10)
loglog(egg_data[:, 0], egg_data[:, 1], 'o', z, 10**k*z**m,'r-')
xlabel('Adult bird mass')
ylabel('Egg mass')
Equation = "egg_mass = %.3f * bird_mass**%.3f"%(10**k, m)
text(300, 0.5, Equation)
savefig('bird_to_egg_scaling.png')
show()