sparse-ir 1.1.6__py3-none-any.whl → 2.0.0a2__py3-none-any.whl

sparse_ir/sampling.py

@@ -1,157 +1,168 @@
-# Copyright (C) 2020-2022 Markus Wallerberger, Hiroshi Shinaoka, and others
-# SPDX-License-Identifier: MIT
-import numpy as np
-from warnings import warn
-
-from . import _util
-
-
-class AbstractSampling:
-    """Base class for sparse sampling.
+"""
+High-level Python classes for sparse sampling
+"""
 
-    Encodes the "basis transformation" of a propagator from the truncated IR
-    basis coefficients ``G_ir[l]`` to time/frequency sampled on sparse points
-    ``G(x[i])`` together with its inverse, a least squares fit::
+import numpy as np
+from ctypes import POINTER, c_double, c_int, byref
+from pylibsparseir.core import c_double_complex, tau_sampling_new, tau_sampling_new_with_matrix, matsubara_sampling_new, matsubara_sampling_new_with_matrix, _lib
+from pylibsparseir.constants import COMPUTATION_SUCCESS, SPIR_ORDER_ROW_MAJOR
+from . import augment
 
-             ________________                   ___________________
-            |                |    evaluate     |                   |
-            |     Basis      |---------------->|     Value on      |
-            |  coefficients  |<----------------|  sampling points  |
-            |________________|      fit        |___________________|
+class TauSampling:
+    """Sparse sampling in imaginary time.
 
+    Allows the transformation between the IR basis and a set of sampling points
+    in (scaled/unscaled) imaginary time.
     """
-    def evaluate(self, al, axis=None, *, points=None):
-        """Evaluate the basis coefficients at sampling points.
-
-        Arguments:
-            al (array):
-                Array where the `l`-th item along `axis` corresponds to the
-                `l`-th basis coefficient
-            axis (integer):
-                Axis or dimension of `al` along which to evaluate the function.
-                Defaults to the last, i.e., rightmost axis.
-            points (vector):
-                Points on which the results should be evaluated.  Defaults
-                to the sampling points for which the sampling objects was
-                created.
-
-                .. versionadded:: 1.1
-
-        Return:
-            Array where the `n`-th item along `axis` corresponds to the
-            value on the `n`-th sampling point (or value on `point[n]`, if
-            given.)
-
-        Note:
-            If `points` is given, a new sampling is created at each invocation,
-            which can result in a performance hit.  Consider caching sampling
-            objects or simply using the `.u()` and `.uhat()` methods of the
-            underlying basis.
+
+    def __init__(self, basis, sampling_points=None, use_positive_taus=False):
         """
-        if points is not None:
-            return self._for_sampling_points(points).evaluate(al, axis)
-
-        return self.matrix.matmul(al, axis)
-
-    def fit(self, ax, axis=None, *, points=None):
-        """Fit the basis coefficients from the sampling points.
-
-        Arguments:
-            ax (array):
-                Array where the `n`-th item along `axis` corresponds to the
-                value on the `n`-th sampling point (or value on `point[n]`, if
-                given.)
-            axis (integer):
-                Axis or dimension of `ax` along which to fit the function.
-                Defaults to the last, i.e., rightmost axis.
-            points (vector):
-                Points on which the `ax` is given.  Defaults to the sampling
-                points for which the sampling objects was created.
-
-                .. versionadded:: 1.1
-
-        Return:
-            Array where the `l`-th item along `axis` corresponds to the
-            `l`-th basis coefficient
-
-        Note:
-            If `points` is given, a new sampling is created at each invocation,
-            which can result in a performance hit.  Consider caching sampling
-            objects.
+        Initialize tau sampling.
+
+        Parameters:
+        -----------
+        basis : FiniteTempBasis
+            Finite temperature basis
+        sampling_points : array_like, optional
+            Tau sampling points. If None, use default.
+        use_positive_taus : bool, optional
+            If `use_positive_taus=False`, the sampling points are within
+            the range [-β/2, β/2] and the distribution is symmetric.
+            If `use_positive_taus=True`, the sampling points are
+            folded to the positive tau domain [0, β), which is
+            the default behavior of sparseir 1.x.x.
         """
-        if points is not None:
-            return self._for_sampling_points(points).fit(ax, axis)
-
-        matrix = self.matrix
-        if self.basis.is_well_conditioned and not (matrix.cond <= 1e8):
-            warn(f"Sampling matrix is poorly conditioned "
-                 f"(kappa = {matrix.cond:.2g})", ConditioningWarning)
-
-        return matrix.lstsq(ax, axis)
-
-    @property
-    def cond(self):
-        """Condition number of the fitting problem"""
-        return self.matrix.cond
+        self.basis = basis
 
-    @property
-    def sampling_points(self):
-        """Set of sampling points"""
-        raise NotImplementedError()
+        if sampling_points is None:
+            self.sampling_points = basis.default_tau_sampling_points()
+            if use_positive_taus:
+                self.sampling_points = np.mod(self.sampling_points, basis.beta)
+        else:
+            self.sampling_points = np.asarray(sampling_points, dtype=np.float64)
 
-    @property
-    def matrix(self):
-        """Evaluation matrix is decomposed form"""
-        raise NotImplementedError()
+        self.sampling_points = np.sort(self.sampling_points)
+        if isinstance(basis, augment.AugmentedBasis):
+            # Create sampling object
+            matrix = basis.u(self.sampling_points).T
+            self._ptr = tau_sampling_new_with_matrix(basis, basis.statistics, self.sampling_points, matrix)
+        else:
+            # Create sampling object
+            self._ptr = tau_sampling_new(basis._ptr, self.sampling_points)
 
     @property
-    def basis(self):
-        """Basis instance"""
-        raise NotImplementedError()
+    def tau(self):
+        """Tau sampling points."""
+        return self.sampling_points
 
-    def _for_sampling_points(self, x):
-        raise RuntimeError("Changing sampling points is not possible")
+    def evaluate(self, al, axis=0):
+        """
+        Transform basis coefficients to sampling points.
+
+        Parameters:
+        -----------
+        al : array_like
+            Basis coefficients
+        axis : int, optional
+            Axis along which to transform
+
+        Returns:
+        --------
+        ndarray
+            Values at sampling points
+        """
+        al = np.ascontiguousarray(al)
+        output_dims = list(al.shape)
+        ndim = len(output_dims)
+        input_dims = np.asarray(al.shape, dtype=np.int32)
+        output_dims[axis] = len(self.sampling_points)
+        if al.dtype.kind == "f":
+            output = np.zeros(output_dims, dtype=np.float64)
+
+            status = _lib.spir_sampling_eval_dd(
+                self._ptr,
+                SPIR_ORDER_ROW_MAJOR,
+                ndim,
+                input_dims.ctypes.data_as(POINTER(c_int)),
+                axis,
+                al.ctypes.data_as(POINTER(c_double)),
+                output.ctypes.data_as(POINTER(c_double))
+            )
+        elif al.dtype.kind == "c":
+            output = np.zeros(output_dims, dtype=c_double_complex)
+
+            status = _lib.spir_sampling_eval_zz(
+                self._ptr,
+                SPIR_ORDER_ROW_MAJOR,
+                ndim,
+                input_dims.ctypes.data_as(POINTER(c_int)),
+                axis,
+                al.ctypes.data_as(POINTER(c_double_complex)),
+                output.ctypes.data_as(POINTER(c_double_complex))
+            )
+            output = output['real'] + 1j * output['imag']
+        else:
+            raise ValueError(f"Unsupported dtype: {al.dtype}")
 
+        if status != COMPUTATION_SUCCESS:
+            raise RuntimeError(f"Failed to evaluate sampling: {status}")
 
-class TauSampling(AbstractSampling):
-    """Sparse sampling in imaginary time.
+        return output
 
-    Allows the transformation between the IR basis and a set of sampling points
-    in (scaled/unscaled) imaginary time.
-    """
-    def __init__(self, basis, sampling_points=None):
-        if sampling_points is None:
-            sampling_points = basis.default_tau_sampling_points()
+    def fit(self, ax, axis=0):
+        """
+        Fit basis coefficients from sampling point values.
+        """
+        ax = np.ascontiguousarray(ax)
+        ndim = len(ax.shape)
+        input_dims = np.asarray(ax.shape, dtype=np.int32)
+        output_dims = list(ax.shape)
+        output_dims[axis] = self.basis.size
+        if ax.dtype.kind == "f":
+            output = np.zeros(output_dims, dtype=np.float64)
+            status = _lib.spir_sampling_fit_dd(
+                self._ptr,
+                SPIR_ORDER_ROW_MAJOR,
+            ndim,
+            input_dims.ctypes.data_as(POINTER(c_int)),
+            axis,
+            ax.ctypes.data_as(POINTER(c_double)),
+            output.ctypes.data_as(POINTER(c_double))
+        )
+        elif ax.dtype.kind == "c":
+            output = np.zeros(output_dims, dtype=c_double_complex)
+            status = _lib.spir_sampling_fit_zz(
+                self._ptr,
+                SPIR_ORDER_ROW_MAJOR,
+                ndim,
+                input_dims.ctypes.data_as(POINTER(c_int)),
+                axis,
+                ax.ctypes.data_as(POINTER(c_double_complex)),
+                output.ctypes.data_as(POINTER(c_double_complex))
+            )
+            output = output['real'] + 1j * output['imag']
         else:
-            sampling_points = np.asarray(sampling_points)
-            if sampling_points.ndim != 1:
-                raise ValueError("sampling points must be vector")
+            raise ValueError(f"Unsupported dtype: {ax.dtype}")
 
-        matrix = basis.u(sampling_points).T
-        self._basis = basis
-        self._sampling_points = sampling_points
-        self._matrix = DecomposedMatrix(matrix)
+        if status != COMPUTATION_SUCCESS:
+            raise RuntimeError(f"Failed to fit sampling: {status}")
 
-    @property
-    def basis(self): return self._basis
-
-    @property
-    def sampling_points(self): return self._sampling_points
+        return output
 
     @property
-    def matrix(self): return self._matrix
-
-    @property
-    def tau(self):
-        """Sampling points in (reduced) imaginary time"""
-        return self._sampling_points
+    def cond(self):
+        """Condition number of the sampling matrix."""
+        cond = c_double()
+        status = _lib.spir_sampling_get_cond_num(self._ptr, byref(cond))
+        if status != COMPUTATION_SUCCESS:
+            raise RuntimeError(f"Failed to get condition number: {status}")
+        return cond.value
 
-    def _for_sampling_points(self, x):
-        x = np.asarray(x)
-        return TauSampling(self._basis, x)
+    def __repr__(self):
+        return f"TauSampling(n_points={len(self.sampling_points)})"
 
 
-class MatsubaraSampling(AbstractSampling):
+class MatsubaraSampling:
     """Sparse sampling in Matsubara frequencies.
 
     Allows the transformation between the IR basis and a set of sampling points
@@ -166,256 +177,134 @@ class MatsubaraSampling(AbstractSampling):
     case, sparse sampling is performed over non-negative frequencies only,
     cutting away half of the necessary sampling space.
     """
-    def __init__(self, basis, sampling_points=None, *, positive_only=False):
+
+    def __init__(self, basis, sampling_points=None, positive_only=False):
+        """
+        Initialize Matsubara sampling.
+
+        Parameters:
+        -----------
+        basis : FiniteTempBasis
+            Finite temperature basis
+        sampling_points : array_like, optional
+            Matsubara frequency indices. If None, use default.
+        positive_only : bool, optional
+            If True, use only positive frequencies
+        """
+        self.basis = basis
+        self.positive_only = positive_only
+
         if sampling_points is None:
-            sampling_points = basis.default_matsubara_sampling_points(
-                                                positive_only=positive_only)
+            self.sampling_points = basis.default_matsubara_sampling_points(positive_only=positive_only)
         else:
-            sampling_points = _util.check_reduced_matsubara(sampling_points)
-            if sampling_points.ndim != 1:
-                raise ValueError("sampling points must be vector")
-            sampling_points = np.sort(sampling_points)
-            if positive_only and not sampling_points[0] >= 0:
-                raise ValueError("invalid negative sampling frequencies")
-
-        matrix = basis.uhat(sampling_points).T
-        self._basis = basis
-        self._sampling_points = sampling_points
-        self._positive_only = positive_only
-
-        if positive_only:
-            ssvd_result = _split_complex(matrix, sampling_points[0] == 0)
-            self._matrix = SplitDecomposedMatrix(matrix, ssvd_result)
+            self.sampling_points = np.asarray(sampling_points, dtype=np.int64)
+
+        if isinstance(basis, augment.AugmentedBasis):
+            # Create sampling object
+            matrix = basis.uhat(self.sampling_points).T
+            matrix = np.ascontiguousarray(matrix, dtype=np.complex128)
+
+            self._ptr = matsubara_sampling_new_with_matrix(
+                basis.statistics,
+                basis.size,
+                positive_only,
+                self.sampling_points,
+                matrix
+            )
         else:
-            self._matrix = DecomposedMatrix(matrix)
-
-    @property
-    def basis(self): return self._basis
-
-    @property
-    def sampling_points(self): return self._sampling_points
-
-    @property
-    def matrix(self): return self._matrix
-
-    @property
-    def positive_only(self):
-        """Sampling is performed only on non-negative sampling frequencies"""
-        return self._positive_only
+            # Create sampling object
+            self._ptr = matsubara_sampling_new(basis._ptr, positive_only, self.sampling_points)
 
     @property
     def wn(self):
-        """Sampling points as (reduced) Matsubara frequencies"""
-        return self._sampling_points
-
-    def _for_sampling_points(self, x):
-        x = np.asarray(x)
-        return MatsubaraSampling(self._basis, x,
-                                 positive_only=self._positive_only)
-
+        """Matsubara frequency indices."""
+        return self.sampling_points
 
-class DecomposedMatrix:
-    """Matrix in SVD decomposed form for fast and accurate fitting.
-
-    Stores a matrix ``A`` together with its thin SVD form::
-
-        A == (u * s) @ vH.
-
-    This allows for fast and accurate least squares fits using ``A.lstsq(x)``.
-    """
-    def __init__(self, a, svd_result=None):
-        a = np.asarray(a)
-        if a.ndim != 2:
-            raise ValueError("a must be of matrix form")
-        if svd_result is None:
-            u, s, vH = np.linalg.svd(a, full_matrices=False)
+    def evaluate(self, al, axis=0):
+        """
+        Transform basis coefficients to sampling points.
+
+        Parameters:
+        -----------
+        al : array_like
+            Basis coefficients
+        axis : int, optional
+            Axis along which to transform
+
+        Returns:
+        --------
+        ndarray
+            Values at Matsubara frequencies (complex)
+        """
+        # For better numerical stability, we need to make the input array contiguous.
+        al = np.ascontiguousarray(al)
+        output_dims = list(al.shape)
+        ndim = len(output_dims)
+        input_dims = np.asarray(al.shape, dtype=np.int32)
+        output_dims[axis] = len(self.sampling_points)
+        output_cdouble_complex = np.zeros(output_dims, dtype=c_double_complex)
+        if al.dtype.kind == "f":
+            status = _lib.spir_sampling_eval_dz(
+                self._ptr,
+                SPIR_ORDER_ROW_MAJOR,
+                ndim,
+                input_dims.ctypes.data_as(POINTER(c_int)),
+                axis,
+                al.ctypes.data_as(POINTER(c_double)),
+                output_cdouble_complex.ctypes.data_as(POINTER(c_double_complex))
+            )
+            output = output_cdouble_complex['real'] + 1j * output_cdouble_complex['imag']
+        elif al.dtype.kind == "c":
+            status = _lib.spir_sampling_eval_zz(
+                self._ptr,
+                SPIR_ORDER_ROW_MAJOR,
+                ndim,
+                input_dims.ctypes.data_as(POINTER(c_int)),
+                axis,
+                al.ctypes.data_as(POINTER(c_double_complex)),
+                output_cdouble_complex.ctypes.data_as(POINTER(c_double_complex))
+            )
+            output = output_cdouble_complex['real'] + 1j * output_cdouble_complex['imag']
         else:
-            u, s, vH = _util.check_svd_result(svd_result, a.shape)
-
-        # Remove singular values which are exactly zero
-        where = s.astype(bool)
-        if not where.all():
-            u, s, vH = u[:, where], s[where], vH[where]
-
-        self._a = a
-        self._uH = np.array(u.conj().T)
-        self._s = s
-        self._v = np.array(vH.conj().T)
-
-    def __matmul__(self, x):
-        """Matrix-matrix multiplication."""
-        return self._a @ x
-
-    def matmul(self, x, axis=None):
-        """Compute ``A @ x`` (optionally along specified axis of x)"""
-        return _matop_along_axis(self._a.__matmul__, x, axis)
-
-    def _lstsq(self, x):
-        r = self._uH @ x
-        r = r / (self._s[:, None] if r.ndim > 1 else self._s)
-        return self._v @ r
-
-    def lstsq(self, x, axis=None):
-        """Return ``y`` such that ``np.linalg.norm(A @ y - x)`` is minimal"""
-        return _matop_along_axis(self._lstsq, x, axis)
-
-    def __array__(self, dtype=""):
-        """Convert to numpy array."""
-        return self._a if dtype == "" else self._a.astype(dtype)
-
-    @property
-    def a(self):
-        """Full matrix"""
-        return self._a
-
-    @property
-    def u(self):
-        """Left singular vectors, aranged column-wise"""
-        return self._uH.conj().T
-
-    @property
-    def s(self):
-        """Most significant, nonzero singular values"""
-        return self._s
-
-    @property
-    def vH(self):
-        """Right singular vectors, transposed"""
-        return self._v.conj().T
-
-    @property
-    def cond(self):
-        """Condition number of matrix"""
-        return self._s[0] / self._s[-1]
-
-
-class SplitDecomposedMatrix:
-    """Matrix in "split" SVD decomposed form for fast and accurate fitting.
+            raise ValueError(f"Unsupported dtype: {al.dtype}")
 
-    Stores a matrix ``A`` together with its "split SVD" form::
+        if status != COMPUTATION_SUCCESS:
+            raise RuntimeError(f"Failed to evaluate sampling: {status}")
 
-        A == u * s @ vT
+        return output
 
-    where `vT` is a real matrix and `u` is a complex matrix.  The "split" SVD
-    form differs from the SVD in that the least squares fit has to be
-    constructed as follows:
-
-        fit(A, x) == vT.T / s @ (u.conj().T * x).real
-
-    This again allows for fast and accurate least squares fits using
-    ``A.lstsq(x)``.  This is useful in the case where.
-    """
-    def __init__(self, a, ssvd_result):
-        a = np.asarray(a)
-        if a.ndim != 2:
-            raise ValueError("a must be of matrix form")
-
-        u, s, vT = _util.check_svd_result(ssvd_result, a.shape)
-        if np.issubdtype(vT.dtype, np.complexfloating):
-            raise ValueError("Split SVD part vT shall be a real matrix")
-        if not np.issubdtype(u.dtype, np.complexfloating):
-            raise ValueError("Split SVD part u shall be a complex matrix")
-
-        # Remove singular values which are exactly zero
-        where = s.astype(bool)
-        if not where.all():
-            u, s, vT = u[:, where], s[where], vT[where]
-
-        self._a = a
-        self._urealT = np.array(u.T.real)
-        self._uimagT = np.array(u.T.imag)
-        self._s = s
-        self._v = np.array(vT.T)
-
-    def __matmul__(self, x):
-        """Matrix-matrix multiplication."""
-        x = np.asarray(x)
-        if np.issubdtype(x.dtype, np.complexfloating):
-            warn("Expecting array of real numbers in expansion", UserWarning, 2)
-        return self._a @ x
-
-    def matmul(self, x, axis=None):
-        """Compute ``A @ x`` (optionally along specified axis of x)"""
-        return _matop_along_axis(self._a.__matmul__, x, axis)
-
-    def _lstsq(self, x):
-        r = self._urealT @ x.real
-        r += self._uimagT @ x.imag
-        r = r / (self._s[:, None] if r.ndim > 1 else self._s)
-        return self._v @ r
-
-    def lstsq(self, x, axis=None):
-        """Return ``y`` such that ``np.linalg.norm(A @ y - x)`` is minimal"""
-        return _matop_along_axis(self._lstsq, x, axis)
-
-    def __array__(self, dtype=""):
-        """Convert to numpy array."""
-        return self._a if dtype == "" else self._a.astype(dtype)
-
-    @property
-    def a(self):
-        """Full matrix"""
-        return self._a
-
-    @property
-    def u(self):
-        """Split left singular vectors, aranged column-wise"""
-        return (self._urealT + 1j * self._uimagT).T
-
-    @property
-    def s(self):
-        """Most significant, nonzero singular values"""
-        return self._s
-
-    @property
-    def vH(self):
-        """Right singular vectors, transposed"""
-        return self._v.conj().T
+    def fit(self, ax, axis=0):
+        """
+        Fit basis coefficients from Matsubara frequency values.
+        """
+        ax = np.ascontiguousarray(ax)
+        ndim = len(ax.shape)
+        input_dims = np.asarray(ax.shape, dtype=np.int32)
+        output_dims = list(ax.shape)
+        output_dims[axis] = self.basis.size
+        output = np.zeros(output_dims, dtype=c_double_complex)
+
+        status = _lib.spir_sampling_fit_zz(
+            self._ptr,
+            SPIR_ORDER_ROW_MAJOR,
+            ndim,
+            input_dims.ctypes.data_as(POINTER(c_int)),
+            axis,
+            ax.ctypes.data_as(POINTER(c_double_complex)),
+            output.ctypes.data_as(POINTER(c_double_complex))
+        )
+        if status != COMPUTATION_SUCCESS:
+            raise RuntimeError(f"Failed to fit sampling: {status}")
+        return output['real']
 
     @property
     def cond(self):
-        """Condition number of matrix"""
-        return self._s[0] / self._s[-1]
-
-
-class ConditioningWarning(RuntimeWarning):
-    """Warns about a poorly conditioned problem.
-
-    This warning is issued if the library detects a poorly conditioned fitting
-    problem.  This essentially means there is a high degree of ambiguity in how
-    to choose the solution.  One must therefore expect to lose significant
-    precision in the parameter values.
-    """
-    pass
-
-
-def _matop_along_axis(op, x, axis=None):
-    if axis is None:
-        return op(x)
-
-    x = np.asarray(x)
-    target_axis = max(x.ndim - 2, 0)
-    x = np.moveaxis(x, axis, target_axis)
-    r = op(x)
-    return np.moveaxis(r, target_axis, axis)
-
-
-def _split_complex(mat, has_zero=False, svd_algo=np.linalg.svd):
-    mat = np.asarray(mat)
-    n, _l = mat.shape
-    if not np.issubdtype(mat.dtype, np.complexfloating):
-        raise ValueError("mat must be complex matrix")
-
-    # split real and imaginary part into separate matrices
-    offset_imag = 1 if has_zero else 0
-    rmat = np.vstack((mat.real, mat[offset_imag:].imag))
-
-    # perform real-valued SVD
-    ur, s, vT = svd_algo(rmat, full_matrices=False)
-
-    # undo the split of the resulting vT vector
-    u = np.empty((n, s.size), mat.dtype)
-    u.real = ur[:n]
-    u[:offset_imag].imag = 0
-    u[offset_imag:].imag = ur[n:]
-    return u, s, vT
+        """Condition number of the sampling matrix."""
+        cond = c_double()
+        status = _lib.spir_sampling_get_cond_num(self._ptr, byref(cond))
+        if status != COMPUTATION_SUCCESS:
+            raise RuntimeError(f"Failed to get condition number: {status}")
+        return cond.value
+
+    def __repr__(self):
+        return f"MatsubaraSampling(n_points={len(self.sampling_points)}, positive_only={self.positive_only})"