From 84a83dc8fba922f7196e9c9257e60314ac380807 Mon Sep 17 00:00:00 2001 From: Hiroshi Shinaoka Date: Mon, 29 Dec 2025 21:21:29 +0900 Subject: [PATCH] feat: add work_dtype support and fix MatsubaraSampling.fit() bug - Add work_dtype parameter to sve.compute() and SVEResult to control working precision during SVE/SVD computation - Add backward-compatible API to FiniteTempBasis.__init__() to support kernel= and max_size= keyword arguments - Fix bug in MatsubaraSampling.fit() that only returned real part, causing loss of imaginary components in fitted coefficients --- src/sparse_ir/basis.py | 88 ++++++++++++++++++++++++++++++++------- src/sparse_ir/sampling.py | 2 +- src/sparse_ir/sve.py | 80 +++++++++++++++++++++++++++++++---- 3 files changed, 146 insertions(+), 24 deletions(-) diff --git a/src/sparse_ir/basis.py b/src/sparse_ir/basis.py index 4b6ce43..bf6143f 100644 --- a/src/sparse_ir/basis.py +++ b/src/sparse_ir/basis.py @@ -5,12 +5,28 @@ """ from typing import Optional import numpy as np -from pylibsparseir.core import basis_new, basis_get_size, basis_get_svals, basis_get_u, basis_get_v, basis_get_uhat, basis_get_default_tau_sampling_points, basis_get_default_omega_sampling_points, basis_get_default_matsubara_sampling_points -from pylibsparseir.constants import SPIR_STATISTICS_FERMIONIC, SPIR_STATISTICS_BOSONIC +from pylibsparseir.core import ( + basis_new, + basis_get_svals, + basis_get_u, + basis_get_v, + basis_get_uhat, + basis_get_default_tau_sampling_points, + basis_get_default_matsubara_sampling_points, +) +from pylibsparseir.constants import ( + SPIR_STATISTICS_FERMIONIC, + SPIR_STATISTICS_BOSONIC, +) from .kernel import LogisticKernel -from .abstract import AbstractBasis +from .abstract import AbstractBasis, AbstractKernel from .sve import SVEResult -from .poly import PiecewiseLegendrePolyVector, PiecewiseLegendrePolyFTVector, FunctionSet, FunctionSetFT +from .poly import ( + PiecewiseLegendrePolyVector, + PiecewiseLegendrePolyFTVector, + FunctionSet, + FunctionSetFT, +) class FiniteTempBasis(AbstractBasis): r"""Intermediate representation (IR) basis for given temperature. @@ -43,7 +59,17 @@ class FiniteTempBasis(AbstractBasis): giw = gl @ basis.uhat([1, 3, 5, 7]) """ - def __init__(self, statistics: str, beta: float, wmax: float, eps: float = np.finfo(np.float64).eps, sve_result: Optional[SVEResult] = None, max_size: int =-1): + def __init__( + self, + statistics: str, + beta: float, + wmax: float, + eps: Optional[float] = None, + *, + max_size: Optional[int] = None, + kernel: Optional[AbstractKernel] = None, + sve_result: Optional[SVEResult] = None, + ): """ Initialize finite temperature basis. @@ -55,31 +81,65 @@ def __init__(self, statistics: str, beta: float, wmax: float, eps: float = np.fi Inverse temperature wmax : float Frequency cutoff - eps : float - Relative truncation threshold for the singular values, - defaulting to the machine epsilon (2.2e-16) + eps : float, optional + Relative truncation threshold for the singular values. + Defaults to machine epsilon (~2.2e-16). + max_size : int, optional + Maximum basis size. If given, only at most the ``max_size`` most + significant singular values and associated basis functions are + retained. + kernel : AbstractKernel, optional + Kernel to use for the basis. If not given, a LogisticKernel is + created from beta * wmax. (Deprecated: kernel is inferred from + statistics.) + sve_result : SVEResult, optional + Precomputed SVE result. If not given, the SVE is computed. """ + if not (beta > 0): + raise ValueError("inverse temperature beta must be positive") + if not (wmax >= 0): + raise ValueError("frequency cutoff must be non-negative") + self._statistics = statistics self._beta = beta self._wmax = wmax self._lambda = beta * wmax + + # Handle eps default + if eps is None: + eps = np.finfo(np.float64).eps self._eps = eps - # Create kernel - if statistics == 'F' or statistics == 'B': + # Handle max_size + if max_size is None: + max_size = -1 + + # Create or use provided kernel + if kernel is not None: + # Backward compatibility: use provided kernel + self._kernel = kernel + elif statistics in ('F', 'B'): self._kernel = LogisticKernel(self._lambda) else: - raise ValueError(f"Invalid statistics: {statistics} expected 'F' or 'B'") + raise ValueError( + f"Invalid statistics: {statistics}, expected 'F' or 'B'" + ) - # Compute SVE + # Compute SVE if not provided if sve_result is None: self._sve = SVEResult(self._kernel, eps) else: self._sve = sve_result # Create basis - stats_int = SPIR_STATISTICS_FERMIONIC if statistics == 'F' else SPIR_STATISTICS_BOSONIC - self._ptr = basis_new(stats_int, self._beta, self._wmax, self._eps, self._kernel._ptr, self._sve._ptr, max_size) + stats_int = ( + SPIR_STATISTICS_FERMIONIC if statistics == 'F' + else SPIR_STATISTICS_BOSONIC + ) + self._ptr = basis_new( + stats_int, self._beta, self._wmax, self._eps, + self._kernel._ptr, self._sve._ptr, max_size + ) u_funcs = FunctionSet(basis_get_u(self._ptr)) v_funcs = FunctionSet(basis_get_v(self._ptr)) diff --git a/src/sparse_ir/sampling.py b/src/sparse_ir/sampling.py index b4e0a3f..9443ebf 100644 --- a/src/sparse_ir/sampling.py +++ b/src/sparse_ir/sampling.py @@ -332,7 +332,7 @@ def fit(self, ax, axis=0): ) if status != COMPUTATION_SUCCESS: raise RuntimeError(f"Failed to fit sampling: {status}") - return output['real'] + return output['real'] + 1j * output['imag'] @property def cond(self): diff --git a/src/sparse_ir/sve.py b/src/sparse_ir/sve.py index 7e79019..759058f 100644 --- a/src/sparse_ir/sve.py +++ b/src/sparse_ir/sve.py @@ -1,4 +1,5 @@ -# Copyright (C) 2020-2025 Satoshi Terasaki, Markus Wallerberger, Hiroshi Shinaoka, and others +# Copyright (C) 2020-2025 Satoshi Terasaki, Markus Wallerberger, +# Hiroshi Shinaoka, and others # SPDX-License-Identifier: MIT """ SVE (Singular Value Expansion) functionality for SparseIR. @@ -7,10 +8,40 @@ """ import numpy as np -from pylibsparseir.core import _lib, sve_result_new, sve_result_get_svals, sve_result_get_size +from pylibsparseir.constants import ( + SPIR_TWORK_FLOAT64, + SPIR_TWORK_FLOAT64X2, +) +from pylibsparseir.core import ( + _lib, + sve_result_new, + sve_result_get_svals, + sve_result_get_size, +) from .abstract import AbstractKernel from .kernel import LogisticKernel, RegularizedBoseKernel + +def _resolve_work_dtype(work_dtype): + if work_dtype is None: + return None + + if isinstance(work_dtype, str): + work_dtype = work_dtype.lower() + if work_dtype in {"float64", "double"}: + return SPIR_TWORK_FLOAT64 + if work_dtype in {"float64x2", "ddouble"}: + return SPIR_TWORK_FLOAT64X2 + raise TypeError(f"unexpected work_dtype string {work_dtype!r}") + + dtype = np.dtype(work_dtype) + if dtype == np.float64: + return SPIR_TWORK_FLOAT64 + if dtype.itemsize > np.dtype(np.float64).itemsize: + return SPIR_TWORK_FLOAT64X2 + raise TypeError(f"unsupported work_dtype {dtype}") + + class SVEResult: """ Result of a singular value expansion (SVE). @@ -19,7 +50,14 @@ class SVEResult: the SVE of an integral kernel. """ - def __init__(self, kernel: AbstractKernel, eps: float, cutoff: float=-1, n_sv: int=-1): + def __init__( + self, + kernel: AbstractKernel, + eps: float, + cutoff: float = -1, + n_sv: int = -1, + work_dtype=None, + ): """ Compute SVE of the given kernel. @@ -37,14 +75,23 @@ def __init__(self, kernel: AbstractKernel, eps: float, cutoff: float=-1, n_sv: i returned. """ if not isinstance(kernel, (LogisticKernel, RegularizedBoseKernel)): - raise TypeError("kernel must be LogisticKernel or RegularizedBoseKernel") + raise TypeError( + "kernel must be LogisticKernel or RegularizedBoseKernel" + ) self._kernel = kernel # Store kernel for later use self._eps = eps self._cutoff = cutoff self._n_sv = n_sv - self._ptr = sve_result_new(kernel._ptr, eps, cutoff=cutoff, lmax=n_sv) + twork = _resolve_work_dtype(work_dtype) + self._ptr = sve_result_new( + kernel._ptr, + eps, + cutoff=cutoff, + lmax=n_sv, + Twork=twork, + ) def __len__(self): return sve_result_get_size(self._ptr) @@ -59,7 +106,12 @@ def __del__(self): _lib.spir_sve_result_release(self._ptr) -def compute(kernel, eps=np.finfo(np.float64).eps, n_sv=-1): +def compute( + kernel, + eps=np.finfo(np.float64).eps, + n_sv=-1, + work_dtype=None, +): """Perform truncated singular value expansion of a kernel. Perform a truncated singular value expansion (SVE) of an integral @@ -85,8 +137,12 @@ def compute(kernel, eps=np.finfo(np.float64).eps, n_sv=-1): n_sv (int): Maximum basis size. If given, only at most the ``n_sv`` most significant singular values and associated singular functions are + returned. Defaults to -1, which means all singular values are returned. - Defaulting to -1, which means all singular values are returned. + work_dtype (dtype-like or str, optional): + Working data type used during the SVE / SVD computation. Accepts + ``numpy.float64``, ``float``, or strings such as ``"float64"`` or + ``"float64x2"``. Defaults to ``float64x2`` for maximal precision. Returns: An ``SVEResult`` containing the truncated singular value expansion. @@ -94,8 +150,14 @@ def compute(kernel, eps=np.finfo(np.float64).eps, n_sv=-1): if eps is None: eps = np.finfo(np.float64).eps - return SVEResult(kernel, eps=eps, cutoff=-1, n_sv=n_sv) + return SVEResult( + kernel, + eps=eps, + cutoff=-1, + n_sv=n_sv, + work_dtype=work_dtype, + ) # Backward compatibility -compute_sve = compute \ No newline at end of file +compute_sve = compute