Initial commit.

grext/buffer.py

+import numpy as np
+from typing import List, Tuple
+
+class CircularBuffer:
+
+    """Circular buffer to store the last `n` matrices."""
+
+    def __init__(self, n: int, shape: Tuple[int, int]):
+        self.n = n
+        self.shape = shape
+        self.buffer = [np.zeros(shape, dtype=np.float64) for _ in range(n)]
+        self.index = 0
+
+    def push(self, data):
+        """Add a new matrix to the buffer."""
+        self.buffer[self.index] = data.copy()
+        self.index = (self.index + 1) % self.n
+
+    def get(self, m) -> List[np.ndarray]:
+        """Get the last `m` matrices."""
+        if m > self.n:
+            raise ValueError("`m` should be less than or equal to the buffer `n`")
+
+        start_idx = (self.index - m) % self.n
+        return [self.buffer[i] for i in range(start_idx, start_idx + m)]

grext/fitting.py

+
+def linear():
+    pass
+
+def quasi_time_reversible():
+    pass

grext/grassmann.py

+import numpy as np
+
+def log_plain(c: np.ndarray, c0: np.ndarray) -> np.ndarray:
+    c0c_inv = np.linalg.inv(c0.T @ c)
+    L = c @ c0c_inv - c0
+    q, s, vt = np.linalg.svd(L, full_matrices=False)
+    arctan_s = np.diag(np.arctan(s))
+    return q @ arctan_s @ vt
+
+def log(c: np.ndarray, c0: np.ndarray) -> np.ndarray:
+    psi, s, rt = np.linalg.svd(c.T @ c0, full_matrices=False)
+    cstar = c @ psi @ rt
+    L = (np.identity(c.shape[0]) - c0 @ c0.T) @ cstar
+    u, s, vt = np.linalg.svd(L, full_matrices=False)
+    arcsin_s = np.diag(np.arcsin(s))
+    return u @ arcsin_s @ vt
+
+def exp(gamma: np.ndarray, c0: np.ndarray) -> np.ndarray:
+    q, s, vt = np.linalg.svd(gamma, full_matrices=False)
+    sin_s = np.diag(np.sin(s))
+    cos_s = np.diag(np.cos(s))
+    return c0 @ vt.T @ cos_s @ vt + q @ sin_s @ vt
+
+def psi(d: np.ndarray, n: np.ndarray) -> np.ndarray:
+    a = d[0:n, 0:n]
+    b = d[n:, 0:n]
+    ainv = np.linalg.inv(a)
+    return b @ ainv
+
+def phi(b: np.ndarray) -> np.ndarray:
+    nb_n, n = b.shape
+    q = np.linalg.inv(np.identity(n) + b.T @ b)
+    l = np.zeros((nb_n + n, n))
+    l[0:n,:] = np.identity(n)
+    l[n:,:] = b
+    return l @ q @ l.T

grext/main.py

+import numpy as np
+from typing import Optional
+
+from . import grassmann
+from .buffer import CircularBuffer
+
+class GrassmannExt:
+
+    """Module for performing Grassmann extrapolations."""
+
+    def __init__(self, nelectrons: int, nbasis: int, natoms: int,
+            nsteps: int = 10):
+
+        self.nelectrons = nelectrons
+        self.nbasis = nbasis
+        self.natoms = natoms
+        self.nsteps = nsteps
+
+        self.gammas = CircularBuffer(self.nsteps, (self.nelectrons, self.nbasis))
+        self.overlaps = CircularBuffer(self.nsteps, (self.nbasis, self.nbasis))
+        self.coords = CircularBuffer(self.nsteps, (self.natoms, 3))
+
+        self.is_tangent_set = False
+
+    def load_data(self, coords: np.ndarray, coeff: np.ndarray,
+            overlap: np.ndarray):
+        """Load a new data point in the extrapolator."""
+        coeff = self._crop_coeff(coeff)
+        coeff = self._normalize(coeff, overlap)
+
+        self.gammas.push(self._grassmann_log(coeff))
+        self.coords.push(coords)
+        self.overlaps.push(overlap)
+
+    def guess(self, coords: np.ndarray, overlap: Optional[np.ndarray]):
+        """Get a new electronic density to be used as a guess."""
+        pass
+
+    def _crop_coeff(self, coeff) -> np.ndarray:
+        """Crop the coefficient matrix to remove the virtual orbitals."""
+        return coeff[:, :self.nelectrons]
+
+    def _normalize(self, coeff: np.ndarray, overlap: np.ndarray) -> np.ndarray:
+        """Normalize the coefficients such that C.T * C = 1, D * D = D."""
+        q, s, vt = np.linalg.svd(overlap, full_matrices=False)
+        sqrt_overlap = q @ np.diag(np.sqrt(s)) @ vt
+        return sqrt_overlap.dot(coeff)
+
+    def _set_tangent(self, c: np.ndarray):
+        """Set the tangent point."""
+        self.is_tangent_set = True
+        self.tangent = c
+
+    def _grassmann_log(self, coeff: np.ndarray):
+        """Map from the manifold to the tangent plane."""
+        return grassmann.log(coeff, self.tangent)
+
+    def _grassmann_exp(self, gamma: np.ndarray):
+        """Map from the tangent plane to the manifold."""
+        return grassmann.exp(gamma, self.tangent)