#! /usr/bin/env python
"""Perform some linear algebra and vector operations

Date: November 14, 2008

Author: Bob Wimmer

"""

from math import *   #This imports all math functions and constants
from scipy.integrate import odeint
from numpy import *
from numpy.linalg import inv,solve
#import Gnuplot, Gnuplot.funcutils
import sys, time


def main():
    """Do something"""
    a = array([[3,1,5],[1,0,8],[2,1,4]])
    inva = inv(a)
    print('inv(a)=', inva)
    one = dot(inva,a)
    print('this should be unity matrix:', one)
    b = array([1.,2.,3.])
    x=solve(a,b)
    print('x=',x)
    bb = dot(a,x)
    print('bb=',bb)

    print('------------------------------------------')
# Now solve some less standard problem:
# Sometimes we have more data than equations or vice versa,
# more equations than data

    from numpy.linalg import pinv,svd,lstsq
    A = array([[1., 3., 5.],[2., 4., 6.]])
    B = array([1., 3.])
    # Question: find x such that ||A*x-b|| is minimal
    # Answer:   x = pinvA * b,  with pinvA the pseudo-inverse of A
    pinvA = pinv(A)
    print('pinvA=', pinvA)
    x = dot(pinvA, B)
    print('x=',x)
    print('A*x=', dot(A,x))

    #and now let's check what this has to do with least squares

    x,resids,rank,s = lstsq(A,B, rcond=None)
    print('x,resids,rank,s = ',x,resids,rank,s)


    #and with singular value decomposition

    print('-------------------------------------')
    
    U,sigma,V = svd(A)              #these are the three parts in A's SVD
    S = zeros_like(A.transpose())   #fill S with zeros with hape of A
    for n in range(len(sigma)): S[n,n] = 1. / sigma[n] #fill diagonal with inverse of sigma (that's the SVD trick!)
    print(A)
    print(pinvA)
    print(dot(V.transpose(), dot(S, U.transpose())))
    




main()