def _set_initial_v(vinit_type,
dtype,
device,
batch_dims,
na,
nguess,
M=None,
mparams=None):
torch.manual_seed(12421)
if vinit_type == "eye":
nbatch = functools.reduce(lambda x, y: x * y, bcast_dims, 1)
V = torch.eye((na, nguess), dtype=dtype,
device=device).unsqueeze(0).repeat(nbatch, 1, 1).reshape(
*batch_dims, na, nguess)
elif vinit_type == "randn":
V = torch.randn((*batch_dims, na, nguess), dtype=dtype, device=device)
elif vinit_type == "random" or vinit_type == "rand":
V = torch.rand((*batch_dims, na, nguess), dtype=dtype, device=device)
else:
raise ValueError("Unknown v_init type: %s" % vinit_type)
# orthogonalize V
if M is not None:
V, R = tallqr(V, MV=M.mm(V))
else:
V, R = tallqr(V)
return V
def _set_initial_v(vinit_type: str,
dtype: torch.dtype, device: torch.device,
batch_dims: Sequence,
na: int,
nguess: int,
M: Optional[LinearOperator] = None) -> torch.Tensor:
torch.manual_seed(12421)
if vinit_type == "eye":
nbatch = functools.reduce(lambda x, y: x * y, batch_dims, 1)
V = torch.eye(na, nguess, dtype=dtype, device=device).unsqueeze(
0).repeat(nbatch, 1, 1).reshape(*batch_dims, na, nguess)
elif vinit_type == "randn":
V = torch.randn((*batch_dims, na, nguess), dtype=dtype, device=device)
elif vinit_type == "random" or vinit_type == "rand":
V = torch.rand((*batch_dims, na, nguess), dtype=dtype, device=device)
else:
raise ValueError("Unknown v_init type: %s" % vinit_type)
# orthogonalize V
if isinstance(M, LinearOperator):
V, R = tallqr(V, MV=M.mm(V))
else:
V, R = tallqr(V)
return V
def davidson(A, params, neig, mode, M=None, mparams=[], **options):
"""
Iterative methods to obtain the `neig` lowest eigenvalues and eigenvectors.
This function is written so that the backpropagation can be done.
It solves the eigendecomposition AV = VME where V are the matrix of eigenvectors,
and E are the diagonal matrix consists of the eigenvalues.
Arguments
---------
* A: LinearOperator instance (*BA, na, na)
The linear operator object on which the eigenpairs are constructed.
* params: list of differentiable torch.tensor of any shapes
List of differentiable torch.tensor to be put to A.
* neig: int
The number of eigenpairs to be retrieved.
* mode: str
Take the `neig` "lowest" or "uppest" of the eigenpairs.
* M: LinearOperator instance (*BM, na, na) or None
The transformation on the right hand side. If None, then M=I.
* mparams: list of differentiable torch.tensor of any shapes
List of differentiable torch.tensor to be put to M.
* **options:
Iterative algorithm options.
Returns
-------
* eigvals: torch.tensor (*BAM, neig)
* eigvecs: torch.tensor (*BAM, na, neig)
The `neig` lowest eigenpairs
"""
# TODO: optimize for large linear operator and strict min_eps
# Ideas:
# (1) use better strategy to get the estimate on eigenvalues
# (2) use restart strategy
config = set_default_option(
{
"max_niter": 1000,
"nguess": neig, # number of initial guess
"min_eps": 1e-6,
"verbose": False,
"eps_cond": 1e-6,
"v_init": "randn",
"max_addition": neig,
},
options)
# get some of the options
nguess = config["nguess"]
max_niter = config["max_niter"]
min_eps = config["min_eps"]
verbose = config["verbose"]
eps_cond = config["eps_cond"]
max_addition = config["max_addition"]
# get the shape of the transformation
na = A.shape[-1]
if M is None:
bcast_dims = A.shape[:-2]
else:
bcast_dims = get_bcasted_dims(A.shape[:-2], M.shape[:-2])
dtype = A.dtype
device = A.device
# TODO: A to use params
prev_eigvals = None
prev_eigvalT = None
stop_reason = "max_niter"
shift_is_eigvalT = False
idx = torch.arange(neig).unsqueeze(-1) # (neig, 1)
with A.uselinopparams(*params), M.uselinopparams(
*mparams) if M is not None else dummy_context_manager():
# set up the initial guess
V = _set_initial_v(config["v_init"].lower(),
dtype,
device,
bcast_dims,
na,
nguess,
M=M,
mparams=mparams) # (*BAM, na, nguess)
# V = V.reshape(*bcast_dims, na, nguess) # (*BAM, na, nguess)
# estimating the lowest eigenvalues
eig_est, rms_eig = _estimate_eigvals(A,
neig,
mode,
bcast_dims=bcast_dims,
na=na,
ntest=20,
dtype=V.dtype,
device=V.device)
best_resid = float("inf")
AV = A.mm(V)
for i in range(max_niter):
VT = V.transpose(-2, -1) # (*BAM,nguess,na)
# Can be optimized by saving AV from the previous iteration and only
# operate AV for the new V. This works because the old V has already
# been orthogonalized, so it will stay the same
# AV = A.mm(V) # (*BAM,na,nguess)
T = torch.matmul(VT, AV) # (*BAM,nguess,nguess)
# eigvals are sorted from the lowest
# eval: (*BAM, nguess), evec: (*BAM, nguess, nguess)
eigvalT, eigvecT = torch.symeig(T, eigenvectors=True)
eigvalT, eigvecT = _take_eigpairs(
eigvalT, eigvecT, neig,
mode) # (*BAM, neig) and (*BAM, nguess, neig)
# calculate the eigenvectors of A
eigvecA = torch.matmul(V, eigvecT) # (*BAM, na, neig)
# calculate the residual
AVs = torch.matmul(AV, eigvecT) # (*BAM, na, neig)
LVs = eigvalT.unsqueeze(-2) * eigvecA # (*BAM, na, neig)
if M is not None:
LVs = M.mm(LVs)
resid = AVs - LVs # (*BAM, na, neig)
# print information and check convergence
max_resid = resid.abs().max()
if prev_eigvalT is not None:
deigval = eigvalT - prev_eigvalT
max_deigval = deigval.abs().max()
if verbose:
print("Iter %3d (guess size: %d): resid: %.3e, devals: %.3e" % \
(i+1, nguess, max_resid, max_deigval))
if max_resid < best_resid:
best_resid = max_resid
best_eigvals = eigvalT
best_eigvecs = eigvecA
if max_resid < min_eps:
break
if AV.shape[-1] == AV.shape[-2]:
break
prev_eigvalT = eigvalT
# apply the preconditioner
# initial guess of the eigenvalues are actually help really much
if not shift_is_eigvalT:
z = eig_est # (*BAM,neig)
else:
z = eigvalT # (*BAM,neig)
# if A.is_precond_set():
# t = A.precond(-resid, *params, biases=z, M=M, mparams=mparams) # (nbatch, na, neig)
# else:
t = -resid # (*BAM, na, neig)
# set the estimate of the eigenvalues
if not shift_is_eigvalT:
eigvalT_pred = eigvalT + torch.einsum(
'...ae,...ae->...e', eigvecA, A.mm(t)) # (*BAM, neig)
diff_eigvalT = (eigvalT - eigvalT_pred) # (*BAM, neig)
if diff_eigvalT.abs().max() < rms_eig * 1e-2:
shift_is_eigvalT = True
else:
change_idx = eig_est > eigvalT
next_value = eigvalT - 2 * diff_eigvalT
eig_est[change_idx] = next_value[change_idx]
# orthogonalize t with the rest of the V
t = to_fortran_order(t)
Vnew = torch.cat((V, t), dim=-1)
if Vnew.shape[-1] > Vnew.shape[-2]:
Vnew = Vnew[..., :Vnew.shape[-2]]
nadd = Vnew.shape[-1] - V.shape[-1]
nguess = nguess + nadd
if M is not None:
MV_ = M.mm(Vnew)
V, R = tallqr(Vnew, MV=MV_)
else:
V, R = tallqr(Vnew)
AVnew = A.mm(V[..., -nadd:]) # (*BAM,na,nadd)
AVnew = to_fortran_order(AVnew)
AV = torch.cat((AV, AVnew), dim=-1)
eigvals = best_eigvals # (*BAM, neig)
eigvecs = best_eigvecs # (*BAM, na, neig)
return eigvals, eigvecs
def davidson(A: LinearOperator, neig: int,
mode: str,
M: Optional[LinearOperator] = None,
max_niter: int = 1000,
nguess: Optional[int] = None,
v_init: str = "randn",
max_addition: Optional[int] = None,
min_eps: float = 1e-6,
verbose: bool = False,
**unused) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Using Davidson method for large sparse matrix eigendecomposition [1]_.
Arguments
---------
max_niter: int
Maximum number of iterations
v_init: str
Mode of the initial guess (``"randn"``, ``"rand"``, ``"eye"``)
max_addition: int or None
Maximum number of new guesses to be added to the collected vectors.
If None, set to ``neig``.
min_eps: float
Minimum residual error to be stopped
verbose: bool
Option to be verbose
References
----------
.. [1] P. Arbenz, "Lecture Notes on Solving Large Scale Eigenvalue Problems"
http://people.inf.ethz.ch/arbenz/ewp/Lnotes/chapter12.pdf
"""
# TODO: optimize for large linear operator and strict min_eps
# Ideas:
# (1) use better strategy to get the estimate on eigenvalues
# (2) use restart strategy
if nguess is None:
nguess = neig
if max_addition is None:
max_addition = neig
# get the shape of the transformation
na = A.shape[-1]
if M is None:
bcast_dims = A.shape[:-2]
else:
bcast_dims = get_bcasted_dims(A.shape[:-2], M.shape[:-2])
dtype = A.dtype
device = A.device
prev_eigvals = None
prev_eigvalT = None
stop_reason = "max_niter"
shift_is_eigvalT = False
idx = torch.arange(neig).unsqueeze(-1) # (neig, 1)
# set up the initial guess
V = _set_initial_v(v_init.lower(), dtype, device,
bcast_dims, na, nguess,
M=M) # (*BAM, na, nguess)
best_resid: Union[float, torch.Tensor] = float("inf")
AV = A.mm(V)
for i in range(max_niter):
VT = V.transpose(-2, -1) # (*BAM,nguess,na)
# Can be optimized by saving AV from the previous iteration and only
# operate AV for the new V. This works because the old V has already
# been orthogonalized, so it will stay the same
# AV = A.mm(V) # (*BAM,na,nguess)
T = torch.matmul(VT, AV) # (*BAM,nguess,nguess)
# eigvals are sorted from the lowest
# eval: (*BAM, nguess), evec: (*BAM, nguess, nguess)
eigvalT, eigvecT = torch.symeig(T, eigenvectors=True)
eigvalT, eigvecT = _take_eigpairs(eigvalT, eigvecT, neig, mode) # (*BAM, neig) and (*BAM, nguess, neig)
# calculate the eigenvectors of A
eigvecA = torch.matmul(V, eigvecT) # (*BAM, na, neig)
# calculate the residual
AVs = torch.matmul(AV, eigvecT) # (*BAM, na, neig)
LVs = eigvalT.unsqueeze(-2) * eigvecA # (*BAM, na, neig)
if M is not None:
LVs = M.mm(LVs)
resid = AVs - LVs # (*BAM, na, neig)
# print information and check convergence
max_resid = resid.abs().max()
if prev_eigvalT is not None:
deigval = eigvalT - prev_eigvalT
max_deigval = deigval.abs().max()
if verbose:
print("Iter %3d (guess size: %d): resid: %.3e, devals: %.3e" %
(i + 1, nguess, max_resid, max_deigval)) # type:ignore
if max_resid < best_resid:
best_resid = max_resid
best_eigvals = eigvalT
best_eigvecs = eigvecA
if max_resid < min_eps:
break
if AV.shape[-1] == AV.shape[-2]:
break
prev_eigvalT = eigvalT
# apply the preconditioner
t = -resid # (*BAM, na, neig)
# orthogonalize t with the rest of the V
t = to_fortran_order(t)
Vnew = torch.cat((V, t), dim=-1)
if Vnew.shape[-1] > Vnew.shape[-2]:
Vnew = Vnew[..., :Vnew.shape[-2]]
nadd = Vnew.shape[-1] - V.shape[-1]
nguess = nguess + nadd
if M is not None:
MV_ = M.mm(Vnew)
V, R = tallqr(Vnew, MV=MV_)
else:
V, R = tallqr(Vnew)
AVnew = A.mm(V[..., -nadd:]) # (*BAM,na,nadd)
AVnew = to_fortran_order(AVnew)
AV = torch.cat((AV, AVnew), dim=-1)
eigvals = best_eigvals # (*BAM, neig)
eigvecs = best_eigvecs # (*BAM, na, neig)
return eigvals, eigvecs
def davidson(A,
params,
neig,
mode,
M=None,
mparams=[],
max_niter=1000,
nguess=None,
v_init="randn",
max_addition=None,
min_eps=1e-6,
verbose=False,
**unused):
"""
Using Davidson method for large sparse matrix eigendecomposition [1]_.
Arguments
---------
max_niter: int
Maximum number of iterations
nguess: int or None
The number of initial guess of the eigenvectors
If None, set to ``neig``.
v_init: str
Mode of the initial guess (``"randn"``, ``"rand"``, ``"eye"``)
max_addition: int or None
Maximum number of new guesses to be added to the collected vectors.
If None, set to ``neig``.
min_eps: float
Minimum residual error to be stopped
verbose: bool
Option to be verbose
References
----------
.. [1] P. Arbenz, "Lecture Notes on Solving Large Scale Eigenvalue Problems"
http://people.inf.ethz.ch/arbenz/ewp/Lnotes/chapter12.pdf
"""
# TODO: optimize for large linear operator and strict min_eps
# Ideas:
# (1) use better strategy to get the estimate on eigenvalues
# (2) use restart strategy
if nguess is None:
nguess = neig
if max_addition is None:
max_addition = neig
# get the shape of the transformation
na = A.shape[-1]
if M is None:
bcast_dims = A.shape[:-2]
else:
bcast_dims = get_bcasted_dims(A.shape[:-2], M.shape[:-2])
dtype = A.dtype
device = A.device
# TODO: A to use params
prev_eigvals = None
prev_eigvalT = None
stop_reason = "max_niter"
shift_is_eigvalT = False
idx = torch.arange(neig).unsqueeze(-1) # (neig, 1)
with A.uselinopparams(*params), M.uselinopparams(
*mparams) if M is not None else dummy_context_manager():
# set up the initial guess
V = _set_initial_v(v_init.lower(),
dtype,
device,
bcast_dims,
na,
nguess,
M=M,
mparams=mparams) # (*BAM, na, nguess)
# V = V.reshape(*bcast_dims, na, nguess) # (*BAM, na, nguess)
# estimating the lowest eigenvalues
eig_est, rms_eig = _estimate_eigvals(A,
neig,
mode,
bcast_dims=bcast_dims,
na=na,
ntest=20,
dtype=V.dtype,
device=V.device)
best_resid = float("inf")
AV = A.mm(V)
for i in range(max_niter):
VT = V.transpose(-2, -1) # (*BAM,nguess,na)
# Can be optimized by saving AV from the previous iteration and only
# operate AV for the new V. This works because the old V has already
# been orthogonalized, so it will stay the same
# AV = A.mm(V) # (*BAM,na,nguess)
T = torch.matmul(VT, AV) # (*BAM,nguess,nguess)
# eigvals are sorted from the lowest
# eval: (*BAM, nguess), evec: (*BAM, nguess, nguess)
eigvalT, eigvecT = torch.symeig(T, eigenvectors=True)
eigvalT, eigvecT = _take_eigpairs(
eigvalT, eigvecT, neig,
mode) # (*BAM, neig) and (*BAM, nguess, neig)
# calculate the eigenvectors of A
eigvecA = torch.matmul(V, eigvecT) # (*BAM, na, neig)
# calculate the residual
AVs = torch.matmul(AV, eigvecT) # (*BAM, na, neig)
LVs = eigvalT.unsqueeze(-2) * eigvecA # (*BAM, na, neig)
if M is not None:
LVs = M.mm(LVs)
resid = AVs - LVs # (*BAM, na, neig)
# print information and check convergence
max_resid = resid.abs().max()
if prev_eigvalT is not None:
deigval = eigvalT - prev_eigvalT
max_deigval = deigval.abs().max()
if verbose:
print("Iter %3d (guess size: %d): resid: %.3e, devals: %.3e" % \
(i+1, nguess, max_resid, max_deigval))
if max_resid < best_resid:
best_resid = max_resid
best_eigvals = eigvalT
best_eigvecs = eigvecA
if max_resid < min_eps:
break
if AV.shape[-1] == AV.shape[-2]:
break
prev_eigvalT = eigvalT
# apply the preconditioner
# initial guess of the eigenvalues are actually help really much
if not shift_is_eigvalT:
z = eig_est # (*BAM,neig)
else:
z = eigvalT # (*BAM,neig)
# if A.is_precond_set():
# t = A.precond(-resid, *params, biases=z, M=M, mparams=mparams) # (nbatch, na, neig)
# else:
t = -resid # (*BAM, na, neig)
# set the estimate of the eigenvalues
if not shift_is_eigvalT:
eigvalT_pred = eigvalT + torch.einsum(
'...ae,...ae->...e', eigvecA, A.mm(t)) # (*BAM, neig)
diff_eigvalT = (eigvalT - eigvalT_pred) # (*BAM, neig)
if diff_eigvalT.abs().max() < rms_eig * 1e-2:
shift_is_eigvalT = True
else:
change_idx = eig_est > eigvalT
next_value = eigvalT - 2 * diff_eigvalT
eig_est[change_idx] = next_value[change_idx]
# orthogonalize t with the rest of the V
t = to_fortran_order(t)
Vnew = torch.cat((V, t), dim=-1)
if Vnew.shape[-1] > Vnew.shape[-2]:
Vnew = Vnew[..., :Vnew.shape[-2]]
nadd = Vnew.shape[-1] - V.shape[-1]
nguess = nguess + nadd
if M is not None:
MV_ = M.mm(Vnew)
V, R = tallqr(Vnew, MV=MV_)
else:
V, R = tallqr(Vnew)
AVnew = A.mm(V[..., -nadd:]) # (*BAM,na,nadd)
AVnew = to_fortran_order(AVnew)
AV = torch.cat((AV, AVnew), dim=-1)
eigvals = best_eigvals # (*BAM, neig)
eigvecs = best_eigvecs # (*BAM, na, neig)
return eigvals, eigvecs