```#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

"""This module is to find lowest eigenvalues with Davidson algorithm."""

import logging
import warnings

import numpy
import numpy.linalg
import scipy
import scipy.linalg
import scipy.sparse
import scipy.sparse.linalg

from openfermion.utils._sparse_tools import get_linear_qubit_operator_diagonal
from openfermion.utils._linear_qubit_operator import \
generate_linear_qubit_operator

class DavidsonError(Exception):
"""Exceptions."""
pass

class DavidsonOptions(object):
"""Davidson algorithm iteration options."""

def __init__(self, max_subspace=100, max_iterations=300, eps=1e-6,
real_only=False):
"""
Args:
max_subspace(int): Max number of vectors in the auxiliary subspace.
max_iterations(int): Max number of iterations.
eps(float): The max error for eigen vector error's elements during
iterations: linear_operator * v - v * lambda.
real_only(bool): Desired eigenvectors are real only or not. When one
specifies the real_only to be true but it only has complex ones,
no matter it converges or not, the returned vectors will be
complex.
"""
if max_subspace <= 2 or max_iterations <= 0 or eps <= 0:
raise ValueError('Invalid values for max_subspace, max_iterations '
'and/ or eps: ({}, {}, {}).'.format(
max_subspace, max_iterations, eps))

self.max_subspace = max_subspace
self.max_iterations = max_iterations
self.eps = eps
self.real_only = real_only

def set_dimension(self, dimension):
"""
Args:
dimension(int): Dimension of the matrix, which sets a upper limit on
the work space.
"""
if dimension <= 0:
raise ValueError('Invalid dimension: {}).'.format(dimension))
self.max_subspace = min(self.max_subspace, dimension + 1)

class Davidson(object):
"""Davidson algorithm to get the n states with smallest eigenvalues."""

def __init__(self, linear_operator, linear_operator_diagonal, options=None):
"""
Args:
linear_operator(scipy.sparse.linalg.LinearOperator): The linear
operator which defines a dot function when applying on a vector.
linear_operator_diagonal(numpy.ndarray): The linear operator's
diagonal elements.
options(DavidsonOptions): Iteration options.
"""
if options is None:
options = DavidsonOptions()

if not isinstance(linear_operator,
(scipy.sparse.linalg.LinearOperator,
scipy.sparse.spmatrix)):
raise ValueError(
'linear_operator is not a LinearOperator: {}.'.format(type(
linear_operator)))

self.linear_operator = linear_operator
self.linear_operator_diagonal = linear_operator_diagonal

self.options = options
self.options.set_dimension(len(linear_operator_diagonal))

def get_lowest_n(self, n_lowest=1, initial_guess=None, max_iterations=None):
"""
Returns `n` smallest eigenvalues and corresponding eigenvectors for
linear_operator.

Args:
n(int):
The number of states corresponding to the smallest eigenvalues
and associated eigenvectors for the linear_operator.
initial_guess(numpy.ndarray[complex]): Initial guess of eigenvectors
associated with the `n` smallest eigenvalues.
max_iterations(int): Max number of iterations when not converging.

Returns:
success(bool): Indicates whether it converged, i.e. max elementwise
error is smaller than eps.
eigen_values(numpy.ndarray[complex]): The smallest n eigenvalues.
eigen_vectors(numpy.ndarray[complex]): The smallest n eigenvectors
corresponding with those eigen values.
"""
# Goes through a few checks and preprocessing before iterative
# diagonalization.

# 1. Checks for number of states desired, should be in the range of
# [0, max_subspace).
if n_lowest <= 0 or n_lowest >= self.options.max_subspace:
raise ValueError('n_lowest {} is supposed to be in [1, {}).'.format(
n_lowest, self.options.max_subspace))

# 2. Checks for initial guess vectors' dimension is the same to that of
# the operator.
if initial_guess is None:
initial_guess = generate_random_vectors(
len(self.linear_operator_diagonal), n_lowest,
real_only=self.options.real_only)
if initial_guess.shape != len(self.linear_operator_diagonal):
raise ValueError('Guess vectors have a different dimension with '
'linear opearator diagonal elements: {} != {}.'
.format(initial_guess.shape,
len(self.linear_operator_diagonal)))

# 3. Makes sure real guess vector if real_only is specified.
if self.options.real_only:
if not numpy.allclose(numpy.real(initial_guess), initial_guess):
warnings.warn('Initial guess is not real only!', RuntimeWarning)
initial_guess = numpy.real(initial_guess)

# 4. Checks for non-trivial (non-zero) initial guesses.
if numpy.max(numpy.abs(initial_guess)) < self.options.eps:
raise ValueError('Guess vectors are all zero! {}'.format(
initial_guess.shape))
initial_guess = scipy.linalg.orth(initial_guess)

# 5. Makes sure number of initial guess vector is at least n_lowest.
if initial_guess.shape < n_lowest:
initial_guess = append_random_vectors(
initial_guess, n_lowest - initial_guess.shape,
real_only=self.options.real_only)

success = False
num_iterations = 0
guess_v = initial_guess
guess_mv = None
max_iterations = max_iterations or self.options.max_iterations
while (num_iterations < max_iterations and not success):
(eigen_values, eigen_vectors, mat_eigen_vectors, max_trial_error,
guess_v, guess_mv) = self._iterate(n_lowest, guess_v, guess_mv)
logging.info("Eigenvalues for iteration %d: %s, error is %f.",
num_iterations, eigen_values, max_trial_error)

if max_trial_error < self.options.eps:
success = True
break

# Make sure it keeps real components only.
if self.options.real_only:
guess_v = numpy.real(guess_v)

# Deals with new directions to make sure they're orthonormal.
# Also makes sure there're new directions added for the next
# iteration, if not, add n_lowest random vectors.
count_mvs = guess_mv.shape
guess_v = orthonormalize(guess_v, count_mvs, self.options.eps)
if guess_v.shape <= count_mvs:
guess_v = append_random_vectors(
guess_v, n_lowest, real_only=self.options.real_only)

# Limits number of vectors to self.options.max_subspace, in this
# case, keep the following:
# 1) first n_lowest eigen_vectors;
# 2) first n_lowest matrix multiplication result for eigen_vectors;
#
# 3) new search directions which will be used for improvement for
#    the next iteration.
if guess_v.shape >= self.options.max_subspace:
guess_v = numpy.hstack([
eigen_vectors,
guess_v[:, count_mvs:],
])
guess_mv = mat_eigen_vectors

if self.options.real_only:
if (not numpy.allclose(numpy.real(guess_v), guess_v) or
not numpy.allclose(numpy.real(guess_mv), guess_mv)):
# Forces recalculation for matrix multiplication with
# vectors.
guess_mv = None

num_iterations += 1

if (self.options.real_only and
not numpy.allclose(numpy.real(eigen_vectors), eigen_vectors)):
warnings.warn('Unable to get real only eigenvectors, return '
'complex vectors instead with success state {}.'
.format(success), RuntimeWarning)
return success, eigen_values, eigen_vectors

def _iterate(self, n_lowest, guess_v, guess_mv=None):
"""One iteration with guess vectors.

Args:
n_lowest(int): The first n_lowest number of eigenvalues and
eigenvectors one is interested in.
guess_v(numpy.ndarray(complex)): Guess eigenvectors associated with
the smallest eigenvalues.
guess_mv(numpy.ndarray(complex)): Matrix applied on guess_v,
therefore they should have the same dimension.

Returns:
trial_lambda(numpy.ndarray(float)): The minimal eigenvalues based on
guess eigenvectors.
trial_v(numpy.ndarray(complex)): New guess eigenvectors.
trial_mv(numpy.ndarray(complex)): New guess eigenvectors' matrix
multiplication result.
max_trial_error(float): The max elementwise error for all guess
vectors.

guess_v(numpy.ndarray(complex)): Cached guess eigenvectors to avoid
recalculation for the next iterations.
guess_mv(numpy.ndarray(complex)): Cached guess vectors which is the
matrix product of linear_operator with guess_v.
"""
if guess_mv is None:
guess_mv = self.linear_operator.dot(guess_v)
dimension = guess_v.shape

# Note that getting guess_mv is the most expensive step.
if guess_mv.shape < dimension:
guess_mv = numpy.hstack([guess_mv, self.linear_operator.dot(
guess_v[:, guess_mv.shape : dimension])])
guess_vmv = numpy.dot(guess_v.conj().T, guess_mv)

# Gets new set of eigenvalues and eigenvectors in the vmv space, with a
# smaller dimension which is the number of vectors in guess_v.
#
# Note that we don't get the eigenvectors directly, instead we only get
# a transformation based on the raw vectors, so that mv don't need to be
# recalculated.
trial_lambda, trial_transformation = numpy.linalg.eigh(guess_vmv)

# Sorts eigenvalues in ascending order.
sorted_index = list(reversed(trial_lambda.argsort()[::-1]))
trial_lambda = trial_lambda[sorted_index]
trial_transformation = trial_transformation[:, sorted_index]

if len(trial_lambda) > n_lowest:
trial_lambda = trial_lambda[:n_lowest]
trial_transformation = trial_transformation[:, :n_lowest]

# Estimates errors based on diagonalization in the smaller space.
trial_v = numpy.dot(guess_v, trial_transformation)
trial_mv = numpy.dot(guess_mv, trial_transformation)
trial_error = trial_mv - trial_v * trial_lambda

new_directions, max_trial_error = self._get_new_directions(
trial_error, trial_lambda, trial_v)
if new_directions:
guess_v = numpy.hstack([guess_v, numpy.stack(new_directions).T])
return (trial_lambda, trial_v, trial_mv, max_trial_error,
guess_v, guess_mv)

def _get_new_directions(self, error_v, trial_lambda, trial_v):
"""Gets new directions from error vectors.

Args:
error_v(numpy.ndarray(complex)): Error vectors from the guess
eigenvalues and associated eigenvectors.
trial_lambda(numpy.ndarray(float)): The n_lowest minimal guess
eigenvalues.
trial_v(numpy.ndarray(complex)): Guess eigenvectors associated with
trial_lambda.

Returns:
new_directions(numpy.ndarray(complex)): New directions for searching
for real eigenvalues and eigenvectors.
max_trial_error(float): The max elementwise error for all guess
vectors.
"""
n_lowest = error_v.shape

max_trial_error = 0
# Adds new guess vectors for the next iteration for the first n_lowest
# directions.
origonal_dimension = error_v.shape

new_directions = []
for i in range(n_lowest):
current_error_v = error_v[:, i]

if numpy.max(numpy.abs(current_error_v)) < self.options.eps:
# Already converged for this eigenvector, no contribution to
# search for new directions.
continue

max_trial_error = max(max_trial_error,
numpy.linalg.norm(current_error_v))
diagonal_inverse = numpy.ones(origonal_dimension)
for j in range(origonal_dimension):
# Makes sure error vectors are bounded.
diff_lambda = self.linear_operator_diagonal[j] - trial_lambda[i]
if numpy.abs(diff_lambda) > self.options.eps:
diagonal_inverse[j] /= diff_lambda
else:
diagonal_inverse[j] /= self.options.eps
diagonal_inverse_error = diagonal_inverse * current_error_v
diagonal_inverse_trial = diagonal_inverse * trial_v[:, i]
new_direction = -current_error_v + (trial_v[:, i] * numpy.dot(
trial_v[:, i].conj(), diagonal_inverse_error) / numpy.dot(
trial_v[:, i].conj(), diagonal_inverse_trial))

new_directions.append(new_direction)
return new_directions, max_trial_error

class QubitDavidson(Davidson):
"""Davidson algorithm applied to a QubitOperator."""

def __init__(self, qubit_operator, n_qubits=None, options=None):
"""
Args:
qubit_operator(QubitOperator): A qubit operator which is a linear
operator as well.
n_qubits(int): Number of qubits.
options(DavidsonOptions): Iteration options.
"""
super(QubitDavidson, self).__init__(
generate_linear_qubit_operator(qubit_operator, n_qubits, options),
get_linear_qubit_operator_diagonal(qubit_operator, n_qubits),
options=options)

class SparseDavidson(Davidson):
"""Davidson algorithm for a sparse matrix."""

def __init__(self, sparse_matrix, options=None):
"""
Args:
sparse_matrix(scipy.sparse.spmatrix): A sparse matrix in scipy.
options(DavidsonOptions): Iteration options.
"""
super(SparseDavidson, self).__init__(
sparse_matrix, sparse_matrix.diagonal(), options=options)

def generate_random_vectors(row, col, real_only=False):
"""Generates orthonormal random vectors with col columns.

Args:
row(int): Number of rows for the vectors.
col(int): Number of columns for the vectors.
real_only(bool): Real vectors or complex ones.

Returns:
random_vectors(numpy.ndarray(complex)): Orthonormal random vectors.
"""
random_vectors = numpy.random.rand(row, col)
if not real_only:
random_vectors = random_vectors + numpy.random.rand(row, col) * 1.0j
random_vectors = scipy.linalg.orth(random_vectors)
return random_vectors

def append_random_vectors(vectors, col, max_trial=3, real_only=False):
"""Appends exactly col orthonormal random vectors for vectors.

Assumes vectors is already orthonormal.

Args:
vectors(numpy.ndarray(complex)): Orthonormal original vectors to be
appended.
col(int): Number of columns to be appended.
real_only(bool): Real vectors or complex ones.

Returns:
vectors(numpy.ndarray(complex)): Orthonormal vectors with n columns.
"""
if col <= 0:
return vectors

vector_columns = vectors.shape
total_columns = min(vector_columns + col, vectors.shape + 1)

num_trial = 0
while vector_columns < total_columns:
num_trial += 1

vectors = numpy.hstack([vectors, generate_random_vectors(
vectors.shape, total_columns - vector_columns, real_only)])
vectors = orthonormalize(vectors, vector_columns)

# Checks whether there are any new vectors added successfully.
if vectors.shape == vector_columns:
if num_trial > max_trial:
warnings.warn('Unable to generate specified number of random '
'vectors {}: returning {} in total.'.format(
col, vector_columns), RuntimeWarning)
break
else:
num_trial = 1
vector_columns = vectors.shape
return vectors

def orthonormalize(vectors, num_orthonormals=1, eps=1e-6):
"""Orthonormalize vectors, so that they're all normalized and orthogoal.

The first vector is the same to that of vectors, while vector_i is
orthogonal to vector_j, where j < i.

Args:
vectors(numpy.ndarray(complex)): Input vectors to be
orthonormalized.
num_orthonormals(int): First `num_orthonormals` columns are already
orthonormal, so that one doesn't need to make any changes.
eps(float): criterion of elements' max absolute value for zero vectors.

Returns:
ortho_normals(numpy.ndarray(complex)): Output orthonormal vectors.
"""
ortho_normals = vectors
count_orthonormals = num_orthonormals
# Skip unchanged ones.
for i in range(num_orthonormals, vectors.shape):
vector_i = vectors[:, i]
# Makes sure vector_i is orthogonal to all processed vectors.
for j in range(i):
vector_i -= ortho_normals[:, j] * numpy.dot(
ortho_normals[:, j].conj(), vector_i)

# Makes sure vector_i is normalized.
if numpy.max(numpy.abs(vector_i)) < eps:
continue
ortho_normals[:, count_orthonormals] = (vector_i /
numpy.linalg.norm(vector_i))
count_orthonormals += 1
return ortho_normals[:, :count_orthonormals]
```