#!/usr/bin/python

# Figure 6.7, page 311.
# Sparse regressor selection.
#
# The problem data are different from the book.

import pylab
from cvxopt import blas, lapack, solvers, matrix, mul
from pickle import load
solvers.options['show_progress'] = 0

data = load(open('regsel.bin','r'))
A, b = data['A'], data['b']
m, n = A.size 

# In the heuristic, set x[k] to zero if abs(x[k]) <= tol * max(abs(x)).
tol = 1e-1

# Data for QP
#
#     minimize    (1/2) ||A*x - b|_2^2 
#     subject to  -y <= x <= y
#                 sum(y) <= alpha

P = matrix(0.0, (2*n,2*n))
P[:n,:n] = A.T*A
q = matrix(0.0, (2*n,1))
q[:n] = -A.T*b
I = matrix(0.0, (n,n)) 
I[::n+1] = 1.0
G = matrix([[I, -I, matrix(0.0, (1,n))], [-I, -I, matrix(1.0, (1,n))]])
h = matrix(0.0, (2*n+1,1))

# Least-norm solution
xln = matrix(0.0, (n,1))
xln[:m] = b
lapack.gels(+A, xln)

nopts = 100
res = [ blas.nrm2(b) ]
card = [ 0 ]
alphas = blas.asum(xln)/(nopts-1) * matrix(range(1,nopts), tc='d')
for alpha in alphas:

    #    minimize    ||A*x-b||_2
    #    subject to  ||x||_1 <= alpha

    h[-1] = alpha
    x = solvers.qp(P, q, G, h)['x'][:n]
    xmax = max(abs(x))
    I = [ k for k in xrange(n) if abs(x[k]) > tol*xmax ]
    if len(I) <= m:
        xs = +b 
        lapack.gels(A[:,I], xs)
        x[:] = 0.0
        x[I] = xs[:len(I)]
        res += [ blas.nrm2(A*x-b) ]
        card += [ len(I) ]

# Eliminate duplicate cardinalities and make staircase plot.
res2, card2 = [], []
for c in xrange(m+1):
    r = [ res[k] for k in xrange(len(res)) if card[k] == c ]
    if r:  
        res2 += [ min(r), min(r) ]
        card2 += [ c, c+1 ]

#pylab.figure(1, facecolor='w')
#pylab.plot( res2[::2], card2[::2], 'o')
#pylab.plot( res2, card2, '-') 
#pylab.xlabel('||A*x-b||_2')
#pylab.ylabel('card(x)')
#pylab.title('Sparse regressor selection (fig 6.7)')
#print "Close figure to start exhaustive search."
#pylab.show()


# Exhaustive search.

def patterns(k,n):
    """
    Generates all 0-1 sequences of length n with exactly k nonzeros.
    """
    if k==0:
        yield n*[0]
    else:
        for x in patterns(k-1,n-1): yield [1] + x
        if k <= n-1:
            for x in patterns(k,n-1): yield [0] + x

             
bestx = matrix(0.0, (n, m))   # best solution for each cardinality
bestres = matrix(blas.nrm2(b), (1, m+1))   # best residual
x = matrix(0.0, (n,1))
for k in xrange(1,m):
    for s in patterns(k,n):
        I = [ i for i in xrange(n) if s[i] ]
        s = [k] + s 
        print ("%d nonzeros: " + n*"%d") %tuple(s)
        x = +b
        lapack.gels(A[:,I], x)
        res = blas.nrm2(b - A[:,I] * x[:k])
        if res < bestres[k]:
            bestres[k] = res 
            bestx[:,k][I] = x[:k] 
bestres[m] = 0.0

pylab.figure(1, facecolor='w')

# heuristic result
pylab.plot( res2[::2], card2[::2], 'o' )
pylab.plot( res2, card2, '-') 

# exhaustive result
res2, card2 = [ bestres[0] ], [ 0 ] 
for k in xrange(1,m+1):
    res2 += [bestres[k-1], bestres[k]]
    card2 += [ k, k]
pylab.plot( bestres.T, range(m+1), 'go')
pylab.plot( res2, card2, 'g-')

pylab.xlabel('||A*x-b||_2')
pylab.ylabel('card(x)')
pylab.title('Sparse regressor selection (fig 6.7)')
pylab.show()