Source code for pysisyphus.linalg

from math import cos, sin, sqrt
from typing import Callable, Literal, Optional, Tuple

import numpy as np
from numpy.typing import NDArray
from scipy.spatial.transform import Rotation
from scipy.linalg.lapack import dpstrf


[docs] def gram_schmidt(vecs, thresh=1e-8): def proj(v1, v2): return v1.dot(v2) / v1.dot(v1) ortho = [ vecs[0], ] for v1 in vecs[1:]: tmp = v1.copy() for v2 in ortho: tmp -= proj(v2, v1) * v2 norm = np.linalg.norm(tmp) # Don't append linear dependent vectors, as their norm will be # near zero. Renormalizing them to unity would lead to numerical # garbage and to erronous results later on, when we orthgonalize # against this 'arbitrary' vector. if norm <= thresh: continue ortho.append(tmp / norm) return np.array(ortho)
[docs] def orthogonalize_against(mat, vecs, max_cycles=5, thresh=1e-10): """Orthogonalize rows of 'mat' against rows in 'vecs' Returns a (modified) copy of mat. """ omat = mat.copy() for _ in range(max_cycles): max_overlap = 0.0 for row in omat: for ovec in vecs: row -= ovec.dot(row) * ovec # Remaining overlap overlap = np.sum(np.dot(vecs, row)) max_overlap = max(overlap, max_overlap) if max_overlap < thresh: break else: raise Exception(f"Orthogonalization did not succeed in {max_cycles} cycles!") return omat
[docs] def perp_comp(vec, along): """Return the perpendicular component of vec along along.""" return vec - vec.dot(along) * along
[docs] def make_unit_vec(vec1, vec2): """Return unit vector pointing from vec2 to vec1.""" diff = vec1 - vec2 return diff / np.linalg.norm(diff)
[docs] def svd_inv(array, thresh, hermitian=False): U, S, Vt = np.linalg.svd(array, hermitian=hermitian) keep = S > thresh S_inv = np.zeros_like(S) S_inv[keep] = 1 / S[keep] return Vt.T.dot(np.diag(S_inv)).dot(U.T)
[docs] def get_rot_mat(abc=None): # Euler angles if abc is None: abc = np.random.rand(3) * np.pi * 2 a, b, c = abc R = np.array( ( ( cos(a) * cos(b) * cos(c) - sin(a) * sin(c), -cos(a) * cos(b) * sin(c) - sin(a) * cos(c), cos(a) * sin(b), ), ( sin(a) * cos(b) * cos(c) + cos(a) * sin(c), -sin(a) * cos(b) * sin(c) + cos(a) * cos(c), sin(a) * sin(b), ), (-sin(b) * cos(c), sin(b) * sin(c), cos(b)), ) ) return R
[docs] def get_rot_mat_for_coords(coords3d_1, coords3d_2): coords3d_1 = coords3d_1.copy() coords3d_2 = coords3d_2.copy() centroid_1 = coords3d_1.mean(axis=0) centroid_2 = coords3d_2.mean(axis=0) coords3d_1 -= centroid_1[None, :] coords3d_2 -= centroid_2[None, :] tmp = coords3d_2.T.dot(coords3d_1) U, W, Vt = np.linalg.svd(tmp) rot_mat = U.dot(Vt) if np.linalg.det(rot_mat) < 0: U[:, -1] *= -1 rot_mat = U.dot(Vt) return coords3d_1, coords3d_2, rot_mat
[docs] def eigvec_grad(w, v, ind, mat_grad): """Gradient of 'ind'-th eigenvector. dv_i / dx_i = (w_i*I - mat)⁻¹ dmat/dx_i v_i """ eigval = w[ind] eigvec = v[:, ind] w_diff = eigval - w w_inv = np.divide( 1.0, w_diff, out=np.zeros_like(w_diff).astype(float), where=w_diff != 0.0 ) assert np.isfinite(w_inv).all() pinv = v.dot(np.diag(w_inv)).dot(v.T) pinv.dot(mat_grad) wh = pinv.dot(mat_grad).dot(eigvec) return wh
[docs] def cross3(a, b): """10x as fast as np.cross for two 1d arrays of size 3.""" return np.array( ( a[1] * b[2] - a[2] * b[1], a[2] * b[0] - a[0] * b[2], a[0] * b[1] - a[1] * b[0], ) )
[docs] def norm3(a): """5x as fas as np.linalg.norm for a 1d array of size 3.""" return sqrt(a[0] * a[0] + a[1] * a[1] + a[2] * a[2])
[docs] def finite_difference_hessian( coords: NDArray[float], grad_func: Callable[[NDArray[float]], NDArray[float]], step_size: float = 1e-2, acc: Literal[2, 4] = 2, callback: Optional[Callable] = None, ) -> NDArray[float]: """Numerical Hessian from central finite gradient differences. See central differences in https://en.wikipedia.org/wiki/Finite_difference_coefficient for the different accuracies. """ if callback is None: def callback(*args): pass accuracies = { 2: ((-0.5, -1), (0.5, 1)), # 2 calculations 4: ((1 / 12, -2), (-2 / 3, -1), (2 / 3, 1), (-1 / 12, 2)), # 4 calculations } accs_avail = list(accuracies.keys()) assert acc in accs_avail size = coords.size fd_hessian = np.zeros((size, size)) zero_step = np.zeros(size) coeffs = accuracies[acc] for i, _ in enumerate(coords): step = zero_step.copy() step[i] = step_size def get_grad(factor, displ, j): displ_coords = coords + step * displ callback(i, j) grad = grad_func(displ_coords) return factor * grad grads = [get_grad(factor, displ, j) for j, (factor, displ) in enumerate(coeffs)] fd = np.sum(grads, axis=0) / step_size fd_hessian[i] = fd # Symmetrize fd_hessian = (fd_hessian + fd_hessian.T) / 2 return fd_hessian
[docs] def rot_quaternion(coords3d, ref_coords3d): # Translate to origin by removing centroid c3d = coords3d - coords3d.mean(axis=0) ref_c3d = ref_coords3d - ref_coords3d.mean(axis=0) # Setup correlation matrix R = c3d.T.dot(ref_c3d) # Setup F matrix, Eq. (6) in [1] F = np.zeros((4, 4)) R11, R12, R13, R21, R22, R23, R31, R32, R33 = R.flatten() # Fill only upper triangular part. F[0, 0] = R11 + R22 + R33 F[0, 1] = R23 - R32 F[0, 2] = R31 - R13 F[0, 3] = R12 - R21 # F[1, 1] = R11 - R22 - R33 F[1, 2] = R12 + R21 F[1, 3] = R13 + R31 # F[2, 2] = -R11 + R22 - R33 F[2, 3] = R23 + R32 # F[3, 3] = -R11 - R22 + R33 # Eigenvalues, eigenvectors of upper triangular part. w, v_ = np.linalg.eigh(F, UPLO="U") # Quaternion corresponds to biggest (last) eigenvalue. # np.linalg.eigh already returns sorted eigenvalues. return w, v_, c3d, ref_c3d
[docs] def quaternion_to_rot_mat(q): q0, q1, q2, q3 = q q_ = q0**2 - (q1**2 + q2**2 + q3**2) R = np.zeros((3, 3)) R[0, 0] = q_ + 2 * q1**2 R[0, 1] = 2 * (q1 * q2 - q0 * q3) R[0, 2] = 2 * (q1 * q3 - q0 * q2) R[1, 0] = 2 * (q1 * q2 + q0 * q3) R[1, 1] = q_ + 2 * q2**2 R[1, 2] = 2 * (q2 * q3 - q0 * q1) R[2, 0] = 2 * (q1 * q3 - q0 * q2) R[2, 1] = 2 * (q2 * q3 + q0 * q1) R[2, 2] = q_ + 2 * q3**2 return R
[docs] def rmsd_grad( coords3d: NDArray[float], ref_coords3d: NDArray[float], offset: float = 1e-9 ) -> Tuple[float, NDArray[float]]: """RMSD and gradient between two sets of coordinates from quaternions. The gradient is given w.r.t. the coordinates of 'coords3d'. Python adaption of ls_rmsd.f90 from the xtb repository of the Grimme group, which in turn implements [1] https://doi.org/10.1002/jcc.20110 . Parameters ---------- coords3d Coordinate array of shape (N, 3) with N denoting the number of atoms. ref_coords3d Reference coordinates. offset Small floating-point number that is added to the RMSD, to avoid division by zero. Returns ------- rmsd RMSD value. rmsd_grad Gradient of the RMSD value w.r.t. to 'coords3d'. """ assert coords3d.shape == ref_coords3d.shape w, v_, c3d, ref_c3d = rot_quaternion(coords3d, ref_coords3d) quat = v_[:, -1] eigval = w[-1] atom_num = coords3d.shape[0] x_norm = np.linalg.norm(c3d) ** 2 y_norm = np.linalg.norm(ref_c3d) ** 2 rmsd = np.sqrt(max(0.0, ((x_norm + y_norm) - 2.0 * eigval)) / (atom_num)) + offset # scipy expects the quaternion # a + b*i + c*j + d*k # as (i, j, k, a), that is the scalar component must come last. # Currently we have (a, i, j, k), so we have to rearrange before passing the # quaternion to scipy. rot = Rotation.from_quat((*quat[1:], quat[0])) U = rot.as_matrix() grad = (c3d - ref_c3d @ U) / (rmsd * atom_num) return rmsd, grad
[docs] def pivoted_cholesky(A: NDArray, tol: float = -1.0): """Cholesky factorization a real symmetric positive semidefinite matrix. Cholesky factorization is carried out with full pivoting. Adapated from PySCF. P.T * A * P = L * L.T Parameters ---------- A Matrix to be factorized. tol User defined tolerance, as outlined in the LAPACK docs. Returns ------- L Lower or upper triangular matrix. piv Pivot vectors, starting at 0. rank Rank of the factoirzed matrix. """ N = A.shape[0] L, piv, rank, info = dpstrf(A, tol=tol, lower=True) piv -= 1 # LAPACK returns 1-based indices assert info >= 0 L[np.triu_indices(N, k=1)] = 0 L[:, rank:] = 0 return L, piv, rank
[docs] def pivoted_cholesky2(A: NDArray[float], tol: Optional[float] = None, m_max: int = 0): """https://doi.org/10.1016/j.apnum.2011.10.001""" R = np.zeros_like(A) m = 0 n = len(A) # Decompose to max rank. if tol is None: tol = 0.0 assert tol >= 0.0, f"{tol=} must be >= 0.0!" if m_max == 0: m_max = n assert 0 <= m_max <= n, f"{m_max=} must fullfil (0 <= m_max <= {n=})!" d = np.diag(A).copy() # Diagonal error = np.linalg.norm(A, ord="nuc") piv = np.arange(n, dtype=int) while True: # Stop when error is sufficiently small or we are at max rank. # while (error > tol) or (m_max and m < m_max): if (error <= tol) or (m_max and (m == m_max)): break i = m + d[piv[m:]].argmax() piv[m], piv[i] = piv[i], piv[m] R[m, piv[m]] = np.sqrt(d[piv[m]]) for i in range(m + 1, n): j = np.arange(m) sum_ = (R[j, piv[m]] * R[j, piv[i]]).sum() R[m, piv[i]] = (A[piv[m], piv[i]] - sum_) / R[m, piv[m]] d[piv[i]] -= R[m, piv[i]] ** 2 error = d[m + 1 :].sum() m += 1 # R.T @ R == A return R, piv, m
[docs] def matrix_power(mat, p, thresh=1e-12, strict=True): w, v = np.linalg.eigh(mat) assert (not strict) or ( w > 0.0 ).all(), "matrix_power must be called with a (semi)-positive-definite matrix!" mask = np.abs(w) > thresh w_pow = w[mask] ** p v_mask = v[:, mask] return v_mask @ np.diag(w_pow) @ v_mask.T
[docs] def sym_mat_from_triu(arr, data): nrows, ncols = arr.shape assert nrows == ncols triu = np.triu_indices(nrows) tril1 = np.tril_indices(nrows, k=-1) arr[triu] = data arr[tril1] = arr.T[tril1]
[docs] def sym_mat_from_tril(arr, data): nrows, ncols = arr.shape assert nrows == ncols tril = np.tril_indices(nrows) arr[tril] = data triu1 = np.triu_indices(nrows, k=1) arr[triu1] = arr.T[triu1]
[docs] def multi_component_sym_mat(arr, dim): target_shape = (dim, dim) shape = arr.shape sym = np.zeros((*target_shape, *shape[1:])) triu = np.triu_indices(dim) triu1 = np.triu_indices(dim, k=1) tril1 = np.tril_indices(dim, k=-1) sym[triu] = arr sym[tril1] = sym[triu1] return sym.reshape(*target_shape, *shape[1:])