#!/usr/bin/env python
# coding: utf-8

# # Nonnegative matrix factorization
# 
# A derivative work by Judson Wilson, 6/2/2014.    
# Adapted from the CVX example of the same name, by Argyris Zymnis, Joelle Skaf and Stephen Boyd
# 
# ## Introduction
# 
# We are given a matrix $A \in \mathbf{\mbox{R}}^{m \times n}$ and are interested in solving the problem:
#     \begin{array}{ll}
#     \mbox{minimize}   & \| A - YX \|_F \\
#     \mbox{subject to} & Y \succeq 0 \\
#                       & X \succeq 0,
#     \end{array}
# where $Y \in \mathbf{\mbox{R}}^{m \times k}$ and $X \in \mathbf{\mbox{R}}^{k \times n}$.
# 
# This example generates a random matrix $A$ and obtains an
# *approximate* solution to the above problem by first generating
# a random initial guess for $Y$ and then alternatively minimizing
# over $X$ and $Y$ for a fixed number of iterations.
# 
# ## Generate problem data

# In[1]:


import cvxpy as cp
import numpy as np

# Ensure repeatably random problem data.
np.random.seed(0)

# Generate random data matrix A.
m = 10
n = 10
k = 5
A = np.random.rand(m, k).dot(np.random.rand(k, n))

# Initialize Y randomly.
Y_init = np.random.rand(m, k)


# ## Perform alternating minimization

# In[2]:


# Ensure same initial random Y, rather than generate new one
# when executing this cell.
Y = Y_init 

# Perform alternating minimization.
MAX_ITERS = 30
residual = np.zeros(MAX_ITERS)
for iter_num in range(1, 1+MAX_ITERS):
    # At the beginning of an iteration, X and Y are NumPy
    # array types, NOT CVXPY variables.

    # For odd iterations, treat Y constant, optimize over X.
    if iter_num % 2 == 1:
        X = cp.Variable(shape=(k, n))
        constraint = [X >= 0]
    # For even iterations, treat X constant, optimize over Y.
    else:
        Y = cp.Variable(shape=(m, k))
        constraint = [Y >= 0]
    
    # Solve the problem.
    # increase max iters otherwise, a few iterations are "OPTIMAL_INACCURATE"
    # (eg a few of the entries in X or Y are negative beyond standard tolerances)
    obj = cp.Minimize(cp.norm(A - Y*X, 'fro'))
    prob = cp.Problem(obj, constraint)
    prob.solve(solver=cp.SCS, max_iters=10000)

    if prob.status != cp.OPTIMAL:
        raise Exception("Solver did not converge!")
    
    print('Iteration {}, residual norm {}'.format(iter_num, prob.value))
    residual[iter_num-1] = prob.value

    # Convert variable to NumPy array constant for next iteration.
    if iter_num % 2 == 1:
        X = X.value
    else:
        Y = Y.value


# ## Output results

# In[3]:


#
# Plot residuals.
#

import matplotlib.pyplot as plt

# Show plot inline in ipython.
get_ipython().run_line_magic('matplotlib', 'inline')

# Set plot properties.
plt.rc('text', usetex=True)
plt.rc('font', family='serif')
font = {'weight' : 'normal',
        'size'   : 16}
plt.rc('font', **font)

# Create the plot.
plt.plot(residual)
plt.xlabel('Iteration Number')
plt.ylabel('Residual Norm')
plt.show()

#
# Print results.
#
print('Original matrix:')
print(A)
print('Left factor Y:')
print(Y)
print('Right factor X:')
print(X)
print('Residual A - Y * X:')
print(A - Y.dot(X))
print('Residual after {} iterations: {}'.format(iter_num, prob.value))